Skip to main content

memra_engine/
lib.rs

1//! memra engine: Stage-1 correctness-first forward-pass kernels + ops, on sm_120 via cudarc.
2
3use cudarc::driver::{
4    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
5};
6use cudarc::nvrtc::Ptx;
7use std::sync::{Arc, Mutex};
8
9#[cfg(debug_assertions)]
10pub(crate) fn debug_assert_tensor_stream_device<T>(
11    tensor: &CudaSlice<T>,
12    stream: &CudaStream,
13    site: &str,
14) {
15    let tensor_dev = tensor.ordinal();
16    let stream_dev = stream.context().ordinal();
17    assert_eq!(
18        tensor_dev, stream_dev,
19        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
20    );
21}
22
23pub use memra_gguf;
24pub use memra_runtime;
25
26pub mod forward;
27pub mod hybrid;
28pub mod hybrid_forward;
29pub mod model;
30pub mod sigrouter_contract;
31pub mod vision;
32pub mod vision_gemma;
33pub mod vision_pre;
34/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
35/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
36pub mod cache {
37    pub use memra_kv::*;
38}
39pub mod decode;
40pub mod decode_batch;
41pub mod dflash;
42pub mod eagle;
43pub mod gemma_spec;
44pub mod graph_update;
45/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
46/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
47/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
48pub mod mla;
49pub mod moesd;
50pub mod parallel;
51pub mod pp;
52pub mod round_stream;
53pub mod spec;
54pub use memra_sampling as sampler;
55
56/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
57/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
58/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
59/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
60/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
61///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
62///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
63///                     stream sync per projection (round-47 ledgered defect).
64///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
65///                     construction, zero syncs, f32 C with the act row-scale folded in.
66/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
67/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
68/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
69/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
70/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
71/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
72///
73/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
74/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
75/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
76/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
77/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
78/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
79/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
80/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
81///
82/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
83/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
84/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
85/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
86/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
87/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
88/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
89///
90/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
91/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
92/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
93/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
94/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
95/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
96/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
97/// the k-quant-only admission survives as the rollback seam, not the default.
98/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
99pub fn moe_f16g_mode() -> u8 {
100    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
101    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
102        Ok("0") => 0,
103        Ok("2") => 2,
104        Ok("3") => 3,
105        Ok(_) => 1,
106        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
107        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
108        Err(_) => 2,
109    })
110}
111/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
112/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
113/// (shape_sel, cross) for the FFI:
114///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
115///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
116///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
117///                         back to 32x64 in-launcher when the device/in_f can't take it).
118///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
119///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
120///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
121///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
122///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
123///                         verdict was stale).
124pub fn moe_f16g_sk_params() -> (i32, i32) {
125    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
126    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
127        Ok("0") => (-1, 0),
128        Ok("32") => (0, i32::MAX),
129        Ok("128") => (0, 1),
130        _ => {
131            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
132                .ok()
133                .and_then(|v| v.parse().ok())
134                .unwrap_or(64);
135            (0, cross)
136        }
137    })
138}
139/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
140/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
141/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
142/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
143/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
144/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
145/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
146/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
147/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
148/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
149pub fn moe_f16g_direct_on(qtype: i32) -> bool {
150    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
151    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
152        Ok("0") => 0,
153        Ok("kq") => 1,
154        _ => 2,
155    });
156    match m {
157        0 => false,
158        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
159        _ => true,
160    }
161}
162/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
163/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
164/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
165/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
166/// stage under q35's routing skew. Bit-identical to every other sk form by construction
167/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
168/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
169/// tail. in_f % 64 != 0 falls back in-launcher.
170pub fn moe_f16g_tail_on() -> bool {
171    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
173}
174
175/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
176/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
177/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
178/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
179/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
180/// still opens this door for A/B.
181pub fn moe_f16g_gemma_on() -> bool {
182    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
183    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
184}
185
186/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
187/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
188/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
189pub fn moe_fuse_actq_on() -> bool {
190    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
192}
193
194/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
195/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
196/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
197/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
198/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
199/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
200/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
201/// verify already use (dispatch parity, one router kernel for every t).
202/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
203pub fn router_prefill_exact_on() -> bool {
204    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
206}
207
208pub fn router_kernel_on() -> bool {
209    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210    *ON.get_or_init(|| {
211        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
212        if !on {
213            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
214        }
215        on
216    })
217}
218
219/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
220/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
221/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
222/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
223/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
224/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
225/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
226/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
227/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
228/// seam, perf-only: bits are equal by the kernel-check gate).
229/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
230/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
231/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
232/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
233pub const ROUTER_BATCH_MIN_T: usize = 8;
234pub fn router_batch_on() -> bool {
235    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
236    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
237}
238mod cpu_experts;
239#[cfg(memra_cutlass)]
240pub mod cutlass_ffi;
241pub mod f16_ffi;
242pub mod fp8_ffi;
243pub mod mmq_ffi;
244pub mod moe_cache;
245pub mod prime_graph;
246pub mod spill;
247mod spill_pread;
248
249// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
250// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
251// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
252// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
253// broke every machine that wasn't the build machine. Same bytes, same module image;
254// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
255const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
256const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
257const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
258const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
259const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
260const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
261/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
262const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
263
264/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
265/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
266/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
267/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
268/// compile-time default (zero behavior change).
269fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
270    assert!(
271        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
272        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
273    );
274    match std::env::var("MEMRA_GEMM_FATBIN") {
275        Ok(path) => std::borrow::Cow::Owned(
276            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
277        ),
278        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
279    }
280}
281
282/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
283/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
284/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
285/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
286/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
287/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
288pub(crate) const fn portable_mma_gated() -> bool {
289    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
290}
291
292/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
293/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
294/// in a pure helper so the dispatch guard can be regression-tested without constructing an
295/// Engine or allocating a GPU tensor.
296const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
297    (!portable_cuda || hopper_mma) && !no_gemm
298}
299
300// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
301// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
302// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
303// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
304// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
305// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
306// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
307const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
308const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
309const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
310const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
311const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
312
313/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
314/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
315pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
316
317/// The flash_attn fatbin matching the selected KV formats.
318fn flash_fatbin_bytes() -> &'static [u8] {
319    match kv_cache_formats() {
320        ("q8_0", "q5_1") => FLASH_FATBIN,
321        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
322        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
323        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
324        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
325        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
326        other => unreachable!("kv_cache_formats returned {other:?}"),
327    }
328}
329
330/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
331/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
332/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
333/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
334/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
335/// defaults (zero behavior change).
336fn k1_launch_override() -> Option<(u32, u32, u32)> {
337    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
338    *K1.get_or_init(|| {
339        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
340        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
341        match p.as_slice() {
342            [bm, bn, w] => Some((*bm, *bn, *w)),
343            _ => None,
344        }
345    })
346}
347
348/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
349/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
350/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
351/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
352/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
353/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
354pub(crate) fn wgmma_gemm_enabled() -> bool {
355    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
357}
358
359/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
360/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
361/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
362/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
363/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
364/// the split count changes the combine's FP summation order, and the spec verify's batched forward
365/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
366/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
367/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
368/// adaptive retries (any retry MUST pass run-spec self-consistency first).
369/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
370/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
371/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
372/// between eager decode and the verify (the spec-exactness law).
373pub const FA_VEC_MIN_TKV: usize = 96;
374/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
375/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
376/// which moves the crossover — sweep per model, adopt per the battery.
377pub fn fa_vec_min_tkv() -> usize {
378    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
379    *V.get_or_init(|| {
380        std::env::var("MEMRA_FA_VEC_MIN")
381            .ok()
382            .and_then(|v| v.parse().ok())
383            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
384    })
385}
386
387/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
388/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
389/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
390///
391/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
392/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
393/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
394/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
395/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
396/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
397pub fn fa_f16pv_on() -> bool {
398    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
399    *ON.get_or_init(|| {
400        std::env::var("MEMRA_FA_F16PV")
401            .map(|v| v != "0")
402            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
403    })
404}
405
406/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
407/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
408/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
409pub fn fa512_hp_on() -> bool {
410    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
411    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
412}
413
414/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
415/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
416/// accumulation. Even n_head and even GQA group required (guarded per call).
417pub fn faw_hp_on() -> bool {
418    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
419    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
420}
421
422/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
423/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
424/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
425pub fn fa512_wide_warps() -> usize {
426    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
427    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
428        Ok("1") => 4,
429        _ => 2,
430    })
431}
432
433/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
434/// and the gemma global-layer rows/parity call sites.
435pub fn fa512_min_tkv() -> usize {
436    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
437    *FA512_MIN.get_or_init(|| {
438        std::env::var("MEMRA_FA512_MIN")
439            .ok()
440            .and_then(|v| v.parse().ok())
441            .unwrap_or(512)
442    })
443}
444/// Per-model crossover default, set at model load BEFORE the first decode (per-model
445/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
446/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
447pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
448    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
449/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
450/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
451/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
452pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
453/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
454/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
455/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
456/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
457/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
458pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
459    std::sync::atomic::AtomicBool::new(false);
460/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
461/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
462/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
463/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
464/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
465/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
466pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
467    std::sync::atomic::AtomicBool::new(true);
468pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
469    std::sync::atomic::AtomicUsize::new(16);
470/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
471/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
472/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
473/// latency-bound at 256 threads — 7us/launch measured).
474pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
475/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
476pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
477/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
478/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
479/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
480/// explicit numerical-form seam. mmq_ffi reads this before the env.
481pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
482/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
483/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
484pub use memra_kv::KV_FP8_FORCE;
485pub(crate) fn rms_block() -> u32 {
486    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
487    *V.get_or_init(|| {
488        std::env::var("MEMRA_RMS_BLOCK")
489            .ok()
490            .and_then(|v| v.parse().ok())
491            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
492    })
493}
494
495pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
496    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
497    if let Some(forced) = *S.get_or_init(|| {
498        std::env::var("MEMRA_FA_SPLIT")
499            .ok()
500            .and_then(|v| v.parse().ok())
501            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
502    }) {
503        return forced;
504    }
505    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
506    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
507    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
508    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
509    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
510    //
511    // SM-AWARE SHORT-CTX RUNG (2026-07-06 g7e): the 32-key rung was tuned on the 82-SM 5090.
512    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
513    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on g7e (N=1 sweep + N=3
514    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
515    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
516    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
517    // rig-divergence law: this branch is measured on 188 SMs only).
518    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
519    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
520    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
521    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
522    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
523        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
524    {
525        return if t_kv <= 8192 {
526            16
527        } else if t_kv <= 16384 {
528            64
529        } else {
530            128
531        };
532    }
533    let big_rig = fa_sm_count() >= 128;
534    if big_rig {
535        let _ = n_head_kv;
536        if t_kv <= 2048 {
537            16
538        } else if t_kv <= 16384 {
539            64
540        } else {
541            128
542        }
543    } else if n_head_kv <= 4 {
544        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
545        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
546        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
547        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
548        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
549        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
550        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
551        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
552        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
553        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
554        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
555        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
556        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
557        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
558        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
559        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
560        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
561        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
562        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
563        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
564        if t_kv <= 512 {
565            8
566        } else if t_kv <= 16384 {
567            64
568        } else {
569            128
570        }
571    } else {
572        if t_kv <= 8192 {
573            32
574        } else if t_kv <= 16384 {
575            64
576        } else {
577            128
578        }
579    }
580}
581
582/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
583/// same attribute Engine::batched_variant reads).
584fn fa_sm_count() -> i32 {
585    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
586    *N.get_or_init(|| {
587        cudarc::driver::result::init().ok();
588        cudarc::driver::result::device::get(0)
589            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
590                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
591            .unwrap_or(82)
592    })
593}
594
595/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
596/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
597/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
598fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
599    match head_dim {
600        256 => Ok(""),
601        128 => Ok("_hd128"),
602        d => Err(format!(
603            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
604                          callers must gate to sdpa_naive"
605        )
606        .into()),
607    }
608}
609
610/// Quant type codes matching qmatvec.cu QType enum.
611pub const QT_Q8_0: i32 = 0;
612pub const QT_Q4_K: i32 = 1;
613pub const QT_Q6_K: i32 = 2;
614pub const QT_Q5_K: i32 = 3;
615pub const QT_Q3_K: i32 = 4;
616pub const QT_IQ4_XS: i32 = 5;
617pub const QT_IQ3_S: i32 = 6;
618pub const QT_NVFP4: i32 = 7;
619/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
620/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
621/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
622/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
623/// — ONE weight copy total, no Q8_0 re-encode duplicate.
624pub const QT_F8_E4M3: i32 = 10;
625/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
626/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
627pub const QT_NVFP4_RP: i32 = 9;
628/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
629pub const QT_F32: i32 = 8;
630pub const QT_BF16: i32 = 11;
631pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
632/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
633/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
634/// dp4a/MMQ implementation exists.
635pub const QT_Q2_K: i32 = 13;
636/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
637/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
638/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
639/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
640/// scalar `scale` field is 1.0 by the layout contract.
641///
642/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
643/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
644/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
645/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
646/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
647/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
648/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
649/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
650/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
651pub const QT_F8_E4M3_BLK: i32 = 14;
652
653/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
654pub struct Engine {
655    pub gpu: memra_runtime::Gpu,
656    module: Arc<CudaModule>,
657    hybrid: Arc<CudaModule>,
658    qmatvec: Arc<CudaModule>,
659    flash: Arc<CudaModule>,
660    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
661    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
662    /// Lazy: loaded on first global-format use; None until then.
663    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
664    gemm: Arc<CudaModule>,
665    router: Arc<CudaModule>,
666    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
667    sample: Arc<CudaModule>,
668    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
669    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
670    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
671    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
672    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
673    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
674    /// the single largest block. The cache still owns every address for its full lifetime.
675    moe_cache_layout: Mutex<Option<Vec<usize>>>,
676    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
677    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
678    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
679    /// verify between replays) reuse their addresses and the replay reads/writes live memory
680    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
681    capture_keep_on: std::sync::atomic::AtomicBool,
682    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
683    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
684    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
685    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
686    verify_exact: std::sync::atomic::AtomicBool,
687    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
688    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
689    pub copy_stream: Arc<CudaStream>,
690    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
691    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
692    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
693    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
694    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
695    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
696    #[cfg(memra_cutlass)]
697    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
698    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
699    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
700    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
701    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
702    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
703    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
704    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
705    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
706    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
707    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
708    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
709    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
710    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
711    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
712    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
713    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
714    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
715    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
716    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
717    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
718    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
719    /// before capture under the generate_graph tracking-off window so it carries no events).
720    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
721    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
722    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
723    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
724    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
725    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
726    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
727    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
728    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
729    router_stage: Mutex<Option<PinnedStage>>,
730}
731
732/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
733/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
734/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
735/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
736/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
737/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
738/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
739/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
740/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
741fn fa_v2_on() -> bool {
742    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
743    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
744    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
745    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
746    // + graph bit-identity green on all three models.
747    std::env::var("MEMRA_FA_V2")
748        .map(|v| v != "0")
749        .unwrap_or(true)
750}
751
752/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
753/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
754/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
755/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
756/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
757/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
758/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
759fn fa_v3_on() -> bool {
760    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
761    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
762    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
763    std::env::var("MEMRA_FA_V3")
764        .map(|v| v != "0")
765        .unwrap_or(true)
766}
767
768/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
769/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
770/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
771/// predicate so the twins can never diverge.
772fn fa_v4_mode() -> &'static str {
773    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
774    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
775}
776fn fa_v4_on() -> bool {
777    fa_v4_mode() != "0"
778} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
779/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
780/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
781/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
782/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
783/// stays kernel-family-identical to decode at the same t_kv.
784/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
785/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
786pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
787    std::sync::atomic::AtomicUsize::new(1024);
788pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
789    std::sync::atomic::AtomicUsize::new(usize::MAX);
790pub fn fa_v4_at_pub(t_kv: usize) -> bool {
791    fa_v4_at(t_kv)
792}
793fn fa_v4_at(t_kv: usize) -> bool {
794    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
795    let mx = *M.get_or_init(|| {
796        std::env::var("MEMRA_FA_V4_MAX")
797            .ok()
798            .and_then(|v| v.parse().ok())
799            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
800    });
801    fa_v4_on() && t_kv < mx
802}
803/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
804/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
805/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
806/// (same split partition, same softmax/accumulation order, same partials/combine) and only
807/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
808/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
809/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
810/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
811/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
812/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
813/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
814/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
815/// within one process (the v2/v3 pattern).
816pub const FA_DEEP_MIN_DEFAULT: usize = 0;
817fn fa_deep_at(t_kv: usize) -> bool {
818    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
819        return false;
820    }
821    let min = std::env::var("MEMRA_FA_DEEP_MIN")
822        .ok()
823        .and_then(|v| v.parse().ok())
824        .unwrap_or(FA_DEEP_MIN_DEFAULT);
825    t_kv >= min
826}
827/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
828pub fn fa_deep_at_pub(t_kv: usize) -> bool {
829    fa_deep_at(t_kv)
830}
831
832fn fa_v3_active(head_dim: usize) -> bool {
833    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
834    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
835    fa_v3_on()
836        && head_dim % 128 == 0
837        && kv_cache_formats() == ("q8_0", "q5_1")
838        && !Engine::kv_fp8_on()
839}
840
841/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
842/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
843/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
844/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
845/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
846/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
847/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
848pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
849    std::env::var("MEMRA_NO_FA_VEC").is_err()
850        && t_kv >= fa_vec_min_tkv()
851        && head_dim == 256
852        && fa_v4_at(t_kv)
853        && !matches!(fa_v4_mode(), "noB3" | "stage")
854        && !Engine::kv_fp8_on()
855}
856/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
857pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
858    fa_split_keys(t_kv, n_head_kv)
859}
860
861/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
862/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
863/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
864/// so we allocate through `result::malloc_host` with flags=0 directly.
865struct PinnedStage {
866    ptr: *mut u8,
867    cap: usize,
868}
869unsafe impl Send for PinnedStage {}
870impl PinnedStage {
871    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
872        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
873        Ok(PinnedStage { ptr, cap })
874    }
875}
876impl Drop for PinnedStage {
877    fn drop(&mut self) {
878        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
879    }
880}
881
882/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
883/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
884pub const ARGMAX_NB: usize = 256;
885
886/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
887pub(crate) use memra_fa3_vl as fa3_vl_raw;
888
889unsafe extern "C" {
890    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
891    fn memra_fa3_prefill(
892        q16: *const core::ffi::c_void,
893        k16: *const core::ffi::c_void,
894        v16: *const core::ffi::c_void,
895        o: *mut f32,
896        t: i32,
897        h: i32,
898        hkv: i32,
899        d: i32,
900        scale: f32,
901        stream: *mut core::ffi::c_void,
902    ) -> i32;
903    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
904    pub(crate) fn memra_fa3_vl(
905        q16s: *const *const core::ffi::c_void,
906        k16s: *const *const core::ffi::c_void,
907        v16s: *const *const core::ffi::c_void,
908        os: *const *mut f32,
909        ts: *const i32,
910        b: i32,
911        h: i32,
912        hkv: i32,
913        d: i32,
914        scale: f32,
915        stream: *mut core::ffi::c_void,
916    ) -> i32;
917}
918
919/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
920/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
921/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
922/// (slots are never re-allocated), so passing raw values is stable across the launch.
923#[repr(C)]
924#[derive(Clone, Copy)]
925pub struct WPtr8(pub [u64; 8]);
926unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
927
928/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
929/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
930/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
931/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
932#[repr(C)]
933#[derive(Clone, Copy, Default)]
934pub struct GdnSeqVl {
935    pub kb16: u64,
936    pub gcum: u64,
937    pub beta: u64,
938    pub u: u64,
939    pub wb16: u64,
940    pub y: u64,
941    pub ssnap: u64,
942    pub state_in: u64,
943    pub state_out: u64,
944    pub q: u64,
945    pub p: u64,
946    pub o: u64,
947    pub k: u64,
948    pub v: u64,
949    pub g: u64,
950    pub a: u64,
951    pub w: u64,
952    pub t: i32,
953    pub nc: i32,
954}
955unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
956#[repr(C)]
957#[derive(Clone, Copy)]
958pub struct GdnVl8(pub [GdnSeqVl; 8]);
959unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
960
961/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
962/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
963#[repr(C)]
964#[derive(Clone, Copy, Default)]
965pub struct GdnWVl {
966    pub qb16: u64,
967    pub pb16: u64,
968}
969unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
970#[repr(C)]
971#[derive(Clone, Copy)]
972pub struct GdnWVl8(pub [GdnWVl; 8]);
973unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
974
975/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
976#[repr(C)]
977#[derive(Clone, Copy, Default)]
978pub struct GdnPrepVl {
979    pub qkv: u64,
980    pub conv_state: u64,
981    pub conv_out: u64,
982    pub q_g: u64,
983    pub k_g: u64,
984    pub v_g: u64,
985    pub q_l2: u64,
986    pub k_l2: u64,
987    pub beta_raw: u64,
988    pub alpha: u64,
989    pub beta: u64,
990    pub g_log: u64,
991    pub o: u64,
992    pub z: u64,
993    pub gn: u64,
994    pub gn16: u64,
995    pub kb16: u64,
996    pub qb16: u64,
997    pub t: i32,
998    pub pad: i32,
999}
1000unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1001#[repr(C)]
1002#[derive(Clone, Copy)]
1003pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1004unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1005
1006/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1007#[repr(C)]
1008#[derive(Clone, Copy, Default)]
1009pub struct FaSeqVl {
1010    pub q: u64,
1011    pub k16: u64,
1012    pub v16: u64,
1013    pub o: u64,
1014    pub kf: u64,
1015    pub vf: u64,
1016    pub t: i32,
1017    pub pad: i32,
1018}
1019unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1020#[repr(C)]
1021#[derive(Clone, Copy)]
1022pub struct FaVl8(pub [FaSeqVl; 8]);
1023unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1024
1025/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1026#[repr(C)]
1027#[derive(Clone, Copy, Default)]
1028pub struct AttnPreVl {
1029    pub qf: u64,
1030    pub kf: u64,
1031    pub vf: u64,
1032    pub q: u64,
1033    pub gate: u64,
1034    pub qn: u64,
1035    pub kn: u64,
1036    pub kc: u64,
1037    pub vc: u64,
1038    pub t: i32,
1039    pub pad: i32,
1040}
1041unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1042#[repr(C)]
1043#[derive(Clone, Copy)]
1044pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1045unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1046
1047/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1048/// varlen K1-K5 chain fills them).
1049pub struct GdnChunkBufs {
1050    pub gcum: CudaSlice<f32>,
1051    pub a: CudaSlice<f32>,
1052    pub p: CudaSlice<f32>,
1053    pub u: CudaSlice<f32>,
1054    pub w: CudaSlice<f32>,
1055    pub kb16: CudaSlice<u8>,
1056    pub wb16: CudaSlice<u8>,
1057    pub y16: CudaSlice<u8>,
1058    pub ssnap16: CudaSlice<u8>,
1059    pub qb16: CudaSlice<u8>,
1060    pub pb16: CudaSlice<u8>,
1061    pub o: CudaSlice<f32>,
1062    pub t: usize,
1063    pub nc: usize,
1064}
1065
1066/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1067#[repr(C)]
1068#[derive(Clone, Copy)]
1069pub struct F32x8(pub [f32; 8]);
1070unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1071
1072/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1073/// process. Bench binaries read it right after the call to print gen-only throughput without the
1074/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1075pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1076
1077impl Engine {
1078    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1079        let gpu = memra_runtime::Gpu::new(ordinal)?;
1080        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1081        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1082        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1083        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1084            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1085            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1086                .and_then(|d| unsafe {
1087                    Ok((
1088                        cudarc::driver::result::device::get_attribute(
1089                            d,
1090                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1091                        )?,
1092                        cudarc::driver::result::device::get_attribute(
1093                            d,
1094                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1095                        )?,
1096                    ))
1097                })
1098                .unwrap_or((0, 0));
1099            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1100            let ok = matches!(
1101                (built, maj, min),
1102                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1103            );
1104            if !ok {
1105                return Err(format!(
1106                    "memra was built for sm_{built} but device {ordinal} reports compute \
1107                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1108                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1109                )
1110                .into());
1111            }
1112        }
1113        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1114        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1115        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1116        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1117        unsafe {
1118            use cudarc::driver::sys;
1119            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1120            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1121            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1122                let mut thresh: u64 = u64::MAX;
1123                let _ = sys::cuMemPoolSetAttribute(
1124                    pool,
1125                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1126                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1127                );
1128            }
1129        }
1130        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1131        let hybrid = gpu
1132            .ctx
1133            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1134        let qmatvec = gpu
1135            .ctx
1136            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1137        let flash = gpu
1138            .ctx
1139            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1140        let gemm = gpu
1141            .ctx
1142            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1143        let router = gpu
1144            .ctx
1145            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1146        let sample = gpu
1147            .ctx
1148            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1149        let copy_stream = gpu.ctx.new_stream()?;
1150        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1151        // cudarc is in multi-stream mode (main stream +
1152        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1153        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1154        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1155        // (~7 ms/tok host time, measured nsys 2026-07-04 g7e), and +4.6% measured on 27B decode —
1156        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1157        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1158        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1159        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1160        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1161        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1162        // implicit event tracking.
1163        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1164        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1165        if std::env::var("MEMRA_EVT")
1166            .map(|v| v == "1")
1167            .unwrap_or(false)
1168        {
1169            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1170        } else {
1171            unsafe {
1172                gpu.ctx.disable_event_tracking();
1173            }
1174        }
1175        Ok(Self {
1176            gpu,
1177            module,
1178            hybrid,
1179            qmatvec,
1180            flash,
1181            flash_g: std::sync::OnceLock::new(),
1182            gemm,
1183            router,
1184            sample,
1185            moe_cache: Mutex::new(None),
1186            moe_cache_layout: Mutex::new(None),
1187            copy_stream,
1188            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1189            verify_exact: std::sync::atomic::AtomicBool::new(false),
1190            capture_keep: Mutex::new(Vec::new()),
1191            argmax_partials: Mutex::new(None),
1192            prime_deqw_ws: Mutex::new(None),
1193            router_stage: Mutex::new(None),
1194            fp8_scratch: Mutex::new(None),
1195            fa_vf16_scratch: Mutex::new(None),
1196            fa_part_pool: Mutex::new(None),
1197            fa_part_retired: Mutex::new(Vec::new()),
1198            fn_cache: Mutex::new(Default::default()),
1199            f16_scratch: Mutex::new(None),
1200            #[cfg(memra_cutlass)]
1201            cutlass_scratch: Mutex::new(None),
1202        })
1203    }
1204
1205    pub fn ctx(&self) -> &Arc<CudaContext> {
1206        &self.gpu.ctx
1207    }
1208
1209    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1210    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1211    ///
1212    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1213    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1214    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1215    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1216    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1217    ///
1218    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1219    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1220    /// under-count headroom does not belong in a gate that queues real work, but the honest
1221    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1222    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1223    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1224    ///
1225    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1226    pub fn pool_cached_bytes(&self) -> usize {
1227        let (reserved, used) = self.pool_reserved_used();
1228        reserved.saturating_sub(used)
1229    }
1230
1231    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1232    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1233    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1234    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1235    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1236    /// (0, 0) if the pool cannot be queried.
1237    pub fn pool_reserved_used(&self) -> (usize, usize) {
1238        use cudarc::driver::sys;
1239        unsafe {
1240            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1241            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1242                != sys::CUresult::CUDA_SUCCESS
1243            {
1244                return (0, 0);
1245            }
1246            let (mut reserved, mut used) = (0u64, 0u64);
1247            if sys::cuMemPoolGetAttribute(
1248                pool,
1249                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1250                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1251            ) != sys::CUresult::CUDA_SUCCESS
1252            {
1253                return (0, 0);
1254            }
1255            if sys::cuMemPoolGetAttribute(
1256                pool,
1257                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1258                &mut used as *mut u64 as *mut core::ffi::c_void,
1259            ) != sys::CUresult::CUDA_SUCCESS
1260            {
1261                return (0, 0);
1262            }
1263            (reserved as usize, used as usize)
1264        }
1265    }
1266
1267    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1268    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1269    pub fn stream(&self) -> Arc<CudaStream> {
1270        self.gpu.stream()
1271    }
1272    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1273    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1274    pub fn gkv_on() -> bool {
1275        memra_kv::gkv_on()
1276    }
1277
1278    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1279    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1280    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1281    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1282    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1283    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1284    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1285    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1286    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1287    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1288    /// ON for both — no acceptance cost measured.
1289    pub fn wkv_on() -> bool {
1290        memra_kv::wkv_on()
1291    }
1292
1293    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1294    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1295    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1296    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1297    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1298    pub fn kv_fp8_on() -> bool {
1299        memra_kv::kv_fp8_on()
1300    }
1301
1302    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1303    /// when the fp8-globals arm is on; everything else from the default flash module.
1304    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1305        if head_dim == 512 && Self::gkv_on() {
1306            self.func_g(name)
1307        } else {
1308            self.func(name)
1309        }
1310    }
1311
1312    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1313    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1314    /// per-format fatbins; fall back to the base modules for those.
1315    fn func_g(&self, name: &str) -> CudaFunction {
1316        let m = self.flash_g.get_or_init(|| {
1317            self.gpu
1318                .ctx
1319                .load_module(cudarc::nvrtc::Ptx::from_binary(
1320                    FLASH_FATBIN_KF8VF8.to_vec(),
1321                ))
1322                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1323        });
1324        let key = format!("g:{name}");
1325        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1326            return f.clone();
1327        }
1328        let f = match m.load_function(name) {
1329            Ok(f) => f,
1330            Err(_) => self.func(name),
1331        };
1332        self.fn_cache.lock().unwrap().insert(key, f.clone());
1333        f
1334    }
1335
1336    fn func(&self, name: &str) -> CudaFunction {
1337        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1338        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1339        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1340            return f.clone();
1341        }
1342        let f = self
1343            .module
1344            .load_function(name)
1345            .or_else(|_| self.hybrid.load_function(name))
1346            .or_else(|_| self.qmatvec.load_function(name))
1347            .or_else(|_| self.flash.load_function(name))
1348            .or_else(|_| self.gemm.load_function(name))
1349            .or_else(|_| self.router.load_function(name))
1350            .or_else(|_| self.sample.load_function(name))
1351            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1352        self.fn_cache
1353            .lock()
1354            .unwrap()
1355            .insert(name.to_string(), f.clone());
1356        f
1357    }
1358
1359    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1360    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1361    pub fn scatter_trim_logits(
1362        &self,
1363        src: &CudaSlice<f32>,
1364        d2t: &CudaSlice<u32>,
1365        dst: &mut CudaSlice<f32>,
1366        d_vocab: usize,
1367        n_vocab: usize,
1368    ) -> Result<(), Box<dyn std::error::Error>> {
1369        let f1 = self.func("scatter_trim_logits_f32");
1370        let f2 = self.func("scatter_trim_logits_pass2_f32");
1371        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1372        let cfg1 = LaunchConfig {
1373            grid_dim: (256, 1, 1),
1374            block_dim: (256, 1, 1),
1375            shared_mem_bytes: 0,
1376        };
1377        let __s_b1 = self.gpu.stream();
1378        let mut b1 = __s_b1.launch_builder(&f1);
1379        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1380        unsafe {
1381            b1.launch(cfg1)?;
1382        }
1383        let cfg2 = LaunchConfig {
1384            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1385            block_dim: (256, 1, 1),
1386            shared_mem_bytes: 0,
1387        };
1388        let __s_b2 = self.gpu.stream();
1389        let mut b2 = __s_b2.launch_builder(&f2);
1390        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1391        unsafe {
1392            b2.launch(cfg2)?;
1393        }
1394        Ok(())
1395    }
1396
1397    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1398    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1399
1400    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1401    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1402    #[allow(clippy::too_many_arguments)]
1403    pub fn filter_stats(
1404        &self,
1405        x: &CudaSlice<f32>,
1406        row_stride: usize,
1407        rows: &CudaSlice<i32>,
1408        out_th: &mut CudaSlice<f32>,
1409        out_z: &mut CudaSlice<f32>,
1410        out_max: &mut CudaSlice<f32>,
1411        n: usize,
1412        nrow: usize,
1413        temp: f32,
1414        top_k: i32,
1415        top_p: f32,
1416        min_p: f32,
1417    ) -> Result<(), Box<dyn std::error::Error>> {
1418        let f = self.func("filter_stats_f32");
1419        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1420        let cfg = LaunchConfig {
1421            grid_dim: (nrow as u32, 1, 1),
1422            block_dim: (1024, 1, 1),
1423            shared_mem_bytes: 0,
1424        };
1425        let __s_b = self.gpu.stream();
1426        let mut b = __s_b.launch_builder(&f);
1427        b.arg(x)
1428            .arg(&rs)
1429            .arg(rows)
1430            .arg(&mut *out_th)
1431            .arg(&mut *out_z)
1432            .arg(&mut *out_max)
1433            .arg(&ni)
1434            .arg(&nr)
1435            .arg(&temp)
1436            .arg(&top_k)
1437            .arg(&top_p)
1438            .arg(&min_p);
1439        unsafe {
1440            b.launch(cfg)?;
1441        }
1442        Ok(())
1443    }
1444
1445    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1446    #[allow(clippy::too_many_arguments)]
1447    pub fn softmax_gather_filtered(
1448        &self,
1449        x: &CudaSlice<f32>,
1450        row_stride: usize,
1451        ids: &CudaSlice<u32>,
1452        rows: &CudaSlice<i32>,
1453        th: &CudaSlice<f32>,
1454        z: &CudaSlice<f32>,
1455        out: &mut CudaSlice<f32>,
1456        n: usize,
1457        npair: usize,
1458        temp: f32,
1459    ) -> Result<(), Box<dyn std::error::Error>> {
1460        let f = self.func("softmax_gather_filtered_f32");
1461        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1462        let cfg = LaunchConfig {
1463            grid_dim: (npair as u32, 1, 1),
1464            block_dim: (256, 1, 1),
1465            shared_mem_bytes: 0,
1466        };
1467        let __s_b = self.gpu.stream();
1468        let mut b = __s_b.launch_builder(&f);
1469        b.arg(x)
1470            .arg(&rs)
1471            .arg(ids)
1472            .arg(rows)
1473            .arg(th)
1474            .arg(z)
1475            .arg(&mut *out)
1476            .arg(&ni)
1477            .arg(&np)
1478            .arg(&temp);
1479        unsafe {
1480            b.launch(cfg)?;
1481        }
1482        Ok(())
1483    }
1484
1485    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1486    #[allow(clippy::too_many_arguments)]
1487    pub fn residual_sample_filtered(
1488        &self,
1489        p: &CudaSlice<f32>,
1490        q: Option<&CudaSlice<f32>>,
1491        n: usize,
1492        temp: f32,
1493        seed: u64,
1494        stream_pos: u32,
1495        p_stats: (f32, f32, f32),
1496        q_stats: (f32, f32, f32),
1497        out_tok: &mut CudaSlice<u32>,
1498    ) -> Result<(), Box<dyn std::error::Error>> {
1499        let f = self.func("residual_sample_filtered_f32");
1500        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1501        let has_q: i32 = q.is_some() as i32;
1502        let qbuf = q.unwrap_or(p);
1503        let (pm, pth, pz) = p_stats;
1504        let (qm, qth, qz) = q_stats;
1505        let cfg = LaunchConfig {
1506            grid_dim: (1, 1, 1),
1507            block_dim: (1024, 1, 1),
1508            shared_mem_bytes: 0,
1509        };
1510        let __s_b = self.gpu.stream();
1511        let mut b = __s_b.launch_builder(&f);
1512        b.arg(p)
1513            .arg(qbuf)
1514            .arg(&has_q)
1515            .arg(&ni)
1516            .arg(&temp)
1517            .arg(&slo)
1518            .arg(&shi)
1519            .arg(&stream_pos)
1520            .arg(&pm)
1521            .arg(&pth)
1522            .arg(&pz)
1523            .arg(&qm)
1524            .arg(&qth)
1525            .arg(&qz)
1526            .arg(&mut *out_tok);
1527        unsafe {
1528            b.launch(cfg)?;
1529        }
1530        Ok(())
1531    }
1532
1533    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1534    #[allow(clippy::too_many_arguments)]
1535    pub fn gumbel_perturb_filtered(
1536        &self,
1537        x: &CudaSlice<f32>,
1538        y: &mut CudaSlice<f32>,
1539        n: usize,
1540        seed: u64,
1541        stream_pos: u32,
1542        temp: f32,
1543        row_max: f32,
1544        th: f32,
1545    ) -> Result<(), Box<dyn std::error::Error>> {
1546        let f = self.func("gumbel_perturb_filtered_f32");
1547        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1548        let cfg = LaunchConfig {
1549            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1550            block_dim: (256, 1, 1),
1551            shared_mem_bytes: 0,
1552        };
1553        let __s_b = self.gpu.stream();
1554        let mut b = __s_b.launch_builder(&f);
1555        b.arg(x)
1556            .arg(&mut *y)
1557            .arg(&ni)
1558            .arg(&slo)
1559            .arg(&shi)
1560            .arg(&stream_pos)
1561            .arg(&temp)
1562            .arg(&row_max)
1563            .arg(&th);
1564        unsafe {
1565            b.launch(cfg)?;
1566        }
1567        Ok(())
1568    }
1569
1570    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1571    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1572    /// filtered rejection sampling exact for the penalized target.
1573    #[allow(clippy::too_many_arguments)]
1574    pub fn penalize_logits(
1575        &self,
1576        x: &mut CudaSlice<f32>,
1577        hist: &CudaSlice<u32>,
1578        n_hist: usize,
1579        rep: f32,
1580        freq: f32,
1581        present: f32,
1582        n: usize,
1583    ) -> Result<(), Box<dyn std::error::Error>> {
1584        if n_hist == 0 {
1585            return Ok(());
1586        }
1587        let f = self.func("penalize_logits_f32");
1588        let (nh, ni) = (n_hist as i32, n as i32);
1589        let cfg = LaunchConfig {
1590            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1591            block_dim: (128, 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(&mut *x)
1597            .arg(hist)
1598            .arg(&nh)
1599            .arg(&rep)
1600            .arg(&freq)
1601            .arg(&present)
1602            .arg(&ni);
1603        unsafe {
1604            b.launch(cfg)?;
1605        }
1606        Ok(())
1607    }
1608
1609    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1610    #[allow(clippy::too_many_arguments)]
1611    pub fn penalize_logits_rows(
1612        &self,
1613        x: &mut CudaSlice<f32>,
1614        hist: &CudaSlice<u32>,
1615        n_hist: usize,
1616        rep: f32,
1617        freq: f32,
1618        present: f32,
1619        n: usize,
1620        nrow: usize,
1621    ) -> Result<(), Box<dyn std::error::Error>> {
1622        if n_hist == 0 || nrow == 0 {
1623            return Ok(());
1624        }
1625        let f = self.func("penalize_logits_rows_f32");
1626        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1627        let cfg = LaunchConfig {
1628            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1629            block_dim: (128, 1, 1),
1630            shared_mem_bytes: 0,
1631        };
1632        let __s_b = self.gpu.stream();
1633        let mut b = __s_b.launch_builder(&f);
1634        b.arg(&mut *x)
1635            .arg(hist)
1636            .arg(&nh)
1637            .arg(&rep)
1638            .arg(&freq)
1639            .arg(&present)
1640            .arg(&ni)
1641            .arg(&nr);
1642        unsafe {
1643            b.launch(cfg)?;
1644        }
1645        Ok(())
1646    }
1647
1648    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1649    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1650    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1651    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1652    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1653    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1654    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1655    pub fn wpf_level() -> u32 {
1656        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1657        *ON.get_or_init(|| {
1658            std::env::var("MEMRA_WPF")
1659                .ok()
1660                .and_then(|v| v.parse().ok())
1661                .unwrap_or(1)
1662        })
1663    }
1664
1665    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1666    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1667    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1668    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1669    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1670    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1671    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1672    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1673    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1674    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1675    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1676    pub fn set_verify_exact(&self, on: bool) {
1677        self.verify_exact
1678            .store(on, std::sync::atomic::Ordering::Relaxed);
1679    }
1680    pub(crate) fn verify_exact_on(&self) -> bool {
1681        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1682    }
1683
1684    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1685    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1686    pub fn qkv_append_on() -> bool {
1687        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1688        *ON.get_or_init(|| {
1689            std::env::var("MEMRA_QKV_APPEND")
1690                .map(|v| v != "0")
1691                .unwrap_or(true)
1692        })
1693    }
1694
1695    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1696    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1697    pub fn pdl_wb_on() -> bool {
1698        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1699        *ON.get_or_init(|| {
1700            std::env::var("MEMRA_PDL_WB")
1701                .map(|v| v != "0")
1702                .unwrap_or(true)
1703        })
1704    }
1705
1706    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1707    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1708    /// per-model no-harm bisect knob.
1709    pub fn pdl_mmvq_on() -> bool {
1710        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1711        *ON.get_or_init(|| {
1712            std::env::var("MEMRA_PDL_MMVQ")
1713                .map(|v| v != "0")
1714                .unwrap_or(true)
1715        })
1716    }
1717
1718    pub fn pdl_on() -> bool {
1719        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1720        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1721    }
1722
1723    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1724    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1725    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1726    /// on the producer before any read), bit-identical by construction.
1727    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1728    pub fn pdl_nvfp4q8_on() -> bool {
1729        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1730        *ON.get_or_init(|| {
1731            std::env::var("MEMRA_PDL_NVFP4")
1732                .map(|v| v != "0")
1733                .unwrap_or(true)
1734        })
1735    }
1736
1737    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1738    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1739    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1740    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1741    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1742    fn q40_mr1_on() -> bool {
1743        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1744        match *Q40MR.get_or_init(|| {
1745            std::env::var("MEMRA_Q40_MR")
1746                .ok()
1747                .and_then(|v| v.parse().ok())
1748        }) {
1749            Some(v) => v == 1,
1750            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1751        }
1752    }
1753
1754    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1755    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1756    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1757    /// writes wrong bytes silently.
1758    fn pdl_func_flash(
1759        &self,
1760        g: bool,
1761        name: &'static str,
1762    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1763        use cudarc::driver::sys as cu;
1764        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1765        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1766        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1767        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1768        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1769        // this engine's CUcontext; single-context runs behave exactly as before.
1770        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1771            std::sync::Mutex::new(None);
1772        static FNS: std::sync::Mutex<
1773            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1774        > = std::sync::Mutex::new(None);
1775        let ctx_key = self.ctx().cu_ctx() as usize;
1776        if let Some(&f) = FNS
1777            .lock()
1778            .unwrap()
1779            .get_or_insert_with(Default::default)
1780            .get(&(ctx_key, g, name))
1781        {
1782            return Ok(f as cu::CUfunction);
1783        }
1784        let module = {
1785            let mut mods = MODS.lock().unwrap();
1786            let map = mods.get_or_insert_with(Default::default);
1787            match map.get(&(ctx_key, g)) {
1788                Some(&m) => m,
1789                None => {
1790                    let m = self.pdl_load_module_in_ctx(if g {
1791                        FLASH_FATBIN_KF8VF8
1792                    } else {
1793                        FLASH_FATBIN
1794                    })?;
1795                    map.insert((ctx_key, g), m);
1796                    m
1797                }
1798            }
1799        };
1800        let cname = std::ffi::CString::new(name)?;
1801        let mut f: cu::CUfunction = std::ptr::null_mut();
1802        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1803        if r != cu::CUresult::CUDA_SUCCESS {
1804            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1805        }
1806        FNS.lock()
1807            .unwrap()
1808            .get_or_insert_with(Default::default)
1809            .insert((ctx_key, g, name), f as usize);
1810        Ok(f)
1811    }
1812
1813    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
1814    /// the module to the thread's CURRENT context — a remote-stage engine must not
1815    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
1816    /// current context before returning.
1817    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1818        use cudarc::driver::sys as cu;
1819        let mut prev: cu::CUcontext = std::ptr::null_mut();
1820        unsafe {
1821            cu::cuCtxGetCurrent(&mut prev).result()?;
1822        }
1823        self.ctx().bind_to_thread()?;
1824        let mut m: cu::CUmodule = std::ptr::null_mut();
1825        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1826        let restore = if prev.is_null() {
1827            cu::CUresult::CUDA_SUCCESS
1828        } else {
1829            unsafe { cu::cuCtxSetCurrent(prev) }
1830        };
1831        if r != cu::CUresult::CUDA_SUCCESS {
1832            return Err(format!("pdl module load: {r:?}").into());
1833        }
1834        if restore != cu::CUresult::CUDA_SUCCESS {
1835            return Err(format!("pdl module load: ctx restore {restore:?}").into());
1836        }
1837        Ok(m as usize)
1838    }
1839
1840    fn pdl_func(
1841        &self,
1842        name: &'static str,
1843    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1844        use cudarc::driver::sys as cu;
1845        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
1846        // are context-scoped; key everything by this engine's CUcontext).
1847        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1848            std::sync::Mutex::new(None);
1849        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
1850        // duplicate module, loaded lazily on the first kernels-module miss.
1851        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1852            std::sync::Mutex::new(None);
1853        static FNS: std::sync::Mutex<
1854            Option<std::collections::HashMap<(usize, &'static str), usize>>,
1855        > = std::sync::Mutex::new(None);
1856        let ctx_key = self.ctx().cu_ctx() as usize;
1857        if let Some(&f) = FNS
1858            .lock()
1859            .unwrap()
1860            .get_or_insert_with(Default::default)
1861            .get(&(ctx_key, name))
1862        {
1863            return Ok(f as cu::CUfunction);
1864        }
1865        let module = {
1866            let mut mods = MODULES.lock().unwrap();
1867            let map = mods.get_or_insert_with(Default::default);
1868            match map.get(&ctx_key) {
1869                Some(&m) => m,
1870                None => {
1871                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
1872                    map.insert(ctx_key, m);
1873                    m
1874                }
1875            }
1876        };
1877        let cname = std::ffi::CString::new(name)?;
1878        let mut f: cu::CUfunction = std::ptr::null_mut();
1879        let mut r =
1880            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1881        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1882            let qmodule = {
1883                let mut mods = QMODULES.lock().unwrap();
1884                let map = mods.get_or_insert_with(Default::default);
1885                match map.get(&ctx_key) {
1886                    Some(&m) => m,
1887                    None => {
1888                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1889                        map.insert(ctx_key, m);
1890                        m
1891                    }
1892                }
1893            };
1894            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1895        }
1896        if r != cu::CUresult::CUDA_SUCCESS {
1897            return Err(format!("pdl_func {name}: {r:?}").into());
1898        }
1899        FNS.lock()
1900            .unwrap()
1901            .get_or_insert_with(Default::default)
1902            .insert((ctx_key, name), f as usize);
1903        Ok(f)
1904    }
1905
1906    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
1907    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
1908    ///
1909    /// # Safety
1910    /// `params` must match the kernel's exact parameter list (order, types, count) —
1911    /// a mismatch corrupts the launch silently.
1912    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
1913    /// builder path's fa_func/func_g choice exactly).
1914    ///
1915    /// # Safety
1916    /// Same contract as `launch_pdl`.
1917    unsafe fn launch_pdl_flash(
1918        &self,
1919        g: bool,
1920        name: &'static str,
1921        grid: (u32, u32, u32),
1922        block: (u32, u32, u32),
1923        smem: u32,
1924        params: &mut [*mut std::ffi::c_void],
1925    ) -> Result<(), Box<dyn std::error::Error>> {
1926        use cudarc::driver::sys as cu;
1927        let f = self.pdl_func_flash(g, name)?;
1928        if smem > 0 {
1929            // mirror the builder path's opt-in ceiling (idempotent host-side set).
1930            let r =
1931                unsafe {
1932                    cu::cuFuncSetAttribute(f,
1933                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1934                smem as i32)
1935                };
1936            if r != cu::CUresult::CUDA_SUCCESS {
1937                return Err(format!("pdl smem attr {name}: {r:?}").into());
1938            }
1939        }
1940        let mut attr = cu::CUlaunchAttribute {
1941            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1942            pad: [0; 4],
1943            value: cu::CUlaunchAttributeValue {
1944                programmaticStreamSerializationAllowed: 1,
1945            },
1946        };
1947        let cfg = cu::CUlaunchConfig {
1948            gridDimX: grid.0,
1949            gridDimY: grid.1,
1950            gridDimZ: grid.2,
1951            blockDimX: block.0,
1952            blockDimY: block.1,
1953            blockDimZ: block.2,
1954            sharedMemBytes: smem,
1955            hStream: self.gpu.stream().cu_stream(),
1956            attrs: &mut attr,
1957            numAttrs: 1,
1958        };
1959        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1960        if r != cu::CUresult::CUDA_SUCCESS {
1961            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
1962        }
1963        Ok(())
1964    }
1965
1966    unsafe fn launch_pdl(
1967        &self,
1968        name: &'static str,
1969        grid: (u32, u32, u32),
1970        block: (u32, u32, u32),
1971        params: &mut [*mut std::ffi::c_void],
1972    ) -> Result<(), Box<dyn std::error::Error>> {
1973        use cudarc::driver::sys as cu;
1974        let f = self.pdl_func(name)?;
1975        let mut attr = cu::CUlaunchAttribute {
1976            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1977            pad: [0; 4],
1978            value: cu::CUlaunchAttributeValue {
1979                programmaticStreamSerializationAllowed: 1,
1980            },
1981        };
1982        let cfg = cu::CUlaunchConfig {
1983            gridDimX: grid.0,
1984            gridDimY: grid.1,
1985            gridDimZ: grid.2,
1986            blockDimX: block.0,
1987            blockDimY: block.1,
1988            blockDimZ: block.2,
1989            sharedMemBytes: 0,
1990            hStream: self.gpu.stream().cu_stream(),
1991            attrs: &mut attr,
1992            numAttrs: 1,
1993        };
1994        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1995        if r != cu::CUresult::CUDA_SUCCESS {
1996            return Err(format!("launch_pdl {name}: {r:?}").into());
1997        }
1998        Ok(())
1999    }
2000
2001    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2002    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2003    pub fn prefetch_weight_l2(
2004        &self,
2005        w: &crate::model::GpuTensor,
2006    ) -> Result<(), Box<dyn std::error::Error>> {
2007        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2008            let p = rp4.as_ref().unwrap_or(bytes);
2009            self.prefetch_l2(p, p.len())?;
2010        }
2011        Ok(())
2012    }
2013
2014    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2015    /// by the DEVICE token id at tok[idx] into f32.
2016    pub fn gather_row_bf16(
2017        &self,
2018        table: &CudaSlice<u8>,
2019        tok: &CudaSlice<u32>,
2020        idx: usize,
2021        dst: &mut CudaSlice<f32>,
2022        ncols: usize,
2023    ) -> Result<(), Box<dyn std::error::Error>> {
2024        let f = self.func("gather_row_bf16_f32");
2025        let cfg = LaunchConfig {
2026            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2027            block_dim: (256, 1, 1),
2028            shared_mem_bytes: 0,
2029        };
2030        let (nc, ix) = (ncols as i32, idx as i32);
2031        let __s_b = self.gpu.stream();
2032        let mut b = __s_b.launch_builder(&f);
2033        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2034        unsafe {
2035            b.launch(cfg)?;
2036        }
2037        Ok(())
2038    }
2039
2040    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2041    pub fn add_row_inplace(
2042        &self,
2043        logits: &mut CudaSlice<f32>,
2044        bias: &CudaSlice<f32>,
2045        n: usize,
2046        row_off: usize,
2047    ) -> Result<(), Box<dyn std::error::Error>> {
2048        let f = self.func("add_row_inplace_f32");
2049        let cfg = LaunchConfig {
2050            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2051            block_dim: (256, 1, 1),
2052            shared_mem_bytes: 0,
2053        };
2054        let (ni, off) = (n as i32, row_off as i64);
2055        let __s_b = self.gpu.stream();
2056        let mut b = __s_b.launch_builder(&f);
2057        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2058        unsafe {
2059            b.launch(cfg)?;
2060        }
2061        Ok(())
2062    }
2063
2064    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2065    pub fn prefetch_l2(
2066        &self,
2067        p: &CudaSlice<u8>,
2068        n: usize,
2069    ) -> Result<(), Box<dyn std::error::Error>> {
2070        let f = self.func("prefetch_l2_bytes");
2071        let lines = n.div_ceil(128);
2072        let ni = n as i64;
2073        let cfg = LaunchConfig {
2074            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2075            block_dim: (256, 1, 1),
2076            shared_mem_bytes: 0,
2077        };
2078        let __s_b = self.gpu.stream();
2079        let mut b = __s_b.launch_builder(&f);
2080        b.arg(p).arg(&ni);
2081        unsafe {
2082            b.launch(cfg)?;
2083        }
2084        Ok(())
2085    }
2086
2087    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2088    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2089    pub fn router_gemv(
2090        &self,
2091        w: &CudaSlice<f32>,
2092        x: &CudaSlice<f32>,
2093        n_embd: usize,
2094        n_experts: usize,
2095        t: usize,
2096    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2097        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2098        // stream differs) — too small to justify a numeric config change; deleted.
2099        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2100        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2101        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2102        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2103            Ok("0") => false,
2104            Ok(_) => true,
2105            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2106        };
2107        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2108        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2109        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2110        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2111        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2112        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2113        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2114        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2115        // (perf-only, bits equal).
2116        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2117        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2118    }
2119
2120    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2121    /// force both forms; `batch` requires `w8`).
2122    pub fn router_gemv_form(
2123        &self,
2124        w: &CudaSlice<f32>,
2125        x: &CudaSlice<f32>,
2126        n_embd: usize,
2127        n_experts: usize,
2128        t: usize,
2129        w8: bool,
2130        batch: bool,
2131    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2132        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2133        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2134        let f = if batch {
2135            self.func("router_gemv_f32_w8_batch")
2136        } else if w8 {
2137            self.func("router_gemv_f32_w8")
2138        } else {
2139            self.func("router_gemv_f32")
2140        };
2141        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2142        let cfg = if batch {
2143            LaunchConfig {
2144                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2145                block_dim: (32, 8, 1),
2146                shared_mem_bytes: 0,
2147            }
2148        } else {
2149            LaunchConfig {
2150                grid_dim: (n_experts as u32, t as u32, 1),
2151                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2152                shared_mem_bytes: 0,
2153            }
2154        };
2155        let __s_b = self.gpu.stream();
2156        let mut b = __s_b.launch_builder(&f);
2157        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2158        unsafe {
2159            b.launch(cfg)?;
2160        }
2161        Ok(y)
2162    }
2163
2164    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2165    pub fn rows_permute(
2166        &self,
2167        src: &CudaSlice<f32>,
2168        idx: &CudaSlice<i32>,
2169        nrows: usize,
2170        ncols: usize,
2171    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2172        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2173        let f = self.func("rows_permute_f32");
2174        let (nc, nr) = (ncols as i32, nrows as i32);
2175        let cfg = LaunchConfig {
2176            grid_dim: (nrows as u32, 1, 1),
2177            block_dim: (256, 1, 1),
2178            shared_mem_bytes: 0,
2179        };
2180        let __s_b = self.gpu.stream();
2181        let mut b = __s_b.launch_builder(&f);
2182        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2183        unsafe {
2184            b.launch(cfg)?;
2185        }
2186        Ok(dst)
2187    }
2188
2189    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2190    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2191    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2192    /// decode chain and the small-t spec-verify chain match per row by construction.
2193    pub fn sigmoid_dot_rows(
2194        &self,
2195        x: &CudaSlice<f32>,
2196        w: &CudaSlice<f32>,
2197        n_embd: usize,
2198        t: usize,
2199    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2200        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2201        // config; same class as MEMRA_ROUTER_V2).
2202        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2203        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2204            let gs = self.linear(x, w, t, n_embd, 1)?;
2205            let mut g = self.uninit(t)?;
2206            self.sigmoid(&gs, &mut g, t)?;
2207            return Ok(g);
2208        }
2209        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2210        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2211        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2212        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2213        // flags doctrine; this per-token form serves every t.
2214        let mut g = self.alloc_uninit::<f32>(t)?;
2215        let f = self.func("sigmoid_dot_rows_f32");
2216        let (ne, ti) = (n_embd as i32, t as i32);
2217        let cfg = LaunchConfig {
2218            grid_dim: (t as u32, 1, 1),
2219            block_dim: (32, 8, 1),
2220            shared_mem_bytes: 0,
2221        };
2222        let __s_b = self.gpu.stream();
2223        let mut b = __s_b.launch_builder(&f);
2224        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2225        unsafe {
2226            b.launch(cfg)?;
2227        }
2228        Ok(g)
2229    }
2230
2231    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2232    pub fn spec_rollback_stream(
2233        &self,
2234        len_ptrs: &CudaSlice<u64>,
2235        pos_start: &CudaSlice<i32>,
2236        acc: &CudaSlice<u32>,
2237        base: usize,
2238        n_rows: usize,
2239    ) -> Result<(), Box<dyn std::error::Error>> {
2240        let f = self.func("spec_rollback_stream");
2241        let (b, nr) = (base as i32, n_rows as i32);
2242        let cfg = LaunchConfig {
2243            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2244            block_dim: (64, 1, 1),
2245            shared_mem_bytes: 0,
2246        };
2247        let __s_bl = self.gpu.stream();
2248        let mut bl = __s_bl.launch_builder(&f);
2249        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2250        unsafe {
2251            bl.launch(cfg)?;
2252        }
2253        Ok(())
2254    }
2255
2256    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2257    pub fn plain_tok_ring(
2258        &self,
2259        vam: &CudaSlice<u32>,
2260        pos_start: &CudaSlice<i32>,
2261        base: usize,
2262        ring: &mut CudaSlice<u32>,
2263    ) -> Result<(), Box<dyn std::error::Error>> {
2264        let f = self.func("plain_tok_ring");
2265        let (b, cap) = (base as i32, ring.len() as i32);
2266        let cfg = LaunchConfig {
2267            grid_dim: (1, 1, 1),
2268            block_dim: (32, 1, 1),
2269            shared_mem_bytes: 0,
2270        };
2271        let __s_bl = self.gpu.stream();
2272        let mut bl = __s_bl.launch_builder(&f);
2273        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2274        unsafe {
2275            bl.launch(cfg)?;
2276        }
2277        Ok(())
2278    }
2279
2280    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2281    pub fn spec_ring_commit(
2282        &self,
2283        vtok: &CudaSlice<u32>,
2284        acc: &CudaSlice<u32>,
2285        brk: &CudaSlice<u32>,
2286        ring: &mut CudaSlice<u32>,
2287        pend: &mut CudaSlice<u32>,
2288    ) -> Result<(), Box<dyn std::error::Error>> {
2289        let f = self.func("spec_ring_commit");
2290        let cfg = LaunchConfig {
2291            grid_dim: (1, 1, 1),
2292            block_dim: (32, 1, 1),
2293            shared_mem_bytes: 0,
2294        };
2295        let __s_b = self.gpu.stream();
2296        let mut b = __s_b.launch_builder(&f);
2297        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2298        unsafe {
2299            b.launch(cfg)?;
2300        }
2301        Ok(())
2302    }
2303    pub fn i32_copy_add(
2304        &self,
2305        src: &CudaSlice<i32>,
2306        dst: &mut CudaSlice<i32>,
2307        delta: i32,
2308    ) -> Result<(), Box<dyn std::error::Error>> {
2309        let f = self.func("i32_copy_add");
2310        let cfg = LaunchConfig {
2311            grid_dim: (1, 1, 1),
2312            block_dim: (32, 1, 1),
2313            shared_mem_bytes: 0,
2314        };
2315        let __s_b = self.gpu.stream();
2316        let mut b = __s_b.launch_builder(&f);
2317        b.arg(src).arg(dst).arg(&delta);
2318        unsafe {
2319            b.launch(cfg)?;
2320        }
2321        Ok(())
2322    }
2323    pub fn u32_copy(
2324        &self,
2325        src: &CudaSlice<u32>,
2326        dst: &mut CudaSlice<u32>,
2327    ) -> Result<(), Box<dyn std::error::Error>> {
2328        let f = self.func("u32_copy");
2329        let cfg = LaunchConfig {
2330            grid_dim: (1, 1, 1),
2331            block_dim: (32, 1, 1),
2332            shared_mem_bytes: 0,
2333        };
2334        let __s_b = self.gpu.stream();
2335        let mut b = __s_b.launch_builder(&f);
2336        b.arg(src).arg(dst);
2337        unsafe {
2338            b.launch(cfg)?;
2339        }
2340        Ok(())
2341    }
2342
2343    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2344    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2345    /// caps acceptance exactly like drafting fewer tokens).
2346    pub fn spec_adapt_k(
2347        &self,
2348        acc: &CudaSlice<u32>,
2349        brk: &mut CudaSlice<u32>,
2350        floor: usize,
2351        cap: usize,
2352    ) -> Result<(), Box<dyn std::error::Error>> {
2353        let f = self.func("spec_adapt_k");
2354        let (fl, cp) = (floor as i32, cap as i32);
2355        let cfg = LaunchConfig {
2356            grid_dim: (1, 1, 1),
2357            block_dim: (32, 1, 1),
2358            shared_mem_bytes: 0,
2359        };
2360        let __s_b = self.gpu.stream();
2361        let mut b = __s_b.launch_builder(&f);
2362        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2363        unsafe {
2364            b.launch(cfg)?;
2365        }
2366        Ok(())
2367    }
2368
2369    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2370    pub fn spec_accept_greedy_dc(
2371        &self,
2372        preds: &CudaSlice<u32>,
2373        vtok: &CudaSlice<u32>,
2374        last_pred: &CudaSlice<u32>,
2375        brk: &CudaSlice<u32>,
2376        out: &mut CudaSlice<u32>,
2377    ) -> Result<(), Box<dyn std::error::Error>> {
2378        let f = self.func("spec_accept_greedy_dc");
2379        let cfg = LaunchConfig {
2380            grid_dim: (1, 1, 1),
2381            block_dim: (32, 1, 1),
2382            shared_mem_bytes: 0,
2383        };
2384        let __s_b = self.gpu.stream();
2385        let mut b = __s_b.launch_builder(&f);
2386        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2387        unsafe {
2388            b.launch(cfg)?;
2389        }
2390        Ok(())
2391    }
2392
2393    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2394    pub fn pos_iota(
2395        &self,
2396        pos0: &CudaSlice<i32>,
2397        out: &mut CudaSlice<i32>,
2398        t: usize,
2399    ) -> Result<(), Box<dyn std::error::Error>> {
2400        let f = self.func("pos_iota_i32");
2401        let ti = t as i32;
2402        let cfg = LaunchConfig {
2403            grid_dim: (1, 1, 1),
2404            block_dim: (t.max(1) as u32, 1, 1),
2405            shared_mem_bytes: 0,
2406        };
2407        let __s_b = self.gpu.stream();
2408        let mut b = __s_b.launch_builder(&f);
2409        b.arg(pos0).arg(out).arg(&ti);
2410        unsafe {
2411            b.launch(cfg)?;
2412        }
2413        Ok(())
2414    }
2415    #[allow(clippy::too_many_arguments)]
2416    pub fn append_kv_quantized_rows_dc(
2417        &self,
2418        k_rows: &CudaSlice<f32>,
2419        v_rows: &CudaSlice<f32>,
2420        kc: &mut CudaSlice<u8>,
2421        vc: &mut CudaSlice<u8>,
2422        t0_dev: &CudaSlice<i32>,
2423        t: usize,
2424        kv_dim_k: usize,
2425        kv_dim_v: usize,
2426        k_tok_bytes: usize,
2427        v_tok_bytes: usize,
2428        g: bool,
2429    ) -> Result<(), Box<dyn std::error::Error>> {
2430        let f = if g {
2431            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2432        } else {
2433            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2434        };
2435        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2436        let cfg = LaunchConfig {
2437            grid_dim: (nblk, t as u32, 1),
2438            block_dim: (32, 1, 1),
2439            shared_mem_bytes: 0,
2440        };
2441        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2442        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2443        let __s_b = self.gpu.stream();
2444        let mut b = __s_b.launch_builder(&f);
2445        b.arg(k_rows)
2446            .arg(v_rows)
2447            .arg(kc)
2448            .arg(vc)
2449            .arg(t0_dev)
2450            .arg(&kdk)
2451            .arg(&kdv)
2452            .arg(&ktb)
2453            .arg(&vtb);
2454        unsafe {
2455            b.launch(cfg)?;
2456        }
2457        Ok(())
2458    }
2459
2460    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2461    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2462    #[allow(clippy::too_many_arguments)]
2463    pub fn append_kv_quantized_row_dc_inc(
2464        &self,
2465        k_row: &CudaSlice<f32>,
2466        v_row: &CudaSlice<f32>,
2467        kc: &mut CudaSlice<u8>,
2468        vc: &mut CudaSlice<u8>,
2469        t0_dev: &mut CudaSlice<i32>,
2470        kv_dim_k: usize,
2471        kv_dim_v: usize,
2472        k_tok_bytes: usize,
2473        v_tok_bytes: usize,
2474        g: bool,
2475    ) -> Result<(), Box<dyn std::error::Error>> {
2476        let f = if g {
2477            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2478        } else {
2479            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2480        };
2481        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2482        let cfg = LaunchConfig {
2483            grid_dim: (1, 1, 1),
2484            block_dim: (nthreads, 1, 1),
2485            shared_mem_bytes: 0,
2486        };
2487        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2488        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2489        let __s_b = self.gpu.stream();
2490        let mut b = __s_b.launch_builder(&f);
2491        b.arg(k_row)
2492            .arg(v_row)
2493            .arg(kc)
2494            .arg(vc)
2495            .arg(t0_dev)
2496            .arg(&kdk)
2497            .arg(&kdv)
2498            .arg(&ktb)
2499            .arg(&vtb);
2500        unsafe {
2501            b.launch(cfg)?;
2502        }
2503        Ok(())
2504    }
2505
2506    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2507    pub fn pack_tok_p(
2508        &self,
2509        tok: &CudaSlice<u32>,
2510        p: &CudaSlice<f32>,
2511        out: &mut CudaSlice<u32>,
2512        slot: usize,
2513    ) -> Result<(), Box<dyn std::error::Error>> {
2514        let f = self.func("pack_tok_p");
2515        let sl = slot as i32;
2516        let cfg = LaunchConfig {
2517            grid_dim: (1, 1, 1),
2518            block_dim: (32, 1, 1),
2519            shared_mem_bytes: 0,
2520        };
2521        let __s_b = self.gpu.stream();
2522        let mut b = __s_b.launch_builder(&f);
2523        b.arg(tok).arg(p).arg(out).arg(&sl);
2524        unsafe {
2525            b.launch(cfg)?;
2526        }
2527        Ok(())
2528    }
2529    pub fn tok_map_u32(
2530        &self,
2531        tok: &mut CudaSlice<u32>,
2532        map: &CudaSlice<u32>,
2533    ) -> Result<(), Box<dyn std::error::Error>> {
2534        let f = self.func("tok_map_u32");
2535        let cfg = LaunchConfig {
2536            grid_dim: (1, 1, 1),
2537            block_dim: (32, 1, 1),
2538            shared_mem_bytes: 0,
2539        };
2540        let __s_b = self.gpu.stream();
2541        let mut b = __s_b.launch_builder(&f);
2542        b.arg(tok).arg(map);
2543        unsafe {
2544            b.launch(cfg)?;
2545        }
2546        Ok(())
2547    }
2548
2549    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2550    #[allow(clippy::too_many_arguments)]
2551    pub fn spec_assemble_verify(
2552        &self,
2553        tokp: &CudaSlice<u32>,
2554        pend: &CudaSlice<u32>,
2555        d2t: Option<&CudaSlice<u32>>,
2556        vtok: &mut CudaSlice<u32>,
2557        brk: &mut CudaSlice<u32>,
2558        p_min: f32,
2559        k: usize,
2560        pmin0: bool,
2561    ) -> Result<(), Box<dyn std::error::Error>> {
2562        let f = self.func("spec_assemble_verify");
2563        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2564        let cfg = LaunchConfig {
2565            grid_dim: (1, 1, 1),
2566            block_dim: (32, 1, 1),
2567            shared_mem_bytes: 0,
2568        };
2569        let __s_b = self.gpu.stream();
2570        let mut b = __s_b.launch_builder(&f);
2571        match d2t {
2572            Some(m) => {
2573                b.arg(tokp)
2574                    .arg(pend)
2575                    .arg(m)
2576                    .arg(vtok)
2577                    .arg(brk)
2578                    .arg(&p_min)
2579                    .arg(&ki)
2580                    .arg(&pm);
2581                unsafe {
2582                    b.launch(cfg)?;
2583                }
2584            }
2585            None => {
2586                let null: u64 = 0;
2587                b.arg(tokp)
2588                    .arg(pend)
2589                    .arg(&null)
2590                    .arg(vtok)
2591                    .arg(brk)
2592                    .arg(&p_min)
2593                    .arg(&ki)
2594                    .arg(&pm);
2595                unsafe {
2596                    b.launch(cfg)?;
2597                }
2598            }
2599        }
2600        Ok(())
2601    }
2602
2603    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
2604    #[allow(clippy::too_many_arguments)]
2605    pub fn ssm_conv_ring_rebuild_dc(
2606        &self,
2607        qkv_tm: &CudaSlice<f32>,
2608        ring_old: &CudaSlice<f32>,
2609        conv_state: &mut CudaSlice<f32>,
2610        conv_dim: usize,
2611        acc: &CudaSlice<u32>,
2612        base: usize,
2613        t_v: usize,
2614        d_conv: usize,
2615    ) -> Result<(), Box<dyn std::error::Error>> {
2616        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2617        let n = conv_dim * (d_conv - 1);
2618        let cfg = LaunchConfig::for_num_elems(n as u32);
2619        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2620        let __s_b = self.gpu.stream();
2621        let mut b = __s_b.launch_builder(&f);
2622        b.arg(qkv_tm)
2623            .arg(ring_old)
2624            .arg(conv_state)
2625            .arg(&cd)
2626            .arg(acc)
2627            .arg(&b0)
2628            .arg(&tv)
2629            .arg(&dc);
2630        unsafe {
2631            b.launch(cfg)?;
2632        }
2633        Ok(())
2634    }
2635    #[allow(clippy::too_many_arguments)]
2636    pub fn gdn_scan_s128_dc(
2637        &self,
2638        q: &CudaSlice<f32>,
2639        k: &CudaSlice<f32>,
2640        v: &CudaSlice<f32>,
2641        g: &CudaSlice<f32>,
2642        beta: &CudaSlice<f32>,
2643        state_in: &CudaSlice<f32>,
2644        state_out: &mut CudaSlice<f32>,
2645        o: &mut CudaSlice<f32>,
2646        n_head: usize,
2647        acc: &CudaSlice<u32>,
2648        base: usize,
2649        t_v: usize,
2650        scale: f32,
2651    ) -> Result<(), Box<dyn std::error::Error>> {
2652        let f = self.func("gdn_scan_s128_dc");
2653        const S_V: u32 = 128;
2654        const WARP: u32 = 32;
2655        const COLS_PER_BLOCK: u32 = 4;
2656        let cfg = LaunchConfig {
2657            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2658            block_dim: (WARP, COLS_PER_BLOCK, 1),
2659            shared_mem_bytes: 0,
2660        };
2661        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2662        let __s_b = self.gpu.stream();
2663        let mut b = __s_b.launch_builder(&f);
2664        b.arg(q)
2665            .arg(k)
2666            .arg(v)
2667            .arg(g)
2668            .arg(beta)
2669            .arg(state_in)
2670            .arg(state_out)
2671            .arg(o)
2672            .arg(&h)
2673            .arg(acc)
2674            .arg(&b0)
2675            .arg(&tv)
2676            .arg(&scale);
2677        unsafe {
2678            b.launch(cfg)?;
2679        }
2680        Ok(())
2681    }
2682
2683    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
2684    pub fn spec_rollback_kv(
2685        &self,
2686        len_ptrs: &CudaSlice<u64>,
2687        saved: &CudaSlice<i32>,
2688        acc: &CudaSlice<u32>,
2689        base: usize,
2690        n_layer: usize,
2691    ) -> Result<(), Box<dyn std::error::Error>> {
2692        let f = self.func("spec_rollback_kv");
2693        let (b, nl) = (base as i32, n_layer as i32);
2694        let cfg = LaunchConfig {
2695            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2696            block_dim: (64, 1, 1),
2697            shared_mem_bytes: 0,
2698        };
2699        let __s_bl = self.gpu.stream();
2700        let mut bl = __s_bl.launch_builder(&f);
2701        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2702        unsafe {
2703            bl.launch(cfg)?;
2704        }
2705        Ok(())
2706    }
2707
2708    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
2709    pub fn spec_fork_valid(
2710        &self,
2711        acc: &CudaSlice<u32>,
2712        optimistic_pending: u32,
2713        valid: &mut CudaSlice<u32>,
2714    ) -> Result<(), Box<dyn std::error::Error>> {
2715        let f = self.func("spec_fork_valid");
2716        let cfg = LaunchConfig {
2717            grid_dim: (1, 1, 1),
2718            block_dim: (1, 1, 1),
2719            shared_mem_bytes: 0,
2720        };
2721        let __s_bl = self.gpu.stream();
2722        let mut bl = __s_bl.launch_builder(&f);
2723        bl.arg(acc).arg(&optimistic_pending).arg(valid);
2724        unsafe {
2725            bl.launch(cfg)?;
2726        }
2727        Ok(())
2728    }
2729
2730    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
2731    pub fn spec_fork_reconcile_kv(
2732        &self,
2733        len_ptrs: &CudaSlice<u64>,
2734        saved: &CudaSlice<i32>,
2735        acc: &CudaSlice<u32>,
2736        valid: &CudaSlice<u32>,
2737        base: usize,
2738        n_layer: usize,
2739    ) -> Result<(), Box<dyn std::error::Error>> {
2740        let f = self.func("spec_fork_reconcile_kv");
2741        let (b, nl) = (base as i32, n_layer as i32);
2742        let cfg = LaunchConfig {
2743            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2744            block_dim: (64, 1, 1),
2745            shared_mem_bytes: 0,
2746        };
2747        let __s_bl = self.gpu.stream();
2748        let mut bl = __s_bl.launch_builder(&f);
2749        bl.arg(len_ptrs)
2750            .arg(saved)
2751            .arg(acc)
2752            .arg(valid)
2753            .arg(&b)
2754            .arg(&nl);
2755        unsafe {
2756            bl.launch(cfg)?;
2757        }
2758        Ok(())
2759    }
2760
2761    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
2762    pub fn spec_fork_restore_f32(
2763        &self,
2764        snapshot: &CudaSlice<f32>,
2765        state: &mut CudaSlice<f32>,
2766        valid: &CudaSlice<u32>,
2767    ) -> Result<(), Box<dyn std::error::Error>> {
2768        assert_eq!(
2769            snapshot.len(),
2770            state.len(),
2771            "fork recurrent snapshot shape mismatch"
2772        );
2773        let f = self.func("spec_fork_restore_f32");
2774        let n = state.len() as i32;
2775        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2776        let cfg = LaunchConfig {
2777            grid_dim: (blocks, 1, 1),
2778            block_dim: (256, 1, 1),
2779            shared_mem_bytes: 0,
2780        };
2781        let __s_bl = self.gpu.stream();
2782        let mut bl = __s_bl.launch_builder(&f);
2783        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2784        unsafe {
2785            bl.launch(cfg)?;
2786        }
2787        Ok(())
2788    }
2789
2790    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
2791    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
2792    pub fn spec_seed_gather(
2793        &self,
2794        vx: &CudaSlice<f32>,
2795        fill_prev: &CudaSlice<f32>,
2796        acc: &CudaSlice<u32>,
2797        h_seed: &mut CudaSlice<f32>,
2798        base: usize,
2799        n_embd: usize,
2800    ) -> Result<(), Box<dyn std::error::Error>> {
2801        let f = self.func("spec_seed_gather");
2802        let (b, ne) = (base as i32, n_embd as i32);
2803        let cfg = LaunchConfig {
2804            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2805            block_dim: (256, 1, 1),
2806            shared_mem_bytes: 0,
2807        };
2808        let __s_bl = self.gpu.stream();
2809        let mut bl = __s_bl.launch_builder(&f);
2810        bl.arg(vx)
2811            .arg(fill_prev)
2812            .arg(acc)
2813            .arg(h_seed)
2814            .arg(&b)
2815            .arg(&ne);
2816        unsafe {
2817            bl.launch(cfg)?;
2818        }
2819        Ok(())
2820    }
2821
2822    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
2823    pub fn spec_accept_greedy(
2824        &self,
2825        preds: &CudaSlice<u32>,
2826        draft: &CudaSlice<u32>,
2827        last_pred: u32,
2828        base: usize,
2829        k_round: usize,
2830        out: &mut CudaSlice<u32>,
2831    ) -> Result<(), Box<dyn std::error::Error>> {
2832        let f = self.func("spec_accept_greedy");
2833        let (b, k) = (base as i32, k_round as i32);
2834        let cfg = LaunchConfig {
2835            grid_dim: (1, 1, 1),
2836            block_dim: (32, 1, 1),
2837            shared_mem_bytes: 0,
2838        };
2839        let __s_bl = self.gpu.stream();
2840        let mut bl = __s_bl.launch_builder(&f);
2841        bl.arg(preds)
2842            .arg(draft)
2843            .arg(&last_pred)
2844            .arg(&b)
2845            .arg(&k)
2846            .arg(out);
2847        unsafe {
2848            bl.launch(cfg)?;
2849        }
2850        Ok(())
2851    }
2852
2853    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
2854    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
2855    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
2856
2857    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
2858    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
2859    pub fn gumbel_perturb(
2860        &self,
2861        x: &CudaSlice<f32>,
2862        y: &mut CudaSlice<f32>,
2863        n: usize,
2864        seed: u64,
2865        stream_pos: u32,
2866        temp: f32,
2867    ) -> Result<(), Box<dyn std::error::Error>> {
2868        let f = self.func("gumbel_perturb_f32");
2869        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2870        let cfg = LaunchConfig {
2871            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2872            block_dim: (256, 1, 1),
2873            shared_mem_bytes: 0,
2874        };
2875        let __s_b = self.gpu.stream();
2876        let mut b = __s_b.launch_builder(&f);
2877        b.arg(x)
2878            .arg(&mut *y)
2879            .arg(&ni)
2880            .arg(&slo)
2881            .arg(&shi)
2882            .arg(&stream_pos)
2883            .arg(&temp);
2884        unsafe {
2885            b.launch(cfg)?;
2886        }
2887        Ok(())
2888    }
2889
2890    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
2891    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
2892    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
2893    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
2894    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
2895    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
2896    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
2897    pub fn mask_logits_col(
2898        &self,
2899        logits: &mut CudaSlice<f32>,
2900        mask: &CudaSlice<u32>,
2901        col: usize,
2902        n: usize,
2903        mask_words: usize,
2904    ) -> Result<(), Box<dyn std::error::Error>> {
2905        let f = self.func("mask_logits_f32");
2906        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2907        let cfg = LaunchConfig {
2908            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2909            block_dim: (256, 1, 1),
2910            shared_mem_bytes: 0,
2911        };
2912        let __s_b = self.gpu.stream();
2913        let mut b = __s_b.launch_builder(&f);
2914        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2915        unsafe {
2916            b.launch(cfg)?;
2917        }
2918        Ok(())
2919    }
2920
2921    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
2922    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
2923    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
2924    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
2925    /// (the lane index is the in-row position; `col` only moves the input pointer). That
2926    /// pointer-invariance IS the serving isolation contract for sampled rows.
2927    pub fn gumbel_perturb_col(
2928        &self,
2929        x: &CudaSlice<f32>,
2930        col: usize,
2931        y: &mut CudaSlice<f32>,
2932        n: usize,
2933        seed: u64,
2934        stream_pos: u32,
2935        temp: f32,
2936    ) -> Result<(), Box<dyn std::error::Error>> {
2937        let f = self.func("gumbel_perturb_f32");
2938        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2939        let col_view = x.slice(col * n..(col + 1) * n);
2940        let cfg = LaunchConfig {
2941            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2942            block_dim: (256, 1, 1),
2943            shared_mem_bytes: 0,
2944        };
2945        let __s_b = self.gpu.stream();
2946        let mut b = __s_b.launch_builder(&f);
2947        b.arg(&col_view)
2948            .arg(&mut *y)
2949            .arg(&ni)
2950            .arg(&slo)
2951            .arg(&shi)
2952            .arg(&stream_pos)
2953            .arg(&temp);
2954        unsafe {
2955            b.launch(cfg)?;
2956        }
2957        Ok(())
2958    }
2959
2960    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
2961    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
2962    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
2963    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
2964    /// the serving isolation contract for sampled rows).
2965    #[allow(clippy::too_many_arguments)]
2966    pub fn gumbel_perturb_filtered_col(
2967        &self,
2968        x: &CudaSlice<f32>,
2969        col: usize,
2970        y: &mut CudaSlice<f32>,
2971        n: usize,
2972        seed: u64,
2973        stream_pos: u32,
2974        temp: f32,
2975        stat_max: &CudaSlice<f32>,
2976        stat_th: &CudaSlice<f32>,
2977        stat_idx: usize,
2978    ) -> Result<(), Box<dyn std::error::Error>> {
2979        let f = self.func("gumbel_perturb_filtered_col_f32");
2980        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2981        let (ci, si) = (col as i32, stat_idx as i32);
2982        let cfg = LaunchConfig {
2983            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2984            block_dim: (256, 1, 1),
2985            shared_mem_bytes: 0,
2986        };
2987        let __s_b = self.gpu.stream();
2988        let mut b = __s_b.launch_builder(&f);
2989        b.arg(x)
2990            .arg(&ci)
2991            .arg(&mut *y)
2992            .arg(&ni)
2993            .arg(&slo)
2994            .arg(&shi)
2995            .arg(&stream_pos)
2996            .arg(&temp)
2997            .arg(stat_max)
2998            .arg(stat_th)
2999            .arg(&si);
3000        unsafe {
3001            b.launch(cfg)?;
3002        }
3003        Ok(())
3004    }
3005
3006    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3007    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3008    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3009    /// reads it (counter is data, not state — graph-replay-safe).
3010    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3011        let f = self.func("memra_sctr_inc");
3012        let cfg = LaunchConfig {
3013            grid_dim: (1, 1, 1),
3014            block_dim: (1, 1, 1),
3015            shared_mem_bytes: 0,
3016        };
3017        let __s_b = self.gpu.stream();
3018        let mut b = __s_b.launch_builder(&f);
3019        b.arg(&mut *ctr);
3020        unsafe {
3021            b.launch(cfg)?;
3022        }
3023        Ok(())
3024    }
3025
3026    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3027    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3028    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3029    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3030    pub fn gumbel_perturb_ctr(
3031        &self,
3032        x: &CudaSlice<f32>,
3033        y: &mut CudaSlice<f32>,
3034        n: usize,
3035        seed: u64,
3036        ctr: &CudaSlice<u32>,
3037        temp: f32,
3038    ) -> Result<(), Box<dyn std::error::Error>> {
3039        let f = self.func("gumbel_perturb_ctr_f32");
3040        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3041        let cfg = LaunchConfig {
3042            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3043            block_dim: (256, 1, 1),
3044            shared_mem_bytes: 0,
3045        };
3046        let __s_b = self.gpu.stream();
3047        let mut b = __s_b.launch_builder(&f);
3048        b.arg(x)
3049            .arg(&mut *y)
3050            .arg(&ni)
3051            .arg(&slo)
3052            .arg(&shi)
3053            .arg(ctr)
3054            .arg(&temp);
3055        unsafe {
3056            b.launch(cfg)?;
3057        }
3058        Ok(())
3059    }
3060
3061    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3062    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3063    /// (smallest-index tie-break — matches the argmax-gate contract).
3064    pub fn softmax_gather(
3065        &self,
3066        x: &CudaSlice<f32>,
3067        row_stride: usize,
3068        ids: &CudaSlice<u32>,
3069        rows: &CudaSlice<i32>,
3070        out: &mut CudaSlice<f32>,
3071        n: usize,
3072        npair: usize,
3073        temp: f32,
3074    ) -> Result<(), Box<dyn std::error::Error>> {
3075        let f = self.func("softmax_gather_f32");
3076        let (ni, rs) = (n as i32, row_stride as i64);
3077        let np = npair as i32;
3078        let cfg = LaunchConfig {
3079            grid_dim: (npair as u32, 1, 1),
3080            block_dim: (256, 1, 1),
3081            shared_mem_bytes: 0,
3082        };
3083        let __s_b = self.gpu.stream();
3084        let mut b = __s_b.launch_builder(&f);
3085        b.arg(x)
3086            .arg(&rs)
3087            .arg(ids)
3088            .arg(rows)
3089            .arg(&mut *out)
3090            .arg(&ni)
3091            .arg(&np)
3092            .arg(&temp);
3093        unsafe {
3094            b.launch(cfg)?;
3095        }
3096        Ok(())
3097    }
3098
3099    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3100    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3101    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3102    pub fn residual_sample(
3103        &self,
3104        p: &CudaSlice<f32>,
3105        q: Option<&CudaSlice<f32>>,
3106        n: usize,
3107        temp: f32,
3108        seed: u64,
3109        stream_pos: u32,
3110        out_tok: &mut CudaSlice<u32>,
3111    ) -> Result<(), Box<dyn std::error::Error>> {
3112        let f = self.func("residual_sample_f32");
3113        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3114        let nth = 1024u32;
3115        let cfg = LaunchConfig {
3116            grid_dim: (1, 1, 1),
3117            block_dim: (nth, 1, 1),
3118            shared_mem_bytes: 0,
3119        };
3120        let has_q: i32 = q.is_some() as i32;
3121        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3122        let __s_b = self.gpu.stream();
3123        let mut b = __s_b.launch_builder(&f);
3124        b.arg(p)
3125            .arg(qbuf)
3126            .arg(&has_q)
3127            .arg(&ni)
3128            .arg(&temp)
3129            .arg(&slo)
3130            .arg(&shi)
3131            .arg(&stream_pos)
3132            .arg(&mut *out_tok);
3133        unsafe {
3134            b.launch(cfg)?;
3135        }
3136        Ok(())
3137    }
3138
3139    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3140    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3141    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3142    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3143    pub fn with_moe_cache<R>(
3144        &self,
3145        max_block_bytes: usize,
3146        f: impl FnOnce(
3147            &mut crate::moe_cache::MoeSlotCache,
3148            &Engine,
3149        ) -> Result<R, Box<dyn std::error::Error>>,
3150    ) -> Result<R, Box<dyn std::error::Error>> {
3151        let mut guard = self.moe_cache.lock().unwrap();
3152        if guard.is_none() {
3153            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3154        }
3155        let cache = guard.as_mut().unwrap();
3156        f(cache, self)
3157    }
3158
3159    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3160    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3161    pub fn freeze_moe_cache(&self) {
3162        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3163            cache.freeze();
3164        }
3165    }
3166
3167    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3168    /// Never constructs a cache.
3169    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3170        self.moe_cache
3171            .lock()
3172            .unwrap()
3173            .as_ref()
3174            .map(crate::moe_cache::MoeSlotCache::export_residency)
3175    }
3176
3177    pub(crate) fn moe_cache_frozen(&self) -> bool {
3178        self.moe_cache
3179            .lock()
3180            .unwrap()
3181            .as_ref()
3182            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3183    }
3184
3185    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3186    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3187    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3188    /// while leaving the profiling warmup's established batched behavior untouched.
3189    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3190    /// tokenwise arm anyway.)
3191    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3192        crate::cpu_experts::configured()
3193            && self.moe_cache_frozen()
3194            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3195    }
3196
3197    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3198    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3199        assert!(
3200            self.moe_cache.lock().unwrap().is_none(),
3201            "MoE cache layout configured after cache construction"
3202        );
3203        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3204    }
3205
3206    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3207        self.moe_cache_layout.lock().unwrap().clone()
3208    }
3209
3210    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3211    pub fn moe_cache_enabled() -> bool {
3212        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3213    }
3214
3215    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3216    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3217    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3218        let guard = self.moe_cache.lock().unwrap();
3219        guard
3220            .as_ref()
3221            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3222    }
3223
3224    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3225    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3226    /// callers compare a before/after snapshot around a decode window.
3227    pub fn cpu_expert_stats(
3228        &self,
3229    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3230        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3231    }
3232
3233    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3234    /// the backend tail that resident-GPU expert work did not hide.
3235    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3236        crate::cpu_experts::predictor_stats()
3237    }
3238
3239    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3240        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3241    }
3242
3243    /// CPU-routed expert selections grouped by how many of their three projections were already
3244    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3245    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3246        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3247    }
3248
3249    /// Positioned-read proof-backend counters:
3250    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3251    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3252        let guard = self.moe_cache.lock().unwrap();
3253        guard
3254            .as_ref()
3255            .and_then(|cache| cache.pread_stats())
3256            .map(|stats| {
3257                (
3258                    stats.reads,
3259                    stats.bytes,
3260                    stats.read_errors,
3261                    stats.short_reads,
3262                    stats.fallbacks,
3263                    stats.buffer_waits,
3264                    stats.ring_full,
3265                )
3266            })
3267    }
3268
3269    /// Spill configuration values that warned and substituted their documented defaults.
3270    pub fn spill_config_fallbacks(&self) -> u64 {
3271        crate::spill_pread::config_fallbacks()
3272    }
3273
3274    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3275    pub fn moe_cache_reset_counters(&self) {
3276        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3277            c.reset_counters();
3278        }
3279    }
3280
3281    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3282        Ok(self.gpu.stream().clone_htod(v)?)
3283    }
3284
3285    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3286    /// past the final q4_0 block through their aligned window — the bytes never reach a
3287    /// result (funnelshift discards them) but must be mapped memory.
3288    pub fn htod_bytes_padded(
3289        &self,
3290        v: &[u8],
3291        pad: usize,
3292    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3293        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3294        {
3295            let mut view = d.slice_mut(0..v.len());
3296            self.gpu.stream().memcpy_htod(v, &mut view)?;
3297        }
3298        Ok(d)
3299    }
3300
3301    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3302    pub fn copy_into(
3303        &self,
3304        dst: &mut CudaSlice<f32>,
3305        off: usize,
3306        src: &CudaSlice<f32>,
3307        len: usize,
3308    ) -> Result<(), Box<dyn std::error::Error>> {
3309        let mut view = dst.slice_mut(off..off + len);
3310        self.gpu
3311            .stream()
3312            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3313        Ok(())
3314    }
3315
3316    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3317    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3318    pub fn copy_u8_into(
3319        &self,
3320        dst: &mut CudaSlice<u8>,
3321        off: usize,
3322        src: &CudaSlice<u8>,
3323        len: usize,
3324    ) -> Result<(), Box<dyn std::error::Error>> {
3325        let mut view = dst.slice_mut(off..off + len);
3326        self.gpu
3327            .stream()
3328            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3329        Ok(())
3330    }
3331
3332    /// D2D byte-range copy with explicit source and destination offsets.
3333    pub fn copy_u8_range_into(
3334        &self,
3335        dst: &mut CudaSlice<u8>,
3336        dst_off: usize,
3337        src: &CudaSlice<u8>,
3338        src_off: usize,
3339        len: usize,
3340    ) -> Result<(), Box<dyn std::error::Error>> {
3341        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3342        self.gpu
3343            .stream()
3344            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3345        Ok(())
3346    }
3347
3348    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3349    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3350    /// keeping the audited attention range contiguous without changing its absolute start.
3351    pub fn prepare_kv_append(
3352        &self,
3353        kv: &mut crate::cache::KvLayer,
3354        retain_from: usize,
3355        append_rows: usize,
3356    ) -> Result<usize, Box<dyn std::error::Error>> {
3357        let Some(plan) = kv
3358            .ring
3359            .as_ref()
3360            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3361            .transpose()?
3362        else {
3363            return Ok(kv.len);
3364        };
3365        match plan {
3366            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3367            crate::cache::KvRingAppend::Rebase {
3368                src_row,
3369                keep_rows,
3370                new_base,
3371                write_row,
3372            } => {
3373                if keep_rows > 0 {
3374                    let k_len = keep_rows * kv.k_tok_bytes;
3375                    let v_len = keep_rows * kv.v_tok_bytes;
3376                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3377                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3378                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3379                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3380                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3381                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3382                }
3383                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3384                Ok(write_row)
3385            }
3386        }
3387    }
3388
3389    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3390    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3391    pub fn htod_u8_into(
3392        &self,
3393        dst: &mut CudaSlice<u8>,
3394        off: usize,
3395        src: &[u8],
3396    ) -> Result<(), Box<dyn std::error::Error>> {
3397        let mut view = dst.slice_mut(off..off + src.len());
3398        self.gpu.stream().memcpy_htod(src, &mut view)?;
3399        Ok(())
3400    }
3401
3402    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3403        b.slice(0..len)
3404    }
3405
3406    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3407    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3408    pub fn view_u8_range<'a>(
3409        &self,
3410        b: &'a CudaSlice<u8>,
3411        start: usize,
3412        end: usize,
3413    ) -> cudarc::driver::CudaView<'a, u8> {
3414        b.slice(start..end)
3415    }
3416    pub fn view_u8<'a>(
3417        &self,
3418        b: &'a CudaSlice<u8>,
3419        len: usize,
3420    ) -> cudarc::driver::CudaView<'a, u8> {
3421        b.slice(0..len)
3422    }
3423
3424    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3425    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3426    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3427    pub fn append_kv_quantized(
3428        &self,
3429        k_row: &CudaSlice<f32>,
3430        v_row: &CudaSlice<f32>,
3431        kc: &mut CudaSlice<u8>,
3432        vc: &mut CudaSlice<u8>,
3433        t: usize,
3434        kv_dim_k: usize,
3435        kv_dim_v: usize,
3436        k_tok_bytes: usize,
3437        v_tok_bytes: usize,
3438        g: bool,
3439    ) -> Result<(), Box<dyn std::error::Error>> {
3440        let f = if g {
3441            self.func_g("append_quantize_kv_q8_0_q5_1")
3442        } else {
3443            self.func("append_quantize_kv_q8_0_q5_1")
3444        };
3445        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3446        let cfg = LaunchConfig {
3447            grid_dim: (nblk, 1, 1),
3448            block_dim: (32, 1, 1),
3449            shared_mem_bytes: 0,
3450        };
3451        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3452        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3453        let __s_b = self.gpu.stream();
3454        let mut b = __s_b.launch_builder(&f);
3455        b.arg(k_row)
3456            .arg(v_row)
3457            .arg(kc)
3458            .arg(vc)
3459            .arg(&ti)
3460            .arg(&kdk)
3461            .arg(&kdv)
3462            .arg(&ktb)
3463            .arg(&vtb);
3464        unsafe {
3465            b.launch(cfg)?;
3466        }
3467        Ok(())
3468    }
3469
3470    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3471    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3472    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3473    pub fn append_kv_quantized_dc(
3474        &self,
3475        k_row: &CudaSlice<f32>,
3476        v_row: &CudaSlice<f32>,
3477        kc: &mut CudaSlice<u8>,
3478        vc: &mut CudaSlice<u8>,
3479        t_dev: &CudaSlice<i32>,
3480        kv_dim_k: usize,
3481        kv_dim_v: usize,
3482        k_tok_bytes: usize,
3483        v_tok_bytes: usize,
3484        g: bool,
3485    ) -> Result<(), Box<dyn std::error::Error>> {
3486        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3487        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3488        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3489        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3490        if Self::pdl_on() && Self::pdl_wb_on() {
3491            use cudarc::driver::{DevicePtr, DevicePtrMut};
3492            let s = &self.gpu.stream();
3493            let (pk, _g0) = k_row.device_ptr(s);
3494            let (pv, _g1) = v_row.device_ptr(s);
3495            let (pkc, _g2) = kc.device_ptr_mut(s);
3496            let (pvc, _g3) = vc.device_ptr_mut(s);
3497            let (pt, _g4) = t_dev.device_ptr(s);
3498            let mut ps = [
3499                &pk as *const _ as *mut std::ffi::c_void,
3500                &pv as *const _ as *mut _,
3501                &pkc as *const _ as *mut _,
3502                &pvc as *const _ as *mut _,
3503                &pt as *const _ as *mut _,
3504                &kdk as *const _ as *mut _,
3505                &kdv as *const _ as *mut _,
3506                &ktb as *const _ as *mut _,
3507                &vtb as *const _ as *mut _,
3508            ];
3509            unsafe {
3510                self.launch_pdl_flash(
3511                    g,
3512                    "append_quantize_kv_q8_0_q5_1_dc",
3513                    (nblk, 1, 1),
3514                    (32, 1, 1),
3515                    0,
3516                    &mut ps,
3517                )?;
3518            }
3519            return Ok(());
3520        }
3521        let f = if g {
3522            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3523        } else {
3524            self.func("append_quantize_kv_q8_0_q5_1_dc")
3525        };
3526        let cfg = LaunchConfig {
3527            grid_dim: (nblk, 1, 1),
3528            block_dim: (32, 1, 1),
3529            shared_mem_bytes: 0,
3530        };
3531        let __s_b = self.gpu.stream();
3532        let mut b = __s_b.launch_builder(&f);
3533        b.arg(k_row)
3534            .arg(v_row)
3535            .arg(kc)
3536            .arg(vc)
3537            .arg(t_dev)
3538            .arg(&kdk)
3539            .arg(&kdv)
3540            .arg(&ktb)
3541            .arg(&vtb);
3542        unsafe {
3543            b.launch(cfg)?;
3544        }
3545        Ok(())
3546    }
3547
3548    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3549    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3550    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3551    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3552    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3553    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3554    #[allow(clippy::too_many_arguments)]
3555    pub fn append_kv_quantized_rows(
3556        &self,
3557        k_rows: &CudaSlice<f32>,
3558        v_rows: &CudaSlice<f32>,
3559        kc: &mut CudaSlice<u8>,
3560        vc: &mut CudaSlice<u8>,
3561        t0: usize,
3562        t: usize,
3563        kv_dim_k: usize,
3564        kv_dim_v: usize,
3565        k_tok_bytes: usize,
3566        v_tok_bytes: usize,
3567        g: bool,
3568    ) -> Result<(), Box<dyn std::error::Error>> {
3569        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3570            for i in 0..t {
3571                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3572                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3573                self.append_kv_quantized_view(
3574                    &k_row,
3575                    &v_row,
3576                    kc,
3577                    vc,
3578                    t0 + i,
3579                    kv_dim_k,
3580                    kv_dim_v,
3581                    k_tok_bytes,
3582                    v_tok_bytes,
3583                    g,
3584                )?;
3585            }
3586            return Ok(());
3587        }
3588        let f = if g {
3589            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3590        } else {
3591            self.func("append_quantize_kv_q8_0_q5_1_rows")
3592        };
3593        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3594        let cfg = LaunchConfig {
3595            grid_dim: (nblk, t as u32, 1),
3596            block_dim: (32, 1, 1),
3597            shared_mem_bytes: 0,
3598        };
3599        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3600        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3601        let __s_b = self.gpu.stream();
3602        let mut b = __s_b.launch_builder(&f);
3603        b.arg(k_rows)
3604            .arg(v_rows)
3605            .arg(kc)
3606            .arg(vc)
3607            .arg(&t0i)
3608            .arg(&kdk)
3609            .arg(&kdv)
3610            .arg(&ktb)
3611            .arg(&vtb);
3612        unsafe {
3613            b.launch(cfg)?;
3614        }
3615        Ok(())
3616    }
3617
3618    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
3619    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
3620    /// later, inside a captured graph) without a host round-trip.
3621    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3622        let f = self.func("inc_i32");
3623        let cfg = LaunchConfig {
3624            grid_dim: (1, 1, 1),
3625            block_dim: (1, 1, 1),
3626            shared_mem_bytes: 0,
3627        };
3628        let __s_b = self.gpu.stream();
3629        let mut b = __s_b.launch_builder(&f);
3630        b.arg(p);
3631        unsafe {
3632            b.launch(cfg)?;
3633        }
3634        Ok(())
3635    }
3636
3637    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
3638    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
3639    pub fn append_kv_quantized_view(
3640        &self,
3641        k_row: &cudarc::driver::CudaView<f32>,
3642        v_row: &cudarc::driver::CudaView<f32>,
3643        kc: &mut CudaSlice<u8>,
3644        vc: &mut CudaSlice<u8>,
3645        t: usize,
3646        kv_dim_k: usize,
3647        kv_dim_v: usize,
3648        k_tok_bytes: usize,
3649        v_tok_bytes: usize,
3650        g: bool,
3651    ) -> Result<(), Box<dyn std::error::Error>> {
3652        let f = if g {
3653            self.func_g("append_quantize_kv_q8_0_q5_1")
3654        } else {
3655            self.func("append_quantize_kv_q8_0_q5_1")
3656        };
3657        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3658        let cfg = LaunchConfig {
3659            grid_dim: (nblk, 1, 1),
3660            block_dim: (32, 1, 1),
3661            shared_mem_bytes: 0,
3662        };
3663        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3664        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3665        let __s_b = self.gpu.stream();
3666        let mut b = __s_b.launch_builder(&f);
3667        b.arg(k_row)
3668            .arg(v_row)
3669            .arg(kc)
3670            .arg(vc)
3671            .arg(&ti)
3672            .arg(&kdk)
3673            .arg(&kdv)
3674            .arg(&ktb)
3675            .arg(&vtb);
3676        unsafe {
3677            b.launch(cfg)?;
3678        }
3679        Ok(())
3680    }
3681
3682    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
3683    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
3684    pub fn copy_view_into(
3685        &self,
3686        dst: &mut CudaSlice<f32>,
3687        off: usize,
3688        src: &cudarc::driver::CudaView<f32>,
3689        len: usize,
3690    ) -> Result<(), Box<dyn std::error::Error>> {
3691        let mut view = dst.slice_mut(off..off + len);
3692        self.gpu
3693            .stream()
3694            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3695        Ok(())
3696    }
3697
3698    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
3699    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
3700    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
3701    pub fn clone_dtod(
3702        &self,
3703        src: &CudaSlice<f32>,
3704    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3705        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3706        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3707        Ok(dst)
3708    }
3709
3710    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
3711    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
3712    pub fn dtod_copy_view(
3713        &self,
3714        src: &cudarc::driver::CudaView<f32>,
3715        dst: &mut CudaSlice<f32>,
3716    ) -> Result<(), Box<dyn std::error::Error>> {
3717        self.gpu.stream().memcpy_dtod(src, dst)?;
3718        Ok(())
3719    }
3720
3721    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
3722    pub fn dtod_copy_view_i8(
3723        &self,
3724        src: &cudarc::driver::CudaView<i8>,
3725        dst: &mut CudaSlice<i8>,
3726    ) -> Result<(), Box<dyn std::error::Error>> {
3727        self.gpu.stream().memcpy_dtod(src, dst)?;
3728        Ok(())
3729    }
3730
3731    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
3732    pub fn dtod_copy_into(
3733        &self,
3734        src: &CudaSlice<f32>,
3735        dst: &mut CudaSlice<f32>,
3736        offset: usize,
3737    ) -> Result<(), Box<dyn std::error::Error>> {
3738        let n = src.len();
3739        let mut dv = dst.slice_mut(offset..offset + n);
3740        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3741        Ok(())
3742    }
3743
3744    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
3745    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3746        self.alloc_uninit::<i8>(n)
3747    }
3748
3749    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
3750    pub fn qmatvec(
3751        &self,
3752        w: &CudaSlice<u8>,
3753        x: &CudaSlice<f32>,
3754        m: usize,
3755        in_f: usize,
3756        out_f: usize,
3757        qtype: i32,
3758        row_bytes: usize,
3759    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3760        let f = self.func("qmatvec_f32");
3761        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
3762        let cfg = LaunchConfig {
3763            grid_dim: (out_f as u32, m as u32, 1),
3764            block_dim: (256, 1, 1),
3765            shared_mem_bytes: 0,
3766        };
3767        let (inf, outf, mi, qt, rb) =
3768            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3769        let __s_b = self.gpu.stream();
3770        let mut b = __s_b.launch_builder(&f);
3771        b.arg(w)
3772            .arg(x)
3773            .arg(&mut y)
3774            .arg(&inf)
3775            .arg(&outf)
3776            .arg(&mi)
3777            .arg(&qt)
3778            .arg(&rb);
3779        unsafe {
3780            b.launch(cfg)?;
3781        }
3782        Ok(y)
3783    }
3784
3785    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
3786    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3787        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3788        self.keep_if_capturing(&s);
3789        Ok(s)
3790    }
3791
3792    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
3793    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
3794    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
3795    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3796        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3797        self.keep_if_capturing(&s);
3798        Ok(s)
3799    }
3800
3801    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
3802    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
3803    pub fn memset_zeros_view(
3804        &self,
3805        dst: &mut cudarc::driver::CudaViewMut<f32>,
3806    ) -> Result<(), Box<dyn std::error::Error>> {
3807        self.gpu.stream().memset_zeros(dst)?;
3808        Ok(())
3809    }
3810
3811    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
3812    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
3813    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
3814    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
3815    /// stream would require an event).
3816    pub fn stage_expert(
3817        &self,
3818        host_bytes: &[u8],
3819        scratch: &mut CudaSlice<u8>,
3820        off: usize,
3821    ) -> Result<(), Box<dyn std::error::Error>> {
3822        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
3823        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
3824        Ok(())
3825    }
3826
3827    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
3828    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
3829    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
3830    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
3831    /// One CTA per token row, 256 threads (one per expert).
3832    pub fn moe_router_topk(
3833        &self,
3834        logits: &CudaSlice<f32>,
3835        t: usize,
3836        n_expert: usize,
3837        n_used: usize,
3838    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3839        let f = self.func("moe_router_topk_f32");
3840        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
3841        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
3842        let cfg = LaunchConfig {
3843            grid_dim: (t as u32, 1, 1),
3844            block_dim: (n_expert as u32, 1, 1),
3845            shared_mem_bytes: 0,
3846        };
3847        let (ne, nu) = (n_expert as i32, n_used as i32);
3848        let __s_b = self.gpu.stream();
3849        let mut b = __s_b.launch_builder(&f);
3850        b.arg(logits)
3851            .arg(&mut sel_idx)
3852            .arg(&mut sel_w)
3853            .arg(&ne)
3854            .arg(&nu);
3855        unsafe {
3856            b.launch(cfg)?;
3857        }
3858        Ok((sel_idx, sel_w))
3859    }
3860
3861    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
3862    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
3863    pub fn moe_router_topk_scaled(
3864        &self,
3865        logits: &CudaSlice<f32>,
3866        t: usize,
3867        n_expert: usize,
3868        n_used: usize,
3869        ex_scale: &CudaSlice<f32>,
3870    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3871        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
3872        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
3873        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
3874        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
3875        let f = self.func("moe_router_topk_scaled_f32");
3876        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3877        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3878        let cfg = LaunchConfig {
3879            grid_dim: (t as u32, 1, 1),
3880            block_dim: (n_expert as u32, 1, 1),
3881            shared_mem_bytes: 0,
3882        };
3883        let (ne, nu) = (n_expert as i32, n_used as i32);
3884        let __s_b = self.gpu.stream();
3885        let mut b = __s_b.launch_builder(&f);
3886        b.arg(logits)
3887            .arg(&mut sel_idx)
3888            .arg(&mut sel_w)
3889            .arg(&ne)
3890            .arg(&nu)
3891            .arg(ex_scale);
3892        unsafe {
3893            b.launch(cfg)?;
3894        }
3895        Ok((sel_idx, sel_w))
3896    }
3897
3898    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
3899    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
3900    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
3901    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
3902    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
3903    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
3904    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
3905    pub fn moe_router_topk_host(
3906        &self,
3907        logits: &CudaSlice<f32>,
3908        t: usize,
3909        n_expert: usize,
3910        n_used: usize,
3911    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3912        let f = self.func("moe_router_topk_f32");
3913        let n = t * n_used;
3914        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
3915        let mut sel_w = self.alloc_uninit::<f32>(n)?;
3916        let cfg = LaunchConfig {
3917            grid_dim: (t as u32, 1, 1),
3918            block_dim: (n_expert as u32, 1, 1),
3919            shared_mem_bytes: 0,
3920        };
3921        let (ne, nu) = (n_expert as i32, n_used as i32);
3922        let __s_b = self.gpu.stream();
3923        let mut b = __s_b.launch_builder(&f);
3924        b.arg(logits)
3925            .arg(&mut sel_idx)
3926            .arg(&mut sel_w)
3927            .arg(&ne)
3928            .arg(&nu);
3929        unsafe {
3930            b.launch(cfg)?;
3931        }
3932        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
3933        let bytes = n * 8;
3934        let mut guard = self.router_stage.lock().unwrap();
3935        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
3936            *guard = Some(PinnedStage::new(bytes.max(4096))?);
3937        }
3938        let stage = guard.as_mut().unwrap();
3939        let (si, sw) = unsafe {
3940            (
3941                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
3942                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
3943            )
3944        };
3945        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
3946        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
3947        self.gpu.stream().synchronize()?; // ONE sync for both
3948        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
3949    }
3950
3951    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
3952    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
3953    /// original expert ids before top-k. Exact key ties choose the smaller original id.
3954    #[allow(clippy::too_many_arguments)]
3955    pub fn moe_router_sigmoid_topk(
3956        &self,
3957        logits: &CudaSlice<f32>,
3958        t: usize,
3959        n_expert: usize,
3960        n_used: usize,
3961        active_count: usize,
3962        correction_bias: &CudaSlice<f32>,
3963        active: &CudaSlice<u8>,
3964        scaling_factor: f32,
3965        route_norm: bool,
3966    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3967        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
3968        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
3969            return Err(format!(
3970                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
3971            )
3972            .into());
3973        }
3974        if logits.len() < t * n_expert
3975            || correction_bias.len() != n_expert
3976            || active.len() != n_expert
3977        {
3978            return Err(format!(
3979                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
3980                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
3981            ).into());
3982        }
3983        let f = self.func("moe_router_sigmoid_topk_f32");
3984        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3985        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3986        let threads = n_expert.div_ceil(32) * 32;
3987        let cfg = LaunchConfig {
3988            grid_dim: (t as u32, 1, 1),
3989            block_dim: (threads as u32, 1, 1),
3990            shared_mem_bytes: 0,
3991        };
3992        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
3993        let __s_b = self.gpu.stream();
3994        let mut b = __s_b.launch_builder(&f);
3995        b.arg(logits)
3996            .arg(correction_bias)
3997            .arg(active)
3998            .arg(&mut sel_idx)
3999            .arg(&mut sel_w)
4000            .arg(&ne)
4001            .arg(&nu)
4002            .arg(&scaling_factor)
4003            .arg(&rn);
4004        unsafe {
4005            b.launch(cfg)?;
4006        }
4007        Ok((sel_idx, sel_w))
4008    }
4009
4010    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4011    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4012    #[allow(clippy::too_many_arguments)]
4013    pub fn moe_router_sigmoid_topk_host(
4014        &self,
4015        logits: &CudaSlice<f32>,
4016        t: usize,
4017        n_expert: usize,
4018        n_used: usize,
4019        active_count: usize,
4020        correction_bias: &CudaSlice<f32>,
4021        active: &CudaSlice<u8>,
4022        scaling_factor: f32,
4023        route_norm: bool,
4024    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4025        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4026            logits,
4027            t,
4028            n_expert,
4029            n_used,
4030            active_count,
4031            correction_bias,
4032            active,
4033            scaling_factor,
4034            route_norm,
4035        )?;
4036        let n = t * n_used;
4037        let bytes = n * 8;
4038        let mut guard = self.router_stage.lock().unwrap();
4039        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4040            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4041        }
4042        let stage = guard.as_mut().unwrap();
4043        let (si, sw) = unsafe {
4044            (
4045                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4046                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4047            )
4048        };
4049        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4050        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4051        self.gpu.stream().synchronize()?;
4052        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4053    }
4054
4055    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4056    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4057    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4058    pub fn stage_expert_async(
4059        &self,
4060        host_bytes: &[u8],
4061        scratch: &mut CudaSlice<u8>,
4062        off: usize,
4063    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4064        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4065        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4066        Ok(self.copy_stream.record_event(None)?)
4067    }
4068
4069    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4070    pub fn compute_wait(
4071        &self,
4072        ev: &cudarc::driver::CudaEvent,
4073    ) -> Result<(), Box<dyn std::error::Error>> {
4074        self.gpu.stream().wait(ev)?;
4075        Ok(())
4076    }
4077
4078    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4079    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4080    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4081    /// CudaView base+offset pointer is honored by the launch arg.
4082    pub fn qmatvec_view(
4083        &self,
4084        w: &CudaSlice<u8>,
4085        range: std::ops::Range<usize>,
4086        x: &cudarc::driver::CudaView<f32>,
4087        m: usize,
4088        in_f: usize,
4089        out_f: usize,
4090        qtype: i32,
4091        row_bytes: usize,
4092    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4093        let f = self.func("qmatvec_f32");
4094        let wv = w.slice(range); // CudaView<u8>, offset honored
4095        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4096        let cfg = LaunchConfig {
4097            grid_dim: (out_f as u32, m as u32, 1),
4098            block_dim: (256, 1, 1),
4099            shared_mem_bytes: 0,
4100        };
4101        let (inf, outf, mi, qt, rb) =
4102            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4103        let __s_b = self.gpu.stream();
4104        let mut b = __s_b.launch_builder(&f);
4105        b.arg(&wv)
4106            .arg(x)
4107            .arg(&mut y)
4108            .arg(&inf)
4109            .arg(&outf)
4110            .arg(&mi)
4111            .arg(&qt)
4112            .arg(&rb);
4113        unsafe {
4114            b.launch(cfg)?;
4115        }
4116        Ok(y)
4117    }
4118
4119    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4120    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4121    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4122    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4123    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4124    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4125    #[allow(clippy::too_many_arguments)]
4126    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4127    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4128    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4129    pub fn moe_gate_up_silu8_q8(
4130        &self,
4131        gp: WPtr8,
4132        up: WPtr8,
4133        aq: &CudaSlice<i8>,
4134        ad: &CudaSlice<f32>,
4135        in_f: usize,
4136        n_ff: usize,
4137        n_used: usize,
4138        qt_g: i32,
4139        qt_u: i32,
4140        rb_g: usize,
4141        rb_u: usize,
4142    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4143        let f = self.func("moe_gate_up_silu8_q8");
4144        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4145        let cfg = LaunchConfig {
4146            grid_dim: (n_ff as u32, n_used as u32, 1),
4147            block_dim: (32, 1, 1),
4148            shared_mem_bytes: 0,
4149        };
4150        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4151        let __s_b = self.gpu.stream();
4152        let mut b = __s_b.launch_builder(&f);
4153        b.arg(&gp)
4154            .arg(&up)
4155            .arg(aq)
4156            .arg(ad)
4157            .arg(&mut act)
4158            .arg(&inf)
4159            .arg(&nff)
4160            .arg(&qt_g)
4161            .arg(&qt_u)
4162            .arg(&rbg)
4163            .arg(&rbu);
4164        unsafe {
4165            b.launch(cfg)?;
4166        }
4167        Ok(act)
4168    }
4169
4170    #[allow(clippy::too_many_arguments)]
4171    pub fn moe_down8_fma_q8(
4172        &self,
4173        dp: WPtr8,
4174        w: F32x8,
4175        aq2: &CudaSlice<i8>,
4176        ad2: &CudaSlice<f32>,
4177        dst: &mut cudarc::driver::CudaViewMut<f32>,
4178        in_f: usize,
4179        out_f: usize,
4180        n_used: usize,
4181        qt: i32,
4182        rb: usize,
4183    ) -> Result<(), Box<dyn std::error::Error>> {
4184        let f = self.func("moe_down8_fma_q8");
4185        let cfg = LaunchConfig {
4186            grid_dim: (out_f as u32, 1, 1),
4187            block_dim: (32, 1, 1),
4188            shared_mem_bytes: 0,
4189        };
4190        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4191        let __s_b = self.gpu.stream();
4192        let mut b = __s_b.launch_builder(&f);
4193        b.arg(&dp)
4194            .arg(&w)
4195            .arg(aq2)
4196            .arg(ad2)
4197            .arg(dst)
4198            .arg(&inf)
4199            .arg(&outf)
4200            .arg(&nu)
4201            .arg(&qt)
4202            .arg(&rbi);
4203        unsafe {
4204            b.launch(cfg)?;
4205        }
4206        Ok(())
4207    }
4208
4209    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4210    pub fn qmatvec_expert_q8(
4211        &self,
4212        w: &CudaSlice<u8>,
4213        range: std::ops::Range<usize>,
4214        aq: &CudaSlice<i8>,
4215        ad: &CudaSlice<f32>,
4216        m: usize,
4217        in_f: usize,
4218        out_f: usize,
4219        qtype: i32,
4220        row_bytes: usize,
4221    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4222        let f = self.func("qmatvec_expert_q8");
4223        let wv = w.slice(range);
4224        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4225        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4226        let cfg = LaunchConfig {
4227            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4228            block_dim: (32, ROWS, 1),
4229            shared_mem_bytes: 0,
4230        };
4231        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4232        let __s_b = self.gpu.stream();
4233        let mut b = __s_b.launch_builder(&f);
4234        b.arg(&wv)
4235            .arg(aq)
4236            .arg(ad)
4237            .arg(&mut y)
4238            .arg(&inf)
4239            .arg(&outf)
4240            .arg(&mi)
4241            .arg(&qtype)
4242            .arg(&rbi);
4243        unsafe {
4244            b.launch(cfg)?;
4245        }
4246        Ok(y)
4247    }
4248
4249    pub fn moe_gate_up_silu8(
4250        &self,
4251        gp: WPtr8,
4252        up: WPtr8,
4253        x: &cudarc::driver::CudaView<f32>,
4254        in_f: usize,
4255        n_ff: usize,
4256        n_used: usize,
4257        qt_g: i32,
4258        qt_u: i32,
4259        rb_g: usize,
4260        rb_u: usize,
4261    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4262        let f = self.func("moe_gate_up_silu8_f32");
4263        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4264        let cfg = LaunchConfig {
4265            grid_dim: (n_ff as u32, n_used as u32, 1),
4266            block_dim: (256, 1, 1),
4267            shared_mem_bytes: 0,
4268        };
4269        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4270        let __s_b = self.gpu.stream();
4271        let mut b = __s_b.launch_builder(&f);
4272        b.arg(&gp)
4273            .arg(&up)
4274            .arg(x)
4275            .arg(&mut act)
4276            .arg(&inf)
4277            .arg(&nff)
4278            .arg(&qt_g)
4279            .arg(&qt_u)
4280            .arg(&rbg)
4281            .arg(&rbu);
4282        unsafe {
4283            b.launch(cfg)?;
4284        }
4285        Ok(act)
4286    }
4287
4288    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4289    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4290    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4291    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4292    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4293    #[allow(clippy::too_many_arguments)]
4294    pub fn moe_down8_fma_into(
4295        &self,
4296        dp: WPtr8,
4297        w: F32x8,
4298        act: &CudaSlice<f32>,
4299        dst: &mut cudarc::driver::CudaViewMut<f32>,
4300        in_f: usize,
4301        out_f: usize,
4302        n_used: usize,
4303        qt: i32,
4304        rb: usize,
4305    ) -> Result<(), Box<dyn std::error::Error>> {
4306        let f = self.func("moe_down8_fma_f32");
4307        let cfg = LaunchConfig {
4308            grid_dim: (out_f as u32, 1, 1),
4309            block_dim: (256, 1, 1),
4310            shared_mem_bytes: 0,
4311        };
4312        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4313        let __s_b = self.gpu.stream();
4314        let mut b = __s_b.launch_builder(&f);
4315        b.arg(&dp)
4316            .arg(&w)
4317            .arg(act)
4318            .arg(dst)
4319            .arg(&inf)
4320            .arg(&outf)
4321            .arg(&nu)
4322            .arg(&qt)
4323            .arg(&rbv);
4324        unsafe {
4325            b.launch(cfg)?;
4326        }
4327        Ok(())
4328    }
4329
4330    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4331    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4332    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4333    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4334    #[allow(clippy::too_many_arguments)]
4335    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4336    ///
4337    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4338    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4339    /// down's FMA chain stays slot-ordered serial). Seams:
4340    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4341    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4342    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4343    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4344    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4345    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4346    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4347    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4348    ///                       only) | w8h2 (h2 x slot-parallel)
4349    #[allow(clippy::too_many_arguments)]
4350    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4351    #[allow(clippy::too_many_arguments)]
4352    pub fn moe_pairs_matvec_q8(
4353        &self,
4354        table: &CudaSlice<u64>,
4355        proj: i32,
4356        pair_tok: &CudaSlice<i32>,
4357        pair_ex: &CudaSlice<i32>,
4358        aq: &CudaSlice<i8>,
4359        ad: &CudaSlice<f32>,
4360        in_f: usize,
4361        out_f: usize,
4362        n_expert: usize,
4363        n_pairs: usize,
4364        qtype: i32,
4365        row_bytes: usize,
4366    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4367        let f = self.func("moe_pairs_matvec_q8");
4368        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4369        const ROWS: u32 = 4;
4370        let cfg = LaunchConfig {
4371            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4372            block_dim: (32, ROWS, 1),
4373            shared_mem_bytes: 0,
4374        };
4375        let (inf, outf, ne, np, rbi) = (
4376            in_f as i32,
4377            out_f as i32,
4378            n_expert as i32,
4379            n_pairs as i32,
4380            row_bytes as i64,
4381        );
4382        let __s_b = self.gpu.stream();
4383        let mut b = __s_b.launch_builder(&f);
4384        b.arg(table)
4385            .arg(&proj)
4386            .arg(pair_tok)
4387            .arg(pair_ex)
4388            .arg(aq)
4389            .arg(ad)
4390            .arg(&mut y)
4391            .arg(&inf)
4392            .arg(&outf)
4393            .arg(&ne)
4394            .arg(&np)
4395            .arg(&qtype)
4396            .arg(&rbi);
4397        unsafe {
4398            b.launch(cfg)?;
4399        }
4400        Ok(y)
4401    }
4402
4403    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4404    #[allow(clippy::too_many_arguments)]
4405    pub fn moe_pairs_matvec_q8_em(
4406        &self,
4407        table: &CudaSlice<u64>,
4408        proj: i32,
4409        ex_ids: &CudaSlice<i32>,
4410        ex_off: &CudaSlice<i32>,
4411        ex_pairs: &CudaSlice<i32>,
4412        pair_tok: &CudaSlice<i32>,
4413        aq: &CudaSlice<i8>,
4414        ad: &CudaSlice<f32>,
4415        in_f: usize,
4416        out_f: usize,
4417        n_expert: usize,
4418        n_active: usize,
4419        n_pairs: usize,
4420        qtype: i32,
4421        row_bytes: usize,
4422    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4423        let f = self.func("moe_pairs_matvec_q8_em");
4424        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4425        const ROWS: u32 = 4;
4426        let cfg = LaunchConfig {
4427            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4428            block_dim: (32, ROWS, 1),
4429            shared_mem_bytes: 0,
4430        };
4431        let (inf, outf, ne, na, rbi) = (
4432            in_f as i32,
4433            out_f as i32,
4434            n_expert as i32,
4435            n_active as i32,
4436            row_bytes as i64,
4437        );
4438        let __s_b = self.gpu.stream();
4439        let mut b = __s_b.launch_builder(&f);
4440        b.arg(table)
4441            .arg(&proj)
4442            .arg(ex_ids)
4443            .arg(ex_off)
4444            .arg(ex_pairs)
4445            .arg(pair_tok)
4446            .arg(aq)
4447            .arg(ad)
4448            .arg(&mut y)
4449            .arg(&inf)
4450            .arg(&outf)
4451            .arg(&ne)
4452            .arg(&na)
4453            .arg(&qtype)
4454            .arg(&rbi);
4455        unsafe {
4456            b.launch(cfg)?;
4457        }
4458        Ok(y)
4459    }
4460
4461    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
4462    // weight group once per (row,group) then dp4a's across the expert's token group.
4463    #[allow(clippy::too_many_arguments)]
4464    pub fn moe_pairs_matvec_q8_dec(
4465        &self,
4466        table: &CudaSlice<u64>,
4467        proj: i32,
4468        ex_ids: &CudaSlice<i32>,
4469        ex_off: &CudaSlice<i32>,
4470        ex_pairs: &CudaSlice<i32>,
4471        pair_tok: &CudaSlice<i32>,
4472        aq: &CudaSlice<i8>,
4473        ad: &CudaSlice<f32>,
4474        in_f: usize,
4475        out_f: usize,
4476        n_expert: usize,
4477        n_active: usize,
4478        n_pairs: usize,
4479        qtype: i32,
4480        row_bytes: usize,
4481    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4482        let f = self.func("moe_pairs_matvec_q8_dec");
4483        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4484        const ROWS: u32 = 4;
4485        let cfg = LaunchConfig {
4486            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4487            block_dim: (32, ROWS, 1),
4488            shared_mem_bytes: 0,
4489        };
4490        let (inf, outf, ne, na, rbi) = (
4491            in_f as i32,
4492            out_f as i32,
4493            n_expert as i32,
4494            n_active as i32,
4495            row_bytes as i64,
4496        );
4497        let __s_b = self.gpu.stream();
4498        let mut b = __s_b.launch_builder(&f);
4499        b.arg(table)
4500            .arg(&proj)
4501            .arg(ex_ids)
4502            .arg(ex_off)
4503            .arg(ex_pairs)
4504            .arg(pair_tok)
4505            .arg(aq)
4506            .arg(ad)
4507            .arg(&mut y)
4508            .arg(&inf)
4509            .arg(&outf)
4510            .arg(&ne)
4511            .arg(&na)
4512            .arg(&qtype)
4513            .arg(&rbi);
4514        unsafe {
4515            b.launch(cfg)?;
4516        }
4517        Ok(y)
4518    }
4519
4520    pub fn moe_pairs_gelu_mul(
4521        &self,
4522        gate: &CudaSlice<f32>,
4523        up: &CudaSlice<f32>,
4524        n: usize,
4525    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4526        let f = self.func("moe_pairs_gelu_mul");
4527        let mut act = self.alloc_uninit::<f32>(n)?;
4528        let cfg = LaunchConfig::for_num_elems(n as u32);
4529        let nl = n as i64;
4530        let __s_b = self.gpu.stream();
4531        let mut b = __s_b.launch_builder(&f);
4532        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4533        unsafe {
4534            b.launch(cfg)?;
4535        }
4536        Ok(act)
4537    }
4538
4539    pub fn moe_pairs_silu_mul(
4540        &self,
4541        gate: &CudaSlice<f32>,
4542        up: &CudaSlice<f32>,
4543        n: usize,
4544    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4545        let f = self.func("moe_pairs_silu_mul");
4546        let mut act = self.alloc_uninit::<f32>(n)?;
4547        let cfg = LaunchConfig::for_num_elems(n as u32);
4548        let nl = n as i64;
4549        let __s_b = self.gpu.stream();
4550        let mut b = __s_b.launch_builder(&f);
4551        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4552        unsafe {
4553            b.launch(cfg)?;
4554        }
4555        Ok(act)
4556    }
4557
4558    #[allow(clippy::too_many_arguments)]
4559    pub fn moe_pairs_scatter(
4560        &self,
4561        y_down: &CudaSlice<f32>,
4562        pair_w: &CudaSlice<f32>,
4563        tok_pair_off: &CudaSlice<i32>,
4564        tok_pair_ids: &CudaSlice<i32>,
4565        moe_out: &mut CudaSlice<f32>,
4566        t: usize,
4567        n_embd: usize,
4568    ) -> Result<(), Box<dyn std::error::Error>> {
4569        let f = self.func("moe_pairs_scatter");
4570        let cfg = LaunchConfig {
4571            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4572            block_dim: (256, 1, 1),
4573            shared_mem_bytes: 0,
4574        };
4575        let ne = n_embd as i32;
4576        let __s_b = self.gpu.stream();
4577        let mut b = __s_b.launch_builder(&f);
4578        b.arg(y_down)
4579            .arg(pair_w)
4580            .arg(tok_pair_off)
4581            .arg(tok_pair_ids)
4582            .arg(moe_out)
4583            .arg(&ne);
4584        unsafe {
4585            b.launch(cfg)?;
4586        }
4587        Ok(())
4588    }
4589
4590    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
4591    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
4592    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
4593    #[allow(clippy::too_many_arguments)]
4594    pub fn moe_gate_up_gelu8_dev_q8(
4595        &self,
4596        table: &CudaSlice<u64>,
4597        sel: &cudarc::driver::CudaView<i32>,
4598        aq: &CudaSlice<i8>,
4599        ad: &CudaSlice<f32>,
4600        in_f: usize,
4601        n_ff: usize,
4602        n_used: usize,
4603        n_expert: usize,
4604        qt_g: i32,
4605        qt_u: i32,
4606        rb_g: usize,
4607        rb_u: usize,
4608    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4609        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4610        let (inf, nff, ne, rbg, rbu) = (
4611            in_f as i32,
4612            n_ff as i32,
4613            n_expert as i32,
4614            rb_g as i64,
4615            rb_u as i64,
4616        );
4617        let f = self.func("moe_gate_up_gelu8_dev_q8");
4618        let cfg = LaunchConfig {
4619            grid_dim: (n_ff as u32, n_used as u32, 1),
4620            block_dim: (32, 1, 1),
4621            shared_mem_bytes: 0,
4622        };
4623        let __s_b = self.gpu.stream();
4624        let mut b = __s_b.launch_builder(&f);
4625        b.arg(table)
4626            .arg(sel)
4627            .arg(aq)
4628            .arg(ad)
4629            .arg(&mut act)
4630            .arg(&inf)
4631            .arg(&nff)
4632            .arg(&ne)
4633            .arg(&qt_g)
4634            .arg(&qt_u)
4635            .arg(&rbg)
4636            .arg(&rbu);
4637        unsafe {
4638            b.launch(cfg)?;
4639        }
4640        Ok(act)
4641    }
4642
4643    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
4644    #[allow(clippy::too_many_arguments)]
4645    pub fn moe_gate_up_gelu8_dev_q8_rows(
4646        &self,
4647        table: &CudaSlice<u64>,
4648        sel: &CudaSlice<i32>,
4649        aq: &CudaSlice<i8>,
4650        ad: &CudaSlice<f32>,
4651        t: usize,
4652        in_f: usize,
4653        n_ff: usize,
4654        n_used: usize,
4655        n_expert: usize,
4656        qt_g: i32,
4657        qt_u: i32,
4658        rb_g: usize,
4659        rb_u: usize,
4660    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4661        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4662        let (inf, nff, ne, rbg, rbu, nu) = (
4663            in_f as i32,
4664            n_ff as i32,
4665            n_expert as i32,
4666            rb_g as i64,
4667            rb_u as i64,
4668            n_used as i32,
4669        );
4670        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4671        let cfg = LaunchConfig {
4672            grid_dim: (n_ff as u32, n_used as u32, t as u32),
4673            block_dim: (32, 1, 1),
4674            shared_mem_bytes: 0,
4675        };
4676        let __s_b = self.gpu.stream();
4677        let mut b = __s_b.launch_builder(&f);
4678        b.arg(table)
4679            .arg(sel)
4680            .arg(aq)
4681            .arg(ad)
4682            .arg(&mut act)
4683            .arg(&inf)
4684            .arg(&nff)
4685            .arg(&ne)
4686            .arg(&qt_g)
4687            .arg(&qt_u)
4688            .arg(&rbg)
4689            .arg(&rbu)
4690            .arg(&nu);
4691        unsafe {
4692            b.launch(cfg)?;
4693        }
4694        Ok(act)
4695    }
4696
4697    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
4698    #[allow(clippy::too_many_arguments)]
4699    pub fn moe_gate_up_gelu8_dev_q8_csr(
4700        &self,
4701        table: &CudaSlice<u64>,
4702        sel: &CudaSlice<i32>,
4703        aq: &CudaSlice<i8>,
4704        ad: &CudaSlice<f32>,
4705        n_pairs: usize,
4706        in_f: usize,
4707        n_ff: usize,
4708        n_used: usize,
4709        n_expert: usize,
4710        qt_g: i32,
4711        qt_u: i32,
4712        rb_g: usize,
4713        rb_u: usize,
4714    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4715        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4716        let (inf, nff, ne, rbg, rbu, nu, npi) = (
4717            in_f as i32,
4718            n_ff as i32,
4719            n_expert as i32,
4720            rb_g as i64,
4721            rb_u as i64,
4722            n_used as i32,
4723            n_pairs as i32,
4724        );
4725        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4726        let cfg = LaunchConfig {
4727            grid_dim: (n_ff as u32, n_pairs as u32, 1),
4728            block_dim: (32, 1, 1),
4729            shared_mem_bytes: 0,
4730        };
4731        let __s_b = self.gpu.stream();
4732        let mut b = __s_b.launch_builder(&f);
4733        b.arg(table)
4734            .arg(sel)
4735            .arg(aq)
4736            .arg(ad)
4737            .arg(&mut act)
4738            .arg(&inf)
4739            .arg(&nff)
4740            .arg(&ne)
4741            .arg(&qt_g)
4742            .arg(&qt_u)
4743            .arg(&rbg)
4744            .arg(&rbu)
4745            .arg(&nu)
4746            .arg(&npi);
4747        unsafe {
4748            b.launch(cfg)?;
4749        }
4750        Ok(act)
4751    }
4752
4753    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
4754    #[allow(clippy::too_many_arguments)]
4755    pub fn moe_down8_fma_dev_q8_rows_g(
4756        &self,
4757        table: &CudaSlice<u64>,
4758        sel: &CudaSlice<i32>,
4759        w: &CudaSlice<f32>,
4760        aq2: &CudaSlice<i8>,
4761        ad2: &CudaSlice<f32>,
4762        dst: &mut CudaSlice<f32>,
4763        t: usize,
4764        in_f: usize,
4765        out_f: usize,
4766        n_used: usize,
4767        n_expert: usize,
4768        qt: i32,
4769        rb: usize,
4770    ) -> Result<(), Box<dyn std::error::Error>> {
4771        let (inf, outf, nu, ne, rbi) = (
4772            in_f as i32,
4773            out_f as i32,
4774            n_used as i32,
4775            n_expert as i32,
4776            rb as i64,
4777        );
4778        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
4779        // eight warps, then replay the original slot-ordered FMA chain. Every
4780        // other shape retains the generic one-warp rows kernel.
4781        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4782        let f = self.func(if step_b1_w8 {
4783            "moe_down8_fma_dev_q8_rows_w8"
4784        } else {
4785            "moe_down8_fma_dev_q8_rows_g"
4786        });
4787        let cfg = LaunchConfig {
4788            grid_dim: (out_f as u32, 1, t as u32),
4789            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4790            shared_mem_bytes: 0,
4791        };
4792        let __s_b = self.gpu.stream();
4793        let mut b = __s_b.launch_builder(&f);
4794        b.arg(table)
4795            .arg(sel)
4796            .arg(w)
4797            .arg(aq2)
4798            .arg(ad2)
4799            .arg(dst)
4800            .arg(&inf)
4801            .arg(&outf)
4802            .arg(&nu)
4803            .arg(&ne)
4804            .arg(&qt)
4805            .arg(&rbi);
4806        unsafe {
4807            b.launch(cfg)?;
4808        }
4809        Ok(())
4810    }
4811
4812    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
4813    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
4814    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
4815    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4816        let (out_f, in_f) = (2048usize, 2816usize);
4817        let nblk = in_f / 32;
4818        let mut seed = 0x9E3779B97F4A7C15u64;
4819        let mut rng = move || {
4820            seed = seed
4821                .wrapping_mul(6364136223846793005)
4822                .wrapping_add(1442695040888963407);
4823            (seed >> 33) as u8
4824        };
4825        let mut w = vec![0u8; out_f * nblk * 18];
4826        for b in w.iter_mut() {
4827            *b = rng();
4828        }
4829        for r in 0..out_f {
4830            for g in 0..nblk {
4831                let off = (r * nblk + g) * 18;
4832                w[off] = 0x00;
4833                w[off + 1] = 0x2C; // sane half d
4834            }
4835        }
4836        let qplane = out_f * nblk * 16;
4837        let mut wrp = vec![0u8; w.len()];
4838        for r in 0..out_f {
4839            for g in 0..nblk {
4840                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4841                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4842                    .copy_from_slice(&src[0..2]);
4843                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4844            }
4845        }
4846        let w_d = self.htod_bytes(&w)?;
4847        let wrp_d = self.htod_bytes(&wrp)?;
4848        let mut aq = vec![0i8; m * in_f];
4849        for v in aq.iter_mut() {
4850            *v = rng() as i8;
4851        }
4852        let aq_d = self.htod_i8(&aq)?;
4853        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
4854        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
4855        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
4856        const RPB: u32 = 4;
4857        let cfg = LaunchConfig {
4858            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
4859            block_dim: (32, RPB, 1),
4860            shared_mem_bytes: 0,
4861        };
4862        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
4863        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
4864        let fb = self.func("qmatvec_q4_0_mmvq_b4");
4865        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
4866        {
4867            let __s_b = self.gpu.stream();
4868            let mut b = __s_b.launch_builder(&fb);
4869            b.arg(&w_d)
4870                .arg(&aq_d)
4871                .arg(&ad_d)
4872                .arg(&mut y0)
4873                .arg(&inf)
4874                .arg(&outf)
4875                .arg(&mi)
4876                .arg(&rb);
4877            unsafe {
4878                b.launch(cfg)?;
4879            }
4880            let __s_b = self.gpu.stream();
4881            let mut b = __s_b.launch_builder(&fr);
4882            b.arg(&wrp_d)
4883                .arg(&aq_d)
4884                .arg(&ad_d)
4885                .arg(&mut y1)
4886                .arg(&inf)
4887                .arg(&outf)
4888                .arg(&mi)
4889                .arg(&qp);
4890            unsafe {
4891                b.launch(cfg)?;
4892            }
4893        }
4894        self.gpu.stream().synchronize()?;
4895        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
4896        let nd = h0
4897            .iter()
4898            .zip(&h1)
4899            .filter(|(a, b)| a.to_bits() != b.to_bits())
4900            .count();
4901        if nd != 0 {
4902            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
4903        }
4904        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
4905            self.gpu.stream().synchronize()?;
4906            let t0 = std::time::Instant::now();
4907            for _ in 0..500 {
4908                if rp {
4909                    let __s_b = self.gpu.stream();
4910                    let mut b = __s_b.launch_builder(&fr);
4911                    b.arg(&wrp_d)
4912                        .arg(&aq_d)
4913                        .arg(&ad_d)
4914                        .arg(&mut y1)
4915                        .arg(&inf)
4916                        .arg(&outf)
4917                        .arg(&mi)
4918                        .arg(&qp);
4919                    unsafe {
4920                        b.launch(cfg)?;
4921                    }
4922                } else {
4923                    let __s_b = self.gpu.stream();
4924                    let mut b = __s_b.launch_builder(&fb);
4925                    b.arg(&w_d)
4926                        .arg(&aq_d)
4927                        .arg(&ad_d)
4928                        .arg(&mut y0)
4929                        .arg(&inf)
4930                        .arg(&outf)
4931                        .arg(&mi)
4932                        .arg(&rb);
4933                    unsafe {
4934                        b.launch(cfg)?;
4935                    }
4936                }
4937            }
4938            self.gpu.stream().synchronize()?;
4939            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
4940        };
4941        let _ = time(false)?;
4942        let _ = time(true)?; // warm
4943        Ok((time(false)?, time(true)?))
4944    }
4945
4946    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
4947    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
4948    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
4949    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
4950    pub fn build_q4_rp4(
4951        &self,
4952        t: &mut crate::model::GpuTensor,
4953    ) -> Result<(), Box<dyn std::error::Error>> {
4954        use crate::model::GpuTensor;
4955        let GpuTensor::Quant {
4956            bytes,
4957            qtype,
4958            row_bytes,
4959            ne,
4960            rp4,
4961            ..
4962        } = t
4963        else {
4964            return Ok(());
4965        };
4966        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
4967            return Ok(());
4968        }
4969        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
4970        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
4971            return Ok(());
4972        }
4973        let nblk = in_f / 32;
4974        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
4975        let f = self.func("q4_0_split_rp_build");
4976        let n = (out_f * nblk) as i32;
4977        let cfg = LaunchConfig {
4978            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
4979            block_dim: (256, 1, 1),
4980            shared_mem_bytes: 0,
4981        };
4982        let (of, nb) = (out_f as i32, nblk as i32);
4983        let _ = n;
4984        let __s_b = self.gpu.stream();
4985        let mut b = __s_b.launch_builder(&f);
4986        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
4987        unsafe {
4988            b.launch(cfg)?;
4989        }
4990        *rp4 = Some(dst);
4991        Ok(())
4992    }
4993
4994    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
4995    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
4996    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
4997    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
4998    pub fn build_q8_rp4(
4999        &self,
5000        t: &mut crate::model::GpuTensor,
5001    ) -> Result<(), Box<dyn std::error::Error>> {
5002        use crate::model::GpuTensor;
5003        let GpuTensor::Quant {
5004            bytes,
5005            qtype,
5006            row_bytes,
5007            ne,
5008            rp4,
5009            ..
5010        } = t
5011        else {
5012            return Ok(());
5013        };
5014        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5015            return Ok(());
5016        }
5017        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5018        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5019            return Ok(());
5020        }
5021        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5022        Ok(())
5023    }
5024
5025    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5026    /// mirror without a GpuTensor (same kernel the loader path above uses).
5027    pub fn build_q8_rp4_raw(
5028        &self,
5029        bytes: &CudaSlice<u8>,
5030        in_f: usize,
5031        out_f: usize,
5032    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5033        assert!(in_f % 32 == 0);
5034        let nblk = in_f / 32;
5035        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5036        let f = self.func("q8_0_split_rp_build");
5037        let cfg = LaunchConfig {
5038            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5039            block_dim: (256, 1, 1),
5040            shared_mem_bytes: 0,
5041        };
5042        let (of, nb) = (out_f as i32, nblk as i32);
5043        let __s_b = self.gpu.stream();
5044        let mut b = __s_b.launch_builder(&f);
5045        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5046        unsafe {
5047            b.launch(cfg)?;
5048        }
5049        Ok(dst)
5050    }
5051
5052    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5053    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5054    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5055    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5056    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5057    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5058    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5059    pub fn build_q4k_rp4(
5060        &self,
5061        t: &mut crate::model::GpuTensor,
5062    ) -> Result<(), Box<dyn std::error::Error>> {
5063        use crate::model::GpuTensor;
5064        let GpuTensor::Quant {
5065            bytes,
5066            qtype,
5067            row_bytes,
5068            ne,
5069            rp4,
5070            ..
5071        } = t
5072        else {
5073            return Ok(());
5074        };
5075        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5076            return Ok(());
5077        }
5078        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5079        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5080            return Ok(());
5081        }
5082        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5083        Ok(())
5084    }
5085
5086    pub fn build_q6k_rp4(
5087        &self,
5088        t: &mut crate::model::GpuTensor,
5089    ) -> Result<(), Box<dyn std::error::Error>> {
5090        use crate::model::GpuTensor;
5091        let GpuTensor::Quant {
5092            bytes,
5093            qtype,
5094            row_bytes,
5095            ne,
5096            rp4,
5097            ..
5098        } = t
5099        else {
5100            return Ok(());
5101        };
5102        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5103            return Ok(());
5104        }
5105        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5106        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5107            return Ok(());
5108        }
5109        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5110        Ok(())
5111    }
5112
5113    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5114    pub fn build_kq_rp4_raw(
5115        &self,
5116        bytes: &CudaSlice<u8>,
5117        in_f: usize,
5118        out_f: usize,
5119        qtype: i32,
5120    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5121        assert!(in_f % 256 == 0);
5122        let nsbk = in_f / 256;
5123        let (sb_bytes, kname) = match qtype {
5124            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5125            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5126            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5127        };
5128        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5129        let f = self.func(kname);
5130        let cfg = LaunchConfig {
5131            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5132            block_dim: (256, 1, 1),
5133            shared_mem_bytes: 0,
5134        };
5135        let (of, nb) = (out_f as i32, nsbk as i32);
5136        let __s_b = self.gpu.stream();
5137        let mut b = __s_b.launch_builder(&f);
5138        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5139        unsafe {
5140            b.launch(cfg)?;
5141        }
5142        Ok(dst)
5143    }
5144
5145    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5146    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5147    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5148    pub fn kqrp_enabled() -> bool {
5149        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5150        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5151            Ok("0") => false,
5152            Ok(_) => true,
5153            Err(_) => cfg!(memra_hopper_mma),
5154        })
5155    }
5156
5157    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5158    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5159    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5160    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5161    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5162    pub fn build_q4_rp_swap(
5163        &self,
5164        t: &mut crate::model::GpuTensor,
5165    ) -> Result<bool, Box<dyn std::error::Error>> {
5166        use crate::model::GpuTensor;
5167        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5168        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5169        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5170        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5171        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5172        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5173        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5174        // this fn's OWN builder serves may ever be swapped; everything else refuses
5175        // here, regardless of walk ordering.
5176        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5177            return Ok(false);
5178        }
5179        self.build_q4_rp4(t)?;
5180        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5181        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5182            return Ok(false);
5183        };
5184        match rp4.take() {
5185            Some(split) => {
5186                *bytes = split; // the GGUF-layout buffer drops here
5187                *rp = true;
5188                Ok(true)
5189            }
5190            None => Ok(false),
5191        }
5192    }
5193
5194    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5195    pub fn q4rp_enabled() -> bool {
5196        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5197        *ON.get_or_init(|| {
5198            std::env::var("MEMRA_Q4RP")
5199                .map(|v| v != "0")
5200                .unwrap_or(true)
5201        })
5202    }
5203
5204    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5205    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5206    pub fn copy_rows_strided(
5207        &self,
5208        src: &CudaSlice<f32>,
5209        dst: &mut CudaSlice<f32>,
5210        row_elems: usize,
5211        n_rows: usize,
5212        src_stride: usize,
5213        src_off: usize,
5214    ) -> Result<(), Box<dyn std::error::Error>> {
5215        let f = self.func("copy_rows_strided_f32");
5216        let cfg = LaunchConfig {
5217            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5218            block_dim: (256, 1, 1),
5219            shared_mem_bytes: 0,
5220        };
5221        let (re, nr) = (row_elems as i32, n_rows as i32);
5222        let (st, off) = (src_stride as i64, src_off as i64);
5223        let __s_b = self.gpu.stream();
5224        let mut b = __s_b.launch_builder(&f);
5225        b.arg(src)
5226            .arg(&mut *dst)
5227            .arg(&re)
5228            .arg(&nr)
5229            .arg(&st)
5230            .arg(&off);
5231        unsafe {
5232            b.launch(cfg)?;
5233        }
5234        Ok(())
5235    }
5236
5237    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5238    pub fn u32_set_k(
5239        &self,
5240        dst: &mut CudaSlice<u32>,
5241        v: u32,
5242        idx: usize,
5243    ) -> Result<(), Box<dyn std::error::Error>> {
5244        let f = self.func("u32_set_k");
5245        let cfg = LaunchConfig {
5246            grid_dim: (1, 1, 1),
5247            block_dim: (1, 1, 1),
5248            shared_mem_bytes: 0,
5249        };
5250        let ii = idx as i32;
5251        let __s_b = self.gpu.stream();
5252        let mut b = __s_b.launch_builder(&f);
5253        b.arg(dst).arg(&v).arg(&ii);
5254        unsafe {
5255            b.launch(cfg)?;
5256        }
5257        Ok(())
5258    }
5259
5260    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5261    pub fn i32_add_k(
5262        &self,
5263        d: &mut CudaSlice<i32>,
5264        v: i32,
5265    ) -> Result<(), Box<dyn std::error::Error>> {
5266        let f = self.func("i32_add_k");
5267        let cfg = LaunchConfig {
5268            grid_dim: (1, 1, 1),
5269            block_dim: (32, 1, 1),
5270            shared_mem_bytes: 0,
5271        };
5272        let __s_b = self.gpu.stream();
5273        let mut b = __s_b.launch_builder(&f);
5274        b.arg(d).arg(&v);
5275        unsafe {
5276            b.launch(cfg)?;
5277        }
5278        Ok(())
5279    }
5280
5281    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5282    pub fn i32_iota_from(
5283        &self,
5284        ctr: &CudaSlice<i32>,
5285        dst: &mut CudaSlice<i32>,
5286        n: usize,
5287    ) -> Result<(), Box<dyn std::error::Error>> {
5288        let f = self.func("i32_iota_from");
5289        let cfg = LaunchConfig::for_num_elems(n as u32);
5290        let ni = n as i32;
5291        let __s_b = self.gpu.stream();
5292        let mut b = __s_b.launch_builder(&f);
5293        b.arg(ctr).arg(dst).arg(&ni);
5294        unsafe {
5295            b.launch(cfg)?;
5296        }
5297        Ok(())
5298    }
5299
5300    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5301    pub fn u32_map_k(
5302        &self,
5303        buf: &mut CudaSlice<u32>,
5304        map: &CudaSlice<u32>,
5305        idx: usize,
5306    ) -> Result<(), Box<dyn std::error::Error>> {
5307        let f = self.func("u32_map_k");
5308        let cfg = LaunchConfig {
5309            grid_dim: (1, 1, 1),
5310            block_dim: (1, 1, 1),
5311            shared_mem_bytes: 0,
5312        };
5313        let ii = idx as i32;
5314        let __s_b = self.gpu.stream();
5315        let mut b = __s_b.launch_builder(&f);
5316        b.arg(buf).arg(map).arg(&ii);
5317        unsafe {
5318            b.launch(cfg)?;
5319        }
5320        Ok(())
5321    }
5322
5323    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5324    #[allow(clippy::too_many_arguments)]
5325    pub fn u32_pack2(
5326        &self,
5327        a: &CudaSlice<u32>,
5328        off_a: usize,
5329        n1: usize,
5330        b_in: &CudaSlice<u32>,
5331        n2: usize,
5332        out: &mut CudaSlice<u32>,
5333    ) -> Result<(), Box<dyn std::error::Error>> {
5334        let f = self.func("u32_pack2");
5335        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5336        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5337        let __s_b = self.gpu.stream();
5338        let mut b = __s_b.launch_builder(&f);
5339        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5340        unsafe {
5341            b.launch(cfg)?;
5342        }
5343        Ok(())
5344    }
5345
5346    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5347    pub fn moe_w_exscale(
5348        &self,
5349        w: &mut CudaSlice<f32>,
5350        sel: &CudaSlice<i32>,
5351        s: &CudaSlice<f32>,
5352        n: usize,
5353    ) -> Result<(), Box<dyn std::error::Error>> {
5354        let f = self.func("moe_w_exscale");
5355        let cfg = LaunchConfig::for_num_elems(n as u32);
5356        let ni = n as i32;
5357        let __s_b = self.gpu.stream();
5358        let mut b = __s_b.launch_builder(&f);
5359        b.arg(w).arg(sel).arg(s).arg(&ni);
5360        unsafe {
5361            b.launch(cfg)?;
5362        }
5363        Ok(())
5364    }
5365
5366    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
5367    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
5368    pub fn moe_w_scale_by_expert(
5369        &self,
5370        w: &mut CudaSlice<f32>,
5371        sel: &CudaSlice<i32>,
5372        macros: &CudaSlice<f32>,
5373        n_expert: usize,
5374        n: usize,
5375    ) -> Result<(), Box<dyn std::error::Error>> {
5376        let f = self.func("moe_w_scale_by_expert");
5377        let cfg = LaunchConfig {
5378            grid_dim: (n.div_ceil(64) as u32, 1, 1),
5379            block_dim: (64, 1, 1),
5380            shared_mem_bytes: 0,
5381        };
5382        let (ne, nn) = (n_expert as i32, n as i32);
5383        let __s_b = self.gpu.stream();
5384        let mut b = __s_b.launch_builder(&f);
5385        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5386        unsafe {
5387            b.launch(cfg)?;
5388        }
5389        Ok(())
5390    }
5391
5392    pub fn moe_gate_up_silu8_dev_q8(
5393        &self,
5394        table: &CudaSlice<u64>,
5395        sel: &cudarc::driver::CudaView<i32>,
5396        aq: &CudaSlice<i8>,
5397        ad: &CudaSlice<f32>,
5398        in_f: usize,
5399        n_ff: usize,
5400        n_used: usize,
5401        n_expert: usize,
5402        qt_g: i32,
5403        qt_u: i32,
5404        rb_g: usize,
5405        rb_u: usize,
5406        macros: &CudaSlice<f32>,
5407    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5408        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5409        let (mode, wpb) = GU.get_or_init(|| {
5410            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5411            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5412                .ok()
5413                .and_then(|v| v.parse().ok())
5414                .unwrap_or(4u32)
5415                .clamp(1, 16);
5416            (mode, wpb)
5417        });
5418        let (mode, wpb) = (mode.as_str(), *wpb);
5419        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5420        let (inf, nff, ne, rbg, rbu) = (
5421            in_f as i32,
5422            n_ff as i32,
5423            n_expert as i32,
5424            rb_g as i64,
5425            rb_u as i64,
5426        );
5427        let (f, cfg) = match mode {
5428            "1" | "2" | "4" => {
5429                let rpw: u32 = mode.parse().unwrap();
5430                let f = self.func(match rpw {
5431                    1 => "moe_gate_up_silu8_dev_q8_r1",
5432                    2 => "moe_gate_up_silu8_dev_q8_r2",
5433                    _ => "moe_gate_up_silu8_dev_q8_r4",
5434                });
5435                let rows_per_block = (rpw * wpb) as usize;
5436                let gx = n_ff.div_ceil(rows_per_block) as u32;
5437                (
5438                    f,
5439                    LaunchConfig {
5440                        grid_dim: (gx, n_used as u32, 1),
5441                        block_dim: (32, wpb, 1),
5442                        shared_mem_bytes: 0,
5443                    },
5444                )
5445            }
5446            "j8" if n_used <= 32 => (
5447                self.func("moe_gate_up_silu8_dev_q8_j8"),
5448                LaunchConfig {
5449                    grid_dim: (n_ff as u32, 1, 1),
5450                    block_dim: (32, n_used as u32, 1),
5451                    shared_mem_bytes: 0,
5452                },
5453            ),
5454            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
5455            "vsm2" => {
5456                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5457                let sh = (rb_g + rb_u) as u32;
5458                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5459                f.set_attribute(
5460                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5461                    sh as i32,
5462                )?;
5463                (
5464                    f,
5465                    LaunchConfig {
5466                        grid_dim: (n_ff as u32, n_used as u32, 1),
5467                        block_dim: (32, 1, 1),
5468                        shared_mem_bytes: sh,
5469                    },
5470                )
5471            }
5472            "vsm" => {
5473                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5474                let sh = (rb_g + rb_u) as u32;
5475                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5476                f.set_attribute(
5477                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5478                    sh as i32,
5479                )?;
5480                (
5481                    f,
5482                    LaunchConfig {
5483                        grid_dim: (n_ff as u32, n_used as u32, 1),
5484                        block_dim: (32, 1, 1),
5485                        shared_mem_bytes: sh,
5486                    },
5487                )
5488            }
5489            "sg" => (
5490                self.func("moe_gate_up_silu8_dev_q8_sg"),
5491                LaunchConfig {
5492                    grid_dim: (n_ff as u32, n_used as u32, 1),
5493                    block_dim: (32, 1, 1),
5494                    shared_mem_bytes: 0,
5495                },
5496            ),
5497            "j8sg" if n_used <= 32 => (
5498                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5499                LaunchConfig {
5500                    grid_dim: (n_ff as u32, 1, 1),
5501                    block_dim: (32, n_used as u32, 1),
5502                    shared_mem_bytes: 0,
5503                },
5504            ),
5505            "u64" if in_f == 2048 => (
5506                self.func("moe_gate_up_silu8_dev_q8_u64"),
5507                LaunchConfig {
5508                    grid_dim: (n_ff as u32, n_used as u32, 1),
5509                    block_dim: (32, 1, 1),
5510                    shared_mem_bytes: 0,
5511                },
5512            ),
5513            "gs4" if in_f == 2048 => (
5514                self.func("moe_gate_up_silu8_dev_q8_gs4"),
5515                LaunchConfig {
5516                    grid_dim: (n_ff as u32, n_used as u32, 1),
5517                    block_dim: (32, 4, 1),
5518                    shared_mem_bytes: 0,
5519                },
5520            ),
5521            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
5522            "v" | "" => (
5523                self.func("moe_gate_up_silu8_dev_q8_v"),
5524                LaunchConfig {
5525                    grid_dim: (n_ff as u32, n_used as u32, 1),
5526                    block_dim: (32, 1, 1),
5527                    shared_mem_bytes: 0,
5528                },
5529            ),
5530            "s2" => (
5531                self.func("moe_gate_up_silu8_dev_q8_s2"),
5532                LaunchConfig {
5533                    grid_dim: (n_ff as u32, n_used as u32, 1),
5534                    block_dim: (32, 2, 1),
5535                    shared_mem_bytes: 0,
5536                },
5537            ),
5538            "s2z" => {
5539                let rz = wpb.min(16); // s2z smem tile is [16][2]
5540                (
5541                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
5542                    LaunchConfig {
5543                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5544                        block_dim: (32, 2, rz),
5545                        shared_mem_bytes: 0,
5546                    },
5547                )
5548            }
5549            _ => (
5550                self.func("moe_gate_up_silu8_dev_q8"),
5551                LaunchConfig {
5552                    grid_dim: (n_ff as u32, n_used as u32, 1),
5553                    block_dim: (32, 1, 1),
5554                    shared_mem_bytes: 0,
5555                },
5556            ),
5557        };
5558        let __s_b = self.gpu.stream();
5559        let mut b = __s_b.launch_builder(&f);
5560        b.arg(table)
5561            .arg(sel)
5562            .arg(aq)
5563            .arg(ad)
5564            .arg(&mut act)
5565            .arg(&inf)
5566            .arg(&nff)
5567            .arg(&ne)
5568            .arg(&qt_g)
5569            .arg(&qt_u)
5570            .arg(&rbg)
5571            .arg(&rbu)
5572            .arg(macros);
5573        unsafe {
5574            b.launch(cfg)?;
5575        }
5576        Ok(act)
5577    }
5578
5579    #[allow(clippy::too_many_arguments)]
5580    pub fn moe_down8_fma_dev_q8(
5581        &self,
5582        table: &CudaSlice<u64>,
5583        sel: &cudarc::driver::CudaView<i32>,
5584        w: &cudarc::driver::CudaView<f32>,
5585        aq2: &CudaSlice<i8>,
5586        ad2: &CudaSlice<f32>,
5587        dst: &mut cudarc::driver::CudaViewMut<f32>,
5588        in_f: usize,
5589        out_f: usize,
5590        n_used: usize,
5591        n_expert: usize,
5592        qt: i32,
5593        rb: usize,
5594    ) -> Result<(), Box<dyn std::error::Error>> {
5595        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5596        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5597        let (inf, outf, nu, ne, rbi) = (
5598            in_f as i32,
5599            out_f as i32,
5600            n_used as i32,
5601            n_expert as i32,
5602            rb as i64,
5603        );
5604        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
5605        // the h2 twins are nsb==16 (in_f==512) shape-gated.
5606        let (f, cfg) = match mode.as_str() {
5607            m @ ("1" | "2" | "4") if n_used <= 8 => {
5608                let rpw: usize = m.parse().unwrap();
5609                let f = self.func(match rpw {
5610                    1 => "moe_down8_fma_dev_q8_w8r1",
5611                    2 => "moe_down8_fma_dev_q8_w8r2",
5612                    _ => "moe_down8_fma_dev_q8_w8r4",
5613                });
5614                (
5615                    f,
5616                    LaunchConfig {
5617                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5618                        block_dim: (32, n_used as u32, 1),
5619                        shared_mem_bytes: 0,
5620                    },
5621                )
5622            }
5623            "h2" if in_f == 512 => (
5624                self.func("moe_down8_fma_dev_q8_h2"),
5625                LaunchConfig {
5626                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5627                    block_dim: (32, 1, 1),
5628                    shared_mem_bytes: 0,
5629                },
5630            ),
5631            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
5632            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
5633            "" if in_f == 704 && n_used <= 8 => (
5634                self.func("moe_down8_fma_dev_q8_w8r2"),
5635                LaunchConfig {
5636                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5637                    block_dim: (32, n_used as u32, 1),
5638                    shared_mem_bytes: 0,
5639                },
5640            ),
5641            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
5642            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
5643            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
5644            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5645                self.func("moe_down8_fma_dev_q8_w8h2v"),
5646                LaunchConfig {
5647                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5648                    block_dim: (32, n_used as u32, 1),
5649                    shared_mem_bytes: 0,
5650                },
5651            ),
5652            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5653                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5654                LaunchConfig {
5655                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5656                    block_dim: (32, n_used as u32, 1),
5657                    shared_mem_bytes: 0,
5658                },
5659            ),
5660            "w8h2r2" if in_f == 512 && n_used <= 8 => (
5661                self.func("moe_down8_fma_dev_q8_w8h2r2"),
5662                LaunchConfig {
5663                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5664                    block_dim: (32, n_used as u32, 1),
5665                    shared_mem_bytes: 0,
5666                },
5667            ),
5668            "w8h2" if in_f == 512 && n_used <= 8 => (
5669                self.func("moe_down8_fma_dev_q8_w8h2"),
5670                LaunchConfig {
5671                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5672                    block_dim: (32, n_used as u32, 1),
5673                    shared_mem_bytes: 0,
5674                },
5675            ),
5676            _ => (
5677                self.func("moe_down8_fma_dev_q8"),
5678                LaunchConfig {
5679                    grid_dim: (out_f as u32, 1, 1),
5680                    block_dim: (32, 1, 1),
5681                    shared_mem_bytes: 0,
5682                },
5683            ),
5684        };
5685        let __s_b = self.gpu.stream();
5686        let mut b = __s_b.launch_builder(&f);
5687        b.arg(table)
5688            .arg(sel)
5689            .arg(w)
5690            .arg(aq2)
5691            .arg(ad2)
5692            .arg(dst)
5693            .arg(&inf)
5694            .arg(&outf)
5695            .arg(&nu)
5696            .arg(&ne)
5697            .arg(&qt)
5698            .arg(&rbi);
5699        unsafe {
5700            b.launch(cfg)?;
5701        }
5702        Ok(())
5703    }
5704
5705    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
5706    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
5707    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
5708    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
5709    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
5710    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
5711    #[allow(clippy::too_many_arguments)]
5712    pub fn moe_gate_up_silu8_dev_q8_rows(
5713        &self,
5714        table: &CudaSlice<u64>,
5715        sel: &CudaSlice<i32>,
5716        aq: &CudaSlice<i8>,
5717        ad: &CudaSlice<f32>,
5718        t: usize,
5719        in_f: usize,
5720        n_ff: usize,
5721        n_used: usize,
5722        n_expert: usize,
5723        qt_g: i32,
5724        qt_u: i32,
5725        rb_g: usize,
5726        rb_u: usize,
5727        macros: &CudaSlice<f32>,
5728    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5729        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5730        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5731        let cfg = LaunchConfig {
5732            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5733            block_dim: (32, 1, 1),
5734            shared_mem_bytes: 0,
5735        };
5736        let (inf, nff, ne, nu, rbg, rbu) = (
5737            in_f as i32,
5738            n_ff as i32,
5739            n_expert as i32,
5740            n_used as i32,
5741            rb_g as i64,
5742            rb_u as i64,
5743        );
5744        let __s_b = self.gpu.stream();
5745        let mut b = __s_b.launch_builder(&f);
5746        b.arg(table)
5747            .arg(sel)
5748            .arg(aq)
5749            .arg(ad)
5750            .arg(&mut act)
5751            .arg(&inf)
5752            .arg(&nff)
5753            .arg(&ne)
5754            .arg(&qt_g)
5755            .arg(&qt_u)
5756            .arg(&rbg)
5757            .arg(&rbu)
5758            .arg(&nu)
5759            .arg(macros);
5760        unsafe {
5761            b.launch(cfg)?;
5762        }
5763        Ok(act)
5764    }
5765
5766    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
5767    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
5768    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
5769    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
5770    #[allow(clippy::too_many_arguments)]
5771    pub fn moe_down8_fma_dev_q8_rows(
5772        &self,
5773        table: &CudaSlice<u64>,
5774        sel: &CudaSlice<i32>,
5775        w: &CudaSlice<f32>,
5776        aq2: &CudaSlice<i8>,
5777        ad2: &CudaSlice<f32>,
5778        dst: &mut CudaSlice<f32>,
5779        t: usize,
5780        in_f: usize,
5781        out_f: usize,
5782        n_used: usize,
5783        n_expert: usize,
5784        qt: i32,
5785        rb: usize,
5786    ) -> Result<(), Box<dyn std::error::Error>> {
5787        assert!(
5788            in_f == 512 && n_used <= 8,
5789            "down rows twin is w8h2v shape-gated"
5790        );
5791        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5792        let cfg = LaunchConfig {
5793            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5794            block_dim: (32, n_used as u32, 1),
5795            shared_mem_bytes: 0,
5796        };
5797        let (inf, outf, nu, ne, rbi) = (
5798            in_f as i32,
5799            out_f as i32,
5800            n_used as i32,
5801            n_expert as i32,
5802            rb as i64,
5803        );
5804        let __s_b = self.gpu.stream();
5805        let mut b = __s_b.launch_builder(&f);
5806        b.arg(table)
5807            .arg(sel)
5808            .arg(w)
5809            .arg(aq2)
5810            .arg(ad2)
5811            .arg(dst)
5812            .arg(&inf)
5813            .arg(&outf)
5814            .arg(&nu)
5815            .arg(&ne)
5816            .arg(&qt)
5817            .arg(&rbi);
5818        unsafe {
5819            b.launch(cfg)?;
5820        }
5821        Ok(())
5822    }
5823
5824    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
5825    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
5826    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
5827    #[allow(clippy::too_many_arguments)]
5828    pub fn moe_gate_up_silu8_dev_q8_csr(
5829        &self,
5830        table: &CudaSlice<u64>,
5831        sel: &CudaSlice<i32>,
5832        aq: &CudaSlice<i8>,
5833        ad: &CudaSlice<f32>,
5834        n_pairs: usize,
5835        in_f: usize,
5836        n_ff: usize,
5837        n_used: usize,
5838        n_expert: usize,
5839        qt_g: i32,
5840        qt_u: i32,
5841        rb_g: usize,
5842        rb_u: usize,
5843    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5844        let f = self.func("moe_gate_up_silu8_dev_q8_csr_iq4");
5845        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5846        let cfg = LaunchConfig {
5847            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5848            block_dim: (32, 1, 1),
5849            shared_mem_bytes: 0,
5850        };
5851        let (inf, nff, ne, nu, npi, rbg, rbu) = (
5852            in_f as i32,
5853            n_ff as i32,
5854            n_expert as i32,
5855            n_used as i32,
5856            n_pairs as i32,
5857            rb_g as i64,
5858            rb_u as i64,
5859        );
5860        let __s_b = self.gpu.stream();
5861        let mut b = __s_b.launch_builder(&f);
5862        b.arg(table)
5863            .arg(sel)
5864            .arg(aq)
5865            .arg(ad)
5866            .arg(&mut act)
5867            .arg(&inf)
5868            .arg(&nff)
5869            .arg(&ne)
5870            .arg(&qt_g)
5871            .arg(&qt_u)
5872            .arg(&rbg)
5873            .arg(&rbu)
5874            .arg(&nu)
5875            .arg(&npi);
5876        unsafe {
5877            b.launch(cfg)?;
5878        }
5879        Ok(act)
5880    }
5881
5882    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
5883    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
5884    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
5885    #[allow(clippy::too_many_arguments)]
5886    pub fn moe_down8_fma_dev_q8_variant(
5887        &self,
5888        variant: &str,
5889        table: &CudaSlice<u64>,
5890        sel: &cudarc::driver::CudaView<i32>,
5891        w: &cudarc::driver::CudaView<f32>,
5892        aq2: &CudaSlice<i8>,
5893        ad2: &CudaSlice<f32>,
5894        dst: &mut cudarc::driver::CudaViewMut<f32>,
5895        in_f: usize,
5896        out_f: usize,
5897        n_used: usize,
5898        n_expert: usize,
5899        qt: i32,
5900        rb: usize,
5901    ) -> Result<(), Box<dyn std::error::Error>> {
5902        let (inf, outf, nu, ne, rbi) = (
5903            in_f as i32,
5904            out_f as i32,
5905            n_used as i32,
5906            n_expert as i32,
5907            rb as i64,
5908        );
5909        let (f, cfg) = match variant {
5910            "w8h2" | "w8h2v" => (
5911                self.func(if variant == "w8h2" {
5912                    "moe_down8_fma_dev_q8_w8h2"
5913                } else {
5914                    "moe_down8_fma_dev_q8_w8h2v"
5915                }),
5916                LaunchConfig {
5917                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5918                    block_dim: (32, n_used as u32, 1),
5919                    shared_mem_bytes: 0,
5920                },
5921            ),
5922            "w8h2r2" | "w8h2r2v" => (
5923                self.func(if variant == "w8h2r2" {
5924                    "moe_down8_fma_dev_q8_w8h2r2"
5925                } else {
5926                    "moe_down8_fma_dev_q8_w8h2r2v"
5927                }),
5928                LaunchConfig {
5929                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5930                    block_dim: (32, n_used as u32, 1),
5931                    shared_mem_bytes: 0,
5932                },
5933            ),
5934            _ => (
5935                self.func("moe_down8_fma_dev_q8"),
5936                LaunchConfig {
5937                    grid_dim: (out_f as u32, 1, 1),
5938                    block_dim: (32, 1, 1),
5939                    shared_mem_bytes: 0,
5940                },
5941            ),
5942        };
5943        let __s_b = self.gpu.stream();
5944        let mut b = __s_b.launch_builder(&f);
5945        b.arg(table)
5946            .arg(sel)
5947            .arg(w)
5948            .arg(aq2)
5949            .arg(ad2)
5950            .arg(dst)
5951            .arg(&inf)
5952            .arg(&outf)
5953            .arg(&nu)
5954            .arg(&ne)
5955            .arg(&qt)
5956            .arg(&rbi);
5957        unsafe {
5958            b.launch(cfg)?;
5959        }
5960        Ok(())
5961    }
5962
5963    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
5964    #[allow(clippy::too_many_arguments)]
5965    pub fn moe_gate_up_silu8_dev_q8_variant(
5966        &self,
5967        variant: &str,
5968        table: &CudaSlice<u64>,
5969        sel: &cudarc::driver::CudaView<i32>,
5970        aq: &CudaSlice<i8>,
5971        ad: &CudaSlice<f32>,
5972        in_f: usize,
5973        n_ff: usize,
5974        n_used: usize,
5975        n_expert: usize,
5976        qt_g: i32,
5977        qt_u: i32,
5978        rb_g: usize,
5979        rb_u: usize,
5980    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5981        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5982        let (inf, nff, ne, rbg, rbu) = (
5983            in_f as i32,
5984            n_ff as i32,
5985            n_expert as i32,
5986            rb_g as i64,
5987            rb_u as i64,
5988        );
5989        let f = self.func(if variant == "v" {
5990            "moe_gate_up_silu8_dev_q8_v"
5991        } else {
5992            "moe_gate_up_silu8_dev_q8"
5993        });
5994        let cfg = LaunchConfig {
5995            grid_dim: (n_ff as u32, n_used as u32, 1),
5996            block_dim: (32, 1, 1),
5997            shared_mem_bytes: 0,
5998        };
5999        let __s_b = self.gpu.stream();
6000        let mut b = __s_b.launch_builder(&f);
6001        b.arg(table)
6002            .arg(sel)
6003            .arg(aq)
6004            .arg(ad)
6005            .arg(&mut act)
6006            .arg(&inf)
6007            .arg(&nff)
6008            .arg(&ne)
6009            .arg(&qt_g)
6010            .arg(&qt_u)
6011            .arg(&rbg)
6012            .arg(&rbu);
6013        unsafe {
6014            b.launch(cfg)?;
6015        }
6016        Ok(act)
6017    }
6018
6019    pub fn moe_gate_up_silu8_dev(
6020        &self,
6021        table: &CudaSlice<u64>,
6022        sel: &cudarc::driver::CudaView<i32>,
6023        x: &cudarc::driver::CudaView<f32>,
6024        in_f: usize,
6025        n_ff: usize,
6026        n_used: usize,
6027        n_expert: usize,
6028        qt_g: i32,
6029        qt_u: i32,
6030        rb_g: usize,
6031        rb_u: usize,
6032        macros: &CudaSlice<f32>,
6033    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6034        let f = self.func("moe_gate_up_silu8_dev");
6035        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6036        let cfg = LaunchConfig {
6037            grid_dim: (n_ff as u32, n_used as u32, 1),
6038            block_dim: (256, 1, 1),
6039            shared_mem_bytes: 0,
6040        };
6041        let (inf, nff, ne, rbg, rbu) = (
6042            in_f as i32,
6043            n_ff as i32,
6044            n_expert as i32,
6045            rb_g as i64,
6046            rb_u as i64,
6047        );
6048        let __s_b = self.gpu.stream();
6049        let mut b = __s_b.launch_builder(&f);
6050        b.arg(table)
6051            .arg(sel)
6052            .arg(x)
6053            .arg(&mut act)
6054            .arg(&inf)
6055            .arg(&nff)
6056            .arg(&ne)
6057            .arg(&qt_g)
6058            .arg(&qt_u)
6059            .arg(&rbg)
6060            .arg(&rbu)
6061            .arg(macros);
6062        unsafe {
6063            b.launch(cfg)?;
6064        }
6065        Ok(act)
6066    }
6067
6068    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6069    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6070    #[allow(clippy::too_many_arguments)]
6071    pub fn moe_down8_fma_dev(
6072        &self,
6073        table: &CudaSlice<u64>,
6074        sel: &cudarc::driver::CudaView<i32>,
6075        w: &cudarc::driver::CudaView<f32>,
6076        act: &CudaSlice<f32>,
6077        dst: &mut cudarc::driver::CudaViewMut<f32>,
6078        in_f: usize,
6079        out_f: usize,
6080        n_used: usize,
6081        n_expert: usize,
6082        qt: i32,
6083        rb: usize,
6084    ) -> Result<(), Box<dyn std::error::Error>> {
6085        let f = self.func("moe_down8_fma_dev");
6086        let cfg = LaunchConfig {
6087            grid_dim: (out_f as u32, 1, 1),
6088            block_dim: (256, 1, 1),
6089            shared_mem_bytes: 0,
6090        };
6091        let (inf, outf, nu, ne, rbv) = (
6092            in_f as i32,
6093            out_f as i32,
6094            n_used as i32,
6095            n_expert as i32,
6096            rb as i64,
6097        );
6098        let __s_b = self.gpu.stream();
6099        let mut b = __s_b.launch_builder(&f);
6100        b.arg(table)
6101            .arg(sel)
6102            .arg(w)
6103            .arg(act)
6104            .arg(dst)
6105            .arg(&inf)
6106            .arg(&outf)
6107            .arg(&nu)
6108            .arg(&ne)
6109            .arg(&qt)
6110            .arg(&rbv);
6111        unsafe {
6112            b.launch(cfg)?;
6113        }
6114        Ok(())
6115    }
6116
6117    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6118    pub fn axpy_into(
6119        &self,
6120        src: &CudaSlice<f32>,
6121        alpha: f32,
6122        dst: &mut cudarc::driver::CudaViewMut<f32>,
6123        n: usize,
6124    ) -> Result<(), Box<dyn std::error::Error>> {
6125        let f = self.func("axpy_f32");
6126        let cfg = LaunchConfig::for_num_elems(n as u32);
6127        let (a, ni) = (alpha, n as i32);
6128        let __s_b = self.gpu.stream();
6129        let mut b = __s_b.launch_builder(&f);
6130        b.arg(src).arg(dst).arg(&a).arg(&ni);
6131        unsafe {
6132            b.launch(cfg)?;
6133        }
6134        Ok(())
6135    }
6136
6137    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6138    pub fn add_scaled_rows(
6139        &self,
6140        src: &CudaSlice<f32>,
6141        scale: &CudaSlice<f32>,
6142        dst: &mut CudaSlice<f32>,
6143        ncols: usize,
6144        nrows: usize,
6145    ) -> Result<(), Box<dyn std::error::Error>> {
6146        let f = self.func("add_scaled_rows_f32");
6147        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6148        let (nc, nr) = (ncols as i32, nrows as i32);
6149        let __s_b = self.gpu.stream();
6150        let mut b = __s_b.launch_builder(&f);
6151        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6152        unsafe {
6153            b.launch(cfg)?;
6154        }
6155        Ok(())
6156    }
6157
6158    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6159
6160    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6161    pub fn gather_rows(
6162        &self,
6163        src: &CudaSlice<f32>,
6164        idx: &CudaSlice<i32>,
6165        dst: &mut CudaSlice<f32>,
6166        ncols: usize,
6167        m_e: usize,
6168    ) -> Result<(), Box<dyn std::error::Error>> {
6169        let f = self.func("gather_rows_f32");
6170        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6171        let (nc, me) = (ncols as i32, m_e as i32);
6172        let __s_b = self.gpu.stream();
6173        let mut b = __s_b.launch_builder(&f);
6174        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6175        unsafe {
6176            b.launch(cfg)?;
6177        }
6178        Ok(())
6179    }
6180
6181    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6182    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6183    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6184    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6185    pub fn scatter_slot(
6186        &self,
6187        src: &CudaSlice<f32>,
6188        tok_idx: &CudaSlice<i32>,
6189        slot_idx: &CudaSlice<i32>,
6190        weight: &CudaSlice<f32>,
6191        dst: &mut CudaSlice<f32>,
6192        wbuf: &mut CudaSlice<f32>,
6193        ncols: usize,
6194        n_used: usize,
6195        m_e: usize,
6196    ) -> Result<(), Box<dyn std::error::Error>> {
6197        let f = self.func("scatter_add_slot_f32");
6198        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6199        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6200        let __s_b = self.gpu.stream();
6201        let mut b = __s_b.launch_builder(&f);
6202        b.arg(src)
6203            .arg(tok_idx)
6204            .arg(slot_idx)
6205            .arg(weight)
6206            .arg(dst)
6207            .arg(wbuf)
6208            .arg(&nc)
6209            .arg(&nu)
6210            .arg(&me);
6211        unsafe {
6212            b.launch(cfg)?;
6213        }
6214        Ok(())
6215    }
6216
6217    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6218    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6219    /// Uses FMA for bit-identity with the sequential axpy path.
6220    pub fn reduce_slots(
6221        &self,
6222        slots: &CudaSlice<f32>,
6223        wbuf: &CudaSlice<f32>,
6224        dst: &mut CudaSlice<f32>,
6225        ncols: usize,
6226        n_used: usize,
6227        t: usize,
6228    ) -> Result<(), Box<dyn std::error::Error>> {
6229        let f = self.func("reduce_slots_f32");
6230        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6231        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6232        let __s_b = self.gpu.stream();
6233        let mut b = __s_b.launch_builder(&f);
6234        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6235        unsafe {
6236            b.launch(cfg)?;
6237        }
6238        Ok(())
6239    }
6240
6241    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6242    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6243    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6244    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6245    /// GPU time, ~half of it redundant re-quantization of the same row.
6246    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6247    pub fn quantize_q8_1_view(
6248        &self,
6249        x: &cudarc::driver::CudaView<f32>,
6250        m: usize,
6251        in_f: usize,
6252    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6253        let f = self.func("quantize_q8_1");
6254        let nblk = in_f / 32;
6255        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6256        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6257        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6258        let (inf, mi) = (in_f as i32, m as i32);
6259        let __s_b = self.gpu.stream();
6260        let mut b = __s_b.launch_builder(&f);
6261        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6262        unsafe {
6263            b.launch(cfg)?;
6264        }
6265        Ok((q, d))
6266    }
6267
6268    pub fn quantize_q8_1(
6269        &self,
6270        x: &CudaSlice<f32>,
6271        m: usize,
6272        in_f: usize,
6273    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6274        let nblk = in_f / 32;
6275        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6276        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6277        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6278        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6279        let (inf, mi) = (in_f as i32, m as i32);
6280        if Self::pdl_on() && Self::pdl_wb_on() {
6281            {
6282                use cudarc::driver::{DevicePtr, DevicePtrMut};
6283                let s = &self.gpu.stream();
6284                let (px, _g0) = x.device_ptr(s);
6285                let (pq, _g1) = q.device_ptr_mut(s);
6286                let (pd, _g2) = d.device_ptr_mut(s);
6287                let mut ps = [
6288                    &px as *const _ as *mut std::ffi::c_void,
6289                    &pq as *const _ as *mut _,
6290                    &pd as *const _ as *mut _,
6291                    &inf as *const _ as *mut _,
6292                    &mi as *const _ as *mut _,
6293                ];
6294                unsafe {
6295                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6296                }
6297            }
6298            return Ok((q, d));
6299        }
6300        let f = self.func("quantize_q8_1");
6301        let __s_b = self.gpu.stream();
6302        let mut b = __s_b.launch_builder(&f);
6303        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6304        unsafe {
6305            b.launch(cfg)?;
6306        }
6307        Ok((q, d))
6308    }
6309
6310    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
6311    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
6312    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
6313    pub fn quantize_fp4_act(
6314        &self,
6315        x: &CudaSlice<f32>,
6316        m: usize,
6317        in_f: usize,
6318    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6319        let f = self.func("quantize_fp4_act");
6320        let nb16 = in_f / 16;
6321        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
6322        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
6323        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6324        let (inf, mi) = (in_f as i32, m as i32);
6325        let __s_b = self.gpu.stream();
6326        let mut b = __s_b.launch_builder(&f);
6327        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6328        unsafe {
6329            b.launch(cfg)?;
6330        }
6331        Ok((aq4, ad4))
6332    }
6333
6334    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
6335    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
6336    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
6337    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
6338    pub fn qmatvec_gemm_nvfp4_fp4(
6339        &self,
6340        bytes: &CudaSlice<u8>,
6341        x: &CudaSlice<f32>,
6342        m: usize,
6343        in_f: usize,
6344        out_f: usize,
6345        row_bytes: usize,
6346        scale: f32,
6347    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6348        assert!(
6349            in_f % 64 == 0,
6350            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6351        );
6352        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6353        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6354        if scale != 1.0 {
6355            self.scale_inplace(&mut y, scale, m * out_f)?;
6356        }
6357        Ok(y)
6358    }
6359
6360    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
6361    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
6362    fn fp4_gemm_launch(
6363        &self,
6364        bytes: &CudaSlice<u8>,
6365        aq4: &CudaSlice<u32>,
6366        ad4: &CudaSlice<u8>,
6367        m: usize,
6368        in_f: usize,
6369        out_f: usize,
6370        row_bytes: usize,
6371    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6372        let f = self.func("qmatvec_gemm_nvfp4_fp4");
6373        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6374        const BM: u32 = 64;
6375        const BN: u32 = 256;
6376        let cfg = LaunchConfig {
6377            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6378            block_dim: (32, 4, 1),
6379            shared_mem_bytes: 0,
6380        };
6381        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6382        let __s_b = self.gpu.stream();
6383        let mut b = __s_b.launch_builder(&f);
6384        b.arg(bytes)
6385            .arg(aq4)
6386            .arg(ad4)
6387            .arg(&mut y)
6388            .arg(&inf)
6389            .arg(&outf)
6390            .arg(&mi)
6391            .arg(&rb);
6392        unsafe {
6393            b.launch(cfg)?;
6394        }
6395        Ok(y)
6396    }
6397
6398    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
6399    pub fn qmatvec_gemm_nvfp4_fp4_raw(
6400        &self,
6401        bytes: &CudaSlice<u8>,
6402        x: &CudaSlice<f32>,
6403        m: usize,
6404        in_f: usize,
6405        out_f: usize,
6406        row_bytes: usize,
6407    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6408        assert!(
6409            in_f % 64 == 0,
6410            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6411        );
6412        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6413        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6414    }
6415
6416    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
6417    pub fn qmatvec_q8_0_fast(
6418        &self,
6419        w: &CudaSlice<u8>,
6420        x: &CudaSlice<f32>,
6421        m: usize,
6422        in_f: usize,
6423        out_f: usize,
6424        row_bytes: usize,
6425    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6426        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6427        let f = self.func("qmatvec_q8_0_dp4a");
6428        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6429        let cfg = LaunchConfig {
6430            grid_dim: (out_f as u32, m as u32, 1),
6431            block_dim: (128, 1, 1),
6432            shared_mem_bytes: 0,
6433        };
6434        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6435        let __s_b = self.gpu.stream();
6436        let mut b = __s_b.launch_builder(&f);
6437        b.arg(w)
6438            .arg(&aq)
6439            .arg(&ad)
6440            .arg(&mut y)
6441            .arg(&inf)
6442            .arg(&outf)
6443            .arg(&mi)
6444            .arg(&rb);
6445        unsafe {
6446            b.launch(cfg)?;
6447        }
6448        Ok(y)
6449    }
6450
6451    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6452    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6453    pub fn qmatvec_q4_K_fast(
6454        &self,
6455        w: &CudaSlice<u8>,
6456        x: &CudaSlice<f32>,
6457        m: usize,
6458        in_f: usize,
6459        out_f: usize,
6460        row_bytes: usize,
6461    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6462        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6463        let f = self.func("qmatvec_q4_K_dp4a");
6464        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6465        let cfg = LaunchConfig {
6466            grid_dim: (out_f as u32, m as u32, 1),
6467            block_dim: (128, 1, 1),
6468            shared_mem_bytes: 0,
6469        };
6470        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6471        let __s_b = self.gpu.stream();
6472        let mut b = __s_b.launch_builder(&f);
6473        b.arg(w)
6474            .arg(&aq)
6475            .arg(&ad)
6476            .arg(&mut y)
6477            .arg(&inf)
6478            .arg(&outf)
6479            .arg(&mi)
6480            .arg(&rb);
6481        unsafe {
6482            b.launch(cfg)?;
6483        }
6484        Ok(y)
6485    }
6486
6487    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6488    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6489    pub fn qmatvec_q6_K_fast(
6490        &self,
6491        w: &CudaSlice<u8>,
6492        x: &CudaSlice<f32>,
6493        m: usize,
6494        in_f: usize,
6495        out_f: usize,
6496        row_bytes: usize,
6497    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6498        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6499        let f = self.func("qmatvec_q6_K_dp4a");
6500        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6501        let cfg = LaunchConfig {
6502            grid_dim: (out_f as u32, m as u32, 1),
6503            block_dim: (128, 1, 1),
6504            shared_mem_bytes: 0,
6505        };
6506        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6507        let __s_b = self.gpu.stream();
6508        let mut b = __s_b.launch_builder(&f);
6509        b.arg(w)
6510            .arg(&aq)
6511            .arg(&ad)
6512            .arg(&mut y)
6513            .arg(&inf)
6514            .arg(&outf)
6515            .arg(&mi)
6516            .arg(&rb);
6517        unsafe {
6518            b.launch(cfg)?;
6519        }
6520        Ok(y)
6521    }
6522
6523    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6524    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6525    pub fn qmatvec_q5_K_fast(
6526        &self,
6527        w: &CudaSlice<u8>,
6528        x: &CudaSlice<f32>,
6529        m: usize,
6530        in_f: usize,
6531        out_f: usize,
6532        row_bytes: usize,
6533    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6534        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6535    }
6536    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6537    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6538    pub fn qmatvec_q3_K_fast(
6539        &self,
6540        w: &CudaSlice<u8>,
6541        x: &CudaSlice<f32>,
6542        m: usize,
6543        in_f: usize,
6544        out_f: usize,
6545        row_bytes: usize,
6546    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6547        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6548    }
6549    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
6550    pub fn qmatvec_nvfp4_fast_rp(
6551        &self,
6552        w: &CudaSlice<u8>,
6553        x: &CudaSlice<f32>,
6554        m: usize,
6555        in_f: usize,
6556        out_f: usize,
6557        row_bytes: usize,
6558    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6559        assert!(
6560            in_f % 64 == 0,
6561            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6562        );
6563        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6564    }
6565    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
6566    pub fn qmatvec_nvfp4_fast(
6567        &self,
6568        w: &CudaSlice<u8>,
6569        x: &CudaSlice<f32>,
6570        m: usize,
6571        in_f: usize,
6572        out_f: usize,
6573        row_bytes: usize,
6574    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6575        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
6576        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
6577        assert!(
6578            in_f % 64 == 0,
6579            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6580        );
6581        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6582    }
6583    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
6584    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6585    pub fn qmatvec_iq4_XS_fast(
6586        &self,
6587        w: &CudaSlice<u8>,
6588        x: &CudaSlice<f32>,
6589        m: usize,
6590        in_f: usize,
6591        out_f: usize,
6592        row_bytes: usize,
6593    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6594        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6595    }
6596
6597    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
6598    fn qmatvec_dp4a_named(
6599        &self,
6600        name: &str,
6601        w: &CudaSlice<u8>,
6602        x: &CudaSlice<f32>,
6603        m: usize,
6604        in_f: usize,
6605        out_f: usize,
6606        row_bytes: usize,
6607    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6608        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6609        let f = self.func(name);
6610        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6611        let cfg = LaunchConfig {
6612            grid_dim: (out_f as u32, m as u32, 1),
6613            block_dim: (128, 1, 1),
6614            shared_mem_bytes: 0,
6615        };
6616        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6617        let __s_b = self.gpu.stream();
6618        let mut b = __s_b.launch_builder(&f);
6619        b.arg(w)
6620            .arg(&aq)
6621            .arg(&ad)
6622            .arg(&mut y)
6623            .arg(&inf)
6624            .arg(&outf)
6625            .arg(&mi)
6626            .arg(&rb);
6627        unsafe {
6628            b.launch(cfg)?;
6629        }
6630        Ok(y)
6631    }
6632
6633    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6634        Ok(self.gpu.stream().clone_htod(v)?)
6635    }
6636    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6637        Ok(self.gpu.stream().clone_htod(v)?)
6638    }
6639    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
6640    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6641        Ok(self.gpu.stream().clone_htod(v)?)
6642    }
6643    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6644        Ok(self.gpu.stream().clone_htod(v)?)
6645    }
6646    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
6647    pub fn dtoh_view(
6648        &self,
6649        d: &cudarc::driver::CudaView<f32>,
6650    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6651        let v = self.gpu.stream().clone_dtoh(d)?;
6652        self.gpu.stream().synchronize()?;
6653        Ok(v)
6654    }
6655    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6656        let v = self.gpu.stream().clone_dtoh(d)?;
6657        self.gpu.stream().synchronize()?;
6658        Ok(v)
6659    }
6660    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
6661    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
6662    /// issuing them together avoids a second stream synchronization in every trunk layer.
6663    pub fn dtoh_pair(
6664        &self,
6665        a: &CudaSlice<f32>,
6666        b: &CudaSlice<f32>,
6667    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6668        let av = self.gpu.stream().clone_dtoh(a)?;
6669        let bv = self.gpu.stream().clone_dtoh(b)?;
6670        self.gpu.stream().synchronize()?;
6671        Ok((av, bv))
6672    }
6673    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
6674    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6675        let v = self.gpu.stream().clone_dtoh(d)?;
6676        self.gpu.stream().synchronize()?;
6677        Ok(v)
6678    }
6679    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
6680    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6681        let v = self.gpu.stream().clone_dtoh(d)?;
6682        self.gpu.stream().synchronize()?;
6683        Ok(v)
6684    }
6685    pub fn dtoh_u8_view(
6686        &self,
6687        d: &cudarc::driver::CudaView<u8>,
6688    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6689        let v = self.gpu.stream().clone_dtoh(d)?;
6690        self.gpu.stream().synchronize()?;
6691        Ok(v)
6692    }
6693    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6694        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6695        self.keep_if_capturing(&s);
6696        Ok(s)
6697    }
6698
6699    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
6700    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
6701    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
6702    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
6703    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
6704    /// back (or kept resident for graph replay). Returns the device token buffer.
6705    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
6706    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
6707    pub fn prob_of_token_device(
6708        &self,
6709        logits: &CudaSlice<f32>,
6710        tok: &CudaSlice<u32>,
6711        n_vocab: usize,
6712    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6713        let nb = ARGMAX_NB;
6714        let mut part = self.alloc_uninit::<f32>(nb)?;
6715        let mut p = self.alloc_uninit::<f32>(1)?;
6716        let f1 = self.func("prob_of_token_partial_f32");
6717        let cfg1 = LaunchConfig {
6718            grid_dim: (nb as u32, 1, 1),
6719            block_dim: (256, 1, 1),
6720            shared_mem_bytes: 0,
6721        };
6722        let nv = n_vocab as i32;
6723        let __s_b1 = self.gpu.stream();
6724        let mut b1 = __s_b1.launch_builder(&f1);
6725        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6726        unsafe {
6727            b1.launch(cfg1)?;
6728        }
6729        let f2 = self.func("prob_of_token_final_f32");
6730        let cfg2 = LaunchConfig {
6731            grid_dim: (1, 1, 1),
6732            block_dim: (256, 1, 1),
6733            shared_mem_bytes: 0,
6734        };
6735        let nbi = nb as i32;
6736        let __s_b2 = self.gpu.stream();
6737        let mut b2 = __s_b2.launch_builder(&f2);
6738        b2.arg(&part).arg(&mut p).arg(&nbi);
6739        unsafe {
6740            b2.launch(cfg2)?;
6741        }
6742        Ok(p)
6743    }
6744
6745    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
6746    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
6747    /// where the host reads the p-min confidence between replays. Same kernels, same math.
6748    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
6749    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
6750    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
6751    pub fn prob_of_token_device_col(
6752        &self,
6753        logits: &CudaSlice<f32>,
6754        tok_all: &CudaSlice<u32>,
6755        tok_idx: usize,
6756        p_out: &mut CudaSlice<f32>,
6757        p_idx: usize,
6758        n_vocab: usize,
6759    ) -> Result<(), Box<dyn std::error::Error>> {
6760        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6761        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6762        let nb = ARGMAX_NB;
6763        let mut part = self.alloc_uninit::<f32>(nb)?;
6764        let f1 = self.func("prob_of_token_partial_f32");
6765        let cfg1 = LaunchConfig {
6766            grid_dim: (nb as u32, 1, 1),
6767            block_dim: (256, 1, 1),
6768            shared_mem_bytes: 0,
6769        };
6770        let nv = n_vocab as i32;
6771        let __s_b1 = self.gpu.stream();
6772        let mut b1 = __s_b1.launch_builder(&f1);
6773        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6774        unsafe {
6775            b1.launch(cfg1)?;
6776        }
6777        let f2 = self.func("prob_of_token_final_f32");
6778        let cfg2 = LaunchConfig {
6779            grid_dim: (1, 1, 1),
6780            block_dim: (256, 1, 1),
6781            shared_mem_bytes: 0,
6782        };
6783        let nbi = nb as i32;
6784        let __s_b2 = self.gpu.stream();
6785        let mut b2 = __s_b2.launch_builder(&f2);
6786        b2.arg(&part).arg(&mut p_v).arg(&nbi);
6787        unsafe {
6788            b2.launch(cfg2)?;
6789        }
6790        Ok(())
6791    }
6792
6793    pub fn prob_of_token_device_into(
6794        &self,
6795        logits: &CudaSlice<f32>,
6796        tok: &CudaSlice<u32>,
6797        p_out: &mut CudaSlice<f32>,
6798        n_vocab: usize,
6799    ) -> Result<(), Box<dyn std::error::Error>> {
6800        let nb = ARGMAX_NB;
6801        let mut part = self.alloc_uninit::<f32>(nb)?;
6802        let f1 = self.func("prob_of_token_partial_f32");
6803        let cfg1 = LaunchConfig {
6804            grid_dim: (nb as u32, 1, 1),
6805            block_dim: (256, 1, 1),
6806            shared_mem_bytes: 0,
6807        };
6808        let nv = n_vocab as i32;
6809        let __s_b1 = self.gpu.stream();
6810        let mut b1 = __s_b1.launch_builder(&f1);
6811        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6812        unsafe {
6813            b1.launch(cfg1)?;
6814        }
6815        let f2 = self.func("prob_of_token_final_f32");
6816        let cfg2 = LaunchConfig {
6817            grid_dim: (1, 1, 1),
6818            block_dim: (256, 1, 1),
6819            shared_mem_bytes: 0,
6820        };
6821        let nbi = nb as i32;
6822        let __s_b2 = self.gpu.stream();
6823        let mut b2 = __s_b2.launch_builder(&f2);
6824        b2.arg(&part).arg(p_out).arg(&nbi);
6825        unsafe {
6826            b2.launch(cfg2)?;
6827        }
6828        Ok(())
6829    }
6830
6831    pub fn argmax_token_device(
6832        &self,
6833        logits: &CudaSlice<f32>,
6834        n_vocab: usize,
6835    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6836        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6837        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6838        Ok(tok)
6839    }
6840    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
6841    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
6842    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
6843    /// pointer is baked once and the token id never round-trips to host inside steady state. The
6844    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
6845    /// captured passes bake fixed addresses.
6846    pub fn argmax_token_device_into(
6847        &self,
6848        logits: &CudaSlice<f32>,
6849        tok: &mut CudaSlice<u32>,
6850        n_vocab: usize,
6851    ) -> Result<(), Box<dyn std::error::Error>> {
6852        let nb = ARGMAX_NB;
6853        let f1 = self.func("argmax_partial_f32");
6854        let f2 = self.func("argmax_final_f32");
6855        let mut guard = self.argmax_partials.lock().unwrap();
6856        if guard.is_none() {
6857            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
6858            // buffers carry no cudarc events (illegal inside capture).
6859            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6860            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6861            *guard = Some((pv, pi));
6862        }
6863        let (part_v, part_i) = guard.as_mut().unwrap();
6864        let nv = n_vocab as i32;
6865        let nbi = nb as i32;
6866        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
6867        let cfg1 = LaunchConfig {
6868            grid_dim: (nb as u32, 1, 1),
6869            block_dim: (256, 1, 1),
6870            shared_mem_bytes: 0,
6871        };
6872        let __s_b1 = self.gpu.stream();
6873        let mut b1 = __s_b1.launch_builder(&f1);
6874        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
6875        unsafe {
6876            b1.launch(cfg1)?;
6877        }
6878        // pass 2: one block reduces NB partials -> token_out[0].
6879        let cfg2 = LaunchConfig {
6880            grid_dim: (1, 1, 1),
6881            block_dim: (256, 1, 1),
6882            shared_mem_bytes: 0,
6883        };
6884        let __s_b2 = self.gpu.stream();
6885        let mut b2 = __s_b2.launch_builder(&f2);
6886        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
6887        unsafe {
6888            b2.launch(cfg2)?;
6889        }
6890        Ok(())
6891    }
6892    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
6893    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
6894    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
6895    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
6896    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
6897    pub fn argmax_token_device_col(
6898        &self,
6899        logits: &CudaSlice<f32>,
6900        col: usize,
6901        n_vocab: usize,
6902        toks: &mut CudaSlice<u32>,
6903        out_idx: usize,
6904    ) -> Result<(), Box<dyn std::error::Error>> {
6905        let nb = ARGMAX_NB;
6906        let f1 = self.func("argmax_partial_f32");
6907        let f2 = self.func("argmax_final_f32");
6908        let mut guard = self.argmax_partials.lock().unwrap();
6909        if guard.is_none() {
6910            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6911            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6912            *guard = Some((pv, pi));
6913        }
6914        let (part_v, part_i) = guard.as_mut().unwrap();
6915        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
6916        let nv = n_vocab as i32;
6917        let nbi = nb as i32;
6918        let cfg1 = LaunchConfig {
6919            grid_dim: (nb as u32, 1, 1),
6920            block_dim: (256, 1, 1),
6921            shared_mem_bytes: 0,
6922        };
6923        let __s_b1 = self.gpu.stream();
6924        let mut b1 = __s_b1.launch_builder(&f1);
6925        b1.arg(&col_view)
6926            .arg(&mut *part_v)
6927            .arg(&mut *part_i)
6928            .arg(&nv);
6929        unsafe {
6930            b1.launch(cfg1)?;
6931        }
6932        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
6933        let cfg2 = LaunchConfig {
6934            grid_dim: (1, 1, 1),
6935            block_dim: (256, 1, 1),
6936            shared_mem_bytes: 0,
6937        };
6938        let __s_b2 = self.gpu.stream();
6939        let mut b2 = __s_b2.launch_builder(&f2);
6940        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
6941        unsafe {
6942            b2.launch(cfg2)?;
6943        }
6944        Ok(())
6945    }
6946    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
6947    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6948        Ok(self.gpu.stream().clone_htod(v)?)
6949    }
6950    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6951        let v = self.gpu.stream().clone_dtoh(d)?;
6952        self.gpu.stream().synchronize()?;
6953        Ok(v)
6954    }
6955    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
6956    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
6957    /// contents change every step, the address must not, so a captured graph can read it).
6958    pub fn htod_u32_into(
6959        &self,
6960        dst: &mut CudaSlice<u32>,
6961        src: &[u32],
6962    ) -> Result<(), Box<dyn std::error::Error>> {
6963        let mut view = dst.slice_mut(0..src.len());
6964        self.gpu.stream().memcpy_htod(src, &mut view)?;
6965        Ok(())
6966    }
6967
6968    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
6969    /// table without changing the device address its reconcile kernel consumes.
6970    pub fn htod_i32_into(
6971        &self,
6972        dst: &mut CudaSlice<i32>,
6973        src: &[i32],
6974    ) -> Result<(), Box<dyn std::error::Error>> {
6975        let mut view = dst.slice_mut(0..src.len());
6976        self.gpu.stream().memcpy_htod(src, &mut view)?;
6977        Ok(())
6978    }
6979
6980    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6981        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
6982        self.keep_if_capturing(&s);
6983        Ok(s)
6984    }
6985    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
6986    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
6987    pub fn embed_gather_device_into(
6988        &self,
6989        embd: &CudaSlice<u8>,
6990        token_d: &CudaSlice<u32>,
6991        x_out: &mut CudaSlice<f32>,
6992        n_embd: usize,
6993        qtype: i32,
6994        row_bytes: usize,
6995    ) -> Result<(), Box<dyn std::error::Error>> {
6996        let f = self.func("embed_gather_u32");
6997        let cfg = LaunchConfig {
6998            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
6999            block_dim: (256, 1, 1),
7000            shared_mem_bytes: 0,
7001        };
7002        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7003        let __s_b = self.gpu.stream();
7004        let mut b = __s_b.launch_builder(&f);
7005        b.arg(embd)
7006            .arg(token_d)
7007            .arg(x_out)
7008            .arg(&ne)
7009            .arg(&qt)
7010            .arg(&rb);
7011        unsafe {
7012            b.launch(cfg)?;
7013        }
7014        Ok(())
7015    }
7016    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
7017    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7018        let v = self.gpu.stream().clone_dtoh(d)?;
7019        self.gpu.stream().synchronize()?;
7020        Ok(v[0])
7021    }
7022    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
7023    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
7024    /// the counter value after the throwaway capture warmups corrupt it.
7025    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
7026    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
7027    /// copy (fine at stream-idle boundaries, poison mid-round).
7028    pub fn i32_set_k(
7029        &self,
7030        dst: &mut CudaSlice<i32>,
7031        v: i32,
7032    ) -> Result<(), Box<dyn std::error::Error>> {
7033        let f = self.func("i32_set_k");
7034        let cfg = LaunchConfig {
7035            grid_dim: (1, 1, 1),
7036            block_dim: (1, 1, 1),
7037            shared_mem_bytes: 0,
7038        };
7039        let idx = 0i32;
7040        let __s_b = self.gpu.stream();
7041        let mut b = __s_b.launch_builder(&f);
7042        b.arg(dst).arg(&v).arg(&idx);
7043        unsafe {
7044            b.launch(cfg)?;
7045        }
7046        Ok(())
7047    }
7048
7049    pub fn set_i32_one(
7050        &self,
7051        d: &mut CudaSlice<i32>,
7052        v: i32,
7053    ) -> Result<(), Box<dyn std::error::Error>> {
7054        self.gpu.stream().memcpy_htod(&[v], d)?;
7055        Ok(())
7056    }
7057    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
7058    /// during priming / capture-state restore.
7059    pub fn set_u32_one(
7060        &self,
7061        d: &mut CudaSlice<u32>,
7062        v: u32,
7063    ) -> Result<(), Box<dyn std::error::Error>> {
7064        self.gpu.stream().memcpy_htod(&[v], d)?;
7065        Ok(())
7066    }
7067    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
7068    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7069        let v = self.gpu.stream().clone_dtoh(d)?;
7070        self.gpu.stream().synchronize()?;
7071        Ok(v[0])
7072    }
7073    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
7074    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7075        Ok(self.gpu.stream().clone_htod(bytes)?)
7076    }
7077    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
7078    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
7079    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
7080    pub fn embed_gather_device(
7081        &self,
7082        embd: &CudaSlice<u8>,
7083        token_d: &CudaSlice<u32>,
7084        n_embd: usize,
7085        qtype: i32,
7086        row_bytes: usize,
7087    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7088        let f = self.func("embed_gather_u32");
7089        let mut x = self.alloc_uninit::<f32>(n_embd)?;
7090        let cfg = LaunchConfig {
7091            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7092            block_dim: (256, 1, 1),
7093            shared_mem_bytes: 0,
7094        };
7095        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7096        let __s_b = self.gpu.stream();
7097        let mut b = __s_b.launch_builder(&f);
7098        b.arg(embd)
7099            .arg(token_d)
7100            .arg(&mut x)
7101            .arg(&ne)
7102            .arg(&qt)
7103            .arg(&rb);
7104        unsafe {
7105            b.launch(cfg)?;
7106        }
7107        Ok(x)
7108    }
7109
7110    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
7111    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
7112    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
7113    pub fn embed_gather_device_t(
7114        &self,
7115        embd: &CudaSlice<u8>,
7116        tokens: &[u32],
7117        n_embd: usize,
7118        qtype: i32,
7119        row_bytes: usize,
7120    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7121        let t = tokens.len();
7122        let tok_d = self.gpu.stream().clone_htod(tokens)?;
7123        let f = self.func("embed_gather_u32_t");
7124        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7125        let cfg = LaunchConfig {
7126            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7127            block_dim: (256, 1, 1),
7128            shared_mem_bytes: 0,
7129        };
7130        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7131        let __s_b = self.gpu.stream();
7132        let mut b = __s_b.launch_builder(&f);
7133        b.arg(embd)
7134            .arg(&tok_d)
7135            .arg(&mut x)
7136            .arg(&ne)
7137            .arg(&qt)
7138            .arg(&rb)
7139            .arg(&ti);
7140        unsafe {
7141            b.launch(cfg)?;
7142        }
7143        Ok(x)
7144    }
7145
7146    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
7147    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
7148    /// as embed_gather_device_t — bit-identical rows.
7149    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
7150    pub fn embed_gather_device_tv(
7151        &self,
7152        embd: &CudaSlice<u8>,
7153        tok_v: &cudarc::driver::CudaView<u32>,
7154        t: usize,
7155        n_embd: usize,
7156        qtype: i32,
7157        row_bytes: usize,
7158    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7159        let f = self.func("embed_gather_u32_t");
7160        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7161        let cfg = LaunchConfig {
7162            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7163            block_dim: (256, 1, 1),
7164            shared_mem_bytes: 0,
7165        };
7166        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7167        let __s_b = self.gpu.stream();
7168        let mut b = __s_b.launch_builder(&f);
7169        b.arg(embd)
7170            .arg(tok_v)
7171            .arg(&mut x)
7172            .arg(&ne)
7173            .arg(&qt)
7174            .arg(&rb)
7175            .arg(&ti);
7176        unsafe {
7177            b.launch(cfg)?;
7178        }
7179        Ok(x)
7180    }
7181
7182    pub fn embed_gather_device_td(
7183        &self,
7184        embd: &CudaSlice<u8>,
7185        tok_d: &CudaSlice<u32>,
7186        t: usize,
7187        n_embd: usize,
7188        qtype: i32,
7189        row_bytes: usize,
7190    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7191        let f = self.func("embed_gather_u32_t");
7192        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7193        let cfg = LaunchConfig {
7194            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7195            block_dim: (256, 1, 1),
7196            shared_mem_bytes: 0,
7197        };
7198        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7199        let __s_b = self.gpu.stream();
7200        let mut b = __s_b.launch_builder(&f);
7201        b.arg(embd)
7202            .arg(tok_d)
7203            .arg(&mut x)
7204            .arg(&ne)
7205            .arg(&qt)
7206            .arg(&rb)
7207            .arg(&ti);
7208        unsafe {
7209            b.launch(cfg)?;
7210        }
7211        Ok(x)
7212    }
7213
7214    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
7215    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
7216    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
7217    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
7218    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
7219    #[inline]
7220    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
7221    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7222        if self
7223            .capture_keep_on
7224            .load(std::sync::atomic::Ordering::Relaxed)
7225        {
7226            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7227        }
7228    }
7229
7230    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7231        &self,
7232        n: usize,
7233    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7234        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7235        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
7236        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
7237        // not cover engine-internal buffers). Debug-only: massive launch overhead.
7238        {
7239            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7240            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7241                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
7242                use cudarc::driver::DevicePtrMut;
7243                let n_bytes = s.len() * std::mem::size_of::<T>();
7244                let stream = self.gpu.stream();
7245                let (p_, _g) = s.device_ptr_mut(&stream);
7246                unsafe {
7247                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7248                        .result()?;
7249                }
7250            }
7251        }
7252        self.keep_if_capturing(&s);
7253        Ok(s)
7254    }
7255
7256    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
7257    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
7258    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
7259    /// consumers alloc through this (m=1 decode arms).
7260    pub fn uninit_q8_pair(
7261        &self,
7262        n: usize,
7263    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7264        Ok((
7265            self.alloc_uninit::<i8>(n)?,
7266            self.alloc_uninit::<f32>(n / 32)?,
7267        ))
7268    }
7269
7270    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7271        self.alloc_uninit::<f32>(n)
7272    }
7273
7274    /// i8 uninitialized scratch (same contract as `uninit`).
7275    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7276        self.alloc_uninit::<i8>(n)
7277    }
7278
7279    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
7280    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
7281    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
7282    #[allow(clippy::too_many_arguments)]
7283    pub fn rms_norm3(
7284        &self,
7285        x: &CudaSlice<f32>,
7286        w0: &CudaSlice<f32>,
7287        w1: &CudaSlice<f32>,
7288        w2: &CudaSlice<f32>,
7289        d0: &mut CudaSlice<f32>,
7290        d1: &mut CudaSlice<f32>,
7291        d2: &mut CudaSlice<f32>,
7292        ncols: usize,
7293        nrows: usize,
7294        eps: f32,
7295    ) -> Result<(), Box<dyn std::error::Error>> {
7296        let f = self.func("rms_norm3_f32");
7297        let cfg = LaunchConfig {
7298            grid_dim: (nrows as u32, 1, 1),
7299            block_dim: (rms_block(), 1, 1),
7300            shared_mem_bytes: 0,
7301        };
7302        let (nc, e) = (ncols as i32, eps);
7303        let __s_b = self.gpu.stream();
7304        let mut b = __s_b.launch_builder(&f);
7305        b.arg(x)
7306            .arg(w0)
7307            .arg(w1)
7308            .arg(w2)
7309            .arg(d0)
7310            .arg(d1)
7311            .arg(d2)
7312            .arg(&nc)
7313            .arg(&e);
7314        unsafe {
7315            b.launch(cfg)?;
7316        }
7317        Ok(())
7318    }
7319
7320    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
7321    #[allow(clippy::too_many_arguments)]
7322    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
7323    /// piggybacks on the same conditions.
7324    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7325        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7326        *WARP_ON.get_or_init(|| {
7327            std::env::var("MEMRA_QKVNORM_W")
7328                .map(|v| v != "0")
7329                .unwrap_or(true)
7330        }) && ncols % 4 == 0
7331            && rows >= 64
7332    }
7333
7334    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
7335    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
7336    #[allow(clippy::too_many_arguments)]
7337    pub fn rms_norm_qkv_w4b(
7338        &self,
7339        q: &CudaSlice<f32>,
7340        k: &CudaSlice<f32>,
7341        v: &CudaSlice<f32>,
7342        wq: &CudaSlice<f32>,
7343        wk: &CudaSlice<f32>,
7344        wv: &CudaSlice<f32>,
7345        dq: &mut CudaSlice<f32>,
7346        dk: &mut CudaSlice<f32>,
7347        dv: &mut CudaSlice<f32>,
7348        dvb: &mut CudaSlice<u8>,
7349        ncols: usize,
7350        rq: usize,
7351        rk: usize,
7352        eps: f32,
7353        vf16: bool,
7354    ) -> Result<(), Box<dyn std::error::Error>> {
7355        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7356        let f = self.func("rms_norm_qkv_w4b_f32");
7357        let rows = (rq + 2 * rk) as u32;
7358        let cfg = LaunchConfig {
7359            grid_dim: (rows.div_ceil(8), 1, 1),
7360            block_dim: (256, 1, 1),
7361            shared_mem_bytes: 0,
7362        };
7363        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7364        let vf = vf16 as i32;
7365        let __s_b = self.gpu.stream();
7366        let mut b = __s_b.launch_builder(&f);
7367        b.arg(q)
7368            .arg(k)
7369            .arg(v)
7370            .arg(wq)
7371            .arg(wk)
7372            .arg(wv)
7373            .arg(dq)
7374            .arg(dk)
7375            .arg(dv)
7376            .arg(&mut *dvb)
7377            .arg(&nc)
7378            .arg(&rqi)
7379            .arg(&rki)
7380            .arg(&rvi)
7381            .arg(&e)
7382            .arg(&vf);
7383        unsafe {
7384            b.launch(cfg)?;
7385        }
7386        Ok(())
7387    }
7388
7389    pub fn rms_norm_qkv(
7390        &self,
7391        q: &CudaSlice<f32>,
7392        k: &CudaSlice<f32>,
7393        v: &CudaSlice<f32>,
7394        wq: &CudaSlice<f32>,
7395        wk: &CudaSlice<f32>,
7396        wv: &CudaSlice<f32>,
7397        dq: &mut CudaSlice<f32>,
7398        dk: &mut CudaSlice<f32>,
7399        dv: &mut CudaSlice<f32>,
7400        ncols: usize,
7401        rq: usize,
7402        rk: usize,
7403        eps: f32,
7404    ) -> Result<(), Box<dyn std::error::Error>> {
7405        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
7406        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
7407        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
7408        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7409        let warp_on = *WARP_ON.get_or_init(|| {
7410            std::env::var("MEMRA_QKVNORM_W")
7411                .map(|v| v != "0")
7412                .unwrap_or(true)
7413        });
7414        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
7415        // replay numerics are untouched on every model; only prefill depth takes the new config.
7416        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7417            let f = self.func("rms_norm_qkv_w4_f32");
7418            let rows = (rq + 2 * rk) as u32;
7419            let cfg = LaunchConfig {
7420                grid_dim: (rows.div_ceil(8), 1, 1),
7421                block_dim: (256, 1, 1),
7422                shared_mem_bytes: 0,
7423            };
7424            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7425            let __s_b = self.gpu.stream();
7426            let mut b = __s_b.launch_builder(&f);
7427            b.arg(q)
7428                .arg(k)
7429                .arg(v)
7430                .arg(wq)
7431                .arg(wk)
7432                .arg(wv)
7433                .arg(dq)
7434                .arg(dk)
7435                .arg(dv)
7436                .arg(&nc)
7437                .arg(&rqi)
7438                .arg(&rki)
7439                .arg(&rvi)
7440                .arg(&e);
7441            unsafe {
7442                b.launch(cfg)?;
7443            }
7444            return Ok(());
7445        }
7446        let f = self.func("rms_norm_qkv_f32");
7447        let grid = (rq + 2 * rk) as u32;
7448        let cfg = LaunchConfig {
7449            grid_dim: (grid, 1, 1),
7450            block_dim: (rms_block(), 1, 1),
7451            shared_mem_bytes: 0,
7452        };
7453        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7454        let __s_b = self.gpu.stream();
7455        let mut b = __s_b.launch_builder(&f);
7456        b.arg(q)
7457            .arg(k)
7458            .arg(v)
7459            .arg(wq)
7460            .arg(wk)
7461            .arg(wv)
7462            .arg(dq)
7463            .arg(dk)
7464            .arg(dv)
7465            .arg(&nc)
7466            .arg(&rqi)
7467            .arg(&rki)
7468            .arg(&e);
7469        unsafe {
7470            b.launch(cfg)?;
7471        }
7472        Ok(())
7473    }
7474
7475    /// gemma4 fused pair of rms_norms over two different inputs (same width).
7476    #[allow(clippy::too_many_arguments)]
7477    pub fn rms_norm2x(
7478        &self,
7479        a: &CudaSlice<f32>,
7480        bb: &CudaSlice<f32>,
7481        wa: &CudaSlice<f32>,
7482        wb: &CudaSlice<f32>,
7483        da: &mut CudaSlice<f32>,
7484        db: &mut CudaSlice<f32>,
7485        ncols: usize,
7486        nrows: usize,
7487        eps: f32,
7488    ) -> Result<(), Box<dyn std::error::Error>> {
7489        let f = self.func("rms_norm2x_f32");
7490        let cfg = LaunchConfig {
7491            grid_dim: (2 * nrows as u32, 1, 1),
7492            block_dim: (rms_block(), 1, 1),
7493            shared_mem_bytes: 0,
7494        };
7495        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7496        let __s_b = self.gpu.stream();
7497        let mut b = __s_b.launch_builder(&f);
7498        b.arg(a)
7499            .arg(bb)
7500            .arg(wa)
7501            .arg(wb)
7502            .arg(da)
7503            .arg(db)
7504            .arg(&nc)
7505            .arg(&nr)
7506            .arg(&e);
7507        unsafe {
7508            b.launch(cfg)?;
7509        }
7510        Ok(())
7511    }
7512
7513    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
7514    pub fn softcap(
7515        &self,
7516        y: &mut CudaSlice<f32>,
7517        cap: f32,
7518        n: usize,
7519    ) -> Result<(), Box<dyn std::error::Error>> {
7520        let f = self.func("softcap_f32");
7521        let cfg = LaunchConfig::for_num_elems(n as u32);
7522        let ni = n as i32;
7523        let __s_b = self.gpu.stream();
7524        let mut b = __s_b.launch_builder(&f);
7525        b.arg(y).arg(&cap).arg(&ni);
7526        unsafe {
7527            b.launch(cfg)?;
7528        }
7529        Ok(())
7530    }
7531
7532    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
7533    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
7534    pub fn mask_ids_rows(
7535        &self,
7536        y: &mut CudaSlice<f32>,
7537        ids: &CudaSlice<i32>,
7538        n_ids: usize,
7539        n_vocab: usize,
7540        t: usize,
7541    ) -> Result<(), Box<dyn std::error::Error>> {
7542        let f = self.func("mask_ids_rows_f32");
7543        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7544        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7545        let __s_b = self.gpu.stream();
7546        let mut b = __s_b.launch_builder(&f);
7547        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7548        unsafe {
7549            b.launch(cfg)?;
7550        }
7551        Ok(())
7552    }
7553
7554    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
7555    #[allow(clippy::too_many_arguments)]
7556    pub fn add_scale_rms_norm(
7557        &self,
7558        a: &CudaSlice<f32>,
7559        b_in: &CudaSlice<f32>,
7560        c: f32,
7561        w: &CudaSlice<f32>,
7562        res: &mut CudaSlice<f32>,
7563        dst: &mut CudaSlice<f32>,
7564        ncols: usize,
7565        nrows: usize,
7566        eps: f32,
7567    ) -> Result<(), Box<dyn std::error::Error>> {
7568        let f = self.func("add_scale_rms_norm_f32");
7569        let cfg = LaunchConfig {
7570            grid_dim: (nrows as u32, 1, 1),
7571            block_dim: (rms_block(), 1, 1),
7572            shared_mem_bytes: 0,
7573        };
7574        let (nc, e2) = (ncols as i32, eps);
7575        let __s_b = self.gpu.stream();
7576        let mut b = __s_b.launch_builder(&f);
7577        b.arg(a)
7578            .arg(b_in)
7579            .arg(&c)
7580            .arg(w)
7581            .arg(res)
7582            .arg(dst)
7583            .arg(&nc)
7584            .arg(&e2);
7585        unsafe {
7586            b.launch(cfg)?;
7587        }
7588        Ok(())
7589    }
7590
7591    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
7592    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
7593    #[allow(clippy::too_many_arguments)]
7594    pub fn add_scale_rms_norm_q8_1(
7595        &self,
7596        a: &CudaSlice<f32>,
7597        b_in: &CudaSlice<f32>,
7598        c: f32,
7599        w: &CudaSlice<f32>,
7600        res: &mut CudaSlice<f32>,
7601        ncols: usize,
7602        nrows: usize,
7603        eps: f32,
7604    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7605        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7606        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7607        let (nc, e2) = (ncols as i32, eps);
7608        if Self::pdl_on() && Self::pdl_wb_on() {
7609            {
7610                use cudarc::driver::{DevicePtr, DevicePtrMut};
7611                let s = &self.gpu.stream();
7612                let (pa, _g0) = a.device_ptr(s);
7613                let (pb, _g1) = b_in.device_ptr(s);
7614                let (pw, _g2) = w.device_ptr(s);
7615                let (pr, _g3) = res.device_ptr_mut(s);
7616                let (pq, _g4) = out_q.device_ptr_mut(s);
7617                let (pd, _g5) = out_d.device_ptr_mut(s);
7618                let mut ps = [
7619                    &pa as *const _ as *mut std::ffi::c_void,
7620                    &pb as *const _ as *mut _,
7621                    &c as *const _ as *mut _,
7622                    &pw as *const _ as *mut _,
7623                    &pr as *const _ as *mut _,
7624                    &pq as *const _ as *mut _,
7625                    &pd as *const _ as *mut _,
7626                    &nc as *const _ as *mut _,
7627                    &e2 as *const _ as *mut _,
7628                ];
7629                unsafe {
7630                    self.launch_pdl(
7631                        "add_scale_rms_norm_q8_1",
7632                        (nrows as u32, 1, 1),
7633                        (rms_block(), 1, 1),
7634                        &mut ps,
7635                    )?;
7636                }
7637            }
7638            return Ok((out_q, out_d));
7639        }
7640        let f = self.func("add_scale_rms_norm_q8_1");
7641        let cfg = LaunchConfig {
7642            grid_dim: (nrows as u32, 1, 1),
7643            block_dim: (rms_block(), 1, 1),
7644            shared_mem_bytes: 0,
7645        };
7646        let __s_b = self.gpu.stream();
7647        let mut b = __s_b.launch_builder(&f);
7648        b.arg(a)
7649            .arg(b_in)
7650            .arg(&c)
7651            .arg(w)
7652            .arg(res)
7653            .arg(&mut out_q)
7654            .arg(&mut out_d)
7655            .arg(&nc)
7656            .arg(&e2);
7657        unsafe {
7658            b.launch(cfg)?;
7659        }
7660        Ok((out_q, out_d))
7661    }
7662
7663    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
7664    #[allow(clippy::too_many_arguments)]
7665    pub fn add_scale_rms_norm_q8_1_into(
7666        &self,
7667        a: &CudaSlice<f32>,
7668        b_in: &CudaSlice<f32>,
7669        c: f32,
7670        w: &CudaSlice<f32>,
7671        res: &mut CudaSlice<f32>,
7672        ncols: usize,
7673        nrows: usize,
7674        eps: f32,
7675        out_q: &mut CudaSlice<i8>,
7676        out_d: &mut CudaSlice<f32>,
7677    ) -> Result<(), Box<dyn std::error::Error>> {
7678        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7679        let (nc, e2) = (ncols as i32, eps);
7680        if Self::pdl_on() && Self::pdl_wb_on() {
7681            use cudarc::driver::{DevicePtr, DevicePtrMut};
7682            let s = &self.gpu.stream();
7683            let (pa, _g0) = a.device_ptr(s);
7684            let (pb, _g1) = b_in.device_ptr(s);
7685            let (pw, _g2) = w.device_ptr(s);
7686            let (pr, _g3) = res.device_ptr_mut(s);
7687            let (pq, _g4) = out_q.device_ptr_mut(s);
7688            let (pd, _g5) = out_d.device_ptr_mut(s);
7689            let mut ps = [
7690                &pa as *const _ as *mut std::ffi::c_void,
7691                &pb as *const _ as *mut _,
7692                &c as *const _ as *mut _,
7693                &pw as *const _ as *mut _,
7694                &pr as *const _ as *mut _,
7695                &pq as *const _ as *mut _,
7696                &pd as *const _ as *mut _,
7697                &nc as *const _ as *mut _,
7698                &e2 as *const _ as *mut _,
7699            ];
7700            unsafe {
7701                self.launch_pdl(
7702                    "add_scale_rms_norm_q8_1",
7703                    (nrows as u32, 1, 1),
7704                    (rms_block(), 1, 1),
7705                    &mut ps,
7706                )?;
7707            }
7708            return Ok(());
7709        }
7710        let f = self.func("add_scale_rms_norm_q8_1");
7711        let cfg = LaunchConfig {
7712            grid_dim: (nrows as u32, 1, 1),
7713            block_dim: (rms_block(), 1, 1),
7714            shared_mem_bytes: 0,
7715        };
7716        let __s_b = self.gpu.stream();
7717        let mut b = __s_b.launch_builder(&f);
7718        b.arg(a)
7719            .arg(b_in)
7720            .arg(&c)
7721            .arg(w)
7722            .arg(res)
7723            .arg(&mut *out_q)
7724            .arg(&mut *out_d)
7725            .arg(&nc)
7726            .arg(&e2);
7727        unsafe {
7728            b.launch(cfg)?;
7729        }
7730        Ok(())
7731    }
7732
7733    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
7734    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
7735    #[allow(clippy::too_many_arguments)]
7736    pub fn rms_pre_add_scale_rms_norm_q8_1(
7737        &self,
7738        a: &CudaSlice<f32>,
7739        wa: &CudaSlice<f32>,
7740        b_in: &CudaSlice<f32>,
7741        c: f32,
7742        w: &CudaSlice<f32>,
7743        res: &mut CudaSlice<f32>,
7744        ncols: usize,
7745        nrows: usize,
7746        eps: f32,
7747    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7748        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7749        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7750        let (nc, e2) = (ncols as i32, eps);
7751        if Self::pdl_on() {
7752            {
7753                use cudarc::driver::{DevicePtr, DevicePtrMut};
7754                let s = &self.gpu.stream();
7755                let (pa, _g0) = a.device_ptr(s);
7756                let (pwa, _g1) = wa.device_ptr(s);
7757                let (pb, _g2) = b_in.device_ptr(s);
7758                let (pw, _g3) = w.device_ptr(s);
7759                let (pr, _g4) = res.device_ptr_mut(s);
7760                let (pq, _g5) = out_q.device_ptr_mut(s);
7761                let (pd, _g6) = out_d.device_ptr_mut(s);
7762                let mut ps = [
7763                    &pa as *const _ as *mut std::ffi::c_void,
7764                    &pwa as *const _ as *mut _,
7765                    &pb as *const _ as *mut _,
7766                    &c as *const _ as *mut _,
7767                    &pw as *const _ as *mut _,
7768                    &pr as *const _ as *mut _,
7769                    &pq as *const _ as *mut _,
7770                    &pd as *const _ as *mut _,
7771                    &nc as *const _ as *mut _,
7772                    &e2 as *const _ as *mut _,
7773                ];
7774                unsafe {
7775                    self.launch_pdl(
7776                        "rms_pre_add_scale_rms_norm_q8_1",
7777                        (nrows as u32, 1, 1),
7778                        (rms_block(), 1, 1),
7779                        &mut ps,
7780                    )?;
7781                }
7782            }
7783            return Ok((out_q, out_d));
7784        }
7785        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7786        let cfg = LaunchConfig {
7787            grid_dim: (nrows as u32, 1, 1),
7788            block_dim: (rms_block(), 1, 1),
7789            shared_mem_bytes: 0,
7790        };
7791        let __s_b = self.gpu.stream();
7792        let mut b = __s_b.launch_builder(&f);
7793        b.arg(a)
7794            .arg(wa)
7795            .arg(b_in)
7796            .arg(&c)
7797            .arg(w)
7798            .arg(res)
7799            .arg(&mut out_q)
7800            .arg(&mut out_d)
7801            .arg(&nc)
7802            .arg(&e2);
7803        unsafe {
7804            b.launch(cfg)?;
7805        }
7806        Ok((out_q, out_d))
7807    }
7808
7809    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
7810    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
7811    pub fn gelu_tanh_mul_q8_1(
7812        &self,
7813        gate: &CudaSlice<f32>,
7814        up: &cudarc::driver::CudaView<f32>,
7815        act: &mut CudaSlice<f32>,
7816        ncols: usize,
7817        nrows: usize,
7818    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7819        debug_assert!(ncols % 128 == 0);
7820        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7821        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7822        let nc = ncols as i32;
7823        if Self::pdl_on() {
7824            {
7825                use cudarc::driver::{DevicePtr, DevicePtrMut};
7826                let s = &self.gpu.stream();
7827                let (pg, _g0) = gate.device_ptr(s);
7828                let (pu, _g1) = up.device_ptr(s);
7829                let (pact, _g2) = act.device_ptr_mut(s);
7830                let (pq, _g3) = out_q.device_ptr_mut(s);
7831                let (pd, _g4) = out_d.device_ptr_mut(s);
7832                let mut ps = [
7833                    &pg as *const _ as *mut std::ffi::c_void,
7834                    &pu as *const _ as *mut _,
7835                    &pact as *const _ as *mut _,
7836                    &pq as *const _ as *mut _,
7837                    &pd as *const _ as *mut _,
7838                    &nc as *const _ as *mut _,
7839                ];
7840                unsafe {
7841                    self.launch_pdl(
7842                        "gelu_tanh_mul_q8_1",
7843                        (nrows as u32, 1, 1),
7844                        (rms_block(), 1, 1),
7845                        &mut ps,
7846                    )?;
7847                }
7848            }
7849            return Ok((out_q, out_d));
7850        }
7851        let f = self.func("gelu_tanh_mul_q8_1");
7852        let cfg = LaunchConfig {
7853            grid_dim: (nrows as u32, 1, 1),
7854            block_dim: (rms_block(), 1, 1),
7855            shared_mem_bytes: 0,
7856        };
7857        let __s_b = self.gpu.stream();
7858        let mut b = __s_b.launch_builder(&f);
7859        b.arg(gate)
7860            .arg(up)
7861            .arg(act)
7862            .arg(&mut out_q)
7863            .arg(&mut out_d)
7864            .arg(&nc);
7865        unsafe {
7866            b.launch(cfg)?;
7867        }
7868        Ok((out_q, out_d))
7869    }
7870
7871    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
7872    #[allow(clippy::too_many_arguments)]
7873    pub fn gelu_tanh_mul_q8_1_into(
7874        &self,
7875        gate: &CudaSlice<f32>,
7876        up: &cudarc::driver::CudaView<f32>,
7877        act: &mut CudaSlice<f32>,
7878        ncols: usize,
7879        nrows: usize,
7880        out_q: &mut CudaSlice<i8>,
7881        out_d: &mut CudaSlice<f32>,
7882    ) -> Result<(), Box<dyn std::error::Error>> {
7883        debug_assert!(ncols % 128 == 0);
7884        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7885        let nc = ncols as i32;
7886        if Self::pdl_on() {
7887            use cudarc::driver::{DevicePtr, DevicePtrMut};
7888            let s = &self.gpu.stream();
7889            let (pg, _g0) = gate.device_ptr(s);
7890            let (pu, _g1) = up.device_ptr(s);
7891            let (pact, _g2) = act.device_ptr_mut(s);
7892            let (pq, _g3) = out_q.device_ptr_mut(s);
7893            let (pd, _g4) = out_d.device_ptr_mut(s);
7894            let mut ps = [
7895                &pg as *const _ as *mut std::ffi::c_void,
7896                &pu as *const _ as *mut _,
7897                &pact as *const _ as *mut _,
7898                &pq as *const _ as *mut _,
7899                &pd as *const _ as *mut _,
7900                &nc as *const _ as *mut _,
7901            ];
7902            unsafe {
7903                self.launch_pdl(
7904                    "gelu_tanh_mul_q8_1",
7905                    (nrows as u32, 1, 1),
7906                    (rms_block(), 1, 1),
7907                    &mut ps,
7908                )?;
7909            }
7910            return Ok(());
7911        }
7912        let f = self.func("gelu_tanh_mul_q8_1");
7913        let cfg = LaunchConfig {
7914            grid_dim: (nrows as u32, 1, 1),
7915            block_dim: (rms_block(), 1, 1),
7916            shared_mem_bytes: 0,
7917        };
7918        let __s_b = self.gpu.stream();
7919        let mut b = __s_b.launch_builder(&f);
7920        b.arg(gate)
7921            .arg(up)
7922            .arg(&mut *act)
7923            .arg(&mut *out_q)
7924            .arg(&mut *out_d)
7925            .arg(&nc);
7926        unsafe {
7927            b.launch(cfg)?;
7928        }
7929        Ok(())
7930    }
7931
7932    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
7933    #[allow(clippy::too_many_arguments)]
7934    pub fn add_rms_norm3_q8z(
7935        &self,
7936        a: &CudaSlice<f32>,
7937        b_in: &CudaSlice<f32>,
7938        w0: &CudaSlice<f32>,
7939        w1: &CudaSlice<f32>,
7940        w2: &CudaSlice<f32>,
7941        res: &mut CudaSlice<f32>,
7942        out1: &mut CudaSlice<f32>,
7943        ncols: usize,
7944        nrows: usize,
7945        eps: f32,
7946    ) -> Result<
7947        (
7948            (CudaSlice<i8>, CudaSlice<f32>),
7949            (CudaSlice<i8>, CudaSlice<f32>),
7950        ),
7951        Box<dyn std::error::Error>,
7952    > {
7953        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
7954        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7955        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
7956        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7957        let f = self.func("add_rms_norm3_q8z_f32");
7958        let cfg = LaunchConfig {
7959            grid_dim: (nrows as u32, 1, 1),
7960            block_dim: (rms_block(), 1, 1),
7961            shared_mem_bytes: 0,
7962        };
7963        let (nc, e2) = (ncols as i32, eps);
7964        let __s_b = self.gpu.stream();
7965        let mut b = __s_b.launch_builder(&f);
7966        b.arg(a)
7967            .arg(b_in)
7968            .arg(w0)
7969            .arg(w1)
7970            .arg(w2)
7971            .arg(res)
7972            .arg(&mut q0)
7973            .arg(&mut d0)
7974            .arg(out1)
7975            .arg(&mut q2)
7976            .arg(&mut d2)
7977            .arg(&nc)
7978            .arg(&e2);
7979        unsafe {
7980            b.launch(cfg)?;
7981        }
7982        Ok(((q0, d0), (q2, d2)))
7983    }
7984
7985    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
7986    #[allow(clippy::too_many_arguments)]
7987    pub fn add_rms_norm3(
7988        &self,
7989        a: &CudaSlice<f32>,
7990        b_in: &CudaSlice<f32>,
7991        w0: &CudaSlice<f32>,
7992        w1: &CudaSlice<f32>,
7993        w2: &CudaSlice<f32>,
7994        res: &mut CudaSlice<f32>,
7995        d0: &mut CudaSlice<f32>,
7996        d1: &mut CudaSlice<f32>,
7997        d2: &mut CudaSlice<f32>,
7998        ncols: usize,
7999        nrows: usize,
8000        eps: f32,
8001    ) -> Result<(), Box<dyn std::error::Error>> {
8002        let f = self.func("add_rms_norm3_f32");
8003        let cfg = LaunchConfig {
8004            grid_dim: (nrows as u32, 1, 1),
8005            block_dim: (rms_block(), 1, 1),
8006            shared_mem_bytes: 0,
8007        };
8008        let (nc, e2) = (ncols as i32, eps);
8009        let __s_b = self.gpu.stream();
8010        let mut b = __s_b.launch_builder(&f);
8011        b.arg(a)
8012            .arg(b_in)
8013            .arg(w0)
8014            .arg(w1)
8015            .arg(w2)
8016            .arg(res)
8017            .arg(d0)
8018            .arg(d1)
8019            .arg(d2)
8020            .arg(&nc)
8021            .arg(&e2);
8022        unsafe {
8023            b.launch(cfg)?;
8024        }
8025        Ok(())
8026    }
8027
8028    /// dst = (a + b) * c (residual add + layer scale, one launch).
8029    pub fn add_scale(
8030        &self,
8031        a: &CudaSlice<f32>,
8032        b_in: &CudaSlice<f32>,
8033        c: f32,
8034        dst: &mut CudaSlice<f32>,
8035        n: usize,
8036    ) -> Result<(), Box<dyn std::error::Error>> {
8037        let f = self.func("add_scale_f32");
8038        let cfg = LaunchConfig::for_num_elems(n as u32);
8039        let ni = n as i32;
8040        let __s_b = self.gpu.stream();
8041        let mut b = __s_b.launch_builder(&f);
8042        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8043        unsafe {
8044            b.launch(cfg)?;
8045        }
8046        Ok(())
8047    }
8048
8049    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
8050    pub fn layer_norm_bias(
8051        &self,
8052        x: &CudaSlice<f32>,
8053        w: &CudaSlice<f32>,
8054        b: &CudaSlice<f32>,
8055        dst: &mut CudaSlice<f32>,
8056        ncols: usize,
8057        nrows: usize,
8058        eps: f32,
8059    ) -> Result<(), Box<dyn std::error::Error>> {
8060        let f = self.func("layer_norm_bias_f32");
8061        let (nc, e) = (ncols as i32, eps);
8062        let cfg = LaunchConfig {
8063            grid_dim: (nrows as u32, 1, 1),
8064            block_dim: (256, 1, 1),
8065            shared_mem_bytes: 0,
8066        };
8067        let __s_b = self.gpu.stream();
8068        let mut lb = __s_b.launch_builder(&f);
8069        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8070        unsafe {
8071            lb.launch(cfg)?;
8072        }
8073        Ok(())
8074    }
8075
8076    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
8077    pub fn gelu_tanh(
8078        &self,
8079        x: &CudaSlice<f32>,
8080        dst: &mut CudaSlice<f32>,
8081        n: usize,
8082    ) -> Result<(), Box<dyn std::error::Error>> {
8083        let f = self.func("gelu_tanh_f32");
8084        let ni = n as i64;
8085        let cfg = LaunchConfig {
8086            grid_dim: (n.div_ceil(256) as u32, 1, 1),
8087            block_dim: (256, 1, 1),
8088            shared_mem_bytes: 0,
8089        };
8090        let __s_b = self.gpu.stream();
8091        let mut lb = __s_b.launch_builder(&f);
8092        lb.arg(x).arg(&mut *dst).arg(&ni);
8093        unsafe {
8094            lb.launch(cfg)?;
8095        }
8096        Ok(())
8097    }
8098
8099    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
8100    pub fn row_softmax(
8101        &self,
8102        x: &mut CudaSlice<f32>,
8103        ncols: usize,
8104        nrows: usize,
8105    ) -> Result<(), Box<dyn std::error::Error>> {
8106        let f = self.func("row_softmax_f32");
8107        let nc = ncols as i32;
8108        let cfg = LaunchConfig {
8109            grid_dim: (nrows as u32, 1, 1),
8110            block_dim: (256, 1, 1),
8111            shared_mem_bytes: 0,
8112        };
8113        let __s_b = self.gpu.stream();
8114        let mut lb = __s_b.launch_builder(&f);
8115        lb.arg(&mut *x).arg(&nc);
8116        unsafe {
8117            lb.launch(cfg)?;
8118        }
8119        Ok(())
8120    }
8121
8122    pub fn rms_norm(
8123        &self,
8124        x: &CudaSlice<f32>,
8125        w: &CudaSlice<f32>,
8126        dst: &mut CudaSlice<f32>,
8127        ncols: usize,
8128        nrows: usize,
8129        eps: f32,
8130    ) -> Result<(), Box<dyn std::error::Error>> {
8131        let (nc, e) = (ncols as i32, eps);
8132        if Self::pdl_on() && Self::pdl_wb_on() {
8133            use cudarc::driver::{DevicePtr, DevicePtrMut};
8134            let s = &self.gpu.stream();
8135            let (px, _g0) = x.device_ptr(s);
8136            let (pw, _g1) = w.device_ptr(s);
8137            let (pd, _g2) = dst.device_ptr_mut(s);
8138            let mut ps = [
8139                &px as *const _ as *mut std::ffi::c_void,
8140                &pw as *const _ as *mut _,
8141                &pd as *const _ as *mut _,
8142                &nc as *const _ as *mut _,
8143                &e as *const _ as *mut _,
8144            ];
8145            unsafe {
8146                self.launch_pdl(
8147                    "rms_norm_f32",
8148                    (nrows as u32, 1, 1),
8149                    (rms_block(), 1, 1),
8150                    &mut ps,
8151                )?;
8152            }
8153            return Ok(());
8154        }
8155        let f = self.func("rms_norm_f32");
8156        let cfg = LaunchConfig {
8157            grid_dim: (nrows as u32, 1, 1),
8158            block_dim: (rms_block(), 1, 1),
8159            shared_mem_bytes: 0,
8160        };
8161        let __s_b = self.gpu.stream();
8162        let mut b = __s_b.launch_builder(&f);
8163        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8164        unsafe {
8165            b.launch(cfg)?;
8166        }
8167        Ok(())
8168    }
8169
8170    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
8171    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
8172    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
8173    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
8174    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
8175    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
8176    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
8177    pub fn rms_norm_decode(
8178        &self,
8179        x: &CudaSlice<f32>,
8180        w: &CudaSlice<f32>,
8181        dst: &mut CudaSlice<f32>,
8182        ncols: usize,
8183        nrows: usize,
8184        eps: f32,
8185    ) -> Result<(), Box<dyn std::error::Error>> {
8186        let f = self.func("rms_norm_f32");
8187        let cfg = LaunchConfig {
8188            grid_dim: (nrows as u32, 1, 1),
8189            block_dim: (1024, 1, 1),
8190            shared_mem_bytes: 0,
8191        };
8192        let (nc, e) = (ncols as i32, eps);
8193        let __s_b = self.gpu.stream();
8194        let mut b = __s_b.launch_builder(&f);
8195        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8196        unsafe {
8197            b.launch(cfg)?;
8198        }
8199        Ok(())
8200    }
8201
8202    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
8203    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
8204    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
8205    pub fn rms_norm_q8_1(
8206        &self,
8207        x: &CudaSlice<f32>,
8208        w: &CudaSlice<f32>,
8209        ncols: usize,
8210        nrows: usize,
8211        eps: f32,
8212    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8213        let nblk = ncols / 32;
8214        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8215        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8216        let (nc, e) = (ncols as i32, eps);
8217        if Self::pdl_on() {
8218            {
8219                use cudarc::driver::{DevicePtr, DevicePtrMut};
8220                let s = &self.gpu.stream();
8221                let (px, _g0) = x.device_ptr(s);
8222                let (pw, _g1) = w.device_ptr(s);
8223                let (pq, _g2) = q.device_ptr_mut(s);
8224                let (pd, _g3) = d.device_ptr_mut(s);
8225                let mut ps = [
8226                    &px as *const _ as *mut std::ffi::c_void,
8227                    &pw as *const _ as *mut _,
8228                    &pq as *const _ as *mut _,
8229                    &pd as *const _ as *mut _,
8230                    &nc as *const _ as *mut _,
8231                    &e as *const _ as *mut _,
8232                ];
8233                unsafe {
8234                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8235                }
8236            }
8237            return Ok((q, d));
8238        }
8239        let f = self.func("rms_norm_q8_1");
8240        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
8241        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
8242        let cfg = LaunchConfig {
8243            grid_dim: (nrows as u32, 1, 1),
8244            block_dim: (1024, 1, 1),
8245            shared_mem_bytes: 0,
8246        };
8247        let __s_b = self.gpu.stream();
8248        let mut b = __s_b.launch_builder(&f);
8249        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8250        unsafe {
8251            b.launch(cfg)?;
8252        }
8253        Ok((q, d))
8254    }
8255
8256    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
8257    /// PDL arm), caller-owned outputs.
8258    pub fn rms_norm_q8_1_into(
8259        &self,
8260        x: &CudaSlice<f32>,
8261        w: &CudaSlice<f32>,
8262        ncols: usize,
8263        nrows: usize,
8264        eps: f32,
8265        q: &mut CudaSlice<i8>,
8266        d: &mut CudaSlice<f32>,
8267    ) -> Result<(), Box<dyn std::error::Error>> {
8268        let nblk = ncols / 32;
8269        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8270        let (nc, e) = (ncols as i32, eps);
8271        if Self::pdl_on() {
8272            use cudarc::driver::{DevicePtr, DevicePtrMut};
8273            let s = &self.gpu.stream();
8274            let (px, _g0) = x.device_ptr(s);
8275            let (pw, _g1) = w.device_ptr(s);
8276            let (pq, _g2) = q.device_ptr_mut(s);
8277            let (pd, _g3) = d.device_ptr_mut(s);
8278            let mut ps = [
8279                &px as *const _ as *mut std::ffi::c_void,
8280                &pw as *const _ as *mut _,
8281                &pq as *const _ as *mut _,
8282                &pd as *const _ as *mut _,
8283                &nc as *const _ as *mut _,
8284                &e as *const _ as *mut _,
8285            ];
8286            unsafe {
8287                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8288            }
8289            return Ok(());
8290        }
8291        let f = self.func("rms_norm_q8_1");
8292        let cfg = LaunchConfig {
8293            grid_dim: (nrows as u32, 1, 1),
8294            block_dim: (1024, 1, 1),
8295            shared_mem_bytes: 0,
8296        };
8297        let __s_b = self.gpu.stream();
8298        let mut b = __s_b.launch_builder(&f);
8299        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8300        unsafe {
8301            b.launch(cfg)?;
8302        }
8303        Ok(())
8304    }
8305
8306    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
8307    pub fn quantize_q8_1_into(
8308        &self,
8309        x: &CudaSlice<f32>,
8310        m: usize,
8311        in_f: usize,
8312        q: &mut CudaSlice<i8>,
8313        d: &mut CudaSlice<f32>,
8314    ) -> Result<(), Box<dyn std::error::Error>> {
8315        let nblk = in_f / 32;
8316        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8317        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8318        let (inf, mi) = (in_f as i32, m as i32);
8319        if Self::pdl_on() && Self::pdl_wb_on() {
8320            use cudarc::driver::{DevicePtr, DevicePtrMut};
8321            let s = &self.gpu.stream();
8322            let (px, _g0) = x.device_ptr(s);
8323            let (pq, _g1) = q.device_ptr_mut(s);
8324            let (pd, _g2) = d.device_ptr_mut(s);
8325            let mut ps = [
8326                &px as *const _ as *mut std::ffi::c_void,
8327                &pq as *const _ as *mut _,
8328                &pd as *const _ as *mut _,
8329                &inf as *const _ as *mut _,
8330                &mi as *const _ as *mut _,
8331            ];
8332            unsafe {
8333                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8334            }
8335            return Ok(());
8336        }
8337        let f = self.func("quantize_q8_1");
8338        let __s_b = self.gpu.stream();
8339        let mut b = __s_b.launch_builder(&f);
8340        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8341        unsafe {
8342            b.launch(cfg)?;
8343        }
8344        Ok(())
8345    }
8346
8347    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
8348    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
8349    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
8350    pub fn add_rms_norm_q8_1(
8351        &self,
8352        a: &CudaSlice<f32>,
8353        b_in: &CudaSlice<f32>,
8354        w: &CudaSlice<f32>,
8355        res: &mut CudaSlice<f32>,
8356        ncols: usize,
8357        nrows: usize,
8358        eps: f32,
8359    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8360        let nblk = ncols / 32;
8361        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8362        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8363        let f = self.func("add_rms_norm_q8_1");
8364        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
8365        let cfg = LaunchConfig {
8366            grid_dim: (nrows as u32, 1, 1),
8367            block_dim: (1024, 1, 1),
8368            shared_mem_bytes: 0,
8369        };
8370        let (nc, e) = (ncols as i32, eps);
8371        let __s_bld = self.gpu.stream();
8372        let mut bld = __s_bld.launch_builder(&f);
8373        bld.arg(a)
8374            .arg(b_in)
8375            .arg(w)
8376            .arg(res)
8377            .arg(&mut q)
8378            .arg(&mut d)
8379            .arg(&nc)
8380            .arg(&e);
8381        unsafe {
8382            bld.launch(cfg)?;
8383        }
8384        Ok((q, d))
8385    }
8386
8387    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
8388    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
8389    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
8390    pub fn add_rms_norm(
8391        &self,
8392        a: &CudaSlice<f32>,
8393        b: &CudaSlice<f32>,
8394        w: &CudaSlice<f32>,
8395        res: &mut CudaSlice<f32>,
8396        dst: &mut CudaSlice<f32>,
8397        ncols: usize,
8398        nrows: usize,
8399        eps: f32,
8400    ) -> Result<(), Box<dyn std::error::Error>> {
8401        let (nc, e) = (ncols as i32, eps);
8402        if Self::pdl_on() && Self::pdl_wb_on() {
8403            use cudarc::driver::{DevicePtr, DevicePtrMut};
8404            let s = &self.gpu.stream();
8405            let (pa, _g0) = a.device_ptr(s);
8406            let (pb, _g1) = b.device_ptr(s);
8407            let (pw, _g2) = w.device_ptr(s);
8408            let (pr, _g3) = res.device_ptr_mut(s);
8409            let (pd, _g4) = dst.device_ptr_mut(s);
8410            let mut ps = [
8411                &pa as *const _ as *mut std::ffi::c_void,
8412                &pb as *const _ as *mut _,
8413                &pw as *const _ as *mut _,
8414                &pr as *const _ as *mut _,
8415                &pd as *const _ as *mut _,
8416                &nc as *const _ as *mut _,
8417                &e as *const _ as *mut _,
8418            ];
8419            unsafe {
8420                self.launch_pdl(
8421                    "add_rms_norm_f32",
8422                    (nrows as u32, 1, 1),
8423                    (rms_block(), 1, 1),
8424                    &mut ps,
8425                )?;
8426            }
8427            return Ok(());
8428        }
8429        let f = self.func("add_rms_norm_f32");
8430        let cfg = LaunchConfig {
8431            grid_dim: (nrows as u32, 1, 1),
8432            block_dim: (rms_block(), 1, 1),
8433            shared_mem_bytes: 0,
8434        };
8435        let __s_b2 = self.gpu.stream();
8436        let mut b2 = __s_b2.launch_builder(&f);
8437        b2.arg(a)
8438            .arg(b)
8439            .arg(w)
8440            .arg(&mut *res)
8441            .arg(&mut *dst)
8442            .arg(&nc)
8443            .arg(&e);
8444        unsafe {
8445            b2.launch(cfg)?;
8446        }
8447        Ok(())
8448    }
8449
8450    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
8451    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
8452    #[allow(clippy::too_many_arguments)]
8453    pub fn rms_pre_add_rms_norm(
8454        &self,
8455        a: &CudaSlice<f32>,
8456        wa: &CudaSlice<f32>,
8457        b: &CudaSlice<f32>,
8458        w: &CudaSlice<f32>,
8459        res: &mut CudaSlice<f32>,
8460        dst: &mut CudaSlice<f32>,
8461        ncols: usize,
8462        nrows: usize,
8463        eps: f32,
8464    ) -> Result<(), Box<dyn std::error::Error>> {
8465        let f = self.func("rms_pre_add_rms_norm_f32");
8466        let cfg = LaunchConfig {
8467            grid_dim: (nrows as u32, 1, 1),
8468            block_dim: (rms_block(), 1, 1),
8469            shared_mem_bytes: 0,
8470        };
8471        let (nc, e) = (ncols as i32, eps);
8472        let __s_b2 = self.gpu.stream();
8473        let mut b2 = __s_b2.launch_builder(&f);
8474        b2.arg(a)
8475            .arg(wa)
8476            .arg(b)
8477            .arg(w)
8478            .arg(&mut *res)
8479            .arg(&mut *dst)
8480            .arg(&nc)
8481            .arg(&e);
8482        unsafe {
8483            b2.launch(cfg)?;
8484        }
8485        Ok(())
8486    }
8487
8488    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
8489    #[allow(clippy::too_many_arguments)]
8490    pub fn rms_pre_add_rms_norm_q8z(
8491        &self,
8492        a: &CudaSlice<f32>,
8493        wa: &CudaSlice<f32>,
8494        b: &CudaSlice<f32>,
8495        w: &CudaSlice<f32>,
8496        res: &mut CudaSlice<f32>,
8497        dst: &mut CudaSlice<f32>,
8498        ncols: usize,
8499        nrows: usize,
8500        eps: f32,
8501    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8502        debug_assert!(ncols % 128 == 0);
8503        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8504        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8505        let (nc, e) = (ncols as i32, eps);
8506        if Self::pdl_on() {
8507            {
8508                use cudarc::driver::{DevicePtr, DevicePtrMut};
8509                let s = &self.gpu.stream();
8510                let (pa, _g0) = a.device_ptr(s);
8511                let (pwa, _g1) = wa.device_ptr(s);
8512                let (pb, _g2) = b.device_ptr(s);
8513                let (pw, _g3) = w.device_ptr(s);
8514                let (pr, _g4) = res.device_ptr_mut(s);
8515                let (pdst, _g5) = dst.device_ptr_mut(s);
8516                let (pq, _g6) = out_q.device_ptr_mut(s);
8517                let (pd, _g7) = out_d.device_ptr_mut(s);
8518                let mut ps = [
8519                    &pa as *const _ as *mut std::ffi::c_void,
8520                    &pwa as *const _ as *mut _,
8521                    &pb as *const _ as *mut _,
8522                    &pw as *const _ as *mut _,
8523                    &pr as *const _ as *mut _,
8524                    &pdst as *const _ as *mut _,
8525                    &pq as *const _ as *mut _,
8526                    &pd as *const _ as *mut _,
8527                    &nc as *const _ as *mut _,
8528                    &e as *const _ as *mut _,
8529                ];
8530                unsafe {
8531                    self.launch_pdl(
8532                        "rms_pre_add_rms_norm_q8z_f32",
8533                        (nrows as u32, 1, 1),
8534                        (rms_block(), 1, 1),
8535                        &mut ps,
8536                    )?;
8537                }
8538            }
8539            return Ok((out_q, out_d));
8540        }
8541        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8542        let cfg = LaunchConfig {
8543            grid_dim: (nrows as u32, 1, 1),
8544            block_dim: (rms_block(), 1, 1),
8545            shared_mem_bytes: 0,
8546        };
8547        let __s_b2 = self.gpu.stream();
8548        let mut b2 = __s_b2.launch_builder(&f);
8549        b2.arg(a)
8550            .arg(wa)
8551            .arg(b)
8552            .arg(w)
8553            .arg(&mut *res)
8554            .arg(&mut *dst)
8555            .arg(&mut out_q)
8556            .arg(&mut out_d)
8557            .arg(&nc)
8558            .arg(&e);
8559        unsafe {
8560            b2.launch(cfg)?;
8561        }
8562        Ok((out_q, out_d))
8563    }
8564
8565    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
8566    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
8567    /// body must stay attribute-free (the fused2_into precedent).
8568    #[allow(clippy::too_many_arguments)]
8569    pub fn rms_pre_add_rms_norm_q8z_into(
8570        &self,
8571        a: &CudaSlice<f32>,
8572        wa: &CudaSlice<f32>,
8573        b: &CudaSlice<f32>,
8574        w: &CudaSlice<f32>,
8575        res: &mut CudaSlice<f32>,
8576        dst: &mut CudaSlice<f32>,
8577        ncols: usize,
8578        nrows: usize,
8579        eps: f32,
8580        out_q: &mut CudaSlice<i8>,
8581        out_d: &mut CudaSlice<f32>,
8582    ) -> Result<(), Box<dyn std::error::Error>> {
8583        debug_assert!(ncols % 128 == 0);
8584        let (nc, e) = (ncols as i32, eps);
8585        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8586        let cfg = LaunchConfig {
8587            grid_dim: (nrows as u32, 1, 1),
8588            block_dim: (rms_block(), 1, 1),
8589            shared_mem_bytes: 0,
8590        };
8591        let __s_b = self.gpu.stream();
8592        let mut b2 = __s_b.launch_builder(&f);
8593        b2.arg(a)
8594            .arg(wa)
8595            .arg(b)
8596            .arg(w)
8597            .arg(&mut *res)
8598            .arg(&mut *dst)
8599            .arg(&mut *out_q)
8600            .arg(&mut *out_d)
8601            .arg(&nc)
8602            .arg(&e);
8603        unsafe {
8604            b2.launch(cfg)?;
8605        }
8606        Ok(())
8607    }
8608
8609    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
8610    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
8611    #[allow(clippy::too_many_arguments)]
8612    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8613        &self,
8614        a: &CudaSlice<f32>,
8615        wa: &CudaSlice<f32>,
8616        b_in: &CudaSlice<f32>,
8617        c: f32,
8618        w: &CudaSlice<f32>,
8619        res: &mut CudaSlice<f32>,
8620        ncols: usize,
8621        nrows: usize,
8622        eps: f32,
8623        out_q: &mut CudaSlice<i8>,
8624        out_d: &mut CudaSlice<f32>,
8625    ) -> Result<(), Box<dyn std::error::Error>> {
8626        debug_assert!(ncols % 128 == 0);
8627        let (nc, e2) = (ncols as i32, eps);
8628        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8629        let cfg = LaunchConfig {
8630            grid_dim: (nrows as u32, 1, 1),
8631            block_dim: (rms_block(), 1, 1),
8632            shared_mem_bytes: 0,
8633        };
8634        let __s_b = self.gpu.stream();
8635        let mut b2 = __s_b.launch_builder(&f);
8636        b2.arg(a)
8637            .arg(wa)
8638            .arg(b_in)
8639            .arg(&c)
8640            .arg(w)
8641            .arg(&mut *res)
8642            .arg(&mut *out_q)
8643            .arg(&mut *out_d)
8644            .arg(&nc)
8645            .arg(&e2);
8646        unsafe {
8647            b2.launch(cfg)?;
8648        }
8649        Ok(())
8650    }
8651
8652    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
8653    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
8654    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
8655    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
8656    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
8657    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
8658    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
8659    pub fn g4_pnfold_on() -> bool {
8660        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8661        *ON.get_or_init(|| {
8662            std::env::var("MEMRA_G4_PNFOLD")
8663                .map(|v| v != "0")
8664                .unwrap_or(true)
8665        })
8666    }
8667
8668    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
8669    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
8670    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
8671    pub fn build_q4_out_concat3(
8672        &self,
8673        w0: &crate::model::GpuTensor,
8674        w1: &crate::model::GpuTensor,
8675        w2: &crate::model::GpuTensor,
8676    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8677        use crate::model::GpuTensor;
8678        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8679            match w {
8680                GpuTensor::Quant {
8681                    qtype,
8682                    row_bytes,
8683                    rp,
8684                    ..
8685                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8686                _ => None,
8687            }
8688        };
8689        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8690        else {
8691            return Ok(None);
8692        };
8693        if rb0 != rb1
8694            || rb0 != rb2
8695            || w0.in_features() != w1.in_features()
8696            || w0.in_features() != w2.in_features()
8697        {
8698            return Ok(None);
8699        }
8700        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8701            match w {
8702                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8703                _ => unreachable!(),
8704            }
8705        }
8706        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8707        let total = rb0 * (o0 + o1 + o2);
8708        let mut cat = self.alloc_u8(total)?;
8709        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8710        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8711        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8712        Ok(Some(GpuTensor::Quant {
8713            bytes: cat,
8714            qtype: QT_Q4_0,
8715            row_bytes: rb0,
8716            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8717            scale: 1.0,
8718            rp: false,
8719            #[cfg(memra_cutlass)]
8720            cutlass: None,
8721            fp8: None,
8722            blk: None,
8723            rp4: None,
8724            f16: None,
8725        }))
8726    }
8727
8728    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
8729    #[allow(clippy::too_many_arguments)]
8730    pub fn rms_norm_qkv_rope_cat(
8731        &self,
8732        qkv: &CudaSlice<f32>,
8733        wq: &CudaSlice<f32>,
8734        wk: &CudaSlice<f32>,
8735        wv: &CudaSlice<f32>,
8736        q: &mut CudaSlice<f32>,
8737        k: &mut CudaSlice<f32>,
8738        v: &mut CudaSlice<f32>,
8739        head_dim: usize,
8740        rq: usize,
8741        rk: usize,
8742        pos: &CudaSlice<i32>,
8743        nh_q: usize,
8744        nh_k: usize,
8745        base: f32,
8746        freq_scale: f32,
8747        ff: Option<&CudaSlice<f32>>,
8748        eps: f32,
8749    ) -> Result<(), Box<dyn std::error::Error>> {
8750        let rows = rq + rk + rk;
8751        let theta_scale = base.powf(-2.0 / head_dim as f32);
8752        let (nc, rqi, rki, nhq, nhk) = (
8753            head_dim as i32,
8754            rq as i32,
8755            rk as i32,
8756            nh_q as i32,
8757            nh_k as i32,
8758        );
8759        if Self::pdl_on() {
8760            use cudarc::driver::{DevicePtr, DevicePtrMut};
8761            let s = &self.gpu.stream();
8762            let (pqkv, _g0) = qkv.device_ptr(s);
8763            let (pwq, _g1) = wq.device_ptr(s);
8764            let (pwk, _g2) = wk.device_ptr(s);
8765            let (pwv, _g3) = wv.device_ptr(s);
8766            let (pq, _g4) = q.device_ptr_mut(s);
8767            let (pk, _g5) = k.device_ptr_mut(s);
8768            let (pv, _g6) = v.device_ptr_mut(s);
8769            let (ppos, _g7) = pos.device_ptr(s);
8770            let (pff, _g8) = match ff {
8771                Some(t) => {
8772                    let (p, g) = t.device_ptr(s);
8773                    (p, Some(g))
8774                }
8775                None => (0, None),
8776            };
8777            let mut ps = [
8778                &pqkv as *const _ as *mut std::ffi::c_void,
8779                &pwq as *const _ as *mut _,
8780                &pwk as *const _ as *mut _,
8781                &pwv as *const _ as *mut _,
8782                &pq as *const _ as *mut _,
8783                &pk as *const _ as *mut _,
8784                &pv as *const _ as *mut _,
8785                &nc as *const _ as *mut _,
8786                &rqi as *const _ as *mut _,
8787                &rki as *const _ as *mut _,
8788                &ppos as *const _ as *mut _,
8789                &nhq as *const _ as *mut _,
8790                &nhk as *const _ as *mut _,
8791                &theta_scale as *const _ as *mut _,
8792                &freq_scale as *const _ as *mut _,
8793                &pff as *const _ as *mut _,
8794                &eps as *const _ as *mut _,
8795            ];
8796            unsafe {
8797                self.launch_pdl(
8798                    "rms_norm_qkv_rope_cat_f32",
8799                    (rows as u32, 1, 1),
8800                    (rms_block(), 1, 1),
8801                    &mut ps,
8802                )?;
8803            }
8804            return Ok(());
8805        }
8806        let f = self.func("rms_norm_qkv_rope_cat_f32");
8807        let cfg = LaunchConfig {
8808            grid_dim: (rows as u32, 1, 1),
8809            block_dim: (rms_block(), 1, 1),
8810            shared_mem_bytes: 0,
8811        };
8812        let __s_b = self.gpu.stream();
8813        let mut b = __s_b.launch_builder(&f);
8814        match ff {
8815            Some(t) => {
8816                b.arg(qkv)
8817                    .arg(wq)
8818                    .arg(wk)
8819                    .arg(wv)
8820                    .arg(&mut *q)
8821                    .arg(&mut *k)
8822                    .arg(&mut *v)
8823                    .arg(&nc)
8824                    .arg(&rqi)
8825                    .arg(&rki)
8826                    .arg(pos)
8827                    .arg(&nhq)
8828                    .arg(&nhk)
8829                    .arg(&theta_scale)
8830                    .arg(&freq_scale)
8831                    .arg(t)
8832                    .arg(&eps);
8833                unsafe {
8834                    b.launch(cfg)?;
8835                }
8836            }
8837            None => {
8838                let null: u64 = 0;
8839                b.arg(qkv)
8840                    .arg(wq)
8841                    .arg(wk)
8842                    .arg(wv)
8843                    .arg(&mut *q)
8844                    .arg(&mut *k)
8845                    .arg(&mut *v)
8846                    .arg(&nc)
8847                    .arg(&rqi)
8848                    .arg(&rki)
8849                    .arg(pos)
8850                    .arg(&nhq)
8851                    .arg(&nhk)
8852                    .arg(&theta_scale)
8853                    .arg(&freq_scale)
8854                    .arg(&null)
8855                    .arg(&eps);
8856                unsafe {
8857                    b.launch(cfg)?;
8858                }
8859            }
8860        }
8861        Ok(())
8862    }
8863
8864    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
8865    #[allow(clippy::too_many_arguments)]
8866    pub fn rms_norm_qkv_rope(
8867        &self,
8868        q0: &CudaSlice<f32>,
8869        k0: &CudaSlice<f32>,
8870        v0: &CudaSlice<f32>,
8871        wq: &CudaSlice<f32>,
8872        wk: &CudaSlice<f32>,
8873        wv: &CudaSlice<f32>,
8874        q: &mut CudaSlice<f32>,
8875        k: &mut CudaSlice<f32>,
8876        v: &mut CudaSlice<f32>,
8877        head_dim: usize,
8878        rq: usize,
8879        rk: usize,
8880        pos: &CudaSlice<i32>,
8881        nh_q: usize,
8882        nh_k: usize,
8883        base: f32,
8884        freq_scale: f32,
8885        ff: Option<&CudaSlice<f32>>,
8886        eps: f32,
8887    ) -> Result<(), Box<dyn std::error::Error>> {
8888        let f = self.func("rms_norm_qkv_rope_f32");
8889        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
8890        let cfg = LaunchConfig {
8891            grid_dim: (rows as u32, 1, 1),
8892            block_dim: (rms_block(), 1, 1),
8893            shared_mem_bytes: 0,
8894        };
8895        let theta_scale = base.powf(-2.0 / head_dim as f32);
8896        let (nc, rqi, rki, nhq, nhk) = (
8897            head_dim as i32,
8898            rq as i32,
8899            rk as i32,
8900            nh_q as i32,
8901            nh_k as i32,
8902        );
8903        let __s_b = self.gpu.stream();
8904        let mut b = __s_b.launch_builder(&f);
8905        match ff {
8906            Some(t) => {
8907                b.arg(q0)
8908                    .arg(k0)
8909                    .arg(v0)
8910                    .arg(wq)
8911                    .arg(wk)
8912                    .arg(wv)
8913                    .arg(&mut *q)
8914                    .arg(&mut *k)
8915                    .arg(&mut *v)
8916                    .arg(&nc)
8917                    .arg(&rqi)
8918                    .arg(&rki)
8919                    .arg(pos)
8920                    .arg(&nhq)
8921                    .arg(&nhk)
8922                    .arg(&theta_scale)
8923                    .arg(&freq_scale)
8924                    .arg(t)
8925                    .arg(&eps);
8926                unsafe {
8927                    b.launch(cfg)?;
8928                }
8929            }
8930            None => {
8931                let null: u64 = 0;
8932                b.arg(q0)
8933                    .arg(k0)
8934                    .arg(v0)
8935                    .arg(wq)
8936                    .arg(wk)
8937                    .arg(wv)
8938                    .arg(&mut *q)
8939                    .arg(&mut *k)
8940                    .arg(&mut *v)
8941                    .arg(&nc)
8942                    .arg(&rqi)
8943                    .arg(&rki)
8944                    .arg(pos)
8945                    .arg(&nhq)
8946                    .arg(&nhk)
8947                    .arg(&theta_scale)
8948                    .arg(&freq_scale)
8949                    .arg(&null)
8950                    .arg(&eps);
8951                unsafe {
8952                    b.launch(cfg)?;
8953                }
8954            }
8955        }
8956        Ok(())
8957    }
8958
8959    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
8960    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
8961    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
8962    #[allow(clippy::too_many_arguments)]
8963    pub fn rms_norm_qkv_rope_append_dc(
8964        &self,
8965        q0: &CudaSlice<f32>,
8966        k0: &CudaSlice<f32>,
8967        v0: &CudaSlice<f32>,
8968        wq: &CudaSlice<f32>,
8969        wk: &CudaSlice<f32>,
8970        wv: &CudaSlice<f32>,
8971        q: &mut CudaSlice<f32>,
8972        k: &mut CudaSlice<f32>,
8973        v: &mut CudaSlice<f32>,
8974        head_dim: usize,
8975        rq: usize,
8976        rk: usize,
8977        pos: &CudaSlice<i32>,
8978        nh_q: usize,
8979        nh_k: usize,
8980        base: f32,
8981        freq_scale: f32,
8982        ff: Option<&CudaSlice<f32>>,
8983        eps: f32,
8984        kc: &mut CudaSlice<u8>,
8985        vc: &mut CudaSlice<u8>,
8986        t_dev: &CudaSlice<i32>,
8987        k_tok_bytes: usize,
8988        v_tok_bytes: usize,
8989        g: bool,
8990    ) -> Result<(), Box<dyn std::error::Error>> {
8991        let rows = rq + rk + rk;
8992        let theta_scale = base.powf(-2.0 / head_dim as f32);
8993        let (nc, rqi, rki, nhq, nhk) = (
8994            head_dim as i32,
8995            rq as i32,
8996            rk as i32,
8997            nh_q as i32,
8998            nh_k as i32,
8999        );
9000        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9001        if Self::pdl_on() && Self::pdl_wb_on() {
9002            use cudarc::driver::{DevicePtr, DevicePtrMut};
9003            let s = &self.gpu.stream();
9004            let (p0, _a0) = q0.device_ptr(s);
9005            let (p1, _a1) = k0.device_ptr(s);
9006            let (p2, _a2) = v0.device_ptr(s);
9007            let (pwq, _a3) = wq.device_ptr(s);
9008            let (pwk, _a4) = wk.device_ptr(s);
9009            let (pwv, _a5) = wv.device_ptr(s);
9010            let (pq, _a6) = q.device_ptr_mut(s);
9011            let (pk, _a7) = k.device_ptr_mut(s);
9012            let (pv, _a8) = v.device_ptr_mut(s);
9013            let (pp, _a9) = pos.device_ptr(s);
9014            let pff: u64 = match ff {
9015                Some(t) => {
9016                    let (p, _gg) = t.device_ptr(s);
9017                    p as u64
9018                }
9019                None => 0,
9020            };
9021            let (pkc, _a10) = kc.device_ptr_mut(s);
9022            let (pvc, _a11) = vc.device_ptr_mut(s);
9023            let (pt, _a12) = t_dev.device_ptr(s);
9024            let mut ps = [
9025                &p0 as *const _ as *mut std::ffi::c_void,
9026                &p1 as *const _ as *mut _,
9027                &p2 as *const _ as *mut _,
9028                &pwq as *const _ as *mut _,
9029                &pwk as *const _ as *mut _,
9030                &pwv as *const _ as *mut _,
9031                &pq as *const _ as *mut _,
9032                &pk as *const _ as *mut _,
9033                &pv as *const _ as *mut _,
9034                &nc as *const _ as *mut _,
9035                &rqi as *const _ as *mut _,
9036                &rki as *const _ as *mut _,
9037                &pp as *const _ as *mut _,
9038                &nhq as *const _ as *mut _,
9039                &nhk as *const _ as *mut _,
9040                &theta_scale as *const _ as *mut _,
9041                &freq_scale as *const _ as *mut _,
9042                &pff as *const _ as *mut _,
9043                &eps as *const _ as *mut _,
9044                &pkc as *const _ as *mut _,
9045                &pvc as *const _ as *mut _,
9046                &pt as *const _ as *mut _,
9047                &ktb as *const _ as *mut _,
9048                &vtb as *const _ as *mut _,
9049            ];
9050            unsafe {
9051                self.launch_pdl_flash(
9052                    g,
9053                    "rms_norm_qkv_rope_append_dc_f32",
9054                    (rows as u32, 1, 1),
9055                    (rms_block(), 1, 1),
9056                    0,
9057                    &mut ps,
9058                )?;
9059            }
9060            return Ok(());
9061        }
9062        let f = if g {
9063            self.func_g("rms_norm_qkv_rope_append_dc_f32")
9064        } else {
9065            self.func("rms_norm_qkv_rope_append_dc_f32")
9066        };
9067        let cfg = LaunchConfig {
9068            grid_dim: (rows as u32, 1, 1),
9069            block_dim: (rms_block(), 1, 1),
9070            shared_mem_bytes: 0,
9071        };
9072        let __s_b = self.gpu.stream();
9073        let mut b = __s_b.launch_builder(&f);
9074        match ff {
9075            Some(t) => {
9076                b.arg(q0)
9077                    .arg(k0)
9078                    .arg(v0)
9079                    .arg(wq)
9080                    .arg(wk)
9081                    .arg(wv)
9082                    .arg(&mut *q)
9083                    .arg(&mut *k)
9084                    .arg(&mut *v)
9085                    .arg(&nc)
9086                    .arg(&rqi)
9087                    .arg(&rki)
9088                    .arg(pos)
9089                    .arg(&nhq)
9090                    .arg(&nhk)
9091                    .arg(&theta_scale)
9092                    .arg(&freq_scale)
9093                    .arg(t)
9094                    .arg(&eps)
9095                    .arg(&mut *kc)
9096                    .arg(&mut *vc)
9097                    .arg(t_dev)
9098                    .arg(&ktb)
9099                    .arg(&vtb);
9100                unsafe {
9101                    b.launch(cfg)?;
9102                }
9103            }
9104            None => {
9105                let null: u64 = 0;
9106                b.arg(q0)
9107                    .arg(k0)
9108                    .arg(v0)
9109                    .arg(wq)
9110                    .arg(wk)
9111                    .arg(wv)
9112                    .arg(&mut *q)
9113                    .arg(&mut *k)
9114                    .arg(&mut *v)
9115                    .arg(&nc)
9116                    .arg(&rqi)
9117                    .arg(&rki)
9118                    .arg(pos)
9119                    .arg(&nhq)
9120                    .arg(&nhk)
9121                    .arg(&theta_scale)
9122                    .arg(&freq_scale)
9123                    .arg(&null)
9124                    .arg(&eps)
9125                    .arg(&mut *kc)
9126                    .arg(&mut *vc)
9127                    .arg(t_dev)
9128                    .arg(&ktb)
9129                    .arg(&vtb);
9130                unsafe {
9131                    b.launch(cfg)?;
9132                }
9133            }
9134        }
9135        Ok(())
9136    }
9137
9138    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
9139    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
9140    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
9141    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
9142    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
9143    /// same law as the dc fold).
9144    #[allow(clippy::too_many_arguments)]
9145    pub fn rms_norm_qkv_rope_append(
9146        &self,
9147        q0: &CudaSlice<f32>,
9148        k0: &CudaSlice<f32>,
9149        v0: &CudaSlice<f32>,
9150        wq: &CudaSlice<f32>,
9151        wk: &CudaSlice<f32>,
9152        wv: &CudaSlice<f32>,
9153        q: &mut CudaSlice<f32>,
9154        k: &mut CudaSlice<f32>,
9155        v: &mut CudaSlice<f32>,
9156        head_dim: usize,
9157        rq: usize,
9158        rk: usize,
9159        pos: &CudaSlice<i32>,
9160        nh_q: usize,
9161        nh_k: usize,
9162        base: f32,
9163        freq_scale: f32,
9164        ff: Option<&CudaSlice<f32>>,
9165        eps: f32,
9166        kc: &mut CudaSlice<u8>,
9167        vc: &mut CudaSlice<u8>,
9168        t: usize,
9169        k_tok_bytes: usize,
9170        v_tok_bytes: usize,
9171        g: bool,
9172    ) -> Result<(), Box<dyn std::error::Error>> {
9173        let rows = rq + rk + rk;
9174        let theta_scale = base.powf(-2.0 / head_dim as f32);
9175        let (nc, rqi, rki, nhq, nhk) = (
9176            head_dim as i32,
9177            rq as i32,
9178            rk as i32,
9179            nh_q as i32,
9180            nh_k as i32,
9181        );
9182        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9183        let ti = t as i32;
9184        if Self::pdl_on() && Self::pdl_wb_on() {
9185            use cudarc::driver::{DevicePtr, DevicePtrMut};
9186            let s = &self.gpu.stream();
9187            let (p0, _a0) = q0.device_ptr(s);
9188            let (p1, _a1) = k0.device_ptr(s);
9189            let (p2, _a2) = v0.device_ptr(s);
9190            let (pwq, _a3) = wq.device_ptr(s);
9191            let (pwk, _a4) = wk.device_ptr(s);
9192            let (pwv, _a5) = wv.device_ptr(s);
9193            let (pq, _a6) = q.device_ptr_mut(s);
9194            let (pk, _a7) = k.device_ptr_mut(s);
9195            let (pv, _a8) = v.device_ptr_mut(s);
9196            let (pp, _a9) = pos.device_ptr(s);
9197            let pff: u64 = match ff {
9198                Some(t) => {
9199                    let (p, _gg) = t.device_ptr(s);
9200                    p as u64
9201                }
9202                None => 0,
9203            };
9204            let (pkc, _a10) = kc.device_ptr_mut(s);
9205            let (pvc, _a11) = vc.device_ptr_mut(s);
9206            let mut ps = [
9207                &p0 as *const _ as *mut std::ffi::c_void,
9208                &p1 as *const _ as *mut _,
9209                &p2 as *const _ as *mut _,
9210                &pwq as *const _ as *mut _,
9211                &pwk as *const _ as *mut _,
9212                &pwv as *const _ as *mut _,
9213                &pq as *const _ as *mut _,
9214                &pk as *const _ as *mut _,
9215                &pv as *const _ as *mut _,
9216                &nc as *const _ as *mut _,
9217                &rqi as *const _ as *mut _,
9218                &rki as *const _ as *mut _,
9219                &pp as *const _ as *mut _,
9220                &nhq as *const _ as *mut _,
9221                &nhk as *const _ as *mut _,
9222                &theta_scale as *const _ as *mut _,
9223                &freq_scale as *const _ as *mut _,
9224                &pff as *const _ as *mut _,
9225                &eps as *const _ as *mut _,
9226                &pkc as *const _ as *mut _,
9227                &pvc as *const _ as *mut _,
9228                &ti as *const _ as *mut _,
9229                &ktb as *const _ as *mut _,
9230                &vtb as *const _ as *mut _,
9231            ];
9232            unsafe {
9233                self.launch_pdl_flash(
9234                    g,
9235                    "rms_norm_qkv_rope_append_f32",
9236                    (rows as u32, 1, 1),
9237                    (rms_block(), 1, 1),
9238                    0,
9239                    &mut ps,
9240                )?;
9241            }
9242            return Ok(());
9243        }
9244        let f = if g {
9245            self.func_g("rms_norm_qkv_rope_append_f32")
9246        } else {
9247            self.func("rms_norm_qkv_rope_append_f32")
9248        };
9249        let cfg = LaunchConfig {
9250            grid_dim: (rows as u32, 1, 1),
9251            block_dim: (rms_block(), 1, 1),
9252            shared_mem_bytes: 0,
9253        };
9254        let __s_b = self.gpu.stream();
9255        let mut b = __s_b.launch_builder(&f);
9256        let null: u64 = 0;
9257        b.arg(q0)
9258            .arg(k0)
9259            .arg(v0)
9260            .arg(wq)
9261            .arg(wk)
9262            .arg(wv)
9263            .arg(&mut *q)
9264            .arg(&mut *k)
9265            .arg(&mut *v)
9266            .arg(&nc)
9267            .arg(&rqi)
9268            .arg(&rki)
9269            .arg(pos)
9270            .arg(&nhq)
9271            .arg(&nhk)
9272            .arg(&theta_scale)
9273            .arg(&freq_scale);
9274        match ff {
9275            Some(t) => {
9276                b.arg(t);
9277            }
9278            None => {
9279                b.arg(&null);
9280            }
9281        }
9282        b.arg(&eps)
9283            .arg(&mut *kc)
9284            .arg(&mut *vc)
9285            .arg(&ti)
9286            .arg(&ktb)
9287            .arg(&vtb);
9288        unsafe {
9289            b.launch(cfg)?;
9290        }
9291        Ok(())
9292    }
9293
9294    pub fn add_q8_1(
9295        &self,
9296        a: &CudaSlice<f32>,
9297        b: &CudaSlice<f32>,
9298        res: &mut CudaSlice<f32>,
9299        ncols: usize,
9300        nrows: usize,
9301    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9302        debug_assert!(ncols % 128 == 0);
9303        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9304        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9305        let f = self.func("add_q8_1_f32");
9306        let cfg = LaunchConfig {
9307            grid_dim: (nrows as u32, 1, 1),
9308            block_dim: (rms_block(), 1, 1),
9309            shared_mem_bytes: 0,
9310        };
9311        let nc = ncols as i32;
9312        let __s_b2 = self.gpu.stream();
9313        let mut b2 = __s_b2.launch_builder(&f);
9314        b2.arg(a)
9315            .arg(b)
9316            .arg(&mut *res)
9317            .arg(&mut out_q)
9318            .arg(&mut out_d)
9319            .arg(&nc);
9320        unsafe {
9321            b2.launch(cfg)?;
9322        }
9323        Ok((out_q, out_d))
9324    }
9325
9326    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
9327    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
9328    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
9329    pub fn rms_pre_add_q8_1(
9330        &self,
9331        a: &CudaSlice<f32>,
9332        wa: &CudaSlice<f32>,
9333        b: &CudaSlice<f32>,
9334        res: &mut CudaSlice<f32>,
9335        ncols: usize,
9336        nrows: usize,
9337        eps: f32,
9338    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9339        debug_assert!(ncols % 128 == 0);
9340        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9341        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9342        let f = self.func("rms_pre_add_q8_1_f32");
9343        let cfg = LaunchConfig {
9344            grid_dim: (nrows as u32, 1, 1),
9345            block_dim: (rms_block(), 1, 1),
9346            shared_mem_bytes: 0,
9347        };
9348        let (nc, ep) = (ncols as i32, eps);
9349        let __s_b2 = self.gpu.stream();
9350        let mut b2 = __s_b2.launch_builder(&f);
9351        b2.arg(a)
9352            .arg(wa)
9353            .arg(b)
9354            .arg(&mut *res)
9355            .arg(&mut out_q)
9356            .arg(&mut out_d)
9357            .arg(&nc)
9358            .arg(&ep);
9359        unsafe {
9360            b2.launch(cfg)?;
9361        }
9362        Ok((out_q, out_d))
9363    }
9364
9365    /// L2 norm per row (head_dim), no weight.
9366    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
9367    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
9368    pub fn l2_v2_on(ncols: usize) -> bool {
9369        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9370    }
9371
9372    pub fn l2_norm_pp(
9373        &self,
9374        x: &CudaSlice<f32>,
9375        dst: &mut CudaSlice<f32>,
9376        dst16: Option<&mut CudaSlice<u8>>,
9377        ncols: usize,
9378        nrows: usize,
9379        eps: f32,
9380    ) -> Result<(), Box<dyn std::error::Error>> {
9381        if Self::l2_v2_on(ncols) {
9382            let f = self.func("l2_norm_pp_v2_f32");
9383            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
9384            let cfg = LaunchConfig {
9385                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9386                block_dim: (256, 1, 1),
9387                shared_mem_bytes: 0,
9388            };
9389            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9390            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
9391            let d16: u64 = match dst16 {
9392                Some(d) => self.addr_u8(d),
9393                None => 0,
9394            };
9395            let __s_b = self.gpu.stream();
9396            let mut b = __s_b.launch_builder(&f);
9397            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9398            unsafe {
9399                b.launch(cfg)?;
9400            }
9401            return Ok(());
9402        }
9403        self.l2_norm(x, dst, ncols, nrows, eps)
9404    }
9405
9406    pub fn l2_norm(
9407        &self,
9408        x: &CudaSlice<f32>,
9409        dst: &mut CudaSlice<f32>,
9410        ncols: usize,
9411        nrows: usize,
9412        eps: f32,
9413    ) -> Result<(), Box<dyn std::error::Error>> {
9414        let f = self.func("l2_norm_f32");
9415        let cfg = LaunchConfig {
9416            grid_dim: (nrows as u32, 1, 1),
9417            block_dim: (256, 1, 1),
9418            shared_mem_bytes: 0,
9419        };
9420        let (nc, e) = (ncols as i32, eps);
9421        let __s_b = self.gpu.stream();
9422        let mut b = __s_b.launch_builder(&f);
9423        b.arg(x).arg(dst).arg(&nc).arg(&e);
9424        unsafe {
9425            b.launch(cfg)?;
9426        }
9427        Ok(())
9428    }
9429
9430    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
9431    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
9432    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
9433    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
9434    /// propagate through gdn_scan and flip argmax on marginal logits.
9435    pub fn l2_norm_decode(
9436        &self,
9437        x: &CudaSlice<f32>,
9438        dst: &mut CudaSlice<f32>,
9439        ncols: usize,
9440        nrows: usize,
9441        eps: f32,
9442    ) -> Result<(), Box<dyn std::error::Error>> {
9443        let f = self.func("l2_norm_f32");
9444        let cfg = LaunchConfig {
9445            grid_dim: (nrows as u32, 1, 1),
9446            block_dim: (32, 1, 1),
9447            shared_mem_bytes: 0,
9448        };
9449        let (nc, e) = (ncols as i32, eps);
9450        let __s_b = self.gpu.stream();
9451        let mut b = __s_b.launch_builder(&f);
9452        b.arg(x).arg(dst).arg(&nc).arg(&e);
9453        unsafe {
9454            b.launch(cfg)?;
9455        }
9456        Ok(())
9457    }
9458
9459    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
9460    pub fn rope_neox(
9461        &self,
9462        x: &mut CudaSlice<f32>,
9463        pos: &CudaSlice<i32>,
9464        head_dim: usize,
9465        n_dims: usize,
9466        n_heads: usize,
9467        n_tokens: usize,
9468        freq_base: f32,
9469        freq_scale: f32,
9470    ) -> Result<(), Box<dyn std::error::Error>> {
9471        let f = self.func("rope_neox_f32");
9472        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9473        let grid = (n_heads * n_tokens) as u32;
9474        let cfg = LaunchConfig {
9475            grid_dim: (grid, 1, 1),
9476            block_dim: ((head_dim / 2) as u32, 1, 1),
9477            shared_mem_bytes: 0,
9478        };
9479        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9480        let __s_b = self.gpu.stream();
9481        let mut b = __s_b.launch_builder(&f);
9482        b.arg(x)
9483            .arg(pos)
9484            .arg(&hd)
9485            .arg(&nd)
9486            .arg(&nh)
9487            .arg(&theta_scale)
9488            .arg(&freq_scale);
9489        unsafe {
9490            b.launch(cfg)?;
9491        }
9492        Ok(())
9493    }
9494
9495    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
9496    pub fn rope_neox_ff(
9497        &self,
9498        x: &mut CudaSlice<f32>,
9499        pos: &CudaSlice<i32>,
9500        head_dim: usize,
9501        n_dims: usize,
9502        n_heads: usize,
9503        n_tokens: usize,
9504        freq_base: f32,
9505        freq_scale: f32,
9506        ff: &CudaSlice<f32>,
9507    ) -> Result<(), Box<dyn std::error::Error>> {
9508        let f = self.func("rope_neox_ff_f32");
9509        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9510        let grid = (n_heads * n_tokens) as u32;
9511        let cfg = LaunchConfig {
9512            grid_dim: (grid, 1, 1),
9513            block_dim: ((head_dim / 2) as u32, 1, 1),
9514            shared_mem_bytes: 0,
9515        };
9516        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9517        let __s_b = self.gpu.stream();
9518        let mut b = __s_b.launch_builder(&f);
9519        b.arg(x)
9520            .arg(pos)
9521            .arg(&hd)
9522            .arg(&nd)
9523            .arg(&nh)
9524            .arg(&theta_scale)
9525            .arg(&freq_scale)
9526            .arg(ff);
9527        unsafe {
9528            b.launch(cfg)?;
9529        }
9530        Ok(())
9531    }
9532
9533    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
9534    #[allow(clippy::too_many_arguments)]
9535    pub fn rope_neox2(
9536        &self,
9537        q: &mut CudaSlice<f32>,
9538        k: &mut CudaSlice<f32>,
9539        pos: &CudaSlice<i32>,
9540        head_dim: usize,
9541        n_dims: usize,
9542        nh_q: usize,
9543        nh_k: usize,
9544        n_tokens: usize,
9545        freq_base: f32,
9546        freq_scale: f32,
9547        ff: Option<&CudaSlice<f32>>,
9548    ) -> Result<(), Box<dyn std::error::Error>> {
9549        let f = self.func("rope_neox2_f32");
9550        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9551        let grid = ((nh_q + nh_k) * n_tokens) as u32;
9552        let cfg = LaunchConfig {
9553            grid_dim: (grid, 1, 1),
9554            block_dim: ((head_dim / 2) as u32, 1, 1),
9555            shared_mem_bytes: 0,
9556        };
9557        let (hd, nd, nq, nk, nt) = (
9558            head_dim as i32,
9559            n_dims as i32,
9560            nh_q as i32,
9561            nh_k as i32,
9562            n_tokens as i32,
9563        );
9564        let __s_b = self.gpu.stream();
9565        let mut b = __s_b.launch_builder(&f);
9566        b.arg(q)
9567            .arg(k)
9568            .arg(pos)
9569            .arg(&hd)
9570            .arg(&nd)
9571            .arg(&nq)
9572            .arg(&nk)
9573            .arg(&nt)
9574            .arg(&theta_scale)
9575            .arg(&freq_scale);
9576        match ff {
9577            Some(ffv) => {
9578                b.arg(ffv);
9579                unsafe {
9580                    b.launch(cfg)?;
9581                }
9582            }
9583            None => {
9584                let null: u64 = 0;
9585                b.arg(&null);
9586                unsafe {
9587                    b.launch(cfg)?;
9588                }
9589            }
9590        }
9591        Ok(())
9592    }
9593
9594    /// gemma4 R1: dst = GELU_tanh(gate) * up.
9595    pub fn gelu_tanh_mul(
9596        &self,
9597        gate: &CudaSlice<f32>,
9598        up: &CudaSlice<f32>,
9599        dst: &mut CudaSlice<f32>,
9600        n: usize,
9601    ) -> Result<(), Box<dyn std::error::Error>> {
9602        let f = self.func("gelu_tanh_mul_f32");
9603        let cfg = LaunchConfig::for_num_elems(n as u32);
9604        let ni = n as i32;
9605        let __s_b = self.gpu.stream();
9606        let mut b = __s_b.launch_builder(&f);
9607        b.arg(gate).arg(up).arg(dst).arg(&ni);
9608        unsafe {
9609            b.launch(cfg)?;
9610        }
9611        Ok(())
9612    }
9613
9614    pub fn silu_mul(
9615        &self,
9616        gate: &CudaSlice<f32>,
9617        up: &CudaSlice<f32>,
9618        dst: &mut CudaSlice<f32>,
9619        n: usize,
9620    ) -> Result<(), Box<dyn std::error::Error>> {
9621        let f = self.func("silu_mul_f32");
9622        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9623        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9624        let ni = n as i32;
9625        let __s_b = self.gpu.stream();
9626        let mut b = __s_b.launch_builder(&f);
9627        b.arg(gate).arg(up).arg(dst).arg(&ni);
9628        unsafe {
9629            b.launch(cfg)?;
9630        }
9631        Ok(())
9632    }
9633
9634    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
9635    /// for the down projection — kills the standalone convert pass. Bit-identical class.
9636    pub fn silu_mul_f16out(
9637        &self,
9638        gate: &CudaSlice<f32>,
9639        up: &CudaSlice<f32>,
9640        dst: &mut CudaSlice<f32>,
9641        dst16: &mut CudaSlice<u8>,
9642        n: usize,
9643    ) -> Result<(), Box<dyn std::error::Error>> {
9644        let f = self.func("silu_mul_f16out_f32");
9645        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9646        let ni = n as i32;
9647        let __s_b = self.gpu.stream();
9648        let mut b = __s_b.launch_builder(&f);
9649        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9650        unsafe {
9651            b.launch(cfg)?;
9652        }
9653        Ok(())
9654    }
9655
9656    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
9657    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
9658    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
9659    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
9660    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
9661    /// launches per dense FFN layer (the gate+up post-matmul scales).
9662    pub fn silu_mul_scaled(
9663        &self,
9664        gate: &CudaSlice<f32>,
9665        up: &CudaSlice<f32>,
9666        gs: f32,
9667        us: f32,
9668        dst: &mut CudaSlice<f32>,
9669        n: usize,
9670    ) -> Result<(), Box<dyn std::error::Error>> {
9671        let f = self.func("silu_mul_scaled_f32");
9672        let cfg = LaunchConfig::for_num_elems(n as u32);
9673        let ni = n as i32;
9674        let (gsf, usf) = (gs, us);
9675        let __s_b = self.gpu.stream();
9676        let mut b = __s_b.launch_builder(&f);
9677        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9678        unsafe {
9679            b.launch(cfg)?;
9680        }
9681        Ok(())
9682    }
9683
9684    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
9685    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
9686    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
9687    #[allow(clippy::too_many_arguments)]
9688    pub fn swigluoai_mul_scaled(
9689        &self,
9690        gate: &CudaSlice<f32>,
9691        up: &CudaSlice<f32>,
9692        gs: f32,
9693        us: f32,
9694        alpha: f32,
9695        limit: f32,
9696        dst: &mut CudaSlice<f32>,
9697        n: usize,
9698    ) -> Result<(), Box<dyn std::error::Error>> {
9699        let f = self.func("swigluoai_mul_scaled_f32");
9700        let cfg = LaunchConfig::for_num_elems(n as u32);
9701        let ni = n as i32;
9702        let __s_b = self.gpu.stream();
9703        let mut b = __s_b.launch_builder(&f);
9704        b.arg(gate)
9705            .arg(up)
9706            .arg(&gs)
9707            .arg(&us)
9708            .arg(&alpha)
9709            .arg(&limit)
9710            .arg(dst)
9711            .arg(&ni);
9712        unsafe {
9713            b.launch(cfg)?;
9714        }
9715        Ok(())
9716    }
9717
9718    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
9719    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
9720    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
9721    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
9722    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
9723    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
9724    /// n must be a multiple of 32 (n_ff always is).
9725    pub fn silu_mul_scaled_q8_1(
9726        &self,
9727        gate: &CudaSlice<f32>,
9728        up: &CudaSlice<f32>,
9729        gs: f32,
9730        us: f32,
9731        n: usize,
9732    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9733        let f = self.func("silu_mul_scaled_q8_1");
9734        let nblk = n / 32;
9735        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
9736        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
9737        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
9738        let cfg = LaunchConfig::for_num_elems(n as u32);
9739        let (gsf, usf, ni) = (gs, us, n as i32);
9740        let __s_b = self.gpu.stream();
9741        let mut b = __s_b.launch_builder(&f);
9742        b.arg(gate)
9743            .arg(up)
9744            .arg(&gsf)
9745            .arg(&usf)
9746            .arg(&mut aq)
9747            .arg(&mut ad)
9748            .arg(&ni);
9749        unsafe {
9750            b.launch(cfg)?;
9751        }
9752        Ok((aq, ad))
9753    }
9754
9755    pub fn add(
9756        &self,
9757        a: &CudaSlice<f32>,
9758        b_in: &CudaSlice<f32>,
9759        dst: &mut CudaSlice<f32>,
9760        n: usize,
9761    ) -> Result<(), Box<dyn std::error::Error>> {
9762        let f = self.func("add_f32");
9763        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9764        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9765        let ni = n as i32;
9766        let __s_bld = self.gpu.stream();
9767        let mut bld = __s_bld.launch_builder(&f);
9768        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9769        unsafe {
9770            bld.launch(cfg)?;
9771        }
9772        Ok(())
9773    }
9774
9775    pub fn mul(
9776        &self,
9777        a: &CudaSlice<f32>,
9778        b_in: &CudaSlice<f32>,
9779        dst: &mut CudaSlice<f32>,
9780        n: usize,
9781    ) -> Result<(), Box<dyn std::error::Error>> {
9782        let f = self.func("mul_f32");
9783        let cfg = LaunchConfig::for_num_elems(n as u32);
9784        let ni = n as i32;
9785        let __s_bld = self.gpu.stream();
9786        let mut bld = __s_bld.launch_builder(&f);
9787        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9788        unsafe {
9789            bld.launch(cfg)?;
9790        }
9791        Ok(())
9792    }
9793
9794    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
9795    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
9796    pub fn matmul(
9797        &self,
9798        w: &crate::model::GpuTensor,
9799        x: &CudaSlice<f32>,
9800        m: usize,
9801    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9802        use crate::model::GpuTensor;
9803        let in_f = w.in_features();
9804        let out_f = w.out_features();
9805        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
9806        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
9807        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
9808        // gives nothing). Quantize the activation once here then call the GEMM.
9809        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
9810        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
9811        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
9812        #[allow(non_snake_case)]
9813        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
9814        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
9815        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
9816            usize::MAX
9817        } else {
9818            16usize
9819        };
9820
9821        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
9822        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
9823        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
9824        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
9825        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
9826        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
9827        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
9828        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
9829        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
9830        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
9831        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
9832        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
9833        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
9834        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
9835        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
9836        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
9837        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
9838        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
9839        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
9840        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
9841        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
9842        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
9843        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
9844        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
9845        if m >= GEMM_M_THRESHOLD {
9846            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
9847                return Ok(y);
9848            }
9849            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
9850            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
9851            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
9852            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
9853            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
9854            // tile defaults differently by operand source.
9855            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
9856                return Ok(y);
9857            }
9858            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
9859            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
9860            if let Some(y) = self.try_f16_gemm(w, x, m)? {
9861                return Ok(y);
9862            }
9863        }
9864        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
9865        // m threshold the rest of this method uses:
9866        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
9867        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
9868        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
9869        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
9870        //     across every tier by construction with no batched twin needed.
9871        //
9872        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
9873        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
9874        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
9875        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
9876        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
9877        // arms is what makes sure it never gets there.
9878        if let GpuTensor::Quant { qtype, .. } = w {
9879            if *qtype == QT_F8_E4M3_BLK {
9880                if m >= GEMM_M_THRESHOLD {
9881                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
9882                        return Ok(y);
9883                    }
9884                }
9885                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9886                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
9887                    return Ok(y);
9888                }
9889            }
9890        }
9891        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
9892            return self.qmatvec_mmq(w, x, m);
9893        }
9894        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
9895            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9896            return self.qmatvec_gemm(w, &aq, &ad, m);
9897        }
9898        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
9899        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
9900        if m >= GEMM_M_THRESHOLD {
9901            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
9902                return Ok(y);
9903            }
9904        }
9905        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
9906        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
9907        // to Stage-A f32-dequant (the correctness oracle path).
9908        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
9909        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
9910        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
9911        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
9912        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
9913        if m == 1 && fast {
9914            if let GpuTensor::Quant {
9915                bytes,
9916                qtype,
9917                row_bytes,
9918                rp,
9919                rp4,
9920                scale,
9921                ..
9922            } = w
9923            {
9924                if self.mmvq_supports(*qtype) {
9925                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
9926                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
9927                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
9928                    let (bytes, rp) = match rp4 {
9929                        Some(m4) => (m4, true),
9930                        None => (bytes, *rp),
9931                    };
9932                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9933                    return self.qmatvec_mmvq(
9934                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
9935                    );
9936                }
9937            }
9938        }
9939        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
9940        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
9941        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
9942        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
9943        // block below. MEMRA_NO_BATCHED -> per-m path.
9944        //
9945        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
9946        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
9947        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
9948        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
9949        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
9950        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
9951        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
9952        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
9953        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
9954        if (2..=16).contains(&m)
9955            && fast
9956            && std::env::var("MEMRA_NO_BATCHED").is_err()
9957            && (m <= 4 || Self::b8_enabled())
9958        {
9959            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
9960            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
9961            // is present (rp4) — the mirror pick below then routes to the _rp family.
9962            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
9963            // because the native e4m3 row layout is already aligned and needs no mirror.
9964            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
9965            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
9966            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
9967            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
9968            let m_ok = m <= 8
9969                || matches!(w, GpuTensor::Quant { qtype, .. }
9970                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
9971                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
9972            if m_ok {
9973                if let GpuTensor::Quant {
9974                    bytes,
9975                    qtype,
9976                    row_bytes,
9977                    rp,
9978                    rp4,
9979                    ..
9980                } = w
9981                {
9982                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
9983                        let (bytes, rp) = match rp4 {
9984                            Some(m4) => (m4, true),
9985                            None => (bytes, *rp),
9986                        };
9987                        let mcols = Self::batched_mcols(m);
9988                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9989                        let mut y = self.qmatvec_mmvq_batched(
9990                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
9991                        )?;
9992                        if let GpuTensor::Quant { scale, .. } = w {
9993                            if *scale != 1.0 {
9994                                self.scale_inplace(&mut y, *scale, m * out_f)?;
9995                            }
9996                        }
9997                        return Ok(y);
9998                    }
9999                }
10000            }
10001        }
10002        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
10003        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
10004        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
10005        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
10006        // for this dtype, so the generic match below must never see it under `fast`.
10007        if fast {
10008            if let GpuTensor::Quant {
10009                bytes,
10010                qtype,
10011                row_bytes,
10012                scale,
10013                ..
10014            } = w
10015            {
10016                if *qtype == QT_F8_E4M3 {
10017                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10018                    return self.qmatvec_mmvq(
10019                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10020                    );
10021                }
10022            }
10023        }
10024        let mut y = match w {
10025            GpuTensor::Quant {
10026                bytes,
10027                qtype,
10028                row_bytes,
10029                ..
10030            } if fast && *qtype == QT_Q8_0 => {
10031                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10032            }
10033            GpuTensor::Quant {
10034                bytes,
10035                qtype,
10036                row_bytes,
10037                ..
10038            } if fast && *qtype == QT_Q4_K => {
10039                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10040            }
10041            GpuTensor::Quant {
10042                bytes,
10043                qtype,
10044                row_bytes,
10045                ..
10046            } if fast && *qtype == QT_Q6_K => {
10047                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10048            }
10049            GpuTensor::Quant {
10050                bytes,
10051                qtype,
10052                row_bytes,
10053                ..
10054            } if fast && *qtype == QT_Q5_K => {
10055                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10056            }
10057            GpuTensor::Quant {
10058                bytes,
10059                qtype,
10060                row_bytes,
10061                ..
10062            } if fast && *qtype == QT_Q3_K => {
10063                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10064            }
10065            GpuTensor::Quant {
10066                bytes,
10067                qtype,
10068                row_bytes,
10069                rp,
10070                ..
10071            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10072                if *rp {
10073                    "qmatvec_nvfp4_dp4a_rp"
10074                } else {
10075                    "qmatvec_nvfp4_dp4a"
10076                },
10077                bytes,
10078                x,
10079                m,
10080                in_f,
10081                out_f,
10082                *row_bytes,
10083            )?,
10084            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
10085            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
10086            // anomaly (research/kat-anomaly-20260802/).
10087            GpuTensor::Quant {
10088                bytes,
10089                qtype,
10090                row_bytes,
10091                ..
10092            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10093                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10094            }
10095            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
10096            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
10097            // without first writing the matching kernel, or func() will panic
10098            // "kernel ... not in any fatbin".
10099            GpuTensor::Quant {
10100                bytes,
10101                qtype,
10102                row_bytes,
10103                rp,
10104                ..
10105            } =>
10106            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
10107            // deq(row,j) form cannot address the planes; same value/product order).
10108            {
10109                self.qmatvec(
10110                    bytes,
10111                    x,
10112                    m,
10113                    in_f,
10114                    out_f,
10115                    if *rp && *qtype == QT_NVFP4 {
10116                        QT_NVFP4_RP
10117                    } else {
10118                        *qtype
10119                    },
10120                    *row_bytes,
10121                )?
10122            }
10123            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10124            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
10125            // cuBLASLt f32 GEMV as the Float arm.
10126            GpuTensor::FloatBf16 { data, .. } => {
10127                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10128            }
10129        };
10130        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
10131        if let GpuTensor::Quant { scale, .. } = w {
10132            if *scale != 1.0 {
10133                self.scale_inplace(&mut y, *scale, m * out_f)?;
10134            }
10135        }
10136        Ok(y)
10137    }
10138
10139    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
10140    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
10141    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10142        use crate::model::GpuTensor;
10143        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10144            return false;
10145        }
10146        match w {
10147            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
10148            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
10149            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
10150            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
10151            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
10152            // block class has no fused twin yet, so each of its projections takes its own launch.
10153            GpuTensor::Quant { qtype, .. } => {
10154                matches!(
10155                    *qtype,
10156                    QT_Q8_0
10157                        | QT_Q4_K
10158                        | QT_Q6_K
10159                        | QT_Q5_K
10160                        | QT_Q3_K
10161                        | QT_NVFP4
10162                        | QT_F8_E4M3
10163                        | QT_F8_E4M3_BLK
10164                        | QT_Q4_0
10165                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10166            }
10167            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10168        }
10169    }
10170
10171    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
10172    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
10173    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
10174    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
10175    pub fn matmul_pre(
10176        &self,
10177        w: &crate::model::GpuTensor,
10178        aq: &CudaSlice<i8>,
10179        ad: &CudaSlice<f32>,
10180        x_fallback: &CudaSlice<f32>,
10181        m: usize,
10182    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10183        use crate::model::GpuTensor;
10184        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
10185        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
10186        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
10187        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
10188        // rc=30013 dig, 2026-07-31).
10189        let x_raw_ok = x_fallback.len() >= m * w.in_features();
10190        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
10191        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
10192        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10193            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10194                return Ok(y);
10195            }
10196            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
10197            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
10198            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10199                return Ok(y);
10200            }
10201            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
10202            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10203                return Ok(y);
10204            }
10205        }
10206        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
10207        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
10208        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
10209        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
10210        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
10211        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10212            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10213                return Ok(y);
10214            }
10215        }
10216        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10217            return Ok(y);
10218        }
10219        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
10220        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
10221        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
10222        // aq/ad.
10223        if m >= 16
10224            && w.out_features() >= 128
10225            && self.mmq_supports(w)
10226            && !self.verify_exact_on()
10227            && x_raw_ok
10228        {
10229            return self.qmatvec_mmq(w, x_fallback, m);
10230        }
10231        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
10232        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
10233        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10234            if let Some(y) =
10235                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10236            {
10237                return Ok(y);
10238            }
10239        }
10240        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
10241        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
10242        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10243            return self.qmatvec_gemm(w, aq, ad, m);
10244        }
10245        if !self.uses_q8_1_fast(w) {
10246            return self.matmul(w, x_fallback, m);
10247        }
10248        let in_f = w.in_features();
10249        let out_f = w.out_features();
10250        let (bytes, qtype, row_bytes, scale, rp) = match w {
10251            GpuTensor::Quant {
10252                bytes,
10253                qtype,
10254                row_bytes,
10255                scale,
10256                rp,
10257                ..
10258            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10259            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10260        };
10261        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
10262        // the dp4a/oracle tails below keep the raw GGUF bytes.
10263        let (mbytes, mrp) = match w {
10264            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10265            _ => (bytes, rp),
10266        };
10267        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
10268        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
10269        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
10270        if m == 1 && self.mmvq_supports(qtype) {
10271            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10272        }
10273        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
10274        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
10275        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
10276        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
10277        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
10278        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
10279        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
10280        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
10281        // m=5..8 on the old per-m path (b8-tier-only seam).
10282        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
10283        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
10284        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
10285        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10286            && std::env::var("MEMRA_NO_BATCHED").is_err()
10287            && (m <= 4 || Self::b8_enabled())
10288            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
10289            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
10290            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
10291            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10292                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10293        {
10294            let mcols = Self::batched_mcols(m);
10295            return self.qmatvec_mmvq_batched(
10296                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10297            );
10298        }
10299        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
10300        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
10301        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
10302        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
10303        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
10304        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10305            let (b2, r2) = if qtype == QT_Q4_0 {
10306                (mbytes, mrp)
10307            } else {
10308                (bytes, rp)
10309            };
10310            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10311        }
10312        let name = match qtype {
10313            QT_Q8_0 => "qmatvec_q8_0_dp4a",
10314            QT_Q4_K => "qmatvec_q4_K_dp4a",
10315            QT_Q6_K => "qmatvec_q6_K_dp4a",
10316            QT_Q5_K => "qmatvec_q5_K_dp4a",
10317            QT_Q3_K => "qmatvec_q3_K_dp4a",
10318            QT_NVFP4 => {
10319                if rp {
10320                    "qmatvec_nvfp4_dp4a_rp"
10321                } else {
10322                    "qmatvec_nvfp4_dp4a"
10323                }
10324            }
10325            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10326            _ => unreachable!(),
10327        };
10328        let f = self.func(name);
10329        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
10330        let cfg = LaunchConfig {
10331            grid_dim: (out_f as u32, m as u32, 1),
10332            block_dim: (128, 1, 1),
10333            shared_mem_bytes: 0,
10334        };
10335        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10336        let __s_b = self.gpu.stream();
10337        let mut b = __s_b.launch_builder(&f);
10338        b.arg(bytes)
10339            .arg(aq)
10340            .arg(ad)
10341            .arg(&mut y)
10342            .arg(&inf)
10343            .arg(&outf)
10344            .arg(&mi)
10345            .arg(&rb);
10346        unsafe {
10347            b.launch(cfg)?;
10348        }
10349        if scale != 1.0 {
10350            self.scale_inplace(&mut y, scale, m * out_f)?;
10351        }
10352        Ok(y)
10353    }
10354
10355    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
10356    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
10357    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
10358    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
10359    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
10360    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
10361    /// reduce as m=1); this method just forces that path unconditionally.
10362    pub fn matmul_decode_exact(
10363        &self,
10364        w: &crate::model::GpuTensor,
10365        x: &CudaSlice<f32>,
10366        m: usize,
10367    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10368        use crate::model::GpuTensor;
10369        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
10370        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
10371        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
10372        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
10373        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
10374        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
10375        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
10376        if let GpuTensor::Float { data, .. } = w {
10377            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10378        }
10379        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
10380        // float linear (same n-independent reduction contract as the Float arm above).
10381        if let GpuTensor::FloatBf16 { data, .. } = w {
10382            let (in_f, out_f) = (w.in_features(), w.out_features());
10383            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10384        }
10385        if !self.uses_q8_1_fast(w) {
10386            return self.matmul(w, x, m);
10387        }
10388        let in_f = w.in_features();
10389        let out_f = w.out_features();
10390        let (bytes, qtype, row_bytes, scale, rp) = match w {
10391            GpuTensor::Quant {
10392                bytes,
10393                qtype,
10394                row_bytes,
10395                scale,
10396                rp,
10397                ..
10398            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10399            _ => return self.matmul(w, x, m),
10400        };
10401        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
10402        // which does its own mirror pick).
10403        let (bytes, rp) = match w {
10404            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10405            _ => (bytes, rp),
10406        };
10407        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10408        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
10409        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
10410        // (token,row) by construction, which is exactly what this method exists to guarantee.
10411        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10412            return Ok(y);
10413        }
10414        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
10415        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
10416        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
10417        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
10418        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
10419        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
10420        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
10421        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
10422        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10423            && std::env::var("MEMRA_NO_BATCHED").is_err()
10424            && (m <= 4 || Self::b8_enabled())
10425            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
10426            // no mirror precondition, `rp` selects the layout only.
10427            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10428                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10429        {
10430            let mcols = Self::batched_mcols(m);
10431            return self.qmatvec_mmvq_batched(
10432                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10433            );
10434        }
10435        if self.mmvq_supports(qtype) {
10436            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
10437            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
10438            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10439        }
10440        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
10441        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
10442        self.matmul_pre(w, &aq, &ad, x, m)
10443    }
10444
10445    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
10446    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
10447    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
10448    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
10449    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
10450    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
10451    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
10452    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
10453    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
10454    pub fn matmul_decode_exact_pre(
10455        &self,
10456        w: &crate::model::GpuTensor,
10457        aq: &CudaSlice<i8>,
10458        ad: &CudaSlice<f32>,
10459        m: usize,
10460    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10461        use crate::model::GpuTensor;
10462        debug_assert!(
10463            self.uses_q8_1_fast(w),
10464            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10465        );
10466        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
10467        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10468            return Ok(y);
10469        }
10470        let in_f = w.in_features();
10471        let out_f = w.out_features();
10472        let (bytes, qtype, row_bytes, scale, rp) = match w {
10473            GpuTensor::Quant {
10474                bytes,
10475                qtype,
10476                row_bytes,
10477                scale,
10478                rp,
10479                ..
10480            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10481            _ => {
10482                return Err(
10483                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10484                );
10485            }
10486        };
10487        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
10488        let (bytes, rp) = match w {
10489            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10490            _ => (bytes, rp),
10491        };
10492        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
10493        if (2..=16).contains(&m)
10494            && self.batched_supports(qtype)
10495            && self.mmvq_supports(qtype)
10496            && std::env::var("MEMRA_NO_BATCHED").is_err()
10497            && (m <= 4 || Self::b8_enabled())
10498            && (m <= 8
10499                || qtype == QT_Q4_0
10500                || qtype == QT_Q6_K
10501                || qtype == QT_F8_E4M3
10502                || qtype == QT_NVFP4
10503                || qtype == QT_Q4_K
10504                || qtype == QT_Q5_K
10505                || qtype == QT_Q8_0)
10506        {
10507            let mcols = Self::batched_mcols(m);
10508            return self.qmatvec_mmvq_batched(
10509                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10510            );
10511        }
10512        if self.mmvq_supports(qtype) {
10513            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10514        }
10515        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
10516        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
10517        let x0 = self.zeros(0)?;
10518        self.matmul_pre(w, aq, ad, &x0, m)
10519    }
10520
10521    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
10522    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
10523    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
10524    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
10525    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
10526    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
10527    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
10528    /// per-tensor path.
10529    pub fn matmul_decode_exact_dual_pre(
10530        &self,
10531        w0: &crate::model::GpuTensor,
10532        w1: &crate::model::GpuTensor,
10533        aq: &CudaSlice<i8>,
10534        ad: &CudaSlice<f32>,
10535        m: usize,
10536    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10537    {
10538        use crate::model::GpuTensor;
10539        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10540        let on = *ON.get_or_init(|| {
10541            std::env::var("MEMRA_SPEC_DUAL_T")
10542                .map(|v| v != "0")
10543                .unwrap_or(true)
10544        });
10545        if !on
10546            || !(2..=7).contains(&m)
10547            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10548            || !self.uses_q8_1_fast(w0)
10549            || !self.uses_q8_1_fast(w1)
10550        {
10551            return Ok(None);
10552        }
10553        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
10554        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
10555        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
10556        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
10557        if !self.mmvq_supports(QT_NVFP4) {
10558            return Ok(None);
10559        }
10560        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10561        if w1.in_features() != in_f || w1.out_features() != out_f {
10562            return Ok(None);
10563        }
10564        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10565            (
10566                GpuTensor::Quant {
10567                    bytes: b0,
10568                    qtype: q0,
10569                    row_bytes: rb0,
10570                    scale: s0,
10571                    rp: rp0,
10572                    rp4: None,
10573                    ..
10574                },
10575                GpuTensor::Quant {
10576                    bytes: b1,
10577                    qtype: q1,
10578                    row_bytes: rb1,
10579                    scale: s1,
10580                    rp: rp1,
10581                    rp4: None,
10582                    ..
10583                },
10584            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10585                (b0, b1, *rb0, *s0, *s1, *rp0)
10586            }
10587            _ => return Ok(None),
10588        };
10589        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
10590        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
10591        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10592        {
10593            return Ok(None);
10594        }
10595        let (y0, y1) =
10596            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10597        Ok(Some(((y0, s0), (y1, s1))))
10598    }
10599
10600    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
10601    /// launch computes both FFN projections of a verify batch — same activation, same shape,
10602    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
10603    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
10604    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
10605    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
10606    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
10607    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
10608    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
10609    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
10610    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
10611    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
10612    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
10613    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
10614    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
10615    pub fn matmul_decode_exact_dual(
10616        &self,
10617        w0: &crate::model::GpuTensor,
10618        w1: &crate::model::GpuTensor,
10619        x: &CudaSlice<f32>,
10620        m: usize,
10621    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10622        use crate::model::GpuTensor;
10623        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10624        let on = *ON.get_or_init(|| {
10625            std::env::var("MEMRA_SPEC_DUAL_T")
10626                .map(|v| v != "0")
10627                .unwrap_or(true)
10628        });
10629        if !on
10630            || !(2..=4).contains(&m)
10631            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10632            || !self.uses_q8_1_fast(w0)
10633            || !self.uses_q8_1_fast(w1)
10634        {
10635            return Ok(None);
10636        }
10637        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
10638        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
10639        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
10640        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
10641        if !self.mmvq_supports(QT_NVFP4) {
10642            return Ok(None);
10643        }
10644        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10645        if w1.in_features() != in_f || w1.out_features() != out_f {
10646            return Ok(None);
10647        }
10648        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10649            (
10650                GpuTensor::Quant {
10651                    bytes: b0,
10652                    qtype: q0,
10653                    row_bytes: rb0,
10654                    scale: s0,
10655                    rp: rp0,
10656                    rp4: None,
10657                    ..
10658                },
10659                GpuTensor::Quant {
10660                    bytes: b1,
10661                    qtype: q1,
10662                    row_bytes: rb1,
10663                    scale: s1,
10664                    rp: rp1,
10665                    rp4: None,
10666                    ..
10667                },
10668            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10669                (b0, b1, *rb0, *s0, *s1, *rp0)
10670            }
10671            _ => return Ok(None),
10672        };
10673        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
10674        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
10675        if std::env::var("MEMRA_DEBUG").is_ok() {
10676            static ONCE: std::sync::Once = std::sync::Once::new();
10677            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10678        }
10679        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10680        let (y0, y1) =
10681            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10682        let mut y0 = y0;
10683        let mut y1 = y1;
10684        if s0 != 1.0 {
10685            self.scale_inplace(&mut y0, s0, m * out_f)?;
10686        }
10687        if s1 != 1.0 {
10688            self.scale_inplace(&mut y1, s1, m * out_f)?;
10689        }
10690        Ok(Some((y0, y1)))
10691    }
10692
10693    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
10694    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
10695    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
10696    /// twins (both buffers must be the repacked layout).
10697    #[allow(clippy::too_many_arguments)]
10698    pub fn qmatvec_batched_dual_raw(
10699        &self,
10700        b0: &CudaSlice<u8>,
10701        b1: &CudaSlice<u8>,
10702        aq: &CudaSlice<i8>,
10703        ad: &CudaSlice<f32>,
10704        m: usize,
10705        in_f: usize,
10706        out_f: usize,
10707        row_bytes: usize,
10708        rp: bool,
10709    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10710        const ROWS_PER_BLOCK: u32 = 4;
10711        let mcols = Self::batched_mcols(m);
10712        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
10713        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
10714        let tiny_rp1 = rp
10715            && mcols == 4
10716            && out_f <= 128
10717            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10718        let (name, rows_per_block) = if tiny_rp1 {
10719            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10720        } else {
10721            match (mcols, rp, m) {
10722                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10723                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10724                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10725                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10726                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10727                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10728                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10729                _ => {
10730                    return Err(
10731                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10732                    );
10733                }
10734            }
10735        };
10736        let f = self.func(name);
10737        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10738        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10739        let cfg = LaunchConfig {
10740            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10741            block_dim: (32, ROWS_PER_BLOCK, 1),
10742            shared_mem_bytes: 0,
10743        };
10744        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10745        let __s_b = self.gpu.stream();
10746        let mut b = __s_b.launch_builder(&f);
10747        b.arg(b0)
10748            .arg(b1)
10749            .arg(aq)
10750            .arg(ad)
10751            .arg(&mut y0)
10752            .arg(&mut y1)
10753            .arg(&inf)
10754            .arg(&outf)
10755            .arg(&mi)
10756            .arg(&rb);
10757        unsafe {
10758            b.launch(cfg)?;
10759        }
10760        Ok((y0, y1))
10761    }
10762
10763    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
10764    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
10765    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
10766    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
10767    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
10768    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
10769    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
10770    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
10771    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
10772    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
10773    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
10774    pub fn matmul_pre_dual_noscale(
10775        &self,
10776        w0: &crate::model::GpuTensor,
10777        w1: &crate::model::GpuTensor,
10778        aq: &CudaSlice<i8>,
10779        ad: &CudaSlice<f32>,
10780        m: usize,
10781    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10782    {
10783        use crate::model::GpuTensor;
10784        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
10785            return Ok(None);
10786        }
10787        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
10788        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
10789        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
10790        // would mix dispatch families across the pair — the exact class `q8_fused_params`
10791        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
10792        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
10793        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
10794        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
10795        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
10796        if !self.mmvq_supports(QT_NVFP4) {
10797            return Ok(None);
10798        }
10799        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10800        if w1.in_features() != in_f || w1.out_features() != out_f {
10801            return Ok(None);
10802        }
10803        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
10804        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
10805        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
10806        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
10807        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
10808        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
10809        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
10810        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
10811        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
10812        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
10813        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
10814        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
10815        let no_mirror =
10816            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
10817        if self.q8_ffn_fuse2_on()
10818            && no_mirror(w0)
10819            && no_mirror(w1)
10820            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
10821        {
10822            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
10823            return Ok(Some(((y0, 1.0), (y1, 1.0))));
10824        }
10825        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
10826        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
10827        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
10828        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
10829        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
10830        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
10831        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
10832        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
10833        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
10834        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
10835            let (y0, y1) =
10836                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
10837            return Ok(Some(((y0, p0.3), (y1, p1.3))));
10838        }
10839        let (b0, q0, rb0, s0, rp0) = match w0 {
10840            GpuTensor::Quant {
10841                bytes,
10842                qtype,
10843                row_bytes,
10844                scale,
10845                rp,
10846                ..
10847            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10848            _ => return Ok(None),
10849        };
10850        let (b1, q1, rb1, s1, rp1) = match w1 {
10851            GpuTensor::Quant {
10852                bytes,
10853                qtype,
10854                row_bytes,
10855                scale,
10856                rp,
10857                ..
10858            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10859            _ => return Ok(None),
10860        };
10861        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
10862            return Ok(None);
10863        }
10864        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
10865        const RPW: u32 = 2;
10866        let rows_per_block = ROWS_PER_BLOCK * RPW;
10867        let f = self.func(if rp0 {
10868            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
10869        } else {
10870            "qmatvec_nvfp4_mmvq_dual_mr2"
10871        });
10872        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
10873        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
10874        let cfg = LaunchConfig {
10875            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10876            block_dim: (32, ROWS_PER_BLOCK, 1),
10877            shared_mem_bytes: 0,
10878        };
10879        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
10880        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
10881        // yscale args stay 1.0 here (they exist for the single-tensor callers).
10882        let one = 1.0f32;
10883        let __s_b = self.gpu.stream();
10884        let mut b = __s_b.launch_builder(&f);
10885        b.arg(b0)
10886            .arg(b1)
10887            .arg(aq)
10888            .arg(ad)
10889            .arg(&mut y0)
10890            .arg(&mut y1)
10891            .arg(&inf)
10892            .arg(&outf)
10893            .arg(&mi)
10894            .arg(&rb)
10895            .arg(&one)
10896            .arg(&one);
10897        unsafe {
10898            b.launch(cfg)?;
10899        }
10900        Ok(Some(((y0, s0), (y1, s1))))
10901    }
10902
10903    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
10904    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
10905    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
10906    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
10907    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
10908    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
10909    /// back to the three singles.
10910    #[allow(clippy::too_many_arguments)]
10911    pub fn matmul_nvfp4_fused3(
10912        &self,
10913        w0: &crate::model::GpuTensor,
10914        w1: &crate::model::GpuTensor,
10915        w2: &crate::model::GpuTensor,
10916        aq: &CudaSlice<i8>,
10917        ad: &CudaSlice<f32>,
10918        m: usize,
10919    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
10920    {
10921        use crate::model::GpuTensor;
10922        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
10923        // read serves all m rows); the fused segments would re-read the weight per row. The
10924        // fusion win is the B=1 decode tick.
10925        if m != 1
10926            || !self.mmvq_supports(QT_NVFP4)
10927            || !self.uses_q8_1_fast(w0)
10928            || !self.uses_q8_1_fast(w1)
10929            || !self.uses_q8_1_fast(w2)
10930        {
10931            return Ok(None);
10932        }
10933        let unpack = |w: &crate::model::GpuTensor| match w {
10934            GpuTensor::Quant {
10935                bytes,
10936                qtype,
10937                scale,
10938                rp,
10939                ..
10940            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
10941            _ => None,
10942        };
10943        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
10944            return Ok(None);
10945        };
10946        let in_f = w0.in_features();
10947        if w1.in_features() != in_f || w2.in_features() != in_f {
10948            return Ok(None);
10949        }
10950        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
10951        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
10952        const RPW: u32 = 2;
10953        let rows_pb = ROWS_PER_BLOCK * RPW;
10954        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
10955        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
10956        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
10957        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
10958        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
10959        let cfg = LaunchConfig {
10960            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
10961            block_dim: (32, ROWS_PER_BLOCK, 1),
10962            shared_mem_bytes: 0,
10963        };
10964        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
10965        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
10966        // only dereferenced for the launch-arg build inside this call.
10967        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
10968        let __s_b = self.gpu.stream();
10969        let mut b = __s_b.launch_builder(&f);
10970        b.arg(b0)
10971            .arg(b1)
10972            .arg(b2)
10973            .arg(aq)
10974            .arg(ad)
10975            .arg(&mut y0)
10976            .arg(&mut y1)
10977            .arg(&mut y2)
10978            .arg(&inf)
10979            .arg(&oi0)
10980            .arg(&oi1)
10981            .arg(&oi2)
10982            .arg(&mi)
10983            .arg(&p0.1)
10984            .arg(&p1.1)
10985            .arg(&p2.1);
10986        unsafe {
10987            b.launch(cfg)?;
10988        }
10989        Ok(Some((y0, y1, y2)))
10990    }
10991
10992    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
10993    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
10994    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
10995    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
10996    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
10997    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
10998    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
10999    /// same-binary interleaved A/B arm.
11000    pub fn matmul_nvfp4_fused2(
11001        &self,
11002        w0: &crate::model::GpuTensor,
11003        w1: &crate::model::GpuTensor,
11004        aq: &CudaSlice<i8>,
11005        ad: &CudaSlice<f32>,
11006        m: usize,
11007    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11008        use crate::model::GpuTensor;
11009        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11010        let off =
11011            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11012        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11013        // read serves all m rows); the fused segments would re-read the weight per row.
11014        if off
11015            || m != 1
11016            || !self.mmvq_supports(QT_NVFP4)
11017            || !self.uses_q8_1_fast(w0)
11018            || !self.uses_q8_1_fast(w1)
11019        {
11020            return Ok(None);
11021        }
11022        let unpack = |w: &crate::model::GpuTensor| match w {
11023            GpuTensor::Quant {
11024                bytes,
11025                qtype,
11026                scale,
11027                rp,
11028                ..
11029            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11030            _ => None,
11031        };
11032        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11033            return Ok(None);
11034        };
11035        let in_f = w0.in_features();
11036        if w1.in_features() != in_f {
11037            return Ok(None);
11038        }
11039        let (o0, o1) = (w0.out_features(), w1.out_features());
11040        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11041        const RPW: u32 = 2;
11042        let rows_pb = ROWS_PER_BLOCK * RPW;
11043        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11044        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11045        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11046        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11047        let cfg = LaunchConfig {
11048            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11049            block_dim: (32, ROWS_PER_BLOCK, 1),
11050            shared_mem_bytes: 0,
11051        };
11052        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11053        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11054        // only dereferenced for the launch-arg build inside this call.
11055        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11056        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
11057        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
11058        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11059            {
11060                use cudarc::driver::{DevicePtr, DevicePtrMut};
11061                let s = &self.gpu.stream();
11062                let (pw0, _g0) = b0.device_ptr(s);
11063                let (pw1, _g1) = b1.device_ptr(s);
11064                let (paq, _g2) = aq.device_ptr(s);
11065                let (pad, _g3) = ad.device_ptr(s);
11066                let (py0, _g4) = y0.device_ptr_mut(s);
11067                let (py1, _g5) = y1.device_ptr_mut(s);
11068                let (s0, s1) = (p0.1, p1.1);
11069                let mut ps = [
11070                    &pw0 as *const _ as *mut std::ffi::c_void,
11071                    &pw1 as *const _ as *mut _,
11072                    &paq as *const _ as *mut _,
11073                    &pad as *const _ as *mut _,
11074                    &py0 as *const _ as *mut _,
11075                    &py1 as *const _ as *mut _,
11076                    &inf as *const _ as *mut _,
11077                    &oi0 as *const _ as *mut _,
11078                    &oi1 as *const _ as *mut _,
11079                    &mi as *const _ as *mut _,
11080                    &s0 as *const _ as *mut _,
11081                    &s1 as *const _ as *mut _,
11082                ];
11083                unsafe {
11084                    self.launch_pdl(
11085                        "qmatvec_nvfp4_mmvq_fused2_rp",
11086                        cfg.grid_dim,
11087                        cfg.block_dim,
11088                        &mut ps,
11089                    )?;
11090                }
11091            }
11092            return Ok(Some((y0, y1)));
11093        }
11094        let __s_b = self.gpu.stream();
11095        let mut b = __s_b.launch_builder(&f);
11096        b.arg(b0)
11097            .arg(b1)
11098            .arg(aq)
11099            .arg(ad)
11100            .arg(&mut y0)
11101            .arg(&mut y1)
11102            .arg(&inf)
11103            .arg(&oi0)
11104            .arg(&oi1)
11105            .arg(&mi)
11106            .arg(&p0.1)
11107            .arg(&p1.1);
11108        unsafe {
11109            b.launch(cfg)?;
11110        }
11111        Ok(Some((y0, y1)))
11112    }
11113
11114    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
11115    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
11116    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
11117    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
11118    pub fn matmul_nvfp4_fused2_into(
11119        &self,
11120        w0: &crate::model::GpuTensor,
11121        w1: &crate::model::GpuTensor,
11122        aq: &CudaSlice<i8>,
11123        ad: &CudaSlice<f32>,
11124        y0: &mut CudaSlice<f32>,
11125        y1: &mut CudaSlice<f32>,
11126    ) -> Result<bool, Box<dyn std::error::Error>> {
11127        use crate::model::GpuTensor;
11128        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11129        let off =
11130            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11131        if off
11132            || !self.mmvq_supports(QT_NVFP4)
11133            || !self.uses_q8_1_fast(w0)
11134            || !self.uses_q8_1_fast(w1)
11135        {
11136            return Ok(false);
11137        }
11138        let unpack = |w: &crate::model::GpuTensor| match w {
11139            GpuTensor::Quant {
11140                bytes,
11141                qtype,
11142                scale,
11143                rp,
11144                ..
11145            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11146            _ => None,
11147        };
11148        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11149            return Ok(false);
11150        };
11151        let in_f = w0.in_features();
11152        if w1.in_features() != in_f {
11153            return Ok(false);
11154        }
11155        let (o0, o1) = (w0.out_features(), w1.out_features());
11156        if y0.len() < o0 || y1.len() < o1 {
11157            return Ok(false);
11158        }
11159        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11160        const RPW: u32 = 2;
11161        let rows_pb = ROWS_PER_BLOCK * RPW;
11162        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11163        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11164        let cfg = LaunchConfig {
11165            grid_dim: (nb(o0) + nb(o1), 1, 1),
11166            block_dim: (32, ROWS_PER_BLOCK, 1),
11167            shared_mem_bytes: 0,
11168        };
11169        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11170        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11171        // only dereferenced for the launch-arg build inside this call.
11172        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11173        let __s_b = self.gpu.stream();
11174        let mut b = __s_b.launch_builder(&f);
11175        b.arg(b0)
11176            .arg(b1)
11177            .arg(aq)
11178            .arg(ad)
11179            .arg(&mut *y0)
11180            .arg(&mut *y1)
11181            .arg(&inf)
11182            .arg(&oi0)
11183            .arg(&oi1)
11184            .arg(&mi)
11185            .arg(&p0.1)
11186            .arg(&p1.1);
11187        unsafe {
11188            b.launch(cfg)?;
11189        }
11190        Ok(true)
11191    }
11192
11193    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
11194    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
11195    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
11196    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
11197    #[allow(clippy::type_complexity)]
11198    pub fn matmul_nvfp4_fused4(
11199        &self,
11200        w0: &crate::model::GpuTensor,
11201        w1: &crate::model::GpuTensor,
11202        w2: &crate::model::GpuTensor,
11203        w3: &crate::model::GpuTensor,
11204        aq: &CudaSlice<i8>,
11205        ad: &CudaSlice<f32>,
11206        m: usize,
11207    ) -> Result<
11208        Option<(
11209            CudaSlice<f32>,
11210            CudaSlice<f32>,
11211            CudaSlice<f32>,
11212            CudaSlice<f32>,
11213        )>,
11214        Box<dyn std::error::Error>,
11215    > {
11216        use crate::model::GpuTensor;
11217        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
11218        if m != 1
11219            || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11220            || !self.mmvq_supports(QT_NVFP4)
11221            || !self.uses_q8_1_fast(w0)
11222            || !self.uses_q8_1_fast(w1)
11223            || !self.uses_q8_1_fast(w2)
11224            || !self.uses_q8_1_fast(w3)
11225        {
11226            return Ok(None);
11227        }
11228        let unpack = |w: &crate::model::GpuTensor| match w {
11229            GpuTensor::Quant {
11230                bytes,
11231                qtype,
11232                scale,
11233                rp,
11234                ..
11235            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11236            _ => None,
11237        };
11238        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11239            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11240        else {
11241            return Ok(None);
11242        };
11243        let in_f = w0.in_features();
11244        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11245            return Ok(None);
11246        }
11247        let (o0, o1, o2, o3) = (
11248            w0.out_features(),
11249            w1.out_features(),
11250            w2.out_features(),
11251            w3.out_features(),
11252        );
11253        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11254        const RPW: u32 = 2;
11255        let rows_pb = ROWS_PER_BLOCK * RPW;
11256        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11257        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11258        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11259        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11260        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11261        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11262        let cfg = LaunchConfig {
11263            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11264            block_dim: (32, ROWS_PER_BLOCK, 1),
11265            shared_mem_bytes: 0,
11266        };
11267        let (inf, oi0, oi1, oi2, oi3, mi) = (
11268            in_f as i32,
11269            o0 as i32,
11270            o1 as i32,
11271            o2 as i32,
11272            o3 as i32,
11273            m as i32,
11274        );
11275        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11276        // only dereferenced for the launch-arg build inside this call.
11277        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11278        let __s_b = self.gpu.stream();
11279        let mut b = __s_b.launch_builder(&f);
11280        b.arg(b0)
11281            .arg(b1)
11282            .arg(b2)
11283            .arg(b3)
11284            .arg(aq)
11285            .arg(ad)
11286            .arg(&mut y0)
11287            .arg(&mut y1)
11288            .arg(&mut y2)
11289            .arg(&mut y3)
11290            .arg(&inf)
11291            .arg(&oi0)
11292            .arg(&oi1)
11293            .arg(&oi2)
11294            .arg(&oi3)
11295            .arg(&mi)
11296            .arg(&p0.1)
11297            .arg(&p1.1)
11298            .arg(&p2.1)
11299            .arg(&p3.1);
11300        unsafe {
11301            b.launch(cfg)?;
11302        }
11303        Ok(Some((y0, y1, y2, y3)))
11304    }
11305
11306    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
11307    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
11308    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
11309    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
11310    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
11311    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
11312    /// back to the per-tensor path.
11313    pub fn matmul_q8_fused2(
11314        &self,
11315        w0: &crate::model::GpuTensor,
11316        w1: &crate::model::GpuTensor,
11317        aq: &CudaSlice<i8>,
11318        ad: &CudaSlice<f32>,
11319    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11320        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
11321        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
11322        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
11323        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
11324        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
11325        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11326            return Ok(Some(self.e4m3_fused2_core(
11327                p0.0,
11328                p1.0,
11329                aq,
11330                ad,
11331                w0.in_features(),
11332                p0.1,
11333                p1.1,
11334                p0.2,
11335                p0.3,
11336                p1.3,
11337            )?));
11338        }
11339        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11340            return Ok(None);
11341        };
11342        Ok(Some(self.q8_fused2_core(
11343            p0.0,
11344            p1.0,
11345            aq,
11346            ad,
11347            w0.in_features(),
11348            p0.1,
11349            p1.1,
11350            p0.2,
11351        )?))
11352    }
11353
11354    #[allow(clippy::too_many_arguments)]
11355    fn q8_fused2_core(
11356        &self,
11357        b0: &CudaSlice<u8>,
11358        b1: &CudaSlice<u8>,
11359        aq: &CudaSlice<i8>,
11360        ad: &CudaSlice<f32>,
11361        in_f: usize,
11362        out0: usize,
11363        out1: usize,
11364        row_bytes: usize,
11365    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11366        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11367        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11368        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11369        let f = self.func("qmatvec_q8_0_mmvq_fused2");
11370        let mut y0 = self.alloc_uninit::<f32>(out0)?;
11371        let mut y1 = self.alloc_uninit::<f32>(out1)?;
11372        let cfg = LaunchConfig {
11373            grid_dim: (nb0 + nb1, 1, 1),
11374            block_dim: (32, ROWS_PER_BLOCK, 1),
11375            shared_mem_bytes: 0,
11376        };
11377        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11378        let __s_b = self.gpu.stream();
11379        let mut b = __s_b.launch_builder(&f);
11380        b.arg(b0)
11381            .arg(b1)
11382            .arg(aq)
11383            .arg(ad)
11384            .arg(&mut y0)
11385            .arg(&mut y1)
11386            .arg(&inf)
11387            .arg(&o0)
11388            .arg(&o1)
11389            .arg(&rbl);
11390        unsafe {
11391            b.launch(cfg)?;
11392        }
11393        Ok((y0, y1))
11394    }
11395
11396    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
11397    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
11398    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
11399    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
11400    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
11401    pub fn matmul_q8_fused2_x(
11402        &self,
11403        w0: &crate::model::GpuTensor,
11404        w1: &crate::model::GpuTensor,
11405        x: &CudaSlice<f32>,
11406    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11407        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11408            return Ok(None);
11409        }
11410        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11411            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11412            return Ok(Some(self.e4m3_fused2_core(
11413                p0.0,
11414                p1.0,
11415                &aq,
11416                &ad,
11417                w0.in_features(),
11418                p0.1,
11419                p1.1,
11420                p0.2,
11421                p0.3,
11422                p1.3,
11423            )?));
11424        }
11425        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11426            return Ok(None);
11427        };
11428        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11429        Ok(Some(self.q8_fused2_core(
11430            p0.0,
11431            p1.0,
11432            &aq,
11433            &ad,
11434            w0.in_features(),
11435            p0.1,
11436            p1.1,
11437            p0.2,
11438        )?))
11439    }
11440
11441    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
11442    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
11443    #[allow(clippy::too_many_arguments)]
11444    pub fn qmatvec_q8_fused2_raw(
11445        &self,
11446        b0: &CudaSlice<u8>,
11447        b1: &CudaSlice<u8>,
11448        x: &CudaSlice<f32>,
11449        in_f: usize,
11450        out0: usize,
11451        out1: usize,
11452        row_bytes: usize,
11453    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11454        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11455        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11456    }
11457
11458    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
11459    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
11460    /// (tensor,row) to three separate m=1 MMVQ launches.
11461    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
11462    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
11463    pub fn matmul_q4_fused3(
11464        &self,
11465        w0: &crate::model::GpuTensor,
11466        w1: &crate::model::GpuTensor,
11467        w2: &crate::model::GpuTensor,
11468        aq: &CudaSlice<i8>,
11469        ad: &CudaSlice<f32>,
11470    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11471    {
11472        use crate::model::GpuTensor;
11473        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11474            match w {
11475                GpuTensor::Quant {
11476                    qtype, row_bytes, ..
11477                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11478                _ => None,
11479            }
11480        };
11481        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11482            return Ok(None);
11483        };
11484        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11485            return Ok(None);
11486        }
11487        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
11488        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
11489        // the separate matvecs (each routes its own rp).
11490        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11491            match w {
11492                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11493                    Some(m) => (m, true),
11494                    None => (bytes, *rp),
11495                },
11496                _ => unreachable!(),
11497            }
11498        }
11499        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11500        if rp0 != rp1 || rp1 != rp2 {
11501            return Ok(None);
11502        }
11503        let rp = rp0;
11504        let rpb: u32 = 4;
11505        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
11506        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
11507        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
11508        let mr1 = rp && Self::q40_mr1_on();
11509        let nb = |o: usize| {
11510            if mr1 {
11511                (o as u32).div_ceil(rpb)
11512            } else {
11513                (o as u32).div_ceil(2).div_ceil(rpb)
11514            }
11515        };
11516        let grid = nb(o0) + nb(o1) + nb(o2);
11517        let mut y0 = self.alloc_uninit::<f32>(o0)?;
11518        let mut y1 = self.alloc_uninit::<f32>(o1)?;
11519        let mut y2 = self.alloc_uninit::<f32>(o2)?;
11520        let f = self.func(if mr1 {
11521            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11522        } else if rp {
11523            "qmatvec_q4_0_mmvq_fused3_rp"
11524        } else {
11525            "qmatvec_q4_0_mmvq_fused3"
11526        });
11527        let cfg = LaunchConfig {
11528            grid_dim: (grid, 1, 1),
11529            block_dim: (32, rpb, 1),
11530            shared_mem_bytes: 0,
11531        };
11532        let inf = w0.in_features() as i32;
11533        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11534        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11535        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
11536        // variant may take the programmatic-serialization launch.
11537        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11538            {
11539                use cudarc::driver::{DevicePtr, DevicePtrMut};
11540                let s = &self.gpu.stream();
11541                let (p0, _g0) = b0.device_ptr(s);
11542                let (p1, _g1) = b1.device_ptr(s);
11543                let (p2, _g2) = b2.device_ptr(s);
11544                let (paq, _g3) = aq.device_ptr(s);
11545                let (pad, _g4) = ad.device_ptr(s);
11546                let (py0, _g5) = y0.device_ptr_mut(s);
11547                let (py1, _g6) = y1.device_ptr_mut(s);
11548                let (py2, _g7) = y2.device_ptr_mut(s);
11549                let mut ps = [
11550                    &p0 as *const _ as *mut std::ffi::c_void,
11551                    &p1 as *const _ as *mut _,
11552                    &p2 as *const _ as *mut _,
11553                    &paq as *const _ as *mut _,
11554                    &pad as *const _ as *mut _,
11555                    &py0 as *const _ as *mut _,
11556                    &py1 as *const _ as *mut _,
11557                    &py2 as *const _ as *mut _,
11558                    &inf as *const _ as *mut _,
11559                    &oo0 as *const _ as *mut _,
11560                    &oo1 as *const _ as *mut _,
11561                    &oo2 as *const _ as *mut _,
11562                    &r0 as *const _ as *mut _,
11563                    &r1 as *const _ as *mut _,
11564                    &r2 as *const _ as *mut _,
11565                ];
11566                unsafe {
11567                    self.launch_pdl(
11568                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11569                        (grid, 1, 1),
11570                        (32, rpb, 1),
11571                        &mut ps,
11572                    )?;
11573                }
11574            }
11575            return Ok(Some((y0, y1, y2)));
11576        }
11577        let __s_b = self.gpu.stream();
11578        let mut b = __s_b.launch_builder(&f);
11579        b.arg(b0)
11580            .arg(b1)
11581            .arg(b2)
11582            .arg(aq)
11583            .arg(ad)
11584            .arg(&mut y0)
11585            .arg(&mut y1)
11586            .arg(&mut y2)
11587            .arg(&inf)
11588            .arg(&oo0)
11589            .arg(&oo1)
11590            .arg(&oo2)
11591            .arg(&r0)
11592            .arg(&r1)
11593            .arg(&r2);
11594        unsafe {
11595            b.launch(cfg)?;
11596        }
11597        Ok(Some((y0, y1, y2)))
11598    }
11599
11600    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
11601    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
11602    #[allow(clippy::too_many_arguments)]
11603    pub fn matmul_q4_fused3_into(
11604        &self,
11605        w0: &crate::model::GpuTensor,
11606        w1: &crate::model::GpuTensor,
11607        w2: &crate::model::GpuTensor,
11608        aq: &CudaSlice<i8>,
11609        ad: &CudaSlice<f32>,
11610        y0: &mut CudaSlice<f32>,
11611        y1: &mut CudaSlice<f32>,
11612        y2: &mut CudaSlice<f32>,
11613    ) -> Result<bool, Box<dyn std::error::Error>> {
11614        use crate::model::GpuTensor;
11615        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11616            match w {
11617                GpuTensor::Quant {
11618                    qtype, row_bytes, ..
11619                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11620                _ => None,
11621            }
11622        };
11623        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11624            return Ok(false);
11625        };
11626        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11627            return Ok(false);
11628        }
11629        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11630            match w {
11631                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11632                    Some(m) => (m, true),
11633                    None => (bytes, *rp),
11634                },
11635                _ => unreachable!(),
11636            }
11637        }
11638        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11639        if rp0 != rp1 || rp1 != rp2 {
11640            return Ok(false);
11641        }
11642        let rp = rp0;
11643        let rpb: u32 = 4;
11644        let mr1 = rp && Self::q40_mr1_on();
11645        let nb = |o: usize| {
11646            if mr1 {
11647                (o as u32).div_ceil(rpb)
11648            } else {
11649                (o as u32).div_ceil(2).div_ceil(rpb)
11650            }
11651        };
11652        let grid = nb(o0) + nb(o1) + nb(o2);
11653        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
11654        let f = self.func(if mr1 {
11655            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11656        } else if rp {
11657            "qmatvec_q4_0_mmvq_fused3_rp"
11658        } else {
11659            "qmatvec_q4_0_mmvq_fused3"
11660        });
11661        let cfg = LaunchConfig {
11662            grid_dim: (grid, 1, 1),
11663            block_dim: (32, rpb, 1),
11664            shared_mem_bytes: 0,
11665        };
11666        let inf = w0.in_features() as i32;
11667        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11668        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11669        // PDL wave-A: identical to the owned twin (capture-lane parity).
11670        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11671            use cudarc::driver::{DevicePtr, DevicePtrMut};
11672            let s = &self.gpu.stream();
11673            let (p0, _g0) = b0.device_ptr(s);
11674            let (p1, _g1) = b1.device_ptr(s);
11675            let (p2, _g2) = b2.device_ptr(s);
11676            let (paq, _g3) = aq.device_ptr(s);
11677            let (pad, _g4) = ad.device_ptr(s);
11678            let (py0, _g5) = y0.device_ptr_mut(s);
11679            let (py1, _g6) = y1.device_ptr_mut(s);
11680            let (py2, _g7) = y2.device_ptr_mut(s);
11681            let mut ps = [
11682                &p0 as *const _ as *mut std::ffi::c_void,
11683                &p1 as *const _ as *mut _,
11684                &p2 as *const _ as *mut _,
11685                &paq as *const _ as *mut _,
11686                &pad as *const _ as *mut _,
11687                &py0 as *const _ as *mut _,
11688                &py1 as *const _ as *mut _,
11689                &py2 as *const _ as *mut _,
11690                &inf as *const _ as *mut _,
11691                &oo0 as *const _ as *mut _,
11692                &oo1 as *const _ as *mut _,
11693                &oo2 as *const _ as *mut _,
11694                &r0 as *const _ as *mut _,
11695                &r1 as *const _ as *mut _,
11696                &r2 as *const _ as *mut _,
11697            ];
11698            unsafe {
11699                self.launch_pdl(
11700                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11701                    (grid, 1, 1),
11702                    (32, rpb, 1),
11703                    &mut ps,
11704                )?;
11705            }
11706            return Ok(true);
11707        }
11708        let __s_b = self.gpu.stream();
11709        let mut b = __s_b.launch_builder(&f);
11710        b.arg(b0)
11711            .arg(b1)
11712            .arg(b2)
11713            .arg(aq)
11714            .arg(ad)
11715            .arg(&mut *y0)
11716            .arg(&mut *y1)
11717            .arg(&mut *y2)
11718            .arg(&inf)
11719            .arg(&oo0)
11720            .arg(&oo1)
11721            .arg(&oo2)
11722            .arg(&r0)
11723            .arg(&r1)
11724            .arg(&r2);
11725        unsafe {
11726            b.launch(cfg)?;
11727        }
11728        Ok(true)
11729    }
11730
11731    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
11732    pub fn matmul_q4_fused2(
11733        &self,
11734        w0: &crate::model::GpuTensor,
11735        w1: &crate::model::GpuTensor,
11736        aq: &CudaSlice<i8>,
11737        ad: &CudaSlice<f32>,
11738    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11739        use crate::model::GpuTensor;
11740        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11741            match w {
11742                GpuTensor::Quant {
11743                    qtype, row_bytes, ..
11744                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11745                _ => None,
11746            }
11747        };
11748        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11749            return Ok(None);
11750        };
11751        if w0.in_features() != w1.in_features() {
11752            return Ok(None);
11753        }
11754        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
11755        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11756            match w {
11757                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11758                    Some(m) => (m, true),
11759                    None => (bytes, *rp),
11760                },
11761                _ => unreachable!(),
11762            }
11763        }
11764        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11765        if rp0 != rp1 {
11766            return Ok(None);
11767        }
11768        let rp = rp0;
11769        let rpb: u32 = 4;
11770        // mr1 twin — see matmul_q4_fused3.
11771        let mr1 = rp && Self::q40_mr1_on();
11772        let nb = |o: usize| {
11773            if mr1 {
11774                (o as u32).div_ceil(rpb)
11775            } else {
11776                (o as u32).div_ceil(2).div_ceil(rpb)
11777            }
11778        };
11779        let grid = nb(o0) + nb(o1);
11780        let mut y0 = self.alloc_uninit::<f32>(o0)?;
11781        let mut y1 = self.alloc_uninit::<f32>(o1)?;
11782        let f = self.func(if mr1 {
11783            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
11784        } else if rp {
11785            "qmatvec_q4_0_mmvq_fused2_rp"
11786        } else {
11787            "qmatvec_q4_0_mmvq_fused2"
11788        });
11789        let cfg = LaunchConfig {
11790            grid_dim: (grid, 1, 1),
11791            block_dim: (32, rpb, 1),
11792            shared_mem_bytes: 0,
11793        };
11794        let inf = w0.in_features() as i32;
11795        let (oo0, oo1) = (o0 as i32, o1 as i32);
11796        let (r0, r1) = (rb0 as i64, rb1 as i64);
11797        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
11798        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11799            {
11800                use cudarc::driver::{DevicePtr, DevicePtrMut};
11801                let s = &self.gpu.stream();
11802                let (p0, _g0) = b0.device_ptr(s);
11803                let (p1, _g1) = b1.device_ptr(s);
11804                let (paq, _g2) = aq.device_ptr(s);
11805                let (pad, _g3) = ad.device_ptr(s);
11806                let (py0, _g4) = y0.device_ptr_mut(s);
11807                let (py1, _g5) = y1.device_ptr_mut(s);
11808                let mut ps = [
11809                    &p0 as *const _ as *mut std::ffi::c_void,
11810                    &p1 as *const _ as *mut _,
11811                    &paq as *const _ as *mut _,
11812                    &pad as *const _ as *mut _,
11813                    &py0 as *const _ as *mut _,
11814                    &py1 as *const _ as *mut _,
11815                    &inf as *const _ as *mut _,
11816                    &oo0 as *const _ as *mut _,
11817                    &oo1 as *const _ as *mut _,
11818                    &r0 as *const _ as *mut _,
11819                    &r1 as *const _ as *mut _,
11820                ];
11821                unsafe {
11822                    self.launch_pdl(
11823                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
11824                        (grid, 1, 1),
11825                        (32, rpb, 1),
11826                        &mut ps,
11827                    )?;
11828                }
11829            }
11830            return Ok(Some((y0, y1)));
11831        }
11832        let __s_b = self.gpu.stream();
11833        let mut b = __s_b.launch_builder(&f);
11834        b.arg(b0)
11835            .arg(b1)
11836            .arg(aq)
11837            .arg(ad)
11838            .arg(&mut y0)
11839            .arg(&mut y1)
11840            .arg(&inf)
11841            .arg(&oo0)
11842            .arg(&oo1)
11843            .arg(&r0)
11844            .arg(&r1);
11845        unsafe {
11846            b.launch(cfg)?;
11847        }
11848        Ok(Some((y0, y1)))
11849    }
11850
11851    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
11852    pub fn matmul_q4_fused2_into(
11853        &self,
11854        w0: &crate::model::GpuTensor,
11855        w1: &crate::model::GpuTensor,
11856        aq: &CudaSlice<i8>,
11857        ad: &CudaSlice<f32>,
11858        y0: &mut CudaSlice<f32>,
11859        y1: &mut CudaSlice<f32>,
11860    ) -> Result<bool, Box<dyn std::error::Error>> {
11861        use crate::model::GpuTensor;
11862        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11863            match w {
11864                GpuTensor::Quant {
11865                    qtype, row_bytes, ..
11866                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11867                _ => None,
11868            }
11869        };
11870        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11871            return Ok(false);
11872        };
11873        if w0.in_features() != w1.in_features() {
11874            return Ok(false);
11875        }
11876        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11877            match w {
11878                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11879                    Some(m) => (m, true),
11880                    None => (bytes, *rp),
11881                },
11882                _ => unreachable!(),
11883            }
11884        }
11885        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11886        if rp0 != rp1 {
11887            return Ok(false);
11888        }
11889        let rp = rp0;
11890        let rpb: u32 = 4;
11891        let mr1 = rp && Self::q40_mr1_on();
11892        let nb = |o: usize| {
11893            if mr1 {
11894                (o as u32).div_ceil(rpb)
11895            } else {
11896                (o as u32).div_ceil(2).div_ceil(rpb)
11897            }
11898        };
11899        let grid = nb(o0) + nb(o1);
11900        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
11901        let f = self.func(if mr1 {
11902            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
11903        } else if rp {
11904            "qmatvec_q4_0_mmvq_fused2_rp"
11905        } else {
11906            "qmatvec_q4_0_mmvq_fused2"
11907        });
11908        let cfg = LaunchConfig {
11909            grid_dim: (grid, 1, 1),
11910            block_dim: (32, rpb, 1),
11911            shared_mem_bytes: 0,
11912        };
11913        let inf = w0.in_features() as i32;
11914        let (oo0, oo1) = (o0 as i32, o1 as i32);
11915        let (r0, r1) = (rb0 as i64, rb1 as i64);
11916        // PDL wave-A: identical to the owned twin (capture-lane parity).
11917        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11918            use cudarc::driver::{DevicePtr, DevicePtrMut};
11919            let s = &self.gpu.stream();
11920            let (p0, _g0) = b0.device_ptr(s);
11921            let (p1, _g1) = b1.device_ptr(s);
11922            let (paq, _g2) = aq.device_ptr(s);
11923            let (pad, _g3) = ad.device_ptr(s);
11924            let (py0, _g4) = y0.device_ptr_mut(s);
11925            let (py1, _g5) = y1.device_ptr_mut(s);
11926            let mut ps = [
11927                &p0 as *const _ as *mut std::ffi::c_void,
11928                &p1 as *const _ as *mut _,
11929                &paq as *const _ as *mut _,
11930                &pad as *const _ as *mut _,
11931                &py0 as *const _ as *mut _,
11932                &py1 as *const _ as *mut _,
11933                &inf as *const _ as *mut _,
11934                &oo0 as *const _ as *mut _,
11935                &oo1 as *const _ as *mut _,
11936                &r0 as *const _ as *mut _,
11937                &r1 as *const _ as *mut _,
11938            ];
11939            unsafe {
11940                self.launch_pdl(
11941                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
11942                    (grid, 1, 1),
11943                    (32, rpb, 1),
11944                    &mut ps,
11945                )?;
11946            }
11947            return Ok(true);
11948        }
11949        let __s_b = self.gpu.stream();
11950        let mut b = __s_b.launch_builder(&f);
11951        b.arg(b0)
11952            .arg(b1)
11953            .arg(aq)
11954            .arg(ad)
11955            .arg(&mut *y0)
11956            .arg(&mut *y1)
11957            .arg(&inf)
11958            .arg(&oo0)
11959            .arg(&oo1)
11960            .arg(&r0)
11961            .arg(&r1);
11962        unsafe {
11963            b.launch(cfg)?;
11964        }
11965        Ok(true)
11966    }
11967
11968    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
11969    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
11970    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
11971    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
11972    pub fn matmul_q4_fused2_batched(
11973        &self,
11974        w0: &crate::model::GpuTensor,
11975        w1: &crate::model::GpuTensor,
11976        aq: &CudaSlice<i8>,
11977        ad: &CudaSlice<f32>,
11978        m: usize,
11979    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11980        use crate::model::GpuTensor;
11981        if m < 2 || m > 8 {
11982            return Ok(None);
11983        }
11984        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11985            match w {
11986                GpuTensor::Quant {
11987                    qtype, row_bytes, ..
11988                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11989                _ => None,
11990            }
11991        };
11992        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
11993            return Ok(None);
11994        };
11995        if w0.in_features() != w1.in_features() {
11996            return Ok(None);
11997        }
11998        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11999            match w {
12000                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12001                    Some(mr) => (mr, true),
12002                    None => (bytes, *rp),
12003                },
12004                _ => unreachable!(),
12005            }
12006        }
12007        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12008        if !rp0 || !rp1 {
12009            return Ok(None);
12010        }
12011        let mcols = Self::batched_mcols(m);
12012        let rpb: u32 = 4;
12013        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12014        let grid = nb(o0) + nb(o1);
12015        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12016        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12017        let f = self.func(match mcols {
12018            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12019            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12020            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12021        });
12022        let cfg = LaunchConfig {
12023            grid_dim: (grid, 1, 1),
12024            block_dim: (32, rpb, 1),
12025            shared_mem_bytes: 0,
12026        };
12027        let inf = w0.in_features() as i32;
12028        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12029        let rb = rb0 as i64;
12030        let __s_b = self.gpu.stream();
12031        let mut b = __s_b.launch_builder(&f);
12032        b.arg(b0)
12033            .arg(b1)
12034            .arg(aq)
12035            .arg(ad)
12036            .arg(&mut y0)
12037            .arg(&mut y1)
12038            .arg(&inf)
12039            .arg(&oo0)
12040            .arg(&oo1)
12041            .arg(&mi)
12042            .arg(&rb);
12043        unsafe {
12044            b.launch(cfg)?;
12045        }
12046        Ok(Some((y0, y1)))
12047    }
12048
12049    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
12050    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
12051    #[allow(clippy::too_many_arguments)]
12052    pub fn matmul_q4_fused3_batched(
12053        &self,
12054        w0: &crate::model::GpuTensor,
12055        w1: &crate::model::GpuTensor,
12056        w2: &crate::model::GpuTensor,
12057        aq: &CudaSlice<i8>,
12058        ad: &CudaSlice<f32>,
12059        m: usize,
12060    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12061    {
12062        use crate::model::GpuTensor;
12063        if m < 2 || m > 8 {
12064            return Ok(None);
12065        }
12066        let q4 = |w: &GpuTensor| -> Option<usize> {
12067            match w {
12068                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12069                _ => None,
12070            }
12071        };
12072        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12073            return Ok(None);
12074        };
12075        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12076            return Ok(None);
12077        }
12078        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12079            match w {
12080                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12081                    Some(mr) => (mr, true),
12082                    None => (bytes, *rp),
12083                },
12084                _ => unreachable!(),
12085            }
12086        }
12087        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12088        if !rp0 || !rp1 || !rp2 {
12089            return Ok(None);
12090        }
12091        let mcols = Self::batched_mcols(m);
12092        let rpb: u32 = 4;
12093        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12094        let grid = nb(o0) + nb(o1) + nb(o2);
12095        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12096        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12097        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12098        let f = self.func(match mcols {
12099            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12100            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12101            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12102        });
12103        let cfg = LaunchConfig {
12104            grid_dim: (grid, 1, 1),
12105            block_dim: (32, rpb, 1),
12106            shared_mem_bytes: 0,
12107        };
12108        let inf = w0.in_features() as i32;
12109        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12110        let rb = 0i64;
12111        let __s_b = self.gpu.stream();
12112        let mut b = __s_b.launch_builder(&f);
12113        b.arg(b0)
12114            .arg(b1)
12115            .arg(b2)
12116            .arg(aq)
12117            .arg(ad)
12118            .arg(&mut y0)
12119            .arg(&mut y1)
12120            .arg(&mut y2)
12121            .arg(&inf)
12122            .arg(&oo0)
12123            .arg(&oo1)
12124            .arg(&oo2)
12125            .arg(&mi)
12126            .arg(&rb);
12127        unsafe {
12128            b.launch(cfg)?;
12129        }
12130        Ok(Some((y0, y1, y2)))
12131    }
12132
12133    pub fn matmul_q8_fused3(
12134        &self,
12135        w0: &crate::model::GpuTensor,
12136        w1: &crate::model::GpuTensor,
12137        w2: &crate::model::GpuTensor,
12138        aq: &CudaSlice<i8>,
12139        ad: &CudaSlice<f32>,
12140    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12141    {
12142        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
12143        // are per-tensor FP8, so native residency without this arm meant three separate launches.
12144        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12145            return Ok(Some(self.e4m3_fused3_core(
12146                p0.0,
12147                p1.0,
12148                p2.0,
12149                aq,
12150                ad,
12151                w0.in_features(),
12152                p0.1,
12153                p1.1,
12154                p2.1,
12155                p0.2,
12156                p0.3,
12157                p1.3,
12158                p2.3,
12159            )?));
12160        }
12161        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12162            return Ok(None);
12163        };
12164        Ok(Some(self.q8_fused3_core(
12165            p0.0,
12166            p1.0,
12167            p2.0,
12168            aq,
12169            ad,
12170            w0.in_features(),
12171            p0.1,
12172            p1.1,
12173            p2.1,
12174            p0.2,
12175        )?))
12176    }
12177
12178    #[allow(clippy::too_many_arguments)]
12179    fn q8_fused3_core(
12180        &self,
12181        b0: &CudaSlice<u8>,
12182        b1: &CudaSlice<u8>,
12183        b2: &CudaSlice<u8>,
12184        aq: &CudaSlice<i8>,
12185        ad: &CudaSlice<f32>,
12186        in_f: usize,
12187        out0: usize,
12188        out1: usize,
12189        out2: usize,
12190        row_bytes: usize,
12191    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12192        const ROWS_PER_BLOCK: u32 = 4;
12193        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12194        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12195        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12196        let f = self.func("qmatvec_q8_0_mmvq_fused3");
12197        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12198        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12199        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12200        let cfg = LaunchConfig {
12201            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12202            block_dim: (32, ROWS_PER_BLOCK, 1),
12203            shared_mem_bytes: 0,
12204        };
12205        let (inf, o0, o1, o2, rbl) = (
12206            in_f as i32,
12207            out0 as i32,
12208            out1 as i32,
12209            out2 as i32,
12210            row_bytes as i64,
12211        );
12212        let __s_b = self.gpu.stream();
12213        let mut b = __s_b.launch_builder(&f);
12214        b.arg(b0)
12215            .arg(b1)
12216            .arg(b2)
12217            .arg(aq)
12218            .arg(ad)
12219            .arg(&mut y0)
12220            .arg(&mut y1)
12221            .arg(&mut y2)
12222            .arg(&inf)
12223            .arg(&o0)
12224            .arg(&o1)
12225            .arg(&o2)
12226            .arg(&rbl);
12227        unsafe {
12228            b.launch(cfg)?;
12229        }
12230        Ok((y0, y1, y2))
12231    }
12232
12233    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
12234    #[allow(clippy::too_many_arguments)]
12235    pub fn qmatvec_q8_fused3_raw(
12236        &self,
12237        b0: &CudaSlice<u8>,
12238        b1: &CudaSlice<u8>,
12239        b2: &CudaSlice<u8>,
12240        x: &CudaSlice<f32>,
12241        in_f: usize,
12242        out0: usize,
12243        out1: usize,
12244        out2: usize,
12245        row_bytes: usize,
12246    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12247        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12248        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12249    }
12250
12251    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
12252    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
12253    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
12254    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
12255    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
12256    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
12257    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
12258    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
12259    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
12260    /// twin must not introduce a batched program the reference path would not run).
12261    pub fn matmul_q8_fused2_t(
12262        &self,
12263        w0: &crate::model::GpuTensor,
12264        w1: &crate::model::GpuTensor,
12265        aq: &CudaSlice<i8>,
12266        ad: &CudaSlice<f32>,
12267        m: usize,
12268    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12269        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
12270        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
12271        // fuses too — same template body, still bit-identical to the two _b8 launches.
12272        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12273            return Ok(None);
12274        }
12275        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
12276        // so the fused b8 launch would introduce a batched program the reference path would not run.
12277        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12278            if m > 4 && !Self::b8_enabled() {
12279                return Ok(None);
12280            }
12281            return Ok(Some(self.e4m3_fused2_t_core(
12282                p0.0,
12283                p1.0,
12284                aq,
12285                ad,
12286                m,
12287                w0.in_features(),
12288                p0.1,
12289                p1.1,
12290                p0.2,
12291                p0.3,
12292                p1.3,
12293            )?));
12294        }
12295        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12296            return Ok(None);
12297        };
12298        Ok(Some(self.q8_fused2_t_core(
12299            p0.0,
12300            p1.0,
12301            aq,
12302            ad,
12303            m,
12304            w0.in_features(),
12305            p0.1,
12306            p1.1,
12307            p0.2,
12308        )?))
12309    }
12310
12311    #[allow(clippy::too_many_arguments)]
12312    fn q8_fused2_t_core(
12313        &self,
12314        b0: &CudaSlice<u8>,
12315        b1: &CudaSlice<u8>,
12316        aq: &CudaSlice<i8>,
12317        ad: &CudaSlice<f32>,
12318        m: usize,
12319        in_f: usize,
12320        out0: usize,
12321        out1: usize,
12322        row_bytes: usize,
12323    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12324        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12325        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12326        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12327        let f = self.func(match Self::batched_mcols(m) {
12328            2 => "qmatvec_q8_0_mmvq_fused2_b2",
12329            4 => "qmatvec_q8_0_mmvq_fused2_b4",
12330            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
12331            _ => "qmatvec_q8_0_mmvq_fused2_b8",
12332        });
12333        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12334        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12335        let cfg = LaunchConfig {
12336            grid_dim: (nb0 + nb1, 1, 1),
12337            block_dim: (32, ROWS_PER_BLOCK, 1),
12338            shared_mem_bytes: 0,
12339        };
12340        let (inf, o0, o1, mi, rbl) = (
12341            in_f as i32,
12342            out0 as i32,
12343            out1 as i32,
12344            m as i32,
12345            row_bytes as i64,
12346        );
12347        let __s_b = self.gpu.stream();
12348        let mut b = __s_b.launch_builder(&f);
12349        b.arg(b0)
12350            .arg(b1)
12351            .arg(aq)
12352            .arg(ad)
12353            .arg(&mut y0)
12354            .arg(&mut y1)
12355            .arg(&inf)
12356            .arg(&o0)
12357            .arg(&o1)
12358            .arg(&mi)
12359            .arg(&rbl);
12360        unsafe {
12361            b.launch(cfg)?;
12362        }
12363        Ok((y0, y1))
12364    }
12365
12366    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
12367    /// q8_1 quant of the [m, in_f] activation), no env gating.
12368    #[allow(clippy::too_many_arguments)]
12369    pub fn qmatvec_q8_fused2_t_raw(
12370        &self,
12371        b0: &CudaSlice<u8>,
12372        b1: &CudaSlice<u8>,
12373        x: &CudaSlice<f32>,
12374        m: usize,
12375        in_f: usize,
12376        out0: usize,
12377        out1: usize,
12378        row_bytes: usize,
12379    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12380        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12381        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12382    }
12383
12384    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
12385    /// `matmul_q8_fused2_t` with three ranges.
12386    #[allow(clippy::too_many_arguments)]
12387    pub fn matmul_q8_fused3_t(
12388        &self,
12389        w0: &crate::model::GpuTensor,
12390        w1: &crate::model::GpuTensor,
12391        w2: &crate::model::GpuTensor,
12392        aq: &CudaSlice<i8>,
12393        ad: &CudaSlice<f32>,
12394        m: usize,
12395    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12396    {
12397        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12398            return Ok(None);
12399        }
12400        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12401            return Ok(Some(self.e4m3_fused3_t_core(
12402                p0.0,
12403                p1.0,
12404                p2.0,
12405                aq,
12406                ad,
12407                m,
12408                w0.in_features(),
12409                p0.1,
12410                p1.1,
12411                p2.1,
12412                p0.2,
12413                p0.3,
12414                p1.3,
12415                p2.3,
12416            )?));
12417        }
12418        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12419            return Ok(None);
12420        };
12421        Ok(Some(self.q8_fused3_t_core(
12422            p0.0,
12423            p1.0,
12424            p2.0,
12425            aq,
12426            ad,
12427            m,
12428            w0.in_features(),
12429            p0.1,
12430            p1.1,
12431            p2.1,
12432            p0.2,
12433        )?))
12434    }
12435
12436    #[allow(clippy::too_many_arguments)]
12437    fn q8_fused3_t_core(
12438        &self,
12439        b0: &CudaSlice<u8>,
12440        b1: &CudaSlice<u8>,
12441        b2: &CudaSlice<u8>,
12442        aq: &CudaSlice<i8>,
12443        ad: &CudaSlice<f32>,
12444        m: usize,
12445        in_f: usize,
12446        out0: usize,
12447        out1: usize,
12448        out2: usize,
12449        row_bytes: usize,
12450    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12451        const ROWS_PER_BLOCK: u32 = 4;
12452        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12453        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12454        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12455        let f = self.func(if Self::batched_mcols(m) == 2 {
12456            "qmatvec_q8_0_mmvq_fused3_b2"
12457        } else {
12458            "qmatvec_q8_0_mmvq_fused3_b4"
12459        });
12460        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12461        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12462        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12463        let cfg = LaunchConfig {
12464            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12465            block_dim: (32, ROWS_PER_BLOCK, 1),
12466            shared_mem_bytes: 0,
12467        };
12468        let (inf, o0, o1, o2, mi, rbl) = (
12469            in_f as i32,
12470            out0 as i32,
12471            out1 as i32,
12472            out2 as i32,
12473            m as i32,
12474            row_bytes as i64,
12475        );
12476        let __s_b = self.gpu.stream();
12477        let mut b = __s_b.launch_builder(&f);
12478        b.arg(b0)
12479            .arg(b1)
12480            .arg(b2)
12481            .arg(aq)
12482            .arg(ad)
12483            .arg(&mut y0)
12484            .arg(&mut y1)
12485            .arg(&mut y2)
12486            .arg(&inf)
12487            .arg(&o0)
12488            .arg(&o1)
12489            .arg(&o2)
12490            .arg(&mi)
12491            .arg(&rbl);
12492        unsafe {
12493            b.launch(cfg)?;
12494        }
12495        Ok((y0, y1, y2))
12496    }
12497
12498    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
12499    #[allow(clippy::too_many_arguments)]
12500    pub fn qmatvec_q8_fused3_t_raw(
12501        &self,
12502        b0: &CudaSlice<u8>,
12503        b1: &CudaSlice<u8>,
12504        b2: &CudaSlice<u8>,
12505        x: &CudaSlice<f32>,
12506        m: usize,
12507        in_f: usize,
12508        out0: usize,
12509        out1: usize,
12510        out2: usize,
12511        row_bytes: usize,
12512    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12513        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12514        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12515    }
12516
12517    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
12518    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
12519    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
12520    pub fn q8_ffn_fuse2_on(&self) -> bool {
12521        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12522        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12523    }
12524
12525    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
12526    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
12527    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
12528    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
12529    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
12530    #[allow(clippy::type_complexity)]
12531    fn q8_fused_params<'w, const N: usize>(
12532        &self,
12533        ws: &[&'w crate::model::GpuTensor; N],
12534    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12535        use crate::model::GpuTensor;
12536        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12537            return None;
12538        }
12539        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12540            return None;
12541        }
12542        let in_f = ws[0].in_features();
12543        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12544        for (i, w) in ws.iter().enumerate() {
12545            match w {
12546                GpuTensor::Quant {
12547                    bytes,
12548                    qtype,
12549                    row_bytes,
12550                    scale,
12551                    ..
12552                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12553                    out[i] = Some((bytes, w.out_features(), *row_bytes))
12554                }
12555                _ => return None,
12556            }
12557        }
12558        Some(out.map(|o| o.unwrap()))
12559    }
12560
12561    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
12562    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
12563    pub fn e4m3_dual_on(&self) -> bool {
12564        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12565        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12566    }
12567
12568    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
12569    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
12570    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
12571    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
12572    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
12573    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
12574    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
12575    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
12576    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
12577    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
12578    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
12579    #[allow(clippy::type_complexity)]
12580    fn e4m3_fused_params<'w, const N: usize>(
12581        &self,
12582        ws: &[&'w crate::model::GpuTensor; N],
12583    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12584        use crate::model::GpuTensor;
12585        if !self.e4m3_dual_on() {
12586            return None;
12587        }
12588        let in_f = ws[0].in_features();
12589        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12590        for (i, w) in ws.iter().enumerate() {
12591            match w {
12592                GpuTensor::Quant {
12593                    bytes,
12594                    qtype,
12595                    row_bytes,
12596                    scale,
12597                    rp,
12598                    rp4,
12599                    ..
12600                } if *qtype == QT_F8_E4M3
12601                    && w.in_features() == in_f
12602                    && *row_bytes == in_f
12603                    && !*rp
12604                    && rp4.is_none() =>
12605                {
12606                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12607                }
12608                _ => return None,
12609            }
12610        }
12611        Some(out.map(|o| o.unwrap()))
12612    }
12613
12614    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
12615    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
12616    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
12617    #[allow(clippy::too_many_arguments)]
12618    fn e4m3_fused2_core(
12619        &self,
12620        b0: &CudaSlice<u8>,
12621        b1: &CudaSlice<u8>,
12622        aq: &CudaSlice<i8>,
12623        ad: &CudaSlice<f32>,
12624        in_f: usize,
12625        out0: usize,
12626        out1: usize,
12627        row_bytes: usize,
12628        ws0: f32,
12629        ws1: f32,
12630    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12631        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12632        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12633        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12634        let f = self.func("qmatvec_e4m3_mmvq_fused2");
12635        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12636        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12637        let cfg = LaunchConfig {
12638            grid_dim: (nb0 + nb1, 1, 1),
12639            block_dim: (32, ROWS_PER_BLOCK, 1),
12640            shared_mem_bytes: 0,
12641        };
12642        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12643        let __s_b = self.gpu.stream();
12644        let mut b = __s_b.launch_builder(&f);
12645        b.arg(b0)
12646            .arg(b1)
12647            .arg(aq)
12648            .arg(ad)
12649            .arg(&mut y0)
12650            .arg(&mut y1)
12651            .arg(&inf)
12652            .arg(&o0)
12653            .arg(&o1)
12654            .arg(&rbl)
12655            .arg(&ws0)
12656            .arg(&ws1);
12657        unsafe {
12658            b.launch(cfg)?;
12659        }
12660        Ok((y0, y1))
12661    }
12662
12663    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
12664    #[allow(clippy::too_many_arguments)]
12665    fn e4m3_fused3_core(
12666        &self,
12667        b0: &CudaSlice<u8>,
12668        b1: &CudaSlice<u8>,
12669        b2: &CudaSlice<u8>,
12670        aq: &CudaSlice<i8>,
12671        ad: &CudaSlice<f32>,
12672        in_f: usize,
12673        out0: usize,
12674        out1: usize,
12675        out2: usize,
12676        row_bytes: usize,
12677        ws0: f32,
12678        ws1: f32,
12679        ws2: f32,
12680    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12681        const ROWS_PER_BLOCK: u32 = 4;
12682        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12683        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12684        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12685        let f = self.func("qmatvec_e4m3_mmvq_fused3");
12686        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12687        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12688        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12689        let cfg = LaunchConfig {
12690            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12691            block_dim: (32, ROWS_PER_BLOCK, 1),
12692            shared_mem_bytes: 0,
12693        };
12694        let (inf, o0, o1, o2, rbl) = (
12695            in_f as i32,
12696            out0 as i32,
12697            out1 as i32,
12698            out2 as i32,
12699            row_bytes as i64,
12700        );
12701        let __s_b = self.gpu.stream();
12702        let mut b = __s_b.launch_builder(&f);
12703        b.arg(b0)
12704            .arg(b1)
12705            .arg(b2)
12706            .arg(aq)
12707            .arg(ad)
12708            .arg(&mut y0)
12709            .arg(&mut y1)
12710            .arg(&mut y2)
12711            .arg(&inf)
12712            .arg(&o0)
12713            .arg(&o1)
12714            .arg(&o2)
12715            .arg(&rbl)
12716            .arg(&ws0)
12717            .arg(&ws1)
12718            .arg(&ws2);
12719        unsafe {
12720            b.launch(cfg)?;
12721        }
12722        Ok((y0, y1, y2))
12723    }
12724
12725    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
12726    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
12727    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
12728    #[allow(clippy::too_many_arguments)]
12729    fn e4m3_fused2_t_core(
12730        &self,
12731        b0: &CudaSlice<u8>,
12732        b1: &CudaSlice<u8>,
12733        aq: &CudaSlice<i8>,
12734        ad: &CudaSlice<f32>,
12735        m: usize,
12736        in_f: usize,
12737        out0: usize,
12738        out1: usize,
12739        row_bytes: usize,
12740        ws0: f32,
12741        ws1: f32,
12742    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12743        const ROWS_PER_BLOCK: u32 = 4;
12744        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12745        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12746        let f = self.func(match Self::batched_mcols(m) {
12747            2 => "qmatvec_e4m3_mmvq_fused2_b2",
12748            4 => "qmatvec_e4m3_mmvq_fused2_b4",
12749            _ => "qmatvec_e4m3_mmvq_fused2_b8",
12750        });
12751        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12752        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12753        let cfg = LaunchConfig {
12754            grid_dim: (nb0 + nb1, 1, 1),
12755            block_dim: (32, ROWS_PER_BLOCK, 1),
12756            shared_mem_bytes: 0,
12757        };
12758        let (inf, o0, o1, mi, rbl) = (
12759            in_f as i32,
12760            out0 as i32,
12761            out1 as i32,
12762            m as i32,
12763            row_bytes as i64,
12764        );
12765        let __s_b = self.gpu.stream();
12766        let mut b = __s_b.launch_builder(&f);
12767        b.arg(b0)
12768            .arg(b1)
12769            .arg(aq)
12770            .arg(ad)
12771            .arg(&mut y0)
12772            .arg(&mut y1)
12773            .arg(&inf)
12774            .arg(&o0)
12775            .arg(&o1)
12776            .arg(&mi)
12777            .arg(&rbl);
12778        unsafe {
12779            b.launch(cfg)?;
12780        }
12781        if ws0 != 1.0 {
12782            self.scale_inplace(&mut y0, ws0, m * out0)?;
12783        }
12784        if ws1 != 1.0 {
12785            self.scale_inplace(&mut y1, ws1, m * out1)?;
12786        }
12787        Ok((y0, y1))
12788    }
12789
12790    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
12791    #[allow(clippy::too_many_arguments)]
12792    fn e4m3_fused3_t_core(
12793        &self,
12794        b0: &CudaSlice<u8>,
12795        b1: &CudaSlice<u8>,
12796        b2: &CudaSlice<u8>,
12797        aq: &CudaSlice<i8>,
12798        ad: &CudaSlice<f32>,
12799        m: usize,
12800        in_f: usize,
12801        out0: usize,
12802        out1: usize,
12803        out2: usize,
12804        row_bytes: usize,
12805        ws0: f32,
12806        ws1: f32,
12807        ws2: f32,
12808    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12809        const ROWS_PER_BLOCK: u32 = 4;
12810        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12811        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12812        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12813        let f = self.func(if Self::batched_mcols(m) == 2 {
12814            "qmatvec_e4m3_mmvq_fused3_b2"
12815        } else {
12816            "qmatvec_e4m3_mmvq_fused3_b4"
12817        });
12818        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12819        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12820        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12821        let cfg = LaunchConfig {
12822            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12823            block_dim: (32, ROWS_PER_BLOCK, 1),
12824            shared_mem_bytes: 0,
12825        };
12826        let (inf, o0, o1, o2, mi, rbl) = (
12827            in_f as i32,
12828            out0 as i32,
12829            out1 as i32,
12830            out2 as i32,
12831            m as i32,
12832            row_bytes as i64,
12833        );
12834        let __s_b = self.gpu.stream();
12835        let mut b = __s_b.launch_builder(&f);
12836        b.arg(b0)
12837            .arg(b1)
12838            .arg(b2)
12839            .arg(aq)
12840            .arg(ad)
12841            .arg(&mut y0)
12842            .arg(&mut y1)
12843            .arg(&mut y2)
12844            .arg(&inf)
12845            .arg(&o0)
12846            .arg(&o1)
12847            .arg(&o2)
12848            .arg(&mi)
12849            .arg(&rbl);
12850        unsafe {
12851            b.launch(cfg)?;
12852        }
12853        if ws0 != 1.0 {
12854            self.scale_inplace(&mut y0, ws0, m * out0)?;
12855        }
12856        if ws1 != 1.0 {
12857            self.scale_inplace(&mut y1, ws1, m * out1)?;
12858        }
12859        if ws2 != 1.0 {
12860            self.scale_inplace(&mut y2, ws2, m * out2)?;
12861        }
12862        Ok((y0, y1, y2))
12863    }
12864
12865    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
12866    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
12867    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
12868    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
12869    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
12870    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
12871    ///
12872    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
12873    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
12874    pub fn qmatvec_e4m3_blk_mmvq(
12875        &self,
12876        bytes: &CudaSlice<u8>,
12877        aq: &CudaSlice<i8>,
12878        ad: &CudaSlice<f32>,
12879        scales: &CudaSlice<f32>,
12880        m: usize,
12881        in_f: usize,
12882        out_f: usize,
12883        row_bytes: usize,
12884        scale_cols: usize,
12885    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12886        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
12887        self.qmatvec_e4m3_blk_mmvq_into(
12888            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
12889        )?;
12890        Ok(y)
12891    }
12892
12893    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
12894    #[allow(clippy::too_many_arguments)]
12895    pub fn qmatvec_e4m3_blk_mmvq_into(
12896        &self,
12897        bytes: &CudaSlice<u8>,
12898        aq: &CudaSlice<i8>,
12899        ad: &CudaSlice<f32>,
12900        scales: &CudaSlice<f32>,
12901        m: usize,
12902        in_f: usize,
12903        out_f: usize,
12904        row_bytes: usize,
12905        scale_cols: usize,
12906        y: &mut CudaSlice<f32>,
12907    ) -> Result<(), Box<dyn std::error::Error>> {
12908        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12909        let f = self.func("qmatvec_e4m3_blk_mmvq");
12910        let cfg = LaunchConfig {
12911            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
12912            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
12913            shared_mem_bytes: 0,                // warp-only reduce
12914        };
12915        let (inf, outf, mi, rb, sc) = (
12916            in_f as i32,
12917            out_f as i32,
12918            m as i32,
12919            row_bytes as i64,
12920            scale_cols as i32,
12921        );
12922        let __s_b = self.gpu.stream();
12923        let mut b = __s_b.launch_builder(&f);
12924        b.arg(bytes)
12925            .arg(aq)
12926            .arg(ad)
12927            .arg(scales)
12928            .arg(&mut *y)
12929            .arg(&inf)
12930            .arg(&outf)
12931            .arg(&mi)
12932            .arg(&rb)
12933            .arg(&sc);
12934        unsafe {
12935            b.launch(cfg)?;
12936        }
12937        Ok(())
12938    }
12939
12940    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
12941    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
12942    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
12943    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
12944    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
12945    #[allow(clippy::too_many_arguments)]
12946    pub fn qmatvec_e4m3_blk_mmvq_batched(
12947        &self,
12948        bytes: &CudaSlice<u8>,
12949        aq: &CudaSlice<i8>,
12950        ad: &CudaSlice<f32>,
12951        scales: &CudaSlice<f32>,
12952        m: usize,
12953        in_f: usize,
12954        out_f: usize,
12955        row_bytes: usize,
12956        scale_cols: usize,
12957        mcols: usize,
12958    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12959        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12960        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
12961        let name = match mcols {
12962            2 => "qmatvec_e4m3_blk_mmvq_b2",
12963            4 => "qmatvec_e4m3_blk_mmvq_b4",
12964            8 => "qmatvec_e4m3_blk_mmvq_b8",
12965            16 => "qmatvec_e4m3_blk_mmvq_b16",
12966            _ => {
12967                return Err(
12968                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
12969                );
12970            }
12971        };
12972        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12973        let f = self.func(name);
12974        let cfg = LaunchConfig {
12975            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
12976            block_dim: (32, ROWS_PER_BLOCK, 1),
12977            shared_mem_bytes: 0,
12978        };
12979        let (inf, outf, mi, rb, sc) = (
12980            in_f as i32,
12981            out_f as i32,
12982            m as i32,
12983            row_bytes as i64,
12984            scale_cols as i32,
12985        );
12986        let __s_b = self.gpu.stream();
12987        let mut b = __s_b.launch_builder(&f);
12988        b.arg(bytes)
12989            .arg(aq)
12990            .arg(ad)
12991            .arg(scales)
12992            .arg(&mut y)
12993            .arg(&inf)
12994            .arg(&outf)
12995            .arg(&mi)
12996            .arg(&rb)
12997            .arg(&sc);
12998        unsafe {
12999            b.launch(cfg)?;
13000        }
13001        Ok(y)
13002    }
13003
13004    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
13005    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
13006    #[allow(clippy::too_many_arguments)]
13007    pub fn qmatvec_e4m3_blk_batched_raw(
13008        &self,
13009        bytes: &CudaSlice<u8>,
13010        x: &CudaSlice<f32>,
13011        scales: &CudaSlice<f32>,
13012        m: usize,
13013        in_f: usize,
13014        out_f: usize,
13015        row_bytes: usize,
13016        scale_cols: usize,
13017        mcols: usize,
13018    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13019        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13020        self.qmatvec_e4m3_blk_mmvq_batched(
13021            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13022        )
13023    }
13024
13025    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
13026    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
13027    #[allow(clippy::too_many_arguments)]
13028    pub fn qmatvec_e4m3_blk_mmvq_raw(
13029        &self,
13030        bytes: &CudaSlice<u8>,
13031        x: &CudaSlice<f32>,
13032        scales: &CudaSlice<f32>,
13033        m: usize,
13034        in_f: usize,
13035        out_f: usize,
13036        row_bytes: usize,
13037        scale_cols: usize,
13038    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13039        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13040        self.qmatvec_e4m3_blk_mmvq(
13041            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13042        )
13043    }
13044
13045    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
13046    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
13047    #[allow(clippy::too_many_arguments)]
13048    pub fn qmatvec_e4m3_fused2_raw(
13049        &self,
13050        b0: &CudaSlice<u8>,
13051        b1: &CudaSlice<u8>,
13052        x: &CudaSlice<f32>,
13053        in_f: usize,
13054        out0: usize,
13055        out1: usize,
13056        row_bytes: usize,
13057        ws0: f32,
13058        ws1: f32,
13059    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13060        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13061        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13062    }
13063
13064    #[allow(clippy::too_many_arguments)]
13065    pub fn qmatvec_e4m3_fused3_raw(
13066        &self,
13067        b0: &CudaSlice<u8>,
13068        b1: &CudaSlice<u8>,
13069        b2: &CudaSlice<u8>,
13070        x: &CudaSlice<f32>,
13071        in_f: usize,
13072        out0: usize,
13073        out1: usize,
13074        out2: usize,
13075        row_bytes: usize,
13076        ws0: f32,
13077        ws1: f32,
13078        ws2: f32,
13079    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13080        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13081        self.e4m3_fused3_core(
13082            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13083        )
13084    }
13085
13086    #[allow(clippy::too_many_arguments)]
13087    pub fn qmatvec_e4m3_fused2_t_raw(
13088        &self,
13089        b0: &CudaSlice<u8>,
13090        b1: &CudaSlice<u8>,
13091        x: &CudaSlice<f32>,
13092        m: usize,
13093        in_f: usize,
13094        out0: usize,
13095        out1: usize,
13096        row_bytes: usize,
13097        ws0: f32,
13098        ws1: f32,
13099    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13100        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13101        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13102    }
13103
13104    #[allow(clippy::too_many_arguments)]
13105    pub fn qmatvec_e4m3_fused3_t_raw(
13106        &self,
13107        b0: &CudaSlice<u8>,
13108        b1: &CudaSlice<u8>,
13109        b2: &CudaSlice<u8>,
13110        x: &CudaSlice<f32>,
13111        m: usize,
13112        in_f: usize,
13113        out0: usize,
13114        out1: usize,
13115        out2: usize,
13116        row_bytes: usize,
13117        ws0: f32,
13118        ws1: f32,
13119        ws2: f32,
13120    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13121        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13122        self.e4m3_fused3_t_core(
13123            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13124        )
13125    }
13126
13127    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
13128    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
13129    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
13130    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
13131    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
13132    ///
13133    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
13134    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
13135    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
13136    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
13137    fn try_e4m3_blk_pre(
13138        &self,
13139        w: &crate::model::GpuTensor,
13140        aq: &CudaSlice<i8>,
13141        ad: &CudaSlice<f32>,
13142        m: usize,
13143    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13144        use crate::model::GpuTensor;
13145        if let GpuTensor::Quant {
13146            bytes,
13147            qtype,
13148            row_bytes,
13149            blk: Some(g),
13150            ..
13151        } = w
13152        {
13153            if *qtype == QT_F8_E4M3_BLK {
13154                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
13155                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
13156                // below, so the decode-exactness contract is preserved at every width. Gated by
13157                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
13158                // one rollback door covers every dtype's batched tier.
13159                if (2..=16).contains(&m)
13160                    && std::env::var("MEMRA_NO_BATCHED").is_err()
13161                    && (m <= 4 || Self::b8_enabled())
13162                {
13163                    let mcols = Self::batched_mcols(m);
13164                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13165                        bytes,
13166                        aq,
13167                        ad,
13168                        &g.scales,
13169                        m,
13170                        w.in_features(),
13171                        w.out_features(),
13172                        *row_bytes,
13173                        g.cols,
13174                        mcols,
13175                    )?));
13176                }
13177                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13178                    bytes,
13179                    aq,
13180                    ad,
13181                    &g.scales,
13182                    m,
13183                    w.in_features(),
13184                    w.out_features(),
13185                    *row_bytes,
13186                    g.cols,
13187                )?));
13188            }
13189        }
13190        Ok(None)
13191    }
13192
13193    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
13194    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
13195    ///
13196    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
13197    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
13198    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
13199    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
13200    /// prefill keeps the floor's arithmetic and the floor's kernels.
13201    ///
13202    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
13203    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
13204    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
13205    /// (projection, prefill call) and frees immediately.
13206    ///
13207    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
13208    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
13209    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
13210    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
13211    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
13212    /// single-variable comparison instead of a two-variable one.
13213    ///
13214    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
13215    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
13216    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
13217    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
13218    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
13219    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
13220    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
13221    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
13222    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
13223    ///
13224    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
13225    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
13226    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
13227    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
13228    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
13229    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
13230    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
13231    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
13232    /// because v2's denominator had its slab already resident while this class's floor must build it
13233    /// every call; same tile, opposite sign, because the question changed.
13234    ///
13235    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
13236    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
13237    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
13238    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
13239    fn try_e4m3_blk_prefill(
13240        &self,
13241        w: &crate::model::GpuTensor,
13242        x: &CudaSlice<f32>,
13243        m: usize,
13244    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13245        use crate::model::GpuTensor;
13246        let GpuTensor::Quant {
13247            bytes,
13248            qtype,
13249            blk: Some(g),
13250            ..
13251        } = w
13252        else {
13253            return Ok(None);
13254        };
13255        if *qtype != QT_F8_E4M3_BLK {
13256            return Ok(None);
13257        }
13258        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
13259        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
13260        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
13261        // through to the dequant below when they do, never silently produce nothing.
13262        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13263            return Ok(Some(y));
13264        }
13265        let (in_f, out_f) = (w.in_features(), w.out_features());
13266        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13267        let tmp = GpuTensor::Quant {
13268            bytes: slab,
13269            qtype: QT_Q8_0,
13270            row_bytes: in_f / 32 * 34,
13271            ne: vec![in_f as u64, out_f as u64],
13272            scale: 1.0,
13273            rp: false,
13274            #[cfg(memra_cutlass)]
13275            cutlass: None,
13276            fp8: None,
13277            blk: None,
13278            f16: None,
13279            rp4: None,
13280        };
13281        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
13282        Ok(Some(self.matmul(&tmp, x, m)?))
13283    }
13284
13285    pub fn matmul_pre_noscale(
13286        &self,
13287        w: &crate::model::GpuTensor,
13288        aq: &CudaSlice<i8>,
13289        ad: &CudaSlice<f32>,
13290        m: usize,
13291    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13292        use crate::model::GpuTensor;
13293        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
13294        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
13295        // rather than let the tail below refuse and cost the caller a re-dispatch.
13296        if m == 1 {
13297            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13298                return Ok(Some((y, 1.0)));
13299            }
13300        }
13301        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
13302        if m != 1 || !self.uses_q8_1_fast(w) {
13303            return Ok(None);
13304        }
13305        let in_f = w.in_features();
13306        let out_f = w.out_features();
13307        let (bytes, qtype, row_bytes, scale, rp) = match w {
13308            GpuTensor::Quant {
13309                bytes,
13310                qtype,
13311                row_bytes,
13312                scale,
13313                rp,
13314                ..
13315            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13316            _ => return Ok(None),
13317        };
13318        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
13319        if self.mmvq_supports(qtype) {
13320            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
13321            let (mbytes, mrp) = match w {
13322                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13323                _ => (bytes, rp),
13324            };
13325            let y = self.qmatvec_mmvq(
13326                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
13327            )?;
13328            return Ok(Some((y, scale)));
13329        }
13330        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
13331        let name = match qtype {
13332            QT_Q8_0 => "qmatvec_q8_0_dp4a",
13333            QT_Q4_K => "qmatvec_q4_K_dp4a",
13334            QT_Q6_K => "qmatvec_q6_K_dp4a",
13335            QT_Q5_K => "qmatvec_q5_K_dp4a",
13336            QT_Q3_K => "qmatvec_q3_K_dp4a",
13337            QT_NVFP4 => {
13338                if rp {
13339                    "qmatvec_nvfp4_dp4a_rp"
13340                } else {
13341                    "qmatvec_nvfp4_dp4a"
13342                }
13343            }
13344            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13345            _ => return Ok(None),
13346        };
13347        let f = self.func(name);
13348        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13349        let cfg = LaunchConfig {
13350            grid_dim: (out_f as u32, m as u32, 1),
13351            block_dim: (128, 1, 1),
13352            shared_mem_bytes: 0,
13353        };
13354        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13355        let __s_b = self.gpu.stream();
13356        let mut b = __s_b.launch_builder(&f);
13357        b.arg(bytes)
13358            .arg(aq)
13359            .arg(ad)
13360            .arg(&mut y)
13361            .arg(&inf)
13362            .arg(&outf)
13363            .arg(&mi)
13364            .arg(&rb);
13365        unsafe {
13366            b.launch(cfg)?;
13367        }
13368        Ok(Some((y, scale)))
13369    }
13370
13371    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
13372    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
13373    pub fn mmvq_supports(&self, qtype: i32) -> bool {
13374        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
13375        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
13376        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
13377        // is a pure function of the dtype — the decode-parity law holds under every env.
13378        if qtype == QT_F8_E4M3 {
13379            return true;
13380        }
13381        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13382            return false;
13383        }
13384        matches!(
13385            qtype,
13386            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13387        )
13388    }
13389
13390    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
13391    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
13392    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
13393    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
13394    pub fn qmatvec_mmvq(
13395        &self,
13396        bytes: &CudaSlice<u8>,
13397        aq: &CudaSlice<i8>,
13398        ad: &CudaSlice<f32>,
13399        m: usize,
13400        in_f: usize,
13401        out_f: usize,
13402        qtype: i32,
13403        row_bytes: usize,
13404        scale: f32,
13405        rp: bool,
13406    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13407        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
13408        self.qmatvec_mmvq_into(
13409            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13410        )?;
13411        Ok(y)
13412    }
13413
13414    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
13415    #[allow(clippy::too_many_arguments)]
13416    pub fn qmatvec_mmvq_into(
13417        &self,
13418        bytes: &CudaSlice<u8>,
13419        aq: &CudaSlice<i8>,
13420        ad: &CudaSlice<f32>,
13421        m: usize,
13422        in_f: usize,
13423        out_f: usize,
13424        qtype: i32,
13425        row_bytes: usize,
13426        scale: f32,
13427        rp: bool,
13428        y: &mut CudaSlice<f32>,
13429    ) -> Result<(), Box<dyn std::error::Error>> {
13430        debug_assert!(y.len() >= m * out_f);
13431        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13432        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
13433        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
13434        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
13435        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
13436        if qtype == QT_Q8_0
13437            && rp
13438            && m == 1
13439            && out_f >= 64
13440            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13441            && {
13442                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13443                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13444            }
13445        {
13446            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13447            let cfg = LaunchConfig {
13448                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13449                block_dim: (32, 2, 1),
13450                shared_mem_bytes: 0,
13451            };
13452            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13453            let __s_b = self.gpu.stream();
13454            let mut b = __s_b.launch_builder(&f);
13455            b.arg(bytes)
13456                .arg(aq)
13457                .arg(ad)
13458                .arg(&mut *y)
13459                .arg(&inf)
13460                .arg(&outf)
13461                .arg(&mi)
13462                .arg(&rb);
13463            unsafe {
13464                b.launch(cfg)?;
13465            }
13466            if scale != 1.0 {
13467                self.scale_inplace(y, scale, out_f)?;
13468            }
13469            return Ok(());
13470        }
13471        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
13472        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
13473        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
13474        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
13475        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
13476        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
13477        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
13478        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
13479        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13480            2
13481        } else {
13482            1
13483        };
13484        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
13485        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
13486        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
13487        // valid-window interleaved, bit-identical per row — same dot program).
13488        if m == 1 && qtype == QT_Q4_0 {
13489            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13490            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
13491            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
13492            mr = *Q40MR.get_or_init(|| {
13493                std::env::var("MEMRA_Q40_MR")
13494                    .ok()
13495                    .and_then(|v| v.parse().ok())
13496                    .unwrap_or(1)
13497            });
13498        }
13499        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
13500        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
13501        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
13502        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
13503        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
13504        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
13505        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
13506        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
13507        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
13508        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
13509        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13510        let q5_force = q5_mode.as_deref() == Some("2");
13511        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
13512        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
13513        let q5_il = qtype == QT_Q5_K
13514            && m == 1
13515            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13516        if q5_il && !q5_force && out_f > 65536 {
13517            mr = 1;
13518        }
13519        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
13520        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
13521        if qtype == QT_Q4_0 && rp && mr != 1 {
13522            mr = 2;
13523        }
13524        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
13525        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
13526        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
13527        if qtype == QT_Q8_0 && rp {
13528            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13529            mr = *Q80MR.get_or_init(|| {
13530                std::env::var("MEMRA_Q80_MR")
13531                    .ok()
13532                    .and_then(|v| v.parse().ok())
13533                    .unwrap_or(1)
13534            });
13535        }
13536        let name = match (qtype, mr, rp) {
13537            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13538            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13539            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13540            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13541            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13542            (QT_Q5_K, 2, _) => {
13543                if q5_il {
13544                    "qmatvec_q5_K_mmvq_mr2_il"
13545                } else {
13546                    "qmatvec_q5_K_mmvq_mr2"
13547                }
13548            }
13549            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13550            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
13551            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
13552            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
13553            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
13554            (QT_Q8_0, _, true)
13555                if in_f % 1024 == 0 && {
13556                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13557                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13558                } =>
13559            {
13560                "qmatvec_q8_0_mmvq_rpca"
13561            }
13562            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13563            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13564            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
13565            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
13566            // reach a GGUF-layout kernel or vice versa.
13567            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13568            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13569            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13570            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13571            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13572            (QT_Q5_K, _, _) => {
13573                if q5_il {
13574                    "qmatvec_q5_K_mmvq_il"
13575                } else {
13576                    "qmatvec_q5_K_mmvq"
13577                }
13578            }
13579            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13580            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13581            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13582            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13583        };
13584        let f = self.func(name);
13585        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
13586        let rows_per_block = ROWS_PER_BLOCK * mr;
13587        let cfg = LaunchConfig {
13588            grid_dim: (
13589                (out_f as u32 + rows_per_block - 1) / rows_per_block,
13590                m as u32,
13591                1,
13592            ),
13593            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
13594            shared_mem_bytes: 0,                // warp-only reduce at m=1
13595        };
13596        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13597        let __s_b = self.gpu.stream();
13598        let mut b = __s_b.launch_builder(&f);
13599        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
13600        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
13601        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
13602        // weight_scale). Other mmvq kernels keep the 8-arg signature.
13603        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13604            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
13605            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
13606            if Self::pdl_on()
13607                && Self::pdl_mmvq_on()
13608                && Self::pdl_nvfp4q8_on()
13609                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13610            {
13611                use cudarc::driver::{DevicePtr, DevicePtrMut};
13612                let s = &self.gpu.stream();
13613                let (pw, _g0) = bytes.device_ptr(s);
13614                let (paq, _g1) = aq.device_ptr(s);
13615                let (pad, _g2) = ad.device_ptr(s);
13616                let (py, _g3) = y.device_ptr_mut(s);
13617                let mut ps = [
13618                    &pw as *const _ as *mut std::ffi::c_void,
13619                    &paq as *const _ as *mut _,
13620                    &pad as *const _ as *mut _,
13621                    &py as *const _ as *mut _,
13622                    &inf as *const _ as *mut _,
13623                    &outf as *const _ as *mut _,
13624                    &mi as *const _ as *mut _,
13625                    &rb as *const _ as *mut _,
13626                    &scale as *const _ as *mut _,
13627                ];
13628                unsafe {
13629                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13630                }
13631                return Ok(());
13632            }
13633            b.arg(bytes)
13634                .arg(aq)
13635                .arg(ad)
13636                .arg(&mut *y)
13637                .arg(&inf)
13638                .arg(&outf)
13639                .arg(&mi)
13640                .arg(&rb)
13641                .arg(&scale);
13642            unsafe {
13643                b.launch(cfg)?;
13644            }
13645        } else if Self::pdl_on()
13646            && Self::pdl_mmvq_on()
13647            && (matches!(
13648                name,
13649                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
13650            ) || (Self::pdl_nvfp4q8_on()
13651                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
13652        {
13653            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
13654            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
13655            // names may take this launch (unmarked kernels would read unordered).
13656            {
13657                use cudarc::driver::{DevicePtr, DevicePtrMut};
13658                let s = &self.gpu.stream();
13659                let (pw, _g0) = bytes.device_ptr(s);
13660                let (paq, _g1) = aq.device_ptr(s);
13661                let (pad, _g2) = ad.device_ptr(s);
13662                let (py, _g3) = y.device_ptr_mut(s);
13663                let mut ps = [
13664                    &pw as *const _ as *mut std::ffi::c_void,
13665                    &paq as *const _ as *mut _,
13666                    &pad as *const _ as *mut _,
13667                    &py as *const _ as *mut _,
13668                    &inf as *const _ as *mut _,
13669                    &outf as *const _ as *mut _,
13670                    &mi as *const _ as *mut _,
13671                    &rb as *const _ as *mut _,
13672                ];
13673                unsafe {
13674                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13675                }
13676            }
13677            if scale != 1.0 {
13678                self.scale_inplace(y, scale, m * out_f)?;
13679            }
13680        } else {
13681            b.arg(bytes)
13682                .arg(aq)
13683                .arg(ad)
13684                .arg(&mut *y)
13685                .arg(&inf)
13686                .arg(&outf)
13687                .arg(&mi)
13688                .arg(&rb);
13689            unsafe {
13690                b.launch(cfg)?;
13691            }
13692            if scale != 1.0 {
13693                self.scale_inplace(y, scale, m * out_f)?;
13694            }
13695        }
13696        Ok(())
13697    }
13698
13699    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
13700    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
13701    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
13702    pub fn qmatvec_mmvq_raw(
13703        &self,
13704        bytes: &CudaSlice<u8>,
13705        x: &CudaSlice<f32>,
13706        m: usize,
13707        in_f: usize,
13708        out_f: usize,
13709        qtype: i32,
13710        row_bytes: usize,
13711        rp: bool,
13712    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13713        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13714        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
13715    }
13716
13717    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
13718    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
13719    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
13720    pub fn batched_supports(&self, qtype: i32) -> bool {
13721        matches!(
13722            qtype,
13723            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
13724        )
13725    }
13726
13727    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
13728    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
13729    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
13730    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
13731    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
13732    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
13733    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
13734    pub fn iq_fast_enabled() -> bool {
13735        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13736        *ON.get_or_init(|| {
13737            std::env::var("MEMRA_IQ_FAST")
13738                .map(|v| v != "0")
13739                .unwrap_or(true)
13740        })
13741    }
13742
13743    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
13744    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
13745    pub fn b8_enabled() -> bool {
13746        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13747        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
13748    }
13749
13750    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
13751    pub fn batched_mcols(m: usize) -> usize {
13752        if m == 2 {
13753            2
13754        } else if m <= 4 {
13755            4
13756        } else if m <= 8 {
13757            8
13758        } else {
13759            16
13760        }
13761    }
13762
13763    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
13764    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
13765    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
13766    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
13767    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
13768        Some(match (qtype, mcols) {
13769            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
13770            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
13771            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
13772            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
13773            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
13774            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
13775            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
13776            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
13777            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
13778            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
13779            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
13780            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
13781            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
13782            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
13783            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
13784            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
13785            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
13786            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
13787            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
13788            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
13789            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
13790            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
13791            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
13792            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
13793            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
13794            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
13795            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
13796            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
13797            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
13798            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
13799            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
13800            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
13801            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
13802            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
13803            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
13804            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
13805            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
13806            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
13807            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
13808            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
13809            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
13810            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
13811            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
13812            _ => return None,
13813        })
13814    }
13815
13816    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
13817    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
13818    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
13819    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
13820    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
13821    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
13822    ///
13823    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
13824    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
13825    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
13826    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
13827    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
13828    /// msweep on all six 27B shapes (2026-07-03):
13829    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
13830    ///          it applies for b4 (-3..-14%), never loses;
13831    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
13832    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
13833    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
13834    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
13835    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
13836    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
13837    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
13838    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
13839    /// b2: in_f>=6144 -> r2, else base.
13840    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
13841    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
13842    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
13843    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
13844    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
13845    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
13846    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
13847    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
13848    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
13849    /// Device SM count (cached) — grid-fill policy input.
13850    pub fn sm_count(&self) -> i32 {
13851        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
13852        *SMS.get_or_init(|| {
13853            use cudarc::driver::sys::CUdevice_attribute_enum as A;
13854            self.gpu
13855                .ctx
13856                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
13857                .unwrap_or(82)
13858        })
13859    }
13860
13861    pub fn batched_variant(
13862        &self,
13863        _m: usize,
13864        in_f: usize,
13865        out_f: usize,
13866        qtype: i32,
13867        row_bytes: usize,
13868        mcols: usize,
13869        rp: bool,
13870    ) -> &'static str {
13871        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
13872        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
13873        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
13874        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
13875        if qtype == QT_Q8_0 {
13876            return if rp { "rp" } else { "base" };
13877        }
13878        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13879        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
13880            Ok("base") => "base",
13881            Ok("pf") => "pf",
13882            Ok("r2") => "r2",
13883            Ok("r2w8") => "r2w8",
13884            Ok("pfr2") => "pfr2",
13885            Ok("ca") => "ca",
13886            Ok("car2") => "car2",
13887            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
13888            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
13889            Ok("rp") => "rp",
13890            Ok("rpr2") => "rpr2",
13891            Ok("rpr2w8") => "rpr2w8",
13892            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
13893            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
13894            Ok("rpca") => "rpca",
13895            Ok("rpcar2") => "rpcar2",
13896            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
13897            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
13898            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
13899            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
13900            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
13901            // bit-identical to the decode path — measurement corpus ONLY, never auto).
13902            Ok("rpsc") => "rpsc",
13903            Ok("rpms") => "rpms",
13904            Ok("rpmsc") => "rpmsc",
13905            Ok("rpks") => "rpks",
13906            Ok("rpksc") => "rpksc",
13907            _ => "auto",
13908        });
13909        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
13910        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
13911        // shapes qualify; anything else falls back to the register variants.
13912        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
13913        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
13914        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
13915        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
13916        // forced MEMRA_MMVQ_BV values still work).
13917        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13918        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
13919        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
13920        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
13921        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
13922        let sms = *SMS.get_or_init(|| {
13923            use cudarc::driver::sys::CUdevice_attribute_enum as A;
13924            self.gpu
13925                .ctx
13926                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
13927                .unwrap_or(82)
13928        });
13929        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
13930        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
13931        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
13932        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
13933        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
13934        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
13935        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
13936        // AUTO RULE = the measured winners table (differs from NVFP4's!):
13937        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
13938        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
13939        //     r2 1258us) — kernels kept behind the force seam for the corpus;
13940        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
13941        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
13942        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
13943        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
13944        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
13945        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
13946        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
13947        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
13948        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
13949        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
13950        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
13951        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13952        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
13953            Ok("base") => "base",
13954            Ok("r2") => "r2",
13955            Ok("r2w8") => "r2w8",
13956            _ => "auto",
13957        });
13958        let variant: &'static str = if qtype == QT_Q4_0 {
13959            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
13960            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
13961            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
13962            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13963            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
13964                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
13965                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
13966                // + syncs cost more than the stalls, bank-pad made no difference);
13967                // register load-ahead flat (nvcc already reorders). The b-tier limiter
13968                // is still unidentified — see the jsonl row.
13969                Ok("base") => "base",
13970                Ok("r2") => "r2",
13971                Ok("ms") => "ms",
13972                Ok("sm") => "sm",
13973                Ok("la") => "la",
13974                _ => "auto",
13975            });
13976            let v = if q40 != "auto" {
13977                q40
13978            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
13979                "r2"
13980            } else {
13981                "base"
13982            };
13983            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
13984            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
13985            // and the limiter is the per-column activation load chain (long_scoreboard
13986            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
13987            if rp {
13988                match v {
13989                    "ms" => "r2ms_rp",
13990                    "sm" => "r2sm_rp",
13991                    "la" => "r2la_rp",
13992                    "r2" => "r2_rp",
13993                    _ => "rp",
13994                }
13995            } else if matches!(v, "ms" | "sm" | "la") {
13996                "r2"
13997            } else {
13998                v
13999            }
14000        } else if qtype != QT_NVFP4 && !kq_r2 {
14001            "base"
14002        } else if kq_r2 && rp {
14003            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
14004            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
14005            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
14006            "rp"
14007        } else if kq_r2 {
14008            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
14009            // mcols != 4 forced r2w8 falls to unbounded r2.
14010            if kq_bv != "auto" {
14011                if kq_bv == "r2w8" && mcols != 4 {
14012                    "r2"
14013                } else {
14014                    kq_bv
14015                }
14016            } else if bv != "auto" {
14017                match bv {
14018                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14019                    "r2w8" | "rpr2w8" => {
14020                        if mcols != 4 {
14021                            "r2"
14022                        } else {
14023                            "r2w8"
14024                        }
14025                    }
14026                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
14027                }
14028            } else {
14029                let blocks = (out_f + 7) / 8;
14030                let waves = blocks as f64 / (7 * sms as usize) as f64;
14031                let filled = blocks >= 4 * sms as usize;
14032                let use_r2 = if qtype == QT_Q4_K {
14033                    filled
14034                } else {
14035                    waves >= 2.0
14036                };
14037                if use_r2 { "r2" } else { "base" }
14038            }
14039        } else if bv != "auto" {
14040            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
14041            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
14042            // unsupported (shape, mcols) combos fall back to pf/r2.
14043            // On rp buffers, forced legacy names map to their rp twins (layout law).
14044            let v = if bv == "r2w8" && mcols == 2 {
14045                "r2"
14046            } else if bv == "ca" && (!ca_ok || mcols == 8) {
14047                "pf"
14048            } else if bv == "car2" && (!ca_ok || mcols == 8) {
14049                "r2"
14050            } else if bv == "pfr2" && mcols == 8 {
14051                "r2"
14052            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14053                "rpr2"
14054            }
14055            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
14056            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14057                if mcols == 8 { "rpr2w8" } else { "rpr2" }
14058            } else if bv == "rpcar2" && mcols == 2 {
14059                "rpca"
14060            }
14061            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
14062            // (rpms has no smem and no alignment need — always valid on rp buffers).
14063            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14064                "rpr2"
14065            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14066                "rpr2"
14067            } else {
14068                bv
14069            };
14070            if rp {
14071                match v {
14072                    "base" | "pf" | "ca" | "rp" => "rp",
14073                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14074                    "r2w8" | "rpr2w8" => {
14075                        if mcols == 2 {
14076                            "rpr2"
14077                        } else {
14078                            "rpr2w8"
14079                        }
14080                    }
14081                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
14082                }
14083            } else {
14084                v
14085            }
14086        } else if mcols == 8 {
14087            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
14088            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
14089            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
14090            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
14091            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
14092            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
14093            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
14094            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
14095            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
14096            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
14097            if rp {
14098                if sc_ok { "rpsc" } else { "rpr2w8" }
14099            } else {
14100                "r2w8"
14101            }
14102        } else if mcols >= 4 {
14103            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
14104            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
14105            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
14106            let blocks = (out_f + 7) / 8;
14107            let r7 = 7 * sms as usize;
14108            let r8 = 8 * sms as usize;
14109            let waves = blocks as f64 / r7 as f64;
14110            let filled = blocks >= 4 * sms as usize;
14111            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
14112            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
14113            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
14114            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14115                // the extra residency drops the INTEGER wave count -> the straggler wave a
14116                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
14117                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
14118                if rp { "rpr2w8" } else { "r2w8" }
14119            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14120                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
14121                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
14122                if rp { "rpr2" } else { "r2" }
14123            } else {
14124                // fractional straggler-wave window with no crossing, or grid too small to fill
14125                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
14126                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
14127                if rp { "rp" } else { "pf" }
14128            }
14129        } else if in_f >= 6144 {
14130            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
14131            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
14132            // stays.
14133            if rp { "rpr2" } else { "r2" }
14134        } else if rp {
14135            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
14136            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
14137            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
14138            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
14139            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14140            if sc_ok && waves >= 0.9 && waves <= 1.1 {
14141                "rpsc"
14142            } else {
14143                "rp"
14144            }
14145        } else {
14146            "base"
14147        };
14148        variant
14149    }
14150
14151    pub fn qmatvec_mmvq_batched(
14152        &self,
14153        bytes: &CudaSlice<u8>,
14154        aq: &CudaSlice<i8>,
14155        ad: &CudaSlice<f32>,
14156        m: usize,
14157        in_f: usize,
14158        out_f: usize,
14159        qtype: i32,
14160        row_bytes: usize,
14161        mcols: usize,
14162        scale: f32,
14163        rp: bool,
14164    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14165        const ROWS_PER_BLOCK: u32 = 4;
14166        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
14167        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
14168        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
14169        // weight keeps its rp-layout kernel family regardless of the override.
14170        let forced: Option<&'static str> = {
14171            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14172            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14173                .as_deref()
14174                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14175        };
14176        let variant = match forced {
14177            Some(v) if !rp || v.contains("rp") => v,
14178            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14179        };
14180        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14181            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14182        })?;
14183        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
14184        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
14185        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
14186        let variant = if mcols == 16 {
14187            if rp { "rp" } else { "base" }
14188        } else {
14189            variant
14190        };
14191        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
14192        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
14193        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
14194        // per-(token,row) chain (columns c >= m never execute in either form) ->
14195        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
14196        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
14197        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14198        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14199        if b567
14200            && qtype == QT_NVFP4
14201            && rp
14202            && mcols == 8
14203            && (5..=7).contains(&m)
14204            && matches!(variant, "rpsc" | "rpr2w8")
14205        {
14206            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14207            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
14208            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14209            let cfg = LaunchConfig {
14210                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14211                block_dim: (32, ROWS_PER_BLOCK, 1),
14212                shared_mem_bytes: 0,
14213            };
14214            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14215            let __s_b = self.gpu.stream();
14216            let mut b = __s_b.launch_builder(&f);
14217            b.arg(bytes)
14218                .arg(aq)
14219                .arg(ad)
14220                .arg(&mut y)
14221                .arg(&inf)
14222                .arg(&outf)
14223                .arg(&mi)
14224                .arg(&rb);
14225            unsafe {
14226                b.launch(cfg)?;
14227            }
14228            if scale != 1.0 {
14229                self.scale_inplace(&mut y, scale, m * out_f)?;
14230            }
14231            return Ok(y);
14232        }
14233        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14234            "base" => (base_name.into(), ROWS_PER_BLOCK),
14235            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14236            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14237            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14238            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
14239            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
14240            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
14241            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14242            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14243            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14244            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14245            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14246            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14247            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14248            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
14249        };
14250        debug_assert!(
14251            !rp || name.contains("_rp"),
14252            "rp weight dispatched to a GGUF-layout kernel"
14253        );
14254        let f = self.func(&name);
14255        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14256        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
14257        let smem = if name.contains("_r2sm_rp") {
14258            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14259        } else {
14260            0
14261        };
14262        let cfg = LaunchConfig {
14263            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14264            block_dim: (32, ROWS_PER_BLOCK, 1),
14265            shared_mem_bytes: smem,
14266        };
14267        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14268        let __s_b = self.gpu.stream();
14269        let mut b = __s_b.launch_builder(&f);
14270        b.arg(bytes)
14271            .arg(aq)
14272            .arg(ad)
14273            .arg(&mut y)
14274            .arg(&inf)
14275            .arg(&outf)
14276            .arg(&mi)
14277            .arg(&rb);
14278        unsafe {
14279            b.launch(cfg)?;
14280        }
14281        if scale != 1.0 {
14282            self.scale_inplace(&mut y, scale, m * out_f)?;
14283        }
14284        Ok(y)
14285    }
14286
14287    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
14288    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
14289    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
14290    pub fn qmatvec_batched_raw(
14291        &self,
14292        bytes: &CudaSlice<u8>,
14293        x: &CudaSlice<f32>,
14294        m: usize,
14295        in_f: usize,
14296        out_f: usize,
14297        qtype: i32,
14298        row_bytes: usize,
14299        mcols: usize,
14300        rp: bool,
14301    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14302        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14303        self.qmatvec_mmvq_batched(
14304            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14305        )
14306    }
14307
14308    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
14309    pub fn qmatvec_nvfp4_batched_raw(
14310        &self,
14311        bytes: &CudaSlice<u8>,
14312        x: &CudaSlice<f32>,
14313        m: usize,
14314        in_f: usize,
14315        out_f: usize,
14316        row_bytes: usize,
14317        mcols: usize,
14318        rp: bool,
14319    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14320        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14321    }
14322
14323    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
14324    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
14325    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
14326    fn try_fp4_gemm(
14327        &self,
14328        w: &crate::model::GpuTensor,
14329        x: &CudaSlice<f32>,
14330        m: usize,
14331        in_f: usize,
14332        out_f: usize,
14333    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14334        use crate::model::GpuTensor;
14335        if cfg!(memra_portable_cuda) {
14336            return Ok(None);
14337        }
14338        if std::env::var("MEMRA_FP4").is_err() {
14339            return Ok(None);
14340        }
14341        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
14342        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
14343        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
14344        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
14345        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
14346        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
14347        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
14348        // for the common no-macro-scale case.
14349        #[cfg(memra_cutlass)]
14350        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14351            if let GpuTensor::Quant {
14352                bytes,
14353                qtype,
14354                scale,
14355                row_bytes,
14356                cutlass,
14357                ..
14358            } = w
14359            {
14360                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14361                    if let Some(cw) = cutlass {
14362                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
14363                        let y = self.cutlass_fp4_gemm(
14364                            &cw.b_packed,
14365                            &cw.sfb_swizzled,
14366                            x,
14367                            *scale,
14368                            m,
14369                            out_f,
14370                            in_f,
14371                        )?;
14372                        return Ok(Some(y));
14373                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14374                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
14375                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
14376                        // (the load-time repack ~doubles it) — needed for models that don't fit the
14377                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
14378                        let (b_packed, sfb_sw) =
14379                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14380                        let y =
14381                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14382                        return Ok(Some(y));
14383                    }
14384                }
14385            }
14386        }
14387        if let GpuTensor::Quant {
14388            bytes,
14389            qtype,
14390            row_bytes,
14391            scale,
14392            rp,
14393            ..
14394        } = w
14395        {
14396            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
14397            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
14398            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14399                let y =
14400                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14401                return Ok(Some(y));
14402            }
14403        }
14404        Ok(None)
14405    }
14406
14407    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
14408    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
14409    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
14410    pub fn rms_norm_f16out(
14411        &self,
14412        x: &CudaSlice<f32>,
14413        w: &CudaSlice<f32>,
14414        dst: &mut CudaSlice<f32>,
14415        dst16: &mut CudaSlice<u8>,
14416        ncols: usize,
14417        nrows: usize,
14418        eps: f32,
14419    ) -> Result<(), Box<dyn std::error::Error>> {
14420        let f = self.func("rms_norm_f16out_f32");
14421        let cfg = LaunchConfig {
14422            grid_dim: (nrows as u32, 1, 1),
14423            block_dim: (rms_block(), 1, 1),
14424            shared_mem_bytes: 0,
14425        };
14426        let (nc, e) = (ncols as i32, eps);
14427        let __s_b = self.gpu.stream();
14428        let mut b = __s_b.launch_builder(&f);
14429        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14430        unsafe {
14431            b.launch(cfg)?;
14432        }
14433        Ok(())
14434    }
14435
14436    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
14437    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
14438    #[allow(clippy::too_many_arguments)]
14439    pub fn add_rms_norm_f16out(
14440        &self,
14441        a: &CudaSlice<f32>,
14442        b: &CudaSlice<f32>,
14443        w: &CudaSlice<f32>,
14444        res: &mut CudaSlice<f32>,
14445        dst: &mut CudaSlice<f32>,
14446        dst16: &mut CudaSlice<u8>,
14447        ncols: usize,
14448        nrows: usize,
14449        eps: f32,
14450    ) -> Result<(), Box<dyn std::error::Error>> {
14451        let f = self.func("add_rms_norm_f16out_f32");
14452        let cfg = LaunchConfig {
14453            grid_dim: (nrows as u32, 1, 1),
14454            block_dim: (rms_block(), 1, 1),
14455            shared_mem_bytes: 0,
14456        };
14457        let (nc, e) = (ncols as i32, eps);
14458        let __s_lb = self.gpu.stream();
14459        let mut lb = __s_lb.launch_builder(&f);
14460        lb.arg(a)
14461            .arg(b)
14462            .arg(w)
14463            .arg(res)
14464            .arg(dst)
14465            .arg(dst16)
14466            .arg(&nc)
14467            .arg(&e);
14468        unsafe {
14469            lb.launch(cfg)?;
14470        }
14471        Ok(())
14472    }
14473
14474    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
14475    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
14476    pub fn matmul_group_xh(
14477        &self,
14478        ws: &[&crate::model::GpuTensor],
14479        x: &CudaSlice<f32>,
14480        xh: &CudaSlice<u8>,
14481        m: usize,
14482    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14483        let mut out = Vec::with_capacity(ws.len());
14484        let in_f = ws[0].in_features();
14485        for w in ws {
14486            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14487                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14488                    out.push(y);
14489                    continue;
14490                }
14491            }
14492            out.push(self.matmul(w, x, m)?);
14493        }
14494        Ok(out)
14495    }
14496
14497    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
14498    /// GDN steps). Layouts [T, H].
14499    pub fn gdn_pad_mask(
14500        &self,
14501        beta: &mut CudaSlice<f32>,
14502        g_log: &mut CudaSlice<f32>,
14503        len_d: &CudaSlice<i32>,
14504        h: usize,
14505        t: usize,
14506    ) -> Result<(), Box<dyn std::error::Error>> {
14507        let f = self.func("gdn_pad_mask_f32");
14508        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14509        let (hi, ti) = (h as i32, t as i32);
14510        let __s_b = self.gpu.stream();
14511        let mut b = __s_b.launch_builder(&f);
14512        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14513        unsafe {
14514            b.launch(cfg)?;
14515        }
14516        Ok(())
14517    }
14518
14519    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
14520    /// gather for the padded prime graph's h_seed/hlast.
14521    pub fn row_gather_dev(
14522        &self,
14523        src: &CudaSlice<f32>,
14524        dst: &mut CudaSlice<f32>,
14525        len_d: &CudaSlice<i32>,
14526        ncols: usize,
14527    ) -> Result<(), Box<dyn std::error::Error>> {
14528        let f = self.func("row_gather_dev_f32");
14529        let cfg = LaunchConfig::for_num_elems(ncols as u32);
14530        let nc = ncols as i32;
14531        let __s_b = self.gpu.stream();
14532        let mut b = __s_b.launch_builder(&f);
14533        b.arg(src).arg(dst).arg(len_d).arg(&nc);
14534        unsafe {
14535            b.launch(cfg)?;
14536        }
14537        Ok(())
14538    }
14539
14540    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
14541    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
14542    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
14543    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
14544    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
14545    /// different in_f) falls back to its own `matmul` — behavior unchanged.
14546    pub fn matmul_group(
14547        &self,
14548        ws: &[&crate::model::GpuTensor],
14549        x: &CudaSlice<f32>,
14550        m: usize,
14551    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14552        use crate::model::GpuTensor;
14553        let mut out = Vec::with_capacity(ws.len());
14554        let any_mirror = ws
14555            .iter()
14556            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14557        if m >= 16 && any_mirror && !self.verify_exact_on() {
14558            let in_f = ws[0].in_features();
14559            let xh = self.f16_act(x, m * in_f, in_f)?;
14560            for w in ws {
14561                if w.in_features() == in_f {
14562                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14563                        out.push(y);
14564                        continue;
14565                    }
14566                }
14567                out.push(self.matmul(w, x, m)?);
14568            }
14569            return Ok(out);
14570        }
14571        for w in ws {
14572            out.push(self.matmul(w, x, m)?);
14573        }
14574        Ok(out)
14575    }
14576
14577    /// Cross-request grouped matmul (task #13): run ONE projection group over the
14578    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
14579    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
14580    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
14581    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
14582    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
14583    pub fn matmul_group_multi(
14584        &self,
14585        ws: &[&crate::model::GpuTensor],
14586        xs: &[&CudaSlice<f32>],
14587        ms: &[usize],
14588    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14589        assert_eq!(xs.len(), ms.len());
14590        let in_f = ws[0].in_features();
14591        let total: usize = ms.iter().sum();
14592        let mut xcat = self.uninit(total * in_f)?;
14593        let mut off = 0usize;
14594        for (x, &m) in xs.iter().zip(ms) {
14595            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14596            off += m;
14597        }
14598        let ys = self.matmul_group(ws, &xcat, total)?;
14599        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14600        for (w, y) in ws.iter().zip(ys) {
14601            let out_f = w.out_features();
14602            let mut off = 0usize;
14603            for (s, &m) in ms.iter().enumerate() {
14604                let mut ys_s = self.uninit(m * out_f)?;
14605                let src = y.slice(off * out_f..(off + m) * out_f);
14606                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14607                out[s].push(ys_s);
14608                off += m;
14609            }
14610        }
14611        Ok(out)
14612    }
14613
14614    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
14615    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
14616    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
14617    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
14618    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
14619    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
14620    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
14621    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
14622    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
14623    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14624        use crate::model::GpuTensor;
14625        if !legacy_quant_gemm_allowed(
14626            cfg!(memra_portable_cuda),
14627            cfg!(memra_hopper_mma),
14628            std::env::var_os("MEMRA_NO_GEMM").is_some(),
14629        ) {
14630            return false;
14631        }
14632        match w {
14633            GpuTensor::Quant { qtype, .. } => {
14634                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14635                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14636            }
14637            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14638        }
14639    }
14640
14641    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
14642    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
14643    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
14644    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
14645    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
14646    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
14647    pub fn qmatvec_gemm(
14648        &self,
14649        w: &crate::model::GpuTensor,
14650        aq: &CudaSlice<i8>,
14651        ad: &CudaSlice<f32>,
14652        m: usize,
14653    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14654        use crate::model::GpuTensor;
14655        let in_f = w.in_features();
14656        let out_f = w.out_features();
14657        let (bytes, qtype, row_bytes, scale, rp) = match w {
14658            GpuTensor::Quant {
14659                bytes,
14660                qtype,
14661                row_bytes,
14662                scale,
14663                rp,
14664                ..
14665            } => (bytes, *qtype, *row_bytes, *scale, *rp),
14666            _ => unreachable!("gemm_supports guaranteed Quant"),
14667        };
14668        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
14669        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
14670        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
14671        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
14672        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
14673        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
14674            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
14675                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
14676                if scale != 1.0 {
14677                    self.scale_inplace(&mut y, scale, m * out_f)?;
14678                }
14679                return Ok(y);
14680            }
14681        }
14682        let name = match qtype {
14683            QT_Q8_0 => "qmatvec_gemm_q8_0",
14684            QT_Q4_K => "qmatvec_gemm_q4_K",
14685            QT_Q4_0 => {
14686                if rp {
14687                    "qmatvec_gemm_q4_0_rp"
14688                } else {
14689                    "qmatvec_gemm_q4_0"
14690                }
14691            }
14692            QT_Q5_K => "qmatvec_gemm_q5_K",
14693            QT_Q6_K => "qmatvec_gemm_q6_K",
14694            QT_NVFP4 => {
14695                if rp {
14696                    "qmatvec_gemm_nvfp4_rp"
14697                } else {
14698                    "qmatvec_gemm_nvfp4"
14699                }
14700            }
14701            _ => unreachable!(),
14702        };
14703        let f = self.func(name);
14704        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14705        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
14706        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
14707        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
14708        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14709        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
14710        let k1_tile = if is_k1 {
14711            k1_launch_override().unwrap_or((128, 128, 8))
14712        } else {
14713            (128, 128, 8)
14714        };
14715        let (bm, bn): (u32, u32) = if is_k1 {
14716            (k1_tile.0, k1_tile.1)
14717        } else {
14718            (64, 256)
14719        };
14720        let warps: u32 = if is_k1 {
14721            k1_tile.2
14722        } else {
14723            match qtype {
14724                QT_NVFP4 => 8,
14725                _ => 4,
14726            }
14727        };
14728        let cfg = LaunchConfig {
14729            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14730            block_dim: (32, warps, 1),
14731            shared_mem_bytes: 0,
14732        };
14733        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14734        let __s_b = self.gpu.stream();
14735        let mut b = __s_b.launch_builder(&f);
14736        b.arg(bytes)
14737            .arg(aq)
14738            .arg(ad)
14739            .arg(&mut y)
14740            .arg(&inf)
14741            .arg(&outf)
14742            .arg(&mi)
14743            .arg(&rb);
14744        unsafe {
14745            b.launch(cfg)?;
14746        }
14747        if scale != 1.0 {
14748            self.scale_inplace(&mut y, scale, m * out_f)?;
14749        }
14750        Ok(y)
14751    }
14752
14753    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
14754    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
14755    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
14756    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
14757    pub fn qmatvec_gemm_raw(
14758        &self,
14759        bytes: &CudaSlice<u8>,
14760        x: &CudaSlice<f32>,
14761        m: usize,
14762        in_f: usize,
14763        out_f: usize,
14764        qtype: i32,
14765        row_bytes: usize,
14766    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14767        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14768        let name = match qtype {
14769            QT_Q8_0 => "qmatvec_gemm_q8_0",
14770            QT_Q4_K => "qmatvec_gemm_q4_K",
14771            QT_Q4_0 => "qmatvec_gemm_q4_0",
14772            QT_Q5_K => "qmatvec_gemm_q5_K",
14773            QT_Q6_K => "qmatvec_gemm_q6_K",
14774            QT_NVFP4 => "qmatvec_gemm_nvfp4",
14775            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
14776            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
14777        };
14778        let f = self.func(name);
14779        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14780        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
14781        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
14782        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14783        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
14784        let k1_tile = if is_k1 {
14785            k1_launch_override().unwrap_or((128, 128, 8))
14786        } else {
14787            (128, 128, 8)
14788        };
14789        let (bm, bn): (u32, u32) = if is_k1 {
14790            (k1_tile.0, k1_tile.1)
14791        } else {
14792            (64, 256)
14793        };
14794        let warps: u32 = if is_k1 {
14795            k1_tile.2
14796        } else {
14797            match qtype {
14798                QT_NVFP4 | QT_NVFP4_RP => 8,
14799                _ => 4,
14800            }
14801        };
14802        let cfg = LaunchConfig {
14803            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14804            block_dim: (32, warps, 1),
14805            shared_mem_bytes: 0,
14806        };
14807        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14808        let __s_b = self.gpu.stream();
14809        let mut b = __s_b.launch_builder(&f);
14810        b.arg(bytes)
14811            .arg(&aq)
14812            .arg(&ad)
14813            .arg(&mut y)
14814            .arg(&inf)
14815            .arg(&outf)
14816            .arg(&mi)
14817            .arg(&rb);
14818        unsafe {
14819            b.launch(cfg)?;
14820        }
14821        Ok(y)
14822    }
14823
14824    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
14825    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
14826    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
14827    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
14828    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
14829    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
14830    pub fn qmatvec_gemm_q8_0_wgmma_raw(
14831        &self,
14832        rp4: &CudaSlice<u8>,
14833        aq: &CudaSlice<i8>,
14834        ad: &CudaSlice<f32>,
14835        m: usize,
14836        in_f: usize,
14837        out_f: usize,
14838    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14839        assert!(
14840            out_f % 64 == 0 && in_f % 32 == 0,
14841            "wgmma GEMM needs out_f%64==0, in_f%32==0"
14842        );
14843        let f = self.func("qmatvec_gemm_q8_0_wgmma");
14844        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
14845        let cfg = LaunchConfig {
14846            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
14847            block_dim: (128, 1, 1),
14848            shared_mem_bytes: 0,
14849        };
14850        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
14851        let __s_b = self.gpu.stream();
14852        let mut b = __s_b.launch_builder(&f);
14853        b.arg(rp4)
14854            .arg(aq)
14855            .arg(ad)
14856            .arg(&mut y)
14857            .arg(&inf)
14858            .arg(&outf)
14859            .arg(&mi);
14860        unsafe {
14861            b.launch(cfg)?;
14862        }
14863        Ok(y)
14864    }
14865
14866    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
14867    pub fn scale_inplace(
14868        &self,
14869        y: &mut CudaSlice<f32>,
14870        s: f32,
14871        n: usize,
14872    ) -> Result<(), Box<dyn std::error::Error>> {
14873        let f = self.func("scale_f32");
14874        let cfg = LaunchConfig::for_num_elems(n as u32);
14875        let (sf, ni) = (s, n as i32);
14876        let __s_b = self.gpu.stream();
14877        let mut b = __s_b.launch_builder(&f);
14878        b.arg(y).arg(&sf).arg(&ni);
14879        unsafe {
14880            b.launch(cfg)?;
14881        }
14882        Ok(())
14883    }
14884
14885    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
14886    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
14887    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
14888    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
14889    pub fn bf16_to_f32(
14890        &self,
14891        data: &cudarc::driver::CudaView<'_, u8>,
14892        n: usize,
14893    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14894        let mut out = self.alloc_uninit::<f32>(n)?;
14895        let f = self.func("bf16_to_f32");
14896        let cfg = LaunchConfig::for_num_elems(n as u32);
14897        let ni = n as i32;
14898        let __s_b = self.gpu.stream();
14899        let mut b = __s_b.launch_builder(&f);
14900        b.arg(data).arg(&mut out).arg(&ni);
14901        unsafe {
14902            b.launch(cfg)?;
14903        }
14904        Ok(out)
14905    }
14906
14907    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
14908    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
14909    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
14910    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
14911    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
14912    /// calls, the spec-verify contract) vs plain linear.
14913    fn linear_bf16_chunked(
14914        &self,
14915        x: &CudaSlice<f32>,
14916        data: &CudaSlice<u8>,
14917        m: usize,
14918        in_f: usize,
14919        out_f: usize,
14920        exact: bool,
14921    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14922        const CHUNK_BYTES: usize = 256 << 20;
14923        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
14924        if chunk_rows >= out_f {
14925            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
14926            return if exact {
14927                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
14928            } else {
14929                self.linear(x, &wf32, m, in_f, out_f)
14930            };
14931        }
14932        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14933        let mut r0 = 0usize;
14934        while r0 < out_f {
14935            let rows = chunk_rows.min(out_f - r0);
14936            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
14937            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
14938            let yc = if exact {
14939                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
14940            } else {
14941                self.linear(x, &wf32, m, in_f, rows)?
14942            };
14943            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
14944            for mi in 0..m {
14945                let src = yc.slice(mi * rows..(mi + 1) * rows);
14946                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
14947                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
14948            }
14949            r0 += rows;
14950        }
14951        Ok(y)
14952    }
14953
14954    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
14955    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
14956    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
14957    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
14958    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
14959    /// router/shexp sites and matmul_decode_exact's Float arm.
14960    pub fn linear_decode_exact(
14961        &self,
14962        x: &CudaSlice<f32>,
14963        w: &CudaSlice<f32>,
14964        m_tokens: usize,
14965        in_f: usize,
14966        out_f: usize,
14967    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14968        if m_tokens == 1 {
14969            return self.linear(x, w, 1, in_f, out_f);
14970        }
14971        let xv = self.view(x, m_tokens * in_f);
14972        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
14973        for t in 0..m_tokens {
14974            let row = xv.slice(t * in_f..(t + 1) * in_f);
14975            let mut xr = self.alloc_uninit::<f32>(in_f)?;
14976            self.copy_view_into(&mut xr, 0, &row, in_f)?;
14977            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
14978            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
14979        }
14980        Ok(y)
14981    }
14982
14983    pub fn linear(
14984        &self,
14985        x: &CudaSlice<f32>,
14986        w: &CudaSlice<f32>,
14987        m_tokens: usize,
14988        in_f: usize,
14989        out_f: usize,
14990    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14991        use cudarc::cublaslt::{Matmul, MatmulConfig};
14992        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
14993        let cfg = MatmulConfig {
14994            transa: true,
14995            transb: false,
14996            transc: false,
14997            m: out_f as u64,
14998            n: m_tokens as u64,
14999            k: in_f as u64,
15000            alpha: 1.0,
15001            lda: in_f as i64,
15002            ldb: in_f as i64,
15003            beta: 0.0,
15004            ldc: out_f as i64,
15005            stride_a: None,
15006            stride_b: None,
15007            stride_c: None,
15008            stride_bias: None,
15009            batch_size: None,
15010        };
15011        unsafe {
15012            self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15013        }
15014        Ok(c)
15015    }
15016
15017    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
15018    pub fn sdpa_naive(
15019        &self,
15020        q: &CudaSlice<f32>,
15021        k: &CudaSlice<f32>,
15022        v: &CudaSlice<f32>,
15023        o: &mut CudaSlice<f32>,
15024        head_dim: usize,
15025        n_head: usize,
15026        n_head_kv: usize,
15027        t: usize,
15028        t_kv: usize,
15029        scale: f32,
15030        causal: bool,
15031    ) -> Result<(), Box<dyn std::error::Error>> {
15032        let f = self.func("sdpa_naive_f32");
15033        let cfg = LaunchConfig {
15034            grid_dim: (n_head as u32, t as u32, 1),
15035            block_dim: (128, 1, 1),
15036            shared_mem_bytes: (t_kv * 4) as u32,
15037        };
15038        let (hd, nh, nhkv, ti, tkvi, cz) = (
15039            head_dim as i32,
15040            n_head as i32,
15041            n_head_kv as i32,
15042            t as i32,
15043            t_kv as i32,
15044            causal as i32,
15045        );
15046        let __s_b = self.gpu.stream();
15047        let mut b = __s_b.launch_builder(&f);
15048        b.arg(q)
15049            .arg(k)
15050            .arg(v)
15051            .arg(o)
15052            .arg(&hd)
15053            .arg(&nh)
15054            .arg(&nhkv)
15055            .arg(&ti)
15056            .arg(&tkvi)
15057            .arg(&scale)
15058            .arg(&cz);
15059        unsafe {
15060            b.launch(cfg)?;
15061        }
15062        Ok(())
15063    }
15064
15065    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
15066    /// bidirectional image islands. `span_id` labels each absolute kv position
15067    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
15068    /// reproducing the reference's non-causal image batch. window 0 = no window.
15069    #[allow(clippy::too_many_arguments)]
15070    pub fn sdpa_naive_island(
15071        &self,
15072        q: &CudaSlice<f32>,
15073        k: &CudaSlice<f32>,
15074        v: &CudaSlice<f32>,
15075        o: &mut CudaSlice<f32>,
15076        span_id: &CudaSlice<i32>,
15077        head_dim: usize,
15078        n_head: usize,
15079        n_head_kv: usize,
15080        t: usize,
15081        t_kv: usize,
15082        scale: f32,
15083        window: usize,
15084    ) -> Result<(), Box<dyn std::error::Error>> {
15085        let f = self.func("sdpa_naive_island_f32");
15086        let cfg = LaunchConfig {
15087            grid_dim: (n_head as u32, t as u32, 1),
15088            block_dim: (128, 1, 1),
15089            shared_mem_bytes: (t_kv * 4) as u32,
15090        };
15091        let (hd, nh, nhkv, ti, tkvi, wi) = (
15092            head_dim as i32,
15093            n_head as i32,
15094            n_head_kv as i32,
15095            t as i32,
15096            t_kv as i32,
15097            window as i32,
15098        );
15099        let __s_b = self.gpu.stream();
15100        let mut b = __s_b.launch_builder(&f);
15101        b.arg(q)
15102            .arg(k)
15103            .arg(v)
15104            .arg(o)
15105            .arg(span_id)
15106            .arg(&hd)
15107            .arg(&nh)
15108            .arg(&nhkv)
15109            .arg(&ti)
15110            .arg(&tkvi)
15111            .arg(&scale)
15112            .arg(&wi);
15113        unsafe {
15114            b.launch(cfg)?;
15115        }
15116        Ok(())
15117    }
15118
15119    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
15120    #[allow(clippy::too_many_arguments)]
15121    pub fn sdpa_naive_w(
15122        &self,
15123        q: &CudaSlice<f32>,
15124        k: &CudaSlice<f32>,
15125        v: &CudaSlice<f32>,
15126        o: &mut CudaSlice<f32>,
15127        head_dim: usize,
15128        n_head: usize,
15129        n_head_kv: usize,
15130        t: usize,
15131        t_kv: usize,
15132        scale: f32,
15133        causal: bool,
15134        window: usize,
15135    ) -> Result<(), Box<dyn std::error::Error>> {
15136        let f = self.func("sdpa_naive_w_f32");
15137        let cfg = LaunchConfig {
15138            grid_dim: (n_head as u32, t as u32, 1),
15139            block_dim: (128, 1, 1),
15140            shared_mem_bytes: (t_kv * 4) as u32,
15141        };
15142        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15143            head_dim as i32,
15144            n_head as i32,
15145            n_head_kv as i32,
15146            t as i32,
15147            t_kv as i32,
15148            causal as i32,
15149            window as i32,
15150        );
15151        let __s_b = self.gpu.stream();
15152        let mut b = __s_b.launch_builder(&f);
15153        b.arg(q)
15154            .arg(k)
15155            .arg(v)
15156            .arg(o)
15157            .arg(&hd)
15158            .arg(&nh)
15159            .arg(&nhkv)
15160            .arg(&ti)
15161            .arg(&tkvi)
15162            .arg(&scale)
15163            .arg(&cz)
15164            .arg(&wi);
15165        unsafe {
15166            b.launch(cfg)?;
15167        }
15168        Ok(())
15169    }
15170
15171    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
15172    pub fn sdpa_naive_view(
15173        &self,
15174        q: &CudaSlice<f32>,
15175        k: &cudarc::driver::CudaView<f32>,
15176        v: &cudarc::driver::CudaView<f32>,
15177        o: &mut CudaSlice<f32>,
15178        head_dim: usize,
15179        n_head: usize,
15180        n_head_kv: usize,
15181        t: usize,
15182        t_kv: usize,
15183        scale: f32,
15184        causal: bool,
15185    ) -> Result<(), Box<dyn std::error::Error>> {
15186        let f = self.func("sdpa_naive_f32");
15187        let cfg = LaunchConfig {
15188            grid_dim: (n_head as u32, t as u32, 1),
15189            block_dim: (128, 1, 1),
15190            shared_mem_bytes: (t_kv * 4) as u32,
15191        };
15192        let (hd, nh, nhkv, ti, tkvi, cz) = (
15193            head_dim as i32,
15194            n_head as i32,
15195            n_head_kv as i32,
15196            t as i32,
15197            t_kv as i32,
15198            causal as i32,
15199        );
15200        let __s_b = self.gpu.stream();
15201        let mut b = __s_b.launch_builder(&f);
15202        b.arg(q)
15203            .arg(k)
15204            .arg(v)
15205            .arg(o)
15206            .arg(&hd)
15207            .arg(&nh)
15208            .arg(&nhkv)
15209            .arg(&ti)
15210            .arg(&tkvi)
15211            .arg(&scale)
15212            .arg(&cz);
15213        unsafe {
15214            b.launch(cfg)?;
15215        }
15216        Ok(())
15217    }
15218
15219    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
15220    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
15221    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
15222    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
15223    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
15224    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
15225    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
15226    #[allow(clippy::too_many_arguments)]
15227    pub fn fa_dequant_kv_view_f32(
15228        &self,
15229        k: &cudarc::driver::CudaView<u8>,
15230        v: &cudarc::driver::CudaView<u8>,
15231        kf: &mut CudaSlice<f32>,
15232        vf: &mut CudaSlice<f32>,
15233        kv_dim_k: usize,
15234        kv_dim_v: usize,
15235        t_kv: usize,
15236        k_tok_bytes: usize,
15237        v_tok_bytes: usize,
15238        g: bool,
15239    ) -> Result<(), Box<dyn std::error::Error>> {
15240        let f = if g {
15241            self.func_g("fa_dequant_kv_ws_f32")
15242        } else {
15243            self.func("fa_dequant_kv_ws_f32")
15244        };
15245        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15246        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15247        let cfg = LaunchConfig {
15248            grid_dim: (nblk.max(1), 1, 1),
15249            block_dim: (256, 1, 1),
15250            shared_mem_bytes: 0,
15251        };
15252        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15253        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15254        let __s_b = self.gpu.stream();
15255        let mut b = __s_b.launch_builder(&f);
15256        b.arg(k)
15257            .arg(v)
15258            .arg(&mut *kf)
15259            .arg(&mut *vf)
15260            .arg(&kdk)
15261            .arg(&kdv)
15262            .arg(&tkvi)
15263            .arg(&ktb)
15264            .arg(&vtb);
15265        unsafe {
15266            b.launch(cfg)?;
15267        }
15268        Ok(())
15269    }
15270
15271    #[allow(clippy::too_many_arguments)]
15272    pub fn sdpa_naive_quantized_view(
15273        &self,
15274        q: &CudaSlice<f32>,
15275        k: &cudarc::driver::CudaView<u8>,
15276        v: &cudarc::driver::CudaView<u8>,
15277        o: &mut CudaSlice<f32>,
15278        head_dim: usize,
15279        n_head: usize,
15280        n_head_kv: usize,
15281        t: usize,
15282        t_kv: usize,
15283        scale: f32,
15284        causal: bool,
15285        k_tok_bytes: usize,
15286        v_tok_bytes: usize,
15287    ) -> Result<(), Box<dyn std::error::Error>> {
15288        let kv_dim = n_head_kv * head_dim;
15289        let mut kf = self.uninit(t_kv * kv_dim)?;
15290        let mut vf = self.uninit(t_kv * kv_dim)?;
15291        let f = self.func("fa_dequant_kv_ws_f32");
15292        let total = (2 * t_kv * kv_dim) as u64;
15293        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15294        let cfg = LaunchConfig {
15295            grid_dim: (nblk.max(1), 1, 1),
15296            block_dim: (256, 1, 1),
15297            shared_mem_bytes: 0,
15298        };
15299        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15300        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15301        let __s_b = self.gpu.stream();
15302        let mut b = __s_b.launch_builder(&f);
15303        b.arg(k)
15304            .arg(v)
15305            .arg(&mut kf)
15306            .arg(&mut vf)
15307            .arg(&kv_dim_i)
15308            .arg(&kv_dim_i)
15309            .arg(&t_kv_i)
15310            .arg(&k_tok_bytes_i)
15311            .arg(&v_tok_bytes_i);
15312        unsafe { b.launch(cfg)? };
15313        self.sdpa_naive(
15314            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15315        )
15316    }
15317
15318    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
15319    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
15320    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
15321    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
15322    /// unwindowed function above and produces bit-identical output at window == 0.
15323    ///
15324    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
15325    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
15326    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
15327    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
15328    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
15329    #[allow(clippy::too_many_arguments)]
15330    pub fn sdpa_naive_w_quantized_view(
15331        &self,
15332        q: &CudaSlice<f32>,
15333        k: &cudarc::driver::CudaView<u8>,
15334        v: &cudarc::driver::CudaView<u8>,
15335        o: &mut CudaSlice<f32>,
15336        head_dim: usize,
15337        n_head: usize,
15338        n_head_kv: usize,
15339        t: usize,
15340        t_kv: usize,
15341        scale: f32,
15342        causal: bool,
15343        window: usize,
15344        k_tok_bytes: usize,
15345        v_tok_bytes: usize,
15346    ) -> Result<(), Box<dyn std::error::Error>> {
15347        let kv_dim = n_head_kv * head_dim;
15348        let mut kf = self.uninit(t_kv * kv_dim)?;
15349        let mut vf = self.uninit(t_kv * kv_dim)?;
15350        let f = self.func("fa_dequant_kv_ws_f32");
15351        let total = (2 * t_kv * kv_dim) as u64;
15352        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15353        let cfg = LaunchConfig {
15354            grid_dim: (nblk.max(1), 1, 1),
15355            block_dim: (256, 1, 1),
15356            shared_mem_bytes: 0,
15357        };
15358        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15359        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15360        let __s_b = self.gpu.stream();
15361        let mut b = __s_b.launch_builder(&f);
15362        b.arg(k)
15363            .arg(v)
15364            .arg(&mut kf)
15365            .arg(&mut vf)
15366            .arg(&kv_dim_i)
15367            .arg(&kv_dim_i)
15368            .arg(&t_kv_i)
15369            .arg(&k_tok_bytes_i)
15370            .arg(&v_tok_bytes_i);
15371        unsafe { b.launch(cfg)? };
15372        self.sdpa_naive_w(
15373            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15374        )
15375    }
15376
15377    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
15378    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
15379    /// Q/K/V/O [head_dim, n_head(_kv), T].
15380    pub fn fa_prefill(
15381        &self,
15382        q: &CudaSlice<f32>,
15383        k: &CudaSlice<f32>,
15384        v: &CudaSlice<f32>,
15385        o: &mut CudaSlice<f32>,
15386        head_dim: usize,
15387        n_head: usize,
15388        n_head_kv: usize,
15389        t: usize,
15390        t_kv: usize,
15391        scale: f32,
15392        causal: bool,
15393    ) -> Result<(), Box<dyn std::error::Error>> {
15394        if portable_mma_gated() {
15395            return self.sdpa_naive(
15396                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15397            );
15398        }
15399        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
15400        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
15401        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
15402        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
15403        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
15404        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
15405        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
15406        let fa3_on = head_dim == 256
15407            && causal
15408            && t == t_kv
15409            && match std::env::var("MEMRA_FA3").as_deref() {
15410                Ok("0") => false,
15411                Ok("1") => true,
15412                _ => cfg!(memra_hopper_mma),
15413            };
15414        if fa3_on {
15415            let n = t * n_head * head_dim;
15416            let nkv = t * n_head_kv * head_dim;
15417            let mut q16 = self.alloc_u8_uninit(n * 2)?;
15418            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15419            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15420            self.f32_to_bf16_into(q, &mut q16, n)?;
15421            self.f32_to_bf16_into(k, &mut k16, nkv)?;
15422            self.f32_to_bf16_into(v, &mut v16, nkv)?;
15423            let rc = {
15424                use cudarc::driver::{DevicePtr, DevicePtrMut};
15425                let stream = self.gpu.stream();
15426                let (qp, _g1) = q16.device_ptr(&stream);
15427                let (kp, _g2) = k16.device_ptr(&stream);
15428                let (vp, _g3) = v16.device_ptr(&stream);
15429                let (op, _g4) = o.device_ptr_mut(&stream);
15430                unsafe {
15431                    memra_fa3_prefill(
15432                        qp as *const core::ffi::c_void,
15433                        kp as *const core::ffi::c_void,
15434                        vp as *const core::ffi::c_void,
15435                        op as *mut f32,
15436                        t as i32,
15437                        n_head as i32,
15438                        n_head_kv as i32,
15439                        head_dim as i32,
15440                        scale,
15441                        stream.cu_stream() as *mut core::ffi::c_void,
15442                    )
15443                }
15444            };
15445            if rc != 0 {
15446                return Err(format!("memra_fa3_prefill rc={rc}").into());
15447            }
15448            return Ok(());
15449        }
15450        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
15451        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
15452        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
15453        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
15454        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15455        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15456        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15457            const BLOCK_Q: usize = 64;
15458            const BKX: usize = 32;
15459            let f = self.func("fa_prefill_bf16_p1");
15460            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15461                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15462            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15463            f.set_attribute(
15464                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15465                shmem as i32,
15466            )?;
15467            let cfg = LaunchConfig {
15468                grid_dim: (
15469                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15470                    n_head as u32,
15471                    1,
15472                ),
15473                block_dim: (32, 4, 1),
15474                shared_mem_bytes: shmem,
15475            };
15476            let (hd, nh, nhkv, ti, tkvi, cz) = (
15477                head_dim as i32,
15478                n_head as i32,
15479                n_head_kv as i32,
15480                t as i32,
15481                t_kv as i32,
15482                causal as i32,
15483            );
15484            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15485            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15486            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15487            let __s_b = self.gpu.stream();
15488            let mut b = __s_b.launch_builder(&f);
15489            b.arg(&qb)
15490                .arg(&kb)
15491                .arg(&vb)
15492                .arg(o)
15493                .arg(&hd)
15494                .arg(&nh)
15495                .arg(&nhkv)
15496                .arg(&ti)
15497                .arg(&tkvi)
15498                .arg(&scale)
15499                .arg(&cz);
15500            unsafe {
15501                b.launch(cfg)?;
15502            }
15503            return Ok(());
15504        }
15505        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
15506        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
15507        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
15508        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
15509        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
15510        const BK: usize = 32;
15511        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
15512        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
15513        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15514        let (block_q, warps, w2_sfx): (usize, u32, &str) =
15515            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15516        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
15517        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
15518        // other head_dims to sdpa_naive before reaching here.
15519        let hd_sfx = fa_hd_suffix(head_dim)?;
15520        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15521        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
15522        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
15523        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
15524        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
15525        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15526        let (kb16, vb16) = if bf16kv {
15527            let n = t_kv * n_head_kv * head_dim;
15528            let mut kb = self.alloc_u8_uninit(n * 2)?;
15529            let mut vb = self.alloc_u8_uninit(n * 2)?;
15530            let fcv = self.func("f32_to_bf16_bulk");
15531            let ni = n as i64;
15532            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15533            let __s_b = self.gpu.stream();
15534            let mut b = __s_b.launch_builder(&fcv);
15535            b.arg(k).arg(&mut kb).arg(&ni);
15536            unsafe {
15537                b.launch(cfgc)?;
15538            }
15539            let __s_b = self.gpu.stream();
15540            let mut b = __s_b.launch_builder(&fcv);
15541            b.arg(v).arg(&mut vb).arg(&ni);
15542            unsafe {
15543                b.launch(cfgc)?;
15544            }
15545            (Some(kb), Some(vb))
15546        } else {
15547            (None, None)
15548        };
15549        let f = self.func(&if bf16kv {
15550            format!("fa_prefill_bf16kv_pp{hd_sfx}")
15551        } else {
15552            format!(
15553                "fa_prefill_f32{}{}{hd_sfx}",
15554                if floor { "" } else { "_pp" },
15555                if floor { "" } else { w2_sfx }
15556            )
15557        });
15558        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
15559        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
15560        let kv_stages = if bf16kv { 2 } else { 1 };
15561        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15562            + 4 * (block_q * BK + 2 * block_q)) as u32;
15563        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15564        f.set_attribute(
15565            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15566            shmem as i32,
15567        )?;
15568        let cfg = LaunchConfig {
15569            grid_dim: (
15570                (t as u32 + block_q as u32 - 1) / block_q as u32,
15571                n_head as u32,
15572                1,
15573            ),
15574            block_dim: (32, warps, 1),
15575            shared_mem_bytes: shmem,
15576        };
15577        let (hd, nh, nhkv, ti, tkvi, cz) = (
15578            head_dim as i32,
15579            n_head as i32,
15580            n_head_kv as i32,
15581            t as i32,
15582            t_kv as i32,
15583            causal as i32,
15584        );
15585        let __s_b = self.gpu.stream();
15586        let mut b = __s_b.launch_builder(&f);
15587        b.arg(q);
15588        match (&kb16, &vb16) {
15589            (Some(kb), Some(vb)) => {
15590                b.arg(kb).arg(vb);
15591            }
15592            _ => {
15593                b.arg(k).arg(v);
15594            }
15595        }
15596        b.arg(o)
15597            .arg(&hd)
15598            .arg(&nh)
15599            .arg(&nhkv)
15600            .arg(&ti)
15601            .arg(&tkvi)
15602            .arg(&scale)
15603            .arg(&cz);
15604        unsafe {
15605            b.launch(cfg)?;
15606        }
15607        Ok(())
15608    }
15609
15610    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
15611    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
15612    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
15613    #[allow(clippy::too_many_arguments)]
15614    pub fn fa_prefill_w(
15615        &self,
15616        q: &CudaSlice<f32>,
15617        k: &CudaSlice<f32>,
15618        v: &CudaSlice<f32>,
15619        o: &mut CudaSlice<f32>,
15620        head_dim: usize,
15621        n_head: usize,
15622        n_head_kv: usize,
15623        t: usize,
15624        t_kv: usize,
15625        scale: f32,
15626        causal: bool,
15627        window: usize,
15628    ) -> Result<(), Box<dyn std::error::Error>> {
15629        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
15630        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
15631        if portable_mma_gated() {
15632            return self.sdpa_naive_w(
15633                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15634            );
15635        }
15636        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
15637        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
15638        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
15639        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15640        let faw_f32 =
15641            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15642        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15643        self.fa_prefill_w_arm(
15644            q,
15645            k,
15646            v,
15647            o,
15648            head_dim,
15649            n_head,
15650            n_head_kv,
15651            t,
15652            t_kv,
15653            scale,
15654            causal,
15655            window,
15656            floor || faw_f32,
15657            floor,
15658        )
15659    }
15660
15661    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
15662    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
15663    #[allow(clippy::too_many_arguments)]
15664    pub fn fa_prefill_w_pre(
15665        &self,
15666        qb: &CudaSlice<u8>,
15667        kb: &CudaSlice<u8>,
15668        vb: &CudaSlice<u8>,
15669        o: &mut CudaSlice<f32>,
15670        head_dim: usize,
15671        n_head: usize,
15672        n_head_kv: usize,
15673        t: usize,
15674        t_kv: usize,
15675        scale: f32,
15676        causal: bool,
15677        window: usize,
15678        v_f16: bool,
15679    ) -> Result<(), Box<dyn std::error::Error>> {
15680        const BLOCK_Q: usize = 64;
15681        const BK: usize = 32;
15682        debug_assert_eq!(head_dim, 256);
15683        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15684        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
15685        if hp {
15686            const BLOCK_QH: usize = 32;
15687            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
15688            // else re-encode through the pooled scratch (stream-ordered reuse).
15689            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
15690            let vh: &CudaSlice<u8> = if v_f16 {
15691                vb
15692            } else {
15693                let n = t_kv * n_head_kv * head_dim;
15694                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
15695                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
15696                }
15697                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
15698                vguard.as_ref().unwrap()
15699            };
15700            let f = self.func("fa_prefill_w_bf16_p1h2");
15701            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15702            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15703            f.set_attribute(
15704                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15705                shmem as i32,
15706            )?;
15707            let cfg = LaunchConfig {
15708                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15709                block_dim: (32, 4, 1),
15710                shared_mem_bytes: shmem,
15711            };
15712            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15713                head_dim as i32,
15714                n_head as i32,
15715                n_head_kv as i32,
15716                t as i32,
15717                t_kv as i32,
15718                causal as i32,
15719                window as i32,
15720            );
15721            let __s_b = self.gpu.stream();
15722            let mut b = __s_b.launch_builder(&f);
15723            b.arg(qb)
15724                .arg(kb)
15725                .arg(vh)
15726                .arg(o)
15727                .arg(&hd)
15728                .arg(&nh)
15729                .arg(&nhkv)
15730                .arg(&ti)
15731                .arg(&tkvi)
15732                .arg(&scale)
15733                .arg(&cz)
15734                .arg(&wi);
15735            unsafe {
15736                b.launch(cfg)?;
15737            }
15738            return Ok(());
15739        }
15740        let f = self.func("fa_prefill_w_bf16_p1");
15741        let shmem =
15742            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15743        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15744        f.set_attribute(
15745            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15746            shmem as i32,
15747        )?;
15748        let cfg = LaunchConfig {
15749            grid_dim: (
15750                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15751                n_head as u32,
15752                1,
15753            ),
15754            block_dim: (32, 4, 1),
15755            shared_mem_bytes: shmem,
15756        };
15757        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15758            head_dim as i32,
15759            n_head as i32,
15760            n_head_kv as i32,
15761            t as i32,
15762            t_kv as i32,
15763            causal as i32,
15764            window as i32,
15765        );
15766        let __s_b = self.gpu.stream();
15767        let mut b = __s_b.launch_builder(&f);
15768        b.arg(qb)
15769            .arg(kb)
15770            .arg(vb)
15771            .arg(o)
15772            .arg(&hd)
15773            .arg(&nh)
15774            .arg(&nhkv)
15775            .arg(&ti)
15776            .arg(&tkvi)
15777            .arg(&scale)
15778            .arg(&cz)
15779            .arg(&wi);
15780        unsafe {
15781            b.launch(cfg)?;
15782        }
15783        Ok(())
15784    }
15785
15786    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
15787    #[allow(clippy::too_many_arguments)]
15788    pub fn fa_prefill_w_arm(
15789        &self,
15790        q: &CudaSlice<f32>,
15791        k: &CudaSlice<f32>,
15792        v: &CudaSlice<f32>,
15793        o: &mut CudaSlice<f32>,
15794        head_dim: usize,
15795        n_head: usize,
15796        n_head_kv: usize,
15797        t: usize,
15798        t_kv: usize,
15799        scale: f32,
15800        causal: bool,
15801        window: usize,
15802        f32_stage: bool,
15803        floor: bool,
15804    ) -> Result<(), Box<dyn std::error::Error>> {
15805        const BLOCK_Q: usize = 64;
15806        const BK: usize = 32;
15807        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
15808        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
15809        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
15810        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
15811        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15812        let p1 = !floor
15813            && !f32_stage
15814            && *P1_ON.get_or_init(|| {
15815                std::env::var("MEMRA_FAW_P1")
15816                    .map(|v| v != "0")
15817                    .unwrap_or(true)
15818            });
15819        let hp =
15820            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15821        if hp {
15822            const BLOCK_QH: usize = 32;
15823            let f = self.func("fa_prefill_w_bf16_p1h2");
15824            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15825            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15826            f.set_attribute(
15827                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15828                shmem as i32,
15829            )?;
15830            let cfg = LaunchConfig {
15831                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15832                block_dim: (32, 4, 1),
15833                shared_mem_bytes: shmem,
15834            };
15835            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15836                head_dim as i32,
15837                n_head as i32,
15838                n_head_kv as i32,
15839                t as i32,
15840                t_kv as i32,
15841                causal as i32,
15842                window as i32,
15843            );
15844            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15845            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15846            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
15847            let __s_b = self.gpu.stream();
15848            let mut b = __s_b.launch_builder(&f);
15849            b.arg(&qb)
15850                .arg(&kb)
15851                .arg(&vh)
15852                .arg(o)
15853                .arg(&hd)
15854                .arg(&nh)
15855                .arg(&nhkv)
15856                .arg(&ti)
15857                .arg(&tkvi)
15858                .arg(&scale)
15859                .arg(&cz)
15860                .arg(&wi);
15861            unsafe {
15862                b.launch(cfg)?;
15863            }
15864            return Ok(());
15865        }
15866        if p1 {
15867            let f = self.func("fa_prefill_w_bf16_p1");
15868            let shmem =
15869                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15870            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15871            f.set_attribute(
15872                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15873                shmem as i32,
15874            )?;
15875            let cfg = LaunchConfig {
15876                grid_dim: (
15877                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15878                    n_head as u32,
15879                    1,
15880                ),
15881                block_dim: (32, 4, 1),
15882                shared_mem_bytes: shmem,
15883            };
15884            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15885                head_dim as i32,
15886                n_head as i32,
15887                n_head_kv as i32,
15888                t as i32,
15889                t_kv as i32,
15890                causal as i32,
15891                window as i32,
15892            );
15893            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15894            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15895            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15896            let __s_b = self.gpu.stream();
15897            let mut b = __s_b.launch_builder(&f);
15898            b.arg(&qb)
15899                .arg(&kb)
15900                .arg(&vb)
15901                .arg(o)
15902                .arg(&hd)
15903                .arg(&nh)
15904                .arg(&nhkv)
15905                .arg(&ti)
15906                .arg(&tkvi)
15907                .arg(&scale)
15908                .arg(&cz)
15909                .arg(&wi);
15910            unsafe {
15911                b.launch(cfg)?;
15912            }
15913            return Ok(());
15914        }
15915        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
15916        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
15917        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15918        let g4 = !floor
15919            && !f32_stage
15920            && n_head_kv == 1
15921            && n_head % 4 == 0
15922            && *G4_ON.get_or_init(|| {
15923                std::env::var("MEMRA_FAW_G4")
15924                    .map(|v| v != "0")
15925                    .unwrap_or(true)
15926            });
15927        if g4 {
15928            const SP_M: usize = 16;
15929            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
15930            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
15931            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15932            let o2 = *O2_ON.get_or_init(|| {
15933                std::env::var("MEMRA_FAW_O2")
15934                    .map(|v| v != "0")
15935                    .unwrap_or(true)
15936            });
15937            let f = self.func(if o2 {
15938                "fa_prefill_w_bf16_g4o2"
15939            } else {
15940                "fa_prefill_w_bf16_g4"
15941            });
15942            let shmem = if o2 {
15943                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
15944            } else {
15945                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
15946                    as u32
15947            };
15948            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15949            f.set_attribute(
15950                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15951                shmem as i32,
15952            )?;
15953            let cfg = LaunchConfig {
15954                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
15955                block_dim: (32, 4, 1),
15956                shared_mem_bytes: shmem,
15957            };
15958            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15959                head_dim as i32,
15960                n_head as i32,
15961                n_head_kv as i32,
15962                t as i32,
15963                t_kv as i32,
15964                causal as i32,
15965                window as i32,
15966            );
15967            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15968            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15969            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15970            let __s_b = self.gpu.stream();
15971            let mut b = __s_b.launch_builder(&f);
15972            b.arg(&qb)
15973                .arg(&kb)
15974                .arg(&vb)
15975                .arg(o)
15976                .arg(&hd)
15977                .arg(&nh)
15978                .arg(&nhkv)
15979                .arg(&ti)
15980                .arg(&tkvi)
15981                .arg(&scale)
15982                .arg(&cz)
15983                .arg(&wi);
15984            unsafe {
15985                b.launch(cfg)?;
15986            }
15987            return Ok(());
15988        }
15989        let f = self.func(if floor {
15990            "fa_prefill_w_f32"
15991        } else if f32_stage {
15992            "fa_prefill_w_f32_pp"
15993        } else {
15994            "fa_prefill_w_bf16_pp"
15995        });
15996        let shmem =
15997            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15998        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15999        f.set_attribute(
16000            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16001            shmem as i32,
16002        )?;
16003        let cfg = LaunchConfig {
16004            grid_dim: (
16005                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16006                n_head as u32,
16007                1,
16008            ),
16009            block_dim: (32, 4, 1),
16010            shared_mem_bytes: shmem,
16011        };
16012        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16013            head_dim as i32,
16014            n_head as i32,
16015            n_head_kv as i32,
16016            t as i32,
16017            t_kv as i32,
16018            causal as i32,
16019            window as i32,
16020        );
16021        if f32_stage {
16022            let __s_b = self.gpu.stream();
16023            let mut b = __s_b.launch_builder(&f);
16024            b.arg(q)
16025                .arg(k)
16026                .arg(v)
16027                .arg(o)
16028                .arg(&hd)
16029                .arg(&nh)
16030                .arg(&nhkv)
16031                .arg(&ti)
16032                .arg(&tkvi)
16033                .arg(&scale)
16034                .arg(&cz)
16035                .arg(&wi);
16036            unsafe {
16037                b.launch(cfg)?;
16038            }
16039        } else {
16040            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16041            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16042            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16043            let __s_b = self.gpu.stream();
16044            let mut b = __s_b.launch_builder(&f);
16045            b.arg(&qb)
16046                .arg(&kb)
16047                .arg(&vb)
16048                .arg(o)
16049                .arg(&hd)
16050                .arg(&nh)
16051                .arg(&nhkv)
16052                .arg(&ti)
16053                .arg(&tkvi)
16054                .arg(&scale)
16055                .arg(&cz)
16056                .arg(&wi);
16057            unsafe {
16058                b.launch(cfg)?;
16059            }
16060        }
16061        Ok(())
16062    }
16063
16064    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
16065    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
16066    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
16067    #[allow(clippy::too_many_arguments)]
16068    pub fn fa_prefill_hd512(
16069        &self,
16070        q: &CudaSlice<f32>,
16071        k: &CudaSlice<f32>,
16072        v: &CudaSlice<f32>,
16073        o: &mut CudaSlice<f32>,
16074        head_dim: usize,
16075        n_head: usize,
16076        n_head_kv: usize,
16077        t: usize,
16078        t_kv: usize,
16079        scale: f32,
16080        causal: bool,
16081    ) -> Result<(), Box<dyn std::error::Error>> {
16082        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
16083        if portable_mma_gated() {
16084            return self.sdpa_naive(
16085                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16086            );
16087        }
16088        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
16089        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
16090        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
16091        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
16092        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
16093        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16094        let f32_stage =
16095            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16096        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
16097        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
16098        // Own numeric config (partial-sum order) — battery-gated.
16099        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16100        let sp = !f32_stage
16101            && *SP_ON.get_or_init(|| {
16102                std::env::var("MEMRA_FA512_SP")
16103                    .map(|v| v != "0")
16104                    .unwrap_or(true)
16105            });
16106        self.fa_prefill_hd512_arm(
16107            q,
16108            k,
16109            v,
16110            o,
16111            head_dim,
16112            n_head,
16113            n_head_kv,
16114            t,
16115            t_kv,
16116            scale,
16117            causal,
16118            f32_stage,
16119            sp,
16120            sp && fa_f16pv_on(),
16121        )
16122    }
16123
16124    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
16125    #[allow(clippy::too_many_arguments)]
16126    pub fn fa_prefill_hd512_pre(
16127        &self,
16128        qb: &CudaSlice<u8>,
16129        kb: &CudaSlice<u8>,
16130        vb: &CudaSlice<u8>,
16131        o: &mut CudaSlice<f32>,
16132        head_dim: usize,
16133        n_head: usize,
16134        n_head_kv: usize,
16135        t: usize,
16136        t_kv: usize,
16137        scale: f32,
16138        causal: bool,
16139        v_f16: bool,
16140    ) -> Result<(), Box<dyn std::error::Error>> {
16141        debug_assert_eq!(head_dim, 512);
16142        const SP_M: usize = 16;
16143        const BKS: usize = 32;
16144        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
16145        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
16146        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
16147        let f16pv = fa_f16pv_on();
16148        let nw = if f16pv { fa512_wide_warps() } else { 2 };
16149        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16150        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16151        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16152        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16153            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
16154            let n = t_kv * n_head_kv * head_dim;
16155            let need = n * 2;
16156            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16157                *vguard = Some(self.alloc_uninit::<u8>(need)?);
16158            }
16159            let dst = vguard.as_mut().unwrap();
16160            self.bf16_to_f16_into(vb, n, dst)?;
16161            vguard.as_ref().unwrap()
16162        } else {
16163            vb
16164        };
16165        let f = self.func(if hp {
16166            "fa_prefill_bf16_hd512_sp16h2"
16167        } else {
16168            match (f16pv, nw) {
16169                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16170                (true, _) => "fa_prefill_bf16_hd512_sp16",
16171                _ => "fa_prefill_bf16_hd512_sp",
16172            }
16173        });
16174        let (nwarp, npart) = if hp {
16175            (4usize, 4usize)
16176        } else if nw > 2 {
16177            (nw, nw)
16178        } else {
16179            (2, 1)
16180        };
16181        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
16182        let shmem = if hp {
16183            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16184                as u32
16185        } else {
16186            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16187                + 4 * (npart * SP_M * BKS + SP_M)) as u32
16188        };
16189        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16190        f.set_attribute(
16191            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16192            shmem as i32,
16193        )?;
16194        let grid_y = if hp {
16195            (n_head / 2) as u32
16196        } else {
16197            n_head as u32
16198        };
16199        let cfg = LaunchConfig {
16200            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16201            block_dim: (32, nwarp as u32, 1),
16202            shared_mem_bytes: shmem,
16203        };
16204        let (hd, nh, nhkv, ti, tkvi, cz) = (
16205            head_dim as i32,
16206            n_head as i32,
16207            n_head_kv as i32,
16208            t as i32,
16209            t_kv as i32,
16210            causal as i32,
16211        );
16212        let __s_b = self.gpu.stream();
16213        let mut b = __s_b.launch_builder(&f);
16214        b.arg(qb)
16215            .arg(kb)
16216            .arg(vref)
16217            .arg(o)
16218            .arg(&hd)
16219            .arg(&nh)
16220            .arg(&nhkv)
16221            .arg(&ti)
16222            .arg(&tkvi)
16223            .arg(&scale)
16224            .arg(&cz);
16225        unsafe {
16226            b.launch(cfg)?;
16227        }
16228        Ok(())
16229    }
16230
16231    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
16232    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
16233    #[allow(clippy::too_many_arguments)]
16234    pub fn fa_prefill_hd512_arm(
16235        &self,
16236        q: &CudaSlice<f32>,
16237        k: &CudaSlice<f32>,
16238        v: &CudaSlice<f32>,
16239        o: &mut CudaSlice<f32>,
16240        head_dim: usize,
16241        n_head: usize,
16242        n_head_kv: usize,
16243        t: usize,
16244        t_kv: usize,
16245        scale: f32,
16246        causal: bool,
16247        f32_stage: bool,
16248        sp: bool,
16249        f16pv: bool,
16250    ) -> Result<(), Box<dyn std::error::Error>> {
16251        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16252        if sp && !f32_stage {
16253            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
16254            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
16255            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
16256            const SP_M: usize = 16;
16257            const BKS: usize = 32;
16258            let nw = if f16pv { fa512_wide_warps() } else { 2 };
16259            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16260            let f = self.func(if hp {
16261                "fa_prefill_bf16_hd512_sp16h2"
16262            } else {
16263                match (f16pv, nw) {
16264                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16265                    (true, _) => "fa_prefill_bf16_hd512_sp16",
16266                    _ => "fa_prefill_bf16_hd512_sp",
16267                }
16268            });
16269            let (nwarp, npart) = if hp {
16270                (4usize, 4usize)
16271            } else if nw > 2 {
16272                (nw, nw)
16273            } else {
16274                (2, 1)
16275            };
16276            let shmem = if hp {
16277                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16278                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16279            } else {
16280                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16281                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
16282            };
16283            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16284            f.set_attribute(
16285                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16286                shmem as i32,
16287            )?;
16288            let grid_y = if hp {
16289                (n_head / 2) as u32
16290            } else {
16291                n_head as u32
16292            };
16293            let cfg = LaunchConfig {
16294                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16295                block_dim: (32, nwarp as u32, 1),
16296                shared_mem_bytes: shmem,
16297            };
16298            let (hd, nh, nhkv, ti, tkvi, cz) = (
16299                head_dim as i32,
16300                n_head as i32,
16301                n_head_kv as i32,
16302                t as i32,
16303                t_kv as i32,
16304                causal as i32,
16305            );
16306            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16307            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16308            let vb = if f16pv {
16309                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16310            } else {
16311                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16312            };
16313            let __s_b = self.gpu.stream();
16314            let mut b = __s_b.launch_builder(&f);
16315            b.arg(&qb)
16316                .arg(&kb)
16317                .arg(&vb)
16318                .arg(o)
16319                .arg(&hd)
16320                .arg(&nh)
16321                .arg(&nhkv)
16322                .arg(&ti)
16323                .arg(&tkvi)
16324                .arg(&scale)
16325                .arg(&cz);
16326            unsafe {
16327                b.launch(cfg)?;
16328            }
16329            return Ok(());
16330        }
16331        const BLOCK_Q: usize = 32;
16332        const BK: usize = 32;
16333        const HALF: usize = 256;
16334        let f = self.func(if f32_stage {
16335            "fa_prefill_f32_hd512"
16336        } else {
16337            "fa_prefill_bf16_hd512"
16338        });
16339        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
16340        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16341            + 4 * BLOCK_Q) as u32;
16342        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16343        f.set_attribute(
16344            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16345            shmem as i32,
16346        )?;
16347        let cfg = LaunchConfig {
16348            grid_dim: (
16349                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16350                n_head as u32,
16351                2,
16352            ),
16353            block_dim: (32, 2, 1),
16354            shared_mem_bytes: shmem,
16355        };
16356        let (hd, nh, nhkv, ti, tkvi, cz) = (
16357            head_dim as i32,
16358            n_head as i32,
16359            n_head_kv as i32,
16360            t as i32,
16361            t_kv as i32,
16362            causal as i32,
16363        );
16364        if f32_stage {
16365            let __s_b = self.gpu.stream();
16366            let mut b = __s_b.launch_builder(&f);
16367            b.arg(q)
16368                .arg(k)
16369                .arg(v)
16370                .arg(o)
16371                .arg(&hd)
16372                .arg(&nh)
16373                .arg(&nhkv)
16374                .arg(&ti)
16375                .arg(&tkvi)
16376                .arg(&scale)
16377                .arg(&cz);
16378            unsafe {
16379                b.launch(cfg)?;
16380            }
16381        } else {
16382            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16383            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16384            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16385            let __s_b = self.gpu.stream();
16386            let mut b = __s_b.launch_builder(&f);
16387            b.arg(&qb)
16388                .arg(&kb)
16389                .arg(&vb)
16390                .arg(o)
16391                .arg(&hd)
16392                .arg(&nh)
16393                .arg(&nhkv)
16394                .arg(&ti)
16395                .arg(&tkvi)
16396                .arg(&scale)
16397                .arg(&cz);
16398            unsafe {
16399                b.launch(cfg)?;
16400            }
16401        }
16402        Ok(())
16403    }
16404
16405    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
16406    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
16407    /// separate f32_to_bf16 the FA entries would run).
16408    #[allow(clippy::too_many_arguments)]
16409    pub fn rope_neox2_bf16e(
16410        &self,
16411        q: &mut CudaSlice<f32>,
16412        k: &mut CudaSlice<f32>,
16413        qb: &mut CudaSlice<u8>,
16414        kb: &mut CudaSlice<u8>,
16415        pos: &CudaSlice<i32>,
16416        head_dim: usize,
16417        n_dims: usize,
16418        nh_q: usize,
16419        nh_k: usize,
16420        n_tokens: usize,
16421        base: f32,
16422        freq_scale: f32,
16423        ff: Option<&CudaSlice<f32>>,
16424    ) -> Result<(), Box<dyn std::error::Error>> {
16425        let f = self.func("rope_neox2_bf16e_f32");
16426        let rows = ((nh_q + nh_k) * n_tokens) as u32;
16427        let cfg = LaunchConfig {
16428            grid_dim: (rows, 1, 1),
16429            block_dim: ((head_dim / 2) as u32, 1, 1),
16430            shared_mem_bytes: 0,
16431        };
16432        let theta_scale = base.powf(-2.0 / n_dims as f32);
16433        let (hd, nd, nhq, nhk, nt) = (
16434            head_dim as i32,
16435            n_dims as i32,
16436            nh_q as i32,
16437            nh_k as i32,
16438            n_tokens as i32,
16439        );
16440        let __s_b = self.gpu.stream();
16441        let mut b = __s_b.launch_builder(&f);
16442        match ff {
16443            Some(t) => {
16444                b.arg(&mut *q)
16445                    .arg(&mut *k)
16446                    .arg(&mut *qb)
16447                    .arg(&mut *kb)
16448                    .arg(pos)
16449                    .arg(&hd)
16450                    .arg(&nd)
16451                    .arg(&nhq)
16452                    .arg(&nhk)
16453                    .arg(&nt)
16454                    .arg(&theta_scale)
16455                    .arg(&freq_scale)
16456                    .arg(t);
16457                unsafe {
16458                    b.launch(cfg)?;
16459                }
16460            }
16461            None => {
16462                let null: u64 = 0;
16463                b.arg(&mut *q)
16464                    .arg(&mut *k)
16465                    .arg(&mut *qb)
16466                    .arg(&mut *kb)
16467                    .arg(pos)
16468                    .arg(&hd)
16469                    .arg(&nd)
16470                    .arg(&nhq)
16471                    .arg(&nhk)
16472                    .arg(&nt)
16473                    .arg(&theta_scale)
16474                    .arg(&freq_scale)
16475                    .arg(&null);
16476                unsafe {
16477                    b.launch(cfg)?;
16478                }
16479            }
16480        }
16481        Ok(())
16482    }
16483
16484    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
16485    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
16486    pub fn f32_to_bf16(
16487        &self,
16488        x: &CudaSlice<f32>,
16489        n: usize,
16490    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16491        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16492        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16493        let f = self.func("f32_to_bf16_flat");
16494        let n_i = n as i64;
16495        let cfg = LaunchConfig {
16496            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16497            block_dim: (256, 1, 1),
16498            shared_mem_bytes: 0,
16499        };
16500        let __s_b = self.gpu.stream();
16501        let mut b = __s_b.launch_builder(&f);
16502        b.arg(x).arg(&mut y).arg(&n_i);
16503        unsafe {
16504            b.launch(cfg)?;
16505        }
16506        Ok(y)
16507    }
16508
16509    pub fn f32_to_f16(
16510        &self,
16511        x: &CudaSlice<f32>,
16512        n: usize,
16513    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16514        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16515        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16516        let f = self.func("f32_to_f16_flat");
16517        let n_i = n as i64;
16518        let cfg = LaunchConfig {
16519            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16520            block_dim: (256, 1, 1),
16521            shared_mem_bytes: 0,
16522        };
16523        let __s_b = self.gpu.stream();
16524        let mut b = __s_b.launch_builder(&f);
16525        b.arg(x).arg(&mut y).arg(&n_i);
16526        unsafe {
16527            b.launch(cfg)?;
16528        }
16529        Ok(y)
16530    }
16531
16532    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
16533    pub fn bf16_to_f16(
16534        &self,
16535        xb: &CudaSlice<u8>,
16536        n: usize,
16537    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16538        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16539        self.bf16_to_f16_into(xb, n, &mut y)?;
16540        Ok(y)
16541    }
16542
16543    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
16544    pub fn bf16_to_f16_into(
16545        &self,
16546        xb: &CudaSlice<u8>,
16547        n: usize,
16548        y: &mut CudaSlice<u8>,
16549    ) -> Result<(), Box<dyn std::error::Error>> {
16550        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16551        assert!(y.len() >= n * 2);
16552        let f = self.func("bf16_to_f16_flat");
16553        let n2 = (n / 2) as i64;
16554        let cfg = LaunchConfig {
16555            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16556            block_dim: (256, 1, 1),
16557            shared_mem_bytes: 0,
16558        };
16559        let __s_b = self.gpu.stream();
16560        let mut b = __s_b.launch_builder(&f);
16561        b.arg(xb).arg(y).arg(&n2);
16562        unsafe {
16563            b.launch(cfg)?;
16564        }
16565        Ok(())
16566    }
16567
16568    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
16569    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
16570    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
16571    /// head_dim in {256, 128}, bf16kv lane on.
16572    #[allow(clippy::too_many_arguments)]
16573    pub fn fa_prefill_vl8(
16574        &self,
16575        seqs: &[FaSeqVl],
16576        head_dim: usize,
16577        n_head: usize,
16578        n_head_kv: usize,
16579        scale: f32,
16580    ) -> Result<(), Box<dyn std::error::Error>> {
16581        const BK: usize = 32;
16582        let b = seqs.len();
16583        assert!(b >= 1 && b <= 8);
16584        let mut packed = [FaSeqVl::default(); 8];
16585        packed[..b].copy_from_slice(seqs);
16586        let v = FaVl8(packed);
16587        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16588        let ept = (n_head_kv * head_dim) as i32;
16589        {
16590            let f = self.func("fa_mirror_vl");
16591            let max_n = (max_t as i64) * ept as i64;
16592            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16593            for which in 0..2i32 {
16594                let cfg = LaunchConfig {
16595                    grid_dim: (blocks, 1, b as u32),
16596                    block_dim: (256, 1, 1),
16597                    shared_mem_bytes: 0,
16598                };
16599                let __s_lb = self.gpu.stream();
16600                let mut lb = __s_lb.launch_builder(&f);
16601                lb.arg(&v).arg(&ept).arg(&which);
16602                unsafe {
16603                    lb.launch(cfg)?;
16604                }
16605            }
16606        }
16607        let hd_sfx = fa_hd_suffix(head_dim)?;
16608        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16609        let block_q = 64usize;
16610        let kv_stages = 2usize;
16611        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16612            + 4 * (block_q * BK + 2 * block_q)) as u32;
16613        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16614        f.set_attribute(
16615            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16616            shmem as i32,
16617        )?;
16618        let cfg = LaunchConfig {
16619            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16620            block_dim: (32, 4, 1),
16621            shared_mem_bytes: shmem,
16622        };
16623        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16624        let __s_lb = self.gpu.stream();
16625        let mut lb = __s_lb.launch_builder(&f);
16626        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16627        unsafe {
16628            lb.launch(cfg)?;
16629        }
16630        Ok(())
16631    }
16632
16633    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
16634    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
16635    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
16636    #[allow(clippy::too_many_arguments)]
16637    pub fn attn_pre_vl8(
16638        &self,
16639        seqs: &[AttnPreVl],
16640        wq: &CudaSlice<f32>,
16641        wk: &CudaSlice<f32>,
16642        head_dim: usize,
16643        rope_dims: usize,
16644        n_head: usize,
16645        n_head_kv: usize,
16646        eps: f32,
16647        freq_base: f32,
16648        freq_scale: f32,
16649        kv_dim_k: usize,
16650        kv_dim_v: usize,
16651        k_tok_bytes: usize,
16652        v_tok_bytes: usize,
16653    ) -> Result<(), Box<dyn std::error::Error>> {
16654        let b = seqs.len();
16655        assert!(b >= 1 && b <= 8);
16656        let mut packed = [AttnPreVl::default(); 8];
16657        packed[..b].copy_from_slice(seqs);
16658        let v = AttnPreVl8(packed);
16659        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16660        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16661        {
16662            let f = self.func("q_gate_split_vl");
16663            let n = max_t * (n_head * head_dim) as u32;
16664            let cfg = LaunchConfig {
16665                grid_dim: (n.div_ceil(256), 1, b as u32),
16666                block_dim: (256, 1, 1),
16667                shared_mem_bytes: 0,
16668            };
16669            let __s_lb = self.gpu.stream();
16670            let mut lb = __s_lb.launch_builder(&f);
16671            lb.arg(&v).arg(&hd).arg(&nh);
16672            unsafe {
16673                lb.launch(cfg)?;
16674            }
16675        }
16676        {
16677            let f = self.func("attn_rms_vl");
16678            let cfg = LaunchConfig {
16679                grid_dim: (max_t * n_head as u32, 2, b as u32),
16680                block_dim: (rms_block(), 1, 1),
16681                shared_mem_bytes: 0,
16682            };
16683            let __s_lb = self.gpu.stream();
16684            let mut lb = __s_lb.launch_builder(&f);
16685            lb.arg(&v)
16686                .arg(wq)
16687                .arg(wk)
16688                .arg(&hd)
16689                .arg(&nh)
16690                .arg(&nhkv)
16691                .arg(&eps);
16692            unsafe {
16693                lb.launch(cfg)?;
16694            }
16695        }
16696        {
16697            let f = self.func("attn_rope_vl");
16698            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
16699            let nd = rope_dims as i32;
16700            let cfg = LaunchConfig {
16701                grid_dim: (max_t * n_head as u32, 2, b as u32),
16702                block_dim: ((head_dim / 2) as u32, 1, 1),
16703                shared_mem_bytes: 0,
16704            };
16705            let __s_lb = self.gpu.stream();
16706            let mut lb = __s_lb.launch_builder(&f);
16707            lb.arg(&v)
16708                .arg(&hd)
16709                .arg(&nd)
16710                .arg(&nh)
16711                .arg(&nhkv)
16712                .arg(&theta_scale)
16713                .arg(&freq_scale);
16714            unsafe {
16715                lb.launch(cfg)?;
16716            }
16717        }
16718        {
16719            let f = self.func("append_kv_vl");
16720            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
16721            let cfg = LaunchConfig {
16722                grid_dim: (nblk, max_t, b as u32),
16723                block_dim: (32, 1, 1),
16724                shared_mem_bytes: 0,
16725            };
16726            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16727            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16728            let __s_lb = self.gpu.stream();
16729            let mut lb = __s_lb.launch_builder(&f);
16730            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
16731            unsafe {
16732                lb.launch(cfg)?;
16733            }
16734        }
16735        Ok(())
16736    }
16737
16738    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
16739    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
16740    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
16741    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
16742    pub fn fa_prefill_view(
16743        &self,
16744        q: &CudaSlice<f32>,
16745        k: &cudarc::driver::CudaView<u8>,
16746        v: &cudarc::driver::CudaView<u8>,
16747        o: &mut CudaSlice<f32>,
16748        head_dim: usize,
16749        n_head: usize,
16750        n_head_kv: usize,
16751        t: usize,
16752        t_kv: usize,
16753        scale: f32,
16754        causal: bool,
16755        k_tok_bytes: usize,
16756        v_tok_bytes: usize,
16757        g: bool,
16758    ) -> Result<(), Box<dyn std::error::Error>> {
16759        if portable_mma_gated() {
16760            return self.sdpa_naive_quantized_view(
16761                q,
16762                k,
16763                v,
16764                o,
16765                head_dim,
16766                n_head,
16767                n_head_kv,
16768                t,
16769                t_kv,
16770                scale,
16771                causal,
16772                k_tok_bytes,
16773                v_tok_bytes,
16774            );
16775        }
16776        const BLOCK_Q: usize = 64;
16777        const BK: usize = 32;
16778        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
16779        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
16780        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
16781        let f = if g {
16782            self.func_g(&name)
16783        } else {
16784            self.func(&name)
16785        };
16786        let shmem =
16787            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16788        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16789        f.set_attribute(
16790            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16791            shmem as i32,
16792        )?;
16793        let cfg = LaunchConfig {
16794            grid_dim: (
16795                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16796                n_head as u32,
16797                1,
16798            ),
16799            block_dim: (32, 4, 1),
16800            shared_mem_bytes: shmem,
16801        };
16802        let (hd, nh, nhkv, ti, tkvi, cz) = (
16803            head_dim as i32,
16804            n_head as i32,
16805            n_head_kv as i32,
16806            t as i32,
16807            t_kv as i32,
16808            causal as i32,
16809        );
16810        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16811        let __s_b = self.gpu.stream();
16812        let mut b = __s_b.launch_builder(&f);
16813        b.arg(q)
16814            .arg(k)
16815            .arg(v)
16816            .arg(o)
16817            .arg(&hd)
16818            .arg(&nh)
16819            .arg(&nhkv)
16820            .arg(&ti)
16821            .arg(&tkvi)
16822            .arg(&scale)
16823            .arg(&cz)
16824            .arg(&ktb)
16825            .arg(&vtb);
16826        unsafe {
16827            b.launch(cfg)?;
16828        }
16829        Ok(())
16830    }
16831
16832    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
16833    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
16834    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
16835    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
16836    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
16837    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
16838    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
16839    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
16840    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
16841    #[allow(clippy::too_many_arguments)]
16842    pub fn fa_prefill_view_ws(
16843        &self,
16844        q: &CudaSlice<f32>,
16845        k: &cudarc::driver::CudaView<u8>,
16846        v: &cudarc::driver::CudaView<u8>,
16847        o: &mut CudaSlice<f32>,
16848        head_dim: usize,
16849        n_head: usize,
16850        n_head_kv: usize,
16851        t: usize,
16852        t_kv: usize,
16853        scale: f32,
16854        causal: bool,
16855        k_tok_bytes: usize,
16856        v_tok_bytes: usize,
16857        g: bool,
16858    ) -> Result<(), Box<dyn std::error::Error>> {
16859        if portable_mma_gated() {
16860            return self.sdpa_naive_quantized_view(
16861                q,
16862                k,
16863                v,
16864                o,
16865                head_dim,
16866                n_head,
16867                n_head_kv,
16868                t,
16869                t_kv,
16870                scale,
16871                causal,
16872                k_tok_bytes,
16873                v_tok_bytes,
16874            );
16875        }
16876        const BLOCK_Q: usize = 64;
16877        const BK: usize = 32;
16878        let kv_dim_k = n_head_kv * head_dim;
16879        let kv_dim_v = n_head_kv * head_dim;
16880        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
16881        let v_ws_bytes = t_kv * kv_dim_v * 2;
16882        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
16883        let mut guard = self.prime_deqw_ws.lock().unwrap();
16884        let need_grow = match guard.as_ref() {
16885            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
16886            None => true,
16887        };
16888        if need_grow {
16889            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
16890            let (ck, cv) = guard
16891                .as_ref()
16892                .map(|(a, b)| (a.len(), b.len()))
16893                .unwrap_or((0, 0));
16894            *guard = Some((
16895                self.alloc_u8(grow(ck, k_ws_bytes))?,
16896                self.alloc_u8(grow(cv, v_ws_bytes))?,
16897            ));
16898        }
16899        let (kw, vw) = guard.as_mut().unwrap();
16900        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
16901        {
16902            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
16903            let f = if g {
16904                self.func_g("fa_dequant_kv_ws_bf16")
16905            } else {
16906                self.func("fa_dequant_kv_ws_bf16")
16907            };
16908            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16909            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16910            let cfg = LaunchConfig {
16911                grid_dim: (nblk.max(1), 1, 1),
16912                block_dim: (256, 1, 1),
16913                shared_mem_bytes: 0,
16914            };
16915            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16916            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16917            let __s_b = self.gpu.stream();
16918            let mut b = __s_b.launch_builder(&f);
16919            b.arg(k)
16920                .arg(v)
16921                .arg(&mut *kw)
16922                .arg(&mut *vw)
16923                .arg(&kdk)
16924                .arg(&kdv)
16925                .arg(&tkvi)
16926                .arg(&ktb)
16927                .arg(&vtb);
16928            unsafe {
16929                b.launch(cfg)?;
16930            }
16931        }
16932        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
16933        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
16934        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
16935        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
16936        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
16937        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
16938        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
16939        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
16940            .map(|v| v != "0")
16941            .unwrap_or(true);
16942        {
16943            let hd_sfx = fa_hd_suffix(head_dim)?;
16944            let f = self.func(&format!(
16945                "fa_prefill_qw{}{hd_sfx}",
16946                if db { "_db" } else { "" }
16947            ));
16948            let shmem = if db {
16949                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
16950                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
16951            } else {
16952                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
16953            };
16954            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16955            f.set_attribute(
16956                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16957                shmem as i32,
16958            )?;
16959            let cfg = LaunchConfig {
16960                grid_dim: (
16961                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16962                    n_head as u32,
16963                    1,
16964                ),
16965                block_dim: (32, 4, 1),
16966                shared_mem_bytes: shmem,
16967            };
16968            let (hd, nh, nhkv, ti, tkvi, cz) = (
16969                head_dim as i32,
16970                n_head as i32,
16971                n_head_kv as i32,
16972                t as i32,
16973                t_kv as i32,
16974                causal as i32,
16975            );
16976            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16977            let __s_b = self.gpu.stream();
16978            let mut b = __s_b.launch_builder(&f);
16979            b.arg(q)
16980                .arg(&*kw)
16981                .arg(&*vw)
16982                .arg(o)
16983                .arg(&hd)
16984                .arg(&nh)
16985                .arg(&nhkv)
16986                .arg(&ti)
16987                .arg(&tkvi)
16988                .arg(&scale)
16989                .arg(&cz)
16990                .arg(&kdk)
16991                .arg(&kdv);
16992            unsafe {
16993                b.launch(cfg)?;
16994            }
16995        }
16996        Ok(())
16997    }
16998
16999    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
17000    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
17001    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
17002    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
17003    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
17004    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
17005    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
17006    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
17007    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
17008    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
17009    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
17010    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
17011    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
17012    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
17013    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
17014    #[allow(clippy::too_many_arguments)]
17015    pub fn fa_prefill_view_ws_w_hd128(
17016        &self,
17017        q: &CudaSlice<f32>,
17018        k: &cudarc::driver::CudaView<u8>,
17019        v: &cudarc::driver::CudaView<u8>,
17020        o: &mut CudaSlice<f32>,
17021        head_dim: usize,
17022        n_head: usize,
17023        n_head_kv: usize,
17024        t: usize,
17025        t_kv: usize,
17026        scale: f32,
17027        causal: bool,
17028        window: usize,
17029        k_tok_bytes: usize,
17030        v_tok_bytes: usize,
17031    ) -> Result<(), Box<dyn std::error::Error>> {
17032        assert_eq!(
17033            head_dim, 128,
17034            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17035        );
17036        if portable_mma_gated() {
17037            return self.sdpa_naive_w_quantized_view(
17038                q,
17039                k,
17040                v,
17041                o,
17042                head_dim,
17043                n_head,
17044                n_head_kv,
17045                t,
17046                t_kv,
17047                scale,
17048                causal,
17049                window,
17050                k_tok_bytes,
17051                v_tok_bytes,
17052            );
17053        }
17054        const BLOCK_Q: usize = 64;
17055        const BK: usize = 32;
17056        let kv_dim_k = n_head_kv * head_dim;
17057        let kv_dim_v = n_head_kv * head_dim;
17058        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17059        let v_ws_bytes = t_kv * kv_dim_v * 2;
17060        let mut guard = self.prime_deqw_ws.lock().unwrap();
17061        let need_grow = match guard.as_ref() {
17062            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17063            None => true,
17064        };
17065        if need_grow {
17066            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17067            let (ck, cv) = guard
17068                .as_ref()
17069                .map(|(a, b)| (a.len(), b.len()))
17070                .unwrap_or((0, 0));
17071            *guard = Some((
17072                self.alloc_u8(grow(ck, k_ws_bytes))?,
17073                self.alloc_u8(grow(cv, v_ws_bytes))?,
17074            ));
17075        }
17076        let (kw, vw) = guard.as_mut().unwrap();
17077        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
17078        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
17079        {
17080            let f = self.func("fa_dequant_kv_ws_bf16");
17081            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17082            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17083            let cfg = LaunchConfig {
17084                grid_dim: (nblk.max(1), 1, 1),
17085                block_dim: (256, 1, 1),
17086                shared_mem_bytes: 0,
17087            };
17088            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17089            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17090            let __s_b = self.gpu.stream();
17091            let mut b = __s_b.launch_builder(&f);
17092            b.arg(k)
17093                .arg(v)
17094                .arg(&mut *kw)
17095                .arg(&mut *vw)
17096                .arg(&kdk)
17097                .arg(&kdv)
17098                .arg(&tkvi)
17099                .arg(&ktb)
17100                .arg(&vtb);
17101            unsafe {
17102                b.launch(cfg)?;
17103            }
17104        }
17105        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
17106        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17107            .map(|v| v != "0")
17108            .unwrap_or(true);
17109        {
17110            let f = self.func(if db {
17111                "fa_prefill_qw_db_w_hd128"
17112            } else {
17113                "fa_prefill_qw_w_hd128"
17114            });
17115            let shmem = if db {
17116                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17117            } else {
17118                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17119            };
17120            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17121            f.set_attribute(
17122                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17123                shmem as i32,
17124            )?;
17125            let cfg = LaunchConfig {
17126                grid_dim: (
17127                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17128                    n_head as u32,
17129                    1,
17130                ),
17131                block_dim: (32, 4, 1),
17132                shared_mem_bytes: shmem,
17133            };
17134            let (hd, nh, nhkv, ti, tkvi, cz) = (
17135                head_dim as i32,
17136                n_head as i32,
17137                n_head_kv as i32,
17138                t as i32,
17139                t_kv as i32,
17140                causal as i32,
17141            );
17142            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17143            let __s_b = self.gpu.stream();
17144            let mut b = __s_b.launch_builder(&f);
17145            b.arg(q)
17146                .arg(&*kw)
17147                .arg(&*vw)
17148                .arg(o)
17149                .arg(&hd)
17150                .arg(&nh)
17151                .arg(&nhkv)
17152                .arg(&ti)
17153                .arg(&tkvi)
17154                .arg(&scale)
17155                .arg(&cz)
17156                .arg(&kdk)
17157                .arg(&kdv)
17158                .arg(&wnd);
17159            unsafe {
17160                b.launch(cfg)?;
17161            }
17162        }
17163        Ok(())
17164    }
17165
17166    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
17167    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
17168    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
17169    pub fn fa_decode(
17170        &self,
17171        q: &CudaSlice<f32>,
17172        k: &cudarc::driver::CudaView<u8>,
17173        v: &cudarc::driver::CudaView<u8>,
17174        o: &mut CudaSlice<f32>,
17175        head_dim: usize,
17176        n_head: usize,
17177        n_head_kv: usize,
17178        t_kv: usize,
17179        scale: f32,
17180        k_tok_bytes: usize,
17181        v_tok_bytes: usize,
17182    ) -> Result<(), Box<dyn std::error::Error>> {
17183        self.fa_decode_kvmod(
17184            q,
17185            k,
17186            v,
17187            o,
17188            head_dim,
17189            n_head,
17190            n_head_kv,
17191            t_kv,
17192            scale,
17193            k_tok_bytes,
17194            v_tok_bytes,
17195            false,
17196        )
17197    }
17198
17199    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
17200    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
17201    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
17202    #[allow(clippy::too_many_arguments)]
17203    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
17204    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
17205    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
17206    #[allow(clippy::too_many_arguments)]
17207    #[allow(clippy::too_many_arguments)]
17208    fn fa_decode_scalar_unified(
17209        &self,
17210        q: &cudarc::driver::CudaView<f32>,
17211        k: &cudarc::driver::CudaView<u8>,
17212        v: &cudarc::driver::CudaView<u8>,
17213        o: &mut cudarc::driver::CudaViewMut<f32>,
17214        head_dim: usize,
17215        n_head: usize,
17216        n_head_kv: usize,
17217        t_kv_host: usize,
17218        t_kv_dev: Option<&CudaSlice<i32>>,
17219        scale: f32,
17220        n_splits: usize,
17221        split_keys: usize,
17222        k_tok_bytes: usize,
17223        v_tok_bytes: usize,
17224        g: bool,
17225        part_o: &mut CudaSlice<f32>,
17226        part_m: &mut CudaSlice<f32>,
17227        part_l: &mut CudaSlice<f32>,
17228        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17229    ) -> Result<(), Box<dyn std::error::Error>> {
17230        let f = if g {
17231            self.func_g("fa_decode_f32")
17232        } else {
17233            self.fa_func("fa_decode_f32", head_dim)
17234        };
17235        let cfg = LaunchConfig {
17236            grid_dim: (n_head as u32, n_splits as u32, 1),
17237            block_dim: (head_dim as u32, 1, 1),
17238            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17239        };
17240        let (hd, nh, nhkv, nsp) = (
17241            head_dim as i32,
17242            n_head as i32,
17243            n_head_kv as i32,
17244            n_splits as i32,
17245        );
17246        let (ktb, vtb, tkvi, ski) = (
17247            k_tok_bytes as i64,
17248            v_tok_bytes as i64,
17249            t_kv_host as i32,
17250            split_keys as i32,
17251        );
17252        let __s_b = self.gpu.stream();
17253        let mut b = __s_b.launch_builder(&f);
17254        match t_kv_dev {
17255            Some(d) => {
17256                b.arg(q)
17257                    .arg(k)
17258                    .arg(v)
17259                    .arg(&mut *part_o)
17260                    .arg(&mut *part_m)
17261                    .arg(&mut *part_l)
17262                    .arg(&hd)
17263                    .arg(&nh)
17264                    .arg(&nhkv)
17265                    .arg(&tkvi)
17266                    .arg(d)
17267                    .arg(&scale)
17268                    .arg(&nsp)
17269                    .arg(&ski)
17270                    .arg(&ktb)
17271                    .arg(&vtb);
17272                unsafe {
17273                    b.launch(cfg)?;
17274                }
17275            }
17276            None => {
17277                let null: u64 = 0;
17278                b.arg(q)
17279                    .arg(k)
17280                    .arg(v)
17281                    .arg(&mut *part_o)
17282                    .arg(&mut *part_m)
17283                    .arg(&mut *part_l)
17284                    .arg(&hd)
17285                    .arg(&nh)
17286                    .arg(&nhkv)
17287                    .arg(&tkvi)
17288                    .arg(&null)
17289                    .arg(&scale)
17290                    .arg(&nsp)
17291                    .arg(&ski)
17292                    .arg(&ktb)
17293                    .arg(&vtb);
17294                unsafe {
17295                    b.launch(cfg)?;
17296                }
17297            }
17298        }
17299        let cfg2 = LaunchConfig {
17300            grid_dim: (n_head as u32, 1, 1),
17301            block_dim: (head_dim as u32, 1, 1),
17302            shared_mem_bytes: 0,
17303        };
17304        if let Some((oq, od)) = q8_out {
17305            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
17306            let fc = if g {
17307                self.func_g("fa_decode_combine_q8_1")
17308            } else {
17309                self.fa_func("fa_decode_combine_q8_1", head_dim)
17310            };
17311            let __s_b2 = self.gpu.stream();
17312            let mut b2 = __s_b2.launch_builder(&fc);
17313            b2.arg(&*part_o)
17314                .arg(&*part_m)
17315                .arg(&*part_l)
17316                .arg(oq)
17317                .arg(od)
17318                .arg(&hd)
17319                .arg(&nh)
17320                .arg(&nsp);
17321            unsafe {
17322                b2.launch(cfg2)?;
17323            }
17324            return Ok(());
17325        }
17326        let fc = if g {
17327            self.func_g("fa_decode_combine_f32")
17328        } else {
17329            self.fa_func("fa_decode_combine_f32", head_dim)
17330        };
17331        let __s_b2 = self.gpu.stream();
17332        let mut b2 = __s_b2.launch_builder(&fc);
17333        b2.arg(&*part_o)
17334            .arg(&*part_m)
17335            .arg(&*part_l)
17336            .arg(o)
17337            .arg(&hd)
17338            .arg(&nh)
17339            .arg(&nsp);
17340        unsafe {
17341            b2.launch(cfg2)?;
17342        }
17343        Ok(())
17344    }
17345
17346    pub fn fa_decode_kvmod(
17347        &self,
17348        q: &CudaSlice<f32>,
17349        k: &cudarc::driver::CudaView<u8>,
17350        v: &cudarc::driver::CudaView<u8>,
17351        o: &mut CudaSlice<f32>,
17352        head_dim: usize,
17353        n_head: usize,
17354        n_head_kv: usize,
17355        t_kv: usize,
17356        scale: f32,
17357        k_tok_bytes: usize,
17358        v_tok_bytes: usize,
17359        g: bool,
17360    ) -> Result<(), Box<dyn std::error::Error>> {
17361        let q_view = q.as_view();
17362        let mut o_view = o.as_view_mut();
17363        self.fa_decode_kvmod_view(
17364            &q_view,
17365            k,
17366            v,
17367            &mut o_view,
17368            head_dim,
17369            n_head,
17370            n_head_kv,
17371            t_kv,
17372            scale,
17373            k_tok_bytes,
17374            v_tok_bytes,
17375            g,
17376        )
17377    }
17378
17379    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
17380    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
17381    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
17382    /// per-session KV view and FA launch.
17383    #[allow(clippy::too_many_arguments)]
17384    pub fn fa_decode_kvmod_view(
17385        &self,
17386        q: &cudarc::driver::CudaView<f32>,
17387        k: &cudarc::driver::CudaView<u8>,
17388        v: &cudarc::driver::CudaView<u8>,
17389        o: &mut cudarc::driver::CudaViewMut<f32>,
17390        head_dim: usize,
17391        n_head: usize,
17392        n_head_kv: usize,
17393        t_kv: usize,
17394        scale: f32,
17395        k_tok_bytes: usize,
17396        v_tok_bytes: usize,
17397        g: bool,
17398    ) -> Result<(), Box<dyn std::error::Error>> {
17399        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
17400        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
17401        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
17402        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
17403        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
17404        //
17405        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
17406        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
17407        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
17408        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
17409        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
17410        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
17411        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
17412        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
17413        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
17414        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
17415        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
17416        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
17417        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
17418        // fall to the exact scalar there instead of the broken register arm.
17419        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17420        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
17421        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
17422        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
17423        if g && head_dim == 256 && !fa_v4_at(t_kv) {
17424            fa_vec = false;
17425        }
17426        let sp = fa_split_keys(t_kv, n_head_kv);
17427        let n_splits = if fa_vec {
17428            ((t_kv + sp - 1) / sp).max(1)
17429        } else {
17430            ((t_kv + 255) / 256).max(1)
17431        };
17432        let o_len = n_head * n_splits * head_dim;
17433        let ml_len = n_head * n_splits;
17434        let mut part_guard = self.fa_part_pool.lock().unwrap();
17435        if part_guard
17436            .as_ref()
17437            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17438            .unwrap_or(true)
17439        {
17440            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
17441            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
17442            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
17443            // later live allocations land at those addresses, and the next graph REPLAY writes
17444            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
17445            // output corruption began the burst after the trunk's t_kv growth first realloc'd
17446            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
17447            // the baked addresses alive (single-stream: eager writes the new buffers, replays
17448            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
17449            // (total retired < final size).
17450            let old = part_guard.take();
17451            let (co, cm) = old
17452                .as_ref()
17453                .map(|pp| (pp.0.len(), pp.1.len()))
17454                .unwrap_or((0, 0));
17455            if let Some(old) = old {
17456                self.fa_part_retired.lock().unwrap().push(old);
17457            }
17458            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17459                eprintln!(
17460                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17461                    co, o_len, cm, ml_len
17462                );
17463            }
17464            *part_guard = Some((
17465                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17466                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17467                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17468            ));
17469        }
17470        let pg = part_guard.as_mut().unwrap();
17471        self.gpu
17472            .stream()
17473            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17474        self.gpu
17475            .stream()
17476            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17477        self.gpu
17478            .stream()
17479            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17480        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17481        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17482        let (hd, nh, nhkv, tkvi, nsp) = (
17483            head_dim as i32,
17484            n_head as i32,
17485            n_head_kv as i32,
17486            t_kv as i32,
17487            n_splits as i32,
17488        );
17489        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17490        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
17491        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
17492        // silently truncating the accumulator.
17493        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17494        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
17495        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
17496        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
17497        // 178.4 -> 173.7 when 512 rode vec unconditionally).
17498        let fa512_min = fa512_min_tkv();
17499        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
17500        // g-module keeps the v4 pick (its class is not the depth-decay class).
17501        let deep = fa_vec
17502            && head_dim == 256
17503            && fa_v4_at(t_kv)
17504            && !g
17505            && fa_deep_at(t_kv)
17506            && !matches!(fa_v4_mode(), "noB3" | "stage");
17507        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17508            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
17509            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
17510            let gqa = (n_head / n_head_kv).max(1) as u32;
17511            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17512            (
17513                fv,
17514                LaunchConfig {
17515                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17516                    block_dim: (32, gqa, 1),
17517                    shared_mem_bytes: 0,
17518                },
17519            )
17520        } else if fa_vec && head_dim <= 256 {
17521            let gqa = (n_head / n_head_kv).max(1) as u32;
17522            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
17523            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
17524            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
17525            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
17526            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
17527            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
17528            // dequant each tile ONCE per block.
17529            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
17530            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
17531            // there by 12x — latency, not bandwidth, rules small KV).
17532            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17533            let smem_tkv = *SMEM_TKV.get_or_init(|| {
17534                std::env::var("MEMRA_FA_SMEM_TKV")
17535                    .ok()
17536                    .and_then(|v| v.parse().ok())
17537                    .unwrap_or_else(|| {
17538                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17539                    })
17540            });
17541            if fa_v4_at(t_kv) && head_dim == 256 {
17542                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
17543                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
17544                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
17545                let v4name = match fa_v4_mode() {
17546                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
17547                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
17548                    _ if deep => "fa_decode_vec_q_v4_deep",
17549                    _ => "fa_decode_vec_q_v4",
17550                };
17551                let fv = if g {
17552                    self.func_g(v4name)
17553                } else {
17554                    self.func(v4name)
17555                };
17556                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
17557                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
17558                let shmem = (if deep { 12160 } else { 11520 }
17559                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17560                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17561                fv.set_attribute(
17562                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17563                    shmem as i32,
17564                )?;
17565                (
17566                    fv,
17567                    LaunchConfig {
17568                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17569                        block_dim: (32, gqa, 1),
17570                        shared_mem_bytes: shmem,
17571                    },
17572                )
17573            } else if fa_v3_active(head_dim) {
17574                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
17575                // smem = sV only (half of v2's).
17576                let fv = if g {
17577                    self.func_g("fa_decode_vec_q_v3")
17578                } else {
17579                    self.func("fa_decode_vec_q_v3")
17580                };
17581                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
17582                (
17583                    fv,
17584                    LaunchConfig {
17585                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17586                        block_dim: (32, gqa, 1),
17587                        shared_mem_bytes: shmem,
17588                    },
17589                )
17590            } else if fa_v2_on() {
17591                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
17592                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
17593                // partials; same 32KB sK+sV tile as the smem twin.
17594                let fv = if g {
17595                    self.func_g("fa_decode_vec_q_v2")
17596                } else {
17597                    self.func("fa_decode_vec_q_v2")
17598                };
17599                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17600                (
17601                    fv,
17602                    LaunchConfig {
17603                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17604                        block_dim: (32, gqa, 1),
17605                        shared_mem_bytes: shmem,
17606                    },
17607                )
17608            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17609            {
17610                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
17611                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
17612                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
17613                let fv = if g {
17614                    self.func_g("fa_decode_vec_q_smem")
17615                } else {
17616                    self.func("fa_decode_vec_q_smem")
17617                };
17618                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17619                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17620                fv.set_attribute(
17621                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17622                    shmem as i32,
17623                )?;
17624                (
17625                    fv,
17626                    LaunchConfig {
17627                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17628                        block_dim: (32, gqa, 1),
17629                        shared_mem_bytes: shmem,
17630                    },
17631                )
17632            } else {
17633                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
17634                // dequant, zero dynamic shared memory.
17635                let fv = if g {
17636                    self.func_g("fa_decode_vec_q")
17637                } else {
17638                    self.func("fa_decode_vec_q")
17639                };
17640                (
17641                    fv,
17642                    LaunchConfig {
17643                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17644                        block_dim: (32, gqa, 1),
17645                        shared_mem_bytes: 0,
17646                    },
17647                )
17648            }
17649        } else {
17650            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
17651            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
17652            return self.fa_decode_scalar_unified(
17653                q,
17654                k,
17655                v,
17656                o,
17657                head_dim,
17658                n_head,
17659                n_head_kv,
17660                t_kv,
17661                None,
17662                scale,
17663                n_splits,
17664                if fa_vec { sp } else { 256 },
17665                k_tok_bytes,
17666                v_tok_bytes,
17667                g,
17668                part_o,
17669                part_m,
17670                part_l,
17671                None,
17672            );
17673        };
17674        let __s_b = self.gpu.stream();
17675        let mut b = __s_b.launch_builder(&f);
17676        b.arg(q)
17677            .arg(k)
17678            .arg(v)
17679            .arg(&mut *part_o)
17680            .arg(&mut *part_m)
17681            .arg(&mut *part_l)
17682            .arg(&hd)
17683            .arg(&nh)
17684            .arg(&nhkv)
17685            .arg(&tkvi)
17686            .arg(&scale)
17687            .arg(&nsp)
17688            .arg(&ktb)
17689            .arg(&vtb);
17690        unsafe {
17691            b.launch(cfg)?;
17692        }
17693        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
17694        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
17695        let (fc, cfg2) = (
17696            if g {
17697                self.func_g("fa_decode_combine_f32")
17698            } else {
17699                self.fa_func("fa_decode_combine_f32", head_dim)
17700            },
17701            LaunchConfig {
17702                grid_dim: (n_head as u32, 1, 1),
17703                block_dim: (head_dim as u32, 1, 1),
17704                shared_mem_bytes: 0,
17705            },
17706        );
17707        let __s_b2 = self.gpu.stream();
17708        let mut b2 = __s_b2.launch_builder(&fc);
17709        b2.arg(&*part_o)
17710            .arg(&*part_m)
17711            .arg(&*part_l)
17712            .arg(o)
17713            .arg(&hd)
17714            .arg(&nh)
17715            .arg(&nsp);
17716        unsafe {
17717            b2.launch(cfg2)?;
17718        }
17719        Ok(())
17720    }
17721
17722    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
17723    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
17724    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
17725    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
17726    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
17727    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
17728    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
17729    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
17730    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
17731    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
17732    #[allow(clippy::too_many_arguments)]
17733    pub fn fa_decode_batch_seqs_v4(
17734        &self,
17735        q: &CudaSlice<f32>,
17736        kv_ptrs: &cudarc::driver::CudaView<u64>,
17737        pos_seq: &CudaSlice<i32>,
17738        o: &mut CudaSlice<f32>,
17739        head_dim: usize,
17740        n_head: usize,
17741        n_head_kv: usize,
17742        b_n: usize,
17743        t_kv_max: usize,
17744        scale: f32,
17745        split_keys: usize,
17746        k_tok_bytes: usize,
17747        v_tok_bytes: usize,
17748    ) -> Result<(), Box<dyn std::error::Error>> {
17749        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
17750        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
17751        let o_len = b_n * n_head * n_splits_max * head_dim;
17752        let ml_len = b_n * n_head * n_splits_max;
17753        let mut part_guard = self.fa_part_pool.lock().unwrap();
17754        if part_guard
17755            .as_ref()
17756            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17757            .unwrap_or(true)
17758        {
17759            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
17760            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
17761            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
17762            // later live allocations land at those addresses, and the next graph REPLAY writes
17763            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
17764            // output corruption began the burst after the trunk's t_kv growth first realloc'd
17765            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
17766            // the baked addresses alive (single-stream: eager writes the new buffers, replays
17767            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
17768            // (total retired < final size).
17769            let old = part_guard.take();
17770            let (co, cm) = old
17771                .as_ref()
17772                .map(|pp| (pp.0.len(), pp.1.len()))
17773                .unwrap_or((0, 0));
17774            if let Some(old) = old {
17775                self.fa_part_retired.lock().unwrap().push(old);
17776            }
17777            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17778                eprintln!(
17779                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17780                    co, o_len, cm, ml_len
17781                );
17782            }
17783            *part_guard = Some((
17784                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17785                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17786                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17787            ));
17788        }
17789        let pg = part_guard.as_mut().unwrap();
17790        self.gpu
17791            .stream()
17792            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17793        self.gpu
17794            .stream()
17795            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17796        self.gpu
17797            .stream()
17798            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17799        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17800        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17801        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
17802        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17803        let gqa = (n_head / n_head_kv).max(1) as u32;
17804        let f = self.func("fa_decode_vec_q_seqs_v4");
17805        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
17806        let shmem = (11520 + 32 * head_dim * 2) as u32;
17807        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17808        f.set_attribute(
17809            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17810            shmem as i32,
17811        )?;
17812        let cfg = LaunchConfig {
17813            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
17814            block_dim: (32, gqa, 1),
17815            shared_mem_bytes: shmem,
17816        };
17817        {
17818            let __s_b = self.gpu.stream();
17819            let mut b = __s_b.launch_builder(&f);
17820            b.arg(q)
17821                .arg(kv_ptrs)
17822                .arg(pos_seq)
17823                .arg(&mut *part_o)
17824                .arg(&mut *part_m)
17825                .arg(&mut *part_l)
17826                .arg(&hd)
17827                .arg(&nh)
17828                .arg(&nhkv)
17829                .arg(&scale)
17830                .arg(&nspm)
17831                .arg(&spk)
17832                .arg(&ktb)
17833                .arg(&vtb);
17834            unsafe {
17835                b.launch(cfg)?;
17836            }
17837        }
17838        let fc = self.func("fa_decode_combine_seqs");
17839        let cfg2 = LaunchConfig {
17840            grid_dim: (n_head as u32, b_n as u32, 1),
17841            block_dim: (head_dim as u32, 1, 1),
17842            shared_mem_bytes: 0,
17843        };
17844        let __s_b2 = self.gpu.stream();
17845        let mut b2 = __s_b2.launch_builder(&fc);
17846        b2.arg(&*part_o)
17847            .arg(&*part_m)
17848            .arg(&*part_l)
17849            .arg(o)
17850            .arg(&hd)
17851            .arg(&nh)
17852            .arg(pos_seq)
17853            .arg(&nspm)
17854            .arg(&spk);
17855        unsafe {
17856            b2.launch(cfg2)?;
17857        }
17858        Ok(())
17859    }
17860
17861    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
17862    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
17863    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
17864    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
17865    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
17866    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
17867    #[allow(clippy::too_many_arguments)]
17868    pub fn append_kv_quantized_seqs(
17869        &self,
17870        k_rows: &CudaSlice<f32>,
17871        v_rows: &CudaSlice<f32>,
17872        kv_ptrs: &cudarc::driver::CudaView<u64>,
17873        pos_seq: &CudaSlice<i32>,
17874        b_n: usize,
17875        kv_dim_k: usize,
17876        kv_dim_v: usize,
17877        k_tok_bytes: usize,
17878        v_tok_bytes: usize,
17879    ) -> Result<(), Box<dyn std::error::Error>> {
17880        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
17881        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17882        let cfg = LaunchConfig {
17883            grid_dim: (nblk, b_n as u32, 1),
17884            block_dim: (32, 1, 1),
17885            shared_mem_bytes: 0,
17886        };
17887        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17888        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17889        let __s_b = self.gpu.stream();
17890        let mut b = __s_b.launch_builder(&f);
17891        b.arg(k_rows)
17892            .arg(v_rows)
17893            .arg(kv_ptrs)
17894            .arg(pos_seq)
17895            .arg(&kdk)
17896            .arg(&kdv)
17897            .arg(&ktb)
17898            .arg(&vtb);
17899        unsafe {
17900            b.launch(cfg)?;
17901        }
17902        Ok(())
17903    }
17904
17905    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
17906    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
17907    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
17908    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
17909    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
17910    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
17911        std::env::var("MEMRA_NO_FA_VEC").is_err()
17912            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
17913            && base_len + 1 >= fa_vec_min_tkv()
17914            && head_dim <= 256
17915            && head_dim % 32 == 0
17916    }
17917
17918    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
17919    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
17920    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
17921    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
17922    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
17923    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
17924    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
17925    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
17926    #[allow(clippy::too_many_arguments)]
17927    pub fn fa_decode_rows(
17928        &self,
17929        q: &CudaSlice<f32>,
17930        k: &cudarc::driver::CudaView<u8>,
17931        v: &cudarc::driver::CudaView<u8>,
17932        o: &mut CudaSlice<f32>,
17933        head_dim: usize,
17934        n_head: usize,
17935        n_head_kv: usize,
17936        base_len: usize,
17937        t: usize,
17938        scale: f32,
17939        k_tok_bytes: usize,
17940        v_tok_bytes: usize,
17941        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
17942        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
17943        // keep the host arg. None is a bug for hd512 (asserted below).
17944        base_dev: Option<(&CudaSlice<i32>, i32)>,
17945        // K and V planes hold the same values (gemma globals, wv:=wk): pick
17946        // the _kv twin — V plane never read, value rides the q8_0 key dq.
17947        kv_shared: bool,
17948        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
17949        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
17950        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
17951        g: bool,
17952        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
17953        // (hd512 path) — the standalone quantize launch folds away.
17954        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17955    ) -> Result<(), Box<dyn std::error::Error>> {
17956        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
17957        let t_kv_max = base_len + t; // LAST row's key bound
17958        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
17959        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
17960        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
17961        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
17962        // (parity law), so the partition is freely tunable — verify and decode move together.
17963        if head_dim == 512 {
17964            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17965            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
17966            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
17967            let v = *SP512.get_or_init(|| {
17968                std::env::var("MEMRA_FA_SP512")
17969                    .ok()
17970                    .and_then(|x| x.parse().ok())
17971                    .unwrap_or(0)
17972            });
17973            sp = if v >= 8 {
17974                v
17975            } else {
17976                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17977            };
17978        }
17979        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17980        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17981        let gqa = (n_head / n_head_kv).max(1) as u32;
17982        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
17983        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
17984        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
17985        // the different partition changes the combine's FP order (greedy tie flips at depth;
17986        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
17987        // consecutive rows by their OWN ladder value and launch once per group — each row then
17988        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
17989        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
17990        // sp override is t_kv-independent by construction).
17991        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
17992        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
17993            groups.push((0, t, sp));
17994        } else {
17995            let mut r0 = 0usize;
17996            while r0 < t {
17997                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
17998                let mut r1 = r0 + 1;
17999                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18000                    r1 += 1;
18001                }
18002                groups.push((r0, r1 - r0, sp_g));
18003                r0 = r1;
18004            }
18005        }
18006        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
18007        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
18008        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
18009        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18010        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18011            std::env::var("MEMRA_FA_SMEM_TKV")
18012                .ok()
18013                .and_then(|v| v.parse().ok())
18014                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18015        });
18016        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18017        let v3 = fa_v3_active(head_dim);
18018        let smem_rows =
18019            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18020        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
18021        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
18022        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
18023        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
18024        let _ = kv_shared;
18025        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
18026        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
18027        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18028        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
18029        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
18030        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
18031        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
18032        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
18033        // (kv_head, split) stages its tile once and loops the rows over it — kills the
18034        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
18035        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
18036        // shared by every hd512 caller through this wrapper (decode+verify flip together;
18037        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
18038        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
18039        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
18040        // not unpack-bound; jsonl 2026-07-14.
18041        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18042        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
18043        let tb512 = head_dim == 512
18044            && sp <= 32
18045            && n_head / n_head_kv.max(1) <= 16
18046            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18047        let fname = if tb512 {
18048            "fa_decode_vec_q_rows_v4_512_tb"
18049        } else if i2 {
18050            "fa_decode_vec_q_rows_dpl16_i2"
18051        } else if head_dim == 512 {
18052            "fa_decode_vec_q_rows_dpl16"
18053        }
18054        // gemma globals (parity law)
18055        else if v4 {
18056            "fa_decode_vec_q_rows_v4"
18057        } else if v3 {
18058            "fa_decode_vec_q_rows_v3"
18059        } else if fa_v2_on() {
18060            "fa_decode_vec_q_rows_v2"
18061        } else if smem_rows {
18062            "fa_decode_vec_q_rows_smem"
18063        } else {
18064            "fa_decode_vec_q_rows"
18065        };
18066        let f = if head_dim == 512 {
18067            self.fa_func(fname, head_dim)
18068        } else if g {
18069            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
18070            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
18071            // g-module rows against decode's g-module v4 — different programs, short-VG
18072            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
18073            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
18074            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
18075            // dq macros are format-aware.
18076            self.func_g(if smem_rows {
18077                "fa_decode_vec_q_rows"
18078            } else {
18079                fname
18080            })
18081        } else {
18082            self.func(fname)
18083        };
18084        let shmem = if tb512 {
18085            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
18086            let gk = Self::gkv_on();
18087            let sh =
18088                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18089            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18090            f.set_attribute(
18091                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18092                sh as i32,
18093            )?;
18094            sh
18095        } else if v4 || v3 || smem_rows || fa_v2_on() {
18096            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
18097            let sh = (if v4 {
18098                11520 + 32 * head_dim * if g { 1 } else { 2 }
18099            } else if v3 {
18100                32 * head_dim * 2
18101            } else {
18102                2 * 32 * head_dim * 2
18103            }) as u32;
18104            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18105            f.set_attribute(
18106                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18107                sh as i32,
18108            )?;
18109            sh
18110        } else {
18111            0
18112        };
18113        // Per-GROUP launches (single group in the common case — identical to the pre-fix
18114        // single launch there): each group gets its own partials (the rows kernel indexes
18115        // partials by its LOCAL grid.z row) and q/o row-offset views.
18116        for &(r0, t_g, sp_g) in &groups {
18117            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18118            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18119            let base_i = (base_len + r0) as i32;
18120            let o_len = t_g * n_head * n_splits_g * head_dim;
18121            let ml_len = t_g * n_head * n_splits_g;
18122            let mut part_guard = self.fa_part_pool.lock().unwrap();
18123            if part_guard
18124                .as_ref()
18125                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18126                .unwrap_or(true)
18127            {
18128                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18129                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18130                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18131                // later live allocations land at those addresses, and the next graph REPLAY writes
18132                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18133                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18134                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18135                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18136                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18137                // (total retired < final size).
18138                let old = part_guard.take();
18139                let (co, cm) = old
18140                    .as_ref()
18141                    .map(|pp| (pp.0.len(), pp.1.len()))
18142                    .unwrap_or((0, 0));
18143                if let Some(old) = old {
18144                    self.fa_part_retired.lock().unwrap().push(old);
18145                }
18146                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18147                    eprintln!(
18148                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18149                        co, o_len, cm, ml_len
18150                    );
18151                }
18152                *part_guard = Some((
18153                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18154                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18155                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18156                ));
18157            }
18158            let pg = part_guard.as_mut().unwrap();
18159            self.gpu
18160                .stream()
18161                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18162            self.gpu
18163                .stream()
18164                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18165            self.gpu
18166                .stream()
18167                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18168            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18169            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18170            let qv = self.view(q, t * n_head * head_dim);
18171            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18172            let cfg = LaunchConfig {
18173                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18174                block_dim: (32, gqa, 1),
18175                shared_mem_bytes: shmem,
18176            };
18177            {
18178                let __s_b = self.gpu.stream();
18179                let mut b = __s_b.launch_builder(&f);
18180                if tb512 {
18181                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
18182                    let (bd, plus) =
18183                        base_dev.expect("hd512 rows twin requires a device base counter");
18184                    let plus_g = plus + r0 as i32;
18185                    let nr = t_g as i32;
18186                    if Self::pdl_on() && Self::pdl_wb_on() {
18187                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
18188                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18189                        let s = &self.gpu.stream();
18190                        let (pq, _b0) = q_g.device_ptr(s);
18191                        let (pk, _b1) = k.device_ptr(s);
18192                        let (pv, _b2) = v.device_ptr(s);
18193                        let (po, _b3) = part_o.device_ptr_mut(s);
18194                        let (pm, _b4) = part_m.device_ptr_mut(s);
18195                        let (pl, _b5) = part_l.device_ptr_mut(s);
18196                        let (pb, _b6) = bd.device_ptr(s);
18197                        let mut ps = [
18198                            &pq as *const _ as *mut std::ffi::c_void,
18199                            &pk as *const _ as *mut _,
18200                            &pv as *const _ as *mut _,
18201                            &po as *const _ as *mut _,
18202                            &pm as *const _ as *mut _,
18203                            &pl as *const _ as *mut _,
18204                            &hd as *const _ as *mut _,
18205                            &nh as *const _ as *mut _,
18206                            &nhkv as *const _ as *mut _,
18207                            &pb as *const _ as *mut _,
18208                            &plus_g as *const _ as *mut _,
18209                            &scale as *const _ as *mut _,
18210                            &nspm as *const _ as *mut _,
18211                            &spk as *const _ as *mut _,
18212                            &ktb as *const _ as *mut _,
18213                            &vtb as *const _ as *mut _,
18214                            &nr as *const _ as *mut _,
18215                        ];
18216                        unsafe {
18217                            self.launch_pdl_flash(
18218                                Self::gkv_on(),
18219                                "fa_decode_vec_q_rows_v4_512_tb",
18220                                (n_head_kv as u32, n_splits_g as u32, 1),
18221                                (32, gqa, 1),
18222                                shmem,
18223                                &mut ps,
18224                            )?;
18225                        }
18226                    } else {
18227                        let cfg_tb = LaunchConfig {
18228                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18229                            block_dim: (32, gqa, 1),
18230                            shared_mem_bytes: shmem,
18231                        };
18232                        b.arg(&q_g)
18233                            .arg(k)
18234                            .arg(v)
18235                            .arg(&mut *part_o)
18236                            .arg(&mut *part_m)
18237                            .arg(&mut *part_l)
18238                            .arg(&hd)
18239                            .arg(&nh)
18240                            .arg(&nhkv)
18241                            .arg(bd)
18242                            .arg(&plus_g)
18243                            .arg(&scale)
18244                            .arg(&nspm)
18245                            .arg(&spk)
18246                            .arg(&ktb)
18247                            .arg(&vtb)
18248                            .arg(&nr);
18249                        unsafe {
18250                            b.launch(cfg_tb)?;
18251                        }
18252                    }
18253                } else if head_dim == 512 {
18254                    let (bd, plus) =
18255                        base_dev.expect("hd512 rows twin requires a device base counter");
18256                    let plus_g = plus + r0 as i32;
18257                    b.arg(&q_g)
18258                        .arg(k)
18259                        .arg(v)
18260                        .arg(&mut *part_o)
18261                        .arg(&mut *part_m)
18262                        .arg(&mut *part_l)
18263                        .arg(&hd)
18264                        .arg(&nh)
18265                        .arg(&nhkv)
18266                        .arg(bd)
18267                        .arg(&plus_g)
18268                        .arg(&scale)
18269                        .arg(&nspm)
18270                        .arg(&spk)
18271                        .arg(&ktb)
18272                        .arg(&vtb);
18273                    unsafe {
18274                        b.launch(cfg)?;
18275                    }
18276                } else {
18277                    b.arg(&q_g)
18278                        .arg(k)
18279                        .arg(v)
18280                        .arg(&mut *part_o)
18281                        .arg(&mut *part_m)
18282                        .arg(&mut *part_l)
18283                        .arg(&hd)
18284                        .arg(&nh)
18285                        .arg(&nhkv)
18286                        .arg(&base_i)
18287                        .arg(&scale)
18288                        .arg(&nspm)
18289                        .arg(&spk)
18290                        .arg(&ktb)
18291                        .arg(&vtb);
18292                    unsafe {
18293                        b.launch(cfg)?;
18294                    }
18295                }
18296            }
18297            let cfg2 = LaunchConfig {
18298                grid_dim: (n_head as u32, t_g as u32, 1),
18299                block_dim: (head_dim as u32, 1, 1),
18300                shared_mem_bytes: 0,
18301            };
18302            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18303            if head_dim == 512 {
18304                // device-len combine (shared by verify/eager/graph — parity by symbol): the
18305                // per-row n_splits derives from the SAME counter the rows kernel read.
18306                let (bd, plus) = base_dev.unwrap();
18307                let plus_g = plus + r0 as i32;
18308                if let Some((oq, od)) = q8_out.as_mut() {
18309                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
18310                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18311                    if Self::pdl_on() && Self::pdl_wb_on() {
18312                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
18313                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18314                        let s = &self.gpu.stream();
18315                        let (po, _g0) = part_o.device_ptr(s);
18316                        let (pm, _g1) = part_m.device_ptr(s);
18317                        let (pl, _g2) = part_l.device_ptr(s);
18318                        let (pq, _g3) = oq.device_ptr_mut(s);
18319                        let (pd, _g4) = od.device_ptr_mut(s);
18320                        let (pb, _g5) = bd.device_ptr(s);
18321                        let mut ps = [
18322                            &po as *const _ as *mut std::ffi::c_void,
18323                            &pm as *const _ as *mut _,
18324                            &pl as *const _ as *mut _,
18325                            &pq as *const _ as *mut _,
18326                            &pd as *const _ as *mut _,
18327                            &hd as *const _ as *mut _,
18328                            &nh as *const _ as *mut _,
18329                            &pb as *const _ as *mut _,
18330                            &plus_g as *const _ as *mut _,
18331                            &nspm as *const _ as *mut _,
18332                            &spk as *const _ as *mut _,
18333                        ];
18334                        unsafe {
18335                            self.launch_pdl_flash(
18336                                Self::gkv_on(),
18337                                "fa_decode_combine_rows_dc_q8_1",
18338                                cfg2.grid_dim,
18339                                cfg2.block_dim,
18340                                0,
18341                                &mut ps,
18342                            )?;
18343                        }
18344                        continue;
18345                    }
18346                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18347                    let __s_b2 = self.gpu.stream();
18348                    let mut b2 = __s_b2.launch_builder(&fc);
18349                    b2.arg(&*part_o)
18350                        .arg(&*part_m)
18351                        .arg(&*part_l)
18352                        .arg(&mut **oq)
18353                        .arg(&mut **od)
18354                        .arg(&hd)
18355                        .arg(&nh)
18356                        .arg(bd)
18357                        .arg(&plus_g)
18358                        .arg(&nspm)
18359                        .arg(&spk);
18360                    unsafe {
18361                        b2.launch(cfg2)?;
18362                    }
18363                    continue;
18364                }
18365                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18366                let __s_b2 = self.gpu.stream();
18367                let mut b2 = __s_b2.launch_builder(&fc);
18368                b2.arg(&*part_o)
18369                    .arg(&*part_m)
18370                    .arg(&*part_l)
18371                    .arg(&mut o_g)
18372                    .arg(&hd)
18373                    .arg(&nh)
18374                    .arg(bd)
18375                    .arg(&plus_g)
18376                    .arg(&nspm)
18377                    .arg(&spk);
18378                unsafe {
18379                    b2.launch(cfg2)?;
18380                }
18381            } else {
18382                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
18383                // leave the caller's pair unwritten (consumer would read garbage).
18384                assert!(
18385                    q8_out.is_none(),
18386                    "rows q8 emit requires the hd512 dc combine"
18387                );
18388                let fc = self.func("fa_decode_combine_rows");
18389                let __s_b2 = self.gpu.stream();
18390                let mut b2 = __s_b2.launch_builder(&fc);
18391                b2.arg(&*part_o)
18392                    .arg(&*part_m)
18393                    .arg(&*part_l)
18394                    .arg(&mut o_g)
18395                    .arg(&hd)
18396                    .arg(&nh)
18397                    .arg(&base_i)
18398                    .arg(&nspm)
18399                    .arg(&spk);
18400                unsafe {
18401                    b2.launch(cfg2)?;
18402                }
18403            }
18404        }
18405        Ok(())
18406    }
18407
18408    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
18409    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
18410    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
18411    #[allow(clippy::too_many_arguments)]
18412    pub fn fa_decode_rows_w(
18413        &self,
18414        q: &CudaSlice<f32>,
18415        k: &cudarc::driver::CudaView<u8>,
18416        v: &cudarc::driver::CudaView<u8>,
18417        o: &mut CudaSlice<f32>,
18418        head_dim: usize,
18419        n_head: usize,
18420        n_head_kv: usize,
18421        base_dev: &CudaSlice<i32>,
18422        base_plus: i32,
18423        t: usize,
18424        scale: f32,
18425        window: usize,
18426        k_tok_bytes: usize,
18427        v_tok_bytes: usize,
18428        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18429    ) -> Result<(), Box<dyn std::error::Error>> {
18430        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
18431        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
18432        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
18433        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
18434        debug_assert!(head_dim == 256);
18435        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
18436        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
18437        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
18438        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
18439        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
18440        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
18441        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
18442        let sp = {
18443            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18444            let v = *SPW.get_or_init(|| {
18445                std::env::var("MEMRA_FA_SPW")
18446                    .ok()
18447                    .and_then(|x| x.parse().ok())
18448                    .unwrap_or(0)
18449            });
18450            if v >= 8 {
18451                v
18452            } else {
18453                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18454            }
18455        };
18456        let n_splits_max = (window + sp - 1) / sp;
18457        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18458        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18459        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18460        let gqa = (n_head / n_head_kv).max(1) as u32;
18461        let o_len = t * n_head * n_splits_max * head_dim;
18462        let ml_len = t * n_head * n_splits_max;
18463        let mut part_guard = self.fa_part_pool.lock().unwrap();
18464        if part_guard
18465            .as_ref()
18466            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18467            .unwrap_or(true)
18468        {
18469            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18470            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18471            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18472            // later live allocations land at those addresses, and the next graph REPLAY writes
18473            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18474            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18475            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18476            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18477            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18478            // (total retired < final size).
18479            let old = part_guard.take();
18480            let (co, cm) = old
18481                .as_ref()
18482                .map(|pp| (pp.0.len(), pp.1.len()))
18483                .unwrap_or((0, 0));
18484            if let Some(old) = old {
18485                self.fa_part_retired.lock().unwrap().push(old);
18486            }
18487            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18488                eprintln!(
18489                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18490                    co, o_len, cm, ml_len
18491                );
18492            }
18493            *part_guard = Some((
18494                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18495                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18496                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18497            ));
18498        }
18499        let pg = part_guard.as_mut().unwrap();
18500        self.gpu
18501            .stream()
18502            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18503        self.gpu
18504            .stream()
18505            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18506        self.gpu
18507            .stream()
18508            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18509        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18510        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
18511        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
18512        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
18513        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
18514        // floor (deep-ctx broadcast win); register twin between.
18515        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18516        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18517            std::env::var("MEMRA_FA_SMEM_TKV")
18518                .ok()
18519                .and_then(|v| v.parse().ok())
18520                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18521        });
18522        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
18523        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
18524        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
18525        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
18526        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
18527        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18528        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
18529        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
18530        // per (lane, format-module) keeps parity structural; the old register-i2 detour
18531        // (-33%) is retired.
18532        let wg = Self::wkv_on();
18533        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
18534        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
18535        let sp2 =
18536            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18537        if sp2 {
18538            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18539            if Self::pdl_on() && Self::pdl_wb_on() {
18540                // wave-B2b: flavor mirrors wg.
18541                use cudarc::driver::{DevicePtr, DevicePtrMut};
18542                let s = &self.gpu.stream();
18543                let (pq, _b0) = q.device_ptr(s);
18544                let (pk, _b1) = k.device_ptr(s);
18545                let (pv, _b2) = v.device_ptr(s);
18546                let (po, _b3) = part_o.device_ptr_mut(s);
18547                let (pm, _b4) = part_m.device_ptr_mut(s);
18548                let (pl, _b5) = part_l.device_ptr_mut(s);
18549                let (pb, _b6) = base_dev.device_ptr(s);
18550                let mut ps = [
18551                    &pq as *const _ as *mut std::ffi::c_void,
18552                    &pk as *const _ as *mut _,
18553                    &pv as *const _ as *mut _,
18554                    &po as *const _ as *mut _,
18555                    &pm as *const _ as *mut _,
18556                    &pl as *const _ as *mut _,
18557                    &hd as *const _ as *mut _,
18558                    &nh as *const _ as *mut _,
18559                    &nhkv as *const _ as *mut _,
18560                    &pb as *const _ as *mut _,
18561                    &base_plus as *const _ as *mut _,
18562                    &scale as *const _ as *mut _,
18563                    &nspm as *const _ as *mut _,
18564                    &spk as *const _ as *mut _,
18565                    &ktb as *const _ as *mut _,
18566                    &vtb as *const _ as *mut _,
18567                    &wini as *const _ as *mut _,
18568                ];
18569                unsafe {
18570                    self.launch_pdl_flash(
18571                        wg,
18572                        "fa_decode_vec_q_rows_v4_w_sp",
18573                        (n_head_kv as u32, n_splits_max as u32, t as u32),
18574                        (32, gqa + 1, 1),
18575                        sh,
18576                        &mut ps,
18577                    )?;
18578                }
18579            } else {
18580                let f = if wg {
18581                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18582                } else {
18583                    self.func("fa_decode_vec_q_rows_v4_w_sp")
18584                };
18585                f.set_attribute(
18586                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18587                    sh as i32,
18588                )?;
18589                let cfg = LaunchConfig {
18590                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18591                    block_dim: (32, gqa + 1, 1),
18592                    shared_mem_bytes: sh,
18593                };
18594                let __s_b = self.gpu.stream();
18595                let mut b = __s_b.launch_builder(&f);
18596                b.arg(q)
18597                    .arg(k)
18598                    .arg(v)
18599                    .arg(&mut *part_o)
18600                    .arg(&mut *part_m)
18601                    .arg(&mut *part_l)
18602                    .arg(&hd)
18603                    .arg(&nh)
18604                    .arg(&nhkv)
18605                    .arg(base_dev)
18606                    .arg(&base_plus)
18607                    .arg(&scale)
18608                    .arg(&nspm)
18609                    .arg(&spk)
18610                    .arg(&ktb)
18611                    .arg(&vtb)
18612                    .arg(&wini);
18613                unsafe {
18614                    b.launch(cfg)?;
18615                }
18616            }
18617        } else {
18618            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18619                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
18620                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18621                use cudarc::driver::{DevicePtr, DevicePtrMut};
18622                let s = &self.gpu.stream();
18623                let (pq, _b0) = q.device_ptr(s);
18624                let (pk, _b1) = k.device_ptr(s);
18625                let (pv, _b2) = v.device_ptr(s);
18626                let (po, _b3) = part_o.device_ptr_mut(s);
18627                let (pm, _b4) = part_m.device_ptr_mut(s);
18628                let (pl, _b5) = part_l.device_ptr_mut(s);
18629                let (pb, _b6) = base_dev.device_ptr(s);
18630                let mut ps = [
18631                    &pq as *const _ as *mut std::ffi::c_void,
18632                    &pk as *const _ as *mut _,
18633                    &pv as *const _ as *mut _,
18634                    &po as *const _ as *mut _,
18635                    &pm as *const _ as *mut _,
18636                    &pl as *const _ as *mut _,
18637                    &hd as *const _ as *mut _,
18638                    &nh as *const _ as *mut _,
18639                    &nhkv as *const _ as *mut _,
18640                    &pb as *const _ as *mut _,
18641                    &base_plus as *const _ as *mut _,
18642                    &scale as *const _ as *mut _,
18643                    &nspm as *const _ as *mut _,
18644                    &spk as *const _ as *mut _,
18645                    &ktb as *const _ as *mut _,
18646                    &vtb as *const _ as *mut _,
18647                    &wini as *const _ as *mut _,
18648                ];
18649                unsafe {
18650                    self.launch_pdl_flash(
18651                        wg,
18652                        "fa_decode_vec_q_rows_v4_w",
18653                        (n_head_kv as u32, n_splits_max as u32, t as u32),
18654                        (32, gqa, 1),
18655                        sh,
18656                        &mut ps,
18657                    )?;
18658                }
18659            } else {
18660                let pick = |name: &str| {
18661                    if wg {
18662                        self.func_g(name)
18663                    } else {
18664                        self.func(name)
18665                    }
18666                };
18667                let (f, sh) = if fa_v4_at(window) {
18668                    let f = pick("fa_decode_vec_q_rows_v4_w");
18669                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
18670                } else if smem_tkv > 0 && window >= smem_tkv {
18671                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
18672                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
18673                    (
18674                        pick("fa_decode_vec_q_rows_smem_w"),
18675                        (2 * 32 * head_dim * 2) as u32,
18676                    )
18677                } else {
18678                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
18679                };
18680                f.set_attribute(
18681                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18682                    sh as i32,
18683                )?;
18684                let cfg = LaunchConfig {
18685                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18686                    block_dim: (32, gqa, 1),
18687                    shared_mem_bytes: sh,
18688                };
18689                let __s_b = self.gpu.stream();
18690                let mut b = __s_b.launch_builder(&f);
18691                b.arg(q)
18692                    .arg(k)
18693                    .arg(v)
18694                    .arg(&mut *part_o)
18695                    .arg(&mut *part_m)
18696                    .arg(&mut *part_l)
18697                    .arg(&hd)
18698                    .arg(&nh)
18699                    .arg(&nhkv)
18700                    .arg(base_dev)
18701                    .arg(&base_plus)
18702                    .arg(&scale)
18703                    .arg(&nspm)
18704                    .arg(&spk)
18705                    .arg(&ktb)
18706                    .arg(&vtb)
18707                    .arg(&wini);
18708                unsafe {
18709                    b.launch(cfg)?;
18710                }
18711            }
18712        }
18713        let cfg2 = LaunchConfig {
18714            grid_dim: (n_head as u32, t as u32, 1),
18715            block_dim: (head_dim as u32, 1, 1),
18716            shared_mem_bytes: 0,
18717        };
18718        if let Some((oq, od)) = q8_out {
18719            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
18720            // consumes the pair directly; the standalone quantize launch folds away.
18721            if Self::pdl_on() && Self::pdl_wb_on() {
18722                // wave-B2: flavor mirrors the builder's wg choice.
18723                use cudarc::driver::{DevicePtr, DevicePtrMut};
18724                let s = &self.gpu.stream();
18725                let (po, _g0) = part_o.device_ptr(s);
18726                let (pm, _g1) = part_m.device_ptr(s);
18727                let (pl, _g2) = part_l.device_ptr(s);
18728                let (pq, _g3) = oq.device_ptr_mut(s);
18729                let (pd, _g4) = od.device_ptr_mut(s);
18730                let mut ps = [
18731                    &po as *const _ as *mut std::ffi::c_void,
18732                    &pm as *const _ as *mut _,
18733                    &pl as *const _ as *mut _,
18734                    &pq as *const _ as *mut _,
18735                    &pd as *const _ as *mut _,
18736                    &hd as *const _ as *mut _,
18737                    &nh as *const _ as *mut _,
18738                    &nspm as *const _ as *mut _,
18739                    &spk as *const _ as *mut _,
18740                    &wini as *const _ as *mut _,
18741                ];
18742                unsafe {
18743                    self.launch_pdl_flash(
18744                        wg,
18745                        "fa_decode_combine_rows_w_q8_1",
18746                        cfg2.grid_dim,
18747                        cfg2.block_dim,
18748                        0,
18749                        &mut ps,
18750                    )?;
18751                }
18752                return Ok(());
18753            }
18754            let fc = if wg {
18755                self.func_g("fa_decode_combine_rows_w_q8_1")
18756            } else {
18757                self.func("fa_decode_combine_rows_w_q8_1")
18758            };
18759            let __s_b2 = self.gpu.stream();
18760            let mut b2 = __s_b2.launch_builder(&fc);
18761            b2.arg(&*part_o)
18762                .arg(&*part_m)
18763                .arg(&*part_l)
18764                .arg(oq)
18765                .arg(od)
18766                .arg(&hd)
18767                .arg(&nh)
18768                .arg(&nspm)
18769                .arg(&spk)
18770                .arg(&wini);
18771            unsafe {
18772                b2.launch(cfg2)?;
18773            }
18774            return Ok(());
18775        }
18776        let fc = if wg {
18777            self.func_g("fa_decode_combine_rows_w")
18778        } else {
18779            self.func("fa_decode_combine_rows_w")
18780        };
18781        let __s_b2 = self.gpu.stream();
18782        let mut b2 = __s_b2.launch_builder(&fc);
18783        b2.arg(&*part_o)
18784            .arg(&*part_m)
18785            .arg(&*part_l)
18786            .arg(o)
18787            .arg(&hd)
18788            .arg(&nh)
18789            .arg(&nspm)
18790            .arg(&spk)
18791            .arg(&wini);
18792        unsafe {
18793            b2.launch(cfg2)?;
18794        }
18795        Ok(())
18796    }
18797
18798    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
18799    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
18800    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
18801    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
18802    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
18803    #[allow(clippy::too_many_arguments)]
18804    pub fn fa_decode_rows_dc(
18805        &self,
18806        q: &CudaSlice<f32>,
18807        k: &cudarc::driver::CudaView<u8>,
18808        v: &cudarc::driver::CudaView<u8>,
18809        o: &mut CudaSlice<f32>,
18810        head_dim: usize,
18811        n_head: usize,
18812        n_head_kv: usize,
18813        base_dev: &CudaSlice<i32>,
18814        t_kv_upper: usize,
18815        t: usize,
18816        scale: f32,
18817        k_tok_bytes: usize,
18818        v_tok_bytes: usize,
18819        base_plus: i32,
18820        g: bool,
18821    ) -> Result<(), Box<dyn std::error::Error>> {
18822        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
18823        assert!(
18824            v4 || fa_v3_active(head_dim),
18825            "stream fa rows requires the v3 or v4 lane"
18826        );
18827        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
18828        if v4 {
18829            let sp = fa_split_keys(t_kv_upper, n_head_kv);
18830            let n_splits_max = (t_kv_upper + sp - 1) / sp;
18831            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18832            let (nspm, spk) = (n_splits_max as i32, sp as i32);
18833            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18834            let gqa = (n_head / n_head_kv).max(1) as u32;
18835            let o_len = t * n_head * n_splits_max * head_dim;
18836            let ml_len = t * n_head * n_splits_max;
18837            let mut part_guard = self.fa_part_pool.lock().unwrap();
18838            if part_guard
18839                .as_ref()
18840                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18841                .unwrap_or(true)
18842            {
18843                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18844                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18845                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18846                // later live allocations land at those addresses, and the next graph REPLAY writes
18847                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18848                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18849                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18850                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18851                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18852                // (total retired < final size).
18853                let old = part_guard.take();
18854                let (co, cm) = old
18855                    .as_ref()
18856                    .map(|pp| (pp.0.len(), pp.1.len()))
18857                    .unwrap_or((0, 0));
18858                if let Some(old) = old {
18859                    self.fa_part_retired.lock().unwrap().push(old);
18860                }
18861                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18862                    eprintln!(
18863                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18864                        co, o_len, cm, ml_len
18865                    );
18866                }
18867                *part_guard = Some((
18868                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18869                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18870                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18871                ));
18872            }
18873            let pg = part_guard.as_mut().unwrap();
18874            self.gpu
18875                .stream()
18876                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18877            self.gpu
18878                .stream()
18879                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18880            self.gpu
18881                .stream()
18882                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18883            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18884            let f = if g {
18885                self.func_g("fa_decode_vec_q_rows_v4_dc")
18886            } else {
18887                self.func("fa_decode_vec_q_rows_v4_dc")
18888            };
18889            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18890            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18891            f.set_attribute(
18892                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18893                sh as i32,
18894            )?;
18895            let cfg = LaunchConfig {
18896                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18897                block_dim: (32, gqa, 1),
18898                shared_mem_bytes: sh,
18899            };
18900            let __s_b = self.gpu.stream();
18901            let mut b = __s_b.launch_builder(&f);
18902            b.arg(q)
18903                .arg(k)
18904                .arg(v)
18905                .arg(&mut *part_o)
18906                .arg(&mut *part_m)
18907                .arg(&mut *part_l)
18908                .arg(&hd)
18909                .arg(&nh)
18910                .arg(&nhkv)
18911                .arg(base_dev)
18912                .arg(&base_plus)
18913                .arg(&scale)
18914                .arg(&nspm)
18915                .arg(&spk)
18916                .arg(&ktb)
18917                .arg(&vtb);
18918            unsafe {
18919                b.launch(cfg)?;
18920            }
18921            let fc = self.func("fa_decode_combine_rows_dc");
18922            let cfg2 = LaunchConfig {
18923                grid_dim: (n_head as u32, t as u32, 1),
18924                block_dim: (head_dim as u32, 1, 1),
18925                shared_mem_bytes: 0,
18926            };
18927            let __s_b2 = self.gpu.stream();
18928            let mut b2 = __s_b2.launch_builder(&fc);
18929            b2.arg(&*part_o)
18930                .arg(&*part_m)
18931                .arg(&*part_l)
18932                .arg(o)
18933                .arg(&hd)
18934                .arg(&nh)
18935                .arg(base_dev)
18936                .arg(&base_plus)
18937                .arg(&nspm)
18938                .arg(&spk);
18939            unsafe {
18940                b2.launch(cfg2)?;
18941            }
18942            return Ok(());
18943        }
18944        let sp = fa_split_keys(t_kv_upper, n_head_kv);
18945        let n_splits_max = (t_kv_upper + sp - 1) / sp;
18946        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18947        let (nspm, spk) = (n_splits_max as i32, sp as i32);
18948        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18949        let gqa = (n_head / n_head_kv).max(1) as u32;
18950        let o_len = t * n_head * n_splits_max * head_dim;
18951        let ml_len = t * n_head * n_splits_max;
18952        let mut part_guard = self.fa_part_pool.lock().unwrap();
18953        if part_guard
18954            .as_ref()
18955            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18956            .unwrap_or(true)
18957        {
18958            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18959            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18960            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18961            // later live allocations land at those addresses, and the next graph REPLAY writes
18962            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18963            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18964            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18965            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18966            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18967            // (total retired < final size).
18968            let old = part_guard.take();
18969            let (co, cm) = old
18970                .as_ref()
18971                .map(|pp| (pp.0.len(), pp.1.len()))
18972                .unwrap_or((0, 0));
18973            if let Some(old) = old {
18974                self.fa_part_retired.lock().unwrap().push(old);
18975            }
18976            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18977                eprintln!(
18978                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18979                    co, o_len, cm, ml_len
18980                );
18981            }
18982            *part_guard = Some((
18983                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18984                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18985                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18986            ));
18987        }
18988        let pg = part_guard.as_mut().unwrap();
18989        self.gpu
18990            .stream()
18991            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18992        self.gpu
18993            .stream()
18994            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18995        self.gpu
18996            .stream()
18997            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18998        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18999        let f = self.func("fa_decode_vec_q_rows_v3_dc");
19000        let sh = (32 * head_dim * 2) as u32;
19001        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19002        f.set_attribute(
19003            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19004            sh as i32,
19005        )?;
19006        let cfg = LaunchConfig {
19007            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19008            block_dim: (32, gqa, 1),
19009            shared_mem_bytes: sh,
19010        };
19011        let __s_b = self.gpu.stream();
19012        let mut b = __s_b.launch_builder(&f);
19013        b.arg(q)
19014            .arg(k)
19015            .arg(v)
19016            .arg(&mut *part_o)
19017            .arg(&mut *part_m)
19018            .arg(&mut *part_l)
19019            .arg(&hd)
19020            .arg(&nh)
19021            .arg(&nhkv)
19022            .arg(base_dev)
19023            .arg(&scale)
19024            .arg(&nspm)
19025            .arg(&spk)
19026            .arg(&ktb)
19027            .arg(&vtb);
19028        unsafe {
19029            b.launch(cfg)?;
19030        }
19031        let fc = self.func("fa_decode_combine_rows_dc");
19032        let cfg2 = LaunchConfig {
19033            grid_dim: (n_head as u32, t as u32, 1),
19034            block_dim: (head_dim as u32, 1, 1),
19035            shared_mem_bytes: 0,
19036        };
19037        let plus0 = 0i32;
19038        let __s_b2 = self.gpu.stream();
19039        let mut b2 = __s_b2.launch_builder(&fc);
19040        b2.arg(&*part_o)
19041            .arg(&*part_m)
19042            .arg(&*part_l)
19043            .arg(o)
19044            .arg(&hd)
19045            .arg(&nh)
19046            .arg(base_dev)
19047            .arg(&plus0)
19048            .arg(&nspm)
19049            .arg(&spk);
19050        unsafe {
19051            b2.launch(cfg2)?;
19052        }
19053        Ok(())
19054    }
19055
19056    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
19057    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
19058    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
19059    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
19060    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
19061    ///
19062    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
19063    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
19064    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
19065    /// grouping (different but mathematically-equal log-sum-exp merge).
19066    pub fn fa_decode_dc(
19067        &self,
19068        q: &CudaSlice<f32>,
19069        k: &cudarc::driver::CudaView<u8>,
19070        v: &cudarc::driver::CudaView<u8>,
19071        o: &mut CudaSlice<f32>,
19072        head_dim: usize,
19073        n_head: usize,
19074        n_head_kv: usize,
19075        t_kv_dev: &CudaSlice<i32>,
19076        bucket_max: usize,
19077        scale: f32,
19078        k_tok_bytes: usize,
19079        v_tok_bytes: usize,
19080        g: bool,
19081    ) -> Result<(), Box<dyn std::error::Error>> {
19082        self.fa_decode_dc_q8(
19083            q,
19084            k,
19085            v,
19086            o,
19087            head_dim,
19088            n_head,
19089            n_head_kv,
19090            t_kv_dev,
19091            bucket_max,
19092            scale,
19093            k_tok_bytes,
19094            v_tok_bytes,
19095            g,
19096            None,
19097        )
19098    }
19099
19100    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
19101    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
19102    #[allow(clippy::too_many_arguments)]
19103    pub fn fa_decode_dc_q8(
19104        &self,
19105        q: &CudaSlice<f32>,
19106        k: &cudarc::driver::CudaView<u8>,
19107        v: &cudarc::driver::CudaView<u8>,
19108        o: &mut CudaSlice<f32>,
19109        head_dim: usize,
19110        n_head: usize,
19111        n_head_kv: usize,
19112        t_kv_dev: &CudaSlice<i32>,
19113        bucket_max: usize,
19114        scale: f32,
19115        k_tok_bytes: usize,
19116        v_tok_bytes: usize,
19117        g: bool,
19118        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19119    ) -> Result<(), Box<dyn std::error::Error>> {
19120        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
19121        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
19122        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
19123        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
19124        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
19125        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
19126        // 2026-07-12).
19127        let mut fa_vec =
19128            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19129        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19130            fa_vec = false;
19131        } // mirror kvmod/geom
19132        let sp = fa_split_keys(bucket_max, n_head_kv);
19133        let n_splits = if fa_vec {
19134            ((bucket_max + sp - 1) / sp).max(1)
19135        } else {
19136            ((bucket_max + 255) / 256).max(1)
19137        };
19138        let o_len = n_head * n_splits * head_dim;
19139        let ml_len = n_head * n_splits;
19140        let mut part_guard = self.fa_part_pool.lock().unwrap();
19141        if part_guard
19142            .as_ref()
19143            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19144            .unwrap_or(true)
19145        {
19146            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19147            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19148            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19149            // later live allocations land at those addresses, and the next graph REPLAY writes
19150            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19151            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19152            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19153            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19154            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19155            // (total retired < final size).
19156            let old = part_guard.take();
19157            let (co, cm) = old
19158                .as_ref()
19159                .map(|pp| (pp.0.len(), pp.1.len()))
19160                .unwrap_or((0, 0));
19161            if let Some(old) = old {
19162                self.fa_part_retired.lock().unwrap().push(old);
19163            }
19164            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19165                eprintln!(
19166                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19167                    co, o_len, cm, ml_len
19168                );
19169            }
19170            *part_guard = Some((
19171                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19172                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19173                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19174            ));
19175        }
19176        let pg = part_guard.as_mut().unwrap();
19177        self.gpu
19178            .stream()
19179            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19180        self.gpu
19181            .stream()
19182            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19183        self.gpu
19184            .stream()
19185            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19186        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19187        let (hd, nh, nhkv, nsp) = (
19188            head_dim as i32,
19189            n_head as i32,
19190            n_head_kv as i32,
19191            n_splits as i32,
19192        );
19193        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19194        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19195        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
19196        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
19197        let deep = fa_vec
19198            && head_dim == 256
19199            && fa_v4_at(bucket_max)
19200            && !g
19201            && fa_deep_at(bucket_max)
19202            && !matches!(fa_v4_mode(), "noB3" | "stage");
19203        let (f, cfg) = if fa_vec
19204            && head_dim == 512
19205            && bucket_max >= {
19206                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19207                *FA512_MIN_DC.get_or_init(|| {
19208                    std::env::var("MEMRA_FA512_MIN")
19209                        .ok()
19210                        .and_then(|v| v.parse().ok())
19211                        .unwrap_or(512)
19212                })
19213            } {
19214            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
19215            let gqa = (n_head / n_head_kv).max(1) as u32;
19216            (
19217                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19218                LaunchConfig {
19219                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19220                    block_dim: (32, gqa, 1),
19221                    shared_mem_bytes: 0,
19222                },
19223            )
19224        } else if fa_vec && head_dim == 512 {
19225            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
19226            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
19227            let q_view = q.as_view();
19228            let mut o_view = o.as_view_mut();
19229            return self.fa_decode_scalar_unified(
19230                &q_view,
19231                k,
19232                v,
19233                &mut o_view,
19234                head_dim,
19235                n_head,
19236                n_head_kv,
19237                0,
19238                Some(t_kv_dev),
19239                scale,
19240                n_splits,
19241                sp,
19242                k_tok_bytes,
19243                v_tok_bytes,
19244                g,
19245                &mut *part_o,
19246                &mut *part_m,
19247                &mut *part_l,
19248                q8_out,
19249            );
19250        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19251            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
19252            // incl the g-module route + raw-e4m3 sV sizing.
19253            let gqa = (n_head / n_head_kv).max(1) as u32;
19254            let fv = if g {
19255                self.func_g("fa_decode_vec_q_v4_dc")
19256            } else if deep {
19257                self.func("fa_decode_vec_q_v4_deep_dc")
19258            } else {
19259                self.func("fa_decode_vec_q_v4_dc")
19260            };
19261            let shmem =
19262                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19263            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19264            fv.set_attribute(
19265                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19266                shmem as i32,
19267            )?;
19268            (
19269                fv,
19270                LaunchConfig {
19271                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19272                    block_dim: (32, gqa, 1),
19273                    shared_mem_bytes: shmem,
19274                },
19275            )
19276        } else if fa_vec && fa_v3_active(head_dim) {
19277            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
19278            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
19279            let gqa = (n_head / n_head_kv).max(1) as u32;
19280            let fv = if g {
19281                self.func_g("fa_decode_vec_q_v3_dc")
19282            } else {
19283                self.func("fa_decode_vec_q_v3_dc")
19284            };
19285            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
19286            (
19287                fv,
19288                LaunchConfig {
19289                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19290                    block_dim: (32, gqa, 1),
19291                    shared_mem_bytes: shmem,
19292                },
19293            )
19294        } else if fa_vec && fa_v2_on() {
19295            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
19296            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
19297            // a numeric config; eager, rows-verify and graph all switch together).
19298            let gqa = (n_head / n_head_kv).max(1) as u32;
19299            let fv = if g {
19300                self.func_g("fa_decode_vec_q_v2_dc")
19301            } else {
19302                self.func("fa_decode_vec_q_v2_dc")
19303            };
19304            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
19305            (
19306                fv,
19307                LaunchConfig {
19308                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19309                    block_dim: (32, gqa, 1),
19310                    shared_mem_bytes: shmem,
19311                },
19312            )
19313        } else if fa_vec {
19314            let gqa = (n_head / n_head_kv).max(1) as u32;
19315            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
19316            let fv = if g {
19317                self.func_g("fa_decode_vec_q_dc")
19318            } else {
19319                self.func("fa_decode_vec_q_dc")
19320            };
19321            (
19322                fv,
19323                LaunchConfig {
19324                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19325                    block_dim: (32, gqa, 1),
19326                    shared_mem_bytes: 0,
19327                },
19328            )
19329        } else {
19330            let q_view = q.as_view();
19331            let mut o_view = o.as_view_mut();
19332            return self.fa_decode_scalar_unified(
19333                &q_view,
19334                k,
19335                v,
19336                &mut o_view,
19337                head_dim,
19338                n_head,
19339                n_head_kv,
19340                0,
19341                Some(t_kv_dev),
19342                scale,
19343                n_splits,
19344                if fa_vec { sp } else { 256 },
19345                k_tok_bytes,
19346                v_tok_bytes,
19347                g,
19348                &mut *part_o,
19349                &mut *part_m,
19350                &mut *part_l,
19351                q8_out,
19352            );
19353        };
19354        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
19355        let __s_b = self.gpu.stream();
19356        let mut b = __s_b.launch_builder(&f);
19357        b.arg(q)
19358            .arg(k)
19359            .arg(v)
19360            .arg(&mut *part_o)
19361            .arg(&mut *part_m)
19362            .arg(&mut *part_l)
19363            .arg(&hd)
19364            .arg(&nh)
19365            .arg(&nhkv)
19366            .arg(t_kv_dev)
19367            .arg(&scale)
19368            .arg(&nsp)
19369            .arg(&ski)
19370            .arg(&ktb)
19371            .arg(&vtb);
19372        unsafe {
19373            b.launch(cfg)?;
19374        }
19375        let cfg2 = LaunchConfig {
19376            grid_dim: (n_head as u32, 1, 1),
19377            block_dim: (head_dim as u32, 1, 1),
19378            shared_mem_bytes: 0,
19379        };
19380        if let Some((oq, od)) = q8_out {
19381            let fc = if g {
19382                self.func_g("fa_decode_combine_q8_1")
19383            } else {
19384                self.fa_func("fa_decode_combine_q8_1", head_dim)
19385            };
19386            let __s_b2 = self.gpu.stream();
19387            let mut b2 = __s_b2.launch_builder(&fc);
19388            b2.arg(&*part_o)
19389                .arg(&*part_m)
19390                .arg(&*part_l)
19391                .arg(oq)
19392                .arg(od)
19393                .arg(&hd)
19394                .arg(&nh)
19395                .arg(&nsp);
19396            unsafe {
19397                b2.launch(cfg2)?;
19398            }
19399            return Ok(());
19400        }
19401        let fc = if g {
19402            self.func_g("fa_decode_combine_f32")
19403        } else {
19404            self.fa_func("fa_decode_combine_f32", head_dim)
19405        };
19406        let __s_b2 = self.gpu.stream();
19407        let mut b2 = __s_b2.launch_builder(&fc);
19408        b2.arg(&*part_o)
19409            .arg(&*part_m)
19410            .arg(&*part_l)
19411            .arg(o)
19412            .arg(&hd)
19413            .arg(&nh)
19414            .arg(&nsp);
19415        unsafe {
19416            b2.launch(cfg2)?;
19417        }
19418        Ok(())
19419    }
19420
19421    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
19422    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
19423    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
19424    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
19425    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
19426    pub fn fa_geom_eager(
19427        &self,
19428        t_kv: usize,
19429        head_dim: usize,
19430        n_head_kv: usize,
19431        g: bool,
19432    ) -> (bool, usize) {
19433        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
19434        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
19435        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
19436        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19437        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
19438        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
19439        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
19440        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
19441        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
19442        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19443        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19444        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
19445        // family; everything else falls to the g-module scalar.
19446        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
19447        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
19448        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
19449        if g && head_dim == 256 && !fa_v4_at(t_kv) {
19450            fa_vec = false;
19451        }
19452        let sp = fa_split_keys(t_kv, n_head_kv);
19453        let n_splits = if fa_vec {
19454            ((t_kv + sp - 1) / sp).max(1)
19455        } else {
19456            ((t_kv + 255) / 256).max(1)
19457        };
19458        (fa_vec, n_splits)
19459    }
19460
19461    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
19462    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
19463    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
19464    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
19465    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
19466    pub fn fa_bucket_key(
19467        &self,
19468        t_kv: usize,
19469        head_dim: usize,
19470        n_head_kv: usize,
19471        g: bool,
19472    ) -> (bool, usize) {
19473        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19474    }
19475
19476    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
19477    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
19478    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
19479    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
19480    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
19481    /// device data) — every per-step varying scalar must come from a device counter. Returns the
19482    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
19483    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
19484    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
19485    /// replays (transients returning to the pool get reused by unrelated work and corrupt
19486    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
19487    pub fn capture_graph_retained<F>(
19488        &self,
19489        step: F,
19490    ) -> Result<
19491        (
19492            cudarc::driver::CudaGraph,
19493            Vec<Box<dyn std::any::Any + Send>>,
19494        ),
19495        Box<dyn std::error::Error>,
19496    >
19497    where
19498        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19499    {
19500        use cudarc::driver::sys::CUgraphInstantiate_flags;
19501        self.capture_graph_retained_flags(
19502            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19503            step,
19504        )
19505    }
19506
19507    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
19508    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
19509    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
19510    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
19511    pub fn capture_graph_retained_flags<F>(
19512        &self,
19513        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19514        mut step: F,
19515    ) -> Result<
19516        (
19517            cudarc::driver::CudaGraph,
19518            Vec<Box<dyn std::any::Any + Send>>,
19519        ),
19520        Box<dyn std::error::Error>,
19521    >
19522    where
19523        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19524    {
19525        use cudarc::driver::sys::CUstreamCaptureMode;
19526        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
19527        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
19528        // while the capture region is open become dead copy NODES replayed every launch
19529        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
19530        // warmup runs allocate the same transient sequence at the same pool addresses, so
19531        // retaining the warmup clones preserves the draft-graph fix without polluting the
19532        // captured graph.
19533        self.capture_keep.lock().unwrap().clear();
19534        let was_tracking = self.gpu.ctx.is_event_tracking();
19535        if was_tracking {
19536            unsafe {
19537                self.gpu.ctx.disable_event_tracking();
19538            }
19539        }
19540        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19541            self.capture_keep_on
19542                .store(true, std::sync::atomic::Ordering::Relaxed);
19543            let w = (|| {
19544                step(self)?;
19545                step(self)
19546            })();
19547            self.capture_keep_on
19548                .store(false, std::sync::atomic::Ordering::Relaxed);
19549            w?;
19550            self.gpu.stream().synchronize()?;
19551            self.gpu
19552                .stream()
19553                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19554            let r = step(self);
19555            let g = self.gpu.stream().end_capture(flags);
19556            r?;
19557            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19558            graph.upload()?;
19559            Ok(graph)
19560        };
19561        let result = run();
19562        self.capture_keep_on
19563            .store(false, std::sync::atomic::Ordering::Relaxed);
19564        if was_tracking {
19565            unsafe {
19566                self.gpu.ctx.enable_event_tracking();
19567            }
19568        }
19569        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19570        Ok((result?, keeper))
19571    }
19572
19573    pub fn capture_graph<F>(
19574        &self,
19575        mut step: F,
19576    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19577    where
19578        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19579    {
19580        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19581        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
19582        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
19583        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
19584        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
19585        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
19586        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
19587        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
19588        let was_tracking = self.gpu.ctx.is_event_tracking();
19589        if was_tracking {
19590            unsafe {
19591                self.gpu.ctx.disable_event_tracking();
19592            }
19593        }
19594        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
19595        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
19596        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
19597        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
19598        // measure that scan's real cost on the generic path. Diagnostic door only; the
19599        // default stays AUTO_FREE until a measured A/B justifies moving it.
19600        let iflag = {
19601            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19602            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19603                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
19604                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
19605                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19606                Ok("priority") => {
19607                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19608                }
19609                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19610            })
19611        };
19612        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
19613        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
19614        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
19615        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
19616        // eager step executions and are node-count-invariant. Printing the split bounds the
19617        // refactor's ceiling instead of assuming it.
19618        let ct = {
19619            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19620            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19621        };
19622        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
19623        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
19624        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
19625        // chased, and node-count-invariant, so no capture-body refactor could touch it.
19626        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
19627        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
19628        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
19629        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
19630        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
19631        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
19632        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
19633        // grow and never frees, resident counters/scratch, cache set in place), and the
19634        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
19635        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
19636        // settling and pool mapping. Arbitrated adversarially, not by taste:
19637        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
19638        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
19639        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
19640        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
19641        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
19642        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
19643        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
19644        let warmups = {
19645            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19646            *W.get_or_init(|| {
19647                std::env::var("MEMRA_GRAPH_WARMUPS")
19648                    .ok()
19649                    .and_then(|v| v.parse().ok())
19650                    .filter(|n| *n >= 1)
19651                    .unwrap_or(1)
19652            })
19653        };
19654        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19655            let t_w = std::time::Instant::now();
19656            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
19657            for _ in 0..warmups {
19658                step(self)?;
19659            }
19660            self.gpu.stream().synchronize()?;
19661            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
19662            // capture the third run.
19663            let t_c = std::time::Instant::now();
19664            self.gpu
19665                .stream()
19666                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19667            // If the body errors mid-capture, end the capture before propagating so the stream isn't
19668            // left in a capturing state.
19669            let r = step(self);
19670            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
19671            let t_i = std::time::Instant::now();
19672            let g = self.gpu.stream().end_capture(iflag);
19673            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
19674            r?;
19675            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19676            let t_u = std::time::Instant::now();
19677            graph.upload()?;
19678            if ct {
19679                println!(
19680                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
19681                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
19682                    t_u.elapsed().as_secs_f64() * 1e3
19683                );
19684            }
19685            Ok(graph)
19686        };
19687        let result = run();
19688        if was_tracking {
19689            unsafe {
19690                self.gpu.ctx.enable_event_tracking();
19691            }
19692        }
19693        result
19694    }
19695
19696    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
19697    pub fn gdn_scan_s128_view(
19698        &self,
19699        q: &CudaSlice<f32>,
19700        k: &CudaSlice<f32>,
19701        v: &CudaSlice<f32>,
19702        g: &CudaSlice<f32>,
19703        beta: &CudaSlice<f32>,
19704        state_in: &cudarc::driver::CudaView<f32>,
19705        state_out: &mut cudarc::driver::CudaViewMut<f32>,
19706        o: &mut CudaSlice<f32>,
19707        n_head: usize,
19708        t: usize,
19709        scale: f32,
19710    ) -> Result<(), Box<dyn std::error::Error>> {
19711        let f = self.func("gdn_scan_s128");
19712        const S_V: u32 = 128;
19713        const WARP: u32 = 32;
19714        const COLS: u32 = 4;
19715        let cfg = LaunchConfig {
19716            grid_dim: (n_head as u32, 1, S_V / COLS),
19717            block_dim: (WARP, COLS, 1),
19718            shared_mem_bytes: 0,
19719        };
19720        let (h, ti) = (n_head as i32, t as i32);
19721        let __s_b = self.gpu.stream();
19722        let mut b = __s_b.launch_builder(&f);
19723        b.arg(q)
19724            .arg(k)
19725            .arg(v)
19726            .arg(g)
19727            .arg(beta)
19728            .arg(state_in)
19729            .arg(state_out)
19730            .arg(o)
19731            .arg(&h)
19732            .arg(&ti)
19733            .arg(&scale);
19734        unsafe {
19735            b.launch(cfg)?;
19736        }
19737        Ok(())
19738    }
19739
19740    /// conv1d where the input is a CudaView (resident conv state assembled in place).
19741    pub fn ssm_conv1d_view(
19742        &self,
19743        x: &cudarc::driver::CudaView<f32>,
19744        w: &CudaSlice<f32>,
19745        y: &mut CudaSlice<f32>,
19746        conv_dim: usize,
19747        t: usize,
19748        d_conv: usize,
19749        silu: bool,
19750    ) -> Result<(), Box<dyn std::error::Error>> {
19751        let f = self.func("ssm_conv1d_silu_f32");
19752        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
19753        let cfg = LaunchConfig {
19754            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
19755            block_dim: (256, 1, 1),
19756            shared_mem_bytes: 0,
19757        };
19758        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
19759        let __s_b = self.gpu.stream();
19760        let mut b = __s_b.launch_builder(&f);
19761        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
19762        unsafe {
19763            b.launch(cfg)?;
19764        }
19765        Ok(())
19766    }
19767
19768    /// Depthwise causal conv1d + optional SiLU.
19769    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
19770    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
19771    /// FUSED prefill conv (token-major input, zero left-state): replaces
19772    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
19773    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
19774    pub fn ssm_conv1d_tm(
19775        &self,
19776        qkv_tm: &CudaSlice<f32>,
19777        w: &CudaSlice<f32>,
19778        y: &mut CudaSlice<f32>,
19779        conv_dim: usize,
19780        t: usize,
19781        d_conv: usize,
19782    ) -> Result<(), Box<dyn std::error::Error>> {
19783        let f = self.func("ssm_conv1d_tm_f32");
19784        let cfg = LaunchConfig {
19785            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19786            block_dim: (256, 1, 1),
19787            shared_mem_bytes: 0,
19788        };
19789        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19790        let __s_b = self.gpu.stream();
19791        let mut b = __s_b.launch_builder(&f);
19792        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
19793        unsafe {
19794            b.launch(cfg)?;
19795        }
19796        Ok(())
19797    }
19798
19799    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
19800    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
19801    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
19802    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
19803    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
19804    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
19805    /// columns; the final ring == what T sequential decode ring rolls leave).
19806    pub fn ssm_conv1d_tm_state(
19807        &self,
19808        qkv_tm: &CudaSlice<f32>,
19809        conv_state: &mut CudaSlice<f32>,
19810        w: &CudaSlice<f32>,
19811        y: &mut CudaSlice<f32>,
19812        conv_dim: usize,
19813        t: usize,
19814        d_conv: usize,
19815    ) -> Result<(), Box<dyn std::error::Error>> {
19816        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
19817    }
19818
19819    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
19820    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
19821    #[allow(clippy::too_many_arguments)]
19822    pub fn ssm_conv1d_tm_state_pad(
19823        &self,
19824        qkv_tm: &CudaSlice<f32>,
19825        conv_state: &mut CudaSlice<f32>,
19826        w: &CudaSlice<f32>,
19827        y: &mut CudaSlice<f32>,
19828        conv_dim: usize,
19829        t: usize,
19830        d_conv: usize,
19831        pad_len: Option<&CudaSlice<i32>>,
19832    ) -> Result<(), Box<dyn std::error::Error>> {
19833        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
19834        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
19835        // the window kernel both read the pre-roll ring; the roll launches after both) — but
19836        // cloning first keeps the ordering trivially correct under any future stream split.
19837        let ring_old = if t < d_conv - 1 {
19838            Some(self.clone_dtod(conv_state)?)
19839        } else {
19840            None
19841        };
19842        {
19843            let f = self.func("ssm_conv1d_tm_state_f32");
19844            let cfg = LaunchConfig {
19845                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19846                block_dim: (256, 1, 1),
19847                shared_mem_bytes: 0,
19848            };
19849            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19850            let __s_b = self.gpu.stream();
19851            let mut b = __s_b.launch_builder(&f);
19852            b.arg(qkv_tm)
19853                .arg(&*conv_state)
19854                .arg(w)
19855                .arg(y)
19856                .arg(&cd)
19857                .arg(&ti)
19858                .arg(&dc);
19859            unsafe {
19860                b.launch(cfg)?;
19861            }
19862        }
19863        match (ring_old, pad_len) {
19864            (None, Some(len_d)) => {
19865                let f = self.func("ssm_conv_ring_update_dev_f32");
19866                let n = conv_dim * (d_conv - 1);
19867                let cfg = LaunchConfig::for_num_elems(n as u32);
19868                let (cd, dc) = (conv_dim as i32, d_conv as i32);
19869                let __s_b = self.gpu.stream();
19870                let mut b = __s_b.launch_builder(&f);
19871                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
19872                unsafe {
19873                    b.launch(cfg)?;
19874                }
19875            }
19876            (None, None) => {
19877                let f = self.func("ssm_conv_ring_update_f32");
19878                let n = conv_dim * (d_conv - 1);
19879                let cfg = LaunchConfig::for_num_elems(n as u32);
19880                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19881                let __s_b = self.gpu.stream();
19882                let mut b = __s_b.launch_builder(&f);
19883                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
19884                unsafe {
19885                    b.launch(cfg)?;
19886                }
19887            }
19888            (Some(old), _) => {
19889                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
19890            }
19891        }
19892        Ok(())
19893    }
19894
19895    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
19896    pub fn ssm_conv1d_tm_state_pad_v(
19897        &self,
19898        qkv_tm: &cudarc::driver::CudaView<f32>,
19899        conv_state: &mut CudaSlice<f32>,
19900        w: &CudaSlice<f32>,
19901        y: &mut CudaSlice<f32>,
19902        conv_dim: usize,
19903        t: usize,
19904        d_conv: usize,
19905        pad_len: Option<&CudaSlice<i32>>,
19906    ) -> Result<(), Box<dyn std::error::Error>> {
19907        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
19908        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
19909        // the window kernel both read the pre-roll ring; the roll launches after both) — but
19910        // cloning first keeps the ordering trivially correct under any future stream split.
19911        let ring_old = if t < d_conv - 1 {
19912            Some(self.clone_dtod(conv_state)?)
19913        } else {
19914            None
19915        };
19916        {
19917            let f = self.func("ssm_conv1d_tm_state_f32");
19918            let cfg = LaunchConfig {
19919                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19920                block_dim: (256, 1, 1),
19921                shared_mem_bytes: 0,
19922            };
19923            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19924            let __s_b = self.gpu.stream();
19925            let mut b = __s_b.launch_builder(&f);
19926            b.arg(qkv_tm)
19927                .arg(&*conv_state)
19928                .arg(w)
19929                .arg(y)
19930                .arg(&cd)
19931                .arg(&ti)
19932                .arg(&dc);
19933            unsafe {
19934                b.launch(cfg)?;
19935            }
19936        }
19937        match (ring_old, pad_len) {
19938            (None, Some(len_d)) => {
19939                let f = self.func("ssm_conv_ring_update_dev_f32");
19940                let n = conv_dim * (d_conv - 1);
19941                let cfg = LaunchConfig::for_num_elems(n as u32);
19942                let (cd, dc) = (conv_dim as i32, d_conv as i32);
19943                let __s_b = self.gpu.stream();
19944                let mut b = __s_b.launch_builder(&f);
19945                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
19946                unsafe {
19947                    b.launch(cfg)?;
19948                }
19949            }
19950            (None, None) => {
19951                let f = self.func("ssm_conv_ring_update_f32");
19952                let n = conv_dim * (d_conv - 1);
19953                let cfg = LaunchConfig::for_num_elems(n as u32);
19954                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19955                let __s_b = self.gpu.stream();
19956                let mut b = __s_b.launch_builder(&f);
19957                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
19958                unsafe {
19959                    b.launch(cfg)?;
19960                }
19961            }
19962            (Some(_), _) => unreachable!(
19963                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
19964            ),
19965        }
19966        Ok(())
19967    }
19968
19969    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
19970    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
19971    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
19972    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
19973    pub fn ssm_conv_ring_rebuild(
19974        &self,
19975        qkv_tm: &CudaSlice<f32>,
19976        ring_old: &CudaSlice<f32>,
19977        conv_state: &mut CudaSlice<f32>,
19978        conv_dim: usize,
19979        tc: usize,
19980        d_conv: usize,
19981    ) -> Result<(), Box<dyn std::error::Error>> {
19982        let f = self.func("ssm_conv_ring_rebuild_f32");
19983        let n = conv_dim * (d_conv - 1);
19984        let cfg = LaunchConfig::for_num_elems(n as u32);
19985        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
19986        let __s_b = self.gpu.stream();
19987        let mut b = __s_b.launch_builder(&f);
19988        b.arg(qkv_tm)
19989            .arg(ring_old)
19990            .arg(conv_state)
19991            .arg(&cd)
19992            .arg(&ti)
19993            .arg(&dc);
19994        unsafe {
19995            b.launch(cfg)?;
19996        }
19997        Ok(())
19998    }
19999
20000    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
20001    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
20002    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
20003    /// the argmax + run-spec gates are the authority.
20004    #[allow(clippy::too_many_arguments)]
20005    pub fn gdn_prep_decode(
20006        &self,
20007        conv_out: &CudaSlice<f32>,
20008        beta_raw: &CudaSlice<f32>,
20009        alpha: &CudaSlice<f32>,
20010        dt_bias: &CudaSlice<f32>,
20011        a: &CudaSlice<f32>,
20012        q_l2: &mut CudaSlice<f32>,
20013        k_l2: &mut CudaSlice<f32>,
20014        v_g: &mut CudaSlice<f32>,
20015        beta: &mut CudaSlice<f32>,
20016        g_log: &mut CudaSlice<f32>,
20017        d_state: usize,
20018        num_v: usize,
20019        num_k: usize,
20020        key_dim: usize,
20021        eps: f32,
20022    ) -> Result<(), Box<dyn std::error::Error>> {
20023        let f = self.func("gdn_prep_decode_f32");
20024        let cfg = LaunchConfig {
20025            grid_dim: (num_v as u32, 1, 1),
20026            block_dim: (32, 4, 1),
20027            shared_mem_bytes: 0,
20028        };
20029        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20030        let __s_b = self.gpu.stream();
20031        let mut b = __s_b.launch_builder(&f);
20032        b.arg(conv_out)
20033            .arg(beta_raw)
20034            .arg(alpha)
20035            .arg(dt_bias)
20036            .arg(a)
20037            .arg(q_l2)
20038            .arg(k_l2)
20039            .arg(v_g)
20040            .arg(beta)
20041            .arg(g_log)
20042            .arg(&ds)
20043            .arg(&nv)
20044            .arg(&nk)
20045            .arg(&kd)
20046            .arg(&eps);
20047        unsafe {
20048            b.launch(cfg)?;
20049        }
20050        Ok(())
20051    }
20052
20053    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
20054    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
20055    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
20056    #[allow(clippy::too_many_arguments)]
20057    pub fn ssm_conv1d_gdn(
20058        &self,
20059        qkv_tm: &CudaSlice<f32>,
20060        w: &CudaSlice<f32>,
20061        q_g: &mut CudaSlice<f32>,
20062        k_g: &mut CudaSlice<f32>,
20063        v_g: &mut CudaSlice<f32>,
20064        conv_dim: usize,
20065        t: usize,
20066        d_conv: usize,
20067        d_state: usize,
20068        num_v: usize,
20069        num_k: usize,
20070        key_dim: usize,
20071    ) -> Result<(), Box<dyn std::error::Error>> {
20072        let f = self.func("ssm_conv1d_gdn_f32");
20073        let cfg = LaunchConfig {
20074            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20075            block_dim: (256, 1, 1),
20076            shared_mem_bytes: 0,
20077        };
20078        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20079        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20080        let __s_b = self.gpu.stream();
20081        let mut b = __s_b.launch_builder(&f);
20082        b.arg(qkv_tm)
20083            .arg(w)
20084            .arg(q_g)
20085            .arg(k_g)
20086            .arg(v_g)
20087            .arg(&cd)
20088            .arg(&ti)
20089            .arg(&dc)
20090            .arg(&ds)
20091            .arg(&nv)
20092            .arg(&nk)
20093            .arg(&kd);
20094        unsafe {
20095            b.launch(cfg)?;
20096        }
20097        Ok(())
20098    }
20099
20100    pub fn ssm_conv1d(
20101        &self,
20102        x: &CudaSlice<f32>,
20103        w: &CudaSlice<f32>,
20104        y: &mut CudaSlice<f32>,
20105        conv_dim: usize,
20106        t: usize,
20107        d_conv: usize,
20108        silu: bool,
20109    ) -> Result<(), Box<dyn std::error::Error>> {
20110        let f = self.func("ssm_conv1d_silu_f32");
20111        let cfg = LaunchConfig {
20112            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20113            block_dim: (256, 1, 1),
20114            shared_mem_bytes: 0,
20115        };
20116        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20117        let __s_b = self.gpu.stream();
20118        let mut b = __s_b.launch_builder(&f);
20119        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20120        unsafe {
20121            b.launch(cfg)?;
20122        }
20123        Ok(())
20124    }
20125
20126    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
20127    /// o:[128,H,T]. Single sequence.
20128    pub fn gdn_scan_s128(
20129        &self,
20130        q: &CudaSlice<f32>,
20131        k: &CudaSlice<f32>,
20132        v: &CudaSlice<f32>,
20133        g: &CudaSlice<f32>,
20134        beta: &CudaSlice<f32>,
20135        state_in: &CudaSlice<f32>,
20136        state_out: &mut CudaSlice<f32>,
20137        o: &mut CudaSlice<f32>,
20138        n_head: usize,
20139        t: usize,
20140        scale: f32,
20141    ) -> Result<(), Box<dyn std::error::Error>> {
20142        let f = self.func("gdn_scan_s128");
20143        const S_V: u32 = 128;
20144        const WARP: u32 = 32;
20145        const COLS_PER_BLOCK: u32 = 4;
20146        let cfg = LaunchConfig {
20147            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20148            block_dim: (WARP, COLS_PER_BLOCK, 1),
20149            shared_mem_bytes: 0,
20150        };
20151        let (h, ti) = (n_head as i32, t as i32);
20152        let __s_b = self.gpu.stream();
20153        let mut b = __s_b.launch_builder(&f);
20154        b.arg(q)
20155            .arg(k)
20156            .arg(v)
20157            .arg(g)
20158            .arg(beta)
20159            .arg(state_in)
20160            .arg(state_out)
20161            .arg(o)
20162            .arg(&h)
20163            .arg(&ti)
20164            .arg(&scale);
20165        unsafe {
20166            b.launch(cfg)?;
20167        }
20168        Ok(())
20169    }
20170
20171    // ==== B2' batched decode state ops (decode_batch.rs) ====
20172    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
20173    // Bodies are the single-seq kernels per sequence — bit-identical per row.
20174
20175    #[allow(clippy::too_many_arguments)]
20176    pub fn ssm_conv1d_fused_decode_b(
20177        &self,
20178        qkv_cols: &CudaSlice<f32>,
20179        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20180        w: &CudaSlice<f32>,
20181        conv_outs: &mut CudaSlice<f32>,
20182        conv_dim: usize,
20183        d_conv: usize,
20184        b_n: usize,
20185    ) -> Result<(), Box<dyn std::error::Error>> {
20186        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20187        let cfg = LaunchConfig {
20188            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20189            block_dim: (256, 1, 1),
20190            shared_mem_bytes: 0,
20191        };
20192        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20193        let __s_b = self.gpu.stream();
20194        let mut b = __s_b.launch_builder(&f);
20195        b.arg(qkv_cols)
20196            .arg(conv_state_ptrs)
20197            .arg(w)
20198            .arg(conv_outs)
20199            .arg(&cd)
20200            .arg(&dc);
20201        unsafe {
20202            b.launch(cfg)?;
20203        }
20204        Ok(())
20205    }
20206
20207    #[allow(clippy::too_many_arguments)]
20208    pub fn gdn_prep_decode_b(
20209        &self,
20210        conv_outs: &CudaSlice<f32>,
20211        beta_raws: &CudaSlice<f32>,
20212        alphas: &CudaSlice<f32>,
20213        dt_bias: &CudaSlice<f32>,
20214        a: &CudaSlice<f32>,
20215        q_l2: &mut CudaSlice<f32>,
20216        k_l2: &mut CudaSlice<f32>,
20217        v_g: &mut CudaSlice<f32>,
20218        beta: &mut CudaSlice<f32>,
20219        g_log: &mut CudaSlice<f32>,
20220        d_state: usize,
20221        num_v: usize,
20222        num_k: usize,
20223        key_dim: usize,
20224        eps: f32,
20225        conv_dim: usize,
20226        b_n: usize,
20227    ) -> Result<(), Box<dyn std::error::Error>> {
20228        let f = self.func("gdn_prep_decode_b_f32");
20229        let cfg = LaunchConfig {
20230            grid_dim: (num_v as u32, 1, b_n as u32),
20231            block_dim: (32, 4, 1),
20232            shared_mem_bytes: 0,
20233        };
20234        let (ds, nv, nk, kd, cd) = (
20235            d_state as i32,
20236            num_v as i32,
20237            num_k as i32,
20238            key_dim as i32,
20239            conv_dim as i32,
20240        );
20241        let __s_b = self.gpu.stream();
20242        let mut b = __s_b.launch_builder(&f);
20243        b.arg(conv_outs)
20244            .arg(beta_raws)
20245            .arg(alphas)
20246            .arg(dt_bias)
20247            .arg(a)
20248            .arg(q_l2)
20249            .arg(k_l2)
20250            .arg(v_g)
20251            .arg(beta)
20252            .arg(g_log)
20253            .arg(&ds)
20254            .arg(&nv)
20255            .arg(&nk)
20256            .arg(&kd)
20257            .arg(&eps)
20258            .arg(&cd);
20259        unsafe {
20260            b.launch(cfg)?;
20261        }
20262        Ok(())
20263    }
20264
20265    #[allow(clippy::too_many_arguments)]
20266    pub fn gdn_scan_s128_batched(
20267        &self,
20268        q: &CudaSlice<f32>,
20269        k: &CudaSlice<f32>,
20270        v: &CudaSlice<f32>,
20271        g: &CudaSlice<f32>,
20272        beta: &CudaSlice<f32>,
20273        state_in_ptrs: &cudarc::driver::CudaView<u64>,
20274        state_out_ptrs: &cudarc::driver::CudaView<u64>,
20275        o: &mut CudaSlice<f32>,
20276        n_head: usize,
20277        b_n: usize,
20278        scale: f32,
20279    ) -> Result<(), Box<dyn std::error::Error>> {
20280        let f = self.func("gdn_scan_s128_b");
20281        const S_V: u32 = 128;
20282        const WARP: u32 = 32;
20283        const COLS_PER_BLOCK: u32 = 4;
20284        let cfg = LaunchConfig {
20285            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20286            block_dim: (WARP, COLS_PER_BLOCK, 1),
20287            shared_mem_bytes: 0,
20288        };
20289        let h = n_head as i32;
20290        let __s_b = self.gpu.stream();
20291        let mut b = __s_b.launch_builder(&f);
20292        b.arg(q)
20293            .arg(k)
20294            .arg(v)
20295            .arg(g)
20296            .arg(beta)
20297            .arg(state_in_ptrs)
20298            .arg(state_out_ptrs)
20299            .arg(o)
20300            .arg(&h)
20301            .arg(&scale);
20302        unsafe {
20303            b.launch(cfg)?;
20304        }
20305        Ok(())
20306    }
20307
20308    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
20309    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
20310    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
20311    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
20312    /// numeric class; only the pointer arithmetic moved host-side.
20313    #[allow(clippy::too_many_arguments)]
20314    pub fn ssm_conv1d_fused_decode_b_view(
20315        &self,
20316        qkv_cols: &cudarc::driver::CudaView<f32>,
20317        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20318        w: &CudaSlice<f32>,
20319        conv_outs: &mut CudaSlice<f32>,
20320        conv_dim: usize,
20321        d_conv: usize,
20322        b_n: usize,
20323    ) -> Result<(), Box<dyn std::error::Error>> {
20324        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20325        let cfg = LaunchConfig {
20326            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20327            block_dim: (256, 1, 1),
20328            shared_mem_bytes: 0,
20329        };
20330        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20331        let __s_b = self.gpu.stream();
20332        let mut b = __s_b.launch_builder(&f);
20333        b.arg(qkv_cols)
20334            .arg(conv_state_ptrs)
20335            .arg(w)
20336            .arg(conv_outs)
20337            .arg(&cd)
20338            .arg(&dc);
20339        unsafe {
20340            b.launch(cfg)?;
20341        }
20342        Ok(())
20343    }
20344
20345    #[allow(clippy::too_many_arguments)]
20346    pub fn gdn_prep_decode_b_view(
20347        &self,
20348        conv_outs: &CudaSlice<f32>,
20349        beta_raws: &cudarc::driver::CudaView<f32>,
20350        alphas: &cudarc::driver::CudaView<f32>,
20351        dt_bias: &CudaSlice<f32>,
20352        a: &CudaSlice<f32>,
20353        q_l2: &mut CudaSlice<f32>,
20354        k_l2: &mut CudaSlice<f32>,
20355        v_g: &mut CudaSlice<f32>,
20356        beta: &mut CudaSlice<f32>,
20357        g_log: &mut CudaSlice<f32>,
20358        d_state: usize,
20359        num_v: usize,
20360        num_k: usize,
20361        key_dim: usize,
20362        eps: f32,
20363        conv_dim: usize,
20364        b_n: usize,
20365    ) -> Result<(), Box<dyn std::error::Error>> {
20366        let f = self.func("gdn_prep_decode_b_f32");
20367        let cfg = LaunchConfig {
20368            grid_dim: (num_v as u32, 1, b_n as u32),
20369            block_dim: (32, 4, 1),
20370            shared_mem_bytes: 0,
20371        };
20372        let (ds, nv, nk, kd, cd) = (
20373            d_state as i32,
20374            num_v as i32,
20375            num_k as i32,
20376            key_dim as i32,
20377            conv_dim as i32,
20378        );
20379        let __s_b = self.gpu.stream();
20380        let mut b = __s_b.launch_builder(&f);
20381        b.arg(conv_outs)
20382            .arg(beta_raws)
20383            .arg(alphas)
20384            .arg(dt_bias)
20385            .arg(a)
20386            .arg(q_l2)
20387            .arg(k_l2)
20388            .arg(v_g)
20389            .arg(beta)
20390            .arg(g_log)
20391            .arg(&ds)
20392            .arg(&nv)
20393            .arg(&nk)
20394            .arg(&kd)
20395            .arg(&eps)
20396            .arg(&cd);
20397        unsafe {
20398            b.launch(cfg)?;
20399        }
20400        Ok(())
20401    }
20402
20403    #[allow(clippy::too_many_arguments)]
20404    pub fn gdn_scan_s128_batched_view(
20405        &self,
20406        q: &CudaSlice<f32>,
20407        k: &CudaSlice<f32>,
20408        v: &CudaSlice<f32>,
20409        g: &CudaSlice<f32>,
20410        beta: &CudaSlice<f32>,
20411        state_in_ptrs: &cudarc::driver::CudaView<u64>,
20412        state_out_ptrs: &cudarc::driver::CudaView<u64>,
20413        o: &mut cudarc::driver::CudaViewMut<f32>,
20414        n_head: usize,
20415        b_n: usize,
20416        scale: f32,
20417    ) -> Result<(), Box<dyn std::error::Error>> {
20418        let f = self.func("gdn_scan_s128_b");
20419        const S_V: u32 = 128;
20420        const WARP: u32 = 32;
20421        const COLS_PER_BLOCK: u32 = 4;
20422        let cfg = LaunchConfig {
20423            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20424            block_dim: (WARP, COLS_PER_BLOCK, 1),
20425            shared_mem_bytes: 0,
20426        };
20427        let h = n_head as i32;
20428        let __s_b = self.gpu.stream();
20429        let mut b = __s_b.launch_builder(&f);
20430        b.arg(q)
20431            .arg(k)
20432            .arg(v)
20433            .arg(g)
20434            .arg(beta)
20435            .arg(state_in_ptrs)
20436            .arg(state_out_ptrs)
20437            .arg(o)
20438            .arg(&h)
20439            .arg(&scale);
20440        unsafe {
20441            b.launch(cfg)?;
20442        }
20443        Ok(())
20444    }
20445
20446    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
20447    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
20448    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
20449    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
20450    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
20451    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
20452    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
20453    /// identity law); prime_cache/forward/forward_last are the only callers.
20454    pub fn gdn_chunked_enabled() -> bool {
20455        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20456        *E.get_or_init(|| {
20457            std::env::var("MEMRA_GDN_CHUNKED")
20458                .map(|v| v != "0")
20459                .unwrap_or(true)
20460        })
20461    }
20462
20463    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
20464    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
20465    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
20466    /// of 32 in [32, 128] (kernel row mappings require it).
20467    pub fn gdn_chunk_size() -> usize {
20468        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20469        *C.get_or_init(|| {
20470            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20471                .ok()
20472                .and_then(|v| v.parse().ok())
20473                .unwrap_or(32);
20474            c.clamp(32, 128) / 32 * 32
20475        })
20476    }
20477
20478    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
20479    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
20480    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
20481    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
20482    #[allow(clippy::too_many_arguments)]
20483    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
20484    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
20485    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20486    #[allow(clippy::too_many_arguments)]
20487    pub fn gdn_chunk_k123(
20488        &self,
20489        q: &CudaSlice<f32>,
20490        k: &CudaSlice<f32>,
20491        v: &CudaSlice<f32>,
20492        g: &CudaSlice<f32>,
20493        beta: &CudaSlice<f32>,
20494        wb16: Option<&mut CudaSlice<u8>>,
20495        n_head: usize,
20496        t: usize,
20497        c: usize,
20498        hk: usize,
20499        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20500    ) -> Result<
20501        (
20502            CudaSlice<f32>,
20503            CudaSlice<f32>,
20504            CudaSlice<f32>,
20505            CudaSlice<f32>,
20506        ),
20507        Box<dyn std::error::Error>,
20508    > {
20509        const D: usize = 128;
20510        let h = n_head;
20511        let nc = (t + c - 1) / c;
20512        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20513        let mut gcum = self.uninit(t * h)?;
20514        let mut a = self.uninit(nc * h * c * c)?;
20515        let mut p = self.uninit(nc * h * c * c)?;
20516        let mut u = self.uninit(nc * h * c * D)?;
20517        let mut w = self.uninit(nc * h * c * D)?;
20518        {
20519            // K1
20520            let f = self.func("gdn_chunk_cumgate_f32");
20521            let cfg = LaunchConfig {
20522                grid_dim: (nc as u32, h as u32, 1),
20523                block_dim: (32, 1, 1),
20524                shared_mem_bytes: 0,
20525            };
20526            let __s_b = self.gpu.stream();
20527            let mut b = __s_b.launch_builder(&f);
20528            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20529            unsafe {
20530                b.launch(cfg)?;
20531            }
20532        }
20533        if let Some((qb, kb, pb)) = k2w {
20534            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
20535            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
20536            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20537            let f = self.func("gdn_k2_wgmma");
20538            let cfg = LaunchConfig {
20539                grid_dim: (nc as u32, h as u32, 1),
20540                block_dim: (128, 1, 1),
20541                shared_mem_bytes: 0,
20542            };
20543            let hki = hk as i32;
20544            let __s_b = self.gpu.stream();
20545            let mut b = __s_b.launch_builder(&f);
20546            b.arg(qb)
20547                .arg(kb)
20548                .arg(&gcum)
20549                .arg(beta)
20550                .arg(&mut a)
20551                .arg(&mut *pb)
20552                .arg(&hi)
20553                .arg(&ti)
20554                .arg(&ci)
20555                .arg(&hki);
20556            unsafe {
20557                b.launch(cfg)?;
20558            }
20559        } else if c <= 64 && !portable_mma_gated() {
20560            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
20561            let f = self.func("gdn_chunk_attn_f32");
20562            let jt = ((c + 31) / 32) as u32;
20563            let cfg = LaunchConfig {
20564                grid_dim: (nc as u32, h as u32, jt),
20565                block_dim: (256, 1, 1),
20566                shared_mem_bytes: 0,
20567            };
20568            let hki = hk as i32;
20569            let __s_b = self.gpu.stream();
20570            let mut b = __s_b.launch_builder(&f);
20571            b.arg(q)
20572                .arg(k)
20573                .arg(&gcum)
20574                .arg(beta)
20575                .arg(&mut a)
20576                .arg(&mut p)
20577                .arg(&hi)
20578                .arg(&ti)
20579                .arg(&ci)
20580                .arg(&hki);
20581            unsafe {
20582                b.launch(cfg)?;
20583            }
20584        } else {
20585            // K2 generic (C = 128, or the portable target's low-smem fallback)
20586            assert!(
20587                hk == h,
20588                "generic K2 is broadcast-only (de-broadcast rides C==32)"
20589            );
20590            let f = self.func("gdn_chunk_attn_g_f32");
20591            let cfg = LaunchConfig {
20592                grid_dim: (nc as u32, h as u32, 1),
20593                block_dim: (32, 8, 1),
20594                shared_mem_bytes: 0,
20595            };
20596            let __s_b = self.gpu.stream();
20597            let mut b = __s_b.launch_builder(&f);
20598            b.arg(q)
20599                .arg(k)
20600                .arg(&gcum)
20601                .arg(beta)
20602                .arg(&mut a)
20603                .arg(&mut p)
20604                .arg(&hi)
20605                .arg(&ti)
20606                .arg(&ci);
20607            unsafe {
20608                b.launch(cfg)?;
20609            }
20610        }
20611        {
20612            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
20613            let cfg = LaunchConfig {
20614                grid_dim: (nc as u32, h as u32, 1),
20615                block_dim: (256, 1, 1),
20616                shared_mem_bytes: 0,
20617            };
20618            match c {
20619                32 | 64 => {
20620                    let f = self.func(if c == 32 {
20621                        "gdn_chunk_solve32_f32"
20622                    } else {
20623                        "gdn_chunk_solve64_f32"
20624                    });
20625                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
20626                    let wb: u64 = match wb16 {
20627                        Some(d) => self.addr_u8(d),
20628                        None => 0,
20629                    };
20630                    let hki = hk as i32;
20631                    let __s_b = self.gpu.stream();
20632                    let mut b = __s_b.launch_builder(&f);
20633                    b.arg(v)
20634                        .arg(k)
20635                        .arg(&a)
20636                        .arg(&gcum)
20637                        .arg(&mut u)
20638                        .arg(&mut w)
20639                        .arg(&wb)
20640                        .arg(&hi)
20641                        .arg(&ti)
20642                        .arg(&hki);
20643                    unsafe {
20644                        b.launch(cfg)?;
20645                    }
20646                }
20647                _ => {
20648                    assert!(hk == h, "generic K3 is broadcast-only");
20649                    let f = self.func("gdn_chunk_solve_f32");
20650                    let __s_b = self.gpu.stream();
20651                    let mut b = __s_b.launch_builder(&f);
20652                    b.arg(v)
20653                        .arg(k)
20654                        .arg(&a)
20655                        .arg(&gcum)
20656                        .arg(&mut u)
20657                        .arg(&mut w)
20658                        .arg(&hi)
20659                        .arg(&ti)
20660                        .arg(&ci);
20661                    unsafe {
20662                        b.launch(cfg)?;
20663                    }
20664                }
20665            }
20666        }
20667        Ok((gcum, p, u, w))
20668    }
20669
20670    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
20671    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
20672    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
20673    pub fn gdn_db_on() -> bool {
20674        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
20675    }
20676
20677    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
20678    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
20679    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
20680        !portable_mma_gated()
20681            && c == 32
20682            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
20683                Ok("1") => true,
20684                Ok("0") => false,
20685                _ => cfg!(memra_hopper_mma),
20686            }
20687    }
20688
20689    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
20690    /// mma config; same per-call env read discipline).
20691    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
20692        self.gdn_mma_enabled(c)
20693            && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
20694                Ok("0") => false,
20695                Ok("1") => true,
20696                _ => cfg!(memra_hopper_mma),
20697            }
20698    }
20699
20700    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
20701    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
20702    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
20703    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
20704    #[allow(clippy::too_many_arguments)]
20705    pub fn ssm_conv1d_gdn_state_pad(
20706        &self,
20707        qkv_tm: &cudarc::driver::CudaView<f32>,
20708        conv_state: &mut CudaSlice<f32>,
20709        w: &CudaSlice<f32>,
20710        q_g: &mut CudaSlice<f32>,
20711        k_g: &mut CudaSlice<f32>,
20712        v_g: &mut CudaSlice<f32>,
20713        conv_dim: usize,
20714        t: usize,
20715        d_conv: usize,
20716        d_state: usize,
20717        num_v: usize,
20718        num_k: usize,
20719        key_dim: usize,
20720        hk: usize,
20721        pad_len: Option<&CudaSlice<i32>>,
20722    ) -> Result<(), Box<dyn std::error::Error>> {
20723        assert!(
20724            t >= d_conv - 1,
20725            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
20726        );
20727        {
20728            let f = self.func("ssm_conv1d_gdn_state_f32");
20729            let cfg = LaunchConfig {
20730                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20731                block_dim: (256, 1, 1),
20732                shared_mem_bytes: 0,
20733            };
20734            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20735            let (ds, nv, nk, kd, hki) = (
20736                d_state as i32,
20737                num_v as i32,
20738                num_k as i32,
20739                key_dim as i32,
20740                hk as i32,
20741            );
20742            let __s_b = self.gpu.stream();
20743            let mut b = __s_b.launch_builder(&f);
20744            b.arg(qkv_tm)
20745                .arg(&*conv_state)
20746                .arg(w)
20747                .arg(q_g)
20748                .arg(k_g)
20749                .arg(v_g)
20750                .arg(&cd)
20751                .arg(&ti)
20752                .arg(&dc)
20753                .arg(&ds)
20754                .arg(&nv)
20755                .arg(&nk)
20756                .arg(&kd)
20757                .arg(&hki);
20758            unsafe {
20759                b.launch(cfg)?;
20760            }
20761        }
20762        match pad_len {
20763            Some(len_d) => {
20764                let f = self.func("ssm_conv_ring_update_dev_f32");
20765                let n = conv_dim * (d_conv - 1);
20766                let cfg = LaunchConfig::for_num_elems(n as u32);
20767                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20768                let __s_b = self.gpu.stream();
20769                let mut b = __s_b.launch_builder(&f);
20770                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20771                unsafe {
20772                    b.launch(cfg)?;
20773                }
20774            }
20775            None => {
20776                let f = self.func("ssm_conv_ring_update_f32");
20777                let n = conv_dim * (d_conv - 1);
20778                let cfg = LaunchConfig::for_num_elems(n as u32);
20779                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20780                let __s_b = self.gpu.stream();
20781                let mut b = __s_b.launch_builder(&f);
20782                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20783                unsafe {
20784                    b.launch(cfg)?;
20785                }
20786            }
20787        }
20788        Ok(())
20789    }
20790
20791    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
20792    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
20793    /// K2/K3 can write them.
20794    pub fn gdn_chunk_alloc(
20795        &self,
20796        n_head: usize,
20797        t: usize,
20798        c: usize,
20799        hk: usize,
20800    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
20801        const D: usize = 128;
20802        assert!(
20803            c == 32,
20804            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
20805        );
20806        let h = n_head;
20807        let nc = (t + c - 1) / c;
20808        Ok(GdnChunkBufs {
20809            gcum: self.uninit(t * h)?,
20810            a: self.uninit(nc * h * c * c)?,
20811            p: self.uninit(nc * h * c * c)?,
20812            u: self.uninit(nc * h * c * D)?,
20813            w: self.uninit(nc * h * c * D)?,
20814            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
20815            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
20816            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
20817            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
20818            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
20819            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
20820            o: self.uninit(D * h * t)?,
20821            t,
20822            nc,
20823        })
20824    }
20825
20826    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
20827    pub fn f32_to_bf16_v(
20828        &self,
20829        x: &cudarc::driver::CudaView<f32>,
20830        dst: &mut CudaSlice<u8>,
20831        n: usize,
20832    ) -> Result<(), Box<dyn std::error::Error>> {
20833        let f = self.func("f32_to_bf16_bulk");
20834        let ni = n as i64;
20835        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20836        let __s_b = self.gpu.stream();
20837        let mut b = __s_b.launch_builder(&f);
20838        b.arg(x).arg(dst).arg(&ni);
20839        unsafe {
20840            b.launch(cfg)?;
20841        }
20842        Ok(())
20843    }
20844
20845    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
20846    pub fn f32_to_bf16_into(
20847        &self,
20848        x: &CudaSlice<f32>,
20849        dst: &mut CudaSlice<u8>,
20850        n: usize,
20851    ) -> Result<(), Box<dyn std::error::Error>> {
20852        let f = self.func("f32_to_bf16_bulk");
20853        let ni = n as i64;
20854        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20855        let __s_b = self.gpu.stream();
20856        let mut b = __s_b.launch_builder(&f);
20857        b.arg(x).arg(dst).arg(&ni);
20858        unsafe {
20859            b.launch(cfg)?;
20860        }
20861        Ok(())
20862    }
20863
20864    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
20865    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
20866    pub fn gdn_chunk_k123_vl8(
20867        &self,
20868        seqs: &[GdnSeqVl],
20869        n_head: usize,
20870        hk: usize,
20871        wq: Option<&GdnWVl8>,
20872    ) -> Result<(), Box<dyn std::error::Error>> {
20873        let b = seqs.len();
20874        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
20875        let mut packed = [GdnSeqVl::default(); 8];
20876        packed[..b].copy_from_slice(seqs);
20877        let v = GdnVl8(packed);
20878        let (hi, ci) = (n_head as i32, 32i32);
20879        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
20880        {
20881            let f = self.func("gdn_chunk_cumgate_vl");
20882            let cfg = LaunchConfig {
20883                grid_dim: (max_nc, n_head as u32, b as u32),
20884                block_dim: (32, 1, 1),
20885                shared_mem_bytes: 0,
20886            };
20887            let __s_lb = self.gpu.stream();
20888            let mut lb = __s_lb.launch_builder(&f);
20889            lb.arg(&v).arg(&hi).arg(&ci);
20890            unsafe {
20891                lb.launch(cfg)?;
20892            }
20893        }
20894        let hki = hk as i32;
20895        if let Some(w) = wq {
20896            // K2-wgmma vl twin (writes A + pre-masked Pb16)
20897            let f = self.func("gdn_k2_wgmma_vl");
20898            let cfg = LaunchConfig {
20899                grid_dim: (max_nc, n_head as u32, b as u32),
20900                block_dim: (128, 1, 1),
20901                shared_mem_bytes: 0,
20902            };
20903            let __s_lb = self.gpu.stream();
20904            let mut lb = __s_lb.launch_builder(&f);
20905            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
20906            unsafe {
20907                lb.launch(cfg)?;
20908            }
20909        } else {
20910            let f = self.func("gdn_chunk_attn_vl");
20911            let cfg = LaunchConfig {
20912                grid_dim: (max_nc, n_head as u32, b as u32),
20913                block_dim: (256, 1, 1),
20914                shared_mem_bytes: 0,
20915            };
20916            let __s_lb = self.gpu.stream();
20917            let mut lb = __s_lb.launch_builder(&f);
20918            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
20919            unsafe {
20920                lb.launch(cfg)?;
20921            }
20922        }
20923        {
20924            let f = self.func("gdn_chunk_solve32_vl");
20925            let cfg = LaunchConfig {
20926                grid_dim: (max_nc, n_head as u32, b as u32),
20927                block_dim: (256, 1, 1),
20928                shared_mem_bytes: 0,
20929            };
20930            let __s_lb = self.gpu.stream();
20931            let mut lb = __s_lb.launch_builder(&f);
20932            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
20933            unsafe {
20934                lb.launch(cfg)?;
20935            }
20936        }
20937        Ok(())
20938    }
20939
20940    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
20941    /// fused gate-prep, 5 launches for every sequence (per-element math identical
20942    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
20943    #[allow(clippy::too_many_arguments)]
20944    pub fn gdn_prep_vl8(
20945        &self,
20946        seqs: &[GdnPrepVl],
20947        conv_w: &CudaSlice<f32>,
20948        dt_bias: &CudaSlice<f32>,
20949        a: &CudaSlice<f32>,
20950        conv_dim: usize,
20951        d_conv: usize,
20952        d_state: usize,
20953        num_v: usize,
20954        num_k: usize,
20955        key_dim: usize,
20956        hk: usize,
20957        eps: f32,
20958    ) -> Result<(), Box<dyn std::error::Error>> {
20959        let b = seqs.len();
20960        assert!(b >= 1 && b <= 8);
20961        let mut packed = [GdnPrepVl::default(); 8];
20962        packed[..b].copy_from_slice(seqs);
20963        let v = GdnPrepVl8(packed);
20964        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20965        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
20966        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
20967        assert!(
20968            conv_fuse || hk == num_v,
20969            "de-broadcast requires the fused conv"
20970        );
20971        if conv_fuse {
20972            let f = self.func("ssm_conv1d_gdn_state_vl");
20973            let cfg = LaunchConfig {
20974                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
20975                block_dim: (256, 1, 1),
20976                shared_mem_bytes: 0,
20977            };
20978            let (dsi, nvi, nki, kdi, hki) = (
20979                d_state as i32,
20980                num_v as i32,
20981                num_k as i32,
20982                key_dim as i32,
20983                hk as i32,
20984            );
20985            let __s_lb = self.gpu.stream();
20986            let mut lb = __s_lb.launch_builder(&f);
20987            lb.arg(&v)
20988                .arg(conv_w)
20989                .arg(&cdi)
20990                .arg(&dci)
20991                .arg(&dsi)
20992                .arg(&nvi)
20993                .arg(&nki)
20994                .arg(&kdi)
20995                .arg(&hki);
20996            unsafe {
20997                lb.launch(cfg)?;
20998            }
20999        } else {
21000            let f = self.func("ssm_conv1d_tm_state_vl");
21001            let cfg = LaunchConfig {
21002                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21003                block_dim: (256, 1, 1),
21004                shared_mem_bytes: 0,
21005            };
21006            let __s_lb = self.gpu.stream();
21007            let mut lb = __s_lb.launch_builder(&f);
21008            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21009            unsafe {
21010                lb.launch(cfg)?;
21011            }
21012        }
21013        {
21014            let f = self.func("ssm_conv_ring_update_vl");
21015            let n = (conv_dim * (d_conv - 1)) as u32;
21016            let cfg = LaunchConfig {
21017                grid_dim: (n.div_ceil(256), 1, b as u32),
21018                block_dim: (256, 1, 1),
21019                shared_mem_bytes: 0,
21020            };
21021            let __s_lb = self.gpu.stream();
21022            let mut lb = __s_lb.launch_builder(&f);
21023            lb.arg(&v).arg(&cdi).arg(&dci);
21024            unsafe {
21025                lb.launch(cfg)?;
21026            }
21027        }
21028        if !conv_fuse {
21029            let f = self.func("qkv_to_gdn_repack_vl");
21030            let n = max_t * (num_v * d_state) as u32;
21031            let cfg = LaunchConfig {
21032                grid_dim: (n.div_ceil(256), 1, b as u32),
21033                block_dim: (256, 1, 1),
21034                shared_mem_bytes: 0,
21035            };
21036            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21037            let __s_lb = self.gpu.stream();
21038            let mut lb = __s_lb.launch_builder(&f);
21039            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21040            unsafe {
21041                lb.launch(cfg)?;
21042            }
21043        }
21044        if Self::l2_v2_on(d_state) {
21045            let f = self.func("gdn_l2_v2_vl");
21046            let cfg = LaunchConfig {
21047                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21048                block_dim: (256, 1, 1),
21049                shared_mem_bytes: 0,
21050            };
21051            let (dsi, nvi) = (d_state as i32, hk as i32);
21052            let __s_lb = self.gpu.stream();
21053            let mut lb = __s_lb.launch_builder(&f);
21054            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21055            unsafe {
21056                lb.launch(cfg)?;
21057            }
21058        } else {
21059            let f = self.func("gdn_l2_vl");
21060            let cfg = LaunchConfig {
21061                grid_dim: (max_t * hk as u32, 2, b as u32),
21062                block_dim: (256, 1, 1),
21063                shared_mem_bytes: 0,
21064            };
21065            let (dsi, nvi) = (d_state as i32, hk as i32);
21066            let __s_lb = self.gpu.stream();
21067            let mut lb = __s_lb.launch_builder(&f);
21068            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21069            unsafe {
21070                lb.launch(cfg)?;
21071            }
21072        }
21073        {
21074            let f = self.func("gdn_gate_prep_vl");
21075            let n = max_t * num_v as u32;
21076            let cfg = LaunchConfig {
21077                grid_dim: (n.div_ceil(256), 1, b as u32),
21078                block_dim: (256, 1, 1),
21079                shared_mem_bytes: 0,
21080            };
21081            let nvi = num_v as i32;
21082            let __s_lb = self.gpu.stream();
21083            let mut lb = __s_lb.launch_builder(&f);
21084            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21085            unsafe {
21086                lb.launch(cfg)?;
21087            }
21088        }
21089        Ok(())
21090    }
21091
21092    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
21093    pub fn gdn_mirror_vl8(
21094        &self,
21095        seqs: &[GdnSeqVl],
21096        n_head: usize,
21097        which: i32,
21098        hk: usize,
21099    ) -> Result<(), Box<dyn std::error::Error>> {
21100        let b = seqs.len();
21101        assert!(b >= 1 && b <= 8);
21102        let mut packed = [GdnSeqVl::default(); 8];
21103        packed[..b].copy_from_slice(seqs);
21104        let v = GdnVl8(packed);
21105        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21106        let max_n = seqs
21107            .iter()
21108            .map(|s| {
21109                if which == 0 {
21110                    s.t as i64 * ept as i64
21111                } else {
21112                    s.nc as i64 * ept as i64 * 32
21113                }
21114            })
21115            .max()
21116            .unwrap();
21117        let f = self.func("gdn_mirror_vl");
21118        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21119        let cfg = LaunchConfig {
21120            grid_dim: (blocks, 1, b as u32),
21121            block_dim: (256, 1, 1),
21122            shared_mem_bytes: 0,
21123        };
21124        let __s_lb = self.gpu.stream();
21125        let mut lb = __s_lb.launch_builder(&f);
21126        lb.arg(&v).arg(&ept).arg(&which);
21127        unsafe {
21128            lb.launch(cfg)?;
21129        }
21130        Ok(())
21131    }
21132
21133    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
21134    pub fn gdn_tail_vl8(
21135        &self,
21136        seqs: &[GdnPrepVl],
21137        norm_w: &CudaSlice<f32>,
21138        d_state: usize,
21139        num_v: usize,
21140        eps: f32,
21141    ) -> Result<(), Box<dyn std::error::Error>> {
21142        let b = seqs.len();
21143        assert!(b >= 1 && b <= 8);
21144        let mut packed = [GdnPrepVl::default(); 8];
21145        packed[..b].copy_from_slice(seqs);
21146        let v = GdnPrepVl8(packed);
21147        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21148        let f = self.func("gated_rmsnorm_f16out_vl");
21149        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21150        let cfg = LaunchConfig {
21151            grid_dim: (max_t * num_v as u32, 1, b as u32),
21152            block_dim: (128, 1, 1),
21153            shared_mem_bytes: 0,
21154        };
21155        let (dsi, nvi) = (d_state as i32, num_v as i32);
21156        let __s_lb = self.gpu.stream();
21157        let mut lb = __s_lb.launch_builder(&f);
21158        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21159        unsafe {
21160            lb.launch(cfg)?;
21161        }
21162        Ok(())
21163    }
21164
21165    /// Raw device address helpers for the varlen by-value arg struct (single-stream
21166    /// launches; every buffer outlives the call — the f16 FFI discipline).
21167    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21168        use cudarc::driver::DevicePtr;
21169        let s = self.gpu.stream();
21170        let (p, _g) = x.device_ptr(&s);
21171        p as u64
21172    }
21173    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21174        use cudarc::driver::DevicePtrMut;
21175        let s = self.gpu.stream();
21176        let (p, _g) = x.device_ptr_mut(&s);
21177        p as u64
21178    }
21179    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21180        use cudarc::driver::DevicePtr;
21181        let s = self.gpu.stream();
21182        let (p, _g) = x.device_ptr(&s);
21183        p as u64
21184    }
21185    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21186        use cudarc::driver::DevicePtr;
21187        let s = self.gpu.stream();
21188        let (p, _g) = x.device_ptr(&s);
21189        p as u64
21190    }
21191
21192    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
21193    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
21194    /// launches, so this is strictly bit-gateable against them).
21195    pub fn gdn_chunk_vl8(
21196        &self,
21197        seqs: &[GdnSeqVl],
21198        n_head: usize,
21199        scale: f32,
21200        hk: usize,
21201        wq: Option<&GdnWVl8>,
21202    ) -> Result<(), Box<dyn std::error::Error>> {
21203        const NSPLIT: u32 = 4;
21204        let b = seqs.len();
21205        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21206        let mut packed = [GdnSeqVl::default(); 8];
21207        packed[..b].copy_from_slice(seqs);
21208        let v = GdnVl8(packed);
21209        let (hi, ci) = (n_head as i32, 32i32);
21210        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21211        let hki = hk as i32;
21212        if let Some(w) = wq {
21213            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
21214            let f = self.func("gdn_k45_wgmma_vl");
21215            let cfg = LaunchConfig {
21216                grid_dim: (n_head as u32, NSPLIT, b as u32),
21217                block_dim: (256, 1, 1),
21218                shared_mem_bytes: 0,
21219            };
21220            let __s_lb = self.gpu.stream();
21221            let mut lb = __s_lb.launch_builder(&f);
21222            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21223            unsafe {
21224                lb.launch(cfg)?;
21225            }
21226            let _ = max_nc;
21227            return Ok(());
21228        }
21229        {
21230            let f = self.func("gdn_chunk_state_mma_vl");
21231            let cfg = LaunchConfig {
21232                grid_dim: (n_head as u32, NSPLIT, b as u32),
21233                block_dim: (256, 1, 1),
21234                shared_mem_bytes: 0,
21235            };
21236            let __s_lb = self.gpu.stream();
21237            let mut lb = __s_lb.launch_builder(&f);
21238            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21239            unsafe {
21240                lb.launch(cfg)?;
21241            }
21242        }
21243        {
21244            let f = self.func("gdn_chunk_output_mma_vl");
21245            let cfg = LaunchConfig {
21246                grid_dim: (max_nc, n_head as u32, b as u32),
21247                block_dim: (256, 1, 1),
21248                shared_mem_bytes: 0,
21249            };
21250            let __s_lb = self.gpu.stream();
21251            let mut lb = __s_lb.launch_builder(&f);
21252            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21253            unsafe {
21254                lb.launch(cfg)?;
21255            }
21256        }
21257        Ok(())
21258    }
21259    pub fn gdn_scan_chunked(
21260        &self,
21261        q: &CudaSlice<f32>,
21262        k: &CudaSlice<f32>,
21263        v: &CudaSlice<f32>,
21264        g: &CudaSlice<f32>,
21265        beta: &CudaSlice<f32>,
21266        kb16_pre: Option<&CudaSlice<u8>>,
21267        qb16_pre: Option<&CudaSlice<u8>>,
21268        state_in: &CudaSlice<f32>,
21269        state_out: &mut CudaSlice<f32>,
21270        o: &mut CudaSlice<f32>,
21271        n_head: usize,
21272        t: usize,
21273        scale: f32,
21274        c: usize,
21275        hk: usize,
21276    ) -> Result<(), Box<dyn std::error::Error>> {
21277        const D: usize = 128;
21278        const NSPLIT: u32 = 4;
21279        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21280        let h = n_head;
21281        let nc = (t + c - 1) / c;
21282        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21283        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
21284        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
21285        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
21286        let gdn_mma_pre = !portable_mma_gated()
21287            && c == 32
21288            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21289                Ok("1") => true,
21290                Ok("0") => false,
21291                _ => cfg!(memra_hopper_mma),
21292            };
21293        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21294            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21295        } else {
21296            None
21297        };
21298        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
21299        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
21300        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
21301        let gdn_wgmma_pre = gdn_mma_pre
21302            && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
21303                Ok("0") => false,
21304                Ok("1") => true,
21305                _ => cfg!(memra_hopper_mma),
21306            };
21307        let nk = t * hk * D;
21308        let mut kb16_local: Option<CudaSlice<u8>> = None;
21309        if gdn_mma_pre && kb16_pre.is_none() {
21310            let mut kb = self.alloc_u8_uninit(nk * 2)?;
21311            let f = self.func("f32_to_bf16_bulk");
21312            let n2 = nk as i64;
21313            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21314            let __s_b = self.gpu.stream();
21315            let mut b = __s_b.launch_builder(&f);
21316            b.arg(k).arg(&mut kb).arg(&n2);
21317            unsafe {
21318                b.launch(cfg2)?;
21319            }
21320            kb16_local = Some(kb);
21321        }
21322        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21323        if let Some(kb) = kb16_pre {
21324            assert!(kb.len() >= nk * 2, "kb16_pre too small");
21325        }
21326        let mut qb16: Option<CudaSlice<u8>> = None;
21327        let mut pb16: Option<CudaSlice<u8>> = None;
21328        if gdn_wgmma_pre {
21329            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
21330            // the standalone bulk cvt only serves callers without the prep mirror.
21331            if qb16_pre.is_none() {
21332                let mut qb = self.alloc_u8_uninit(nk * 2)?;
21333                let f = self.func("f32_to_bf16_bulk");
21334                let n2 = nk as i64;
21335                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21336                let __s_b = self.gpu.stream();
21337                let mut b = __s_b.launch_builder(&f);
21338                b.arg(q).arg(&mut qb).arg(&n2);
21339                unsafe {
21340                    b.launch(cfg2)?;
21341                }
21342                qb16 = Some(qb);
21343            } else if let Some(qb) = qb16_pre {
21344                assert!(qb.len() >= nk * 2, "qb16_pre too small");
21345            }
21346            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21347        }
21348        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21349        let k2w = if gdn_wgmma_pre {
21350            Some((
21351                *qb16_ref0.as_ref().unwrap(),
21352                *kb16_ref0.as_ref().unwrap(),
21353                pb16.as_mut().unwrap(),
21354            ))
21355        } else {
21356            None
21357        };
21358        let (gcum, p, u, w) =
21359            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21360        let _ = &w;
21361        let mut y = self.uninit(nc * h * c * D)?;
21362        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
21363        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
21364        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
21365        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
21366        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
21367        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
21368        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
21369        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
21370        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
21371        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
21372        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
21373        let gdn_mma = !portable_mma_gated()
21374            && c == 32
21375            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21376                Ok("1") => true,
21377                Ok("0") => false,
21378                _ => cfg!(memra_hopper_mma),
21379            };
21380        if gdn_mma {
21381            let wb16 = wb16_pre
21382                .take()
21383                .expect("mma path pre-allocates wb16 (K3 store fold)");
21384            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21385            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
21386            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
21387            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
21388            // pass runs inside the persistent-M kernel; Y and Ssnap are never
21389            // materialized. New numeric class (gk folds into k^T instead of ys) —
21390            // explicit opt-in until the state-carry battery promotes it. Env read per
21391            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
21392            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
21393            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
21394            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
21395            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
21396            if gdn_wgmma_pre {
21397                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
21398                let qb16 = qb16_ref0.unwrap();
21399                let pb16 = pb16.as_ref().unwrap();
21400                {
21401                    let f = self.func("gdn_k45_wgmma");
21402                    let cfg = LaunchConfig {
21403                        grid_dim: (h as u32, 4, 1),
21404                        block_dim: (256, 1, 1),
21405                        shared_mem_bytes: 0,
21406                    };
21407                    let hki = hk as i32;
21408                    let __s_b = self.gpu.stream();
21409                    let mut b = __s_b.launch_builder(&f);
21410                    b.arg(kb16_ref)
21411                        .arg(&gcum)
21412                        .arg(beta)
21413                        .arg(&u)
21414                        .arg(&wb16)
21415                        .arg(qb16)
21416                        .arg(pb16)
21417                        .arg(o)
21418                        .arg(&scale)
21419                        .arg(state_in)
21420                        .arg(&mut *state_out)
21421                        .arg(&hi)
21422                        .arg(&ti)
21423                        .arg(&ci)
21424                        .arg(&hki);
21425                    unsafe {
21426                        b.launch(cfg)?;
21427                    }
21428                }
21429                return Ok(());
21430            }
21431            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
21432            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
21433            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
21434            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21435            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21436            {
21437                let f = self.func("gdn_chunk_state_mma");
21438                let cfg = LaunchConfig {
21439                    grid_dim: (h as u32, NSPLIT, 1),
21440                    block_dim: (256, 1, 1),
21441                    shared_mem_bytes: 0,
21442                };
21443                let hki = hk as i32;
21444                let __s_b = self.gpu.stream();
21445                let mut b = __s_b.launch_builder(&f);
21446                b.arg(kb16_ref)
21447                    .arg(&gcum)
21448                    .arg(beta)
21449                    .arg(&u)
21450                    .arg(&wb16)
21451                    .arg(&mut y16)
21452                    .arg(&mut ssnap16)
21453                    .arg(state_in)
21454                    .arg(&mut *state_out)
21455                    .arg(&hi)
21456                    .arg(&ti)
21457                    .arg(&ci)
21458                    .arg(&hki);
21459                unsafe {
21460                    b.launch(cfg)?;
21461                }
21462            }
21463            {
21464                // K5-mma (bf16 St/Y consumers)
21465                let f = self.func("gdn_chunk_output_mma");
21466                let jt = ((c + 31) / 32) as u32;
21467                let cfg = LaunchConfig {
21468                    grid_dim: (nc as u32, h as u32, jt),
21469                    block_dim: (256, 1, 1),
21470                    shared_mem_bytes: 0,
21471                };
21472                let hki = hk as i32;
21473                let __s_b = self.gpu.stream();
21474                let mut b = __s_b.launch_builder(&f);
21475                b.arg(q)
21476                    .arg(&gcum)
21477                    .arg(&p)
21478                    .arg(&y16)
21479                    .arg(&ssnap16)
21480                    .arg(o)
21481                    .arg(&hi)
21482                    .arg(&ti)
21483                    .arg(&ci)
21484                    .arg(&scale)
21485                    .arg(&hki);
21486                unsafe {
21487                    b.launch(cfg)?;
21488                }
21489            }
21490            return Ok(());
21491        }
21492        {
21493            // K4 (sequential over chunks inside; blocks col-partition the state)
21494            let f = self.func("gdn_chunk_state_f32");
21495            let cfg = LaunchConfig {
21496                grid_dim: (h as u32, NSPLIT, 1),
21497                block_dim: (256, 1, 1),
21498                shared_mem_bytes: 0,
21499            };
21500            let __s_b = self.gpu.stream();
21501            let mut b = __s_b.launch_builder(&f);
21502            b.arg(k)
21503                .arg(&gcum)
21504                .arg(beta)
21505                .arg(&u)
21506                .arg(&w)
21507                .arg(&mut y)
21508                .arg(&mut ssnap)
21509                .arg(state_in)
21510                .arg(&mut *state_out)
21511                .arg(&hi)
21512                .arg(&ti)
21513                .arg(&ci);
21514            unsafe {
21515                b.launch(cfg)?;
21516            }
21517        }
21518        {
21519            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
21520            let f = self.func("gdn_chunk_output_f32");
21521            let jt = ((c + 31) / 32) as u32;
21522            let cfg = LaunchConfig {
21523                grid_dim: (nc as u32, h as u32, jt),
21524                block_dim: (256, 1, 1),
21525                shared_mem_bytes: 0,
21526            };
21527            let __s_b = self.gpu.stream();
21528            let mut b = __s_b.launch_builder(&f);
21529            b.arg(q)
21530                .arg(&gcum)
21531                .arg(&p)
21532                .arg(&y)
21533                .arg(&ssnap)
21534                .arg(o)
21535                .arg(&hi)
21536                .arg(&ti)
21537                .arg(&ci)
21538                .arg(&scale);
21539            unsafe {
21540                b.launch(cfg)?;
21541            }
21542        }
21543        Ok(())
21544    }
21545
21546    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
21547    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
21548    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
21549    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
21550    ///
21551    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
21552    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
21553    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
21554    #[allow(clippy::too_many_arguments)]
21555    #[allow(clippy::too_many_arguments)]
21556    pub fn gdn_scan_prefill(
21557        &self,
21558        q: &CudaSlice<f32>,
21559        k: &CudaSlice<f32>,
21560        v: &CudaSlice<f32>,
21561        g: &CudaSlice<f32>,
21562        beta: &CudaSlice<f32>,
21563        kb16_pre: Option<&CudaSlice<u8>>,
21564        qb16_pre: Option<&CudaSlice<u8>>,
21565        state_in: &CudaSlice<f32>,
21566        state_out: &mut CudaSlice<f32>,
21567        o: &mut CudaSlice<f32>,
21568        n_head: usize,
21569        t: usize,
21570        scale: f32,
21571        hk: usize,
21572    ) -> Result<(), Box<dyn std::error::Error>> {
21573        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21574            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21575            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21576        }
21577        if Self::gdn_chunked_enabled() && t >= 16 {
21578            self.gdn_scan_chunked(
21579                q,
21580                k,
21581                v,
21582                g,
21583                beta,
21584                kb16_pre,
21585                qb16_pre,
21586                state_in,
21587                state_out,
21588                o,
21589                n_head,
21590                t,
21591                scale,
21592                Self::gdn_chunk_size(),
21593                hk,
21594            )
21595        } else {
21596            assert!(
21597                hk == n_head,
21598                "s128 scan is broadcast-only (prep guarantees by predicate)"
21599            );
21600            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21601        }
21602    }
21603
21604    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
21605    #[allow(clippy::too_many_arguments)]
21606    fn gdn_scan_diff(
21607        &self,
21608        q: &CudaSlice<f32>,
21609        k: &CudaSlice<f32>,
21610        v: &CudaSlice<f32>,
21611        g: &CudaSlice<f32>,
21612        beta: &CudaSlice<f32>,
21613        state_in: &CudaSlice<f32>,
21614        state_out: &mut CudaSlice<f32>,
21615        o: &mut CudaSlice<f32>,
21616        n_head: usize,
21617        t: usize,
21618        scale: f32,
21619    ) -> Result<(), Box<dyn std::error::Error>> {
21620        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21621        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21622        let mut o_c = self.uninit(o.len())?;
21623        let mut st_c = self.uninit(state_out.len())?;
21624        self.gdn_scan_chunked(
21625            q,
21626            k,
21627            v,
21628            g,
21629            beta,
21630            None,
21631            None,
21632            state_in,
21633            &mut st_c,
21634            &mut o_c,
21635            n_head,
21636            t,
21637            scale,
21638            Self::gdn_chunk_size(),
21639            n_head,
21640        )?;
21641        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
21642        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
21643        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
21644        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
21645            let mut max_abs = 0f32;
21646            let mut max_rel = 0f32;
21647            let mut sum_rel = 0f64;
21648            for (x, y) in a.iter().zip(b) {
21649                let ad = (x - y).abs();
21650                let rel = ad / x.abs().max(y.abs()).max(1e-3);
21651                if ad > max_abs {
21652                    max_abs = ad;
21653                }
21654                if rel > max_rel {
21655                    max_rel = rel;
21656                }
21657                sum_rel += rel as f64;
21658            }
21659            (max_abs, max_rel, sum_rel / a.len() as f64)
21660        };
21661        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
21662        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
21663        println!(
21664            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
21665                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
21666            Self::gdn_chunk_size()
21667        );
21668        Ok(())
21669    }
21670
21671    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
21672    pub fn gdn_glog(
21673        &self,
21674        alpha: &CudaSlice<f32>,
21675        dt_bias: &CudaSlice<f32>,
21676        a: &CudaSlice<f32>,
21677        g_log: &mut CudaSlice<f32>,
21678        n_head: usize,
21679        t: usize,
21680    ) -> Result<(), Box<dyn std::error::Error>> {
21681        let f = self.func("gdn_glog_f32");
21682        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21683        let (h, ti) = (n_head as i32, t as i32);
21684        let __s_b = self.gpu.stream();
21685        let mut b = __s_b.launch_builder(&f);
21686        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21687        unsafe {
21688            b.launch(cfg)?;
21689        }
21690        Ok(())
21691    }
21692
21693    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
21694    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
21695    pub fn sigmoid_v(
21696        &self,
21697        x: &cudarc::driver::CudaView<f32>,
21698        y: &mut CudaSlice<f32>,
21699        n: usize,
21700    ) -> Result<(), Box<dyn std::error::Error>> {
21701        let f = self.func("sigmoid_f32");
21702        let cfg = LaunchConfig::for_num_elems(n as u32);
21703        let ni = n as i32;
21704        let __s_b = self.gpu.stream();
21705        let mut b = __s_b.launch_builder(&f);
21706        b.arg(x).arg(y).arg(&ni);
21707        unsafe {
21708            b.launch(cfg)?;
21709        }
21710        Ok(())
21711    }
21712
21713    pub fn gdn_glog_v(
21714        &self,
21715        alpha: &cudarc::driver::CudaView<f32>,
21716        dt_bias: &CudaSlice<f32>,
21717        a: &CudaSlice<f32>,
21718        g_log: &mut CudaSlice<f32>,
21719        n_head: usize,
21720        t: usize,
21721    ) -> Result<(), Box<dyn std::error::Error>> {
21722        let f = self.func("gdn_glog_f32");
21723        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21724        let (h, ti) = (n_head as i32, t as i32);
21725        let __s_b = self.gpu.stream();
21726        let mut b = __s_b.launch_builder(&f);
21727        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21728        unsafe {
21729            b.launch(cfg)?;
21730        }
21731        Ok(())
21732    }
21733
21734    pub fn sigmoid(
21735        &self,
21736        x: &CudaSlice<f32>,
21737        y: &mut CudaSlice<f32>,
21738        n: usize,
21739    ) -> Result<(), Box<dyn std::error::Error>> {
21740        let f = self.func("sigmoid_f32");
21741        let cfg = LaunchConfig::for_num_elems(n as u32);
21742        let ni = n as i32;
21743        let __s_b = self.gpu.stream();
21744        let mut b = __s_b.launch_builder(&f);
21745        b.arg(x).arg(y).arg(&ni);
21746        unsafe {
21747            b.launch(cfg)?;
21748        }
21749        Ok(())
21750    }
21751
21752    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
21753    /// (replaces sigmoid + mul + convert). Bit-identical class.
21754    pub fn sig_mul_f16out(
21755        &self,
21756        a: &CudaSlice<f32>,
21757        g: &CudaSlice<f32>,
21758        dst: &mut CudaSlice<f32>,
21759        dst16: &mut CudaSlice<u8>,
21760        n: usize,
21761    ) -> Result<(), Box<dyn std::error::Error>> {
21762        let f = self.func("sig_mul_f16out_f32");
21763        let cfg = LaunchConfig::for_num_elems(n as u32);
21764        let ni = n as i32;
21765        let __s_b = self.gpu.stream();
21766        let mut b = __s_b.launch_builder(&f);
21767        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
21768        unsafe {
21769            b.launch(cfg)?;
21770        }
21771        Ok(())
21772    }
21773
21774    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
21775    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
21776    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
21777    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
21778    ///
21779    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
21780    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
21781    /// applies the wrong number of distinct gate values.
21782    #[allow(clippy::too_many_arguments)]
21783    pub fn attn_head_gate(
21784        &self,
21785        a: &CudaSlice<f32>,
21786        g: &CudaSlice<f32>,
21787        dst: &mut CudaSlice<f32>,
21788        dst16: Option<&mut CudaSlice<u8>>,
21789        head_dim: usize,
21790        n_head: usize,
21791        t: usize,
21792    ) -> Result<(), Box<dyn std::error::Error>> {
21793        let f = self.func("attn_head_gate_f32");
21794        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
21795        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
21796        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
21797        let d16: u64 = match dst16 {
21798            Some(d) => self.addr_u8(d),
21799            None => 0,
21800        };
21801        let __s_b = self.gpu.stream();
21802        let mut b = __s_b.launch_builder(&f);
21803        b.arg(a)
21804            .arg(g)
21805            .arg(dst)
21806            .arg(&d16)
21807            .arg(&hd)
21808            .arg(&nh)
21809            .arg(&ti);
21810        unsafe {
21811            b.launch(cfg)?;
21812        }
21813        Ok(())
21814    }
21815
21816    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
21817    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
21818    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
21819    ///
21820    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
21821    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
21822    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
21823    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
21824    #[allow(clippy::too_many_arguments)]
21825    pub fn swiglu_clamped_mul_scaled(
21826        &self,
21827        gate: &CudaSlice<f32>,
21828        up: &CudaSlice<f32>,
21829        gs: f32,
21830        us: f32,
21831        limit: f32,
21832        dst: &mut CudaSlice<f32>,
21833        n: usize,
21834    ) -> Result<(), Box<dyn std::error::Error>> {
21835        debug_assert!(
21836            limit > 1e-6,
21837            "swiglu_clamped needs a live limit; use silu_mul_scaled"
21838        );
21839        let f = self.func("swiglu_clamped_mul_scaled_f32");
21840        let cfg = LaunchConfig::for_num_elems(n as u32);
21841        let ni = n as i32;
21842        let __s_b = self.gpu.stream();
21843        let mut b = __s_b.launch_builder(&f);
21844        b.arg(gate)
21845            .arg(up)
21846            .arg(&gs)
21847            .arg(&us)
21848            .arg(&limit)
21849            .arg(dst)
21850            .arg(&ni);
21851        unsafe {
21852            b.launch(cfg)?;
21853        }
21854        Ok(())
21855    }
21856
21857    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
21858    pub fn gated_rmsnorm(
21859        &self,
21860        o: &CudaSlice<f32>,
21861        w: &CudaSlice<f32>,
21862        z: &CudaSlice<f32>,
21863        dst: &mut CudaSlice<f32>,
21864        ncols: usize,
21865        nrows: usize,
21866        eps: f32,
21867    ) -> Result<(), Box<dyn std::error::Error>> {
21868        let f = self.func("gated_rmsnorm_f32");
21869        let cfg = LaunchConfig {
21870            grid_dim: (nrows as u32, 1, 1),
21871            block_dim: (128, 1, 1),
21872            shared_mem_bytes: 0,
21873        };
21874        let (nc, e) = (ncols as i32, eps);
21875        let __s_b = self.gpu.stream();
21876        let mut b = __s_b.launch_builder(&f);
21877        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
21878        unsafe {
21879            b.launch(cfg)?;
21880        }
21881        Ok(())
21882    }
21883
21884    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
21885    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
21886    pub fn gated_rmsnorm_f16out(
21887        &self,
21888        o: &CudaSlice<f32>,
21889        w: &CudaSlice<f32>,
21890        z: &CudaSlice<f32>,
21891        dst: &mut CudaSlice<f32>,
21892        dst16: &mut CudaSlice<u8>,
21893        ncols: usize,
21894        nrows: usize,
21895        eps: f32,
21896    ) -> Result<(), Box<dyn std::error::Error>> {
21897        let f = self.func("gated_rmsnorm_f16out_f32");
21898        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21899        let cfg = LaunchConfig {
21900            grid_dim: (nrows as u32, 1, 1),
21901            block_dim: (128, 1, 1),
21902            shared_mem_bytes: 0,
21903        };
21904        let (nc, e) = (ncols as i32, eps);
21905        let __s_b = self.gpu.stream();
21906        let mut b = __s_b.launch_builder(&f);
21907        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
21908        unsafe {
21909            b.launch(cfg)?;
21910        }
21911        Ok(())
21912    }
21913
21914    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
21915    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
21916    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
21917    #[allow(clippy::too_many_arguments)]
21918    pub fn add_rms_norm_zq8(
21919        &self,
21920        a: &CudaSlice<f32>,
21921        b_in: &CudaSlice<f32>,
21922        w: &CudaSlice<f32>,
21923        res: &mut CudaSlice<f32>,
21924        z: &mut CudaSlice<f32>,
21925        ncols: usize,
21926        nrows: usize,
21927        eps: f32,
21928    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21929        assert!(ncols % 32 == 0);
21930        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
21931        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
21932        let f = self.func("add_rms_norm_zq8");
21933        let cfg = LaunchConfig {
21934            grid_dim: (nrows as u32, 1, 1),
21935            block_dim: (1024, 1, 1),
21936            shared_mem_bytes: 0,
21937        };
21938        let (nc, ep) = (ncols as i32, eps);
21939        let __s_b = self.gpu.stream();
21940        let mut b = __s_b.launch_builder(&f);
21941        b.arg(a)
21942            .arg(b_in)
21943            .arg(w)
21944            .arg(res)
21945            .arg(z)
21946            .arg(&mut q)
21947            .arg(&mut d)
21948            .arg(&nc)
21949            .arg(&ep);
21950        unsafe {
21951            b.launch(cfg)?;
21952        }
21953        Ok((q, d))
21954    }
21955
21956    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
21957    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
21958    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
21959    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
21960    pub fn gated_rmsnorm_zv(
21961        &self,
21962        o: &CudaSlice<f32>,
21963        w: &CudaSlice<f32>,
21964        z: &cudarc::driver::CudaView<f32>,
21965        dst: &mut CudaSlice<f32>,
21966        ncols: usize,
21967        nrows: usize,
21968        eps: f32,
21969    ) -> Result<(), Box<dyn std::error::Error>> {
21970        let f = self.func("gated_rmsnorm_f32");
21971        let cfg = LaunchConfig {
21972            grid_dim: (nrows as u32, 1, 1),
21973            block_dim: (128, 1, 1),
21974            shared_mem_bytes: 0,
21975        };
21976        let (nc, e) = (ncols as i32, eps);
21977        let __s_b = self.gpu.stream();
21978        let mut b = __s_b.launch_builder(&f);
21979        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
21980        unsafe {
21981            b.launch(cfg)?;
21982        }
21983        Ok(())
21984    }
21985
21986    pub fn gated_rmsnorm_f16out_zv(
21987        &self,
21988        o: &CudaSlice<f32>,
21989        w: &CudaSlice<f32>,
21990        z: &cudarc::driver::CudaView<f32>,
21991        dst: &mut CudaSlice<f32>,
21992        dst16: &mut CudaSlice<u8>,
21993        ncols: usize,
21994        nrows: usize,
21995        eps: f32,
21996    ) -> Result<(), Box<dyn std::error::Error>> {
21997        let f = self.func("gated_rmsnorm_f16out_f32");
21998        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21999        let cfg = LaunchConfig {
22000            grid_dim: (nrows as u32, 1, 1),
22001            block_dim: (128, 1, 1),
22002            shared_mem_bytes: 0,
22003        };
22004        let (nc, e) = (ncols as i32, eps);
22005        let __s_b = self.gpu.stream();
22006        let mut b = __s_b.launch_builder(&f);
22007        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22008        unsafe {
22009            b.launch(cfg)?;
22010        }
22011        Ok(())
22012    }
22013
22014    pub fn gated_rmsnorm_q8_1(
22015        &self,
22016        o: &CudaSlice<f32>,
22017        w: &CudaSlice<f32>,
22018        z: &CudaSlice<f32>,
22019        ncols: usize,
22020        nrows: usize,
22021        eps: f32,
22022    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22023        assert!(ncols % 32 == 0);
22024        let f = self.func("gated_rmsnorm_q8_1");
22025        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22026        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22027        let cfg = LaunchConfig {
22028            grid_dim: (nrows as u32, 1, 1),
22029            block_dim: (128, 1, 1),
22030            shared_mem_bytes: 0,
22031        };
22032        let (nc, ep) = (ncols as i32, eps);
22033        let __s_b = self.gpu.stream();
22034        let mut b = __s_b.launch_builder(&f);
22035        b.arg(o)
22036            .arg(w)
22037            .arg(z)
22038            .arg(&mut out_q)
22039            .arg(&mut out_d)
22040            .arg(&nc)
22041            .arg(&ep);
22042        unsafe {
22043            b.launch(cfg)?;
22044        }
22045        Ok((out_q, out_d))
22046    }
22047
22048    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
22049    pub fn transpose(
22050        &self,
22051        inp: &CudaSlice<f32>,
22052        rows: usize,
22053        cols: usize,
22054    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22055        let f = self.func("transpose_f32");
22056        let mut out = self.zeros(rows * cols)?;
22057        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22058        let (r, c) = (rows as i32, cols as i32);
22059        let __s_b = self.gpu.stream();
22060        let mut b = __s_b.launch_builder(&f);
22061        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22062        unsafe {
22063            b.launch(cfg)?;
22064        }
22065        Ok(out)
22066    }
22067
22068    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
22069    pub fn repeat_heads(
22070        &self,
22071        inp: &CudaSlice<f32>,
22072        out: &mut CudaSlice<f32>,
22073        head_dim: usize,
22074        n_in: usize,
22075        n_out: usize,
22076        t: usize,
22077    ) -> Result<(), Box<dyn std::error::Error>> {
22078        let f = self.func("repeat_heads_f32");
22079        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22080        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22081        let __s_b = self.gpu.stream();
22082        let mut b = __s_b.launch_builder(&f);
22083        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22084        unsafe {
22085            b.launch(cfg)?;
22086        }
22087        Ok(())
22088    }
22089
22090    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
22091    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
22092    pub fn q_gate_split(
22093        &self,
22094        qf: &CudaSlice<f32>,
22095        q_out: &mut CudaSlice<f32>,
22096        gate_out: &mut CudaSlice<f32>,
22097        head_dim: usize,
22098        n_head: usize,
22099        t: usize,
22100    ) -> Result<(), Box<dyn std::error::Error>> {
22101        let f = self.func("q_gate_split_f32");
22102        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22103        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22104        let __s_b = self.gpu.stream();
22105        let mut b = __s_b.launch_builder(&f);
22106        b.arg(qf)
22107            .arg(q_out)
22108            .arg(gate_out)
22109            .arg(&hd)
22110            .arg(&nh)
22111            .arg(&ti);
22112        unsafe {
22113            b.launch(cfg)?;
22114        }
22115        Ok(())
22116    }
22117
22118    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
22119    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
22120    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
22121    pub fn qkv_to_gdn_repack(
22122        &self,
22123        conv_out: &CudaSlice<f32>,
22124        q_g: &mut CudaSlice<f32>,
22125        k_g: &mut CudaSlice<f32>,
22126        v_g: &mut CudaSlice<f32>,
22127        d_state: usize,
22128        num_v: usize,
22129        num_k: usize,
22130        key_dim: usize,
22131        t: usize,
22132    ) -> Result<(), Box<dyn std::error::Error>> {
22133        let f = self.func("qkv_to_gdn_repack_f32");
22134        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22135        let (ds, nv, nk, kd, ti) = (
22136            d_state as i32,
22137            num_v as i32,
22138            num_k as i32,
22139            key_dim as i32,
22140            t as i32,
22141        );
22142        let __s_b = self.gpu.stream();
22143        let mut b = __s_b.launch_builder(&f);
22144        b.arg(conv_out)
22145            .arg(q_g)
22146            .arg(k_g)
22147            .arg(v_g)
22148            .arg(&ds)
22149            .arg(&nv)
22150            .arg(&nk)
22151            .arg(&kd)
22152            .arg(&ti);
22153        unsafe {
22154            b.launch(cfg)?;
22155        }
22156        Ok(())
22157    }
22158
22159    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
22160    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
22161    pub fn conv_left_pad(
22162        &self,
22163        src: &CudaSlice<f32>,
22164        dst: &mut CudaSlice<f32>,
22165        conv_dim: usize,
22166        t: usize,
22167        pad: usize,
22168    ) -> Result<(), Box<dyn std::error::Error>> {
22169        let f = self.func("conv_left_pad_f32");
22170        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22171        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22172        let __s_b = self.gpu.stream();
22173        let mut b = __s_b.launch_builder(&f);
22174        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22175        unsafe {
22176            b.launch(cfg)?;
22177        }
22178        Ok(())
22179    }
22180
22181    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
22182    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
22183    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
22184    pub fn conv_assemble_and_roll(
22185        &self,
22186        qkv_col: &CudaSlice<f32>,
22187        conv_state: &mut CudaSlice<f32>,
22188        conv_in: &mut CudaSlice<f32>,
22189        conv_dim: usize,
22190        pad: usize,
22191    ) -> Result<(), Box<dyn std::error::Error>> {
22192        let f = self.func("conv_assemble_and_roll_f32");
22193        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22194        let (cd, p) = (conv_dim as i32, pad as i32);
22195        let __s_b = self.gpu.stream();
22196        let mut b = __s_b.launch_builder(&f);
22197        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22198        unsafe {
22199            b.launch(cfg)?;
22200        }
22201        Ok(())
22202    }
22203
22204    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
22205    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
22206    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
22207    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
22208    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
22209    pub fn ssm_conv1d_fused_decode(
22210        &self,
22211        qkv_col: &CudaSlice<f32>,
22212        conv_state: &mut CudaSlice<f32>,
22213        w: &CudaSlice<f32>,
22214        conv_out: &mut CudaSlice<f32>,
22215        conv_dim: usize,
22216        d_conv: usize,
22217    ) -> Result<(), Box<dyn std::error::Error>> {
22218        let f = self.func("ssm_conv1d_fused_decode_f32");
22219        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22220        let (cd, dc) = (conv_dim as i32, d_conv as i32);
22221        let __s_b = self.gpu.stream();
22222        let mut b = __s_b.launch_builder(&f);
22223        b.arg(qkv_col)
22224            .arg(conv_state)
22225            .arg(w)
22226            .arg(conv_out)
22227            .arg(&cd)
22228            .arg(&dc);
22229        unsafe {
22230            b.launch(cfg)?;
22231        }
22232        Ok(())
22233    }
22234
22235    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
22236    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
22237    pub fn slice_range(
22238        &self,
22239        src: &CudaSlice<f32>,
22240        start: usize,
22241        len: usize,
22242    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22243        let host = self.gpu.stream().clone_dtoh(src)?;
22244        self.gpu.stream().synchronize()?;
22245        Ok(self.htod(&host[start..start + len])?)
22246    }
22247}
22248
22249#[cfg(test)]
22250mod target_dispatch_tests {
22251    use super::legacy_quant_gemm_allowed;
22252
22253    #[test]
22254    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22255        // sm_120a native lane
22256        assert!(legacy_quant_gemm_allowed(false, false, false));
22257        assert!(!legacy_quant_gemm_allowed(false, false, true));
22258        // pure portable lane (sm_89): gated
22259        assert!(!legacy_quant_gemm_allowed(true, false, false));
22260        assert!(!legacy_quant_gemm_allowed(true, false, true));
22261        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
22262        assert!(legacy_quant_gemm_allowed(true, true, false));
22263        assert!(!legacy_quant_gemm_allowed(true, true, true));
22264    }
22265
22266    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22267    #[test]
22268    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22269        assert!(!legacy_quant_gemm_allowed(
22270            cfg!(memra_portable_cuda),
22271            cfg!(memra_hopper_mma),
22272            false
22273        ));
22274    }
22275
22276    #[cfg(memra_hopper_mma)]
22277    #[test]
22278    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22279        assert!(legacy_quant_gemm_allowed(
22280            cfg!(memra_portable_cuda),
22281            cfg!(memra_hopper_mma),
22282            false
22283        ));
22284        assert!(super::portable_mma_gated() == false);
22285    }
22286}
22287
22288/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
22289/// inherent methods (inherent methods win name resolution, so no recursion).
22290impl memra_kv::KvDev for Engine {
22291    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22292        Engine::zeros(self, n)
22293    }
22294    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22295        Engine::uninit(self, n)
22296    }
22297    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22298        Engine::alloc_u8(self, n)
22299    }
22300    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22301        Engine::htod_i32(self, v)
22302    }
22303    fn clone_dtod(
22304        &self,
22305        src: &CudaSlice<f32>,
22306    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22307        Engine::clone_dtod(self, src)
22308    }
22309    fn copy_into(
22310        &self,
22311        dst: &mut CudaSlice<f32>,
22312        off: usize,
22313        src: &CudaSlice<f32>,
22314        len: usize,
22315    ) -> Result<(), Box<dyn std::error::Error>> {
22316        Engine::copy_into(self, dst, off, src, len)
22317    }
22318    fn set_i32_one(
22319        &self,
22320        d: &mut CudaSlice<i32>,
22321        v: i32,
22322    ) -> Result<(), Box<dyn std::error::Error>> {
22323        Engine::set_i32_one(self, d, v)
22324    }
22325}