Skip to main content

memra_engine/
lib.rs

1//! memra engine: Stage-1 correctness-first forward-pass kernels + ops, on sm_120 via cudarc.
2
3use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceSlice, LaunchConfig,
6    PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13#[cfg(debug_assertions)]
14pub(crate) fn debug_assert_tensor_stream_device<T>(
15    tensor: &CudaSlice<T>,
16    stream: &CudaStream,
17    site: &str,
18) {
19    let tensor_dev = tensor.ordinal();
20    let stream_dev = stream.context().ordinal();
21    assert_eq!(
22        tensor_dev, stream_dev,
23        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
24    );
25}
26
27fn ensure_tensor_stream_device<T>(
28    tensor: &impl DeviceSlice<T>,
29    stream: &CudaStream,
30    site: &str,
31) -> Result<(), Box<dyn std::error::Error>> {
32    let tensor_dev = tensor.stream().context().ordinal();
33    let stream_dev = stream.context().ordinal();
34    if tensor_dev != stream_dev {
35        return Err(format!(
36            "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
37             stream on dev{stream_dev}"
38        )
39        .into());
40    }
41    Ok(())
42}
43
44pub use memra_gguf;
45pub use memra_runtime;
46
47pub mod forward;
48pub mod hybrid;
49pub mod hybrid_forward;
50pub mod model;
51pub mod sigrouter_contract;
52pub mod vision;
53pub mod vision_gemma;
54pub mod vision_pre;
55/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
56/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
57pub mod cache {
58    pub use memra_kv::*;
59}
60pub mod decode;
61pub mod decode_batch;
62pub mod dflash;
63pub mod eagle;
64pub mod gemma_spec;
65pub mod graph_update;
66/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
67/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
68/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
69pub mod mla;
70pub mod moesd;
71pub mod parallel;
72pub mod plan_backend;
73pub mod pp;
74pub mod round_stream;
75pub mod spec;
76pub mod tp;
77pub use memra_sampling as sampler;
78
79/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
80/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
81/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
82/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
83/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
84///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
85///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
86///                     stream sync per projection (round-47 ledgered defect).
87///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
88///                     construction, zero syncs, f32 C with the act row-scale folded in.
89/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
90/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
91/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
92/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
93/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
94/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
95///
96/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
97/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
98/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
99/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
100/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
101/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
102/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
103/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
104///
105/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
106/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
107/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
108/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
109/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
110/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
111/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
112///
113/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
114/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
115/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
116/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
117/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
118/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
119/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
120/// the k-quant-only admission survives as the rollback seam, not the default.
121/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
122pub fn moe_f16g_mode() -> u8 {
123    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
124    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
125        Ok("0") => 0,
126        Ok("2") => 2,
127        Ok("3") => 3,
128        Ok(_) => 1,
129        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
130        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
131        Err(_) => 2,
132    })
133}
134/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
135/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
136/// (shape_sel, cross) for the FFI:
137///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
138///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
139///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
140///                         back to 32x64 in-launcher when the device/in_f can't take it).
141///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
142///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
143///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
144///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
145///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
146///                         verdict was stale).
147pub fn moe_f16g_sk_params() -> (i32, i32) {
148    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
149    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
150        Ok("0") => (-1, 0),
151        Ok("32") => (0, i32::MAX),
152        Ok("128") => (0, 1),
153        _ => {
154            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
155                .ok()
156                .and_then(|v| v.parse().ok())
157                .unwrap_or(64);
158            (0, cross)
159        }
160    })
161}
162/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
163/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
164/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
165/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
166/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
167/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
168/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
169/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
170/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
171/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
172pub fn moe_f16g_direct_on(qtype: i32) -> bool {
173    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
174    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
175        Ok("0") => 0,
176        Ok("kq") => 1,
177        _ => 2,
178    });
179    match m {
180        0 => false,
181        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
182        _ => true,
183    }
184}
185/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
186/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
187/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
188/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
189/// stage under q35's routing skew. Bit-identical to every other sk form by construction
190/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
191/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
192/// tail. in_f % 64 != 0 falls back in-launcher.
193pub fn moe_f16g_tail_on() -> bool {
194    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
196}
197
198/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
199/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
200/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
201/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
202/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
203/// still opens this door for A/B.
204pub fn moe_f16g_gemma_on() -> bool {
205    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
207}
208
209/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
210/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
211/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
212pub fn moe_fuse_actq_on() -> bool {
213    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
214    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
215}
216
217/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
218/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
219/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
220/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
221/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
222/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
223/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
224/// verify already use (dispatch parity, one router kernel for every t).
225/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
226pub fn router_prefill_exact_on() -> bool {
227    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
229}
230
231pub fn router_kernel_on() -> bool {
232    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
233    *ON.get_or_init(|| {
234        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
235        if !on {
236            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
237        }
238        on
239    })
240}
241
242/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
243/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
244/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
245/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
246/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
247/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
248/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
249/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
250/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
251/// seam, perf-only: bits are equal by the kernel-check gate).
252/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
253/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
254/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
255/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
256pub const ROUTER_BATCH_MIN_T: usize = 8;
257pub fn router_batch_on() -> bool {
258    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
259    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
260}
261mod cpu_experts;
262#[cfg(memra_cutlass)]
263pub mod cutlass_ffi;
264pub mod dsv4_ffi;
265pub mod dsv4_gpu;
266pub mod f16_ffi;
267pub mod fp8_ffi;
268pub mod mmq_ffi;
269pub mod moe_cache;
270pub mod prime_graph;
271pub mod spill;
272mod spill_pread;
273
274// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
275// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
276// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
277// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
278// broke every machine that wasn't the build machine. Same bytes, same module image;
279// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
280const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
281const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
282const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
283const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
284const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
285const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
286/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
287const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
288
289/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
290/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
291/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
292/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
293/// compile-time default (zero behavior change).
294fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
295    assert!(
296        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
297        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
298    );
299    match std::env::var("MEMRA_GEMM_FATBIN") {
300        Ok(path) => std::borrow::Cow::Owned(
301            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
302        ),
303        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
304    }
305}
306
307/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
308/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
309/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
310/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
311/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
312/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
313pub(crate) const fn portable_mma_gated() -> bool {
314    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
315}
316
317/// Refuse an env force that would reach a kernel THIS BUILD DOES NOT CONTAIN.
318///
319/// Doors of the shape `MEMRA_X=1 => true` are arch-blind: they were written so an operator could
320/// force a promoted path on, and the default arm (`cfg!(memra_hopper_mma)` or similar) is the only
321/// thing that consulted the arch. On a portable build the forced path then reaches
322/// `Engine::func`, which resolves lazily and ends in `panic!("kernel {name} not in any fatbin")` —
323/// a confusing crash naming a kernel the operator never heard of, several frames from the switch
324/// they actually flipped.
325///
326/// Found 2026-08-23 by tools/fatbin-lookup-census.py, which listed 20 looked-up kernels absent
327/// from the sm_89 fatbins. 18 of those turned out to be correctly unreachable (the GDN varlen
328/// chain is gated through `gdn_mma_enabled`, which starts with `!portable_mma_gated()`); these
329/// env doors were the two that were genuinely reachable, and only by explicit operator action.
330///
331/// Same shape and same message style as `gemm_fatbin_bytes`'s refusal above — one idiom for
332/// "this switch cannot work on this build", so it fails at the switch instead of at the lookup.
333#[track_caller]
334pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
335    assert!(
336        !portable_mma_gated(),
337        "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
338         portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
339    );
340}
341
342/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
343/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
344/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
345/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
346/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
347/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
348/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
349/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
350pub(crate) const fn gdn_mma_default_on() -> bool {
351    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
352}
353
354/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
355const fn konst_eq(a: &str, b: &str) -> bool {
356    let (a, b) = (a.as_bytes(), b.as_bytes());
357    if a.len() != b.len() {
358        return false;
359    }
360    let mut i = 0;
361    while i < a.len() {
362        if a[i] != b[i] {
363            return false;
364        }
365        i += 1;
366    }
367    true
368}
369
370/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
371/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
372/// in a pure helper so the dispatch guard can be regression-tested without constructing an
373/// Engine or allocating a GPU tensor.
374const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
375    (!portable_cuda || hopper_mma) && !no_gemm
376}
377
378// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
379// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
380// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
381// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
382// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
383// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
384// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
385const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
386const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
387const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
388const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
389const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
390
391/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
392/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
393pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
394
395/// The flash_attn fatbin matching the selected KV formats.
396fn flash_fatbin_bytes() -> &'static [u8] {
397    match kv_cache_formats() {
398        ("q8_0", "q5_1") => FLASH_FATBIN,
399        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
400        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
401        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
402        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
403        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
404        other => unreachable!("kv_cache_formats returned {other:?}"),
405    }
406}
407
408/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
409/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
410/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
411/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
412/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
413/// defaults (zero behavior change).
414fn k1_launch_override() -> Option<(u32, u32, u32)> {
415    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
416    *K1.get_or_init(|| {
417        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
418        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
419        match p.as_slice() {
420            [bm, bn, w] => Some((*bm, *bn, *w)),
421            _ => None,
422        }
423    })
424}
425
426/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
427/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
428/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
429/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
430/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
431/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
432pub(crate) fn wgmma_gemm_enabled() -> bool {
433    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
434    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
435}
436
437/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
438/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
439/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
440/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
441/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
442/// the split count changes the combine's FP summation order, and the spec verify's batched forward
443/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
444/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
445/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
446/// adaptive retries (any retry MUST pass run-spec self-consistency first).
447/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
448/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
449/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
450/// between eager decode and the verify (the spec-exactness law).
451pub const FA_VEC_MIN_TKV: usize = 96;
452/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
453/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
454/// which moves the crossover — sweep per model, adopt per the battery.
455pub fn fa_vec_min_tkv() -> usize {
456    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
457    *V.get_or_init(|| {
458        std::env::var("MEMRA_FA_VEC_MIN")
459            .ok()
460            .and_then(|v| v.parse().ok())
461            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
462    })
463}
464
465/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
466/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
467/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
468///
469/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
470/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
471/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
472/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
473/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
474/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
475pub fn fa_f16pv_on() -> bool {
476    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
477    *ON.get_or_init(|| {
478        std::env::var("MEMRA_FA_F16PV")
479            .map(|v| v != "0")
480            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
481    })
482}
483
484/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
485/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
486/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
487pub fn fa512_hp_on() -> bool {
488    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
489    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
490}
491
492/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
493/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
494/// accumulation. Even n_head and even GQA group required (guarded per call).
495pub fn faw_hp_on() -> bool {
496    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
497    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
498}
499
500/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
501/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
502/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
503pub fn fa512_wide_warps() -> usize {
504    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
505    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
506        Ok("1") => 4,
507        _ => 2,
508    })
509}
510
511/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
512/// and the gemma global-layer rows/parity call sites.
513pub fn fa512_min_tkv() -> usize {
514    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
515    *FA512_MIN.get_or_init(|| {
516        std::env::var("MEMRA_FA512_MIN")
517            .ok()
518            .and_then(|v| v.parse().ok())
519            .unwrap_or(512)
520    })
521}
522/// Per-model crossover default, set at model load BEFORE the first decode (per-model
523/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
524/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
525pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
526    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
527/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
528/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
529/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
530pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
531/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
532/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
533/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
534/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
535/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
536pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
537    std::sync::atomic::AtomicBool::new(false);
538/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
539/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
540/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
541/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
542/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
543/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
544pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
545    std::sync::atomic::AtomicBool::new(true);
546pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
547    std::sync::atomic::AtomicUsize::new(16);
548/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
549/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
550/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
551/// latency-bound at 256 threads — 7us/launch measured).
552pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
553/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
554pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
555/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
556/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
557/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
558/// explicit numerical-form seam. mmq_ffi reads this before the env.
559pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
560/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
561/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
562pub use memra_kv::KV_FP8_FORCE;
563/// bf16 matvec family block size (MEMRA_MMV_BLOCK, default 128, clamped to [64, 256] and a
564/// multiple of 32 — the f32acc twin's shared reduce caps at 256). NUMERIC-CLASS knob: the
565/// per-thread stride and reduction order change with the block, same acceptance class as
566/// MEMRA_RMS_BLOCK (fresh-tape identity + battery at the pinned value).
567pub(crate) fn mmv_block() -> u32 {
568    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
569    *V.get_or_init(|| {
570        std::env::var("MEMRA_MMV_BLOCK")
571            .ok()
572            .and_then(|v| v.parse().ok())
573            .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
574            .unwrap_or(128)
575    })
576}
577
578/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
579/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
580/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
581/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
582pub(crate) fn sig_expf_dev_on() -> bool {
583    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
584    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
585}
586
587pub(crate) fn topk_fast_on() -> bool {
588    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
589    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
590}
591
592pub(crate) fn rms_block() -> u32 {
593    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
594    *V.get_or_init(|| {
595        std::env::var("MEMRA_RMS_BLOCK")
596            .ok()
597            .and_then(|v| v.parse().ok())
598            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
599    })
600}
601
602pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
603    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
604    if let Some(forced) = *S.get_or_init(|| {
605        std::env::var("MEMRA_FA_SPLIT")
606            .ok()
607            .and_then(|v| v.parse().ok())
608            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
609    }) {
610        return forced;
611    }
612    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
613    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
614    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
615    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
616    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
617    //
618    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
619    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
620    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
621    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
622    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
623    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
624    // rig-divergence law: this branch is measured on 188 SMs only).
625    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
626    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
627    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
628    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
629    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
630        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
631    {
632        return if t_kv <= 8192 {
633            16
634        } else if t_kv <= 16384 {
635            64
636        } else {
637            128
638        };
639    }
640    let big_rig = fa_sm_count() >= 128;
641    if big_rig {
642        let _ = n_head_kv;
643        if t_kv <= 2048 {
644            // MEMRA_FA_SP_SHORT=N: the SHORT rung only (the SWA layers' capped t_kv lands
645            // here on step37: 33 of 45 layers at t_kv=512). At 16 the tile loop runs
646            // HALF-EMPTY (FA_DEC_TILE=32 -> nt=16 per split), so the V staging pass moves a
647            // half tile per iteration and the combine carries 2x the partials; 32 makes each
648            // split exactly one full tile. A global MEMRA_FA_SPLIT cannot isolate this — it
649            // moves the deep-ctx rung too, where more splits measured worse.
650            // NUMERIC-CLASS door (key partition -> different per-split partials/combine):
651            // new tape + battery, exactly like every other split-ladder change.
652            static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
653            if let Some(sp) = *SHORT.get_or_init(|| {
654                std::env::var("MEMRA_FA_SP_SHORT")
655                    .ok()
656                    .and_then(|v| v.parse().ok())
657                    .filter(|&s: &usize| s >= 8 && s % 8 == 0)
658            }) {
659                return sp;
660            }
661            16
662        } else if t_kv <= 16384 {
663            64
664        } else {
665            128
666        }
667    } else if n_head_kv <= 4 {
668        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
669        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
670        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
671        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
672        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
673        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
674        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
675        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
676        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
677        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
678        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
679        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
680        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
681        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
682        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
683        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
684        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
685        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
686        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
687        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
688        if t_kv <= 512 {
689            8
690        } else if t_kv <= 16384 {
691            64
692        } else {
693            128
694        }
695    } else {
696        if t_kv <= 8192 {
697            32
698        } else if t_kv <= 16384 {
699            64
700        } else {
701            128
702        }
703    }
704}
705
706/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
707/// same attribute Engine::batched_variant reads).
708fn fa_sm_count() -> i32 {
709    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
710    *N.get_or_init(|| {
711        cudarc::driver::result::init().ok();
712        cudarc::driver::result::device::get(0)
713            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
714                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
715            .unwrap_or(82)
716    })
717}
718
719/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
720/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
721/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
722fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
723    match head_dim {
724        256 => Ok(""),
725        128 => Ok("_hd128"),
726        d => Err(format!(
727            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
728                          callers must gate to sdpa_naive"
729        )
730        .into()),
731    }
732}
733
734/// Quant type codes matching qmatvec.cu QType enum.
735pub const QT_Q8_0: i32 = 0;
736pub const QT_Q4_K: i32 = 1;
737pub const QT_Q6_K: i32 = 2;
738pub const QT_Q5_K: i32 = 3;
739pub const QT_Q3_K: i32 = 4;
740pub const QT_IQ4_XS: i32 = 5;
741pub const QT_IQ3_S: i32 = 6;
742pub const QT_NVFP4: i32 = 7;
743/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
744/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
745/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
746/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
747/// — ONE weight copy total, no Q8_0 re-encode duplicate.
748pub const QT_F8_E4M3: i32 = 10;
749/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
750/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
751pub const QT_NVFP4_RP: i32 = 9;
752/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
753pub const QT_F32: i32 = 8;
754pub const QT_BF16: i32 = 11;
755pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
756/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
757/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
758/// dp4a/MMQ implementation exists.
759pub const QT_Q2_K: i32 = 13;
760/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
761/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
762/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
763/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
764/// scalar `scale` field is 1.0 by the layout contract.
765///
766/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
767/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
768/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
769/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
770/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
771/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
772/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
773/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
774/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
775pub const QT_F8_E4M3_BLK: i32 = 14;
776
777/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
778pub struct Engine {
779    pub gpu: memra_runtime::Gpu,
780    module: Arc<CudaModule>,
781    hybrid: Arc<CudaModule>,
782    qmatvec: Arc<CudaModule>,
783    flash: Arc<CudaModule>,
784    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
785    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
786    /// Lazy: loaded on first global-format use; None until then.
787    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
788    gemm: Arc<CudaModule>,
789    router: Arc<CudaModule>,
790    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
791    sample: Arc<CudaModule>,
792    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
793    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
794    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
795    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
796    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
797    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
798    /// the single largest block. The cache still owns every address for its full lifetime.
799    moe_cache_layout: Mutex<Option<Vec<usize>>>,
800    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
801    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
802    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
803    /// verify between replays) reuse their addresses and the replay reads/writes live memory
804    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
805    capture_keep_on: std::sync::atomic::AtomicBool,
806    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
807    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
808    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
809    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
810    verify_exact: std::sync::atomic::AtomicBool,
811    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
812    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
813    pub copy_stream: Arc<CudaStream>,
814    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
815    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
816    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
817    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
818    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
819    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
820    #[cfg(memra_cutlass)]
821    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
822    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
823    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
824    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
825    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
826    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
827    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
828    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
829    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
830    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
831    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
832    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
833    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
834    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
835    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
836    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
837    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
838    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
839    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
840    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
841    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
842    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
843    /// before capture under the generate_graph tracking-off window so it carries no events).
844    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
845    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
846    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
847    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
848    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
849    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
850    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
851    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
852    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
853    router_stage: Mutex<Option<PinnedStage>>,
854}
855
856/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
857/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
858/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
859/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
860/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
861/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
862/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
863/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
864/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
865fn fa_v2_on() -> bool {
866    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
867    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
868    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
869    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
870    // + graph bit-identity green on all three models.
871    std::env::var("MEMRA_FA_V2")
872        .map(|v| v != "0")
873        .unwrap_or(true)
874}
875
876/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
877/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
878/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
879/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
880/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
881/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
882/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
883fn fa_v3_on() -> bool {
884    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
885    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
886    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
887    std::env::var("MEMRA_FA_V3")
888        .map(|v| v != "0")
889        .unwrap_or(true)
890}
891
892/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
893/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
894/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
895/// predicate so the twins can never diverge.
896fn fa_v4_mode() -> &'static str {
897    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
898    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
899}
900fn fa_v4_on() -> bool {
901    fa_v4_mode() != "0"
902} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
903/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
904/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
905/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
906/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
907/// stays kernel-family-identical to decode at the same t_kv.
908/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
909/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
910pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
911    std::sync::atomic::AtomicUsize::new(1024);
912pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
913    std::sync::atomic::AtomicUsize::new(usize::MAX);
914pub fn fa_v4_at_pub(t_kv: usize) -> bool {
915    fa_v4_at(t_kv)
916}
917fn fa_v4_at(t_kv: usize) -> bool {
918    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
919    let mx = *M.get_or_init(|| {
920        std::env::var("MEMRA_FA_V4_MAX")
921            .ok()
922            .and_then(|v| v.parse().ok())
923            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
924    });
925    fa_v4_on() && t_kv < mx
926}
927/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
928/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
929/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
930/// (same split partition, same softmax/accumulation order, same partials/combine) and only
931/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
932/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
933/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
934/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
935/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
936/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
937/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
938/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
939/// within one process (the v2/v3 pattern).
940pub const FA_DEEP_MIN_DEFAULT: usize = 0;
941fn fa_deep_at(t_kv: usize) -> bool {
942    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
943        return false;
944    }
945    let min = std::env::var("MEMRA_FA_DEEP_MIN")
946        .ok()
947        .and_then(|v| v.parse().ok())
948        .unwrap_or(FA_DEEP_MIN_DEFAULT);
949    t_kv >= min
950}
951/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
952pub fn fa_deep_at_pub(t_kv: usize) -> bool {
953    fa_deep_at(t_kv)
954}
955
956fn fa_v3_active(head_dim: usize) -> bool {
957    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
958    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
959    fa_v3_on()
960        && head_dim % 128 == 0
961        && kv_cache_formats() == ("q8_0", "q5_1")
962        && !Engine::kv_fp8_on()
963}
964
965/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
966/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
967/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
968/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
969/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
970/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
971/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
972pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
973    std::env::var("MEMRA_NO_FA_VEC").is_err()
974        && t_kv >= fa_vec_min_tkv()
975        && head_dim == 256
976        && fa_v4_at(t_kv)
977        && !matches!(fa_v4_mode(), "noB3" | "stage")
978        && !Engine::kv_fp8_on()
979}
980/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
981pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
982    fa_split_keys(t_kv, n_head_kv)
983}
984
985/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
986/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
987/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
988/// so we allocate through `result::malloc_host` with flags=0 directly.
989struct PinnedStage {
990    ptr: *mut u8,
991    cap: usize,
992}
993unsafe impl Send for PinnedStage {}
994impl PinnedStage {
995    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
996        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
997        Ok(PinnedStage { ptr, cap })
998    }
999}
1000impl Drop for PinnedStage {
1001    fn drop(&mut self) {
1002        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1003    }
1004}
1005
1006/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
1007/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
1008pub const ARGMAX_NB: usize = 256;
1009
1010/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
1011pub(crate) use memra_fa3_vl as fa3_vl_raw;
1012
1013unsafe extern "C" {
1014    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
1015    fn memra_fa3_prefill(
1016        q16: *const core::ffi::c_void,
1017        k16: *const core::ffi::c_void,
1018        v16: *const core::ffi::c_void,
1019        o: *mut f32,
1020        t: i32,
1021        h: i32,
1022        hkv: i32,
1023        d: i32,
1024        scale: f32,
1025        stream: *mut core::ffi::c_void,
1026    ) -> i32;
1027    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
1028    pub(crate) fn memra_fa3_vl(
1029        q16s: *const *const core::ffi::c_void,
1030        k16s: *const *const core::ffi::c_void,
1031        v16s: *const *const core::ffi::c_void,
1032        os: *const *mut f32,
1033        ts: *const i32,
1034        b: i32,
1035        h: i32,
1036        hkv: i32,
1037        d: i32,
1038        scale: f32,
1039        stream: *mut core::ffi::c_void,
1040    ) -> i32;
1041}
1042
1043/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1044/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1045/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1046/// (slots are never re-allocated), so passing raw values is stable across the launch.
1047#[repr(C)]
1048#[derive(Clone, Copy)]
1049pub struct WPtr8(pub [u64; 8]);
1050unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1051
1052/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1053/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1054/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1055/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1056#[repr(C)]
1057#[derive(Clone, Copy, Default)]
1058pub struct GdnSeqVl {
1059    pub kb16: u64,
1060    pub gcum: u64,
1061    pub beta: u64,
1062    pub u: u64,
1063    pub wb16: u64,
1064    pub y: u64,
1065    pub ssnap: u64,
1066    pub state_in: u64,
1067    pub state_out: u64,
1068    pub q: u64,
1069    pub p: u64,
1070    pub o: u64,
1071    pub k: u64,
1072    pub v: u64,
1073    pub g: u64,
1074    pub a: u64,
1075    pub w: u64,
1076    pub t: i32,
1077    pub nc: i32,
1078}
1079unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1080#[repr(C)]
1081#[derive(Clone, Copy)]
1082pub struct GdnVl8(pub [GdnSeqVl; 8]);
1083unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1084
1085/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1086/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1087#[repr(C)]
1088#[derive(Clone, Copy, Default)]
1089pub struct GdnWVl {
1090    pub qb16: u64,
1091    pub pb16: u64,
1092}
1093unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1094#[repr(C)]
1095#[derive(Clone, Copy)]
1096pub struct GdnWVl8(pub [GdnWVl; 8]);
1097unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1098
1099/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1100#[repr(C)]
1101#[derive(Clone, Copy, Default)]
1102pub struct GdnPrepVl {
1103    pub qkv: u64,
1104    pub conv_state: u64,
1105    pub conv_out: u64,
1106    pub q_g: u64,
1107    pub k_g: u64,
1108    pub v_g: u64,
1109    pub q_l2: u64,
1110    pub k_l2: u64,
1111    pub beta_raw: u64,
1112    pub alpha: u64,
1113    pub beta: u64,
1114    pub g_log: u64,
1115    pub o: u64,
1116    pub z: u64,
1117    pub gn: u64,
1118    pub gn16: u64,
1119    pub kb16: u64,
1120    pub qb16: u64,
1121    pub t: i32,
1122    pub pad: i32,
1123}
1124unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1125#[repr(C)]
1126#[derive(Clone, Copy)]
1127pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1128unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1129
1130/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1131#[repr(C)]
1132#[derive(Clone, Copy, Default)]
1133pub struct FaSeqVl {
1134    pub q: u64,
1135    pub k16: u64,
1136    pub v16: u64,
1137    pub o: u64,
1138    pub kf: u64,
1139    pub vf: u64,
1140    pub t: i32,
1141    pub pad: i32,
1142}
1143unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1144#[repr(C)]
1145#[derive(Clone, Copy)]
1146pub struct FaVl8(pub [FaSeqVl; 8]);
1147unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1148
1149/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1150#[repr(C)]
1151#[derive(Clone, Copy, Default)]
1152pub struct AttnPreVl {
1153    pub qf: u64,
1154    pub kf: u64,
1155    pub vf: u64,
1156    pub q: u64,
1157    pub gate: u64,
1158    pub qn: u64,
1159    pub kn: u64,
1160    pub kc: u64,
1161    pub vc: u64,
1162    pub t: i32,
1163    pub pad: i32,
1164}
1165unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1166#[repr(C)]
1167#[derive(Clone, Copy)]
1168pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1169unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1170
1171/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1172/// varlen K1-K5 chain fills them).
1173pub struct GdnChunkBufs {
1174    pub gcum: CudaSlice<f32>,
1175    pub a: CudaSlice<f32>,
1176    pub p: CudaSlice<f32>,
1177    pub u: CudaSlice<f32>,
1178    pub w: CudaSlice<f32>,
1179    pub kb16: CudaSlice<u8>,
1180    pub wb16: CudaSlice<u8>,
1181    pub y16: CudaSlice<u8>,
1182    pub ssnap16: CudaSlice<u8>,
1183    pub qb16: CudaSlice<u8>,
1184    pub pb16: CudaSlice<u8>,
1185    pub o: CudaSlice<f32>,
1186    pub t: usize,
1187    pub nc: usize,
1188}
1189
1190/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1191#[repr(C)]
1192#[derive(Clone, Copy)]
1193pub struct F32x8(pub [f32; 8]);
1194unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1195
1196/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1197/// process. Bench binaries read it right after the call to print gen-only throughput without the
1198/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1199pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1200
1201impl Engine {
1202    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1203        let gpu = memra_runtime::Gpu::new(ordinal)?;
1204        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1205        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1206        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1207        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1208            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1209            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1210                .and_then(|d| unsafe {
1211                    Ok((
1212                        cudarc::driver::result::device::get_attribute(
1213                            d,
1214                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1215                        )?,
1216                        cudarc::driver::result::device::get_attribute(
1217                            d,
1218                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1219                        )?,
1220                    ))
1221                })
1222                .unwrap_or((0, 0));
1223            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1224            let ok = matches!(
1225                (built, maj, min),
1226                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1227            );
1228            if !ok {
1229                return Err(format!(
1230                    "memra was built for sm_{built} but device {ordinal} reports compute \
1231                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1232                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1233                )
1234                .into());
1235            }
1236        }
1237        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1238        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1239        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1240        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1241        unsafe {
1242            use cudarc::driver::sys;
1243            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1244            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1245            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1246                let mut thresh: u64 = u64::MAX;
1247                let _ = sys::cuMemPoolSetAttribute(
1248                    pool,
1249                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1250                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1251                );
1252            }
1253        }
1254        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1255        let hybrid = gpu
1256            .ctx
1257            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1258        let qmatvec = gpu
1259            .ctx
1260            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1261        let flash = gpu
1262            .ctx
1263            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1264        let gemm = gpu
1265            .ctx
1266            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1267        let router = gpu
1268            .ctx
1269            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1270        let sample = gpu
1271            .ctx
1272            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1273        let copy_stream = gpu.ctx.new_stream()?;
1274        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1275        // cudarc is in multi-stream mode (main stream +
1276        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1277        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1278        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1279        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
1280        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1281        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1282        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1283        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1284        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1285        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1286        // implicit event tracking.
1287        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1288        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1289        if std::env::var("MEMRA_EVT")
1290            .map(|v| v == "1")
1291            .unwrap_or(false)
1292        {
1293            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1294        } else {
1295            unsafe {
1296                gpu.ctx.disable_event_tracking();
1297            }
1298        }
1299        Ok(Self {
1300            gpu,
1301            module,
1302            hybrid,
1303            qmatvec,
1304            flash,
1305            flash_g: std::sync::OnceLock::new(),
1306            gemm,
1307            router,
1308            sample,
1309            moe_cache: Mutex::new(None),
1310            moe_cache_layout: Mutex::new(None),
1311            copy_stream,
1312            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1313            verify_exact: std::sync::atomic::AtomicBool::new(false),
1314            capture_keep: Mutex::new(Vec::new()),
1315            argmax_partials: Mutex::new(None),
1316            prime_deqw_ws: Mutex::new(None),
1317            router_stage: Mutex::new(None),
1318            fp8_scratch: Mutex::new(None),
1319            fa_vf16_scratch: Mutex::new(None),
1320            fa_part_pool: Mutex::new(None),
1321            fa_part_retired: Mutex::new(Vec::new()),
1322            fn_cache: Mutex::new(Default::default()),
1323            f16_scratch: Mutex::new(None),
1324            #[cfg(memra_cutlass)]
1325            cutlass_scratch: Mutex::new(None),
1326        })
1327    }
1328
1329    pub fn ctx(&self) -> &Arc<CudaContext> {
1330        &self.gpu.ctx
1331    }
1332
1333    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1334    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1335    ///
1336    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1337    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1338    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1339    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1340    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1341    ///
1342    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1343    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1344    /// under-count headroom does not belong in a gate that queues real work, but the honest
1345    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1346    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1347    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1348    ///
1349    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1350    pub fn pool_cached_bytes(&self) -> usize {
1351        let (reserved, used) = self.pool_reserved_used();
1352        reserved.saturating_sub(used)
1353    }
1354
1355    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1356    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1357    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1358    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1359    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1360    /// (0, 0) if the pool cannot be queried.
1361    pub fn pool_reserved_used(&self) -> (usize, usize) {
1362        use cudarc::driver::sys;
1363        unsafe {
1364            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1365            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1366                != sys::CUresult::CUDA_SUCCESS
1367            {
1368                return (0, 0);
1369            }
1370            let (mut reserved, mut used) = (0u64, 0u64);
1371            if sys::cuMemPoolGetAttribute(
1372                pool,
1373                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1374                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1375            ) != sys::CUresult::CUDA_SUCCESS
1376            {
1377                return (0, 0);
1378            }
1379            if sys::cuMemPoolGetAttribute(
1380                pool,
1381                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1382                &mut used as *mut u64 as *mut core::ffi::c_void,
1383            ) != sys::CUresult::CUDA_SUCCESS
1384            {
1385                return (0, 0);
1386            }
1387            (reserved as usize, used as usize)
1388        }
1389    }
1390
1391    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1392    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1393    pub fn stream(&self) -> Arc<CudaStream> {
1394        self.gpu.stream()
1395    }
1396    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1397    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1398    pub fn gkv_on() -> bool {
1399        memra_kv::gkv_on()
1400    }
1401
1402    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1403    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1404    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1405    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1406    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1407    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1408    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1409    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1410    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1411    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1412    /// ON for both — no acceptance cost measured.
1413    pub fn wkv_on() -> bool {
1414        memra_kv::wkv_on()
1415    }
1416
1417    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1418    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1419    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1420    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1421    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1422    pub fn kv_fp8_on() -> bool {
1423        memra_kv::kv_fp8_on()
1424    }
1425
1426    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1427    /// when the fp8-globals arm is on; everything else from the default flash module.
1428    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1429        if head_dim == 512 && Self::gkv_on() {
1430            self.func_g(name)
1431        } else {
1432            self.func(name)
1433        }
1434    }
1435
1436    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1437    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1438    /// per-format fatbins; fall back to the base modules for those.
1439    fn func_g(&self, name: &str) -> CudaFunction {
1440        let m = self.flash_g.get_or_init(|| {
1441            self.gpu
1442                .ctx
1443                .load_module(cudarc::nvrtc::Ptx::from_binary(
1444                    FLASH_FATBIN_KF8VF8.to_vec(),
1445                ))
1446                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1447        });
1448        let key = format!("g:{name}");
1449        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1450            return f.clone();
1451        }
1452        let f = match m.load_function(name) {
1453            Ok(f) => f,
1454            Err(_) => self.func(name),
1455        };
1456        self.fn_cache.lock().unwrap().insert(key, f.clone());
1457        f
1458    }
1459
1460    fn func(&self, name: &str) -> CudaFunction {
1461        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1462        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1463        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1464            return f.clone();
1465        }
1466        let f = self
1467            .module
1468            .load_function(name)
1469            .or_else(|_| self.hybrid.load_function(name))
1470            .or_else(|_| self.qmatvec.load_function(name))
1471            .or_else(|_| self.flash.load_function(name))
1472            .or_else(|_| self.gemm.load_function(name))
1473            .or_else(|_| self.router.load_function(name))
1474            .or_else(|_| self.sample.load_function(name))
1475            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1476        self.fn_cache
1477            .lock()
1478            .unwrap()
1479            .insert(name.to_string(), f.clone());
1480        f
1481    }
1482
1483    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1484    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1485    pub fn scatter_trim_logits(
1486        &self,
1487        src: &CudaSlice<f32>,
1488        d2t: &CudaSlice<u32>,
1489        dst: &mut CudaSlice<f32>,
1490        d_vocab: usize,
1491        n_vocab: usize,
1492    ) -> Result<(), Box<dyn std::error::Error>> {
1493        let f1 = self.func("scatter_trim_logits_f32");
1494        let f2 = self.func("scatter_trim_logits_pass2_f32");
1495        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1496        let cfg1 = LaunchConfig {
1497            grid_dim: (256, 1, 1),
1498            block_dim: (256, 1, 1),
1499            shared_mem_bytes: 0,
1500        };
1501        let __s_b1 = self.gpu.stream();
1502        let mut b1 = __s_b1.launch_builder(&f1);
1503        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1504        unsafe {
1505            b1.launch(cfg1)?;
1506        }
1507        let cfg2 = LaunchConfig {
1508            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1509            block_dim: (256, 1, 1),
1510            shared_mem_bytes: 0,
1511        };
1512        let __s_b2 = self.gpu.stream();
1513        let mut b2 = __s_b2.launch_builder(&f2);
1514        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1515        unsafe {
1516            b2.launch(cfg2)?;
1517        }
1518        Ok(())
1519    }
1520
1521    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1522    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1523
1524    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1525    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1526    #[allow(clippy::too_many_arguments)]
1527    pub fn filter_stats(
1528        &self,
1529        x: &CudaSlice<f32>,
1530        row_stride: usize,
1531        rows: &CudaSlice<i32>,
1532        out_th: &mut CudaSlice<f32>,
1533        out_z: &mut CudaSlice<f32>,
1534        out_max: &mut CudaSlice<f32>,
1535        n: usize,
1536        nrow: usize,
1537        temp: f32,
1538        top_k: i32,
1539        top_p: f32,
1540        min_p: f32,
1541    ) -> Result<(), Box<dyn std::error::Error>> {
1542        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1543        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1544        // L2-resident, so the extra passes are near-free while the per-thread selection list
1545        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1546        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1547        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1548        //
1549        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1550        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1551        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1552        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1553        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1554        // Admission: cooperative grid must co-reside (16*nrow blocks vs SM count).
1555        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1556        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1557        let coop_on =
1558            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1559        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1560        if coop_on && 16 * nrow <= self.sm_count() as usize {
1561            let f = self.func("filter_stats_coop_f32");
1562            let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1563            let cfg = LaunchConfig {
1564                grid_dim: (16, nrow as u32, 1),
1565                block_dim: (512, 1, 1),
1566                shared_mem_bytes: 0,
1567            };
1568            let __s_b = self.gpu.stream();
1569            let mut b = __s_b.launch_builder(&f);
1570            b.arg(x)
1571                .arg(&rs)
1572                .arg(rows)
1573                .arg(&mut *out_th)
1574                .arg(&mut *out_z)
1575                .arg(&mut *out_max)
1576                .arg(&mut ws)
1577                .arg(&ni)
1578                .arg(&nr)
1579                .arg(&temp)
1580                .arg(&top_k)
1581                .arg(&top_p)
1582                .arg(&min_p);
1583            unsafe {
1584                b.launch_cooperative(cfg)?;
1585            }
1586            return Ok(());
1587        }
1588        let f = self.func("filter_stats_f32");
1589        let cfg = LaunchConfig {
1590            grid_dim: (nrow as u32, 1, 1),
1591            block_dim: (1024, 1, 1),
1592            shared_mem_bytes: 0,
1593        };
1594        let __s_b = self.gpu.stream();
1595        let mut b = __s_b.launch_builder(&f);
1596        b.arg(x)
1597            .arg(&rs)
1598            .arg(rows)
1599            .arg(&mut *out_th)
1600            .arg(&mut *out_z)
1601            .arg(&mut *out_max)
1602            .arg(&ni)
1603            .arg(&nr)
1604            .arg(&temp)
1605            .arg(&top_k)
1606            .arg(&top_p)
1607            .arg(&min_p);
1608        unsafe {
1609            b.launch(cfg)?;
1610        }
1611        Ok(())
1612    }
1613
1614    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1615    #[allow(clippy::too_many_arguments)]
1616    pub fn softmax_gather_filtered(
1617        &self,
1618        x: &CudaSlice<f32>,
1619        row_stride: usize,
1620        ids: &CudaSlice<u32>,
1621        rows: &CudaSlice<i32>,
1622        th: &CudaSlice<f32>,
1623        z: &CudaSlice<f32>,
1624        out: &mut CudaSlice<f32>,
1625        n: usize,
1626        npair: usize,
1627        temp: f32,
1628    ) -> Result<(), Box<dyn std::error::Error>> {
1629        let f = self.func("softmax_gather_filtered_f32");
1630        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1631        let cfg = LaunchConfig {
1632            grid_dim: (npair as u32, 1, 1),
1633            block_dim: (256, 1, 1),
1634            shared_mem_bytes: 0,
1635        };
1636        let __s_b = self.gpu.stream();
1637        let mut b = __s_b.launch_builder(&f);
1638        b.arg(x)
1639            .arg(&rs)
1640            .arg(ids)
1641            .arg(rows)
1642            .arg(th)
1643            .arg(z)
1644            .arg(&mut *out)
1645            .arg(&ni)
1646            .arg(&np)
1647            .arg(&temp);
1648        unsafe {
1649            b.launch(cfg)?;
1650        }
1651        Ok(())
1652    }
1653
1654    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1655    #[allow(clippy::too_many_arguments)]
1656    pub fn residual_sample_filtered(
1657        &self,
1658        p: &CudaSlice<f32>,
1659        q: Option<&CudaSlice<f32>>,
1660        n: usize,
1661        temp: f32,
1662        seed: u64,
1663        stream_pos: u32,
1664        p_stats: (f32, f32, f32),
1665        q_stats: (f32, f32, f32),
1666        out_tok: &mut CudaSlice<u32>,
1667    ) -> Result<(), Box<dyn std::error::Error>> {
1668        let f = self.func("residual_sample_filtered_f32");
1669        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1670        let has_q: i32 = q.is_some() as i32;
1671        let qbuf = q.unwrap_or(p);
1672        let (pm, pth, pz) = p_stats;
1673        let (qm, qth, qz) = q_stats;
1674        let cfg = LaunchConfig {
1675            grid_dim: (1, 1, 1),
1676            block_dim: (1024, 1, 1),
1677            shared_mem_bytes: 0,
1678        };
1679        let __s_b = self.gpu.stream();
1680        let mut b = __s_b.launch_builder(&f);
1681        b.arg(p)
1682            .arg(qbuf)
1683            .arg(&has_q)
1684            .arg(&ni)
1685            .arg(&temp)
1686            .arg(&slo)
1687            .arg(&shi)
1688            .arg(&stream_pos)
1689            .arg(&pm)
1690            .arg(&pth)
1691            .arg(&pz)
1692            .arg(&qm)
1693            .arg(&qth)
1694            .arg(&qz)
1695            .arg(&mut *out_tok);
1696        unsafe {
1697            b.launch(cfg)?;
1698        }
1699        Ok(())
1700    }
1701
1702    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1703    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1704    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1705    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1706    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1707    #[allow(clippy::too_many_arguments)]
1708    pub fn residual_sample_sparse_q(
1709        &self,
1710        p: &CudaSlice<f32>,
1711        cand_ids: &CudaSlice<u32>,
1712        q_probs: &CudaSlice<f32>,
1713        n_cand: usize,
1714        n: usize,
1715        temp: f32,
1716        seed: u64,
1717        stream_pos: u32,
1718        p_stats: (f32, f32, f32),
1719        out_tok: &mut CudaSlice<u32>,
1720    ) -> Result<(), Box<dyn std::error::Error>> {
1721        assert!(
1722            n_cand >= 1 && n_cand <= 32,
1723            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1724        );
1725        let f = self.func("residual_sample_sparse_q_f32");
1726        let (ni, nc) = (n as i32, n_cand as i32);
1727        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1728        let (pm, pth, pz) = p_stats;
1729        let cfg = LaunchConfig {
1730            grid_dim: (1, 1, 1),
1731            block_dim: (1024, 1, 1),
1732            shared_mem_bytes: 0,
1733        };
1734        let __s_b = self.gpu.stream();
1735        let mut b = __s_b.launch_builder(&f);
1736        b.arg(p)
1737            .arg(cand_ids)
1738            .arg(q_probs)
1739            .arg(&nc)
1740            .arg(&ni)
1741            .arg(&temp)
1742            .arg(&slo)
1743            .arg(&shi)
1744            .arg(&stream_pos)
1745            .arg(&pm)
1746            .arg(&pth)
1747            .arg(&pz)
1748            .arg(&mut *out_tok);
1749        unsafe {
1750            b.launch(cfg)?;
1751        }
1752        Ok(())
1753    }
1754
1755    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1756    #[allow(clippy::too_many_arguments)]
1757    pub fn gumbel_perturb_filtered(
1758        &self,
1759        x: &CudaSlice<f32>,
1760        y: &mut CudaSlice<f32>,
1761        n: usize,
1762        seed: u64,
1763        stream_pos: u32,
1764        temp: f32,
1765        row_max: f32,
1766        th: f32,
1767    ) -> Result<(), Box<dyn std::error::Error>> {
1768        let f = self.func("gumbel_perturb_filtered_f32");
1769        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1770        let cfg = LaunchConfig {
1771            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1772            block_dim: (256, 1, 1),
1773            shared_mem_bytes: 0,
1774        };
1775        let __s_b = self.gpu.stream();
1776        let mut b = __s_b.launch_builder(&f);
1777        b.arg(x)
1778            .arg(&mut *y)
1779            .arg(&ni)
1780            .arg(&slo)
1781            .arg(&shi)
1782            .arg(&stream_pos)
1783            .arg(&temp)
1784            .arg(&row_max)
1785            .arg(&th);
1786        unsafe {
1787            b.launch(cfg)?;
1788        }
1789        Ok(())
1790    }
1791
1792    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1793    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1794    /// filtered rejection sampling exact for the penalized target.
1795    #[allow(clippy::too_many_arguments)]
1796    pub fn penalize_logits(
1797        &self,
1798        x: &mut CudaSlice<f32>,
1799        hist: &CudaSlice<u32>,
1800        n_hist: usize,
1801        rep: f32,
1802        freq: f32,
1803        present: f32,
1804        n: usize,
1805    ) -> Result<(), Box<dyn std::error::Error>> {
1806        if n_hist == 0 {
1807            return Ok(());
1808        }
1809        let f = self.func("penalize_logits_f32");
1810        let (nh, ni) = (n_hist as i32, n as i32);
1811        let cfg = LaunchConfig {
1812            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1813            block_dim: (128, 1, 1),
1814            shared_mem_bytes: 0,
1815        };
1816        let __s_b = self.gpu.stream();
1817        let mut b = __s_b.launch_builder(&f);
1818        b.arg(&mut *x)
1819            .arg(hist)
1820            .arg(&nh)
1821            .arg(&rep)
1822            .arg(&freq)
1823            .arg(&present)
1824            .arg(&ni);
1825        unsafe {
1826            b.launch(cfg)?;
1827        }
1828        Ok(())
1829    }
1830
1831    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1832    #[allow(clippy::too_many_arguments)]
1833    pub fn penalize_logits_rows(
1834        &self,
1835        x: &mut CudaSlice<f32>,
1836        hist: &CudaSlice<u32>,
1837        n_hist: usize,
1838        rep: f32,
1839        freq: f32,
1840        present: f32,
1841        n: usize,
1842        nrow: usize,
1843    ) -> Result<(), Box<dyn std::error::Error>> {
1844        if n_hist == 0 || nrow == 0 {
1845            return Ok(());
1846        }
1847        let f = self.func("penalize_logits_rows_f32");
1848        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1849        let cfg = LaunchConfig {
1850            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1851            block_dim: (128, 1, 1),
1852            shared_mem_bytes: 0,
1853        };
1854        let __s_b = self.gpu.stream();
1855        let mut b = __s_b.launch_builder(&f);
1856        b.arg(&mut *x)
1857            .arg(hist)
1858            .arg(&nh)
1859            .arg(&rep)
1860            .arg(&freq)
1861            .arg(&present)
1862            .arg(&ni)
1863            .arg(&nr);
1864        unsafe {
1865            b.launch(cfg)?;
1866        }
1867        Ok(())
1868    }
1869
1870    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
1871    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
1872    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
1873    /// is the within-round evolving penalty state block drafting needs: verify row r's
1874    /// target is penalized by every token committed before it INCLUDING same-round
1875    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
1876    /// approximation this exists to replace on the dspark route.
1877    #[allow(clippy::too_many_arguments)]
1878    pub fn penalize_logits_rows_inc(
1879        &self,
1880        x: &mut CudaSlice<f32>,
1881        hist: &CudaSlice<u32>,
1882        n_hist0: usize,
1883        rep: f32,
1884        freq: f32,
1885        present: f32,
1886        n: usize,
1887        nrow: usize,
1888        win: usize,
1889    ) -> Result<(), Box<dyn std::error::Error>> {
1890        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1891            return Ok(());
1892        }
1893        debug_assert!(
1894            hist.len() >= n_hist0 + nrow - 1,
1895            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1896        );
1897        let f = self.func("penalize_logits_rows_inc_f32");
1898        let max_len = win.min(n_hist0 + nrow - 1).max(1);
1899        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1900        let cfg = LaunchConfig {
1901            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1902            block_dim: (128, 1, 1),
1903            shared_mem_bytes: 0,
1904        };
1905        let __s_b = self.gpu.stream();
1906        let mut b = __s_b.launch_builder(&f);
1907        b.arg(&mut *x)
1908            .arg(hist)
1909            .arg(&nh)
1910            .arg(&rep)
1911            .arg(&freq)
1912            .arg(&present)
1913            .arg(&ni)
1914            .arg(&nr)
1915            .arg(&wi);
1916        unsafe {
1917            b.launch(cfg)?;
1918        }
1919        Ok(())
1920    }
1921
1922    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1923    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1924    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1925    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1926    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1927    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1928    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1929    pub fn wpf_level() -> u32 {
1930        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1931        *ON.get_or_init(|| {
1932            std::env::var("MEMRA_WPF")
1933                .ok()
1934                .and_then(|v| v.parse().ok())
1935                .unwrap_or(1)
1936        })
1937    }
1938
1939    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1940    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1941    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1942    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1943    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1944    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1945    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1946    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1947    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1948    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1949    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1950    pub fn set_verify_exact(&self, on: bool) {
1951        self.verify_exact
1952            .store(on, std::sync::atomic::Ordering::Relaxed);
1953    }
1954    pub(crate) fn verify_exact_on(&self) -> bool {
1955        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1956    }
1957
1958    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1959    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1960    pub fn qkv_append_on() -> bool {
1961        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1962        *ON.get_or_init(|| {
1963            std::env::var("MEMRA_QKV_APPEND")
1964                .map(|v| v != "0")
1965                .unwrap_or(true)
1966        })
1967    }
1968
1969    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1970    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1971    pub fn pdl_wb_on() -> bool {
1972        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1973        *ON.get_or_init(|| {
1974            std::env::var("MEMRA_PDL_WB")
1975                .map(|v| v != "0")
1976                .unwrap_or(true)
1977        })
1978    }
1979
1980    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
1981    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
1982    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
1983    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
1984    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
1985    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
1986    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
1987    pub fn norm_ilp_on() -> bool {
1988        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1989        *ON.get_or_init(|| {
1990            std::env::var("MEMRA_NORM_ILP")
1991                .map(|v| v != "0")
1992                .unwrap_or(true)
1993        })
1994    }
1995
1996    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
1997    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
1998    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
1999    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
2000    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
2001    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
2002    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
2003    pub fn tk_ffn_dual_on() -> bool {
2004        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2005        *ON.get_or_init(|| {
2006            std::env::var("MEMRA_TK_FFN_DUAL")
2007                .map(|v| v != "0")
2008                .unwrap_or(true)
2009        })
2010    }
2011
2012    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
2013    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
2014    /// per-model no-harm bisect knob.
2015    pub fn pdl_mmvq_on() -> bool {
2016        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2017        *ON.get_or_init(|| {
2018            std::env::var("MEMRA_PDL_MMVQ")
2019                .map(|v| v != "0")
2020                .unwrap_or(true)
2021        })
2022    }
2023
2024    pub fn pdl_on() -> bool {
2025        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2026        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2027    }
2028
2029    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
2030    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
2031    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
2032    /// on the producer before any read), bit-identical by construction.
2033    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
2034    pub fn pdl_nvfp4q8_on() -> bool {
2035        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2036        *ON.get_or_init(|| {
2037            std::env::var("MEMRA_PDL_NVFP4")
2038                .map(|v| v != "0")
2039                .unwrap_or(true)
2040        })
2041    }
2042
2043    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
2044    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
2045    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
2046    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
2047    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
2048    fn q40_mr1_on() -> bool {
2049        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2050        match *Q40MR.get_or_init(|| {
2051            std::env::var("MEMRA_Q40_MR")
2052                .ok()
2053                .and_then(|v| v.parse().ok())
2054        }) {
2055            Some(v) => v == 1,
2056            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2057        }
2058    }
2059
2060    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
2061    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
2062    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
2063    /// writes wrong bytes silently.
2064    fn pdl_func_flash(
2065        &self,
2066        g: bool,
2067        name: &'static str,
2068    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2069        use cudarc::driver::sys as cu;
2070        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
2071        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
2072        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
2073        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
2074        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
2075        // this engine's CUcontext; single-context runs behave exactly as before.
2076        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2077            std::sync::Mutex::new(None);
2078        static FNS: std::sync::Mutex<
2079            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2080        > = std::sync::Mutex::new(None);
2081        let ctx_key = self.ctx().cu_ctx() as usize;
2082        if let Some(&f) = FNS
2083            .lock()
2084            .unwrap()
2085            .get_or_insert_with(Default::default)
2086            .get(&(ctx_key, g, name))
2087        {
2088            return Ok(f as cu::CUfunction);
2089        }
2090        let module = {
2091            let mut mods = MODS.lock().unwrap();
2092            let map = mods.get_or_insert_with(Default::default);
2093            match map.get(&(ctx_key, g)) {
2094                Some(&m) => m,
2095                None => {
2096                    let m = self.pdl_load_module_in_ctx(if g {
2097                        FLASH_FATBIN_KF8VF8
2098                    } else {
2099                        FLASH_FATBIN
2100                    })?;
2101                    map.insert((ctx_key, g), m);
2102                    m
2103                }
2104            }
2105        };
2106        let cname = std::ffi::CString::new(name)?;
2107        let mut f: cu::CUfunction = std::ptr::null_mut();
2108        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2109        if r != cu::CUresult::CUDA_SUCCESS {
2110            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2111        }
2112        FNS.lock()
2113            .unwrap()
2114            .get_or_insert_with(Default::default)
2115            .insert((ctx_key, g, name), f as usize);
2116        Ok(f)
2117    }
2118
2119    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2120    /// the module to the thread's CURRENT context — a remote-stage engine must not
2121    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2122    /// current context before returning.
2123    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2124        use cudarc::driver::sys as cu;
2125        let mut prev: cu::CUcontext = std::ptr::null_mut();
2126        unsafe {
2127            cu::cuCtxGetCurrent(&mut prev).result()?;
2128        }
2129        self.ctx().bind_to_thread()?;
2130        let mut m: cu::CUmodule = std::ptr::null_mut();
2131        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2132        let restore = if prev.is_null() {
2133            cu::CUresult::CUDA_SUCCESS
2134        } else {
2135            unsafe { cu::cuCtxSetCurrent(prev) }
2136        };
2137        if r != cu::CUresult::CUDA_SUCCESS {
2138            return Err(format!("pdl module load: {r:?}").into());
2139        }
2140        if restore != cu::CUresult::CUDA_SUCCESS {
2141            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2142        }
2143        Ok(m as usize)
2144    }
2145
2146    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
2147    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
2148    pub fn raw_kernel_function(
2149        &self,
2150        name: &'static str,
2151    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2152        self.pdl_func(name)
2153    }
2154
2155    fn pdl_func(
2156        &self,
2157        name: &'static str,
2158    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2159        use cudarc::driver::sys as cu;
2160        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2161        // are context-scoped; key everything by this engine's CUcontext).
2162        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2163            std::sync::Mutex::new(None);
2164        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2165        // duplicate module, loaded lazily on the first kernels-module miss.
2166        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2167            std::sync::Mutex::new(None);
2168        static FNS: std::sync::Mutex<
2169            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2170        > = std::sync::Mutex::new(None);
2171        let ctx_key = self.ctx().cu_ctx() as usize;
2172        if let Some(&f) = FNS
2173            .lock()
2174            .unwrap()
2175            .get_or_insert_with(Default::default)
2176            .get(&(ctx_key, name))
2177        {
2178            return Ok(f as cu::CUfunction);
2179        }
2180        let module = {
2181            let mut mods = MODULES.lock().unwrap();
2182            let map = mods.get_or_insert_with(Default::default);
2183            match map.get(&ctx_key) {
2184                Some(&m) => m,
2185                None => {
2186                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2187                    map.insert(ctx_key, m);
2188                    m
2189                }
2190            }
2191        };
2192        let cname = std::ffi::CString::new(name)?;
2193        let mut f: cu::CUfunction = std::ptr::null_mut();
2194        let mut r =
2195            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2196        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2197            let qmodule = {
2198                let mut mods = QMODULES.lock().unwrap();
2199                let map = mods.get_or_insert_with(Default::default);
2200                match map.get(&ctx_key) {
2201                    Some(&m) => m,
2202                    None => {
2203                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2204                        map.insert(ctx_key, m);
2205                        m
2206                    }
2207                }
2208            };
2209            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2210        }
2211        if r != cu::CUresult::CUDA_SUCCESS {
2212            return Err(format!("pdl_func {name}: {r:?}").into());
2213        }
2214        FNS.lock()
2215            .unwrap()
2216            .get_or_insert_with(Default::default)
2217            .insert((ctx_key, name), f as usize);
2218        Ok(f)
2219    }
2220
2221    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2222    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2223    ///
2224    /// # Safety
2225    /// `params` must match the kernel's exact parameter list (order, types, count) —
2226    /// a mismatch corrupts the launch silently.
2227    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2228    /// builder path's fa_func/func_g choice exactly).
2229    ///
2230    /// # Safety
2231    /// Same contract as `launch_pdl`.
2232    unsafe fn launch_pdl_flash(
2233        &self,
2234        g: bool,
2235        name: &'static str,
2236        grid: (u32, u32, u32),
2237        block: (u32, u32, u32),
2238        smem: u32,
2239        params: &mut [*mut std::ffi::c_void],
2240    ) -> Result<(), Box<dyn std::error::Error>> {
2241        use cudarc::driver::sys as cu;
2242        let f = self.pdl_func_flash(g, name)?;
2243        if smem > 0 {
2244            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2245            let r =
2246                unsafe {
2247                    cu::cuFuncSetAttribute(f,
2248                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2249                smem as i32)
2250                };
2251            if r != cu::CUresult::CUDA_SUCCESS {
2252                return Err(format!("pdl smem attr {name}: {r:?}").into());
2253            }
2254        }
2255        let mut attr = cu::CUlaunchAttribute {
2256            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2257            pad: [0; 4],
2258            value: cu::CUlaunchAttributeValue {
2259                programmaticStreamSerializationAllowed: 1,
2260            },
2261        };
2262        let cfg = cu::CUlaunchConfig {
2263            gridDimX: grid.0,
2264            gridDimY: grid.1,
2265            gridDimZ: grid.2,
2266            blockDimX: block.0,
2267            blockDimY: block.1,
2268            blockDimZ: block.2,
2269            sharedMemBytes: smem,
2270            hStream: self.gpu.stream().cu_stream(),
2271            attrs: &mut attr,
2272            numAttrs: 1,
2273        };
2274        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2275        if r != cu::CUresult::CUDA_SUCCESS {
2276            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2277        }
2278        Ok(())
2279    }
2280
2281    unsafe fn launch_pdl(
2282        &self,
2283        name: &'static str,
2284        grid: (u32, u32, u32),
2285        block: (u32, u32, u32),
2286        params: &mut [*mut std::ffi::c_void],
2287    ) -> Result<(), Box<dyn std::error::Error>> {
2288        use cudarc::driver::sys as cu;
2289        let f = self.pdl_func(name)?;
2290        let mut attr = cu::CUlaunchAttribute {
2291            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2292            pad: [0; 4],
2293            value: cu::CUlaunchAttributeValue {
2294                programmaticStreamSerializationAllowed: 1,
2295            },
2296        };
2297        let cfg = cu::CUlaunchConfig {
2298            gridDimX: grid.0,
2299            gridDimY: grid.1,
2300            gridDimZ: grid.2,
2301            blockDimX: block.0,
2302            blockDimY: block.1,
2303            blockDimZ: block.2,
2304            sharedMemBytes: 0,
2305            hStream: self.gpu.stream().cu_stream(),
2306            attrs: &mut attr,
2307            numAttrs: 1,
2308        };
2309        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2310        if r != cu::CUresult::CUDA_SUCCESS {
2311            return Err(format!("launch_pdl {name}: {r:?}").into());
2312        }
2313        Ok(())
2314    }
2315
2316    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2317    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2318    pub fn prefetch_weight_l2(
2319        &self,
2320        w: &crate::model::GpuTensor,
2321    ) -> Result<(), Box<dyn std::error::Error>> {
2322        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2323            let p = rp4.as_ref().unwrap_or(bytes);
2324            self.prefetch_l2(p, p.len())?;
2325        }
2326        Ok(())
2327    }
2328
2329    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2330    /// by the DEVICE token id at tok[idx] into f32.
2331    pub fn gather_row_bf16(
2332        &self,
2333        table: &CudaSlice<u8>,
2334        tok: &CudaSlice<u32>,
2335        idx: usize,
2336        dst: &mut CudaSlice<f32>,
2337        ncols: usize,
2338    ) -> Result<(), Box<dyn std::error::Error>> {
2339        let f = self.func("gather_row_bf16_f32");
2340        let cfg = LaunchConfig {
2341            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2342            block_dim: (256, 1, 1),
2343            shared_mem_bytes: 0,
2344        };
2345        let (nc, ix) = (ncols as i32, idx as i32);
2346        let __s_b = self.gpu.stream();
2347        let mut b = __s_b.launch_builder(&f);
2348        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2349        unsafe {
2350            b.launch(cfg)?;
2351        }
2352        Ok(())
2353    }
2354
2355    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2356    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2357    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2358    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2359    /// finish(1).
2360    #[allow(clippy::too_many_arguments)]
2361    pub fn dflash2_dynconv(
2362        &self,
2363        x: &CudaSlice<f32>,
2364        dyn_: &CudaSlice<f32>,
2365        base: &CudaSlice<f32>,
2366        out: &mut CudaSlice<f32>,
2367        rows: usize,
2368        hidden: usize,
2369        group_size: usize,
2370        ksize: usize,
2371        half: usize,
2372    ) -> Result<(), Box<dyn std::error::Error>> {
2373        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2374        let f = self.func("dflash2_dynconv_f32");
2375        let n = rows * hidden;
2376        let cfg = LaunchConfig {
2377            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2378            block_dim: (256, 1, 1),
2379            shared_mem_bytes: 0,
2380        };
2381        let (ri, hi, gi, ki, hf) = (
2382            rows as i32,
2383            hidden as i32,
2384            group_size as i32,
2385            ksize as i32,
2386            half as i32,
2387        );
2388        let __s_b = self.gpu.stream();
2389        let mut b = __s_b.launch_builder(&f);
2390        b.arg(x)
2391            .arg(dyn_)
2392            .arg(base)
2393            .arg(out)
2394            .arg(&ri)
2395            .arg(&hi)
2396            .arg(&gi)
2397            .arg(&ki)
2398            .arg(&hf);
2399        unsafe {
2400            b.launch(cfg)?;
2401        }
2402        Ok(())
2403    }
2404
2405    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2406    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2407    /// value-descending, ties to the lower index.
2408    pub fn topk_rows(
2409        &self,
2410        logits: &CudaSlice<f32>,
2411        n_rows: usize,
2412        n_cols: usize,
2413        k: usize,
2414    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2415        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2416        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2417        let f = self.func("topk_rows_f32");
2418        let nth = 256usize;
2419        let mut vals = self.uninit(n_rows * k)?;
2420        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2421        let cfg = LaunchConfig {
2422            grid_dim: (n_rows as u32, 1, 1),
2423            block_dim: (nth as u32, 1, 1),
2424            shared_mem_bytes: (nth * k * 8) as u32,
2425        };
2426        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2427        let __s_b = self.gpu.stream();
2428        let mut b = __s_b.launch_builder(&f);
2429        b.arg(logits)
2430            .arg(&nr)
2431            .arg(&nc)
2432            .arg(&ki)
2433            .arg(&mut vals)
2434            .arg(&mut idxs);
2435        unsafe {
2436            b.launch(cfg)?;
2437        }
2438        Ok((vals, idxs))
2439    }
2440
2441    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2442    pub fn add_row_inplace(
2443        &self,
2444        logits: &mut CudaSlice<f32>,
2445        bias: &CudaSlice<f32>,
2446        n: usize,
2447        row_off: usize,
2448    ) -> Result<(), Box<dyn std::error::Error>> {
2449        let f = self.func("add_row_inplace_f32");
2450        let cfg = LaunchConfig {
2451            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2452            block_dim: (256, 1, 1),
2453            shared_mem_bytes: 0,
2454        };
2455        let (ni, off) = (n as i32, row_off as i64);
2456        let __s_b = self.gpu.stream();
2457        let mut b = __s_b.launch_builder(&f);
2458        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2459        unsafe {
2460            b.launch(cfg)?;
2461        }
2462        Ok(())
2463    }
2464
2465    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2466    pub fn prefetch_l2(
2467        &self,
2468        p: &CudaSlice<u8>,
2469        n: usize,
2470    ) -> Result<(), Box<dyn std::error::Error>> {
2471        let f = self.func("prefetch_l2_bytes");
2472        let lines = n.div_ceil(128);
2473        let ni = n as i64;
2474        let cfg = LaunchConfig {
2475            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2476            block_dim: (256, 1, 1),
2477            shared_mem_bytes: 0,
2478        };
2479        let __s_b = self.gpu.stream();
2480        let mut b = __s_b.launch_builder(&f);
2481        b.arg(p).arg(&ni);
2482        unsafe {
2483            b.launch(cfg)?;
2484        }
2485        Ok(())
2486    }
2487
2488    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2489    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2490    pub fn router_gemv(
2491        &self,
2492        w: &CudaSlice<f32>,
2493        x: &CudaSlice<f32>,
2494        n_embd: usize,
2495        n_experts: usize,
2496        t: usize,
2497    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2498        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2499        // stream differs) — too small to justify a numeric config change; deleted.
2500        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2501        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2502        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2503        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2504            Ok("0") => false,
2505            Ok(_) => true,
2506            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2507        };
2508        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2509        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2510        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2511        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2512        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2513        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2514        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2515        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2516        // (perf-only, bits equal).
2517        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2518        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2519    }
2520
2521    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2522    /// force both forms; `batch` requires `w8`).
2523    pub fn router_gemv_form(
2524        &self,
2525        w: &CudaSlice<f32>,
2526        x: &CudaSlice<f32>,
2527        n_embd: usize,
2528        n_experts: usize,
2529        t: usize,
2530        w8: bool,
2531        batch: bool,
2532    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2533        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2534        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2535        let f = if batch {
2536            self.func("router_gemv_f32_w8_batch")
2537        } else if w8 {
2538            self.func("router_gemv_f32_w8")
2539        } else {
2540            self.func("router_gemv_f32")
2541        };
2542        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2543        let cfg = if batch {
2544            LaunchConfig {
2545                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2546                block_dim: (32, 8, 1),
2547                shared_mem_bytes: 0,
2548            }
2549        } else {
2550            LaunchConfig {
2551                grid_dim: (n_experts as u32, t as u32, 1),
2552                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2553                shared_mem_bytes: 0,
2554            }
2555        };
2556        let __s_b = self.gpu.stream();
2557        let mut b = __s_b.launch_builder(&f);
2558        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2559        unsafe {
2560            b.launch(cfg)?;
2561        }
2562        Ok(y)
2563    }
2564
2565    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
2566    /// buffer — token-graph alloc-free.
2567    pub fn router_gemv_into(
2568        &self,
2569        w: &CudaSlice<f32>,
2570        x: &CudaSlice<f32>,
2571        y: &mut CudaSlice<f32>,
2572        n_embd: usize,
2573        n_experts: usize,
2574        t: usize,
2575    ) -> Result<(), Box<dyn std::error::Error>> {
2576        if y.len() < t * n_experts {
2577            return Err("router_gemv_into output too small".into());
2578        }
2579        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2580            Ok("0") => false,
2581            Ok(_) => true,
2582            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2583        };
2584        let f = if w8 {
2585            self.func("router_gemv_f32_w8")
2586        } else {
2587            self.func("router_gemv_f32")
2588        };
2589        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2590        let cfg = LaunchConfig {
2591            grid_dim: (n_experts as u32, t as u32, 1),
2592            block_dim: (32, if w8 { 8 } else { 1 }, 1),
2593            shared_mem_bytes: 0,
2594        };
2595        let __s_b = self.gpu.stream();
2596        let mut b = __s_b.launch_builder(&f);
2597        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2598        unsafe {
2599            b.launch(cfg)?;
2600        }
2601        Ok(())
2602    }
2603
2604    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2605    pub fn rows_permute(
2606        &self,
2607        src: &CudaSlice<f32>,
2608        idx: &CudaSlice<i32>,
2609        nrows: usize,
2610        ncols: usize,
2611    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2612        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2613        let f = self.func("rows_permute_f32");
2614        let (nc, nr) = (ncols as i32, nrows as i32);
2615        let cfg = LaunchConfig {
2616            grid_dim: (nrows as u32, 1, 1),
2617            block_dim: (256, 1, 1),
2618            shared_mem_bytes: 0,
2619        };
2620        let __s_b = self.gpu.stream();
2621        let mut b = __s_b.launch_builder(&f);
2622        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2623        unsafe {
2624            b.launch(cfg)?;
2625        }
2626        Ok(dst)
2627    }
2628
2629    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2630    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2631    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2632    /// decode chain and the small-t spec-verify chain match per row by construction.
2633    pub fn sigmoid_dot_rows(
2634        &self,
2635        x: &CudaSlice<f32>,
2636        w: &CudaSlice<f32>,
2637        n_embd: usize,
2638        t: usize,
2639    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2640        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2641        // config; same class as MEMRA_ROUTER_V2).
2642        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2643        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2644            let gs = self.linear(x, w, t, n_embd, 1)?;
2645            let mut g = self.uninit(t)?;
2646            self.sigmoid(&gs, &mut g, t)?;
2647            return Ok(g);
2648        }
2649        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2650        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2651        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2652        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2653        // flags doctrine; this per-token form serves every t.
2654        let mut g = self.alloc_uninit::<f32>(t)?;
2655        let f = self.func("sigmoid_dot_rows_f32");
2656        let (ne, ti) = (n_embd as i32, t as i32);
2657        let cfg = LaunchConfig {
2658            grid_dim: (t as u32, 1, 1),
2659            block_dim: (32, 8, 1),
2660            shared_mem_bytes: 0,
2661        };
2662        let __s_b = self.gpu.stream();
2663        let mut b = __s_b.launch_builder(&f);
2664        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2665        unsafe {
2666            b.launch(cfg)?;
2667        }
2668        Ok(g)
2669    }
2670
2671    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
2672    pub fn sigmoid_dot_rows_into(
2673        &self,
2674        x: &CudaSlice<f32>,
2675        w: &CudaSlice<f32>,
2676        g: &mut CudaSlice<f32>,
2677        n_embd: usize,
2678        t: usize,
2679    ) -> Result<(), Box<dyn std::error::Error>> {
2680        if g.len() < t {
2681            return Err("sigmoid_dot_rows_into output too small".into());
2682        }
2683        let f = self.func("sigmoid_dot_rows_f32");
2684        let (ne, ti) = (n_embd as i32, t as i32);
2685        let cfg = LaunchConfig {
2686            grid_dim: (t as u32, 1, 1),
2687            block_dim: (32, 8, 1),
2688            shared_mem_bytes: 0,
2689        };
2690        let __s_b = self.gpu.stream();
2691        let mut b = __s_b.launch_builder(&f);
2692        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2693        unsafe {
2694            b.launch(cfg)?;
2695        }
2696        Ok(())
2697    }
2698
2699    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2700    pub fn spec_rollback_stream(
2701        &self,
2702        len_ptrs: &CudaSlice<u64>,
2703        pos_start: &CudaSlice<i32>,
2704        acc: &CudaSlice<u32>,
2705        base: usize,
2706        n_rows: usize,
2707    ) -> Result<(), Box<dyn std::error::Error>> {
2708        let f = self.func("spec_rollback_stream");
2709        let (b, nr) = (base as i32, n_rows as i32);
2710        let cfg = LaunchConfig {
2711            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2712            block_dim: (64, 1, 1),
2713            shared_mem_bytes: 0,
2714        };
2715        let __s_bl = self.gpu.stream();
2716        let mut bl = __s_bl.launch_builder(&f);
2717        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2718        unsafe {
2719            bl.launch(cfg)?;
2720        }
2721        Ok(())
2722    }
2723
2724    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2725    pub fn plain_tok_ring(
2726        &self,
2727        vam: &CudaSlice<u32>,
2728        pos_start: &CudaSlice<i32>,
2729        base: usize,
2730        ring: &mut CudaSlice<u32>,
2731    ) -> Result<(), Box<dyn std::error::Error>> {
2732        let f = self.func("plain_tok_ring");
2733        let (b, cap) = (base as i32, ring.len() as i32);
2734        let cfg = LaunchConfig {
2735            grid_dim: (1, 1, 1),
2736            block_dim: (32, 1, 1),
2737            shared_mem_bytes: 0,
2738        };
2739        let __s_bl = self.gpu.stream();
2740        let mut bl = __s_bl.launch_builder(&f);
2741        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2742        unsafe {
2743            bl.launch(cfg)?;
2744        }
2745        Ok(())
2746    }
2747
2748    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2749    pub fn spec_ring_commit(
2750        &self,
2751        vtok: &CudaSlice<u32>,
2752        acc: &CudaSlice<u32>,
2753        brk: &CudaSlice<u32>,
2754        ring: &mut CudaSlice<u32>,
2755        pend: &mut CudaSlice<u32>,
2756    ) -> Result<(), Box<dyn std::error::Error>> {
2757        let f = self.func("spec_ring_commit");
2758        let cfg = LaunchConfig {
2759            grid_dim: (1, 1, 1),
2760            block_dim: (32, 1, 1),
2761            shared_mem_bytes: 0,
2762        };
2763        let __s_b = self.gpu.stream();
2764        let mut b = __s_b.launch_builder(&f);
2765        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2766        unsafe {
2767            b.launch(cfg)?;
2768        }
2769        Ok(())
2770    }
2771    pub fn i32_copy_add(
2772        &self,
2773        src: &CudaSlice<i32>,
2774        dst: &mut CudaSlice<i32>,
2775        delta: i32,
2776    ) -> Result<(), Box<dyn std::error::Error>> {
2777        let f = self.func("i32_copy_add");
2778        let cfg = LaunchConfig {
2779            grid_dim: (1, 1, 1),
2780            block_dim: (32, 1, 1),
2781            shared_mem_bytes: 0,
2782        };
2783        let __s_b = self.gpu.stream();
2784        let mut b = __s_b.launch_builder(&f);
2785        b.arg(src).arg(dst).arg(&delta);
2786        unsafe {
2787            b.launch(cfg)?;
2788        }
2789        Ok(())
2790    }
2791    pub fn u32_copy(
2792        &self,
2793        src: &CudaSlice<u32>,
2794        dst: &mut CudaSlice<u32>,
2795    ) -> Result<(), Box<dyn std::error::Error>> {
2796        let f = self.func("u32_copy");
2797        let cfg = LaunchConfig {
2798            grid_dim: (1, 1, 1),
2799            block_dim: (32, 1, 1),
2800            shared_mem_bytes: 0,
2801        };
2802        let __s_b = self.gpu.stream();
2803        let mut b = __s_b.launch_builder(&f);
2804        b.arg(src).arg(dst);
2805        unsafe {
2806            b.launch(cfg)?;
2807        }
2808        Ok(())
2809    }
2810
2811    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2812    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2813    /// caps acceptance exactly like drafting fewer tokens).
2814    pub fn spec_adapt_k(
2815        &self,
2816        acc: &CudaSlice<u32>,
2817        brk: &mut CudaSlice<u32>,
2818        floor: usize,
2819        cap: usize,
2820    ) -> Result<(), Box<dyn std::error::Error>> {
2821        let f = self.func("spec_adapt_k");
2822        let (fl, cp) = (floor as i32, cap as i32);
2823        let cfg = LaunchConfig {
2824            grid_dim: (1, 1, 1),
2825            block_dim: (32, 1, 1),
2826            shared_mem_bytes: 0,
2827        };
2828        let __s_b = self.gpu.stream();
2829        let mut b = __s_b.launch_builder(&f);
2830        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2831        unsafe {
2832            b.launch(cfg)?;
2833        }
2834        Ok(())
2835    }
2836
2837    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2838    pub fn spec_accept_greedy_dc(
2839        &self,
2840        preds: &CudaSlice<u32>,
2841        vtok: &CudaSlice<u32>,
2842        last_pred: &CudaSlice<u32>,
2843        brk: &CudaSlice<u32>,
2844        out: &mut CudaSlice<u32>,
2845    ) -> Result<(), Box<dyn std::error::Error>> {
2846        let f = self.func("spec_accept_greedy_dc");
2847        let cfg = LaunchConfig {
2848            grid_dim: (1, 1, 1),
2849            block_dim: (32, 1, 1),
2850            shared_mem_bytes: 0,
2851        };
2852        let __s_b = self.gpu.stream();
2853        let mut b = __s_b.launch_builder(&f);
2854        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2855        unsafe {
2856            b.launch(cfg)?;
2857        }
2858        Ok(())
2859    }
2860
2861    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2862    pub fn pos_iota(
2863        &self,
2864        pos0: &CudaSlice<i32>,
2865        out: &mut CudaSlice<i32>,
2866        t: usize,
2867    ) -> Result<(), Box<dyn std::error::Error>> {
2868        let f = self.func("pos_iota_i32");
2869        let ti = t as i32;
2870        let cfg = LaunchConfig {
2871            grid_dim: (1, 1, 1),
2872            block_dim: (t.max(1) as u32, 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(pos0).arg(out).arg(&ti);
2878        unsafe {
2879            b.launch(cfg)?;
2880        }
2881        Ok(())
2882    }
2883    #[allow(clippy::too_many_arguments)]
2884    pub fn append_kv_quantized_rows_dc(
2885        &self,
2886        k_rows: &CudaSlice<f32>,
2887        v_rows: &CudaSlice<f32>,
2888        kc: &mut CudaSlice<u8>,
2889        vc: &mut CudaSlice<u8>,
2890        t0_dev: &CudaSlice<i32>,
2891        t: usize,
2892        kv_dim_k: usize,
2893        kv_dim_v: usize,
2894        k_tok_bytes: usize,
2895        v_tok_bytes: usize,
2896        g: bool,
2897    ) -> Result<(), Box<dyn std::error::Error>> {
2898        let f = if g {
2899            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2900        } else {
2901            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2902        };
2903        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2904        let cfg = LaunchConfig {
2905            grid_dim: (nblk, t as u32, 1),
2906            block_dim: (32, 1, 1),
2907            shared_mem_bytes: 0,
2908        };
2909        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2910        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2911        let __s_b = self.gpu.stream();
2912        let mut b = __s_b.launch_builder(&f);
2913        b.arg(k_rows)
2914            .arg(v_rows)
2915            .arg(kc)
2916            .arg(vc)
2917            .arg(t0_dev)
2918            .arg(&kdk)
2919            .arg(&kdv)
2920            .arg(&ktb)
2921            .arg(&vtb);
2922        unsafe {
2923            b.launch(cfg)?;
2924        }
2925        Ok(())
2926    }
2927
2928    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2929    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2930    #[allow(clippy::too_many_arguments)]
2931    pub fn append_kv_quantized_row_dc_inc(
2932        &self,
2933        k_row: &CudaSlice<f32>,
2934        v_row: &CudaSlice<f32>,
2935        kc: &mut CudaSlice<u8>,
2936        vc: &mut CudaSlice<u8>,
2937        t0_dev: &mut CudaSlice<i32>,
2938        kv_dim_k: usize,
2939        kv_dim_v: usize,
2940        k_tok_bytes: usize,
2941        v_tok_bytes: usize,
2942        g: bool,
2943    ) -> Result<(), Box<dyn std::error::Error>> {
2944        let f = if g {
2945            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2946        } else {
2947            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2948        };
2949        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2950        let cfg = LaunchConfig {
2951            grid_dim: (1, 1, 1),
2952            block_dim: (nthreads, 1, 1),
2953            shared_mem_bytes: 0,
2954        };
2955        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2956        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2957        let __s_b = self.gpu.stream();
2958        let mut b = __s_b.launch_builder(&f);
2959        b.arg(k_row)
2960            .arg(v_row)
2961            .arg(kc)
2962            .arg(vc)
2963            .arg(t0_dev)
2964            .arg(&kdk)
2965            .arg(&kdv)
2966            .arg(&ktb)
2967            .arg(&vtb);
2968        unsafe {
2969            b.launch(cfg)?;
2970        }
2971        Ok(())
2972    }
2973
2974    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2975    pub fn pack_tok_p(
2976        &self,
2977        tok: &CudaSlice<u32>,
2978        p: &CudaSlice<f32>,
2979        out: &mut CudaSlice<u32>,
2980        slot: usize,
2981    ) -> Result<(), Box<dyn std::error::Error>> {
2982        let f = self.func("pack_tok_p");
2983        let sl = slot as i32;
2984        let cfg = LaunchConfig {
2985            grid_dim: (1, 1, 1),
2986            block_dim: (32, 1, 1),
2987            shared_mem_bytes: 0,
2988        };
2989        let __s_b = self.gpu.stream();
2990        let mut b = __s_b.launch_builder(&f);
2991        b.arg(tok).arg(p).arg(out).arg(&sl);
2992        unsafe {
2993            b.launch(cfg)?;
2994        }
2995        Ok(())
2996    }
2997    pub fn tok_map_u32(
2998        &self,
2999        tok: &mut CudaSlice<u32>,
3000        map: &CudaSlice<u32>,
3001    ) -> Result<(), Box<dyn std::error::Error>> {
3002        let f = self.func("tok_map_u32");
3003        let cfg = LaunchConfig {
3004            grid_dim: (1, 1, 1),
3005            block_dim: (32, 1, 1),
3006            shared_mem_bytes: 0,
3007        };
3008        let __s_b = self.gpu.stream();
3009        let mut b = __s_b.launch_builder(&f);
3010        b.arg(tok).arg(map);
3011        unsafe {
3012            b.launch(cfg)?;
3013        }
3014        Ok(())
3015    }
3016
3017    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
3018    #[allow(clippy::too_many_arguments)]
3019    pub fn spec_assemble_verify(
3020        &self,
3021        tokp: &CudaSlice<u32>,
3022        pend: &CudaSlice<u32>,
3023        d2t: Option<&CudaSlice<u32>>,
3024        vtok: &mut CudaSlice<u32>,
3025        brk: &mut CudaSlice<u32>,
3026        p_min: f32,
3027        k: usize,
3028        pmin0: bool,
3029    ) -> Result<(), Box<dyn std::error::Error>> {
3030        let f = self.func("spec_assemble_verify");
3031        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3032        let cfg = LaunchConfig {
3033            grid_dim: (1, 1, 1),
3034            block_dim: (32, 1, 1),
3035            shared_mem_bytes: 0,
3036        };
3037        let __s_b = self.gpu.stream();
3038        let mut b = __s_b.launch_builder(&f);
3039        match d2t {
3040            Some(m) => {
3041                b.arg(tokp)
3042                    .arg(pend)
3043                    .arg(m)
3044                    .arg(vtok)
3045                    .arg(brk)
3046                    .arg(&p_min)
3047                    .arg(&ki)
3048                    .arg(&pm);
3049                unsafe {
3050                    b.launch(cfg)?;
3051                }
3052            }
3053            None => {
3054                let null: u64 = 0;
3055                b.arg(tokp)
3056                    .arg(pend)
3057                    .arg(&null)
3058                    .arg(vtok)
3059                    .arg(brk)
3060                    .arg(&p_min)
3061                    .arg(&ki)
3062                    .arg(&pm);
3063                unsafe {
3064                    b.launch(cfg)?;
3065                }
3066            }
3067        }
3068        Ok(())
3069    }
3070
3071    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
3072    #[allow(clippy::too_many_arguments)]
3073    pub fn ssm_conv_ring_rebuild_dc(
3074        &self,
3075        qkv_tm: &CudaSlice<f32>,
3076        ring_old: &CudaSlice<f32>,
3077        conv_state: &mut CudaSlice<f32>,
3078        conv_dim: usize,
3079        acc: &CudaSlice<u32>,
3080        base: usize,
3081        t_v: usize,
3082        d_conv: usize,
3083    ) -> Result<(), Box<dyn std::error::Error>> {
3084        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3085        let n = conv_dim * (d_conv - 1);
3086        let cfg = LaunchConfig::for_num_elems(n as u32);
3087        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3088        let __s_b = self.gpu.stream();
3089        let mut b = __s_b.launch_builder(&f);
3090        b.arg(qkv_tm)
3091            .arg(ring_old)
3092            .arg(conv_state)
3093            .arg(&cd)
3094            .arg(acc)
3095            .arg(&b0)
3096            .arg(&tv)
3097            .arg(&dc);
3098        unsafe {
3099            b.launch(cfg)?;
3100        }
3101        Ok(())
3102    }
3103    #[allow(clippy::too_many_arguments)]
3104    pub fn gdn_scan_s128_dc(
3105        &self,
3106        q: &CudaSlice<f32>,
3107        k: &CudaSlice<f32>,
3108        v: &CudaSlice<f32>,
3109        g: &CudaSlice<f32>,
3110        beta: &CudaSlice<f32>,
3111        state_in: &CudaSlice<f32>,
3112        state_out: &mut CudaSlice<f32>,
3113        o: &mut CudaSlice<f32>,
3114        n_head: usize,
3115        acc: &CudaSlice<u32>,
3116        base: usize,
3117        t_v: usize,
3118        scale: f32,
3119    ) -> Result<(), Box<dyn std::error::Error>> {
3120        let f = self.func("gdn_scan_s128_dc");
3121        const S_V: u32 = 128;
3122        const WARP: u32 = 32;
3123        const COLS_PER_BLOCK: u32 = 4;
3124        let cfg = LaunchConfig {
3125            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3126            block_dim: (WARP, COLS_PER_BLOCK, 1),
3127            shared_mem_bytes: 0,
3128        };
3129        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3130        let __s_b = self.gpu.stream();
3131        let mut b = __s_b.launch_builder(&f);
3132        b.arg(q)
3133            .arg(k)
3134            .arg(v)
3135            .arg(g)
3136            .arg(beta)
3137            .arg(state_in)
3138            .arg(state_out)
3139            .arg(o)
3140            .arg(&h)
3141            .arg(acc)
3142            .arg(&b0)
3143            .arg(&tv)
3144            .arg(&scale);
3145        unsafe {
3146            b.launch(cfg)?;
3147        }
3148        Ok(())
3149    }
3150
3151    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
3152    pub fn spec_rollback_kv(
3153        &self,
3154        len_ptrs: &CudaSlice<u64>,
3155        saved: &CudaSlice<i32>,
3156        acc: &CudaSlice<u32>,
3157        base: usize,
3158        n_layer: usize,
3159    ) -> Result<(), Box<dyn std::error::Error>> {
3160        let f = self.func("spec_rollback_kv");
3161        let (b, nl) = (base as i32, n_layer as i32);
3162        let cfg = LaunchConfig {
3163            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3164            block_dim: (64, 1, 1),
3165            shared_mem_bytes: 0,
3166        };
3167        let __s_bl = self.gpu.stream();
3168        let mut bl = __s_bl.launch_builder(&f);
3169        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3170        unsafe {
3171            bl.launch(cfg)?;
3172        }
3173        Ok(())
3174    }
3175
3176    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3177    pub fn spec_fork_valid(
3178        &self,
3179        acc: &CudaSlice<u32>,
3180        optimistic_pending: u32,
3181        valid: &mut CudaSlice<u32>,
3182    ) -> Result<(), Box<dyn std::error::Error>> {
3183        let f = self.func("spec_fork_valid");
3184        let cfg = LaunchConfig {
3185            grid_dim: (1, 1, 1),
3186            block_dim: (1, 1, 1),
3187            shared_mem_bytes: 0,
3188        };
3189        let __s_bl = self.gpu.stream();
3190        let mut bl = __s_bl.launch_builder(&f);
3191        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3192        unsafe {
3193            bl.launch(cfg)?;
3194        }
3195        Ok(())
3196    }
3197
3198    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3199    pub fn spec_fork_reconcile_kv(
3200        &self,
3201        len_ptrs: &CudaSlice<u64>,
3202        saved: &CudaSlice<i32>,
3203        acc: &CudaSlice<u32>,
3204        valid: &CudaSlice<u32>,
3205        base: usize,
3206        n_layer: usize,
3207    ) -> Result<(), Box<dyn std::error::Error>> {
3208        let f = self.func("spec_fork_reconcile_kv");
3209        let (b, nl) = (base as i32, n_layer as i32);
3210        let cfg = LaunchConfig {
3211            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3212            block_dim: (64, 1, 1),
3213            shared_mem_bytes: 0,
3214        };
3215        let __s_bl = self.gpu.stream();
3216        let mut bl = __s_bl.launch_builder(&f);
3217        bl.arg(len_ptrs)
3218            .arg(saved)
3219            .arg(acc)
3220            .arg(valid)
3221            .arg(&b)
3222            .arg(&nl);
3223        unsafe {
3224            bl.launch(cfg)?;
3225        }
3226        Ok(())
3227    }
3228
3229    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3230    pub fn spec_fork_restore_f32(
3231        &self,
3232        snapshot: &CudaSlice<f32>,
3233        state: &mut CudaSlice<f32>,
3234        valid: &CudaSlice<u32>,
3235    ) -> Result<(), Box<dyn std::error::Error>> {
3236        assert_eq!(
3237            snapshot.len(),
3238            state.len(),
3239            "fork recurrent snapshot shape mismatch"
3240        );
3241        let f = self.func("spec_fork_restore_f32");
3242        let n = state.len() as i32;
3243        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3244        let cfg = LaunchConfig {
3245            grid_dim: (blocks, 1, 1),
3246            block_dim: (256, 1, 1),
3247            shared_mem_bytes: 0,
3248        };
3249        let __s_bl = self.gpu.stream();
3250        let mut bl = __s_bl.launch_builder(&f);
3251        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3252        unsafe {
3253            bl.launch(cfg)?;
3254        }
3255        Ok(())
3256    }
3257
3258    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3259    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3260    pub fn spec_seed_gather(
3261        &self,
3262        vx: &CudaSlice<f32>,
3263        fill_prev: &CudaSlice<f32>,
3264        acc: &CudaSlice<u32>,
3265        h_seed: &mut CudaSlice<f32>,
3266        base: usize,
3267        n_embd: usize,
3268    ) -> Result<(), Box<dyn std::error::Error>> {
3269        let f = self.func("spec_seed_gather");
3270        let (b, ne) = (base as i32, n_embd as i32);
3271        let cfg = LaunchConfig {
3272            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3273            block_dim: (256, 1, 1),
3274            shared_mem_bytes: 0,
3275        };
3276        let __s_bl = self.gpu.stream();
3277        let mut bl = __s_bl.launch_builder(&f);
3278        bl.arg(vx)
3279            .arg(fill_prev)
3280            .arg(acc)
3281            .arg(h_seed)
3282            .arg(&b)
3283            .arg(&ne);
3284        unsafe {
3285            bl.launch(cfg)?;
3286        }
3287        Ok(())
3288    }
3289
3290    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3291    pub fn spec_accept_greedy(
3292        &self,
3293        preds: &CudaSlice<u32>,
3294        draft: &CudaSlice<u32>,
3295        last_pred: u32,
3296        base: usize,
3297        k_round: usize,
3298        out: &mut CudaSlice<u32>,
3299    ) -> Result<(), Box<dyn std::error::Error>> {
3300        let f = self.func("spec_accept_greedy");
3301        let (b, k) = (base as i32, k_round as i32);
3302        let cfg = LaunchConfig {
3303            grid_dim: (1, 1, 1),
3304            block_dim: (32, 1, 1),
3305            shared_mem_bytes: 0,
3306        };
3307        let __s_bl = self.gpu.stream();
3308        let mut bl = __s_bl.launch_builder(&f);
3309        bl.arg(preds)
3310            .arg(draft)
3311            .arg(&last_pred)
3312            .arg(&b)
3313            .arg(&k)
3314            .arg(out);
3315        unsafe {
3316            bl.launch(cfg)?;
3317        }
3318        Ok(())
3319    }
3320
3321    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3322    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3323    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3324
3325    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3326    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3327    pub fn gumbel_perturb(
3328        &self,
3329        x: &CudaSlice<f32>,
3330        y: &mut CudaSlice<f32>,
3331        n: usize,
3332        seed: u64,
3333        stream_pos: u32,
3334        temp: f32,
3335    ) -> Result<(), Box<dyn std::error::Error>> {
3336        let f = self.func("gumbel_perturb_f32");
3337        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3338        let cfg = LaunchConfig {
3339            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3340            block_dim: (256, 1, 1),
3341            shared_mem_bytes: 0,
3342        };
3343        let __s_b = self.gpu.stream();
3344        let mut b = __s_b.launch_builder(&f);
3345        b.arg(x)
3346            .arg(&mut *y)
3347            .arg(&ni)
3348            .arg(&slo)
3349            .arg(&shi)
3350            .arg(&stream_pos)
3351            .arg(&temp);
3352        unsafe {
3353            b.launch(cfg)?;
3354        }
3355        Ok(())
3356    }
3357
3358    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3359    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3360    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3361    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3362    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3363    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3364    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3365    pub fn mask_logits_col(
3366        &self,
3367        logits: &mut CudaSlice<f32>,
3368        mask: &CudaSlice<u32>,
3369        col: usize,
3370        n: usize,
3371        mask_words: usize,
3372    ) -> Result<(), Box<dyn std::error::Error>> {
3373        let f = self.func("mask_logits_f32");
3374        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3375        let cfg = LaunchConfig {
3376            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3377            block_dim: (256, 1, 1),
3378            shared_mem_bytes: 0,
3379        };
3380        let __s_b = self.gpu.stream();
3381        let mut b = __s_b.launch_builder(&f);
3382        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3383        unsafe {
3384            b.launch(cfg)?;
3385        }
3386        Ok(())
3387    }
3388
3389    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3390    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3391    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3392    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3393    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3394    /// pointer-invariance IS the serving isolation contract for sampled rows.
3395    pub fn gumbel_perturb_col(
3396        &self,
3397        x: &CudaSlice<f32>,
3398        col: usize,
3399        y: &mut CudaSlice<f32>,
3400        n: usize,
3401        seed: u64,
3402        stream_pos: u32,
3403        temp: f32,
3404    ) -> Result<(), Box<dyn std::error::Error>> {
3405        let f = self.func("gumbel_perturb_f32");
3406        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3407        let col_view = x.slice(col * n..(col + 1) * n);
3408        let cfg = LaunchConfig {
3409            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3410            block_dim: (256, 1, 1),
3411            shared_mem_bytes: 0,
3412        };
3413        let __s_b = self.gpu.stream();
3414        let mut b = __s_b.launch_builder(&f);
3415        b.arg(&col_view)
3416            .arg(&mut *y)
3417            .arg(&ni)
3418            .arg(&slo)
3419            .arg(&shi)
3420            .arg(&stream_pos)
3421            .arg(&temp);
3422        unsafe {
3423            b.launch(cfg)?;
3424        }
3425        Ok(())
3426    }
3427
3428    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3429    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3430    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3431    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3432    /// the serving isolation contract for sampled rows).
3433    #[allow(clippy::too_many_arguments)]
3434    pub fn gumbel_perturb_filtered_col(
3435        &self,
3436        x: &CudaSlice<f32>,
3437        col: usize,
3438        y: &mut CudaSlice<f32>,
3439        n: usize,
3440        seed: u64,
3441        stream_pos: u32,
3442        temp: f32,
3443        stat_max: &CudaSlice<f32>,
3444        stat_th: &CudaSlice<f32>,
3445        stat_idx: usize,
3446    ) -> Result<(), Box<dyn std::error::Error>> {
3447        let f = self.func("gumbel_perturb_filtered_col_f32");
3448        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3449        let (ci, si) = (col as i32, stat_idx as i32);
3450        let cfg = LaunchConfig {
3451            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3452            block_dim: (256, 1, 1),
3453            shared_mem_bytes: 0,
3454        };
3455        let __s_b = self.gpu.stream();
3456        let mut b = __s_b.launch_builder(&f);
3457        b.arg(x)
3458            .arg(&ci)
3459            .arg(&mut *y)
3460            .arg(&ni)
3461            .arg(&slo)
3462            .arg(&shi)
3463            .arg(&stream_pos)
3464            .arg(&temp)
3465            .arg(stat_max)
3466            .arg(stat_th)
3467            .arg(&si);
3468        unsafe {
3469            b.launch(cfg)?;
3470        }
3471        Ok(())
3472    }
3473
3474    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3475    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3476    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3477    /// reads it (counter is data, not state — graph-replay-safe).
3478    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3479        let f = self.func("memra_sctr_inc");
3480        let cfg = LaunchConfig {
3481            grid_dim: (1, 1, 1),
3482            block_dim: (1, 1, 1),
3483            shared_mem_bytes: 0,
3484        };
3485        let __s_b = self.gpu.stream();
3486        let mut b = __s_b.launch_builder(&f);
3487        b.arg(&mut *ctr);
3488        unsafe {
3489            b.launch(cfg)?;
3490        }
3491        Ok(())
3492    }
3493
3494    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3495    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3496    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3497    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3498    pub fn gumbel_perturb_ctr(
3499        &self,
3500        x: &CudaSlice<f32>,
3501        y: &mut CudaSlice<f32>,
3502        n: usize,
3503        seed: u64,
3504        ctr: &CudaSlice<u32>,
3505        temp: f32,
3506    ) -> Result<(), Box<dyn std::error::Error>> {
3507        let f = self.func("gumbel_perturb_ctr_f32");
3508        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3509        let cfg = LaunchConfig {
3510            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3511            block_dim: (256, 1, 1),
3512            shared_mem_bytes: 0,
3513        };
3514        let __s_b = self.gpu.stream();
3515        let mut b = __s_b.launch_builder(&f);
3516        b.arg(x)
3517            .arg(&mut *y)
3518            .arg(&ni)
3519            .arg(&slo)
3520            .arg(&shi)
3521            .arg(ctr)
3522            .arg(&temp);
3523        unsafe {
3524            b.launch(cfg)?;
3525        }
3526        Ok(())
3527    }
3528
3529    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3530    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3531    /// (smallest-index tie-break — matches the argmax-gate contract).
3532    pub fn softmax_gather(
3533        &self,
3534        x: &CudaSlice<f32>,
3535        row_stride: usize,
3536        ids: &CudaSlice<u32>,
3537        rows: &CudaSlice<i32>,
3538        out: &mut CudaSlice<f32>,
3539        n: usize,
3540        npair: usize,
3541        temp: f32,
3542    ) -> Result<(), Box<dyn std::error::Error>> {
3543        let f = self.func("softmax_gather_f32");
3544        let (ni, rs) = (n as i32, row_stride as i64);
3545        let np = npair as i32;
3546        let cfg = LaunchConfig {
3547            grid_dim: (npair as u32, 1, 1),
3548            block_dim: (256, 1, 1),
3549            shared_mem_bytes: 0,
3550        };
3551        let __s_b = self.gpu.stream();
3552        let mut b = __s_b.launch_builder(&f);
3553        b.arg(x)
3554            .arg(&rs)
3555            .arg(ids)
3556            .arg(rows)
3557            .arg(&mut *out)
3558            .arg(&ni)
3559            .arg(&np)
3560            .arg(&temp);
3561        unsafe {
3562            b.launch(cfg)?;
3563        }
3564        Ok(())
3565    }
3566
3567    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3568    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3569    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3570    pub fn residual_sample(
3571        &self,
3572        p: &CudaSlice<f32>,
3573        q: Option<&CudaSlice<f32>>,
3574        n: usize,
3575        temp: f32,
3576        seed: u64,
3577        stream_pos: u32,
3578        out_tok: &mut CudaSlice<u32>,
3579    ) -> Result<(), Box<dyn std::error::Error>> {
3580        let f = self.func("residual_sample_f32");
3581        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3582        let nth = 1024u32;
3583        let cfg = LaunchConfig {
3584            grid_dim: (1, 1, 1),
3585            block_dim: (nth, 1, 1),
3586            shared_mem_bytes: 0,
3587        };
3588        let has_q: i32 = q.is_some() as i32;
3589        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3590        let __s_b = self.gpu.stream();
3591        let mut b = __s_b.launch_builder(&f);
3592        b.arg(p)
3593            .arg(qbuf)
3594            .arg(&has_q)
3595            .arg(&ni)
3596            .arg(&temp)
3597            .arg(&slo)
3598            .arg(&shi)
3599            .arg(&stream_pos)
3600            .arg(&mut *out_tok);
3601        unsafe {
3602            b.launch(cfg)?;
3603        }
3604        Ok(())
3605    }
3606
3607    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3608    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3609    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3610    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3611    pub fn with_moe_cache<R>(
3612        &self,
3613        max_block_bytes: usize,
3614        f: impl FnOnce(
3615            &mut crate::moe_cache::MoeSlotCache,
3616            &Engine,
3617        ) -> Result<R, Box<dyn std::error::Error>>,
3618    ) -> Result<R, Box<dyn std::error::Error>> {
3619        let mut guard = self.moe_cache.lock().unwrap();
3620        if guard.is_none() {
3621            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3622        }
3623        let cache = guard.as_mut().unwrap();
3624        f(cache, self)
3625    }
3626
3627    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3628    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3629    pub fn freeze_moe_cache(&self) {
3630        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3631            cache.freeze();
3632        }
3633    }
3634
3635    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3636    /// Never constructs a cache.
3637    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3638        self.moe_cache
3639            .lock()
3640            .unwrap()
3641            .as_ref()
3642            .map(crate::moe_cache::MoeSlotCache::export_residency)
3643    }
3644
3645    pub(crate) fn moe_cache_frozen(&self) -> bool {
3646        self.moe_cache
3647            .lock()
3648            .unwrap()
3649            .as_ref()
3650            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3651    }
3652
3653    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3654    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3655    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3656    /// while leaving the profiling warmup's established batched behavior untouched.
3657    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3658    /// tokenwise arm anyway.)
3659    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3660        crate::cpu_experts::configured()
3661            && self.moe_cache_frozen()
3662            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3663    }
3664
3665    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3666    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3667        assert!(
3668            self.moe_cache.lock().unwrap().is_none(),
3669            "MoE cache layout configured after cache construction"
3670        );
3671        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3672    }
3673
3674    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3675        self.moe_cache_layout.lock().unwrap().clone()
3676    }
3677
3678    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3679    pub fn moe_cache_enabled() -> bool {
3680        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3681    }
3682
3683    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3684    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3685    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3686        let guard = self.moe_cache.lock().unwrap();
3687        guard
3688            .as_ref()
3689            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3690    }
3691
3692    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3693    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3694    /// callers compare a before/after snapshot around a decode window.
3695    pub fn cpu_expert_stats(
3696        &self,
3697    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3698        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3699    }
3700
3701    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3702    /// the backend tail that resident-GPU expert work did not hide.
3703    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3704        crate::cpu_experts::predictor_stats()
3705    }
3706
3707    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3708        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3709    }
3710
3711    /// CPU-routed expert selections grouped by how many of their three projections were already
3712    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3713    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3714        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3715    }
3716
3717    /// Positioned-read proof-backend counters:
3718    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3719    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3720        let guard = self.moe_cache.lock().unwrap();
3721        guard
3722            .as_ref()
3723            .and_then(|cache| cache.pread_stats())
3724            .map(|stats| {
3725                (
3726                    stats.reads,
3727                    stats.bytes,
3728                    stats.read_errors,
3729                    stats.short_reads,
3730                    stats.fallbacks,
3731                    stats.buffer_waits,
3732                    stats.ring_full,
3733                )
3734            })
3735    }
3736
3737    /// Spill configuration values that warned and substituted their documented defaults.
3738    pub fn spill_config_fallbacks(&self) -> u64 {
3739        crate::spill_pread::config_fallbacks()
3740    }
3741
3742    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3743    pub fn moe_cache_reset_counters(&self) {
3744        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3745            c.reset_counters();
3746        }
3747    }
3748
3749    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3750        Ok(self.gpu.stream().clone_htod(v)?)
3751    }
3752
3753    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3754    /// past the final q4_0 block through their aligned window — the bytes never reach a
3755    /// result (funnelshift discards them) but must be mapped memory.
3756    pub fn htod_bytes_padded(
3757        &self,
3758        v: &[u8],
3759        pad: usize,
3760    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3761        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3762        {
3763            let mut view = d.slice_mut(0..v.len());
3764            self.gpu.stream().memcpy_htod(v, &mut view)?;
3765        }
3766        Ok(d)
3767    }
3768
3769    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3770    pub fn copy_into(
3771        &self,
3772        dst: &mut CudaSlice<f32>,
3773        off: usize,
3774        src: &CudaSlice<f32>,
3775        len: usize,
3776    ) -> Result<(), Box<dyn std::error::Error>> {
3777        let mut view = dst.slice_mut(off..off + len);
3778        self.gpu
3779            .stream()
3780            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3781        Ok(())
3782    }
3783
3784    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3785    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3786    pub fn copy_u8_into(
3787        &self,
3788        dst: &mut CudaSlice<u8>,
3789        off: usize,
3790        src: &CudaSlice<u8>,
3791        len: usize,
3792    ) -> Result<(), Box<dyn std::error::Error>> {
3793        let mut view = dst.slice_mut(off..off + len);
3794        self.gpu
3795            .stream()
3796            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3797        Ok(())
3798    }
3799
3800    /// D2D byte-range copy with explicit source and destination offsets.
3801    pub fn copy_u8_range_into(
3802        &self,
3803        dst: &mut CudaSlice<u8>,
3804        dst_off: usize,
3805        src: &CudaSlice<u8>,
3806        src_off: usize,
3807        len: usize,
3808    ) -> Result<(), Box<dyn std::error::Error>> {
3809        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3810        self.gpu
3811            .stream()
3812            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3813        Ok(())
3814    }
3815
3816    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3817    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3818    /// keeping the audited attention range contiguous without changing its absolute start.
3819    pub fn prepare_kv_append(
3820        &self,
3821        kv: &mut crate::cache::KvLayer,
3822        retain_from: usize,
3823        append_rows: usize,
3824    ) -> Result<usize, Box<dyn std::error::Error>> {
3825        let Some(plan) = kv
3826            .ring
3827            .as_ref()
3828            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3829            .transpose()?
3830        else {
3831            return Ok(kv.len);
3832        };
3833        match plan {
3834            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3835            crate::cache::KvRingAppend::Rebase {
3836                src_row,
3837                keep_rows,
3838                new_base,
3839                write_row,
3840            } => {
3841                if keep_rows > 0 {
3842                    let k_len = keep_rows * kv.k_tok_bytes;
3843                    let v_len = keep_rows * kv.v_tok_bytes;
3844                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3845                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3846                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3847                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3848                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3849                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3850                }
3851                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3852                Ok(write_row)
3853            }
3854        }
3855    }
3856
3857    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3858    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3859    pub fn htod_u8_into(
3860        &self,
3861        dst: &mut CudaSlice<u8>,
3862        off: usize,
3863        src: &[u8],
3864    ) -> Result<(), Box<dyn std::error::Error>> {
3865        let mut view = dst.slice_mut(off..off + src.len());
3866        self.gpu.stream().memcpy_htod(src, &mut view)?;
3867        Ok(())
3868    }
3869
3870    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3871        b.slice(0..len)
3872    }
3873
3874    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3875    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3876    pub fn view_u8_range<'a>(
3877        &self,
3878        b: &'a CudaSlice<u8>,
3879        start: usize,
3880        end: usize,
3881    ) -> cudarc::driver::CudaView<'a, u8> {
3882        b.slice(start..end)
3883    }
3884    pub fn view_u8<'a>(
3885        &self,
3886        b: &'a CudaSlice<u8>,
3887        len: usize,
3888    ) -> cudarc::driver::CudaView<'a, u8> {
3889        b.slice(0..len)
3890    }
3891
3892    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3893    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3894    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3895    pub fn append_kv_quantized(
3896        &self,
3897        k_row: &CudaSlice<f32>,
3898        v_row: &CudaSlice<f32>,
3899        kc: &mut CudaSlice<u8>,
3900        vc: &mut CudaSlice<u8>,
3901        t: usize,
3902        kv_dim_k: usize,
3903        kv_dim_v: usize,
3904        k_tok_bytes: usize,
3905        v_tok_bytes: usize,
3906        g: bool,
3907    ) -> Result<(), Box<dyn std::error::Error>> {
3908        let f = if g {
3909            self.func_g("append_quantize_kv_q8_0_q5_1")
3910        } else {
3911            self.func("append_quantize_kv_q8_0_q5_1")
3912        };
3913        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3914        let cfg = LaunchConfig {
3915            grid_dim: (nblk, 1, 1),
3916            block_dim: (32, 1, 1),
3917            shared_mem_bytes: 0,
3918        };
3919        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3920        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3921        let __s_b = self.gpu.stream();
3922        let mut b = __s_b.launch_builder(&f);
3923        b.arg(k_row)
3924            .arg(v_row)
3925            .arg(kc)
3926            .arg(vc)
3927            .arg(&ti)
3928            .arg(&kdk)
3929            .arg(&kdv)
3930            .arg(&ktb)
3931            .arg(&vtb);
3932        unsafe {
3933            b.launch(cfg)?;
3934        }
3935        Ok(())
3936    }
3937
3938    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3939    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3940    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3941    pub fn append_kv_quantized_dc(
3942        &self,
3943        k_row: &CudaSlice<f32>,
3944        v_row: &CudaSlice<f32>,
3945        kc: &mut CudaSlice<u8>,
3946        vc: &mut CudaSlice<u8>,
3947        t_dev: &CudaSlice<i32>,
3948        kv_dim_k: usize,
3949        kv_dim_v: usize,
3950        k_tok_bytes: usize,
3951        v_tok_bytes: usize,
3952        g: bool,
3953    ) -> Result<(), Box<dyn std::error::Error>> {
3954        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3955        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3956        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3957        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3958        if Self::pdl_on() && Self::pdl_wb_on() {
3959            use cudarc::driver::{DevicePtr, DevicePtrMut};
3960            let s = &self.gpu.stream();
3961            let (pk, _g0) = k_row.device_ptr(s);
3962            let (pv, _g1) = v_row.device_ptr(s);
3963            let (pkc, _g2) = kc.device_ptr_mut(s);
3964            let (pvc, _g3) = vc.device_ptr_mut(s);
3965            let (pt, _g4) = t_dev.device_ptr(s);
3966            let mut ps = [
3967                &pk as *const _ as *mut std::ffi::c_void,
3968                &pv as *const _ as *mut _,
3969                &pkc as *const _ as *mut _,
3970                &pvc as *const _ as *mut _,
3971                &pt as *const _ as *mut _,
3972                &kdk as *const _ as *mut _,
3973                &kdv as *const _ as *mut _,
3974                &ktb as *const _ as *mut _,
3975                &vtb as *const _ as *mut _,
3976            ];
3977            unsafe {
3978                self.launch_pdl_flash(
3979                    g,
3980                    "append_quantize_kv_q8_0_q5_1_dc",
3981                    (nblk, 1, 1),
3982                    (32, 1, 1),
3983                    0,
3984                    &mut ps,
3985                )?;
3986            }
3987            return Ok(());
3988        }
3989        let f = if g {
3990            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3991        } else {
3992            self.func("append_quantize_kv_q8_0_q5_1_dc")
3993        };
3994        let cfg = LaunchConfig {
3995            grid_dim: (nblk, 1, 1),
3996            block_dim: (32, 1, 1),
3997            shared_mem_bytes: 0,
3998        };
3999        let __s_b = self.gpu.stream();
4000        let mut b = __s_b.launch_builder(&f);
4001        b.arg(k_row)
4002            .arg(v_row)
4003            .arg(kc)
4004            .arg(vc)
4005            .arg(t_dev)
4006            .arg(&kdk)
4007            .arg(&kdv)
4008            .arg(&ktb)
4009            .arg(&vtb);
4010        unsafe {
4011            b.launch(cfg)?;
4012        }
4013        Ok(())
4014    }
4015
4016    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
4017    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
4018    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
4019    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
4020    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
4021    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
4022    #[allow(clippy::too_many_arguments)]
4023    pub fn append_kv_quantized_rows(
4024        &self,
4025        k_rows: &CudaSlice<f32>,
4026        v_rows: &CudaSlice<f32>,
4027        kc: &mut CudaSlice<u8>,
4028        vc: &mut CudaSlice<u8>,
4029        t0: usize,
4030        t: usize,
4031        kv_dim_k: usize,
4032        kv_dim_v: usize,
4033        k_tok_bytes: usize,
4034        v_tok_bytes: usize,
4035        g: bool,
4036    ) -> Result<(), Box<dyn std::error::Error>> {
4037        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4038            for i in 0..t {
4039                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4040                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4041                self.append_kv_quantized_view(
4042                    &k_row,
4043                    &v_row,
4044                    kc,
4045                    vc,
4046                    t0 + i,
4047                    kv_dim_k,
4048                    kv_dim_v,
4049                    k_tok_bytes,
4050                    v_tok_bytes,
4051                    g,
4052                )?;
4053            }
4054            return Ok(());
4055        }
4056        let f = if g {
4057            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4058        } else {
4059            self.func("append_quantize_kv_q8_0_q5_1_rows")
4060        };
4061        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4062        let cfg = LaunchConfig {
4063            grid_dim: (nblk, t as u32, 1),
4064            block_dim: (32, 1, 1),
4065            shared_mem_bytes: 0,
4066        };
4067        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4068        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4069        let __s_b = self.gpu.stream();
4070        let mut b = __s_b.launch_builder(&f);
4071        b.arg(k_rows)
4072            .arg(v_rows)
4073            .arg(kc)
4074            .arg(vc)
4075            .arg(&t0i)
4076            .arg(&kdk)
4077            .arg(&kdv)
4078            .arg(&ktb)
4079            .arg(&vtb);
4080        unsafe {
4081            b.launch(cfg)?;
4082        }
4083        Ok(())
4084    }
4085
4086    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
4087    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
4088    /// later, inside a captured graph) without a host round-trip.
4089    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4090        let f = self.func("inc_i32");
4091        let cfg = LaunchConfig {
4092            grid_dim: (1, 1, 1),
4093            block_dim: (1, 1, 1),
4094            shared_mem_bytes: 0,
4095        };
4096        let __s_b = self.gpu.stream();
4097        let mut b = __s_b.launch_builder(&f);
4098        b.arg(p);
4099        unsafe {
4100            b.launch(cfg)?;
4101        }
4102        Ok(())
4103    }
4104
4105    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
4106    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
4107    pub fn append_kv_quantized_view(
4108        &self,
4109        k_row: &cudarc::driver::CudaView<f32>,
4110        v_row: &cudarc::driver::CudaView<f32>,
4111        kc: &mut CudaSlice<u8>,
4112        vc: &mut CudaSlice<u8>,
4113        t: usize,
4114        kv_dim_k: usize,
4115        kv_dim_v: usize,
4116        k_tok_bytes: usize,
4117        v_tok_bytes: usize,
4118        g: bool,
4119    ) -> Result<(), Box<dyn std::error::Error>> {
4120        let stream = self.gpu.stream();
4121        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4122        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4123        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4124        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4125        let f = if g {
4126            self.func_g("append_quantize_kv_q8_0_q5_1")
4127        } else {
4128            self.func("append_quantize_kv_q8_0_q5_1")
4129        };
4130        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4131        let cfg = LaunchConfig {
4132            grid_dim: (nblk, 1, 1),
4133            block_dim: (32, 1, 1),
4134            shared_mem_bytes: 0,
4135        };
4136        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4137        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4138        let mut b = stream.launch_builder(&f);
4139        b.arg(k_row)
4140            .arg(v_row)
4141            .arg(kc)
4142            .arg(vc)
4143            .arg(&ti)
4144            .arg(&kdk)
4145            .arg(&kdv)
4146            .arg(&ktb)
4147            .arg(&vtb);
4148        unsafe {
4149            b.launch(cfg)?;
4150        }
4151        Ok(())
4152    }
4153
4154    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
4155    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
4156    pub fn copy_view_into(
4157        &self,
4158        dst: &mut CudaSlice<f32>,
4159        off: usize,
4160        src: &cudarc::driver::CudaView<f32>,
4161        len: usize,
4162    ) -> Result<(), Box<dyn std::error::Error>> {
4163        let mut view = dst.slice_mut(off..off + len);
4164        self.gpu
4165            .stream()
4166            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4167        Ok(())
4168    }
4169
4170    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
4171    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
4172    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
4173    pub fn clone_dtod(
4174        &self,
4175        src: &CudaSlice<f32>,
4176    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4177        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4178        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4179        Ok(dst)
4180    }
4181
4182    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4183    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4184    pub fn dtod_copy_view(
4185        &self,
4186        src: &cudarc::driver::CudaView<f32>,
4187        dst: &mut CudaSlice<f32>,
4188    ) -> Result<(), Box<dyn std::error::Error>> {
4189        self.gpu.stream().memcpy_dtod(src, dst)?;
4190        Ok(())
4191    }
4192
4193    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4194    pub fn dtod_copy_view_i8(
4195        &self,
4196        src: &cudarc::driver::CudaView<i8>,
4197        dst: &mut CudaSlice<i8>,
4198    ) -> Result<(), Box<dyn std::error::Error>> {
4199        self.gpu.stream().memcpy_dtod(src, dst)?;
4200        Ok(())
4201    }
4202
4203    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4204    pub fn dtod_copy_into(
4205        &self,
4206        src: &CudaSlice<f32>,
4207        dst: &mut CudaSlice<f32>,
4208        offset: usize,
4209    ) -> Result<(), Box<dyn std::error::Error>> {
4210        let n = src.len();
4211        let mut dv = dst.slice_mut(offset..offset + n);
4212        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4213        Ok(())
4214    }
4215
4216    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4217    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4218    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4219    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4220    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4221    pub fn copy_batch_uniform_f32(
4222        &self,
4223        table: &CudaSlice<u64>,
4224        n: usize,
4225        words: usize,
4226    ) -> Result<(), Box<dyn std::error::Error>> {
4227        if n == 0 || words == 0 {
4228            return Ok(());
4229        }
4230        debug_assert!(
4231            table.len() >= 2 * n,
4232            "pointer table must hold n srcs + n dsts"
4233        );
4234        let f = self.func("copy_batch_uniform_f32");
4235        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4236        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4237        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4238        let (ni, wi) = (n as i32, words as i32);
4239        let cfg = LaunchConfig {
4240            grid_dim: (chunks, n as u32, 1),
4241            block_dim: (256, 1, 1),
4242            shared_mem_bytes: 0,
4243        };
4244        let __s = self.gpu.stream();
4245        let mut b = __s.launch_builder(&f);
4246        b.arg(table).arg(&ni).arg(&wi);
4247        unsafe {
4248            b.launch(cfg)?;
4249        }
4250        Ok(())
4251    }
4252
4253    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4254    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4255    pub fn htod_u64_into(
4256        &self,
4257        v: &[u64],
4258        dst: &mut CudaSlice<u64>,
4259    ) -> Result<(), Box<dyn std::error::Error>> {
4260        let mut view = dst.slice_mut(0..v.len());
4261        self.gpu.stream().memcpy_htod(v, &mut view)?;
4262        Ok(())
4263    }
4264
4265    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4266    /// device pointer-table entry at run time, so a captured graph follows the gdn
4267    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4268    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4269    pub fn copy_indirect_src_f32(
4270        &self,
4271        src_entry: &cudarc::driver::CudaView<u64>,
4272        dst: &mut CudaSlice<f32>,
4273        dst_off: usize,
4274        words: usize,
4275    ) -> Result<(), Box<dyn std::error::Error>> {
4276        let f = self.func("copy_indirect_src_f32");
4277        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4278        let wi = words as i32;
4279        let cfg = LaunchConfig {
4280            grid_dim: (chunks, 1, 1),
4281            block_dim: (256, 1, 1),
4282            shared_mem_bytes: 0,
4283        };
4284        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4285        let __s = self.gpu.stream();
4286        let mut b = __s.launch_builder(&f);
4287        b.arg(src_entry).arg(&mut dv).arg(&wi);
4288        unsafe {
4289            b.launch(cfg)?;
4290        }
4291        Ok(())
4292    }
4293
4294    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4295    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4296        self.alloc_uninit::<i8>(n)
4297    }
4298
4299    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4300    pub fn qmatvec(
4301        &self,
4302        w: &CudaSlice<u8>,
4303        x: &CudaSlice<f32>,
4304        m: usize,
4305        in_f: usize,
4306        out_f: usize,
4307        qtype: i32,
4308        row_bytes: usize,
4309    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4310        let f = self.func("qmatvec_f32");
4311        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4312        let cfg = LaunchConfig {
4313            grid_dim: (out_f as u32, m as u32, 1),
4314            block_dim: (256, 1, 1),
4315            shared_mem_bytes: 0,
4316        };
4317        let (inf, outf, mi, qt, rb) =
4318            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4319        let __s_b = self.gpu.stream();
4320        let mut b = __s_b.launch_builder(&f);
4321        b.arg(w)
4322            .arg(x)
4323            .arg(&mut y)
4324            .arg(&inf)
4325            .arg(&outf)
4326            .arg(&mi)
4327            .arg(&qt)
4328            .arg(&rb);
4329        unsafe {
4330            b.launch(cfg)?;
4331        }
4332        Ok(y)
4333    }
4334
4335    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4336    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4337        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4338        self.keep_if_capturing(&s);
4339        Ok(s)
4340    }
4341
4342    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4343    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4344    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4345    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4346        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4347        self.keep_if_capturing(&s);
4348        Ok(s)
4349    }
4350
4351    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4352    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4353    pub fn memset_zeros_view(
4354        &self,
4355        dst: &mut cudarc::driver::CudaViewMut<f32>,
4356    ) -> Result<(), Box<dyn std::error::Error>> {
4357        self.gpu.stream().memset_zeros(dst)?;
4358        Ok(())
4359    }
4360
4361    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4362    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4363    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4364    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4365    /// stream would require an event).
4366    pub fn stage_expert(
4367        &self,
4368        host_bytes: &[u8],
4369        scratch: &mut CudaSlice<u8>,
4370        off: usize,
4371    ) -> Result<(), Box<dyn std::error::Error>> {
4372        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4373        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4374        Ok(())
4375    }
4376
4377    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4378    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4379    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4380    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4381    /// One CTA per token row, 256 threads (one per expert).
4382    pub fn moe_router_topk(
4383        &self,
4384        logits: &CudaSlice<f32>,
4385        t: usize,
4386        n_expert: usize,
4387        n_used: usize,
4388    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4389        let f = self.func("moe_router_topk_f32");
4390        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4391        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4392        let cfg = LaunchConfig {
4393            grid_dim: (t as u32, 1, 1),
4394            block_dim: (n_expert as u32, 1, 1),
4395            shared_mem_bytes: 0,
4396        };
4397        let (ne, nu) = (n_expert as i32, n_used as i32);
4398        let __s_b = self.gpu.stream();
4399        let mut b = __s_b.launch_builder(&f);
4400        b.arg(logits)
4401            .arg(&mut sel_idx)
4402            .arg(&mut sel_w)
4403            .arg(&ne)
4404            .arg(&nu);
4405        unsafe {
4406            b.launch(cfg)?;
4407        }
4408        Ok((sel_idx, sel_w))
4409    }
4410
4411    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4412    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4413    pub fn moe_router_topk_scaled(
4414        &self,
4415        logits: &CudaSlice<f32>,
4416        t: usize,
4417        n_expert: usize,
4418        n_used: usize,
4419        ex_scale: &CudaSlice<f32>,
4420    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4421        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4422        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4423        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4424        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4425        let f = self.func("moe_router_topk_scaled_f32");
4426        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4427        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4428        let cfg = LaunchConfig {
4429            grid_dim: (t as u32, 1, 1),
4430            block_dim: (n_expert as u32, 1, 1),
4431            shared_mem_bytes: 0,
4432        };
4433        let (ne, nu) = (n_expert as i32, n_used as i32);
4434        let __s_b = self.gpu.stream();
4435        let mut b = __s_b.launch_builder(&f);
4436        b.arg(logits)
4437            .arg(&mut sel_idx)
4438            .arg(&mut sel_w)
4439            .arg(&ne)
4440            .arg(&nu)
4441            .arg(ex_scale);
4442        unsafe {
4443            b.launch(cfg)?;
4444        }
4445        Ok((sel_idx, sel_w))
4446    }
4447
4448    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4449    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4450    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4451    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4452    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4453    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4454    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4455    pub fn moe_router_topk_host(
4456        &self,
4457        logits: &CudaSlice<f32>,
4458        t: usize,
4459        n_expert: usize,
4460        n_used: usize,
4461    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4462        let f = self.func("moe_router_topk_f32");
4463        let n = t * n_used;
4464        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4465        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4466        let cfg = LaunchConfig {
4467            grid_dim: (t as u32, 1, 1),
4468            block_dim: (n_expert as u32, 1, 1),
4469            shared_mem_bytes: 0,
4470        };
4471        let (ne, nu) = (n_expert as i32, n_used as i32);
4472        let __s_b = self.gpu.stream();
4473        let mut b = __s_b.launch_builder(&f);
4474        b.arg(logits)
4475            .arg(&mut sel_idx)
4476            .arg(&mut sel_w)
4477            .arg(&ne)
4478            .arg(&nu);
4479        unsafe {
4480            b.launch(cfg)?;
4481        }
4482        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4483        let bytes = n * 8;
4484        let mut guard = self.router_stage.lock().unwrap();
4485        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4486            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4487        }
4488        let stage = guard.as_mut().unwrap();
4489        let (si, sw) = unsafe {
4490            (
4491                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4492                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4493            )
4494        };
4495        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4496        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4497        self.gpu.stream().synchronize()?; // ONE sync for both
4498        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4499    }
4500
4501    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4502    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4503    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4504    #[allow(clippy::too_many_arguments)]
4505    pub fn moe_router_sigmoid_topk(
4506        &self,
4507        logits: &CudaSlice<f32>,
4508        t: usize,
4509        n_expert: usize,
4510        n_used: usize,
4511        active_count: usize,
4512        correction_bias: &CudaSlice<f32>,
4513        active: &CudaSlice<u8>,
4514        scaling_factor: f32,
4515        route_norm: bool,
4516    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4517        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4518        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4519            return Err(format!(
4520                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4521            )
4522            .into());
4523        }
4524        if logits.len() < t * n_expert
4525            || correction_bias.len() != n_expert
4526            || active.len() != n_expert
4527        {
4528            return Err(format!(
4529                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4530                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4531            ).into());
4532        }
4533        let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4534            // Latency twin of the dexp arm (identical outputs): barrier-lean top-k
4535            // over the dexp scoring class. Composes the two doors it rides.
4536            self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4537        } else if crate::sig_expf_dev_on() {
4538            self.func("moe_router_sigmoid_topk_f32_dexp")
4539        } else if crate::topk_fast_on() {
4540            self.func("moe_router_sigmoid_topk_f32_fast")
4541        } else {
4542            self.func("moe_router_sigmoid_topk_f32")
4543        };
4544        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4545        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4546        let threads = n_expert.div_ceil(32) * 32;
4547        let cfg = LaunchConfig {
4548            grid_dim: (t as u32, 1, 1),
4549            block_dim: (threads as u32, 1, 1),
4550            shared_mem_bytes: 0,
4551        };
4552        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4553        let __s_b = self.gpu.stream();
4554        let mut b = __s_b.launch_builder(&f);
4555        b.arg(logits)
4556            .arg(correction_bias)
4557            .arg(active)
4558            .arg(&mut sel_idx)
4559            .arg(&mut sel_w)
4560            .arg(&ne)
4561            .arg(&nu)
4562            .arg(&scaling_factor)
4563            .arg(&rn);
4564        unsafe {
4565            b.launch(cfg)?;
4566        }
4567        Ok((sel_idx, sel_w))
4568    }
4569
4570    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
4571    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
4572    #[allow(clippy::too_many_arguments)]
4573    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
4574    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
4575    /// the model engine can wait on it with a same-device stream memop.
4576    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4577        if ptr == 0 {
4578            return Err("ring_flag_raw: unarmed flag".into());
4579        }
4580        let f = self.func("memra_ring_flag");
4581        let cfg = LaunchConfig {
4582            grid_dim: (1, 1, 1),
4583            block_dim: (32, 1, 1),
4584            shared_mem_bytes: 0,
4585        };
4586        let __s_b = self.gpu.stream();
4587        let mut b = __s_b.launch_builder(&f);
4588        b.arg(&ptr).arg(&value);
4589        unsafe {
4590            b.launch(cfg)?;
4591        }
4592        Ok(())
4593    }
4594
4595    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
4596    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
4597    pub fn moe_sel_w_mirror(
4598        &self,
4599        sel_src: &CudaSlice<i32>,
4600        w_src: &CudaSlice<f32>,
4601        sel_dst: &mut CudaSlice<i32>,
4602        w_dst: &mut CudaSlice<f32>,
4603        n: usize,
4604    ) -> Result<(), Box<dyn std::error::Error>> {
4605        if n == 0
4606            || n > 32
4607            || sel_src.len() < n
4608            || w_src.len() < n
4609            || sel_dst.len() < n
4610            || w_dst.len() < n
4611        {
4612            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4613        }
4614        let f = self.func("moe_sel_w_mirror");
4615        let cfg = LaunchConfig {
4616            grid_dim: (1, 1, 1),
4617            block_dim: (32, 1, 1),
4618            shared_mem_bytes: 0,
4619        };
4620        let ni = n as i32;
4621        let __s_b = self.gpu.stream();
4622        let mut b = __s_b.launch_builder(&f);
4623        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
4624        unsafe {
4625            b.launch(cfg)?;
4626        }
4627        Ok(())
4628    }
4629
4630    pub fn moe_router_sigmoid_topk_into(
4631        &self,
4632        logits: &CudaSlice<f32>,
4633        t: usize,
4634        n_expert: usize,
4635        n_used: usize,
4636        active_count: usize,
4637        correction_bias: &CudaSlice<f32>,
4638        active: &CudaSlice<u8>,
4639        scaling_factor: f32,
4640        route_norm: bool,
4641        sel_idx: &mut CudaSlice<i32>,
4642        sel_w: &mut CudaSlice<f32>,
4643    ) -> Result<(), Box<dyn std::error::Error>> {
4644        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4645        if n_expert == 0
4646            || n_expert > 1024
4647            || n_used == 0
4648            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
4649            || n_used > n_expert
4650            || logits.len() < t * n_expert
4651            || correction_bias.len() != n_expert
4652            || active.len() != n_expert
4653            || sel_idx.len() < t * n_used
4654            || sel_w.len() < t * n_used
4655        {
4656            return Err("sigmoid router _into geometry mismatch".into());
4657        }
4658        let f = if crate::sig_expf_dev_on() {
4659            self.func("moe_router_sigmoid_topk_f32_dexp")
4660        } else if crate::topk_fast_on() {
4661            self.func("moe_router_sigmoid_topk_f32_fast")
4662        } else {
4663            self.func("moe_router_sigmoid_topk_f32")
4664        };
4665        let threads = n_expert.div_ceil(32) * 32;
4666        let cfg = LaunchConfig {
4667            grid_dim: (t as u32, 1, 1),
4668            block_dim: (threads as u32, 1, 1),
4669            shared_mem_bytes: 0,
4670        };
4671        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4672        let __s_b = self.gpu.stream();
4673        let mut b = __s_b.launch_builder(&f);
4674        b.arg(logits)
4675            .arg(correction_bias)
4676            .arg(active)
4677            .arg(&mut *sel_idx)
4678            .arg(&mut *sel_w)
4679            .arg(&ne)
4680            .arg(&nu)
4681            .arg(&scaling_factor)
4682            .arg(&rn);
4683        unsafe {
4684            b.launch(cfg)?;
4685        }
4686        Ok(())
4687    }
4688
4689    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4690    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4691    #[allow(clippy::too_many_arguments)]
4692    pub fn moe_router_sigmoid_topk_host(
4693        &self,
4694        logits: &CudaSlice<f32>,
4695        t: usize,
4696        n_expert: usize,
4697        n_used: usize,
4698        active_count: usize,
4699        correction_bias: &CudaSlice<f32>,
4700        active: &CudaSlice<u8>,
4701        scaling_factor: f32,
4702        route_norm: bool,
4703    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4704        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4705            logits,
4706            t,
4707            n_expert,
4708            n_used,
4709            active_count,
4710            correction_bias,
4711            active,
4712            scaling_factor,
4713            route_norm,
4714        )?;
4715        let n = t * n_used;
4716        let bytes = n * 8;
4717        let mut guard = self.router_stage.lock().unwrap();
4718        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4719            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4720        }
4721        let stage = guard.as_mut().unwrap();
4722        let (si, sw) = unsafe {
4723            (
4724                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4725                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4726            )
4727        };
4728        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4729        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4730        self.gpu.stream().synchronize()?;
4731        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4732    }
4733
4734    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4735    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4736    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4737    pub fn stage_expert_async(
4738        &self,
4739        host_bytes: &[u8],
4740        scratch: &mut CudaSlice<u8>,
4741        off: usize,
4742    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4743        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4744        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4745        Ok(self.copy_stream.record_event(None)?)
4746    }
4747
4748    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4749    pub fn compute_wait(
4750        &self,
4751        ev: &cudarc::driver::CudaEvent,
4752    ) -> Result<(), Box<dyn std::error::Error>> {
4753        self.gpu.stream().wait(ev)?;
4754        Ok(())
4755    }
4756
4757    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4758    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4759    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4760    /// CudaView base+offset pointer is honored by the launch arg.
4761    pub fn qmatvec_view(
4762        &self,
4763        w: &CudaSlice<u8>,
4764        range: std::ops::Range<usize>,
4765        x: &cudarc::driver::CudaView<f32>,
4766        m: usize,
4767        in_f: usize,
4768        out_f: usize,
4769        qtype: i32,
4770        row_bytes: usize,
4771    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4772        let f = self.func("qmatvec_f32");
4773        let wv = w.slice(range); // CudaView<u8>, offset honored
4774        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4775        let cfg = LaunchConfig {
4776            grid_dim: (out_f as u32, m as u32, 1),
4777            block_dim: (256, 1, 1),
4778            shared_mem_bytes: 0,
4779        };
4780        let (inf, outf, mi, qt, rb) =
4781            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4782        let __s_b = self.gpu.stream();
4783        let mut b = __s_b.launch_builder(&f);
4784        b.arg(&wv)
4785            .arg(x)
4786            .arg(&mut y)
4787            .arg(&inf)
4788            .arg(&outf)
4789            .arg(&mi)
4790            .arg(&qt)
4791            .arg(&rb);
4792        unsafe {
4793            b.launch(cfg)?;
4794        }
4795        Ok(y)
4796    }
4797
4798    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4799    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4800    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4801    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4802    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4803    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4804    #[allow(clippy::too_many_arguments)]
4805    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4806    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4807    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4808    pub fn moe_gate_up_silu8_q8(
4809        &self,
4810        gp: WPtr8,
4811        up: WPtr8,
4812        aq: &CudaSlice<i8>,
4813        ad: &CudaSlice<f32>,
4814        in_f: usize,
4815        n_ff: usize,
4816        n_used: usize,
4817        qt_g: i32,
4818        qt_u: i32,
4819        rb_g: usize,
4820        rb_u: usize,
4821    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4822        let f = self.func("moe_gate_up_silu8_q8");
4823        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4824        let cfg = LaunchConfig {
4825            grid_dim: (n_ff as u32, n_used as u32, 1),
4826            block_dim: (32, 1, 1),
4827            shared_mem_bytes: 0,
4828        };
4829        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4830        let __s_b = self.gpu.stream();
4831        let mut b = __s_b.launch_builder(&f);
4832        b.arg(&gp)
4833            .arg(&up)
4834            .arg(aq)
4835            .arg(ad)
4836            .arg(&mut act)
4837            .arg(&inf)
4838            .arg(&nff)
4839            .arg(&qt_g)
4840            .arg(&qt_u)
4841            .arg(&rbg)
4842            .arg(&rbu);
4843        unsafe {
4844            b.launch(cfg)?;
4845        }
4846        Ok(act)
4847    }
4848
4849    #[allow(clippy::too_many_arguments)]
4850    pub fn moe_down8_fma_q8(
4851        &self,
4852        dp: WPtr8,
4853        w: F32x8,
4854        aq2: &CudaSlice<i8>,
4855        ad2: &CudaSlice<f32>,
4856        dst: &mut cudarc::driver::CudaViewMut<f32>,
4857        in_f: usize,
4858        out_f: usize,
4859        n_used: usize,
4860        qt: i32,
4861        rb: usize,
4862    ) -> Result<(), Box<dyn std::error::Error>> {
4863        let f = self.func("moe_down8_fma_q8");
4864        let cfg = LaunchConfig {
4865            grid_dim: (out_f as u32, 1, 1),
4866            block_dim: (32, 1, 1),
4867            shared_mem_bytes: 0,
4868        };
4869        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4870        let __s_b = self.gpu.stream();
4871        let mut b = __s_b.launch_builder(&f);
4872        b.arg(&dp)
4873            .arg(&w)
4874            .arg(aq2)
4875            .arg(ad2)
4876            .arg(dst)
4877            .arg(&inf)
4878            .arg(&outf)
4879            .arg(&nu)
4880            .arg(&qt)
4881            .arg(&rbi);
4882        unsafe {
4883            b.launch(cfg)?;
4884        }
4885        Ok(())
4886    }
4887
4888    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4889    pub fn qmatvec_expert_q8(
4890        &self,
4891        w: &CudaSlice<u8>,
4892        range: std::ops::Range<usize>,
4893        aq: &CudaSlice<i8>,
4894        ad: &CudaSlice<f32>,
4895        m: usize,
4896        in_f: usize,
4897        out_f: usize,
4898        qtype: i32,
4899        row_bytes: usize,
4900    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4901        let f = self.func("qmatvec_expert_q8");
4902        let wv = w.slice(range);
4903        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4904        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4905        let cfg = LaunchConfig {
4906            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4907            block_dim: (32, ROWS, 1),
4908            shared_mem_bytes: 0,
4909        };
4910        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4911        let __s_b = self.gpu.stream();
4912        let mut b = __s_b.launch_builder(&f);
4913        b.arg(&wv)
4914            .arg(aq)
4915            .arg(ad)
4916            .arg(&mut y)
4917            .arg(&inf)
4918            .arg(&outf)
4919            .arg(&mi)
4920            .arg(&qtype)
4921            .arg(&rbi);
4922        unsafe {
4923            b.launch(cfg)?;
4924        }
4925        Ok(y)
4926    }
4927
4928    pub fn moe_gate_up_silu8(
4929        &self,
4930        gp: WPtr8,
4931        up: WPtr8,
4932        x: &cudarc::driver::CudaView<f32>,
4933        in_f: usize,
4934        n_ff: usize,
4935        n_used: usize,
4936        qt_g: i32,
4937        qt_u: i32,
4938        rb_g: usize,
4939        rb_u: usize,
4940    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4941        let f = self.func("moe_gate_up_silu8_f32");
4942        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4943        let cfg = LaunchConfig {
4944            grid_dim: (n_ff as u32, n_used as u32, 1),
4945            block_dim: (256, 1, 1),
4946            shared_mem_bytes: 0,
4947        };
4948        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4949        let __s_b = self.gpu.stream();
4950        let mut b = __s_b.launch_builder(&f);
4951        b.arg(&gp)
4952            .arg(&up)
4953            .arg(x)
4954            .arg(&mut act)
4955            .arg(&inf)
4956            .arg(&nff)
4957            .arg(&qt_g)
4958            .arg(&qt_u)
4959            .arg(&rbg)
4960            .arg(&rbu);
4961        unsafe {
4962            b.launch(cfg)?;
4963        }
4964        Ok(act)
4965    }
4966
4967    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4968    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4969    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4970    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4971    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4972    #[allow(clippy::too_many_arguments)]
4973    pub fn moe_down8_fma_into(
4974        &self,
4975        dp: WPtr8,
4976        w: F32x8,
4977        act: &CudaSlice<f32>,
4978        dst: &mut cudarc::driver::CudaViewMut<f32>,
4979        in_f: usize,
4980        out_f: usize,
4981        n_used: usize,
4982        qt: i32,
4983        rb: usize,
4984    ) -> Result<(), Box<dyn std::error::Error>> {
4985        let f = self.func("moe_down8_fma_f32");
4986        let cfg = LaunchConfig {
4987            grid_dim: (out_f as u32, 1, 1),
4988            block_dim: (256, 1, 1),
4989            shared_mem_bytes: 0,
4990        };
4991        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4992        let __s_b = self.gpu.stream();
4993        let mut b = __s_b.launch_builder(&f);
4994        b.arg(&dp)
4995            .arg(&w)
4996            .arg(act)
4997            .arg(dst)
4998            .arg(&inf)
4999            .arg(&outf)
5000            .arg(&nu)
5001            .arg(&qt)
5002            .arg(&rbv);
5003        unsafe {
5004            b.launch(cfg)?;
5005        }
5006        Ok(())
5007    }
5008
5009    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
5010    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
5011    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
5012    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
5013    #[allow(clippy::too_many_arguments)]
5014    /// dp4a q8 twin of the _dev pair (resident-experts arc).
5015    ///
5016    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
5017    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
5018    /// down's FMA chain stays slot-ordered serial). Seams:
5019    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
5020    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
5021    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
5022    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
5023    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
5024    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
5025    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
5026    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
5027    ///                       only) | w8h2 (h2 x slot-parallel)
5028    #[allow(clippy::too_many_arguments)]
5029    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
5030    #[allow(clippy::too_many_arguments)]
5031    pub fn moe_pairs_matvec_q8(
5032        &self,
5033        table: &CudaSlice<u64>,
5034        proj: i32,
5035        pair_tok: &CudaSlice<i32>,
5036        pair_ex: &CudaSlice<i32>,
5037        aq: &CudaSlice<i8>,
5038        ad: &CudaSlice<f32>,
5039        in_f: usize,
5040        out_f: usize,
5041        n_expert: usize,
5042        n_pairs: usize,
5043        qtype: i32,
5044        row_bytes: usize,
5045    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5046        let f = self.func("moe_pairs_matvec_q8");
5047        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5048        const ROWS: u32 = 4;
5049        let cfg = LaunchConfig {
5050            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5051            block_dim: (32, ROWS, 1),
5052            shared_mem_bytes: 0,
5053        };
5054        let (inf, outf, ne, np, rbi) = (
5055            in_f as i32,
5056            out_f as i32,
5057            n_expert as i32,
5058            n_pairs as i32,
5059            row_bytes as i64,
5060        );
5061        let __s_b = self.gpu.stream();
5062        let mut b = __s_b.launch_builder(&f);
5063        b.arg(table)
5064            .arg(&proj)
5065            .arg(pair_tok)
5066            .arg(pair_ex)
5067            .arg(aq)
5068            .arg(ad)
5069            .arg(&mut y)
5070            .arg(&inf)
5071            .arg(&outf)
5072            .arg(&ne)
5073            .arg(&np)
5074            .arg(&qtype)
5075            .arg(&rbi);
5076        unsafe {
5077            b.launch(cfg)?;
5078        }
5079        Ok(y)
5080    }
5081
5082    /// Expert-major pair matvec (weight-reuse across each expert's token group).
5083    #[allow(clippy::too_many_arguments)]
5084    pub fn moe_pairs_matvec_q8_em(
5085        &self,
5086        table: &CudaSlice<u64>,
5087        proj: i32,
5088        ex_ids: &CudaSlice<i32>,
5089        ex_off: &CudaSlice<i32>,
5090        ex_pairs: &CudaSlice<i32>,
5091        pair_tok: &CudaSlice<i32>,
5092        aq: &CudaSlice<i8>,
5093        ad: &CudaSlice<f32>,
5094        in_f: usize,
5095        out_f: usize,
5096        n_expert: usize,
5097        n_active: usize,
5098        n_pairs: usize,
5099        qtype: i32,
5100        row_bytes: usize,
5101    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5102        let f = self.func("moe_pairs_matvec_q8_em");
5103        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5104        const ROWS: u32 = 4;
5105        let cfg = LaunchConfig {
5106            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5107            block_dim: (32, ROWS, 1),
5108            shared_mem_bytes: 0,
5109        };
5110        let (inf, outf, ne, na, rbi) = (
5111            in_f as i32,
5112            out_f as i32,
5113            n_expert as i32,
5114            n_active as i32,
5115            row_bytes as i64,
5116        );
5117        let __s_b = self.gpu.stream();
5118        let mut b = __s_b.launch_builder(&f);
5119        b.arg(table)
5120            .arg(&proj)
5121            .arg(ex_ids)
5122            .arg(ex_off)
5123            .arg(ex_pairs)
5124            .arg(pair_tok)
5125            .arg(aq)
5126            .arg(ad)
5127            .arg(&mut y)
5128            .arg(&inf)
5129            .arg(&outf)
5130            .arg(&ne)
5131            .arg(&na)
5132            .arg(&qtype)
5133            .arg(&rbi);
5134        unsafe {
5135            b.launch(cfg)?;
5136        }
5137        Ok(y)
5138    }
5139
5140    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
5141    // weight group once per (row,group) then dp4a's across the expert's token group.
5142    #[allow(clippy::too_many_arguments)]
5143    pub fn moe_pairs_matvec_q8_dec(
5144        &self,
5145        table: &CudaSlice<u64>,
5146        proj: i32,
5147        ex_ids: &CudaSlice<i32>,
5148        ex_off: &CudaSlice<i32>,
5149        ex_pairs: &CudaSlice<i32>,
5150        pair_tok: &CudaSlice<i32>,
5151        aq: &CudaSlice<i8>,
5152        ad: &CudaSlice<f32>,
5153        in_f: usize,
5154        out_f: usize,
5155        n_expert: usize,
5156        n_active: usize,
5157        n_pairs: usize,
5158        qtype: i32,
5159        row_bytes: usize,
5160    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5161        let f = self.func("moe_pairs_matvec_q8_dec");
5162        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5163        const ROWS: u32 = 4;
5164        let cfg = LaunchConfig {
5165            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5166            block_dim: (32, ROWS, 1),
5167            shared_mem_bytes: 0,
5168        };
5169        let (inf, outf, ne, na, rbi) = (
5170            in_f as i32,
5171            out_f as i32,
5172            n_expert as i32,
5173            n_active as i32,
5174            row_bytes as i64,
5175        );
5176        let __s_b = self.gpu.stream();
5177        let mut b = __s_b.launch_builder(&f);
5178        b.arg(table)
5179            .arg(&proj)
5180            .arg(ex_ids)
5181            .arg(ex_off)
5182            .arg(ex_pairs)
5183            .arg(pair_tok)
5184            .arg(aq)
5185            .arg(ad)
5186            .arg(&mut y)
5187            .arg(&inf)
5188            .arg(&outf)
5189            .arg(&ne)
5190            .arg(&na)
5191            .arg(&qtype)
5192            .arg(&rbi);
5193        unsafe {
5194            b.launch(cfg)?;
5195        }
5196        Ok(y)
5197    }
5198
5199    pub fn moe_pairs_gelu_mul(
5200        &self,
5201        gate: &CudaSlice<f32>,
5202        up: &CudaSlice<f32>,
5203        n: usize,
5204    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5205        let f = self.func("moe_pairs_gelu_mul");
5206        let mut act = self.alloc_uninit::<f32>(n)?;
5207        let cfg = LaunchConfig::for_num_elems(n as u32);
5208        let nl = n as i64;
5209        let __s_b = self.gpu.stream();
5210        let mut b = __s_b.launch_builder(&f);
5211        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5212        unsafe {
5213            b.launch(cfg)?;
5214        }
5215        Ok(act)
5216    }
5217
5218    pub fn moe_pairs_silu_mul(
5219        &self,
5220        gate: &CudaSlice<f32>,
5221        up: &CudaSlice<f32>,
5222        n: usize,
5223    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5224        let f = self.func("moe_pairs_silu_mul");
5225        let mut act = self.alloc_uninit::<f32>(n)?;
5226        let cfg = LaunchConfig::for_num_elems(n as u32);
5227        let nl = n as i64;
5228        let __s_b = self.gpu.stream();
5229        let mut b = __s_b.launch_builder(&f);
5230        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5231        unsafe {
5232            b.launch(cfg)?;
5233        }
5234        Ok(act)
5235    }
5236
5237    #[allow(clippy::too_many_arguments)]
5238    pub fn moe_pairs_scatter(
5239        &self,
5240        y_down: &CudaSlice<f32>,
5241        pair_w: &CudaSlice<f32>,
5242        tok_pair_off: &CudaSlice<i32>,
5243        tok_pair_ids: &CudaSlice<i32>,
5244        moe_out: &mut CudaSlice<f32>,
5245        t: usize,
5246        n_embd: usize,
5247    ) -> Result<(), Box<dyn std::error::Error>> {
5248        let f = self.func("moe_pairs_scatter");
5249        let cfg = LaunchConfig {
5250            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5251            block_dim: (256, 1, 1),
5252            shared_mem_bytes: 0,
5253        };
5254        let ne = n_embd as i32;
5255        let __s_b = self.gpu.stream();
5256        let mut b = __s_b.launch_builder(&f);
5257        b.arg(y_down)
5258            .arg(pair_w)
5259            .arg(tok_pair_off)
5260            .arg(tok_pair_ids)
5261            .arg(moe_out)
5262            .arg(&ne);
5263        unsafe {
5264            b.launch(cfg)?;
5265        }
5266        Ok(())
5267    }
5268
5269    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
5270    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
5271    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
5272    #[allow(clippy::too_many_arguments)]
5273    pub fn moe_gate_up_gelu8_dev_q8(
5274        &self,
5275        table: &CudaSlice<u64>,
5276        sel: &cudarc::driver::CudaView<i32>,
5277        aq: &CudaSlice<i8>,
5278        ad: &CudaSlice<f32>,
5279        in_f: usize,
5280        n_ff: usize,
5281        n_used: usize,
5282        n_expert: usize,
5283        qt_g: i32,
5284        qt_u: i32,
5285        rb_g: usize,
5286        rb_u: usize,
5287    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5288        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5289        let (inf, nff, ne, rbg, rbu) = (
5290            in_f as i32,
5291            n_ff as i32,
5292            n_expert as i32,
5293            rb_g as i64,
5294            rb_u as i64,
5295        );
5296        let f = self.func("moe_gate_up_gelu8_dev_q8");
5297        let cfg = LaunchConfig {
5298            grid_dim: (n_ff as u32, n_used as u32, 1),
5299            block_dim: (32, 1, 1),
5300            shared_mem_bytes: 0,
5301        };
5302        let __s_b = self.gpu.stream();
5303        let mut b = __s_b.launch_builder(&f);
5304        b.arg(table)
5305            .arg(sel)
5306            .arg(aq)
5307            .arg(ad)
5308            .arg(&mut act)
5309            .arg(&inf)
5310            .arg(&nff)
5311            .arg(&ne)
5312            .arg(&qt_g)
5313            .arg(&qt_u)
5314            .arg(&rbg)
5315            .arg(&rbu);
5316        unsafe {
5317            b.launch(cfg)?;
5318        }
5319        Ok(act)
5320    }
5321
5322    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5323    #[allow(clippy::too_many_arguments)]
5324    pub fn moe_gate_up_gelu8_dev_q8_rows(
5325        &self,
5326        table: &CudaSlice<u64>,
5327        sel: &CudaSlice<i32>,
5328        aq: &CudaSlice<i8>,
5329        ad: &CudaSlice<f32>,
5330        t: usize,
5331        in_f: usize,
5332        n_ff: usize,
5333        n_used: usize,
5334        n_expert: usize,
5335        qt_g: i32,
5336        qt_u: i32,
5337        rb_g: usize,
5338        rb_u: usize,
5339    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5340        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5341        let (inf, nff, ne, rbg, rbu, nu) = (
5342            in_f as i32,
5343            n_ff as i32,
5344            n_expert as i32,
5345            rb_g as i64,
5346            rb_u as i64,
5347            n_used as i32,
5348        );
5349        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5350        let cfg = LaunchConfig {
5351            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5352            block_dim: (32, 1, 1),
5353            shared_mem_bytes: 0,
5354        };
5355        let __s_b = self.gpu.stream();
5356        let mut b = __s_b.launch_builder(&f);
5357        b.arg(table)
5358            .arg(sel)
5359            .arg(aq)
5360            .arg(ad)
5361            .arg(&mut act)
5362            .arg(&inf)
5363            .arg(&nff)
5364            .arg(&ne)
5365            .arg(&qt_g)
5366            .arg(&qt_u)
5367            .arg(&rbg)
5368            .arg(&rbu)
5369            .arg(&nu);
5370        unsafe {
5371            b.launch(cfg)?;
5372        }
5373        Ok(act)
5374    }
5375
5376    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5377    #[allow(clippy::too_many_arguments)]
5378    pub fn moe_gate_up_gelu8_dev_q8_csr(
5379        &self,
5380        table: &CudaSlice<u64>,
5381        sel: &CudaSlice<i32>,
5382        aq: &CudaSlice<i8>,
5383        ad: &CudaSlice<f32>,
5384        n_pairs: usize,
5385        in_f: usize,
5386        n_ff: usize,
5387        n_used: usize,
5388        n_expert: usize,
5389        qt_g: i32,
5390        qt_u: i32,
5391        rb_g: usize,
5392        rb_u: usize,
5393    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5394        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5395        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5396            in_f as i32,
5397            n_ff as i32,
5398            n_expert as i32,
5399            rb_g as i64,
5400            rb_u as i64,
5401            n_used as i32,
5402            n_pairs as i32,
5403        );
5404        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5405        let cfg = LaunchConfig {
5406            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5407            block_dim: (32, 1, 1),
5408            shared_mem_bytes: 0,
5409        };
5410        let __s_b = self.gpu.stream();
5411        let mut b = __s_b.launch_builder(&f);
5412        b.arg(table)
5413            .arg(sel)
5414            .arg(aq)
5415            .arg(ad)
5416            .arg(&mut act)
5417            .arg(&inf)
5418            .arg(&nff)
5419            .arg(&ne)
5420            .arg(&qt_g)
5421            .arg(&qt_u)
5422            .arg(&rbg)
5423            .arg(&rbu)
5424            .arg(&nu)
5425            .arg(&npi);
5426        unsafe {
5427            b.launch(cfg)?;
5428        }
5429        Ok(act)
5430    }
5431
5432    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5433    #[allow(clippy::too_many_arguments)]
5434    pub fn moe_down8_fma_dev_q8_rows_g(
5435        &self,
5436        table: &CudaSlice<u64>,
5437        sel: &CudaSlice<i32>,
5438        w: &CudaSlice<f32>,
5439        aq2: &CudaSlice<i8>,
5440        ad2: &CudaSlice<f32>,
5441        dst: &mut CudaSlice<f32>,
5442        t: usize,
5443        in_f: usize,
5444        out_f: usize,
5445        n_used: usize,
5446        n_expert: usize,
5447        qt: i32,
5448        rb: usize,
5449    ) -> Result<(), Box<dyn std::error::Error>> {
5450        let (inf, outf, nu, ne, rbi) = (
5451            in_f as i32,
5452            out_f as i32,
5453            n_used as i32,
5454            n_expert as i32,
5455            rb as i64,
5456        );
5457        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5458        // eight warps, then replay the original slot-ordered FMA chain. Every
5459        // other shape retains the generic one-warp rows kernel.
5460        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5461        let f = self.func(if step_b1_w8 {
5462            "moe_down8_fma_dev_q8_rows_w8"
5463        } else {
5464            "moe_down8_fma_dev_q8_rows_g"
5465        });
5466        let cfg = LaunchConfig {
5467            grid_dim: (out_f as u32, 1, t as u32),
5468            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5469            shared_mem_bytes: 0,
5470        };
5471        let __s_b = self.gpu.stream();
5472        let mut b = __s_b.launch_builder(&f);
5473        b.arg(table)
5474            .arg(sel)
5475            .arg(w)
5476            .arg(aq2)
5477            .arg(ad2)
5478            .arg(dst)
5479            .arg(&inf)
5480            .arg(&outf)
5481            .arg(&nu)
5482            .arg(&ne)
5483            .arg(&qt)
5484            .arg(&rbi);
5485        unsafe {
5486            b.launch(cfg)?;
5487        }
5488        Ok(())
5489    }
5490
5491    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5492    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5493    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5494    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5495        let (out_f, in_f) = (2048usize, 2816usize);
5496        let nblk = in_f / 32;
5497        let mut seed = 0x9E3779B97F4A7C15u64;
5498        let mut rng = move || {
5499            seed = seed
5500                .wrapping_mul(6364136223846793005)
5501                .wrapping_add(1442695040888963407);
5502            (seed >> 33) as u8
5503        };
5504        let mut w = vec![0u8; out_f * nblk * 18];
5505        for b in w.iter_mut() {
5506            *b = rng();
5507        }
5508        for r in 0..out_f {
5509            for g in 0..nblk {
5510                let off = (r * nblk + g) * 18;
5511                w[off] = 0x00;
5512                w[off + 1] = 0x2C; // sane half d
5513            }
5514        }
5515        let qplane = out_f * nblk * 16;
5516        let mut wrp = vec![0u8; w.len()];
5517        for r in 0..out_f {
5518            for g in 0..nblk {
5519                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5520                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5521                    .copy_from_slice(&src[0..2]);
5522                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5523            }
5524        }
5525        let w_d = self.htod_bytes(&w)?;
5526        let wrp_d = self.htod_bytes(&wrp)?;
5527        let mut aq = vec![0i8; m * in_f];
5528        for v in aq.iter_mut() {
5529            *v = rng() as i8;
5530        }
5531        let aq_d = self.htod_i8(&aq)?;
5532        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5533        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5534        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5535        const RPB: u32 = 4;
5536        let cfg = LaunchConfig {
5537            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5538            block_dim: (32, RPB, 1),
5539            shared_mem_bytes: 0,
5540        };
5541        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5542        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5543        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5544        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5545        {
5546            let __s_b = self.gpu.stream();
5547            let mut b = __s_b.launch_builder(&fb);
5548            b.arg(&w_d)
5549                .arg(&aq_d)
5550                .arg(&ad_d)
5551                .arg(&mut y0)
5552                .arg(&inf)
5553                .arg(&outf)
5554                .arg(&mi)
5555                .arg(&rb);
5556            unsafe {
5557                b.launch(cfg)?;
5558            }
5559            let __s_b = self.gpu.stream();
5560            let mut b = __s_b.launch_builder(&fr);
5561            b.arg(&wrp_d)
5562                .arg(&aq_d)
5563                .arg(&ad_d)
5564                .arg(&mut y1)
5565                .arg(&inf)
5566                .arg(&outf)
5567                .arg(&mi)
5568                .arg(&qp);
5569            unsafe {
5570                b.launch(cfg)?;
5571            }
5572        }
5573        self.gpu.stream().synchronize()?;
5574        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5575        let nd = h0
5576            .iter()
5577            .zip(&h1)
5578            .filter(|(a, b)| a.to_bits() != b.to_bits())
5579            .count();
5580        if nd != 0 {
5581            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5582        }
5583        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5584            self.gpu.stream().synchronize()?;
5585            let t0 = std::time::Instant::now();
5586            for _ in 0..500 {
5587                if rp {
5588                    let __s_b = self.gpu.stream();
5589                    let mut b = __s_b.launch_builder(&fr);
5590                    b.arg(&wrp_d)
5591                        .arg(&aq_d)
5592                        .arg(&ad_d)
5593                        .arg(&mut y1)
5594                        .arg(&inf)
5595                        .arg(&outf)
5596                        .arg(&mi)
5597                        .arg(&qp);
5598                    unsafe {
5599                        b.launch(cfg)?;
5600                    }
5601                } else {
5602                    let __s_b = self.gpu.stream();
5603                    let mut b = __s_b.launch_builder(&fb);
5604                    b.arg(&w_d)
5605                        .arg(&aq_d)
5606                        .arg(&ad_d)
5607                        .arg(&mut y0)
5608                        .arg(&inf)
5609                        .arg(&outf)
5610                        .arg(&mi)
5611                        .arg(&rb);
5612                    unsafe {
5613                        b.launch(cfg)?;
5614                    }
5615                }
5616            }
5617            self.gpu.stream().synchronize()?;
5618            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5619        };
5620        let _ = time(false)?;
5621        let _ = time(true)?; // warm
5622        Ok((time(false)?, time(true)?))
5623    }
5624
5625    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5626    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5627    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5628    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5629    pub fn build_q4_rp4(
5630        &self,
5631        t: &mut crate::model::GpuTensor,
5632    ) -> Result<(), Box<dyn std::error::Error>> {
5633        use crate::model::GpuTensor;
5634        let GpuTensor::Quant {
5635            bytes,
5636            qtype,
5637            row_bytes,
5638            ne,
5639            rp4,
5640            ..
5641        } = t
5642        else {
5643            return Ok(());
5644        };
5645        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5646            return Ok(());
5647        }
5648        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5649        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5650            return Ok(());
5651        }
5652        let nblk = in_f / 32;
5653        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5654        let f = self.func("q4_0_split_rp_build");
5655        let n = (out_f * nblk) as i32;
5656        let cfg = LaunchConfig {
5657            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5658            block_dim: (256, 1, 1),
5659            shared_mem_bytes: 0,
5660        };
5661        let (of, nb) = (out_f as i32, nblk as i32);
5662        let _ = n;
5663        let __s_b = self.gpu.stream();
5664        let mut b = __s_b.launch_builder(&f);
5665        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5666        unsafe {
5667            b.launch(cfg)?;
5668        }
5669        *rp4 = Some(dst);
5670        Ok(())
5671    }
5672
5673    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5674    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5675    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5676    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5677    pub fn build_q8_rp4(
5678        &self,
5679        t: &mut crate::model::GpuTensor,
5680    ) -> Result<(), Box<dyn std::error::Error>> {
5681        use crate::model::GpuTensor;
5682        let GpuTensor::Quant {
5683            bytes,
5684            qtype,
5685            row_bytes,
5686            ne,
5687            rp4,
5688            ..
5689        } = t
5690        else {
5691            return Ok(());
5692        };
5693        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5694            return Ok(());
5695        }
5696        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5697        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5698            return Ok(());
5699        }
5700        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5701        Ok(())
5702    }
5703
5704    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5705    /// mirror without a GpuTensor (same kernel the loader path above uses).
5706    pub fn build_q8_rp4_raw(
5707        &self,
5708        bytes: &CudaSlice<u8>,
5709        in_f: usize,
5710        out_f: usize,
5711    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5712        assert!(in_f % 32 == 0);
5713        let nblk = in_f / 32;
5714        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5715        let f = self.func("q8_0_split_rp_build");
5716        let cfg = LaunchConfig {
5717            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5718            block_dim: (256, 1, 1),
5719            shared_mem_bytes: 0,
5720        };
5721        let (of, nb) = (out_f as i32, nblk as i32);
5722        let __s_b = self.gpu.stream();
5723        let mut b = __s_b.launch_builder(&f);
5724        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5725        unsafe {
5726            b.launch(cfg)?;
5727        }
5728        Ok(dst)
5729    }
5730
5731    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5732    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5733    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5734    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5735    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5736    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5737    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5738    pub fn build_q4k_rp4(
5739        &self,
5740        t: &mut crate::model::GpuTensor,
5741    ) -> Result<(), Box<dyn std::error::Error>> {
5742        use crate::model::GpuTensor;
5743        let GpuTensor::Quant {
5744            bytes,
5745            qtype,
5746            row_bytes,
5747            ne,
5748            rp4,
5749            ..
5750        } = t
5751        else {
5752            return Ok(());
5753        };
5754        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5755            return Ok(());
5756        }
5757        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5758        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5759            return Ok(());
5760        }
5761        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5762        Ok(())
5763    }
5764
5765    pub fn build_q6k_rp4(
5766        &self,
5767        t: &mut crate::model::GpuTensor,
5768    ) -> Result<(), Box<dyn std::error::Error>> {
5769        use crate::model::GpuTensor;
5770        let GpuTensor::Quant {
5771            bytes,
5772            qtype,
5773            row_bytes,
5774            ne,
5775            rp4,
5776            ..
5777        } = t
5778        else {
5779            return Ok(());
5780        };
5781        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5782            return Ok(());
5783        }
5784        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5785        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5786            return Ok(());
5787        }
5788        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5789        Ok(())
5790    }
5791
5792    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5793    pub fn build_kq_rp4_raw(
5794        &self,
5795        bytes: &CudaSlice<u8>,
5796        in_f: usize,
5797        out_f: usize,
5798        qtype: i32,
5799    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5800        assert!(in_f % 256 == 0);
5801        let nsbk = in_f / 256;
5802        let (sb_bytes, kname) = match qtype {
5803            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5804            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5805            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5806        };
5807        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5808        let f = self.func(kname);
5809        let cfg = LaunchConfig {
5810            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5811            block_dim: (256, 1, 1),
5812            shared_mem_bytes: 0,
5813        };
5814        let (of, nb) = (out_f as i32, nsbk as i32);
5815        let __s_b = self.gpu.stream();
5816        let mut b = __s_b.launch_builder(&f);
5817        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5818        unsafe {
5819            b.launch(cfg)?;
5820        }
5821        Ok(dst)
5822    }
5823
5824    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5825    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5826    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5827    pub fn kqrp_enabled() -> bool {
5828        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5829        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5830            Ok("0") => false,
5831            Ok(_) => true,
5832            Err(_) => cfg!(memra_hopper_mma),
5833        })
5834    }
5835
5836    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5837    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5838    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5839    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5840    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5841    pub fn build_q4_rp_swap(
5842        &self,
5843        t: &mut crate::model::GpuTensor,
5844    ) -> Result<bool, Box<dyn std::error::Error>> {
5845        use crate::model::GpuTensor;
5846        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5847        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5848        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5849        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5850        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5851        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5852        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5853        // this fn's OWN builder serves may ever be swapped; everything else refuses
5854        // here, regardless of walk ordering.
5855        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5856            return Ok(false);
5857        }
5858        self.build_q4_rp4(t)?;
5859        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5860        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5861            return Ok(false);
5862        };
5863        match rp4.take() {
5864            Some(split) => {
5865                *bytes = split; // the GGUF-layout buffer drops here
5866                *rp = true;
5867                Ok(true)
5868            }
5869            None => Ok(false),
5870        }
5871    }
5872
5873    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5874    pub fn q4rp_enabled() -> bool {
5875        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5876        *ON.get_or_init(|| {
5877            std::env::var("MEMRA_Q4RP")
5878                .map(|v| v != "0")
5879                .unwrap_or(true)
5880        })
5881    }
5882
5883    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5884    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5885    pub fn copy_rows_strided(
5886        &self,
5887        src: &CudaSlice<f32>,
5888        dst: &mut CudaSlice<f32>,
5889        row_elems: usize,
5890        n_rows: usize,
5891        src_stride: usize,
5892        src_off: usize,
5893    ) -> Result<(), Box<dyn std::error::Error>> {
5894        let f = self.func("copy_rows_strided_f32");
5895        let cfg = LaunchConfig {
5896            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5897            block_dim: (256, 1, 1),
5898            shared_mem_bytes: 0,
5899        };
5900        let (re, nr) = (row_elems as i32, n_rows as i32);
5901        let (st, off) = (src_stride as i64, src_off as i64);
5902        let __s_b = self.gpu.stream();
5903        let mut b = __s_b.launch_builder(&f);
5904        b.arg(src)
5905            .arg(&mut *dst)
5906            .arg(&re)
5907            .arg(&nr)
5908            .arg(&st)
5909            .arg(&off);
5910        unsafe {
5911            b.launch(cfg)?;
5912        }
5913        Ok(())
5914    }
5915
5916    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
5917    ///
5918    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
5919    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
5920    /// one peer copy per token.
5921    pub fn place_rows_strided(
5922        &self,
5923        src: &CudaSlice<f32>,
5924        dst: &mut CudaSlice<f32>,
5925        row_elems: usize,
5926        n_rows: usize,
5927        dst_stride: usize,
5928        dst_off: usize,
5929    ) -> Result<(), Box<dyn std::error::Error>> {
5930        if row_elems == 0 || n_rows == 0 {
5931            return Err("strided row placement requires nonzero rows and row width".into());
5932        }
5933        let src_len = n_rows
5934            .checked_mul(row_elems)
5935            .ok_or("strided row placement source size overflow")?;
5936        let dst_len = n_rows
5937            .checked_sub(1)
5938            .and_then(|rows| rows.checked_mul(dst_stride))
5939            .and_then(|base| base.checked_add(dst_off))
5940            .and_then(|base| base.checked_add(row_elems))
5941            .ok_or("strided row placement destination size overflow")?;
5942        let row_end = dst_off
5943            .checked_add(row_elems)
5944            .ok_or("strided row placement row size overflow")?;
5945        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
5946            return Err(format!(
5947                "strided row placement geometry mismatch: src={} need_src={src_len} \
5948                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
5949                 dst_stride={dst_stride} dst_off={dst_off}",
5950                src.len(),
5951                dst.len(),
5952            )
5953            .into());
5954        }
5955        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
5956            return Err("strided row placement exceeds CUDA kernel geometry".into());
5957        }
5958        let f = self.func("place_rows_strided_f32");
5959        let cfg = LaunchConfig {
5960            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5961            block_dim: (256, 1, 1),
5962            shared_mem_bytes: 0,
5963        };
5964        let (re, nr) = (row_elems as i32, n_rows as i32);
5965        let (st, off) = (dst_stride as i64, dst_off as i64);
5966        let __s_b = self.gpu.stream();
5967        let mut b = __s_b.launch_builder(&f);
5968        b.arg(src)
5969            .arg(&mut *dst)
5970            .arg(&re)
5971            .arg(&nr)
5972            .arg(&st)
5973            .arg(&off);
5974        unsafe {
5975            b.launch(cfg)?;
5976        }
5977        Ok(())
5978    }
5979
5980    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5981    pub fn u32_set_k(
5982        &self,
5983        dst: &mut CudaSlice<u32>,
5984        v: u32,
5985        idx: usize,
5986    ) -> Result<(), Box<dyn std::error::Error>> {
5987        let f = self.func("u32_set_k");
5988        let cfg = LaunchConfig {
5989            grid_dim: (1, 1, 1),
5990            block_dim: (1, 1, 1),
5991            shared_mem_bytes: 0,
5992        };
5993        let ii = idx as i32;
5994        let __s_b = self.gpu.stream();
5995        let mut b = __s_b.launch_builder(&f);
5996        b.arg(dst).arg(&v).arg(&ii);
5997        unsafe {
5998            b.launch(cfg)?;
5999        }
6000        Ok(())
6001    }
6002
6003    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
6004    pub fn i32_add_k(
6005        &self,
6006        d: &mut CudaSlice<i32>,
6007        v: i32,
6008    ) -> Result<(), Box<dyn std::error::Error>> {
6009        let f = self.func("i32_add_k");
6010        let cfg = LaunchConfig {
6011            grid_dim: (1, 1, 1),
6012            block_dim: (32, 1, 1),
6013            shared_mem_bytes: 0,
6014        };
6015        let __s_b = self.gpu.stream();
6016        let mut b = __s_b.launch_builder(&f);
6017        b.arg(d).arg(&v);
6018        unsafe {
6019            b.launch(cfg)?;
6020        }
6021        Ok(())
6022    }
6023
6024    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
6025    pub fn i32_iota_from(
6026        &self,
6027        ctr: &CudaSlice<i32>,
6028        dst: &mut CudaSlice<i32>,
6029        n: usize,
6030    ) -> Result<(), Box<dyn std::error::Error>> {
6031        let f = self.func("i32_iota_from");
6032        let cfg = LaunchConfig::for_num_elems(n as u32);
6033        let ni = n as i32;
6034        let __s_b = self.gpu.stream();
6035        let mut b = __s_b.launch_builder(&f);
6036        b.arg(ctr).arg(dst).arg(&ni);
6037        unsafe {
6038            b.launch(cfg)?;
6039        }
6040        Ok(())
6041    }
6042
6043    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
6044    pub fn u32_map_k(
6045        &self,
6046        buf: &mut CudaSlice<u32>,
6047        map: &CudaSlice<u32>,
6048        idx: usize,
6049    ) -> Result<(), Box<dyn std::error::Error>> {
6050        let f = self.func("u32_map_k");
6051        let cfg = LaunchConfig {
6052            grid_dim: (1, 1, 1),
6053            block_dim: (1, 1, 1),
6054            shared_mem_bytes: 0,
6055        };
6056        let ii = idx as i32;
6057        let __s_b = self.gpu.stream();
6058        let mut b = __s_b.launch_builder(&f);
6059        b.arg(buf).arg(map).arg(&ii);
6060        unsafe {
6061            b.launch(cfg)?;
6062        }
6063        Ok(())
6064    }
6065
6066    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
6067    #[allow(clippy::too_many_arguments)]
6068    pub fn u32_pack2(
6069        &self,
6070        a: &CudaSlice<u32>,
6071        off_a: usize,
6072        n1: usize,
6073        b_in: &CudaSlice<u32>,
6074        n2: usize,
6075        out: &mut CudaSlice<u32>,
6076    ) -> Result<(), Box<dyn std::error::Error>> {
6077        let f = self.func("u32_pack2");
6078        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6079        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6080        let __s_b = self.gpu.stream();
6081        let mut b = __s_b.launch_builder(&f);
6082        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6083        unsafe {
6084            b.launch(cfg)?;
6085        }
6086        Ok(())
6087    }
6088
6089    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
6090    pub fn moe_w_exscale(
6091        &self,
6092        w: &mut CudaSlice<f32>,
6093        sel: &CudaSlice<i32>,
6094        s: &CudaSlice<f32>,
6095        n: usize,
6096    ) -> Result<(), Box<dyn std::error::Error>> {
6097        let f = self.func("moe_w_exscale");
6098        let cfg = LaunchConfig::for_num_elems(n as u32);
6099        let ni = n as i32;
6100        let __s_b = self.gpu.stream();
6101        let mut b = __s_b.launch_builder(&f);
6102        b.arg(w).arg(sel).arg(s).arg(&ni);
6103        unsafe {
6104            b.launch(cfg)?;
6105        }
6106        Ok(())
6107    }
6108
6109    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
6110    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
6111    pub fn moe_w_scale_by_expert(
6112        &self,
6113        w: &mut CudaSlice<f32>,
6114        sel: &CudaSlice<i32>,
6115        macros: &CudaSlice<f32>,
6116        n_expert: usize,
6117        n: usize,
6118    ) -> Result<(), Box<dyn std::error::Error>> {
6119        let f = self.func("moe_w_scale_by_expert");
6120        let cfg = LaunchConfig {
6121            grid_dim: (n.div_ceil(64) as u32, 1, 1),
6122            block_dim: (64, 1, 1),
6123            shared_mem_bytes: 0,
6124        };
6125        let (ne, nn) = (n_expert as i32, n as i32);
6126        let __s_b = self.gpu.stream();
6127        let mut b = __s_b.launch_builder(&f);
6128        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6129        unsafe {
6130            b.launch(cfg)?;
6131        }
6132        Ok(())
6133    }
6134
6135    pub fn moe_gate_up_silu8_dev_q8(
6136        &self,
6137        table: &CudaSlice<u64>,
6138        sel: &cudarc::driver::CudaView<i32>,
6139        aq: &CudaSlice<i8>,
6140        ad: &CudaSlice<f32>,
6141        in_f: usize,
6142        n_ff: usize,
6143        n_used: usize,
6144        n_expert: usize,
6145        qt_g: i32,
6146        qt_u: i32,
6147        rb_g: usize,
6148        rb_u: usize,
6149        macros: &CudaSlice<f32>,
6150    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6151        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6152        let (mode, wpb) = GU.get_or_init(|| {
6153            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6154            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6155                .ok()
6156                .and_then(|v| v.parse().ok())
6157                .unwrap_or(4u32)
6158                .clamp(1, 16);
6159            (mode, wpb)
6160        });
6161        let (mode, wpb) = (mode.as_str(), *wpb);
6162        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6163        let (inf, nff, ne, rbg, rbu) = (
6164            in_f as i32,
6165            n_ff as i32,
6166            n_expert as i32,
6167            rb_g as i64,
6168            rb_u as i64,
6169        );
6170        let (f, cfg) = match mode {
6171            "1" | "2" | "4" => {
6172                let rpw: u32 = mode.parse().unwrap();
6173                let f = self.func(match rpw {
6174                    1 => "moe_gate_up_silu8_dev_q8_r1",
6175                    2 => "moe_gate_up_silu8_dev_q8_r2",
6176                    _ => "moe_gate_up_silu8_dev_q8_r4",
6177                });
6178                let rows_per_block = (rpw * wpb) as usize;
6179                let gx = n_ff.div_ceil(rows_per_block) as u32;
6180                (
6181                    f,
6182                    LaunchConfig {
6183                        grid_dim: (gx, n_used as u32, 1),
6184                        block_dim: (32, wpb, 1),
6185                        shared_mem_bytes: 0,
6186                    },
6187                )
6188            }
6189            "j8" if n_used <= 32 => (
6190                self.func("moe_gate_up_silu8_dev_q8_j8"),
6191                LaunchConfig {
6192                    grid_dim: (n_ff as u32, 1, 1),
6193                    block_dim: (32, n_used as u32, 1),
6194                    shared_mem_bytes: 0,
6195                },
6196            ),
6197            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
6198            "vsm2" => {
6199                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6200                let sh = (rb_g + rb_u) as u32;
6201                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6202                f.set_attribute(
6203                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6204                    sh as i32,
6205                )?;
6206                (
6207                    f,
6208                    LaunchConfig {
6209                        grid_dim: (n_ff as u32, n_used as u32, 1),
6210                        block_dim: (32, 1, 1),
6211                        shared_mem_bytes: sh,
6212                    },
6213                )
6214            }
6215            "vsm" => {
6216                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6217                let sh = (rb_g + rb_u) as u32;
6218                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6219                f.set_attribute(
6220                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6221                    sh as i32,
6222                )?;
6223                (
6224                    f,
6225                    LaunchConfig {
6226                        grid_dim: (n_ff as u32, n_used as u32, 1),
6227                        block_dim: (32, 1, 1),
6228                        shared_mem_bytes: sh,
6229                    },
6230                )
6231            }
6232            "sg" => (
6233                self.func("moe_gate_up_silu8_dev_q8_sg"),
6234                LaunchConfig {
6235                    grid_dim: (n_ff as u32, n_used as u32, 1),
6236                    block_dim: (32, 1, 1),
6237                    shared_mem_bytes: 0,
6238                },
6239            ),
6240            "j8sg" if n_used <= 32 => (
6241                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6242                LaunchConfig {
6243                    grid_dim: (n_ff as u32, 1, 1),
6244                    block_dim: (32, n_used as u32, 1),
6245                    shared_mem_bytes: 0,
6246                },
6247            ),
6248            "u64" if in_f == 2048 => (
6249                self.func("moe_gate_up_silu8_dev_q8_u64"),
6250                LaunchConfig {
6251                    grid_dim: (n_ff as u32, n_used as u32, 1),
6252                    block_dim: (32, 1, 1),
6253                    shared_mem_bytes: 0,
6254                },
6255            ),
6256            "gs4" if in_f == 2048 => (
6257                self.func("moe_gate_up_silu8_dev_q8_gs4"),
6258                LaunchConfig {
6259                    grid_dim: (n_ff as u32, n_used as u32, 1),
6260                    block_dim: (32, 4, 1),
6261                    shared_mem_bytes: 0,
6262                },
6263            ),
6264            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
6265            "v" | "" => (
6266                self.func("moe_gate_up_silu8_dev_q8_v"),
6267                LaunchConfig {
6268                    grid_dim: (n_ff as u32, n_used as u32, 1),
6269                    block_dim: (32, 1, 1),
6270                    shared_mem_bytes: 0,
6271                },
6272            ),
6273            "s2" => (
6274                self.func("moe_gate_up_silu8_dev_q8_s2"),
6275                LaunchConfig {
6276                    grid_dim: (n_ff as u32, n_used as u32, 1),
6277                    block_dim: (32, 2, 1),
6278                    shared_mem_bytes: 0,
6279                },
6280            ),
6281            "s2z" => {
6282                let rz = wpb.min(16); // s2z smem tile is [16][2]
6283                (
6284                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
6285                    LaunchConfig {
6286                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6287                        block_dim: (32, 2, rz),
6288                        shared_mem_bytes: 0,
6289                    },
6290                )
6291            }
6292            _ => (
6293                self.func("moe_gate_up_silu8_dev_q8"),
6294                LaunchConfig {
6295                    grid_dim: (n_ff as u32, n_used as u32, 1),
6296                    block_dim: (32, 1, 1),
6297                    shared_mem_bytes: 0,
6298                },
6299            ),
6300        };
6301        let __s_b = self.gpu.stream();
6302        let mut b = __s_b.launch_builder(&f);
6303        b.arg(table)
6304            .arg(sel)
6305            .arg(aq)
6306            .arg(ad)
6307            .arg(&mut act)
6308            .arg(&inf)
6309            .arg(&nff)
6310            .arg(&ne)
6311            .arg(&qt_g)
6312            .arg(&qt_u)
6313            .arg(&rbg)
6314            .arg(&rbu)
6315            .arg(macros);
6316        unsafe {
6317            b.launch(cfg)?;
6318        }
6319        Ok(act)
6320    }
6321
6322    #[allow(clippy::too_many_arguments)]
6323    pub fn moe_down8_fma_dev_q8(
6324        &self,
6325        table: &CudaSlice<u64>,
6326        sel: &cudarc::driver::CudaView<i32>,
6327        w: &cudarc::driver::CudaView<f32>,
6328        aq2: &CudaSlice<i8>,
6329        ad2: &CudaSlice<f32>,
6330        dst: &mut cudarc::driver::CudaViewMut<f32>,
6331        in_f: usize,
6332        out_f: usize,
6333        n_used: usize,
6334        n_expert: usize,
6335        qt: i32,
6336        rb: usize,
6337    ) -> Result<(), Box<dyn std::error::Error>> {
6338        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6339        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6340        let (inf, outf, nu, ne, rbi) = (
6341            in_f as i32,
6342            out_f as i32,
6343            n_used as i32,
6344            n_expert as i32,
6345            rb as i64,
6346        );
6347        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
6348        // the h2 twins are nsb==16 (in_f==512) shape-gated.
6349        let (f, cfg) = match mode.as_str() {
6350            m @ ("1" | "2" | "4") if n_used <= 8 => {
6351                let rpw: usize = m.parse().unwrap();
6352                let f = self.func(match rpw {
6353                    1 => "moe_down8_fma_dev_q8_w8r1",
6354                    2 => "moe_down8_fma_dev_q8_w8r2",
6355                    _ => "moe_down8_fma_dev_q8_w8r4",
6356                });
6357                (
6358                    f,
6359                    LaunchConfig {
6360                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6361                        block_dim: (32, n_used as u32, 1),
6362                        shared_mem_bytes: 0,
6363                    },
6364                )
6365            }
6366            "h2" if in_f == 512 => (
6367                self.func("moe_down8_fma_dev_q8_h2"),
6368                LaunchConfig {
6369                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6370                    block_dim: (32, 1, 1),
6371                    shared_mem_bytes: 0,
6372                },
6373            ),
6374            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6375            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6376            "" if in_f == 704 && n_used <= 8 => (
6377                self.func("moe_down8_fma_dev_q8_w8r2"),
6378                LaunchConfig {
6379                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6380                    block_dim: (32, n_used as u32, 1),
6381                    shared_mem_bytes: 0,
6382                },
6383            ),
6384            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6385            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6386            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6387            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6388                self.func("moe_down8_fma_dev_q8_w8h2v"),
6389                LaunchConfig {
6390                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6391                    block_dim: (32, n_used as u32, 1),
6392                    shared_mem_bytes: 0,
6393                },
6394            ),
6395            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6396                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6397                LaunchConfig {
6398                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6399                    block_dim: (32, n_used as u32, 1),
6400                    shared_mem_bytes: 0,
6401                },
6402            ),
6403            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6404                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6405                LaunchConfig {
6406                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6407                    block_dim: (32, n_used as u32, 1),
6408                    shared_mem_bytes: 0,
6409                },
6410            ),
6411            "w8h2" if in_f == 512 && n_used <= 8 => (
6412                self.func("moe_down8_fma_dev_q8_w8h2"),
6413                LaunchConfig {
6414                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6415                    block_dim: (32, n_used as u32, 1),
6416                    shared_mem_bytes: 0,
6417                },
6418            ),
6419            _ => (
6420                self.func("moe_down8_fma_dev_q8"),
6421                LaunchConfig {
6422                    grid_dim: (out_f as u32, 1, 1),
6423                    block_dim: (32, 1, 1),
6424                    shared_mem_bytes: 0,
6425                },
6426            ),
6427        };
6428        let __s_b = self.gpu.stream();
6429        let mut b = __s_b.launch_builder(&f);
6430        b.arg(table)
6431            .arg(sel)
6432            .arg(w)
6433            .arg(aq2)
6434            .arg(ad2)
6435            .arg(dst)
6436            .arg(&inf)
6437            .arg(&outf)
6438            .arg(&nu)
6439            .arg(&ne)
6440            .arg(&qt)
6441            .arg(&rbi);
6442        unsafe {
6443            b.launch(cfg)?;
6444        }
6445        Ok(())
6446    }
6447
6448    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6449    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6450    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6451    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6452    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6453    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6454    #[allow(clippy::too_many_arguments)]
6455    pub fn moe_gate_up_silu8_dev_q8_rows(
6456        &self,
6457        table: &CudaSlice<u64>,
6458        sel: &CudaSlice<i32>,
6459        aq: &CudaSlice<i8>,
6460        ad: &CudaSlice<f32>,
6461        t: usize,
6462        in_f: usize,
6463        n_ff: usize,
6464        n_used: usize,
6465        n_expert: usize,
6466        qt_g: i32,
6467        qt_u: i32,
6468        rb_g: usize,
6469        rb_u: usize,
6470        macros: &CudaSlice<f32>,
6471    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6472        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6473        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6474        let cfg = LaunchConfig {
6475            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6476            block_dim: (32, 1, 1),
6477            shared_mem_bytes: 0,
6478        };
6479        let (inf, nff, ne, nu, rbg, rbu) = (
6480            in_f as i32,
6481            n_ff as i32,
6482            n_expert as i32,
6483            n_used as i32,
6484            rb_g as i64,
6485            rb_u as i64,
6486        );
6487        let __s_b = self.gpu.stream();
6488        let mut b = __s_b.launch_builder(&f);
6489        b.arg(table)
6490            .arg(sel)
6491            .arg(aq)
6492            .arg(ad)
6493            .arg(&mut act)
6494            .arg(&inf)
6495            .arg(&nff)
6496            .arg(&ne)
6497            .arg(&qt_g)
6498            .arg(&qt_u)
6499            .arg(&rbg)
6500            .arg(&rbu)
6501            .arg(&nu)
6502            .arg(macros);
6503        unsafe {
6504            b.launch(cfg)?;
6505        }
6506        Ok(act)
6507    }
6508
6509    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6510    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6511    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6512    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6513    #[allow(clippy::too_many_arguments)]
6514    pub fn moe_down8_fma_dev_q8_rows(
6515        &self,
6516        table: &CudaSlice<u64>,
6517        sel: &CudaSlice<i32>,
6518        w: &CudaSlice<f32>,
6519        aq2: &CudaSlice<i8>,
6520        ad2: &CudaSlice<f32>,
6521        dst: &mut CudaSlice<f32>,
6522        t: usize,
6523        in_f: usize,
6524        out_f: usize,
6525        n_used: usize,
6526        n_expert: usize,
6527        qt: i32,
6528        rb: usize,
6529    ) -> Result<(), Box<dyn std::error::Error>> {
6530        assert!(
6531            in_f == 512 && n_used <= 8,
6532            "down rows twin is w8h2v shape-gated"
6533        );
6534        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6535        let cfg = LaunchConfig {
6536            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6537            block_dim: (32, n_used as u32, 1),
6538            shared_mem_bytes: 0,
6539        };
6540        let (inf, outf, nu, ne, rbi) = (
6541            in_f as i32,
6542            out_f as i32,
6543            n_used as i32,
6544            n_expert as i32,
6545            rb as i64,
6546        );
6547        let __s_b = self.gpu.stream();
6548        let mut b = __s_b.launch_builder(&f);
6549        b.arg(table)
6550            .arg(sel)
6551            .arg(w)
6552            .arg(aq2)
6553            .arg(ad2)
6554            .arg(dst)
6555            .arg(&inf)
6556            .arg(&outf)
6557            .arg(&nu)
6558            .arg(&ne)
6559            .arg(&qt)
6560            .arg(&rbi);
6561        unsafe {
6562            b.launch(cfg)?;
6563        }
6564        Ok(())
6565    }
6566
6567    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6568    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6569    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6570    #[allow(clippy::too_many_arguments)]
6571    pub fn moe_gate_up_silu8_dev_q8_csr(
6572        &self,
6573        table: &CudaSlice<u64>,
6574        sel: &CudaSlice<i32>,
6575        aq: &CudaSlice<i8>,
6576        ad: &CudaSlice<f32>,
6577        n_pairs: usize,
6578        in_f: usize,
6579        n_ff: usize,
6580        n_used: usize,
6581        n_expert: usize,
6582        qt_g: i32,
6583        qt_u: i32,
6584        rb_g: usize,
6585        rb_u: usize,
6586    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6587        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6588        // host gate guarantees qt_g == qt_u within a supported class.
6589        let f = if qt_g == crate::QT_NVFP4 {
6590            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6591        } else {
6592            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6593        };
6594        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6595        let cfg = LaunchConfig {
6596            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6597            block_dim: (32, 1, 1),
6598            shared_mem_bytes: 0,
6599        };
6600        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6601            in_f as i32,
6602            n_ff as i32,
6603            n_expert as i32,
6604            n_used as i32,
6605            n_pairs as i32,
6606            rb_g as i64,
6607            rb_u as i64,
6608        );
6609        let __s_b = self.gpu.stream();
6610        let mut b = __s_b.launch_builder(&f);
6611        b.arg(table)
6612            .arg(sel)
6613            .arg(aq)
6614            .arg(ad)
6615            .arg(&mut act)
6616            .arg(&inf)
6617            .arg(&nff)
6618            .arg(&ne)
6619            .arg(&qt_g)
6620            .arg(&qt_u)
6621            .arg(&rbg)
6622            .arg(&rbu)
6623            .arg(&nu)
6624            .arg(&npi);
6625        unsafe {
6626            b.launch(cfg)?;
6627        }
6628        Ok(act)
6629    }
6630
6631    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6632    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6633    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6634    #[allow(clippy::too_many_arguments)]
6635    pub fn moe_down8_fma_dev_q8_variant(
6636        &self,
6637        variant: &str,
6638        table: &CudaSlice<u64>,
6639        sel: &cudarc::driver::CudaView<i32>,
6640        w: &cudarc::driver::CudaView<f32>,
6641        aq2: &CudaSlice<i8>,
6642        ad2: &CudaSlice<f32>,
6643        dst: &mut cudarc::driver::CudaViewMut<f32>,
6644        in_f: usize,
6645        out_f: usize,
6646        n_used: usize,
6647        n_expert: usize,
6648        qt: i32,
6649        rb: usize,
6650    ) -> Result<(), Box<dyn std::error::Error>> {
6651        let (inf, outf, nu, ne, rbi) = (
6652            in_f as i32,
6653            out_f as i32,
6654            n_used as i32,
6655            n_expert as i32,
6656            rb as i64,
6657        );
6658        let (f, cfg) = match variant {
6659            "w8h2" | "w8h2v" => (
6660                self.func(if variant == "w8h2" {
6661                    "moe_down8_fma_dev_q8_w8h2"
6662                } else {
6663                    "moe_down8_fma_dev_q8_w8h2v"
6664                }),
6665                LaunchConfig {
6666                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6667                    block_dim: (32, n_used as u32, 1),
6668                    shared_mem_bytes: 0,
6669                },
6670            ),
6671            "w8h2r2" | "w8h2r2v" => (
6672                self.func(if variant == "w8h2r2" {
6673                    "moe_down8_fma_dev_q8_w8h2r2"
6674                } else {
6675                    "moe_down8_fma_dev_q8_w8h2r2v"
6676                }),
6677                LaunchConfig {
6678                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6679                    block_dim: (32, n_used as u32, 1),
6680                    shared_mem_bytes: 0,
6681                },
6682            ),
6683            _ => (
6684                self.func("moe_down8_fma_dev_q8"),
6685                LaunchConfig {
6686                    grid_dim: (out_f as u32, 1, 1),
6687                    block_dim: (32, 1, 1),
6688                    shared_mem_bytes: 0,
6689                },
6690            ),
6691        };
6692        let __s_b = self.gpu.stream();
6693        let mut b = __s_b.launch_builder(&f);
6694        b.arg(table)
6695            .arg(sel)
6696            .arg(w)
6697            .arg(aq2)
6698            .arg(ad2)
6699            .arg(dst)
6700            .arg(&inf)
6701            .arg(&outf)
6702            .arg(&nu)
6703            .arg(&ne)
6704            .arg(&qt)
6705            .arg(&rbi);
6706        unsafe {
6707            b.launch(cfg)?;
6708        }
6709        Ok(())
6710    }
6711
6712    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6713    #[allow(clippy::too_many_arguments)]
6714    pub fn moe_gate_up_silu8_dev_q8_variant(
6715        &self,
6716        variant: &str,
6717        table: &CudaSlice<u64>,
6718        sel: &cudarc::driver::CudaView<i32>,
6719        aq: &CudaSlice<i8>,
6720        ad: &CudaSlice<f32>,
6721        in_f: usize,
6722        n_ff: usize,
6723        n_used: usize,
6724        n_expert: usize,
6725        qt_g: i32,
6726        qt_u: i32,
6727        rb_g: usize,
6728        rb_u: usize,
6729    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6730        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6731        let (inf, nff, ne, rbg, rbu) = (
6732            in_f as i32,
6733            n_ff as i32,
6734            n_expert as i32,
6735            rb_g as i64,
6736            rb_u as i64,
6737        );
6738        let f = self.func(if variant == "v" {
6739            "moe_gate_up_silu8_dev_q8_v"
6740        } else {
6741            "moe_gate_up_silu8_dev_q8"
6742        });
6743        let cfg = LaunchConfig {
6744            grid_dim: (n_ff as u32, n_used as u32, 1),
6745            block_dim: (32, 1, 1),
6746            shared_mem_bytes: 0,
6747        };
6748        let __s_b = self.gpu.stream();
6749        let mut b = __s_b.launch_builder(&f);
6750        b.arg(table)
6751            .arg(sel)
6752            .arg(aq)
6753            .arg(ad)
6754            .arg(&mut act)
6755            .arg(&inf)
6756            .arg(&nff)
6757            .arg(&ne)
6758            .arg(&qt_g)
6759            .arg(&qt_u)
6760            .arg(&rbg)
6761            .arg(&rbu);
6762        unsafe {
6763            b.launch(cfg)?;
6764        }
6765        Ok(act)
6766    }
6767
6768    pub fn moe_gate_up_silu8_dev(
6769        &self,
6770        table: &CudaSlice<u64>,
6771        sel: &cudarc::driver::CudaView<i32>,
6772        x: &cudarc::driver::CudaView<f32>,
6773        in_f: usize,
6774        n_ff: usize,
6775        n_used: usize,
6776        n_expert: usize,
6777        qt_g: i32,
6778        qt_u: i32,
6779        rb_g: usize,
6780        rb_u: usize,
6781        macros: &CudaSlice<f32>,
6782    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6783        let f = self.func("moe_gate_up_silu8_dev");
6784        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6785        let cfg = LaunchConfig {
6786            grid_dim: (n_ff as u32, n_used as u32, 1),
6787            block_dim: (256, 1, 1),
6788            shared_mem_bytes: 0,
6789        };
6790        let (inf, nff, ne, rbg, rbu) = (
6791            in_f as i32,
6792            n_ff as i32,
6793            n_expert as i32,
6794            rb_g as i64,
6795            rb_u as i64,
6796        );
6797        let __s_b = self.gpu.stream();
6798        let mut b = __s_b.launch_builder(&f);
6799        b.arg(table)
6800            .arg(sel)
6801            .arg(x)
6802            .arg(&mut act)
6803            .arg(&inf)
6804            .arg(&nff)
6805            .arg(&ne)
6806            .arg(&qt_g)
6807            .arg(&qt_u)
6808            .arg(&rbg)
6809            .arg(&rbu)
6810            .arg(macros);
6811        unsafe {
6812            b.launch(cfg)?;
6813        }
6814        Ok(act)
6815    }
6816
6817    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6818    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6819    #[allow(clippy::too_many_arguments)]
6820    pub fn moe_down8_fma_dev(
6821        &self,
6822        table: &CudaSlice<u64>,
6823        sel: &cudarc::driver::CudaView<i32>,
6824        w: &cudarc::driver::CudaView<f32>,
6825        act: &CudaSlice<f32>,
6826        dst: &mut cudarc::driver::CudaViewMut<f32>,
6827        in_f: usize,
6828        out_f: usize,
6829        n_used: usize,
6830        n_expert: usize,
6831        qt: i32,
6832        rb: usize,
6833    ) -> Result<(), Box<dyn std::error::Error>> {
6834        let f = self.func("moe_down8_fma_dev");
6835        let cfg = LaunchConfig {
6836            grid_dim: (out_f as u32, 1, 1),
6837            block_dim: (256, 1, 1),
6838            shared_mem_bytes: 0,
6839        };
6840        let (inf, outf, nu, ne, rbv) = (
6841            in_f as i32,
6842            out_f as i32,
6843            n_used as i32,
6844            n_expert as i32,
6845            rb as i64,
6846        );
6847        let __s_b = self.gpu.stream();
6848        let mut b = __s_b.launch_builder(&f);
6849        b.arg(table)
6850            .arg(sel)
6851            .arg(w)
6852            .arg(act)
6853            .arg(dst)
6854            .arg(&inf)
6855            .arg(&outf)
6856            .arg(&nu)
6857            .arg(&ne)
6858            .arg(&qt)
6859            .arg(&rbv);
6860        unsafe {
6861            b.launch(cfg)?;
6862        }
6863        Ok(())
6864    }
6865
6866    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6867    pub fn axpy_into(
6868        &self,
6869        src: &CudaSlice<f32>,
6870        alpha: f32,
6871        dst: &mut cudarc::driver::CudaViewMut<f32>,
6872        n: usize,
6873    ) -> Result<(), Box<dyn std::error::Error>> {
6874        let f = self.func("axpy_f32");
6875        let cfg = LaunchConfig::for_num_elems(n as u32);
6876        let (a, ni) = (alpha, n as i32);
6877        let __s_b = self.gpu.stream();
6878        let mut b = __s_b.launch_builder(&f);
6879        b.arg(src).arg(dst).arg(&a).arg(&ni);
6880        unsafe {
6881            b.launch(cfg)?;
6882        }
6883        Ok(())
6884    }
6885
6886    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
6887    pub fn axpy_host_into(
6888        &self,
6889        src: &cudarc::driver::CudaView<'_, f32>,
6890        alpha: f32,
6891        dst: &mut cudarc::driver::CudaViewMut<f32>,
6892        n: usize,
6893    ) -> Result<(), Box<dyn std::error::Error>> {
6894        let f = self.func("axpy_host_f32");
6895        let cfg = LaunchConfig::for_num_elems(n as u32);
6896        let (a, ni) = (alpha, n as i32);
6897        let __s_b = self.gpu.stream();
6898        let mut b = __s_b.launch_builder(&f);
6899        b.arg(src).arg(dst).arg(&a).arg(&ni);
6900        unsafe {
6901            b.launch(cfg)?;
6902        }
6903        Ok(())
6904    }
6905
6906    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6907    pub fn add_scaled_rows(
6908        &self,
6909        src: &CudaSlice<f32>,
6910        scale: &CudaSlice<f32>,
6911        dst: &mut CudaSlice<f32>,
6912        ncols: usize,
6913        nrows: usize,
6914    ) -> Result<(), Box<dyn std::error::Error>> {
6915        let f = self.func("add_scaled_rows_f32");
6916        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6917        let (nc, nr) = (ncols as i32, nrows as i32);
6918        let __s_b = self.gpu.stream();
6919        let mut b = __s_b.launch_builder(&f);
6920        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6921        unsafe {
6922            b.launch(cfg)?;
6923        }
6924        Ok(())
6925    }
6926
6927    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6928
6929    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6930    pub fn gather_rows(
6931        &self,
6932        src: &CudaSlice<f32>,
6933        idx: &CudaSlice<i32>,
6934        dst: &mut CudaSlice<f32>,
6935        ncols: usize,
6936        m_e: usize,
6937    ) -> Result<(), Box<dyn std::error::Error>> {
6938        let f = self.func("gather_rows_f32");
6939        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6940        let (nc, me) = (ncols as i32, m_e as i32);
6941        let __s_b = self.gpu.stream();
6942        let mut b = __s_b.launch_builder(&f);
6943        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6944        unsafe {
6945            b.launch(cfg)?;
6946        }
6947        Ok(())
6948    }
6949
6950    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6951    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6952    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6953    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6954    pub fn scatter_slot(
6955        &self,
6956        src: &CudaSlice<f32>,
6957        tok_idx: &CudaSlice<i32>,
6958        slot_idx: &CudaSlice<i32>,
6959        weight: &CudaSlice<f32>,
6960        dst: &mut CudaSlice<f32>,
6961        wbuf: &mut CudaSlice<f32>,
6962        ncols: usize,
6963        n_used: usize,
6964        m_e: usize,
6965    ) -> Result<(), Box<dyn std::error::Error>> {
6966        let f = self.func("scatter_add_slot_f32");
6967        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6968        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6969        let __s_b = self.gpu.stream();
6970        let mut b = __s_b.launch_builder(&f);
6971        b.arg(src)
6972            .arg(tok_idx)
6973            .arg(slot_idx)
6974            .arg(weight)
6975            .arg(dst)
6976            .arg(wbuf)
6977            .arg(&nc)
6978            .arg(&nu)
6979            .arg(&me);
6980        unsafe {
6981            b.launch(cfg)?;
6982        }
6983        Ok(())
6984    }
6985
6986    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6987    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6988    /// Uses FMA for bit-identity with the sequential axpy path.
6989    pub fn reduce_slots(
6990        &self,
6991        slots: &CudaSlice<f32>,
6992        wbuf: &CudaSlice<f32>,
6993        dst: &mut CudaSlice<f32>,
6994        ncols: usize,
6995        n_used: usize,
6996        t: usize,
6997    ) -> Result<(), Box<dyn std::error::Error>> {
6998        let f = self.func("reduce_slots_f32");
6999        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7000        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7001        let __s_b = self.gpu.stream();
7002        let mut b = __s_b.launch_builder(&f);
7003        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7004        unsafe {
7005            b.launch(cfg)?;
7006        }
7007        Ok(())
7008    }
7009
7010    /// Canonical slot-order reduction with separately rounded multiply and add.
7011    ///
7012    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
7013    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
7014    pub fn reduce_slots_host(
7015        &self,
7016        slots: &CudaSlice<f32>,
7017        wbuf: &CudaSlice<f32>,
7018        dst: &mut CudaSlice<f32>,
7019        ncols: usize,
7020        n_used: usize,
7021        t: usize,
7022    ) -> Result<(), Box<dyn std::error::Error>> {
7023        let f = self.func("reduce_slots_host_f32");
7024        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7025        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7026        let __s_b = self.gpu.stream();
7027        let mut b = __s_b.launch_builder(&f);
7028        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7029        unsafe {
7030            b.launch(cfg)?;
7031        }
7032        Ok(())
7033    }
7034
7035    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
7036    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
7037    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
7038    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
7039    /// GPU time, ~half of it redundant re-quantization of the same row.
7040    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
7041    pub fn quantize_q8_1_view(
7042        &self,
7043        x: &cudarc::driver::CudaView<f32>,
7044        m: usize,
7045        in_f: usize,
7046    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7047        let f = self.func("quantize_q8_1");
7048        let nblk = in_f / 32;
7049        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7050        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7051        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7052        let (inf, mi) = (in_f as i32, m as i32);
7053        let __s_b = self.gpu.stream();
7054        let mut b = __s_b.launch_builder(&f);
7055        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7056        unsafe {
7057            b.launch(cfg)?;
7058        }
7059        Ok((q, d))
7060    }
7061
7062    pub fn quantize_q8_1(
7063        &self,
7064        x: &CudaSlice<f32>,
7065        m: usize,
7066        in_f: usize,
7067    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7068        let nblk = in_f / 32;
7069        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
7070        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
7071        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
7072        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7073        let (inf, mi) = (in_f as i32, m as i32);
7074        if Self::pdl_on() && Self::pdl_wb_on() {
7075            {
7076                use cudarc::driver::{DevicePtr, DevicePtrMut};
7077                let s = &self.gpu.stream();
7078                let (px, _g0) = x.device_ptr(s);
7079                let (pq, _g1) = q.device_ptr_mut(s);
7080                let (pd, _g2) = d.device_ptr_mut(s);
7081                let mut ps = [
7082                    &px as *const _ as *mut std::ffi::c_void,
7083                    &pq as *const _ as *mut _,
7084                    &pd as *const _ as *mut _,
7085                    &inf as *const _ as *mut _,
7086                    &mi as *const _ as *mut _,
7087                ];
7088                unsafe {
7089                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7090                }
7091            }
7092            return Ok((q, d));
7093        }
7094        let f = self.func("quantize_q8_1");
7095        let __s_b = self.gpu.stream();
7096        let mut b = __s_b.launch_builder(&f);
7097        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7098        unsafe {
7099            b.launch(cfg)?;
7100        }
7101        Ok((q, d))
7102    }
7103
7104    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
7105    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
7106    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
7107    pub fn quantize_fp4_act(
7108        &self,
7109        x: &CudaSlice<f32>,
7110        m: usize,
7111        in_f: usize,
7112    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7113        let f = self.func("quantize_fp4_act");
7114        let nb16 = in_f / 16;
7115        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
7116        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
7117        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
7118        let (inf, mi) = (in_f as i32, m as i32);
7119        let __s_b = self.gpu.stream();
7120        let mut b = __s_b.launch_builder(&f);
7121        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7122        unsafe {
7123            b.launch(cfg)?;
7124        }
7125        Ok((aq4, ad4))
7126    }
7127
7128    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
7129    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
7130    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
7131    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
7132    pub fn qmatvec_gemm_nvfp4_fp4(
7133        &self,
7134        bytes: &CudaSlice<u8>,
7135        x: &CudaSlice<f32>,
7136        m: usize,
7137        in_f: usize,
7138        out_f: usize,
7139        row_bytes: usize,
7140        scale: f32,
7141    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7142        assert!(
7143            in_f % 64 == 0,
7144            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7145        );
7146        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7147        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7148        if scale != 1.0 {
7149            self.scale_inplace(&mut y, scale, m * out_f)?;
7150        }
7151        Ok(y)
7152    }
7153
7154    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
7155    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
7156    fn fp4_gemm_launch(
7157        &self,
7158        bytes: &CudaSlice<u8>,
7159        aq4: &CudaSlice<u32>,
7160        ad4: &CudaSlice<u8>,
7161        m: usize,
7162        in_f: usize,
7163        out_f: usize,
7164        row_bytes: usize,
7165    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7166        let f = self.func("qmatvec_gemm_nvfp4_fp4");
7167        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7168        const BM: u32 = 64;
7169        const BN: u32 = 256;
7170        let cfg = LaunchConfig {
7171            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7172            block_dim: (32, 4, 1),
7173            shared_mem_bytes: 0,
7174        };
7175        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7176        let __s_b = self.gpu.stream();
7177        let mut b = __s_b.launch_builder(&f);
7178        b.arg(bytes)
7179            .arg(aq4)
7180            .arg(ad4)
7181            .arg(&mut y)
7182            .arg(&inf)
7183            .arg(&outf)
7184            .arg(&mi)
7185            .arg(&rb);
7186        unsafe {
7187            b.launch(cfg)?;
7188        }
7189        Ok(y)
7190    }
7191
7192    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
7193    pub fn qmatvec_gemm_nvfp4_fp4_raw(
7194        &self,
7195        bytes: &CudaSlice<u8>,
7196        x: &CudaSlice<f32>,
7197        m: usize,
7198        in_f: usize,
7199        out_f: usize,
7200        row_bytes: usize,
7201    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7202        assert!(
7203            in_f % 64 == 0,
7204            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7205        );
7206        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7207        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7208    }
7209
7210    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
7211    pub fn qmatvec_q8_0_fast(
7212        &self,
7213        w: &CudaSlice<u8>,
7214        x: &CudaSlice<f32>,
7215        m: usize,
7216        in_f: usize,
7217        out_f: usize,
7218        row_bytes: usize,
7219    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7220        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7221        let f = self.func("qmatvec_q8_0_dp4a");
7222        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7223        let cfg = LaunchConfig {
7224            grid_dim: (out_f as u32, m as u32, 1),
7225            block_dim: (128, 1, 1),
7226            shared_mem_bytes: 0,
7227        };
7228        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7229        let __s_b = self.gpu.stream();
7230        let mut b = __s_b.launch_builder(&f);
7231        b.arg(w)
7232            .arg(&aq)
7233            .arg(&ad)
7234            .arg(&mut y)
7235            .arg(&inf)
7236            .arg(&outf)
7237            .arg(&mi)
7238            .arg(&rb);
7239        unsafe {
7240            b.launch(cfg)?;
7241        }
7242        Ok(y)
7243    }
7244
7245    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7246    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7247    pub fn qmatvec_q4_K_fast(
7248        &self,
7249        w: &CudaSlice<u8>,
7250        x: &CudaSlice<f32>,
7251        m: usize,
7252        in_f: usize,
7253        out_f: usize,
7254        row_bytes: usize,
7255    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7256        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7257        let f = self.func("qmatvec_q4_K_dp4a");
7258        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7259        let cfg = LaunchConfig {
7260            grid_dim: (out_f as u32, m as u32, 1),
7261            block_dim: (128, 1, 1),
7262            shared_mem_bytes: 0,
7263        };
7264        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7265        let __s_b = self.gpu.stream();
7266        let mut b = __s_b.launch_builder(&f);
7267        b.arg(w)
7268            .arg(&aq)
7269            .arg(&ad)
7270            .arg(&mut y)
7271            .arg(&inf)
7272            .arg(&outf)
7273            .arg(&mi)
7274            .arg(&rb);
7275        unsafe {
7276            b.launch(cfg)?;
7277        }
7278        Ok(y)
7279    }
7280
7281    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7282    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7283    pub fn qmatvec_q6_K_fast(
7284        &self,
7285        w: &CudaSlice<u8>,
7286        x: &CudaSlice<f32>,
7287        m: usize,
7288        in_f: usize,
7289        out_f: usize,
7290        row_bytes: usize,
7291    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7292        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7293        let f = self.func("qmatvec_q6_K_dp4a");
7294        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7295        let cfg = LaunchConfig {
7296            grid_dim: (out_f as u32, m as u32, 1),
7297            block_dim: (128, 1, 1),
7298            shared_mem_bytes: 0,
7299        };
7300        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7301        let __s_b = self.gpu.stream();
7302        let mut b = __s_b.launch_builder(&f);
7303        b.arg(w)
7304            .arg(&aq)
7305            .arg(&ad)
7306            .arg(&mut y)
7307            .arg(&inf)
7308            .arg(&outf)
7309            .arg(&mi)
7310            .arg(&rb);
7311        unsafe {
7312            b.launch(cfg)?;
7313        }
7314        Ok(y)
7315    }
7316
7317    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7318    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7319    pub fn qmatvec_q5_K_fast(
7320        &self,
7321        w: &CudaSlice<u8>,
7322        x: &CudaSlice<f32>,
7323        m: usize,
7324        in_f: usize,
7325        out_f: usize,
7326        row_bytes: usize,
7327    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7328        self.qmatvec_dp4a_named(
7329            "qmatvec_q5_K_dp4a",
7330            &w.slice(0..w.len()),
7331            x,
7332            m,
7333            in_f,
7334            out_f,
7335            row_bytes,
7336        )
7337    }
7338    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7339    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7340    pub fn qmatvec_q3_K_fast(
7341        &self,
7342        w: &CudaSlice<u8>,
7343        x: &CudaSlice<f32>,
7344        m: usize,
7345        in_f: usize,
7346        out_f: usize,
7347        row_bytes: usize,
7348    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7349        self.qmatvec_dp4a_named(
7350            "qmatvec_q3_K_dp4a",
7351            &w.slice(0..w.len()),
7352            x,
7353            m,
7354            in_f,
7355            out_f,
7356            row_bytes,
7357        )
7358    }
7359    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
7360    pub fn qmatvec_nvfp4_fast_rp(
7361        &self,
7362        w: &CudaSlice<u8>,
7363        x: &CudaSlice<f32>,
7364        m: usize,
7365        in_f: usize,
7366        out_f: usize,
7367        row_bytes: usize,
7368    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7369        assert!(
7370            in_f % 64 == 0,
7371            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7372        );
7373        self.qmatvec_dp4a_named(
7374            "qmatvec_nvfp4_dp4a_rp",
7375            &w.slice(0..w.len()),
7376            x,
7377            m,
7378            in_f,
7379            out_f,
7380            row_bytes,
7381        )
7382    }
7383    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
7384    pub fn qmatvec_nvfp4_fast(
7385        &self,
7386        w: &cudarc::driver::CudaView<'_, u8>,
7387        x: &CudaSlice<f32>,
7388        m: usize,
7389        in_f: usize,
7390        out_f: usize,
7391        row_bytes: usize,
7392    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7393        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
7394        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
7395        assert!(
7396            in_f % 64 == 0,
7397            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7398        );
7399        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7400    }
7401    /// v2-layout twin of `qmatvec_nvfp4_fast` for the slot-major expert banks
7402    /// (MEMRA_NVFP4_BANK_V2) — bit-identical per row, coalesced reads.
7403    pub fn qmatvec_nvfp4_fast_v2(
7404        &self,
7405        w: &cudarc::driver::CudaView<'_, u8>,
7406        x: &CudaSlice<f32>,
7407        m: usize,
7408        in_f: usize,
7409        out_f: usize,
7410        row_bytes: usize,
7411    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7412        assert!(
7413            in_f % 64 == 0,
7414            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7415        );
7416        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7417    }
7418    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
7419    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7420    pub fn qmatvec_iq4_XS_fast(
7421        &self,
7422        w: &CudaSlice<u8>,
7423        x: &CudaSlice<f32>,
7424        m: usize,
7425        in_f: usize,
7426        out_f: usize,
7427        row_bytes: usize,
7428    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7429        self.qmatvec_dp4a_named(
7430            "qmatvec_iq4_XS_dp4a",
7431            &w.slice(0..w.len()),
7432            x,
7433            m,
7434            in_f,
7435            out_f,
7436            row_bytes,
7437        )
7438    }
7439
7440    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
7441    fn qmatvec_dp4a_named(
7442        &self,
7443        name: &str,
7444        w: &cudarc::driver::CudaView<'_, u8>,
7445        x: &CudaSlice<f32>,
7446        m: usize,
7447        in_f: usize,
7448        out_f: usize,
7449        row_bytes: usize,
7450    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7451        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7452        let f = self.func(name);
7453        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7454        let cfg = LaunchConfig {
7455            grid_dim: (out_f as u32, m as u32, 1),
7456            block_dim: (128, 1, 1),
7457            shared_mem_bytes: 0,
7458        };
7459        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7460        let __s_b = self.gpu.stream();
7461        let mut b = __s_b.launch_builder(&f);
7462        b.arg(w)
7463            .arg(&aq)
7464            .arg(&ad)
7465            .arg(&mut y)
7466            .arg(&inf)
7467            .arg(&outf)
7468            .arg(&mi)
7469            .arg(&rb);
7470        unsafe {
7471            b.launch(cfg)?;
7472        }
7473        Ok(y)
7474    }
7475
7476    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
7477    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
7478    /// its output); this entry exists so a routed-expert program can quantize one activation
7479    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
7480    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
7481    #[allow(clippy::too_many_arguments)]
7482    pub fn qmatvec_nvfp4_fast_prequant_into(
7483        &self,
7484        w: &CudaSlice<u8>,
7485        aq: &CudaSlice<i8>,
7486        ad: &CudaSlice<f32>,
7487        y: &mut CudaSlice<f32>,
7488        m: usize,
7489        in_f: usize,
7490        out_f: usize,
7491        row_bytes: usize,
7492    ) -> Result<(), Box<dyn std::error::Error>> {
7493        assert!(
7494            in_f % 64 == 0,
7495            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7496        );
7497        if y.len() < m * out_f {
7498            return Err(format!(
7499                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7500                y.len()
7501            )
7502            .into());
7503        }
7504        let f = self.func("qmatvec_nvfp4_dp4a");
7505        let cfg = LaunchConfig {
7506            grid_dim: (out_f as u32, m as u32, 1),
7507            block_dim: (128, 1, 1),
7508            shared_mem_bytes: 0,
7509        };
7510        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7511        let __s_b = self.gpu.stream();
7512        let mut b = __s_b.launch_builder(&f);
7513        b.arg(w)
7514            .arg(aq)
7515            .arg(ad)
7516            .arg(y)
7517            .arg(&inf)
7518            .arg(&outf)
7519            .arg(&mi)
7520            .arg(&rb);
7521        unsafe {
7522            b.launch(cfg)?;
7523        }
7524        Ok(())
7525    }
7526
7527    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
7528    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
7529    #[allow(clippy::too_many_arguments)]
7530    pub fn matvec_f32_qkv_into(
7531        &self,
7532        wq: &CudaSlice<f32>,
7533        wk: &CudaSlice<f32>,
7534        wv: &CudaSlice<f32>,
7535        wg: &CudaSlice<f32>,
7536        x: &CudaSlice<f32>,
7537        yq: &mut CudaSlice<f32>,
7538        yk: &mut CudaSlice<f32>,
7539        yv: &mut CudaSlice<f32>,
7540        yg: &mut CudaSlice<f32>,
7541        in_f: usize,
7542        out_q: usize,
7543        out_kv: usize,
7544        out_g: usize,
7545    ) -> Result<(), Box<dyn std::error::Error>> {
7546        if in_f % 4 != 0
7547            || wq.len() != out_q * in_f
7548            || wk.len() != out_kv * in_f
7549            || wv.len() != out_kv * in_f
7550            || wg.len() < out_g * in_f
7551            || x.len() < in_f
7552            || yq.len() < out_q
7553            || yk.len() < out_kv
7554            || yv.len() < out_kv
7555            || (out_g > 0 && yg.len() < out_g)
7556        {
7557            return Err(format!(
7558                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7559                 wq={} wk={} wv={} wg={}",
7560                wq.len(),
7561                wk.len(),
7562                wv.len(),
7563                wg.len()
7564            )
7565            .into());
7566        }
7567        let f = self.func("matvec_f32_qkv");
7568        let cfg = LaunchConfig {
7569            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7570            block_dim: (128, 1, 1),
7571            shared_mem_bytes: 0,
7572        };
7573        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7574        let __s_b = self.gpu.stream();
7575        let mut b = __s_b.launch_builder(&f);
7576        b.arg(wq)
7577            .arg(wk)
7578            .arg(wv)
7579            .arg(wg)
7580            .arg(x)
7581            .arg(yq)
7582            .arg(yk)
7583            .arg(yv)
7584            .arg(yg)
7585            .arg(&inf)
7586            .arg(&oq)
7587            .arg(&okv)
7588            .arg(&og);
7589        unsafe {
7590            b.launch(cfg)?;
7591        }
7592        Ok(())
7593    }
7594
7595    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
7596    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
7597    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
7598    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
7599    /// kernel — the batching only removes host launch latency.
7600    #[allow(clippy::too_many_arguments)]
7601    /// FUSION #2a: gate+up sweeps in one launch (v2 banks only; identical geometry both
7602    /// banks, caller-guarded). Per-row bit-identical to two qmatvec_nvfp4_sel_into calls.
7603    #[allow(clippy::too_many_arguments)]
7604    pub fn qmatvec_nvfp4_sel_gu_into(
7605        &self,
7606        gate_bank: &CudaSlice<u8>,
7607        up_bank: &CudaSlice<u8>,
7608        sel: &CudaSlice<i32>,
7609        aq: &CudaSlice<i8>,
7610        ad: &CudaSlice<f32>,
7611        yg: &mut CudaSlice<f32>,
7612        yu: &mut CudaSlice<f32>,
7613        n_sel: usize,
7614        in_f: usize,
7615        out_f: usize,
7616        row_bytes: usize,
7617        expert_stride: usize,
7618    ) -> Result<(), Box<dyn std::error::Error>> {
7619        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7620        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7621            return Err("NVFP4 gu sel geometry".into());
7622        }
7623        // MEMRA_SEL_GU_RPW=2|4: multirow twin (activation group read once, reused across
7624        // RPW rows' gate+up dots) — bit-identical per row, one block per RPW rows.
7625        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7626        let rpw = *RPW.get_or_init(|| {
7627            std::env::var("MEMRA_SEL_GU_RPW")
7628                .ok()
7629                .and_then(|v| v.parse().ok())
7630                .filter(|r| *r == 2 || *r == 4)
7631                .unwrap_or(1)
7632        });
7633        let rpw = if out_f % rpw == 0 { rpw } else { 1 };
7634        // MEMRA_SEL_GU_WPR=1: warp-per-row (NUMERIC-CLASS — per-row reduction order changes;
7635        // acceptance is the argmax gate + battery, the QKV_FUSED/BF16_MMV class).
7636        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7637        let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
7638        let f = self.func(match (wpr, rpw) {
7639            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
7640            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
7641            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
7642            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
7643        });
7644        let cfg = LaunchConfig {
7645            grid_dim: if wpr {
7646                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
7647            } else if rpw == 1 {
7648                ((2 * out_f) as u32, n_sel as u32, 1)
7649            } else {
7650                ((out_f / rpw) as u32, n_sel as u32, 1)
7651            },
7652            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
7653            shared_mem_bytes: 0,
7654        };
7655        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7656        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7657        let (ars, adrs) = (0i64, 0i64);
7658        let __s_b = self.gpu.stream();
7659        let mut b = __s_b.launch_builder(&f);
7660        b.arg(gate_bank)
7661            .arg(up_bank)
7662            .arg(sel)
7663            .arg(aq)
7664            .arg(ad)
7665            .arg(yg)
7666            .arg(yu)
7667            .arg(&inf)
7668            .arg(&outf)
7669            .arg(&ns)
7670            .arg(&rb)
7671            .arg(&es)
7672            .arg(&ars)
7673            .arg(&adrs);
7674        unsafe {
7675            b.launch(cfg)?;
7676        }
7677        Ok(())
7678    }
7679
7680    /// MEMRA_SEL_DOWN8=1: the DOWN sweep and the route-weight combine in ONE launch
7681    /// (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm ported to the
7682    /// NVFP4 banks). Block = (32, n_sel): one warp per slot instead of one warp per
7683    /// (row, slot), and the n_sel x out_f partial buffer disappears. Bit-identical to
7684    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce
7685    /// tree, same slot-ordered chain. Requires the v2 banks and nsb <= 32 (the fit-block
7686    /// class the reduce identity is argued at).
7687    #[allow(clippy::too_many_arguments)]
7688    pub fn qmatvec_nvfp4_sel_down8_into(
7689        &self,
7690        bank: &CudaSlice<u8>,
7691        sel: &CudaSlice<i32>,
7692        aq: &CudaSlice<i8>,
7693        ad: &CudaSlice<f32>,
7694        route_w: &CudaSlice<f32>,
7695        md: &CudaSlice<f32>,
7696        dst: &mut CudaSlice<f32>,
7697        n_sel: usize,
7698        in_f: usize,
7699        out_f: usize,
7700        row_bytes: usize,
7701        expert_stride: usize,
7702        act_row_stride: usize,
7703        ad_row_stride: usize,
7704    ) -> Result<(), Box<dyn std::error::Error>> {
7705        if in_f % 64 != 0
7706            || n_sel == 0
7707            || n_sel > 8
7708            || (in_f >> 5) > 32
7709            || dst.len() < out_f
7710            || sel.len() < n_sel
7711            || route_w.len() < n_sel
7712        {
7713            return Err(format!(
7714                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
7715                dst.len()
7716            )
7717            .into());
7718        }
7719        if !crate::tp::nvfp4_bank_v2_on() {
7720            return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
7721        }
7722        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
7723        let cfg = LaunchConfig {
7724            grid_dim: (out_f as u32, 1, 1),
7725            block_dim: (32, n_sel as u32, 1),
7726            shared_mem_bytes: 0,
7727        };
7728        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7729        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7730        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7731        let __s_b = self.gpu.stream();
7732        let mut b = __s_b.launch_builder(&f);
7733        b.arg(bank)
7734            .arg(sel)
7735            .arg(aq)
7736            .arg(ad)
7737            .arg(route_w)
7738            .arg(md)
7739            .arg(dst)
7740            .arg(&inf)
7741            .arg(&outf)
7742            .arg(&ns)
7743            .arg(&rb)
7744            .arg(&es)
7745            .arg(&ars)
7746            .arg(&adrs);
7747        unsafe {
7748            b.launch(cfg)?;
7749        }
7750        Ok(())
7751    }
7752
7753    pub fn qmatvec_nvfp4_sel_into(
7754        &self,
7755        bank: &CudaSlice<u8>,
7756        sel: &CudaSlice<i32>,
7757        aq: &CudaSlice<i8>,
7758        ad: &CudaSlice<f32>,
7759        y: &mut CudaSlice<f32>,
7760        n_sel: usize,
7761        in_f: usize,
7762        out_f: usize,
7763        row_bytes: usize,
7764        expert_stride: usize,
7765        act_row_stride: usize,
7766        ad_row_stride: usize,
7767    ) -> Result<(), Box<dyn std::error::Error>> {
7768        assert!(
7769            in_f % 64 == 0,
7770            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7771        );
7772        if y.len() < n_sel * out_f || sel.len() < n_sel {
7773            return Err(format!(
7774                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7775                y.len(),
7776                sel.len()
7777            )
7778            .into());
7779        }
7780        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
7781        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
7782        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
7783        // sequential-rows variant was flat). Default stays the single-row form.
7784        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
7785        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
7786        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7787        let mode = *MR.get_or_init(|| {
7788            if crate::tp::nvfp4_bank_v2_on() {
7789                3
7790            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7791                2
7792            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7793                1
7794            } else {
7795                0
7796            }
7797        });
7798        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7799        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
7800        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
7801        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
7802        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7803        let v2s = mode == 3
7804            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7805            && row_bytes % 16 == 0
7806            && in_f <= 4096;
7807        let f = match (mode, v2s) {
7808            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7809            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7810            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7811            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7812            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7813        };
7814        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
7815        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
7816        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
7817        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
7818        let nsb = in_f >> 5;
7819        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7820            32
7821        } else if mode == 1 {
7822            512
7823        } else {
7824            128
7825        };
7826        let cfg = LaunchConfig {
7827            grid_dim: (
7828                if v2s {
7829                    (out_f as u32).div_ceil(8)
7830                } else {
7831                    match mode {
7832                        2 => (out_f as u32).div_ceil(16),
7833                        1 => (out_f as u32).div_ceil(4),
7834                        _ => out_f as u32,
7835                    }
7836                },
7837                n_sel as u32,
7838                1,
7839            ),
7840            block_dim: (fit_block, 1, 1),
7841            shared_mem_bytes: 0,
7842        };
7843        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7844        let (rb, es, ars, adrs) = (
7845            row_bytes as i64,
7846            expert_stride as i64,
7847            act_row_stride as i64,
7848            ad_row_stride as i64,
7849        );
7850        let __s_b = self.gpu.stream();
7851        let mut b = __s_b.launch_builder(&f);
7852        b.arg(bank)
7853            .arg(sel)
7854            .arg(aq)
7855            .arg(ad)
7856            .arg(y)
7857            .arg(&inf)
7858            .arg(&outf)
7859            .arg(&ns)
7860            .arg(&rb)
7861            .arg(&es)
7862            .arg(&ars)
7863            .arg(&adrs);
7864        unsafe {
7865            b.launch(cfg)?;
7866        }
7867        Ok(())
7868    }
7869
7870    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
7871    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
7872    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
7873    /// takes the plain SiLU kernel.
7874    #[allow(clippy::too_many_arguments)]
7875    pub fn silu_mul_scaled_q8_1_sel_into(
7876        &self,
7877        gate: &CudaSlice<f32>,
7878        up: &CudaSlice<f32>,
7879        gmac: &CudaSlice<f32>,
7880        umac: &CudaSlice<f32>,
7881        sel: &CudaSlice<i32>,
7882        limit: Option<f32>,
7883        out_q: &mut CudaSlice<i8>,
7884        out_d: &mut CudaSlice<f32>,
7885        n_per: usize,
7886        n_sel: usize,
7887    ) -> Result<(), Box<dyn std::error::Error>> {
7888        let n = n_per * n_sel;
7889        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
7890            return Err(format!(
7891                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
7892                out_q.len(),
7893                out_d.len()
7894            )
7895            .into());
7896        }
7897        if let Some(limit) = limit {
7898            if limit <= 1e-6 {
7899                return Err(format!(
7900                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
7901                )
7902                .into());
7903            }
7904            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
7905            let cfg = LaunchConfig::for_num_elems(n as u32);
7906            let (np, ns) = (n_per as i32, n_sel as i32);
7907            let __s_b = self.gpu.stream();
7908            let mut b = __s_b.launch_builder(&f);
7909            b.arg(gate)
7910                .arg(up)
7911                .arg(gmac)
7912                .arg(umac)
7913                .arg(sel)
7914                .arg(&limit)
7915                .arg(out_q)
7916                .arg(out_d)
7917                .arg(&np)
7918                .arg(&ns);
7919            unsafe {
7920                b.launch(cfg)?;
7921            }
7922            return Ok(());
7923        }
7924        let f = self.func("silu_mul_scaled_q8_1_sel");
7925        let cfg = LaunchConfig::for_num_elems(n as u32);
7926        let (np, ns) = (n_per as i32, n_sel as i32);
7927        let __s_b = self.gpu.stream();
7928        let mut b = __s_b.launch_builder(&f);
7929        b.arg(gate)
7930            .arg(up)
7931            .arg(gmac)
7932            .arg(umac)
7933            .arg(sel)
7934            .arg(out_q)
7935            .arg(out_d)
7936            .arg(&np)
7937            .arg(&ns);
7938        unsafe {
7939            b.launch(cfg)?;
7940        }
7941        Ok(())
7942    }
7943
7944    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7945        Ok(self.gpu.stream().clone_htod(v)?)
7946    }
7947    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7948        Ok(self.gpu.stream().clone_htod(v)?)
7949    }
7950    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
7951    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7952        Ok(self.gpu.stream().clone_htod(v)?)
7953    }
7954    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7955        Ok(self.gpu.stream().clone_htod(v)?)
7956    }
7957    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
7958    pub fn dtoh_view(
7959        &self,
7960        d: &cudarc::driver::CudaView<f32>,
7961    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7962        let v = self.gpu.stream().clone_dtoh(d)?;
7963        self.gpu.stream().synchronize()?;
7964        Ok(v)
7965    }
7966    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7967        let v = self.gpu.stream().clone_dtoh(d)?;
7968        self.gpu.stream().synchronize()?;
7969        Ok(v)
7970    }
7971    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
7972    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
7973    /// issuing them together avoids a second stream synchronization in every trunk layer.
7974    pub fn dtoh_pair(
7975        &self,
7976        a: &CudaSlice<f32>,
7977        b: &CudaSlice<f32>,
7978    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7979        let av = self.gpu.stream().clone_dtoh(a)?;
7980        let bv = self.gpu.stream().clone_dtoh(b)?;
7981        self.gpu.stream().synchronize()?;
7982        Ok((av, bv))
7983    }
7984    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
7985    /// cross a shape-sensitive host boundary.
7986    pub fn dtoh_pair_views(
7987        &self,
7988        a: &cudarc::driver::CudaView<f32>,
7989        b: &cudarc::driver::CudaView<f32>,
7990    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7991        let av = self.gpu.stream().clone_dtoh(a)?;
7992        let bv = self.gpu.stream().clone_dtoh(b)?;
7993        self.gpu.stream().synchronize()?;
7994        Ok((av, bv))
7995    }
7996    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
7997    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7998        let v = self.gpu.stream().clone_dtoh(d)?;
7999        self.gpu.stream().synchronize()?;
8000        Ok(v)
8001    }
8002    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
8003    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8004        let v = self.gpu.stream().clone_dtoh(d)?;
8005        self.gpu.stream().synchronize()?;
8006        Ok(v)
8007    }
8008    pub fn dtoh_u8_view(
8009        &self,
8010        d: &cudarc::driver::CudaView<u8>,
8011    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8012        let v = self.gpu.stream().clone_dtoh(d)?;
8013        self.gpu.stream().synchronize()?;
8014        Ok(v)
8015    }
8016    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8017        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8018        self.keep_if_capturing(&s);
8019        Ok(s)
8020    }
8021
8022    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
8023    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
8024    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
8025    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
8026    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
8027    /// back (or kept resident for graph replay). Returns the device token buffer.
8028    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
8029    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
8030    pub fn prob_of_token_device(
8031        &self,
8032        logits: &CudaSlice<f32>,
8033        tok: &CudaSlice<u32>,
8034        n_vocab: usize,
8035    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8036        let nb = ARGMAX_NB;
8037        let mut part = self.alloc_uninit::<f32>(nb)?;
8038        let mut p = self.alloc_uninit::<f32>(1)?;
8039        let f1 = self.func("prob_of_token_partial_f32");
8040        let cfg1 = LaunchConfig {
8041            grid_dim: (nb as u32, 1, 1),
8042            block_dim: (256, 1, 1),
8043            shared_mem_bytes: 0,
8044        };
8045        let nv = n_vocab as i32;
8046        let __s_b1 = self.gpu.stream();
8047        let mut b1 = __s_b1.launch_builder(&f1);
8048        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8049        unsafe {
8050            b1.launch(cfg1)?;
8051        }
8052        let f2 = self.func("prob_of_token_final_f32");
8053        let cfg2 = LaunchConfig {
8054            grid_dim: (1, 1, 1),
8055            block_dim: (256, 1, 1),
8056            shared_mem_bytes: 0,
8057        };
8058        let nbi = nb as i32;
8059        let __s_b2 = self.gpu.stream();
8060        let mut b2 = __s_b2.launch_builder(&f2);
8061        b2.arg(&part).arg(&mut p).arg(&nbi);
8062        unsafe {
8063            b2.launch(cfg2)?;
8064        }
8065        Ok(p)
8066    }
8067
8068    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
8069    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
8070    /// where the host reads the p-min confidence between replays. Same kernels, same math.
8071    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
8072    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
8073    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
8074    pub fn prob_of_token_device_col(
8075        &self,
8076        logits: &CudaSlice<f32>,
8077        tok_all: &CudaSlice<u32>,
8078        tok_idx: usize,
8079        p_out: &mut CudaSlice<f32>,
8080        p_idx: usize,
8081        n_vocab: usize,
8082    ) -> Result<(), Box<dyn std::error::Error>> {
8083        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8084        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8085        let nb = ARGMAX_NB;
8086        let mut part = self.alloc_uninit::<f32>(nb)?;
8087        let f1 = self.func("prob_of_token_partial_f32");
8088        let cfg1 = LaunchConfig {
8089            grid_dim: (nb as u32, 1, 1),
8090            block_dim: (256, 1, 1),
8091            shared_mem_bytes: 0,
8092        };
8093        let nv = n_vocab as i32;
8094        let __s_b1 = self.gpu.stream();
8095        let mut b1 = __s_b1.launch_builder(&f1);
8096        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8097        unsafe {
8098            b1.launch(cfg1)?;
8099        }
8100        let f2 = self.func("prob_of_token_final_f32");
8101        let cfg2 = LaunchConfig {
8102            grid_dim: (1, 1, 1),
8103            block_dim: (256, 1, 1),
8104            shared_mem_bytes: 0,
8105        };
8106        let nbi = nb as i32;
8107        let __s_b2 = self.gpu.stream();
8108        let mut b2 = __s_b2.launch_builder(&f2);
8109        b2.arg(&part).arg(&mut p_v).arg(&nbi);
8110        unsafe {
8111            b2.launch(cfg2)?;
8112        }
8113        Ok(())
8114    }
8115
8116    pub fn prob_of_token_device_into(
8117        &self,
8118        logits: &CudaSlice<f32>,
8119        tok: &CudaSlice<u32>,
8120        p_out: &mut CudaSlice<f32>,
8121        n_vocab: usize,
8122    ) -> Result<(), Box<dyn std::error::Error>> {
8123        let nb = ARGMAX_NB;
8124        let mut part = self.alloc_uninit::<f32>(nb)?;
8125        let f1 = self.func("prob_of_token_partial_f32");
8126        let cfg1 = LaunchConfig {
8127            grid_dim: (nb as u32, 1, 1),
8128            block_dim: (256, 1, 1),
8129            shared_mem_bytes: 0,
8130        };
8131        let nv = n_vocab as i32;
8132        let __s_b1 = self.gpu.stream();
8133        let mut b1 = __s_b1.launch_builder(&f1);
8134        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8135        unsafe {
8136            b1.launch(cfg1)?;
8137        }
8138        let f2 = self.func("prob_of_token_final_f32");
8139        let cfg2 = LaunchConfig {
8140            grid_dim: (1, 1, 1),
8141            block_dim: (256, 1, 1),
8142            shared_mem_bytes: 0,
8143        };
8144        let nbi = nb as i32;
8145        let __s_b2 = self.gpu.stream();
8146        let mut b2 = __s_b2.launch_builder(&f2);
8147        b2.arg(&part).arg(p_out).arg(&nbi);
8148        unsafe {
8149            b2.launch(cfg2)?;
8150        }
8151        Ok(())
8152    }
8153
8154    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
8155    /// (graph-constant params, device-varying index). Capture-safe.
8156    pub fn u32_hist_append(
8157        &self,
8158        tok: &CudaSlice<u32>,
8159        hist: &mut CudaSlice<u32>,
8160        idx: &mut CudaSlice<i32>,
8161    ) -> Result<(), Box<dyn std::error::Error>> {
8162        let f = self.func("u32_hist_append");
8163        let cfg = LaunchConfig {
8164            grid_dim: (1, 1, 1),
8165            block_dim: (32, 1, 1),
8166            shared_mem_bytes: 0,
8167        };
8168        let __s_b = self.gpu.stream();
8169        let mut b = __s_b.launch_builder(&f);
8170        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8171        unsafe {
8172            b.launch(cfg)?;
8173        }
8174        Ok(())
8175    }
8176
8177    pub fn argmax_token_device(
8178        &self,
8179        logits: &CudaSlice<f32>,
8180        n_vocab: usize,
8181    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8182        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8183        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8184        Ok(tok)
8185    }
8186    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
8187    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
8188    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
8189    /// pointer is baked once and the token id never round-trips to host inside steady state. The
8190    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
8191    /// captured passes bake fixed addresses.
8192    pub fn argmax_token_device_into(
8193        &self,
8194        logits: &CudaSlice<f32>,
8195        tok: &mut CudaSlice<u32>,
8196        n_vocab: usize,
8197    ) -> Result<(), Box<dyn std::error::Error>> {
8198        let nb = ARGMAX_NB;
8199        let f1 = self.func("argmax_partial_f32");
8200        let f2 = self.func("argmax_final_f32");
8201        let mut guard = self.argmax_partials.lock().unwrap();
8202        if guard.is_none() {
8203            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8204            // buffers carry no cudarc events (illegal inside capture).
8205            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8206            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8207            *guard = Some((pv, pi));
8208        }
8209        let (part_v, part_i) = guard.as_mut().unwrap();
8210        let nv = n_vocab as i32;
8211        let nbi = nb as i32;
8212        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8213        let cfg1 = LaunchConfig {
8214            grid_dim: (nb as u32, 1, 1),
8215            block_dim: (256, 1, 1),
8216            shared_mem_bytes: 0,
8217        };
8218        let __s_b1 = self.gpu.stream();
8219        let mut b1 = __s_b1.launch_builder(&f1);
8220        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8221        unsafe {
8222            b1.launch(cfg1)?;
8223        }
8224        // pass 2: one block reduces NB partials -> token_out[0].
8225        let cfg2 = LaunchConfig {
8226            grid_dim: (1, 1, 1),
8227            block_dim: (256, 1, 1),
8228            shared_mem_bytes: 0,
8229        };
8230        let __s_b2 = self.gpu.stream();
8231        let mut b2 = __s_b2.launch_builder(&f2);
8232        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8233        unsafe {
8234            b2.launch(cfg2)?;
8235        }
8236        Ok(())
8237    }
8238    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8239    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8240    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8241    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8242    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8243    pub fn argmax_token_device_col(
8244        &self,
8245        logits: &CudaSlice<f32>,
8246        col: usize,
8247        n_vocab: usize,
8248        toks: &mut CudaSlice<u32>,
8249        out_idx: usize,
8250    ) -> Result<(), Box<dyn std::error::Error>> {
8251        let nb = ARGMAX_NB;
8252        let f1 = self.func("argmax_partial_f32");
8253        let f2 = self.func("argmax_final_f32");
8254        let mut guard = self.argmax_partials.lock().unwrap();
8255        if guard.is_none() {
8256            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8257            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8258            *guard = Some((pv, pi));
8259        }
8260        let (part_v, part_i) = guard.as_mut().unwrap();
8261        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8262        let nv = n_vocab as i32;
8263        let nbi = nb as i32;
8264        let cfg1 = LaunchConfig {
8265            grid_dim: (nb as u32, 1, 1),
8266            block_dim: (256, 1, 1),
8267            shared_mem_bytes: 0,
8268        };
8269        let __s_b1 = self.gpu.stream();
8270        let mut b1 = __s_b1.launch_builder(&f1);
8271        b1.arg(&col_view)
8272            .arg(&mut *part_v)
8273            .arg(&mut *part_i)
8274            .arg(&nv);
8275        unsafe {
8276            b1.launch(cfg1)?;
8277        }
8278        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8279        let cfg2 = LaunchConfig {
8280            grid_dim: (1, 1, 1),
8281            block_dim: (256, 1, 1),
8282            shared_mem_bytes: 0,
8283        };
8284        let __s_b2 = self.gpu.stream();
8285        let mut b2 = __s_b2.launch_builder(&f2);
8286        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8287        unsafe {
8288            b2.launch(cfg2)?;
8289        }
8290        Ok(())
8291    }
8292    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8293    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8294        Ok(self.gpu.stream().clone_htod(v)?)
8295    }
8296    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8297        let v = self.gpu.stream().clone_dtoh(d)?;
8298        self.gpu.stream().synchronize()?;
8299        Ok(v)
8300    }
8301
8302    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8303        let v = self.gpu.stream().clone_dtoh(d)?;
8304        self.gpu.stream().synchronize()?;
8305        Ok(v)
8306    }
8307    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8308    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8309    /// contents change every step, the address must not, so a captured graph can read it).
8310    pub fn htod_u32_into(
8311        &self,
8312        dst: &mut CudaSlice<u32>,
8313        src: &[u32],
8314    ) -> Result<(), Box<dyn std::error::Error>> {
8315        let mut view = dst.slice_mut(0..src.len());
8316        self.gpu.stream().memcpy_htod(src, &mut view)?;
8317        Ok(())
8318    }
8319
8320    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8321    /// table without changing the device address its reconcile kernel consumes.
8322    pub fn htod_i32_into(
8323        &self,
8324        dst: &mut CudaSlice<i32>,
8325        src: &[i32],
8326    ) -> Result<(), Box<dyn std::error::Error>> {
8327        let mut view = dst.slice_mut(0..src.len());
8328        self.gpu.stream().memcpy_htod(src, &mut view)?;
8329        Ok(())
8330    }
8331
8332    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8333        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8334        self.keep_if_capturing(&s);
8335        Ok(s)
8336    }
8337    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8338    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8339    pub fn embed_gather_device_into(
8340        &self,
8341        embd: &CudaSlice<u8>,
8342        token_d: &CudaSlice<u32>,
8343        x_out: &mut CudaSlice<f32>,
8344        n_embd: usize,
8345        qtype: i32,
8346        row_bytes: usize,
8347    ) -> Result<(), Box<dyn std::error::Error>> {
8348        let f = self.func("embed_gather_u32");
8349        let cfg = LaunchConfig {
8350            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8351            block_dim: (256, 1, 1),
8352            shared_mem_bytes: 0,
8353        };
8354        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8355        let __s_b = self.gpu.stream();
8356        let mut b = __s_b.launch_builder(&f);
8357        b.arg(embd)
8358            .arg(token_d)
8359            .arg(x_out)
8360            .arg(&ne)
8361            .arg(&qt)
8362            .arg(&rb);
8363        unsafe {
8364            b.launch(cfg)?;
8365        }
8366        Ok(())
8367    }
8368    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
8369    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8370        let v = self.gpu.stream().clone_dtoh(d)?;
8371        self.gpu.stream().synchronize()?;
8372        Ok(v[0])
8373    }
8374    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
8375    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
8376    /// the counter value after the throwaway capture warmups corrupt it.
8377    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
8378    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
8379    /// copy (fine at stream-idle boundaries, poison mid-round).
8380    pub fn i32_set_k(
8381        &self,
8382        dst: &mut CudaSlice<i32>,
8383        v: i32,
8384    ) -> Result<(), Box<dyn std::error::Error>> {
8385        let f = self.func("i32_set_k");
8386        let cfg = LaunchConfig {
8387            grid_dim: (1, 1, 1),
8388            block_dim: (1, 1, 1),
8389            shared_mem_bytes: 0,
8390        };
8391        let idx = 0i32;
8392        let __s_b = self.gpu.stream();
8393        let mut b = __s_b.launch_builder(&f);
8394        b.arg(dst).arg(&v).arg(&idx);
8395        unsafe {
8396            b.launch(cfg)?;
8397        }
8398        Ok(())
8399    }
8400
8401    pub fn set_i32_one(
8402        &self,
8403        d: &mut CudaSlice<i32>,
8404        v: i32,
8405    ) -> Result<(), Box<dyn std::error::Error>> {
8406        self.gpu.stream().memcpy_htod(&[v], d)?;
8407        Ok(())
8408    }
8409    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
8410    /// during priming / capture-state restore.
8411    pub fn set_u32_one(
8412        &self,
8413        d: &mut CudaSlice<u32>,
8414        v: u32,
8415    ) -> Result<(), Box<dyn std::error::Error>> {
8416        self.gpu.stream().memcpy_htod(&[v], d)?;
8417        Ok(())
8418    }
8419    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
8420    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8421        let v = self.gpu.stream().clone_dtoh(d)?;
8422        self.gpu.stream().synchronize()?;
8423        Ok(v[0])
8424    }
8425    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
8426    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8427        Ok(self.gpu.stream().clone_htod(bytes)?)
8428    }
8429    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
8430    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
8431    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
8432    pub fn embed_gather_device(
8433        &self,
8434        embd: &CudaSlice<u8>,
8435        token_d: &CudaSlice<u32>,
8436        n_embd: usize,
8437        qtype: i32,
8438        row_bytes: usize,
8439    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8440        let f = self.func("embed_gather_u32");
8441        let mut x = self.alloc_uninit::<f32>(n_embd)?;
8442        let cfg = LaunchConfig {
8443            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8444            block_dim: (256, 1, 1),
8445            shared_mem_bytes: 0,
8446        };
8447        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8448        let __s_b = self.gpu.stream();
8449        let mut b = __s_b.launch_builder(&f);
8450        b.arg(embd)
8451            .arg(token_d)
8452            .arg(&mut x)
8453            .arg(&ne)
8454            .arg(&qt)
8455            .arg(&rb);
8456        unsafe {
8457            b.launch(cfg)?;
8458        }
8459        Ok(x)
8460    }
8461
8462    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
8463    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
8464    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
8465    pub fn embed_gather_device_t(
8466        &self,
8467        embd: &CudaSlice<u8>,
8468        tokens: &[u32],
8469        n_embd: usize,
8470        qtype: i32,
8471        row_bytes: usize,
8472    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8473        let t = tokens.len();
8474        let tok_d = self.gpu.stream().clone_htod(tokens)?;
8475        let f = self.func("embed_gather_u32_t");
8476        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8477        let cfg = LaunchConfig {
8478            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8479            block_dim: (256, 1, 1),
8480            shared_mem_bytes: 0,
8481        };
8482        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8483        let __s_b = self.gpu.stream();
8484        let mut b = __s_b.launch_builder(&f);
8485        b.arg(embd)
8486            .arg(&tok_d)
8487            .arg(&mut x)
8488            .arg(&ne)
8489            .arg(&qt)
8490            .arg(&rb)
8491            .arg(&ti);
8492        unsafe {
8493            b.launch(cfg)?;
8494        }
8495        Ok(x)
8496    }
8497
8498    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
8499    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
8500    /// as embed_gather_device_t — bit-identical rows.
8501    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
8502    pub fn embed_gather_device_tv(
8503        &self,
8504        embd: &CudaSlice<u8>,
8505        tok_v: &cudarc::driver::CudaView<u32>,
8506        t: usize,
8507        n_embd: usize,
8508        qtype: i32,
8509        row_bytes: usize,
8510    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8511        let f = self.func("embed_gather_u32_t");
8512        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8513        let cfg = LaunchConfig {
8514            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8515            block_dim: (256, 1, 1),
8516            shared_mem_bytes: 0,
8517        };
8518        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8519        let __s_b = self.gpu.stream();
8520        let mut b = __s_b.launch_builder(&f);
8521        b.arg(embd)
8522            .arg(tok_v)
8523            .arg(&mut x)
8524            .arg(&ne)
8525            .arg(&qt)
8526            .arg(&rb)
8527            .arg(&ti);
8528        unsafe {
8529            b.launch(cfg)?;
8530        }
8531        Ok(x)
8532    }
8533
8534    pub fn embed_gather_device_td(
8535        &self,
8536        embd: &CudaSlice<u8>,
8537        tok_d: &CudaSlice<u32>,
8538        t: usize,
8539        n_embd: usize,
8540        qtype: i32,
8541        row_bytes: usize,
8542    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8543        let f = self.func("embed_gather_u32_t");
8544        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8545        let cfg = LaunchConfig {
8546            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8547            block_dim: (256, 1, 1),
8548            shared_mem_bytes: 0,
8549        };
8550        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8551        let __s_b = self.gpu.stream();
8552        let mut b = __s_b.launch_builder(&f);
8553        b.arg(embd)
8554            .arg(tok_d)
8555            .arg(&mut x)
8556            .arg(&ne)
8557            .arg(&qt)
8558            .arg(&rb)
8559            .arg(&ti);
8560        unsafe {
8561            b.launch(cfg)?;
8562        }
8563        Ok(x)
8564    }
8565
8566    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
8567    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
8568    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
8569    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
8570    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
8571    #[inline]
8572    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
8573    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8574        if self
8575            .capture_keep_on
8576            .load(std::sync::atomic::Ordering::Relaxed)
8577        {
8578            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8579        }
8580    }
8581
8582    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8583        &self,
8584        n: usize,
8585    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8586        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8587        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
8588        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
8589        // not cover engine-internal buffers). Debug-only: massive launch overhead.
8590        {
8591            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8592            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8593                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
8594                use cudarc::driver::DevicePtrMut;
8595                let n_bytes = s.len() * std::mem::size_of::<T>();
8596                let stream = self.gpu.stream();
8597                let (p_, _g) = s.device_ptr_mut(&stream);
8598                unsafe {
8599                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8600                        .result()?;
8601                }
8602            }
8603        }
8604        self.keep_if_capturing(&s);
8605        Ok(s)
8606    }
8607
8608    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
8609    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
8610    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
8611    /// consumers alloc through this (m=1 decode arms).
8612    pub fn uninit_q8_pair(
8613        &self,
8614        n: usize,
8615    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8616        Ok((
8617            self.alloc_uninit::<i8>(n)?,
8618            self.alloc_uninit::<f32>(n / 32)?,
8619        ))
8620    }
8621
8622    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8623        self.alloc_uninit::<f32>(n)
8624    }
8625
8626    /// i8 uninitialized scratch (same contract as `uninit`).
8627    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8628        self.alloc_uninit::<i8>(n)
8629    }
8630
8631    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
8632    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
8633    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
8634    #[allow(clippy::too_many_arguments)]
8635    pub fn rms_norm3(
8636        &self,
8637        x: &CudaSlice<f32>,
8638        w0: &CudaSlice<f32>,
8639        w1: &CudaSlice<f32>,
8640        w2: &CudaSlice<f32>,
8641        d0: &mut CudaSlice<f32>,
8642        d1: &mut CudaSlice<f32>,
8643        d2: &mut CudaSlice<f32>,
8644        ncols: usize,
8645        nrows: usize,
8646        eps: f32,
8647    ) -> Result<(), Box<dyn std::error::Error>> {
8648        let f = self.func("rms_norm3_f32");
8649        let cfg = LaunchConfig {
8650            grid_dim: (nrows as u32, 1, 1),
8651            block_dim: (rms_block(), 1, 1),
8652            shared_mem_bytes: 0,
8653        };
8654        let (nc, e) = (ncols as i32, eps);
8655        let __s_b = self.gpu.stream();
8656        let mut b = __s_b.launch_builder(&f);
8657        b.arg(x)
8658            .arg(w0)
8659            .arg(w1)
8660            .arg(w2)
8661            .arg(d0)
8662            .arg(d1)
8663            .arg(d2)
8664            .arg(&nc)
8665            .arg(&e);
8666        unsafe {
8667            b.launch(cfg)?;
8668        }
8669        Ok(())
8670    }
8671
8672    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
8673    #[allow(clippy::too_many_arguments)]
8674    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
8675    /// piggybacks on the same conditions.
8676    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8677        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8678        *WARP_ON.get_or_init(|| {
8679            std::env::var("MEMRA_QKVNORM_W")
8680                .map(|v| v != "0")
8681                .unwrap_or(true)
8682        }) && ncols % 4 == 0
8683            && rows >= 64
8684    }
8685
8686    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
8687    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
8688    #[allow(clippy::too_many_arguments)]
8689    pub fn rms_norm_qkv_w4b(
8690        &self,
8691        q: &CudaSlice<f32>,
8692        k: &CudaSlice<f32>,
8693        v: &CudaSlice<f32>,
8694        wq: &CudaSlice<f32>,
8695        wk: &CudaSlice<f32>,
8696        wv: &CudaSlice<f32>,
8697        dq: &mut CudaSlice<f32>,
8698        dk: &mut CudaSlice<f32>,
8699        dv: &mut CudaSlice<f32>,
8700        dvb: &mut CudaSlice<u8>,
8701        ncols: usize,
8702        rq: usize,
8703        rk: usize,
8704        eps: f32,
8705        vf16: bool,
8706    ) -> Result<(), Box<dyn std::error::Error>> {
8707        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8708        let f = self.func("rms_norm_qkv_w4b_f32");
8709        let rows = (rq + 2 * rk) as u32;
8710        let cfg = LaunchConfig {
8711            grid_dim: (rows.div_ceil(8), 1, 1),
8712            block_dim: (256, 1, 1),
8713            shared_mem_bytes: 0,
8714        };
8715        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8716        let vf = vf16 as i32;
8717        let __s_b = self.gpu.stream();
8718        let mut b = __s_b.launch_builder(&f);
8719        b.arg(q)
8720            .arg(k)
8721            .arg(v)
8722            .arg(wq)
8723            .arg(wk)
8724            .arg(wv)
8725            .arg(dq)
8726            .arg(dk)
8727            .arg(dv)
8728            .arg(&mut *dvb)
8729            .arg(&nc)
8730            .arg(&rqi)
8731            .arg(&rki)
8732            .arg(&rvi)
8733            .arg(&e)
8734            .arg(&vf);
8735        unsafe {
8736            b.launch(cfg)?;
8737        }
8738        Ok(())
8739    }
8740
8741    pub fn rms_norm_qkv(
8742        &self,
8743        q: &CudaSlice<f32>,
8744        k: &CudaSlice<f32>,
8745        v: &CudaSlice<f32>,
8746        wq: &CudaSlice<f32>,
8747        wk: &CudaSlice<f32>,
8748        wv: &CudaSlice<f32>,
8749        dq: &mut CudaSlice<f32>,
8750        dk: &mut CudaSlice<f32>,
8751        dv: &mut CudaSlice<f32>,
8752        ncols: usize,
8753        rq: usize,
8754        rk: usize,
8755        eps: f32,
8756    ) -> Result<(), Box<dyn std::error::Error>> {
8757        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
8758        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
8759        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
8760        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8761        let warp_on = *WARP_ON.get_or_init(|| {
8762            std::env::var("MEMRA_QKVNORM_W")
8763                .map(|v| v != "0")
8764                .unwrap_or(true)
8765        });
8766        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
8767        // replay numerics are untouched on every model; only prefill depth takes the new config.
8768        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8769            let f = self.func("rms_norm_qkv_w4_f32");
8770            let rows = (rq + 2 * rk) as u32;
8771            let cfg = LaunchConfig {
8772                grid_dim: (rows.div_ceil(8), 1, 1),
8773                block_dim: (256, 1, 1),
8774                shared_mem_bytes: 0,
8775            };
8776            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8777            let __s_b = self.gpu.stream();
8778            let mut b = __s_b.launch_builder(&f);
8779            b.arg(q)
8780                .arg(k)
8781                .arg(v)
8782                .arg(wq)
8783                .arg(wk)
8784                .arg(wv)
8785                .arg(dq)
8786                .arg(dk)
8787                .arg(dv)
8788                .arg(&nc)
8789                .arg(&rqi)
8790                .arg(&rki)
8791                .arg(&rvi)
8792                .arg(&e);
8793            unsafe {
8794                b.launch(cfg)?;
8795            }
8796            return Ok(());
8797        }
8798        let f = self.func("rms_norm_qkv_f32");
8799        let grid = (rq + 2 * rk) as u32;
8800        let cfg = LaunchConfig {
8801            grid_dim: (grid, 1, 1),
8802            block_dim: (rms_block(), 1, 1),
8803            shared_mem_bytes: 0,
8804        };
8805        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8806        let __s_b = self.gpu.stream();
8807        let mut b = __s_b.launch_builder(&f);
8808        b.arg(q)
8809            .arg(k)
8810            .arg(v)
8811            .arg(wq)
8812            .arg(wk)
8813            .arg(wv)
8814            .arg(dq)
8815            .arg(dk)
8816            .arg(dv)
8817            .arg(&nc)
8818            .arg(&rqi)
8819            .arg(&rki)
8820            .arg(&e);
8821        unsafe {
8822            b.launch(cfg)?;
8823        }
8824        Ok(())
8825    }
8826
8827    /// gemma4 fused pair of rms_norms over two different inputs (same width).
8828    #[allow(clippy::too_many_arguments)]
8829    pub fn rms_norm2x(
8830        &self,
8831        a: &CudaSlice<f32>,
8832        bb: &CudaSlice<f32>,
8833        wa: &CudaSlice<f32>,
8834        wb: &CudaSlice<f32>,
8835        da: &mut CudaSlice<f32>,
8836        db: &mut CudaSlice<f32>,
8837        ncols: usize,
8838        nrows: usize,
8839        eps: f32,
8840    ) -> Result<(), Box<dyn std::error::Error>> {
8841        let f = self.func("rms_norm2x_f32");
8842        let cfg = LaunchConfig {
8843            grid_dim: (2 * nrows as u32, 1, 1),
8844            block_dim: (rms_block(), 1, 1),
8845            shared_mem_bytes: 0,
8846        };
8847        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
8848        let __s_b = self.gpu.stream();
8849        let mut b = __s_b.launch_builder(&f);
8850        b.arg(a)
8851            .arg(bb)
8852            .arg(wa)
8853            .arg(wb)
8854            .arg(da)
8855            .arg(db)
8856            .arg(&nc)
8857            .arg(&nr)
8858            .arg(&e);
8859        unsafe {
8860            b.launch(cfg)?;
8861        }
8862        Ok(())
8863    }
8864
8865    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
8866    pub fn softcap(
8867        &self,
8868        y: &mut CudaSlice<f32>,
8869        cap: f32,
8870        n: usize,
8871    ) -> Result<(), Box<dyn std::error::Error>> {
8872        let f = self.func("softcap_f32");
8873        let cfg = LaunchConfig::for_num_elems(n as u32);
8874        let ni = n as i32;
8875        let __s_b = self.gpu.stream();
8876        let mut b = __s_b.launch_builder(&f);
8877        b.arg(y).arg(&cap).arg(&ni);
8878        unsafe {
8879            b.launch(cfg)?;
8880        }
8881        Ok(())
8882    }
8883
8884    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
8885    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
8886    pub fn mask_ids_rows(
8887        &self,
8888        y: &mut CudaSlice<f32>,
8889        ids: &CudaSlice<i32>,
8890        n_ids: usize,
8891        n_vocab: usize,
8892        t: usize,
8893    ) -> Result<(), Box<dyn std::error::Error>> {
8894        let f = self.func("mask_ids_rows_f32");
8895        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
8896        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
8897        let __s_b = self.gpu.stream();
8898        let mut b = __s_b.launch_builder(&f);
8899        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
8900        unsafe {
8901            b.launch(cfg)?;
8902        }
8903        Ok(())
8904    }
8905
8906    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
8907    #[allow(clippy::too_many_arguments)]
8908    pub fn add_scale_rms_norm(
8909        &self,
8910        a: &CudaSlice<f32>,
8911        b_in: &CudaSlice<f32>,
8912        c: f32,
8913        w: &CudaSlice<f32>,
8914        res: &mut CudaSlice<f32>,
8915        dst: &mut CudaSlice<f32>,
8916        ncols: usize,
8917        nrows: usize,
8918        eps: f32,
8919    ) -> Result<(), Box<dyn std::error::Error>> {
8920        let f = self.func("add_scale_rms_norm_f32");
8921        let cfg = LaunchConfig {
8922            grid_dim: (nrows as u32, 1, 1),
8923            block_dim: (rms_block(), 1, 1),
8924            shared_mem_bytes: 0,
8925        };
8926        let (nc, e2) = (ncols as i32, eps);
8927        let __s_b = self.gpu.stream();
8928        let mut b = __s_b.launch_builder(&f);
8929        b.arg(a)
8930            .arg(b_in)
8931            .arg(&c)
8932            .arg(w)
8933            .arg(res)
8934            .arg(dst)
8935            .arg(&nc)
8936            .arg(&e2);
8937        unsafe {
8938            b.launch(cfg)?;
8939        }
8940        Ok(())
8941    }
8942
8943    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
8944    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
8945    #[allow(clippy::too_many_arguments)]
8946    pub fn add_scale_rms_norm_q8_1(
8947        &self,
8948        a: &CudaSlice<f32>,
8949        b_in: &CudaSlice<f32>,
8950        c: f32,
8951        w: &CudaSlice<f32>,
8952        res: &mut CudaSlice<f32>,
8953        ncols: usize,
8954        nrows: usize,
8955        eps: f32,
8956    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8957        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8958        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8959        let (nc, e2) = (ncols as i32, eps);
8960        if Self::pdl_on() && Self::pdl_wb_on() {
8961            {
8962                use cudarc::driver::{DevicePtr, DevicePtrMut};
8963                let s = &self.gpu.stream();
8964                let (pa, _g0) = a.device_ptr(s);
8965                let (pb, _g1) = b_in.device_ptr(s);
8966                let (pw, _g2) = w.device_ptr(s);
8967                let (pr, _g3) = res.device_ptr_mut(s);
8968                let (pq, _g4) = out_q.device_ptr_mut(s);
8969                let (pd, _g5) = out_d.device_ptr_mut(s);
8970                let mut ps = [
8971                    &pa as *const _ as *mut std::ffi::c_void,
8972                    &pb as *const _ as *mut _,
8973                    &c as *const _ as *mut _,
8974                    &pw as *const _ as *mut _,
8975                    &pr as *const _ as *mut _,
8976                    &pq as *const _ as *mut _,
8977                    &pd as *const _ as *mut _,
8978                    &nc as *const _ as *mut _,
8979                    &e2 as *const _ as *mut _,
8980                ];
8981                unsafe {
8982                    self.launch_pdl(
8983                        "add_scale_rms_norm_q8_1",
8984                        (nrows as u32, 1, 1),
8985                        (rms_block(), 1, 1),
8986                        &mut ps,
8987                    )?;
8988                }
8989            }
8990            return Ok((out_q, out_d));
8991        }
8992        let f = self.func("add_scale_rms_norm_q8_1");
8993        let cfg = LaunchConfig {
8994            grid_dim: (nrows as u32, 1, 1),
8995            block_dim: (rms_block(), 1, 1),
8996            shared_mem_bytes: 0,
8997        };
8998        let __s_b = self.gpu.stream();
8999        let mut b = __s_b.launch_builder(&f);
9000        b.arg(a)
9001            .arg(b_in)
9002            .arg(&c)
9003            .arg(w)
9004            .arg(res)
9005            .arg(&mut out_q)
9006            .arg(&mut out_d)
9007            .arg(&nc)
9008            .arg(&e2);
9009        unsafe {
9010            b.launch(cfg)?;
9011        }
9012        Ok((out_q, out_d))
9013    }
9014
9015    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
9016    #[allow(clippy::too_many_arguments)]
9017    pub fn add_scale_rms_norm_q8_1_into(
9018        &self,
9019        a: &CudaSlice<f32>,
9020        b_in: &CudaSlice<f32>,
9021        c: f32,
9022        w: &CudaSlice<f32>,
9023        res: &mut CudaSlice<f32>,
9024        ncols: usize,
9025        nrows: usize,
9026        eps: f32,
9027        out_q: &mut CudaSlice<i8>,
9028        out_d: &mut CudaSlice<f32>,
9029    ) -> Result<(), Box<dyn std::error::Error>> {
9030        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9031        let (nc, e2) = (ncols as i32, eps);
9032        if Self::pdl_on() && Self::pdl_wb_on() {
9033            use cudarc::driver::{DevicePtr, DevicePtrMut};
9034            let s = &self.gpu.stream();
9035            let (pa, _g0) = a.device_ptr(s);
9036            let (pb, _g1) = b_in.device_ptr(s);
9037            let (pw, _g2) = w.device_ptr(s);
9038            let (pr, _g3) = res.device_ptr_mut(s);
9039            let (pq, _g4) = out_q.device_ptr_mut(s);
9040            let (pd, _g5) = out_d.device_ptr_mut(s);
9041            let mut ps = [
9042                &pa as *const _ as *mut std::ffi::c_void,
9043                &pb as *const _ as *mut _,
9044                &c as *const _ as *mut _,
9045                &pw as *const _ as *mut _,
9046                &pr as *const _ as *mut _,
9047                &pq as *const _ as *mut _,
9048                &pd as *const _ as *mut _,
9049                &nc as *const _ as *mut _,
9050                &e2 as *const _ as *mut _,
9051            ];
9052            unsafe {
9053                self.launch_pdl(
9054                    "add_scale_rms_norm_q8_1",
9055                    (nrows as u32, 1, 1),
9056                    (rms_block(), 1, 1),
9057                    &mut ps,
9058                )?;
9059            }
9060            return Ok(());
9061        }
9062        let f = self.func("add_scale_rms_norm_q8_1");
9063        let cfg = LaunchConfig {
9064            grid_dim: (nrows as u32, 1, 1),
9065            block_dim: (rms_block(), 1, 1),
9066            shared_mem_bytes: 0,
9067        };
9068        let __s_b = self.gpu.stream();
9069        let mut b = __s_b.launch_builder(&f);
9070        b.arg(a)
9071            .arg(b_in)
9072            .arg(&c)
9073            .arg(w)
9074            .arg(res)
9075            .arg(&mut *out_q)
9076            .arg(&mut *out_d)
9077            .arg(&nc)
9078            .arg(&e2);
9079        unsafe {
9080            b.launch(cfg)?;
9081        }
9082        Ok(())
9083    }
9084
9085    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
9086    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
9087    #[allow(clippy::too_many_arguments)]
9088    pub fn rms_pre_add_scale_rms_norm_q8_1(
9089        &self,
9090        a: &CudaSlice<f32>,
9091        wa: &CudaSlice<f32>,
9092        b_in: &CudaSlice<f32>,
9093        c: f32,
9094        w: &CudaSlice<f32>,
9095        res: &mut CudaSlice<f32>,
9096        ncols: usize,
9097        nrows: usize,
9098        eps: f32,
9099    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9100        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9101        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9102        let (nc, e2) = (ncols as i32, eps);
9103        if Self::pdl_on() {
9104            {
9105                use cudarc::driver::{DevicePtr, DevicePtrMut};
9106                let s = &self.gpu.stream();
9107                let (pa, _g0) = a.device_ptr(s);
9108                let (pwa, _g1) = wa.device_ptr(s);
9109                let (pb, _g2) = b_in.device_ptr(s);
9110                let (pw, _g3) = w.device_ptr(s);
9111                let (pr, _g4) = res.device_ptr_mut(s);
9112                let (pq, _g5) = out_q.device_ptr_mut(s);
9113                let (pd, _g6) = out_d.device_ptr_mut(s);
9114                let mut ps = [
9115                    &pa as *const _ as *mut std::ffi::c_void,
9116                    &pwa as *const _ as *mut _,
9117                    &pb as *const _ as *mut _,
9118                    &c as *const _ as *mut _,
9119                    &pw as *const _ as *mut _,
9120                    &pr as *const _ as *mut _,
9121                    &pq as *const _ as *mut _,
9122                    &pd as *const _ as *mut _,
9123                    &nc as *const _ as *mut _,
9124                    &e2 as *const _ as *mut _,
9125                ];
9126                unsafe {
9127                    self.launch_pdl(
9128                        "rms_pre_add_scale_rms_norm_q8_1",
9129                        (nrows as u32, 1, 1),
9130                        (rms_block(), 1, 1),
9131                        &mut ps,
9132                    )?;
9133                }
9134            }
9135            return Ok((out_q, out_d));
9136        }
9137        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9138        let cfg = LaunchConfig {
9139            grid_dim: (nrows as u32, 1, 1),
9140            block_dim: (rms_block(), 1, 1),
9141            shared_mem_bytes: 0,
9142        };
9143        let __s_b = self.gpu.stream();
9144        let mut b = __s_b.launch_builder(&f);
9145        b.arg(a)
9146            .arg(wa)
9147            .arg(b_in)
9148            .arg(&c)
9149            .arg(w)
9150            .arg(res)
9151            .arg(&mut out_q)
9152            .arg(&mut out_d)
9153            .arg(&nc)
9154            .arg(&e2);
9155        unsafe {
9156            b.launch(cfg)?;
9157        }
9158        Ok((out_q, out_d))
9159    }
9160
9161    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
9162    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
9163    pub fn gelu_tanh_mul_q8_1(
9164        &self,
9165        gate: &CudaSlice<f32>,
9166        up: &cudarc::driver::CudaView<f32>,
9167        act: &mut CudaSlice<f32>,
9168        ncols: usize,
9169        nrows: usize,
9170    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9171        debug_assert!(ncols % 128 == 0);
9172        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9173        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9174        let nc = ncols as i32;
9175        if Self::pdl_on() {
9176            {
9177                use cudarc::driver::{DevicePtr, DevicePtrMut};
9178                let s = &self.gpu.stream();
9179                let (pg, _g0) = gate.device_ptr(s);
9180                let (pu, _g1) = up.device_ptr(s);
9181                let (pact, _g2) = act.device_ptr_mut(s);
9182                let (pq, _g3) = out_q.device_ptr_mut(s);
9183                let (pd, _g4) = out_d.device_ptr_mut(s);
9184                let mut ps = [
9185                    &pg as *const _ as *mut std::ffi::c_void,
9186                    &pu as *const _ as *mut _,
9187                    &pact as *const _ as *mut _,
9188                    &pq as *const _ as *mut _,
9189                    &pd as *const _ as *mut _,
9190                    &nc as *const _ as *mut _,
9191                ];
9192                unsafe {
9193                    self.launch_pdl(
9194                        "gelu_tanh_mul_q8_1",
9195                        (nrows as u32, 1, 1),
9196                        (rms_block(), 1, 1),
9197                        &mut ps,
9198                    )?;
9199                }
9200            }
9201            return Ok((out_q, out_d));
9202        }
9203        let f = self.func("gelu_tanh_mul_q8_1");
9204        let cfg = LaunchConfig {
9205            grid_dim: (nrows as u32, 1, 1),
9206            block_dim: (rms_block(), 1, 1),
9207            shared_mem_bytes: 0,
9208        };
9209        let __s_b = self.gpu.stream();
9210        let mut b = __s_b.launch_builder(&f);
9211        b.arg(gate)
9212            .arg(up)
9213            .arg(act)
9214            .arg(&mut out_q)
9215            .arg(&mut out_d)
9216            .arg(&nc);
9217        unsafe {
9218            b.launch(cfg)?;
9219        }
9220        Ok((out_q, out_d))
9221    }
9222
9223    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9224    #[allow(clippy::too_many_arguments)]
9225    pub fn gelu_tanh_mul_q8_1_into(
9226        &self,
9227        gate: &CudaSlice<f32>,
9228        up: &cudarc::driver::CudaView<f32>,
9229        act: &mut CudaSlice<f32>,
9230        ncols: usize,
9231        nrows: usize,
9232        out_q: &mut CudaSlice<i8>,
9233        out_d: &mut CudaSlice<f32>,
9234    ) -> Result<(), Box<dyn std::error::Error>> {
9235        debug_assert!(ncols % 128 == 0);
9236        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9237        let nc = ncols as i32;
9238        if Self::pdl_on() {
9239            use cudarc::driver::{DevicePtr, DevicePtrMut};
9240            let s = &self.gpu.stream();
9241            let (pg, _g0) = gate.device_ptr(s);
9242            let (pu, _g1) = up.device_ptr(s);
9243            let (pact, _g2) = act.device_ptr_mut(s);
9244            let (pq, _g3) = out_q.device_ptr_mut(s);
9245            let (pd, _g4) = out_d.device_ptr_mut(s);
9246            let mut ps = [
9247                &pg as *const _ as *mut std::ffi::c_void,
9248                &pu as *const _ as *mut _,
9249                &pact as *const _ as *mut _,
9250                &pq as *const _ as *mut _,
9251                &pd as *const _ as *mut _,
9252                &nc as *const _ as *mut _,
9253            ];
9254            unsafe {
9255                self.launch_pdl(
9256                    "gelu_tanh_mul_q8_1",
9257                    (nrows as u32, 1, 1),
9258                    (rms_block(), 1, 1),
9259                    &mut ps,
9260                )?;
9261            }
9262            return Ok(());
9263        }
9264        let f = self.func("gelu_tanh_mul_q8_1");
9265        let cfg = LaunchConfig {
9266            grid_dim: (nrows as u32, 1, 1),
9267            block_dim: (rms_block(), 1, 1),
9268            shared_mem_bytes: 0,
9269        };
9270        let __s_b = self.gpu.stream();
9271        let mut b = __s_b.launch_builder(&f);
9272        b.arg(gate)
9273            .arg(up)
9274            .arg(&mut *act)
9275            .arg(&mut *out_q)
9276            .arg(&mut *out_d)
9277            .arg(&nc);
9278        unsafe {
9279            b.launch(cfg)?;
9280        }
9281        Ok(())
9282    }
9283
9284    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9285    #[allow(clippy::too_many_arguments)]
9286    pub fn add_rms_norm3_q8z(
9287        &self,
9288        a: &CudaSlice<f32>,
9289        b_in: &CudaSlice<f32>,
9290        w0: &CudaSlice<f32>,
9291        w1: &CudaSlice<f32>,
9292        w2: &CudaSlice<f32>,
9293        res: &mut CudaSlice<f32>,
9294        out1: &mut CudaSlice<f32>,
9295        ncols: usize,
9296        nrows: usize,
9297        eps: f32,
9298    ) -> Result<
9299        (
9300            (CudaSlice<i8>, CudaSlice<f32>),
9301            (CudaSlice<i8>, CudaSlice<f32>),
9302        ),
9303        Box<dyn std::error::Error>,
9304    > {
9305        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9306        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9307        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9308        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9309        let f = self.func("add_rms_norm3_q8z_f32");
9310        let cfg = LaunchConfig {
9311            grid_dim: (nrows as u32, 1, 1),
9312            block_dim: (rms_block(), 1, 1),
9313            shared_mem_bytes: 0,
9314        };
9315        let (nc, e2) = (ncols as i32, eps);
9316        let __s_b = self.gpu.stream();
9317        let mut b = __s_b.launch_builder(&f);
9318        b.arg(a)
9319            .arg(b_in)
9320            .arg(w0)
9321            .arg(w1)
9322            .arg(w2)
9323            .arg(res)
9324            .arg(&mut q0)
9325            .arg(&mut d0)
9326            .arg(out1)
9327            .arg(&mut q2)
9328            .arg(&mut d2)
9329            .arg(&nc)
9330            .arg(&e2);
9331        unsafe {
9332            b.launch(cfg)?;
9333        }
9334        Ok(((q0, d0), (q2, d2)))
9335    }
9336
9337    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9338    #[allow(clippy::too_many_arguments)]
9339    pub fn add_rms_norm3(
9340        &self,
9341        a: &CudaSlice<f32>,
9342        b_in: &CudaSlice<f32>,
9343        w0: &CudaSlice<f32>,
9344        w1: &CudaSlice<f32>,
9345        w2: &CudaSlice<f32>,
9346        res: &mut CudaSlice<f32>,
9347        d0: &mut CudaSlice<f32>,
9348        d1: &mut CudaSlice<f32>,
9349        d2: &mut CudaSlice<f32>,
9350        ncols: usize,
9351        nrows: usize,
9352        eps: f32,
9353    ) -> Result<(), Box<dyn std::error::Error>> {
9354        let f = self.func("add_rms_norm3_f32");
9355        let cfg = LaunchConfig {
9356            grid_dim: (nrows as u32, 1, 1),
9357            block_dim: (rms_block(), 1, 1),
9358            shared_mem_bytes: 0,
9359        };
9360        let (nc, e2) = (ncols as i32, eps);
9361        let __s_b = self.gpu.stream();
9362        let mut b = __s_b.launch_builder(&f);
9363        b.arg(a)
9364            .arg(b_in)
9365            .arg(w0)
9366            .arg(w1)
9367            .arg(w2)
9368            .arg(res)
9369            .arg(d0)
9370            .arg(d1)
9371            .arg(d2)
9372            .arg(&nc)
9373            .arg(&e2);
9374        unsafe {
9375            b.launch(cfg)?;
9376        }
9377        Ok(())
9378    }
9379
9380    /// dst = (a + b) * c (residual add + layer scale, one launch).
9381    pub fn add_scale(
9382        &self,
9383        a: &CudaSlice<f32>,
9384        b_in: &CudaSlice<f32>,
9385        c: f32,
9386        dst: &mut CudaSlice<f32>,
9387        n: usize,
9388    ) -> Result<(), Box<dyn std::error::Error>> {
9389        let f = self.func("add_scale_f32");
9390        let cfg = LaunchConfig::for_num_elems(n as u32);
9391        let ni = n as i32;
9392        let __s_b = self.gpu.stream();
9393        let mut b = __s_b.launch_builder(&f);
9394        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9395        unsafe {
9396            b.launch(cfg)?;
9397        }
9398        Ok(())
9399    }
9400
9401    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
9402    pub fn layer_norm_bias(
9403        &self,
9404        x: &CudaSlice<f32>,
9405        w: &CudaSlice<f32>,
9406        b: &CudaSlice<f32>,
9407        dst: &mut CudaSlice<f32>,
9408        ncols: usize,
9409        nrows: usize,
9410        eps: f32,
9411    ) -> Result<(), Box<dyn std::error::Error>> {
9412        let f = self.func("layer_norm_bias_f32");
9413        let (nc, e) = (ncols as i32, eps);
9414        let cfg = LaunchConfig {
9415            grid_dim: (nrows as u32, 1, 1),
9416            block_dim: (256, 1, 1),
9417            shared_mem_bytes: 0,
9418        };
9419        let __s_b = self.gpu.stream();
9420        let mut lb = __s_b.launch_builder(&f);
9421        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9422        unsafe {
9423            lb.launch(cfg)?;
9424        }
9425        Ok(())
9426    }
9427
9428    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
9429    pub fn gelu_tanh(
9430        &self,
9431        x: &CudaSlice<f32>,
9432        dst: &mut CudaSlice<f32>,
9433        n: usize,
9434    ) -> Result<(), Box<dyn std::error::Error>> {
9435        let f = self.func("gelu_tanh_f32");
9436        let ni = n as i64;
9437        let cfg = LaunchConfig {
9438            grid_dim: (n.div_ceil(256) as u32, 1, 1),
9439            block_dim: (256, 1, 1),
9440            shared_mem_bytes: 0,
9441        };
9442        let __s_b = self.gpu.stream();
9443        let mut lb = __s_b.launch_builder(&f);
9444        lb.arg(x).arg(&mut *dst).arg(&ni);
9445        unsafe {
9446            lb.launch(cfg)?;
9447        }
9448        Ok(())
9449    }
9450
9451    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
9452    pub fn row_softmax(
9453        &self,
9454        x: &mut CudaSlice<f32>,
9455        ncols: usize,
9456        nrows: usize,
9457    ) -> Result<(), Box<dyn std::error::Error>> {
9458        let f = self.func("row_softmax_f32");
9459        let nc = ncols as i32;
9460        let cfg = LaunchConfig {
9461            grid_dim: (nrows as u32, 1, 1),
9462            block_dim: (256, 1, 1),
9463            shared_mem_bytes: 0,
9464        };
9465        let __s_b = self.gpu.stream();
9466        let mut lb = __s_b.launch_builder(&f);
9467        lb.arg(&mut *x).arg(&nc);
9468        unsafe {
9469            lb.launch(cfg)?;
9470        }
9471        Ok(())
9472    }
9473
9474    pub fn rms_norm(
9475        &self,
9476        x: &CudaSlice<f32>,
9477        w: &CudaSlice<f32>,
9478        dst: &mut CudaSlice<f32>,
9479        ncols: usize,
9480        nrows: usize,
9481        eps: f32,
9482    ) -> Result<(), Box<dyn std::error::Error>> {
9483        let (nc, e) = (ncols as i32, eps);
9484        let kname = if Self::norm_ilp_on() {
9485            "rms_norm_f32_v2"
9486        } else {
9487            "rms_norm_f32"
9488        };
9489        if Self::pdl_on() && Self::pdl_wb_on() {
9490            use cudarc::driver::{DevicePtr, DevicePtrMut};
9491            let s = &self.gpu.stream();
9492            let (px, _g0) = x.device_ptr(s);
9493            let (pw, _g1) = w.device_ptr(s);
9494            let (pd, _g2) = dst.device_ptr_mut(s);
9495            let mut ps = [
9496                &px as *const _ as *mut std::ffi::c_void,
9497                &pw as *const _ as *mut _,
9498                &pd as *const _ as *mut _,
9499                &nc as *const _ as *mut _,
9500                &e as *const _ as *mut _,
9501            ];
9502            unsafe {
9503                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9504            }
9505            return Ok(());
9506        }
9507        let f = self.func(kname);
9508        let cfg = LaunchConfig {
9509            grid_dim: (nrows as u32, 1, 1),
9510            block_dim: (rms_block(), 1, 1),
9511            shared_mem_bytes: 0,
9512        };
9513        let __s_b = self.gpu.stream();
9514        let mut b = __s_b.launch_builder(&f);
9515        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9516        unsafe {
9517            b.launch(cfg)?;
9518        }
9519        Ok(())
9520    }
9521
9522    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
9523    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
9524    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
9525    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
9526    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
9527    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
9528    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
9529    pub fn rms_norm_decode(
9530        &self,
9531        x: &CudaSlice<f32>,
9532        w: &CudaSlice<f32>,
9533        dst: &mut CudaSlice<f32>,
9534        ncols: usize,
9535        nrows: usize,
9536        eps: f32,
9537    ) -> Result<(), Box<dyn std::error::Error>> {
9538        let f = self.func(if Self::norm_ilp_on() {
9539            "rms_norm_f32_v2"
9540        } else {
9541            "rms_norm_f32"
9542        });
9543        let cfg = LaunchConfig {
9544            grid_dim: (nrows as u32, 1, 1),
9545            block_dim: (1024, 1, 1),
9546            shared_mem_bytes: 0,
9547        };
9548        let (nc, e) = (ncols as i32, eps);
9549        let __s_b = self.gpu.stream();
9550        let mut b = __s_b.launch_builder(&f);
9551        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9552        unsafe {
9553            b.launch(cfg)?;
9554        }
9555        Ok(())
9556    }
9557
9558    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
9559    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
9560    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
9561    pub fn rms_norm_q8_1(
9562        &self,
9563        x: &CudaSlice<f32>,
9564        w: &CudaSlice<f32>,
9565        ncols: usize,
9566        nrows: usize,
9567        eps: f32,
9568    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9569        let nblk = ncols / 32;
9570        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9571        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9572        let (nc, e) = (ncols as i32, eps);
9573        if Self::pdl_on() {
9574            {
9575                use cudarc::driver::{DevicePtr, DevicePtrMut};
9576                let s = &self.gpu.stream();
9577                let (px, _g0) = x.device_ptr(s);
9578                let (pw, _g1) = w.device_ptr(s);
9579                let (pq, _g2) = q.device_ptr_mut(s);
9580                let (pd, _g3) = d.device_ptr_mut(s);
9581                let mut ps = [
9582                    &px as *const _ as *mut std::ffi::c_void,
9583                    &pw as *const _ as *mut _,
9584                    &pq as *const _ as *mut _,
9585                    &pd as *const _ as *mut _,
9586                    &nc as *const _ as *mut _,
9587                    &e as *const _ as *mut _,
9588                ];
9589                unsafe {
9590                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9591                }
9592            }
9593            return Ok((q, d));
9594        }
9595        let f = self.func("rms_norm_q8_1");
9596        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
9597        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
9598        let cfg = LaunchConfig {
9599            grid_dim: (nrows as u32, 1, 1),
9600            block_dim: (1024, 1, 1),
9601            shared_mem_bytes: 0,
9602        };
9603        let __s_b = self.gpu.stream();
9604        let mut b = __s_b.launch_builder(&f);
9605        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9606        unsafe {
9607            b.launch(cfg)?;
9608        }
9609        Ok((q, d))
9610    }
9611
9612    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
9613    /// PDL arm), caller-owned outputs.
9614    pub fn rms_norm_q8_1_into(
9615        &self,
9616        x: &CudaSlice<f32>,
9617        w: &CudaSlice<f32>,
9618        ncols: usize,
9619        nrows: usize,
9620        eps: f32,
9621        q: &mut CudaSlice<i8>,
9622        d: &mut CudaSlice<f32>,
9623    ) -> Result<(), Box<dyn std::error::Error>> {
9624        let nblk = ncols / 32;
9625        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9626        let (nc, e) = (ncols as i32, eps);
9627        if Self::pdl_on() {
9628            use cudarc::driver::{DevicePtr, DevicePtrMut};
9629            let s = &self.gpu.stream();
9630            let (px, _g0) = x.device_ptr(s);
9631            let (pw, _g1) = w.device_ptr(s);
9632            let (pq, _g2) = q.device_ptr_mut(s);
9633            let (pd, _g3) = d.device_ptr_mut(s);
9634            let mut ps = [
9635                &px as *const _ as *mut std::ffi::c_void,
9636                &pw as *const _ as *mut _,
9637                &pq as *const _ as *mut _,
9638                &pd as *const _ as *mut _,
9639                &nc as *const _ as *mut _,
9640                &e as *const _ as *mut _,
9641            ];
9642            unsafe {
9643                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9644            }
9645            return Ok(());
9646        }
9647        let f = self.func("rms_norm_q8_1");
9648        let cfg = LaunchConfig {
9649            grid_dim: (nrows as u32, 1, 1),
9650            block_dim: (1024, 1, 1),
9651            shared_mem_bytes: 0,
9652        };
9653        let __s_b = self.gpu.stream();
9654        let mut b = __s_b.launch_builder(&f);
9655        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9656        unsafe {
9657            b.launch(cfg)?;
9658        }
9659        Ok(())
9660    }
9661
9662    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
9663    pub fn quantize_q8_1_into(
9664        &self,
9665        x: &CudaSlice<f32>,
9666        m: usize,
9667        in_f: usize,
9668        q: &mut CudaSlice<i8>,
9669        d: &mut CudaSlice<f32>,
9670    ) -> Result<(), Box<dyn std::error::Error>> {
9671        let nblk = in_f / 32;
9672        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9673        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9674        let (inf, mi) = (in_f as i32, m as i32);
9675        if Self::pdl_on() && Self::pdl_wb_on() {
9676            use cudarc::driver::{DevicePtr, DevicePtrMut};
9677            let s = &self.gpu.stream();
9678            let (px, _g0) = x.device_ptr(s);
9679            let (pq, _g1) = q.device_ptr_mut(s);
9680            let (pd, _g2) = d.device_ptr_mut(s);
9681            let mut ps = [
9682                &px as *const _ as *mut std::ffi::c_void,
9683                &pq as *const _ as *mut _,
9684                &pd as *const _ as *mut _,
9685                &inf as *const _ as *mut _,
9686                &mi as *const _ as *mut _,
9687            ];
9688            unsafe {
9689                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9690            }
9691            return Ok(());
9692        }
9693        let f = self.func("quantize_q8_1");
9694        let __s_b = self.gpu.stream();
9695        let mut b = __s_b.launch_builder(&f);
9696        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9697        unsafe {
9698            b.launch(cfg)?;
9699        }
9700        Ok(())
9701    }
9702
9703    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
9704    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
9705    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
9706    pub fn add_rms_norm_q8_1(
9707        &self,
9708        a: &CudaSlice<f32>,
9709        b_in: &CudaSlice<f32>,
9710        w: &CudaSlice<f32>,
9711        res: &mut CudaSlice<f32>,
9712        ncols: usize,
9713        nrows: usize,
9714        eps: f32,
9715    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9716        let nblk = ncols / 32;
9717        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9718        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9719        let f = self.func("add_rms_norm_q8_1");
9720        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
9721        let cfg = LaunchConfig {
9722            grid_dim: (nrows as u32, 1, 1),
9723            block_dim: (1024, 1, 1),
9724            shared_mem_bytes: 0,
9725        };
9726        let (nc, e) = (ncols as i32, eps);
9727        let __s_bld = self.gpu.stream();
9728        let mut bld = __s_bld.launch_builder(&f);
9729        bld.arg(a)
9730            .arg(b_in)
9731            .arg(w)
9732            .arg(res)
9733            .arg(&mut q)
9734            .arg(&mut d)
9735            .arg(&nc)
9736            .arg(&e);
9737        unsafe {
9738            bld.launch(cfg)?;
9739        }
9740        Ok((q, d))
9741    }
9742
9743    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
9744    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
9745    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
9746    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
9747    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
9748    #[allow(clippy::too_many_arguments)]
9749    pub fn join_add_rms_norm_raw(
9750        &self,
9751        a0_raw: u64,
9752        a1_raw: u64,
9753        x: &CudaSlice<f32>,
9754        w: &CudaSlice<f32>,
9755        res: &mut CudaSlice<f32>,
9756        dst: &mut CudaSlice<f32>,
9757        ncols: usize,
9758        eps: f32,
9759    ) -> Result<(), Box<dyn std::error::Error>> {
9760        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9761            return Err("join_add_rms_norm geometry".into());
9762        }
9763        let f = self.func("join_add_rms_norm_f32");
9764        let cfg = LaunchConfig {
9765            grid_dim: (1, 1, 1),
9766            block_dim: (rms_block(), 1, 1),
9767            shared_mem_bytes: 0,
9768        };
9769        let (nc, e) = (ncols as i32, eps);
9770        let __s_b = self.gpu.stream();
9771        let mut b = __s_b.launch_builder(&f);
9772        b.arg(&a0_raw)
9773            .arg(&a1_raw)
9774            .arg(x)
9775            .arg(w)
9776            .arg(&mut *res)
9777            .arg(&mut *dst)
9778            .arg(&nc)
9779            .arg(&e);
9780        unsafe {
9781            b.launch(cfg)?;
9782        }
9783        Ok(())
9784    }
9785
9786    pub fn add_rms_norm(
9787        &self,
9788        a: &CudaSlice<f32>,
9789        b: &CudaSlice<f32>,
9790        w: &CudaSlice<f32>,
9791        res: &mut CudaSlice<f32>,
9792        dst: &mut CudaSlice<f32>,
9793        ncols: usize,
9794        nrows: usize,
9795        eps: f32,
9796    ) -> Result<(), Box<dyn std::error::Error>> {
9797        let (nc, e) = (ncols as i32, eps);
9798        let kname = if Self::norm_ilp_on() {
9799            "add_rms_norm_f32_v2"
9800        } else {
9801            "add_rms_norm_f32"
9802        };
9803        if Self::pdl_on() && Self::pdl_wb_on() {
9804            use cudarc::driver::{DevicePtr, DevicePtrMut};
9805            let s = &self.gpu.stream();
9806            let (pa, _g0) = a.device_ptr(s);
9807            let (pb, _g1) = b.device_ptr(s);
9808            let (pw, _g2) = w.device_ptr(s);
9809            let (pr, _g3) = res.device_ptr_mut(s);
9810            let (pd, _g4) = dst.device_ptr_mut(s);
9811            let mut ps = [
9812                &pa as *const _ as *mut std::ffi::c_void,
9813                &pb as *const _ as *mut _,
9814                &pw as *const _ as *mut _,
9815                &pr as *const _ as *mut _,
9816                &pd as *const _ as *mut _,
9817                &nc as *const _ as *mut _,
9818                &e as *const _ as *mut _,
9819            ];
9820            unsafe {
9821                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9822            }
9823            return Ok(());
9824        }
9825        let f = self.func(kname);
9826        let cfg = LaunchConfig {
9827            grid_dim: (nrows as u32, 1, 1),
9828            block_dim: (rms_block(), 1, 1),
9829            shared_mem_bytes: 0,
9830        };
9831        let __s_b2 = self.gpu.stream();
9832        let mut b2 = __s_b2.launch_builder(&f);
9833        b2.arg(a)
9834            .arg(b)
9835            .arg(w)
9836            .arg(&mut *res)
9837            .arg(&mut *dst)
9838            .arg(&nc)
9839            .arg(&e);
9840        unsafe {
9841            b2.launch(cfg)?;
9842        }
9843        Ok(())
9844    }
9845
9846    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
9847    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
9848    #[allow(clippy::too_many_arguments)]
9849    pub fn rms_pre_add_rms_norm(
9850        &self,
9851        a: &CudaSlice<f32>,
9852        wa: &CudaSlice<f32>,
9853        b: &CudaSlice<f32>,
9854        w: &CudaSlice<f32>,
9855        res: &mut CudaSlice<f32>,
9856        dst: &mut CudaSlice<f32>,
9857        ncols: usize,
9858        nrows: usize,
9859        eps: f32,
9860    ) -> Result<(), Box<dyn std::error::Error>> {
9861        let f = self.func("rms_pre_add_rms_norm_f32");
9862        let cfg = LaunchConfig {
9863            grid_dim: (nrows as u32, 1, 1),
9864            block_dim: (rms_block(), 1, 1),
9865            shared_mem_bytes: 0,
9866        };
9867        let (nc, e) = (ncols as i32, eps);
9868        let __s_b2 = self.gpu.stream();
9869        let mut b2 = __s_b2.launch_builder(&f);
9870        b2.arg(a)
9871            .arg(wa)
9872            .arg(b)
9873            .arg(w)
9874            .arg(&mut *res)
9875            .arg(&mut *dst)
9876            .arg(&nc)
9877            .arg(&e);
9878        unsafe {
9879            b2.launch(cfg)?;
9880        }
9881        Ok(())
9882    }
9883
9884    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
9885    #[allow(clippy::too_many_arguments)]
9886    pub fn rms_pre_add_rms_norm_q8z(
9887        &self,
9888        a: &CudaSlice<f32>,
9889        wa: &CudaSlice<f32>,
9890        b: &CudaSlice<f32>,
9891        w: &CudaSlice<f32>,
9892        res: &mut CudaSlice<f32>,
9893        dst: &mut CudaSlice<f32>,
9894        ncols: usize,
9895        nrows: usize,
9896        eps: f32,
9897    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9898        debug_assert!(ncols % 128 == 0);
9899        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9900        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9901        let (nc, e) = (ncols as i32, eps);
9902        if Self::pdl_on() {
9903            {
9904                use cudarc::driver::{DevicePtr, DevicePtrMut};
9905                let s = &self.gpu.stream();
9906                let (pa, _g0) = a.device_ptr(s);
9907                let (pwa, _g1) = wa.device_ptr(s);
9908                let (pb, _g2) = b.device_ptr(s);
9909                let (pw, _g3) = w.device_ptr(s);
9910                let (pr, _g4) = res.device_ptr_mut(s);
9911                let (pdst, _g5) = dst.device_ptr_mut(s);
9912                let (pq, _g6) = out_q.device_ptr_mut(s);
9913                let (pd, _g7) = out_d.device_ptr_mut(s);
9914                let mut ps = [
9915                    &pa as *const _ as *mut std::ffi::c_void,
9916                    &pwa as *const _ as *mut _,
9917                    &pb as *const _ as *mut _,
9918                    &pw as *const _ as *mut _,
9919                    &pr as *const _ as *mut _,
9920                    &pdst as *const _ as *mut _,
9921                    &pq as *const _ as *mut _,
9922                    &pd as *const _ as *mut _,
9923                    &nc as *const _ as *mut _,
9924                    &e as *const _ as *mut _,
9925                ];
9926                unsafe {
9927                    self.launch_pdl(
9928                        "rms_pre_add_rms_norm_q8z_f32",
9929                        (nrows as u32, 1, 1),
9930                        (rms_block(), 1, 1),
9931                        &mut ps,
9932                    )?;
9933                }
9934            }
9935            return Ok((out_q, out_d));
9936        }
9937        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9938        let cfg = LaunchConfig {
9939            grid_dim: (nrows as u32, 1, 1),
9940            block_dim: (rms_block(), 1, 1),
9941            shared_mem_bytes: 0,
9942        };
9943        let __s_b2 = self.gpu.stream();
9944        let mut b2 = __s_b2.launch_builder(&f);
9945        b2.arg(a)
9946            .arg(wa)
9947            .arg(b)
9948            .arg(w)
9949            .arg(&mut *res)
9950            .arg(&mut *dst)
9951            .arg(&mut out_q)
9952            .arg(&mut out_d)
9953            .arg(&nc)
9954            .arg(&e);
9955        unsafe {
9956            b2.launch(cfg)?;
9957        }
9958        Ok((out_q, out_d))
9959    }
9960
9961    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
9962    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
9963    /// body must stay attribute-free (the fused2_into precedent).
9964    #[allow(clippy::too_many_arguments)]
9965    pub fn rms_pre_add_rms_norm_q8z_into(
9966        &self,
9967        a: &CudaSlice<f32>,
9968        wa: &CudaSlice<f32>,
9969        b: &CudaSlice<f32>,
9970        w: &CudaSlice<f32>,
9971        res: &mut CudaSlice<f32>,
9972        dst: &mut CudaSlice<f32>,
9973        ncols: usize,
9974        nrows: usize,
9975        eps: f32,
9976        out_q: &mut CudaSlice<i8>,
9977        out_d: &mut CudaSlice<f32>,
9978    ) -> Result<(), Box<dyn std::error::Error>> {
9979        debug_assert!(ncols % 128 == 0);
9980        let (nc, e) = (ncols as i32, eps);
9981        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9982        let cfg = LaunchConfig {
9983            grid_dim: (nrows as u32, 1, 1),
9984            block_dim: (rms_block(), 1, 1),
9985            shared_mem_bytes: 0,
9986        };
9987        let __s_b = self.gpu.stream();
9988        let mut b2 = __s_b.launch_builder(&f);
9989        b2.arg(a)
9990            .arg(wa)
9991            .arg(b)
9992            .arg(w)
9993            .arg(&mut *res)
9994            .arg(&mut *dst)
9995            .arg(&mut *out_q)
9996            .arg(&mut *out_d)
9997            .arg(&nc)
9998            .arg(&e);
9999        unsafe {
10000            b2.launch(cfg)?;
10001        }
10002        Ok(())
10003    }
10004
10005    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
10006    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
10007    #[allow(clippy::too_many_arguments)]
10008    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10009        &self,
10010        a: &CudaSlice<f32>,
10011        wa: &CudaSlice<f32>,
10012        b_in: &CudaSlice<f32>,
10013        c: f32,
10014        w: &CudaSlice<f32>,
10015        res: &mut CudaSlice<f32>,
10016        ncols: usize,
10017        nrows: usize,
10018        eps: f32,
10019        out_q: &mut CudaSlice<i8>,
10020        out_d: &mut CudaSlice<f32>,
10021    ) -> Result<(), Box<dyn std::error::Error>> {
10022        debug_assert!(ncols % 128 == 0);
10023        let (nc, e2) = (ncols as i32, eps);
10024        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10025        let cfg = LaunchConfig {
10026            grid_dim: (nrows as u32, 1, 1),
10027            block_dim: (rms_block(), 1, 1),
10028            shared_mem_bytes: 0,
10029        };
10030        let __s_b = self.gpu.stream();
10031        let mut b2 = __s_b.launch_builder(&f);
10032        b2.arg(a)
10033            .arg(wa)
10034            .arg(b_in)
10035            .arg(&c)
10036            .arg(w)
10037            .arg(&mut *res)
10038            .arg(&mut *out_q)
10039            .arg(&mut *out_d)
10040            .arg(&nc)
10041            .arg(&e2);
10042        unsafe {
10043            b2.launch(cfg)?;
10044        }
10045        Ok(())
10046    }
10047
10048    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
10049    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
10050    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
10051    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
10052    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
10053    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
10054    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
10055    pub fn g4_pnfold_on() -> bool {
10056        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10057        *ON.get_or_init(|| {
10058            std::env::var("MEMRA_G4_PNFOLD")
10059                .map(|v| v != "0")
10060                .unwrap_or(true)
10061        })
10062    }
10063
10064    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
10065    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
10066    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
10067    pub fn build_q4_out_concat3(
10068        &self,
10069        w0: &crate::model::GpuTensor,
10070        w1: &crate::model::GpuTensor,
10071        w2: &crate::model::GpuTensor,
10072    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10073        use crate::model::GpuTensor;
10074        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10075            match w {
10076                GpuTensor::Quant {
10077                    qtype,
10078                    row_bytes,
10079                    rp,
10080                    ..
10081                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10082                _ => None,
10083            }
10084        };
10085        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10086        else {
10087            return Ok(None);
10088        };
10089        if rb0 != rb1
10090            || rb0 != rb2
10091            || w0.in_features() != w1.in_features()
10092            || w0.in_features() != w2.in_features()
10093        {
10094            return Ok(None);
10095        }
10096        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10097            match w {
10098                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10099                _ => unreachable!(),
10100            }
10101        }
10102        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10103        let total = rb0 * (o0 + o1 + o2);
10104        let mut cat = self.alloc_u8(total)?;
10105        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10106        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10107        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10108        Ok(Some(GpuTensor::Quant {
10109            bytes: cat,
10110            qtype: QT_Q4_0,
10111            row_bytes: rb0,
10112            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10113            scale: 1.0,
10114            rp: false,
10115            #[cfg(memra_cutlass)]
10116            cutlass: None,
10117            fp8: None,
10118            blk: None,
10119            rp4: None,
10120            f16: None,
10121        }))
10122    }
10123
10124    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
10125    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
10126    ///
10127    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
10128    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
10129    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
10130    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
10131    ///
10132    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
10133    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
10134    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
10135    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
10136    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
10137    ///
10138    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
10139    /// width. A future partial-rotary caller fails at its first launch with the geometry named
10140    /// instead of serving quietly wrong logits.
10141    fn full_width_rope_only(
10142        kernel: &str,
10143        n_rot: usize,
10144        head_dim: usize,
10145    ) -> Result<(), Box<dyn std::error::Error>> {
10146        if n_rot == head_dim {
10147            return Ok(());
10148        }
10149        Err(format!(
10150            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10151             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10152             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10153             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10154             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10155        )
10156        .into())
10157    }
10158
10159    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
10160    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10161    /// ([`Engine::full_width_rope_only`]).
10162    #[allow(clippy::too_many_arguments)]
10163    pub fn rms_norm_qkv_rope_cat(
10164        &self,
10165        qkv: &CudaSlice<f32>,
10166        wq: &CudaSlice<f32>,
10167        wk: &CudaSlice<f32>,
10168        wv: &CudaSlice<f32>,
10169        q: &mut CudaSlice<f32>,
10170        k: &mut CudaSlice<f32>,
10171        v: &mut CudaSlice<f32>,
10172        head_dim: usize,
10173        n_rot: usize,
10174        rq: usize,
10175        rk: usize,
10176        pos: &CudaSlice<i32>,
10177        nh_q: usize,
10178        nh_k: usize,
10179        base: f32,
10180        freq_scale: f32,
10181        ff: Option<&CudaSlice<f32>>,
10182        eps: f32,
10183    ) -> Result<(), Box<dyn std::error::Error>> {
10184        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10185        let rows = rq + rk + rk;
10186        let theta_scale = base.powf(-2.0 / head_dim as f32);
10187        let (nc, rqi, rki, nhq, nhk) = (
10188            head_dim as i32,
10189            rq as i32,
10190            rk as i32,
10191            nh_q as i32,
10192            nh_k as i32,
10193        );
10194        if Self::pdl_on() {
10195            use cudarc::driver::{DevicePtr, DevicePtrMut};
10196            let s = &self.gpu.stream();
10197            let (pqkv, _g0) = qkv.device_ptr(s);
10198            let (pwq, _g1) = wq.device_ptr(s);
10199            let (pwk, _g2) = wk.device_ptr(s);
10200            let (pwv, _g3) = wv.device_ptr(s);
10201            let (pq, _g4) = q.device_ptr_mut(s);
10202            let (pk, _g5) = k.device_ptr_mut(s);
10203            let (pv, _g6) = v.device_ptr_mut(s);
10204            let (ppos, _g7) = pos.device_ptr(s);
10205            let (pff, _g8) = match ff {
10206                Some(t) => {
10207                    let (p, g) = t.device_ptr(s);
10208                    (p, Some(g))
10209                }
10210                None => (0, None),
10211            };
10212            let mut ps = [
10213                &pqkv as *const _ as *mut std::ffi::c_void,
10214                &pwq as *const _ as *mut _,
10215                &pwk as *const _ as *mut _,
10216                &pwv as *const _ as *mut _,
10217                &pq as *const _ as *mut _,
10218                &pk as *const _ as *mut _,
10219                &pv as *const _ as *mut _,
10220                &nc as *const _ as *mut _,
10221                &rqi as *const _ as *mut _,
10222                &rki as *const _ as *mut _,
10223                &ppos as *const _ as *mut _,
10224                &nhq as *const _ as *mut _,
10225                &nhk as *const _ as *mut _,
10226                &theta_scale as *const _ as *mut _,
10227                &freq_scale as *const _ as *mut _,
10228                &pff as *const _ as *mut _,
10229                &eps as *const _ as *mut _,
10230            ];
10231            unsafe {
10232                self.launch_pdl(
10233                    "rms_norm_qkv_rope_cat_f32",
10234                    (rows as u32, 1, 1),
10235                    (rms_block(), 1, 1),
10236                    &mut ps,
10237                )?;
10238            }
10239            return Ok(());
10240        }
10241        let f = self.func("rms_norm_qkv_rope_cat_f32");
10242        let cfg = LaunchConfig {
10243            grid_dim: (rows as u32, 1, 1),
10244            block_dim: (rms_block(), 1, 1),
10245            shared_mem_bytes: 0,
10246        };
10247        let __s_b = self.gpu.stream();
10248        let mut b = __s_b.launch_builder(&f);
10249        match ff {
10250            Some(t) => {
10251                b.arg(qkv)
10252                    .arg(wq)
10253                    .arg(wk)
10254                    .arg(wv)
10255                    .arg(&mut *q)
10256                    .arg(&mut *k)
10257                    .arg(&mut *v)
10258                    .arg(&nc)
10259                    .arg(&rqi)
10260                    .arg(&rki)
10261                    .arg(pos)
10262                    .arg(&nhq)
10263                    .arg(&nhk)
10264                    .arg(&theta_scale)
10265                    .arg(&freq_scale)
10266                    .arg(t)
10267                    .arg(&eps);
10268                unsafe {
10269                    b.launch(cfg)?;
10270                }
10271            }
10272            None => {
10273                let null: u64 = 0;
10274                b.arg(qkv)
10275                    .arg(wq)
10276                    .arg(wk)
10277                    .arg(wv)
10278                    .arg(&mut *q)
10279                    .arg(&mut *k)
10280                    .arg(&mut *v)
10281                    .arg(&nc)
10282                    .arg(&rqi)
10283                    .arg(&rki)
10284                    .arg(pos)
10285                    .arg(&nhq)
10286                    .arg(&nhk)
10287                    .arg(&theta_scale)
10288                    .arg(&freq_scale)
10289                    .arg(&null)
10290                    .arg(&eps);
10291                unsafe {
10292                    b.launch(cfg)?;
10293                }
10294            }
10295        }
10296        Ok(())
10297    }
10298
10299    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10300    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10301    /// ([`Engine::full_width_rope_only`]).
10302    #[allow(clippy::too_many_arguments)]
10303    pub fn rms_norm_qkv_rope(
10304        &self,
10305        q0: &CudaSlice<f32>,
10306        k0: &CudaSlice<f32>,
10307        v0: &CudaSlice<f32>,
10308        wq: &CudaSlice<f32>,
10309        wk: &CudaSlice<f32>,
10310        wv: &CudaSlice<f32>,
10311        q: &mut CudaSlice<f32>,
10312        k: &mut CudaSlice<f32>,
10313        v: &mut CudaSlice<f32>,
10314        head_dim: usize,
10315        n_rot: usize,
10316        rq: usize,
10317        rk: usize,
10318        pos: &CudaSlice<i32>,
10319        nh_q: usize,
10320        nh_k: usize,
10321        base: f32,
10322        freq_scale: f32,
10323        ff: Option<&CudaSlice<f32>>,
10324        eps: f32,
10325    ) -> Result<(), Box<dyn std::error::Error>> {
10326        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10327        let f = self.func("rms_norm_qkv_rope_f32");
10328        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10329        let cfg = LaunchConfig {
10330            grid_dim: (rows as u32, 1, 1),
10331            block_dim: (rms_block(), 1, 1),
10332            shared_mem_bytes: 0,
10333        };
10334        let theta_scale = base.powf(-2.0 / head_dim as f32);
10335        let (nc, rqi, rki, nhq, nhk) = (
10336            head_dim as i32,
10337            rq as i32,
10338            rk as i32,
10339            nh_q as i32,
10340            nh_k as i32,
10341        );
10342        let __s_b = self.gpu.stream();
10343        let mut b = __s_b.launch_builder(&f);
10344        match ff {
10345            Some(t) => {
10346                b.arg(q0)
10347                    .arg(k0)
10348                    .arg(v0)
10349                    .arg(wq)
10350                    .arg(wk)
10351                    .arg(wv)
10352                    .arg(&mut *q)
10353                    .arg(&mut *k)
10354                    .arg(&mut *v)
10355                    .arg(&nc)
10356                    .arg(&rqi)
10357                    .arg(&rki)
10358                    .arg(pos)
10359                    .arg(&nhq)
10360                    .arg(&nhk)
10361                    .arg(&theta_scale)
10362                    .arg(&freq_scale)
10363                    .arg(t)
10364                    .arg(&eps);
10365                unsafe {
10366                    b.launch(cfg)?;
10367                }
10368            }
10369            None => {
10370                let null: u64 = 0;
10371                b.arg(q0)
10372                    .arg(k0)
10373                    .arg(v0)
10374                    .arg(wq)
10375                    .arg(wk)
10376                    .arg(wv)
10377                    .arg(&mut *q)
10378                    .arg(&mut *k)
10379                    .arg(&mut *v)
10380                    .arg(&nc)
10381                    .arg(&rqi)
10382                    .arg(&rki)
10383                    .arg(pos)
10384                    .arg(&nhq)
10385                    .arg(&nhk)
10386                    .arg(&theta_scale)
10387                    .arg(&freq_scale)
10388                    .arg(&null)
10389                    .arg(&eps);
10390                unsafe {
10391                    b.launch(cfg)?;
10392                }
10393            }
10394        }
10395        Ok(())
10396    }
10397
10398    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
10399    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
10400    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
10401    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10402    /// ([`Engine::full_width_rope_only`]).
10403    #[allow(clippy::too_many_arguments)]
10404    pub fn rms_norm_qkv_rope_append_dc(
10405        &self,
10406        q0: &CudaSlice<f32>,
10407        k0: &CudaSlice<f32>,
10408        v0: &CudaSlice<f32>,
10409        wq: &CudaSlice<f32>,
10410        wk: &CudaSlice<f32>,
10411        wv: &CudaSlice<f32>,
10412        q: &mut CudaSlice<f32>,
10413        k: &mut CudaSlice<f32>,
10414        v: &mut CudaSlice<f32>,
10415        head_dim: usize,
10416        n_rot: usize,
10417        rq: usize,
10418        rk: usize,
10419        pos: &CudaSlice<i32>,
10420        nh_q: usize,
10421        nh_k: usize,
10422        base: f32,
10423        freq_scale: f32,
10424        ff: Option<&CudaSlice<f32>>,
10425        eps: f32,
10426        kc: &mut CudaSlice<u8>,
10427        vc: &mut CudaSlice<u8>,
10428        t_dev: &CudaSlice<i32>,
10429        k_tok_bytes: usize,
10430        v_tok_bytes: usize,
10431        g: bool,
10432    ) -> Result<(), Box<dyn std::error::Error>> {
10433        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10434        let rows = rq + rk + rk;
10435        let theta_scale = base.powf(-2.0 / head_dim as f32);
10436        let (nc, rqi, rki, nhq, nhk) = (
10437            head_dim as i32,
10438            rq as i32,
10439            rk as i32,
10440            nh_q as i32,
10441            nh_k as i32,
10442        );
10443        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10444        if Self::pdl_on() && Self::pdl_wb_on() {
10445            use cudarc::driver::{DevicePtr, DevicePtrMut};
10446            let s = &self.gpu.stream();
10447            let (p0, _a0) = q0.device_ptr(s);
10448            let (p1, _a1) = k0.device_ptr(s);
10449            let (p2, _a2) = v0.device_ptr(s);
10450            let (pwq, _a3) = wq.device_ptr(s);
10451            let (pwk, _a4) = wk.device_ptr(s);
10452            let (pwv, _a5) = wv.device_ptr(s);
10453            let (pq, _a6) = q.device_ptr_mut(s);
10454            let (pk, _a7) = k.device_ptr_mut(s);
10455            let (pv, _a8) = v.device_ptr_mut(s);
10456            let (pp, _a9) = pos.device_ptr(s);
10457            let pff: u64 = match ff {
10458                Some(t) => {
10459                    let (p, _gg) = t.device_ptr(s);
10460                    p as u64
10461                }
10462                None => 0,
10463            };
10464            let (pkc, _a10) = kc.device_ptr_mut(s);
10465            let (pvc, _a11) = vc.device_ptr_mut(s);
10466            let (pt, _a12) = t_dev.device_ptr(s);
10467            let mut ps = [
10468                &p0 as *const _ as *mut std::ffi::c_void,
10469                &p1 as *const _ as *mut _,
10470                &p2 as *const _ as *mut _,
10471                &pwq as *const _ as *mut _,
10472                &pwk as *const _ as *mut _,
10473                &pwv as *const _ as *mut _,
10474                &pq as *const _ as *mut _,
10475                &pk as *const _ as *mut _,
10476                &pv as *const _ as *mut _,
10477                &nc as *const _ as *mut _,
10478                &rqi as *const _ as *mut _,
10479                &rki as *const _ as *mut _,
10480                &pp as *const _ as *mut _,
10481                &nhq as *const _ as *mut _,
10482                &nhk as *const _ as *mut _,
10483                &theta_scale as *const _ as *mut _,
10484                &freq_scale as *const _ as *mut _,
10485                &pff as *const _ as *mut _,
10486                &eps as *const _ as *mut _,
10487                &pkc as *const _ as *mut _,
10488                &pvc as *const _ as *mut _,
10489                &pt as *const _ as *mut _,
10490                &ktb as *const _ as *mut _,
10491                &vtb as *const _ as *mut _,
10492            ];
10493            unsafe {
10494                self.launch_pdl_flash(
10495                    g,
10496                    "rms_norm_qkv_rope_append_dc_f32",
10497                    (rows as u32, 1, 1),
10498                    (rms_block(), 1, 1),
10499                    0,
10500                    &mut ps,
10501                )?;
10502            }
10503            return Ok(());
10504        }
10505        let f = if g {
10506            self.func_g("rms_norm_qkv_rope_append_dc_f32")
10507        } else {
10508            self.func("rms_norm_qkv_rope_append_dc_f32")
10509        };
10510        let cfg = LaunchConfig {
10511            grid_dim: (rows as u32, 1, 1),
10512            block_dim: (rms_block(), 1, 1),
10513            shared_mem_bytes: 0,
10514        };
10515        let __s_b = self.gpu.stream();
10516        let mut b = __s_b.launch_builder(&f);
10517        match ff {
10518            Some(t) => {
10519                b.arg(q0)
10520                    .arg(k0)
10521                    .arg(v0)
10522                    .arg(wq)
10523                    .arg(wk)
10524                    .arg(wv)
10525                    .arg(&mut *q)
10526                    .arg(&mut *k)
10527                    .arg(&mut *v)
10528                    .arg(&nc)
10529                    .arg(&rqi)
10530                    .arg(&rki)
10531                    .arg(pos)
10532                    .arg(&nhq)
10533                    .arg(&nhk)
10534                    .arg(&theta_scale)
10535                    .arg(&freq_scale)
10536                    .arg(t)
10537                    .arg(&eps)
10538                    .arg(&mut *kc)
10539                    .arg(&mut *vc)
10540                    .arg(t_dev)
10541                    .arg(&ktb)
10542                    .arg(&vtb);
10543                unsafe {
10544                    b.launch(cfg)?;
10545                }
10546            }
10547            None => {
10548                let null: u64 = 0;
10549                b.arg(q0)
10550                    .arg(k0)
10551                    .arg(v0)
10552                    .arg(wq)
10553                    .arg(wk)
10554                    .arg(wv)
10555                    .arg(&mut *q)
10556                    .arg(&mut *k)
10557                    .arg(&mut *v)
10558                    .arg(&nc)
10559                    .arg(&rqi)
10560                    .arg(&rki)
10561                    .arg(pos)
10562                    .arg(&nhq)
10563                    .arg(&nhk)
10564                    .arg(&theta_scale)
10565                    .arg(&freq_scale)
10566                    .arg(&null)
10567                    .arg(&eps)
10568                    .arg(&mut *kc)
10569                    .arg(&mut *vc)
10570                    .arg(t_dev)
10571                    .arg(&ktb)
10572                    .arg(&vtb);
10573                unsafe {
10574                    b.launch(cfg)?;
10575                }
10576            }
10577        }
10578        Ok(())
10579    }
10580
10581    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
10582    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
10583    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
10584    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
10585    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
10586    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
10587    /// `head_dim` ([`Engine::full_width_rope_only`]).
10588    #[allow(clippy::too_many_arguments)]
10589    pub fn rms_norm_qkv_rope_append(
10590        &self,
10591        q0: &CudaSlice<f32>,
10592        k0: &CudaSlice<f32>,
10593        v0: &CudaSlice<f32>,
10594        wq: &CudaSlice<f32>,
10595        wk: &CudaSlice<f32>,
10596        wv: &CudaSlice<f32>,
10597        q: &mut CudaSlice<f32>,
10598        k: &mut CudaSlice<f32>,
10599        v: &mut CudaSlice<f32>,
10600        head_dim: usize,
10601        n_rot: usize,
10602        rq: usize,
10603        rk: usize,
10604        pos: &CudaSlice<i32>,
10605        nh_q: usize,
10606        nh_k: usize,
10607        base: f32,
10608        freq_scale: f32,
10609        ff: Option<&CudaSlice<f32>>,
10610        eps: f32,
10611        kc: &mut CudaSlice<u8>,
10612        vc: &mut CudaSlice<u8>,
10613        t: usize,
10614        k_tok_bytes: usize,
10615        v_tok_bytes: usize,
10616        g: bool,
10617    ) -> Result<(), Box<dyn std::error::Error>> {
10618        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10619        let rows = rq + rk + rk;
10620        let theta_scale = base.powf(-2.0 / head_dim as f32);
10621        let (nc, rqi, rki, nhq, nhk) = (
10622            head_dim as i32,
10623            rq as i32,
10624            rk as i32,
10625            nh_q as i32,
10626            nh_k as i32,
10627        );
10628        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10629        let ti = t as i32;
10630        if Self::pdl_on() && Self::pdl_wb_on() {
10631            use cudarc::driver::{DevicePtr, DevicePtrMut};
10632            let s = &self.gpu.stream();
10633            let (p0, _a0) = q0.device_ptr(s);
10634            let (p1, _a1) = k0.device_ptr(s);
10635            let (p2, _a2) = v0.device_ptr(s);
10636            let (pwq, _a3) = wq.device_ptr(s);
10637            let (pwk, _a4) = wk.device_ptr(s);
10638            let (pwv, _a5) = wv.device_ptr(s);
10639            let (pq, _a6) = q.device_ptr_mut(s);
10640            let (pk, _a7) = k.device_ptr_mut(s);
10641            let (pv, _a8) = v.device_ptr_mut(s);
10642            let (pp, _a9) = pos.device_ptr(s);
10643            let pff: u64 = match ff {
10644                Some(t) => {
10645                    let (p, _gg) = t.device_ptr(s);
10646                    p as u64
10647                }
10648                None => 0,
10649            };
10650            let (pkc, _a10) = kc.device_ptr_mut(s);
10651            let (pvc, _a11) = vc.device_ptr_mut(s);
10652            let mut ps = [
10653                &p0 as *const _ as *mut std::ffi::c_void,
10654                &p1 as *const _ as *mut _,
10655                &p2 as *const _ as *mut _,
10656                &pwq as *const _ as *mut _,
10657                &pwk as *const _ as *mut _,
10658                &pwv as *const _ as *mut _,
10659                &pq as *const _ as *mut _,
10660                &pk as *const _ as *mut _,
10661                &pv as *const _ as *mut _,
10662                &nc as *const _ as *mut _,
10663                &rqi as *const _ as *mut _,
10664                &rki as *const _ as *mut _,
10665                &pp as *const _ as *mut _,
10666                &nhq as *const _ as *mut _,
10667                &nhk as *const _ as *mut _,
10668                &theta_scale as *const _ as *mut _,
10669                &freq_scale as *const _ as *mut _,
10670                &pff as *const _ as *mut _,
10671                &eps as *const _ as *mut _,
10672                &pkc as *const _ as *mut _,
10673                &pvc as *const _ as *mut _,
10674                &ti as *const _ as *mut _,
10675                &ktb as *const _ as *mut _,
10676                &vtb as *const _ as *mut _,
10677            ];
10678            unsafe {
10679                self.launch_pdl_flash(
10680                    g,
10681                    "rms_norm_qkv_rope_append_f32",
10682                    (rows as u32, 1, 1),
10683                    (rms_block(), 1, 1),
10684                    0,
10685                    &mut ps,
10686                )?;
10687            }
10688            return Ok(());
10689        }
10690        let f = if g {
10691            self.func_g("rms_norm_qkv_rope_append_f32")
10692        } else {
10693            self.func("rms_norm_qkv_rope_append_f32")
10694        };
10695        let cfg = LaunchConfig {
10696            grid_dim: (rows as u32, 1, 1),
10697            block_dim: (rms_block(), 1, 1),
10698            shared_mem_bytes: 0,
10699        };
10700        let __s_b = self.gpu.stream();
10701        let mut b = __s_b.launch_builder(&f);
10702        let null: u64 = 0;
10703        b.arg(q0)
10704            .arg(k0)
10705            .arg(v0)
10706            .arg(wq)
10707            .arg(wk)
10708            .arg(wv)
10709            .arg(&mut *q)
10710            .arg(&mut *k)
10711            .arg(&mut *v)
10712            .arg(&nc)
10713            .arg(&rqi)
10714            .arg(&rki)
10715            .arg(pos)
10716            .arg(&nhq)
10717            .arg(&nhk)
10718            .arg(&theta_scale)
10719            .arg(&freq_scale);
10720        match ff {
10721            Some(t) => {
10722                b.arg(t);
10723            }
10724            None => {
10725                b.arg(&null);
10726            }
10727        }
10728        b.arg(&eps)
10729            .arg(&mut *kc)
10730            .arg(&mut *vc)
10731            .arg(&ti)
10732            .arg(&ktb)
10733            .arg(&vtb);
10734        unsafe {
10735            b.launch(cfg)?;
10736        }
10737        Ok(())
10738    }
10739
10740    pub fn add_q8_1(
10741        &self,
10742        a: &CudaSlice<f32>,
10743        b: &CudaSlice<f32>,
10744        res: &mut CudaSlice<f32>,
10745        ncols: usize,
10746        nrows: usize,
10747    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10748        debug_assert!(ncols % 128 == 0);
10749        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10750        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10751        let f = self.func("add_q8_1_f32");
10752        let cfg = LaunchConfig {
10753            grid_dim: (nrows as u32, 1, 1),
10754            block_dim: (rms_block(), 1, 1),
10755            shared_mem_bytes: 0,
10756        };
10757        let nc = ncols as i32;
10758        let __s_b2 = self.gpu.stream();
10759        let mut b2 = __s_b2.launch_builder(&f);
10760        b2.arg(a)
10761            .arg(b)
10762            .arg(&mut *res)
10763            .arg(&mut out_q)
10764            .arg(&mut out_d)
10765            .arg(&nc);
10766        unsafe {
10767            b2.launch(cfg)?;
10768        }
10769        Ok((out_q, out_d))
10770    }
10771
10772    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
10773    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
10774    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
10775    pub fn rms_pre_add_q8_1(
10776        &self,
10777        a: &CudaSlice<f32>,
10778        wa: &CudaSlice<f32>,
10779        b: &CudaSlice<f32>,
10780        res: &mut CudaSlice<f32>,
10781        ncols: usize,
10782        nrows: usize,
10783        eps: f32,
10784    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10785        debug_assert!(ncols % 128 == 0);
10786        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10787        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10788        let f = self.func("rms_pre_add_q8_1_f32");
10789        let cfg = LaunchConfig {
10790            grid_dim: (nrows as u32, 1, 1),
10791            block_dim: (rms_block(), 1, 1),
10792            shared_mem_bytes: 0,
10793        };
10794        let (nc, ep) = (ncols as i32, eps);
10795        let __s_b2 = self.gpu.stream();
10796        let mut b2 = __s_b2.launch_builder(&f);
10797        b2.arg(a)
10798            .arg(wa)
10799            .arg(b)
10800            .arg(&mut *res)
10801            .arg(&mut out_q)
10802            .arg(&mut out_d)
10803            .arg(&nc)
10804            .arg(&ep);
10805        unsafe {
10806            b2.launch(cfg)?;
10807        }
10808        Ok((out_q, out_d))
10809    }
10810
10811    /// L2 norm per row (head_dim), no weight.
10812    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
10813    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
10814    pub fn l2_v2_on(ncols: usize) -> bool {
10815        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10816    }
10817
10818    pub fn l2_norm_pp(
10819        &self,
10820        x: &CudaSlice<f32>,
10821        dst: &mut CudaSlice<f32>,
10822        dst16: Option<&mut CudaSlice<u8>>,
10823        ncols: usize,
10824        nrows: usize,
10825        eps: f32,
10826    ) -> Result<(), Box<dyn std::error::Error>> {
10827        if Self::l2_v2_on(ncols) {
10828            let f = self.func("l2_norm_pp_v2_f32");
10829            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
10830            let cfg = LaunchConfig {
10831                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10832                block_dim: (256, 1, 1),
10833                shared_mem_bytes: 0,
10834            };
10835            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10836            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
10837            let d16: u64 = match dst16 {
10838                Some(d) => self.addr_u8(d),
10839                None => 0,
10840            };
10841            let __s_b = self.gpu.stream();
10842            let mut b = __s_b.launch_builder(&f);
10843            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
10844            unsafe {
10845                b.launch(cfg)?;
10846            }
10847            return Ok(());
10848        }
10849        self.l2_norm(x, dst, ncols, nrows, eps)
10850    }
10851
10852    pub fn l2_norm(
10853        &self,
10854        x: &CudaSlice<f32>,
10855        dst: &mut CudaSlice<f32>,
10856        ncols: usize,
10857        nrows: usize,
10858        eps: f32,
10859    ) -> Result<(), Box<dyn std::error::Error>> {
10860        let f = self.func("l2_norm_f32");
10861        let cfg = LaunchConfig {
10862            grid_dim: (nrows as u32, 1, 1),
10863            block_dim: (256, 1, 1),
10864            shared_mem_bytes: 0,
10865        };
10866        let (nc, e) = (ncols as i32, eps);
10867        let __s_b = self.gpu.stream();
10868        let mut b = __s_b.launch_builder(&f);
10869        b.arg(x).arg(dst).arg(&nc).arg(&e);
10870        unsafe {
10871            b.launch(cfg)?;
10872        }
10873        Ok(())
10874    }
10875
10876    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
10877    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
10878    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
10879    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
10880    /// propagate through gdn_scan and flip argmax on marginal logits.
10881    pub fn l2_norm_decode(
10882        &self,
10883        x: &CudaSlice<f32>,
10884        dst: &mut CudaSlice<f32>,
10885        ncols: usize,
10886        nrows: usize,
10887        eps: f32,
10888    ) -> Result<(), Box<dyn std::error::Error>> {
10889        let f = self.func("l2_norm_f32");
10890        let cfg = LaunchConfig {
10891            grid_dim: (nrows as u32, 1, 1),
10892            block_dim: (32, 1, 1),
10893            shared_mem_bytes: 0,
10894        };
10895        let (nc, e) = (ncols as i32, eps);
10896        let __s_b = self.gpu.stream();
10897        let mut b = __s_b.launch_builder(&f);
10898        b.arg(x).arg(dst).arg(&nc).arg(&e);
10899        unsafe {
10900            b.launch(cfg)?;
10901        }
10902        Ok(())
10903    }
10904
10905    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
10906    pub fn rope_neox(
10907        &self,
10908        x: &mut CudaSlice<f32>,
10909        pos: &CudaSlice<i32>,
10910        head_dim: usize,
10911        n_dims: usize,
10912        n_heads: usize,
10913        n_tokens: usize,
10914        freq_base: f32,
10915        freq_scale: f32,
10916    ) -> Result<(), Box<dyn std::error::Error>> {
10917        let f = self.func("rope_neox_f32");
10918        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10919        let grid = (n_heads * n_tokens) as u32;
10920        let cfg = LaunchConfig {
10921            grid_dim: (grid, 1, 1),
10922            block_dim: ((head_dim / 2) as u32, 1, 1),
10923            shared_mem_bytes: 0,
10924        };
10925        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10926        let __s_b = self.gpu.stream();
10927        let mut b = __s_b.launch_builder(&f);
10928        b.arg(x)
10929            .arg(pos)
10930            .arg(&hd)
10931            .arg(&nd)
10932            .arg(&nh)
10933            .arg(&theta_scale)
10934            .arg(&freq_scale);
10935        unsafe {
10936            b.launch(cfg)?;
10937        }
10938        Ok(())
10939    }
10940
10941    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
10942    pub fn rope_neox_ff(
10943        &self,
10944        x: &mut CudaSlice<f32>,
10945        pos: &CudaSlice<i32>,
10946        head_dim: usize,
10947        n_dims: usize,
10948        n_heads: usize,
10949        n_tokens: usize,
10950        freq_base: f32,
10951        freq_scale: f32,
10952        ff: &CudaSlice<f32>,
10953    ) -> Result<(), Box<dyn std::error::Error>> {
10954        let f = self.func("rope_neox_ff_f32");
10955        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10956        let grid = (n_heads * n_tokens) as u32;
10957        let cfg = LaunchConfig {
10958            grid_dim: (grid, 1, 1),
10959            block_dim: ((head_dim / 2) as u32, 1, 1),
10960            shared_mem_bytes: 0,
10961        };
10962        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10963        let __s_b = self.gpu.stream();
10964        let mut b = __s_b.launch_builder(&f);
10965        b.arg(x)
10966            .arg(pos)
10967            .arg(&hd)
10968            .arg(&nd)
10969            .arg(&nh)
10970            .arg(&theta_scale)
10971            .arg(&freq_scale)
10972            .arg(ff);
10973        unsafe {
10974            b.launch(cfg)?;
10975        }
10976        Ok(())
10977    }
10978
10979    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
10980    #[allow(clippy::too_many_arguments)]
10981    pub fn rope_neox2(
10982        &self,
10983        q: &mut CudaSlice<f32>,
10984        k: &mut CudaSlice<f32>,
10985        pos: &CudaSlice<i32>,
10986        head_dim: usize,
10987        n_dims: usize,
10988        nh_q: usize,
10989        nh_k: usize,
10990        n_tokens: usize,
10991        freq_base: f32,
10992        freq_scale: f32,
10993        ff: Option<&CudaSlice<f32>>,
10994    ) -> Result<(), Box<dyn std::error::Error>> {
10995        let f = self.func("rope_neox2_f32");
10996        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10997        let grid = ((nh_q + nh_k) * n_tokens) as u32;
10998        let cfg = LaunchConfig {
10999            grid_dim: (grid, 1, 1),
11000            block_dim: ((head_dim / 2) as u32, 1, 1),
11001            shared_mem_bytes: 0,
11002        };
11003        let (hd, nd, nq, nk, nt) = (
11004            head_dim as i32,
11005            n_dims as i32,
11006            nh_q as i32,
11007            nh_k as i32,
11008            n_tokens as i32,
11009        );
11010        let __s_b = self.gpu.stream();
11011        let mut b = __s_b.launch_builder(&f);
11012        b.arg(q)
11013            .arg(k)
11014            .arg(pos)
11015            .arg(&hd)
11016            .arg(&nd)
11017            .arg(&nq)
11018            .arg(&nk)
11019            .arg(&nt)
11020            .arg(&theta_scale)
11021            .arg(&freq_scale);
11022        match ff {
11023            Some(ffv) => {
11024                b.arg(ffv);
11025                unsafe {
11026                    b.launch(cfg)?;
11027                }
11028            }
11029            None => {
11030                let null: u64 = 0;
11031                b.arg(&null);
11032                unsafe {
11033                    b.launch(cfg)?;
11034                }
11035            }
11036        }
11037        Ok(())
11038    }
11039
11040    /// gemma4 R1: dst = GELU_tanh(gate) * up.
11041    pub fn gelu_tanh_mul(
11042        &self,
11043        gate: &CudaSlice<f32>,
11044        up: &CudaSlice<f32>,
11045        dst: &mut CudaSlice<f32>,
11046        n: usize,
11047    ) -> Result<(), Box<dyn std::error::Error>> {
11048        let f = self.func("gelu_tanh_mul_f32");
11049        let cfg = LaunchConfig::for_num_elems(n as u32);
11050        let ni = n as i32;
11051        let __s_b = self.gpu.stream();
11052        let mut b = __s_b.launch_builder(&f);
11053        b.arg(gate).arg(up).arg(dst).arg(&ni);
11054        unsafe {
11055            b.launch(cfg)?;
11056        }
11057        Ok(())
11058    }
11059
11060    pub fn silu_mul(
11061        &self,
11062        gate: &CudaSlice<f32>,
11063        up: &CudaSlice<f32>,
11064        dst: &mut CudaSlice<f32>,
11065        n: usize,
11066    ) -> Result<(), Box<dyn std::error::Error>> {
11067        let f = self.func("silu_mul_f32");
11068        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11069        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11070        let ni = n as i32;
11071        let __s_b = self.gpu.stream();
11072        let mut b = __s_b.launch_builder(&f);
11073        b.arg(gate).arg(up).arg(dst).arg(&ni);
11074        unsafe {
11075            b.launch(cfg)?;
11076        }
11077        Ok(())
11078    }
11079
11080    /// SwiGLU twin using Memra's host-matching expf transcription.
11081    pub fn silu_mul_host_expf(
11082        &self,
11083        gate: &CudaSlice<f32>,
11084        up: &CudaSlice<f32>,
11085        dst: &mut CudaSlice<f32>,
11086        n: usize,
11087    ) -> Result<(), Box<dyn std::error::Error>> {
11088        let f = self.func("silu_mul_host_expf_f32");
11089        let cfg = LaunchConfig::for_num_elems(n as u32);
11090        let ni = n as i32;
11091        let __s_b = self.gpu.stream();
11092        let mut b = __s_b.launch_builder(&f);
11093        b.arg(gate).arg(up).arg(dst).arg(&ni);
11094        unsafe {
11095            b.launch(cfg)?;
11096        }
11097        Ok(())
11098    }
11099
11100    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
11101    pub fn silu_clamped_mul_host_expf(
11102        &self,
11103        gate: &CudaSlice<f32>,
11104        up: &CudaSlice<f32>,
11105        limit: f32,
11106        dst: &mut CudaSlice<f32>,
11107        n: usize,
11108    ) -> Result<(), Box<dyn std::error::Error>> {
11109        if !limit.is_finite() || limit <= 0.0 {
11110            return Err(
11111                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11112            );
11113        }
11114        let f = self.func("silu_clamped_mul_host_expf_f32");
11115        let cfg = LaunchConfig::for_num_elems(n as u32);
11116        let ni = n as i32;
11117        let __s_b = self.gpu.stream();
11118        let mut b = __s_b.launch_builder(&f);
11119        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11120        unsafe {
11121            b.launch(cfg)?;
11122        }
11123        Ok(())
11124    }
11125
11126    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
11127    /// for the down projection — kills the standalone convert pass. Bit-identical class.
11128    pub fn silu_mul_f16out(
11129        &self,
11130        gate: &CudaSlice<f32>,
11131        up: &CudaSlice<f32>,
11132        dst: &mut CudaSlice<f32>,
11133        dst16: &mut CudaSlice<u8>,
11134        n: usize,
11135    ) -> Result<(), Box<dyn std::error::Error>> {
11136        let f = self.func("silu_mul_f16out_f32");
11137        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11138        let ni = n as i32;
11139        let __s_b = self.gpu.stream();
11140        let mut b = __s_b.launch_builder(&f);
11141        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11142        unsafe {
11143            b.launch(cfg)?;
11144        }
11145        Ok(())
11146    }
11147
11148    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
11149    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
11150    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
11151    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
11152    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
11153    /// launches per dense FFN layer (the gate+up post-matmul scales).
11154    pub fn silu_mul_scaled(
11155        &self,
11156        gate: &CudaSlice<f32>,
11157        up: &CudaSlice<f32>,
11158        gs: f32,
11159        us: f32,
11160        dst: &mut CudaSlice<f32>,
11161        n: usize,
11162    ) -> Result<(), Box<dyn std::error::Error>> {
11163        let f = self.func("silu_mul_scaled_f32");
11164        let cfg = LaunchConfig::for_num_elems(n as u32);
11165        let ni = n as i32;
11166        let (gsf, usf) = (gs, us);
11167        let __s_b = self.gpu.stream();
11168        let mut b = __s_b.launch_builder(&f);
11169        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11170        unsafe {
11171            b.launch(cfg)?;
11172        }
11173        Ok(())
11174    }
11175
11176    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
11177    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
11178    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
11179    #[allow(clippy::too_many_arguments)]
11180    pub fn swigluoai_mul_scaled(
11181        &self,
11182        gate: &CudaSlice<f32>,
11183        up: &CudaSlice<f32>,
11184        gs: f32,
11185        us: f32,
11186        alpha: f32,
11187        limit: f32,
11188        dst: &mut CudaSlice<f32>,
11189        n: usize,
11190    ) -> Result<(), Box<dyn std::error::Error>> {
11191        let f = self.func("swigluoai_mul_scaled_f32");
11192        let cfg = LaunchConfig::for_num_elems(n as u32);
11193        let ni = n as i32;
11194        let __s_b = self.gpu.stream();
11195        let mut b = __s_b.launch_builder(&f);
11196        b.arg(gate)
11197            .arg(up)
11198            .arg(&gs)
11199            .arg(&us)
11200            .arg(&alpha)
11201            .arg(&limit)
11202            .arg(dst)
11203            .arg(&ni);
11204        unsafe {
11205            b.launch(cfg)?;
11206        }
11207        Ok(())
11208    }
11209
11210    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11211    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11212    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11213    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11214    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11215    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11216    /// n must be a multiple of 32 (n_ff always is).
11217    pub fn silu_mul_scaled_q8_1(
11218        &self,
11219        gate: &CudaSlice<f32>,
11220        up: &CudaSlice<f32>,
11221        gs: f32,
11222        us: f32,
11223        n: usize,
11224    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11225        let f = self.func("silu_mul_scaled_q8_1");
11226        let nblk = n / 32;
11227        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11228        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11229        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11230        let cfg = LaunchConfig::for_num_elems(n as u32);
11231        let (gsf, usf, ni) = (gs, us, n as i32);
11232        let __s_b = self.gpu.stream();
11233        let mut b = __s_b.launch_builder(&f);
11234        b.arg(gate)
11235            .arg(up)
11236            .arg(&gsf)
11237            .arg(&usf)
11238            .arg(&mut aq)
11239            .arg(&mut ad)
11240            .arg(&ni);
11241        unsafe {
11242            b.launch(cfg)?;
11243        }
11244        Ok((aq, ad))
11245    }
11246
11247    pub fn add(
11248        &self,
11249        a: &CudaSlice<f32>,
11250        b_in: &CudaSlice<f32>,
11251        dst: &mut CudaSlice<f32>,
11252        n: usize,
11253    ) -> Result<(), Box<dyn std::error::Error>> {
11254        let f = self.func("add_f32");
11255        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11256        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11257        let ni = n as i32;
11258        let __s_bld = self.gpu.stream();
11259        let mut bld = __s_bld.launch_builder(&f);
11260        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11261        unsafe {
11262            bld.launch(cfg)?;
11263        }
11264        Ok(())
11265    }
11266
11267    pub fn mul(
11268        &self,
11269        a: &CudaSlice<f32>,
11270        b_in: &CudaSlice<f32>,
11271        dst: &mut CudaSlice<f32>,
11272        n: usize,
11273    ) -> Result<(), Box<dyn std::error::Error>> {
11274        let f = self.func("mul_f32");
11275        let cfg = LaunchConfig::for_num_elems(n as u32);
11276        let ni = n as i32;
11277        let __s_bld = self.gpu.stream();
11278        let mut bld = __s_bld.launch_builder(&f);
11279        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11280        unsafe {
11281            bld.launch(cfg)?;
11282        }
11283        Ok(())
11284    }
11285
11286    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11287    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11288    pub fn matmul(
11289        &self,
11290        w: &crate::model::GpuTensor,
11291        x: &CudaSlice<f32>,
11292        m: usize,
11293    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11294        use crate::model::GpuTensor;
11295        let in_f = w.in_features();
11296        let out_f = w.out_features();
11297        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11298        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11299        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11300        // gives nothing). Quantize the activation once here then call the GEMM.
11301        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11302        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11303        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11304        #[allow(non_snake_case)]
11305        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11306        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11307        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11308            usize::MAX
11309        } else {
11310            16usize
11311        };
11312
11313        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11314        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11315        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11316        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11317        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11318        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11319        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11320        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11321        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11322        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11323        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11324        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11325        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11326        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11327        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11328        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11329        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11330        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11331        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11332        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11333        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11334        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11335        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11336        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11337        if m >= GEMM_M_THRESHOLD {
11338            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11339                return Ok(y);
11340            }
11341            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
11342            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
11343            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
11344            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
11345            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
11346            // tile defaults differently by operand source.
11347            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11348                return Ok(y);
11349            }
11350            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
11351            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
11352            if let Some(y) = self.try_f16_gemm(w, x, m)? {
11353                return Ok(y);
11354            }
11355        }
11356        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
11357        // m threshold the rest of this method uses:
11358        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
11359        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
11360        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
11361        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
11362        //     across every tier by construction with no batched twin needed.
11363        //
11364        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
11365        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
11366        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
11367        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
11368        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
11369        // arms is what makes sure it never gets there.
11370        if let GpuTensor::Quant { qtype, .. } = w {
11371            if *qtype == QT_F8_E4M3_BLK {
11372                if m >= GEMM_M_THRESHOLD {
11373                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11374                        return Ok(y);
11375                    }
11376                }
11377                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11378                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11379                    return Ok(y);
11380                }
11381            }
11382        }
11383        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11384            return self.qmatvec_mmq(w, x, m);
11385        }
11386        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11387            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11388            return self.qmatvec_gemm(w, &aq, &ad, m);
11389        }
11390        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
11391        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
11392        if m >= GEMM_M_THRESHOLD {
11393            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11394                return Ok(y);
11395            }
11396        }
11397        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
11398        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
11399        // to Stage-A f32-dequant (the correctness oracle path).
11400        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11401        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
11402        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
11403        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
11404        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
11405        if m == 1 && fast {
11406            if let GpuTensor::Quant {
11407                bytes,
11408                qtype,
11409                row_bytes,
11410                rp,
11411                rp4,
11412                scale,
11413                ..
11414            } = w
11415            {
11416                if self.mmvq_supports(*qtype) {
11417                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
11418                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
11419                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
11420                    let (bytes, rp) = match rp4 {
11421                        Some(m4) => (m4, true),
11422                        None => (bytes, *rp),
11423                    };
11424                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11425                    return self.qmatvec_mmvq(
11426                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11427                    );
11428                }
11429            }
11430        }
11431        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
11432        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
11433        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
11434        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
11435        // block below. MEMRA_NO_BATCHED -> per-m path.
11436        //
11437        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
11438        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
11439        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
11440        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
11441        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
11442        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
11443        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
11444        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
11445        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
11446        if (2..=16).contains(&m)
11447            && fast
11448            && std::env::var("MEMRA_NO_BATCHED").is_err()
11449            && (m <= 4 || Self::b8_enabled())
11450        {
11451            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
11452            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
11453            // is present (rp4) — the mirror pick below then routes to the _rp family.
11454            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
11455            // because the native e4m3 row layout is already aligned and needs no mirror.
11456            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
11457            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
11458            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
11459            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
11460            let m_ok = m <= 8
11461                || matches!(w, GpuTensor::Quant { qtype, .. }
11462                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11463                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11464            if m_ok {
11465                if let GpuTensor::Quant {
11466                    bytes,
11467                    qtype,
11468                    row_bytes,
11469                    rp,
11470                    rp4,
11471                    ..
11472                } = w
11473                {
11474                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11475                        let (bytes, rp) = match rp4 {
11476                            Some(m4) => (m4, true),
11477                            None => (bytes, *rp),
11478                        };
11479                        let mcols = Self::batched_mcols(m);
11480                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11481                        let mut y = self.qmatvec_mmvq_batched(
11482                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11483                        )?;
11484                        if let GpuTensor::Quant { scale, .. } = w {
11485                            if *scale != 1.0 {
11486                                self.scale_inplace(&mut y, *scale, m * out_f)?;
11487                            }
11488                        }
11489                        return Ok(y);
11490                    }
11491                }
11492            }
11493        }
11494        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
11495        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
11496        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
11497        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
11498        // for this dtype, so the generic match below must never see it under `fast`.
11499        if fast {
11500            if let GpuTensor::Quant {
11501                bytes,
11502                qtype,
11503                row_bytes,
11504                scale,
11505                ..
11506            } = w
11507            {
11508                if *qtype == QT_F8_E4M3 {
11509                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11510                    return self.qmatvec_mmvq(
11511                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11512                    );
11513                }
11514            }
11515        }
11516        let mut y = match w {
11517            GpuTensor::Quant {
11518                bytes,
11519                qtype,
11520                row_bytes,
11521                ..
11522            } if fast && *qtype == QT_Q8_0 => {
11523                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11524            }
11525            GpuTensor::Quant {
11526                bytes,
11527                qtype,
11528                row_bytes,
11529                ..
11530            } if fast && *qtype == QT_Q4_K => {
11531                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11532            }
11533            GpuTensor::Quant {
11534                bytes,
11535                qtype,
11536                row_bytes,
11537                ..
11538            } if fast && *qtype == QT_Q6_K => {
11539                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11540            }
11541            GpuTensor::Quant {
11542                bytes,
11543                qtype,
11544                row_bytes,
11545                ..
11546            } if fast && *qtype == QT_Q5_K => {
11547                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11548            }
11549            GpuTensor::Quant {
11550                bytes,
11551                qtype,
11552                row_bytes,
11553                ..
11554            } if fast && *qtype == QT_Q3_K => {
11555                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11556            }
11557            GpuTensor::Quant {
11558                bytes,
11559                qtype,
11560                row_bytes,
11561                rp,
11562                ..
11563            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11564                if *rp {
11565                    "qmatvec_nvfp4_dp4a_rp"
11566                } else {
11567                    "qmatvec_nvfp4_dp4a"
11568                },
11569                &bytes.slice(0..bytes.len()),
11570                x,
11571                m,
11572                in_f,
11573                out_f,
11574                *row_bytes,
11575            )?,
11576            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
11577            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
11578            // anomaly (research/kat-anomaly-20260802/).
11579            GpuTensor::Quant {
11580                bytes,
11581                qtype,
11582                row_bytes,
11583                ..
11584            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11585                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11586            }
11587            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
11588            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
11589            // without first writing the matching kernel, or func() will panic
11590            // "kernel ... not in any fatbin".
11591            GpuTensor::Quant {
11592                bytes,
11593                qtype,
11594                row_bytes,
11595                rp,
11596                ..
11597            } =>
11598            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
11599            // deq(row,j) form cannot address the planes; same value/product order).
11600            {
11601                self.qmatvec(
11602                    bytes,
11603                    x,
11604                    m,
11605                    in_f,
11606                    out_f,
11607                    if *rp && *qtype == QT_NVFP4 {
11608                        QT_NVFP4_RP
11609                    } else {
11610                        *qtype
11611                    },
11612                    *row_bytes,
11613                )?
11614            }
11615            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11616            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
11617            // cuBLASLt f32 GEMV as the Float arm.
11618            GpuTensor::FloatBf16 { data, .. } => {
11619                self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11620            }
11621        };
11622        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
11623        if let GpuTensor::Quant { scale, .. } = w {
11624            if *scale != 1.0 {
11625                self.scale_inplace(&mut y, *scale, m * out_f)?;
11626            }
11627        }
11628        Ok(y)
11629    }
11630
11631    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
11632    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
11633    ///
11634    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
11635    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
11636    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
11637    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
11638    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
11639    /// path must not pay an env lookup for a flag that is off.
11640    pub fn stage_a_raw_needed() -> bool {
11641        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11642        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11643    }
11644
11645    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
11646    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
11647    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11648        use crate::model::GpuTensor;
11649        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11650            return false;
11651        }
11652        match w {
11653            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
11654            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
11655            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
11656            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
11657            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
11658            // block class has no fused twin yet, so each of its projections takes its own launch.
11659            GpuTensor::Quant { qtype, .. } => {
11660                matches!(
11661                    *qtype,
11662                    QT_Q8_0
11663                        | QT_Q4_K
11664                        | QT_Q6_K
11665                        | QT_Q5_K
11666                        | QT_Q3_K
11667                        | QT_NVFP4
11668                        | QT_F8_E4M3
11669                        | QT_F8_E4M3_BLK
11670                        | QT_Q4_0
11671                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11672            }
11673            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11674        }
11675    }
11676
11677    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
11678    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
11679    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
11680    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
11681    pub fn matmul_pre(
11682        &self,
11683        w: &crate::model::GpuTensor,
11684        aq: &CudaSlice<i8>,
11685        ad: &CudaSlice<f32>,
11686        x_fallback: &CudaSlice<f32>,
11687        m: usize,
11688    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11689        use crate::model::GpuTensor;
11690        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
11691        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
11692        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
11693        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
11694        // rc=30013 dig, 2026-07-31).
11695        let x_raw_ok = x_fallback.len() >= m * w.in_features();
11696        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
11697        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
11698        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11699            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11700                return Ok(y);
11701            }
11702            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
11703            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
11704            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11705                return Ok(y);
11706            }
11707            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
11708            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11709                return Ok(y);
11710            }
11711        }
11712        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
11713        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
11714        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
11715        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
11716        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
11717        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11718            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11719                return Ok(y);
11720            }
11721        }
11722        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11723            return Ok(y);
11724        }
11725        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
11726        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
11727        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
11728        // aq/ad.
11729        if m >= 16
11730            && w.out_features() >= 128
11731            && self.mmq_supports(w)
11732            && !self.verify_exact_on()
11733            && x_raw_ok
11734        {
11735            return self.qmatvec_mmq(w, x_fallback, m);
11736        }
11737        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
11738        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
11739        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11740            if let Some(y) =
11741                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11742            {
11743                return Ok(y);
11744            }
11745        }
11746        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
11747        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
11748        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11749            return self.qmatvec_gemm(w, aq, ad, m);
11750        }
11751        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
11752        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
11753        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
11754        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
11755        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
11756        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
11757        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
11758        // which reads `m * in_f` floats out of a 0-byte allocation ->
11759        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
11760        // it poisons the context, so every LATER request in that process fails with an unrelated
11761        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
11762        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
11763        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
11764        // dense artifact and left the arm with no working truth instrument.
11765        //
11766        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
11767        // strictly better than an illegal address surfacing later at an unrelated sync point, and
11768        // an oracle that cannot run must say so rather than corrupt the context it runs in.
11769        if !self.uses_q8_1_fast(w) {
11770            if !x_raw_ok {
11771                return Err(format!(
11772                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11773                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11774                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
11775                     activation (see Engine::rms_norm_decode, which is bit-identical to \
11776                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11777                    x_fallback.len(),
11778                    m,
11779                    w.in_features(),
11780                    m * w.in_features()
11781                )
11782                .into());
11783            }
11784            return self.matmul(w, x_fallback, m);
11785        }
11786        let in_f = w.in_features();
11787        let out_f = w.out_features();
11788        let (bytes, qtype, row_bytes, scale, rp) = match w {
11789            GpuTensor::Quant {
11790                bytes,
11791                qtype,
11792                row_bytes,
11793                scale,
11794                rp,
11795                ..
11796            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11797            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11798        };
11799        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
11800        // the dp4a/oracle tails below keep the raw GGUF bytes.
11801        let (mbytes, mrp) = match w {
11802            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11803            _ => (bytes, rp),
11804        };
11805        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
11806        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
11807        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
11808        if m == 1 && self.mmvq_supports(qtype) {
11809            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11810        }
11811        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
11812        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
11813        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
11814        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
11815        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
11816        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
11817        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
11818        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
11819        // m=5..8 on the old per-m path (b8-tier-only seam).
11820        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
11821        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
11822        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
11823        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11824            && std::env::var("MEMRA_NO_BATCHED").is_err()
11825            && (m <= 4 || Self::b8_enabled())
11826            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
11827            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
11828            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
11829            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11830                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11831        {
11832            let mcols = Self::batched_mcols(m);
11833            return self.qmatvec_mmvq_batched(
11834                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11835            );
11836        }
11837        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
11838        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
11839        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
11840        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
11841        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
11842        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
11843            let (b2, r2) = if qtype == QT_Q4_0 {
11844                (mbytes, mrp)
11845            } else {
11846                (bytes, rp)
11847            };
11848            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
11849        }
11850        let name = match qtype {
11851            QT_Q8_0 => "qmatvec_q8_0_dp4a",
11852            QT_Q4_K => "qmatvec_q4_K_dp4a",
11853            QT_Q6_K => "qmatvec_q6_K_dp4a",
11854            QT_Q5_K => "qmatvec_q5_K_dp4a",
11855            QT_Q3_K => "qmatvec_q3_K_dp4a",
11856            QT_NVFP4 => {
11857                if rp {
11858                    "qmatvec_nvfp4_dp4a_rp"
11859                } else {
11860                    "qmatvec_nvfp4_dp4a"
11861                }
11862            }
11863            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
11864            _ => unreachable!(),
11865        };
11866        let f = self.func(name);
11867        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
11868        let cfg = LaunchConfig {
11869            grid_dim: (out_f as u32, m as u32, 1),
11870            block_dim: (128, 1, 1),
11871            shared_mem_bytes: 0,
11872        };
11873        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11874        let __s_b = self.gpu.stream();
11875        let mut b = __s_b.launch_builder(&f);
11876        b.arg(bytes)
11877            .arg(aq)
11878            .arg(ad)
11879            .arg(&mut y)
11880            .arg(&inf)
11881            .arg(&outf)
11882            .arg(&mi)
11883            .arg(&rb);
11884        unsafe {
11885            b.launch(cfg)?;
11886        }
11887        if scale != 1.0 {
11888            self.scale_inplace(&mut y, scale, m * out_f)?;
11889        }
11890        Ok(y)
11891    }
11892
11893    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
11894    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
11895    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
11896    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
11897    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
11898    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
11899    /// reduce as m=1); this method just forces that path unconditionally.
11900    pub fn matmul_decode_exact(
11901        &self,
11902        w: &crate::model::GpuTensor,
11903        x: &CudaSlice<f32>,
11904        m: usize,
11905    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11906        use crate::model::GpuTensor;
11907        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
11908        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
11909        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
11910        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
11911        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
11912        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
11913        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
11914        if let GpuTensor::Float { data, .. } = w {
11915            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
11916        }
11917        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
11918        // float linear (same n-independent reduction contract as the Float arm above).
11919        if let GpuTensor::FloatBf16 { data, .. } = w {
11920            let (in_f, out_f) = (w.in_features(), w.out_features());
11921            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
11922        }
11923        if !self.uses_q8_1_fast(w) {
11924            return self.matmul(w, x, m);
11925        }
11926        let in_f = w.in_features();
11927        let out_f = w.out_features();
11928        let (bytes, qtype, row_bytes, scale, rp) = match w {
11929            GpuTensor::Quant {
11930                bytes,
11931                qtype,
11932                row_bytes,
11933                scale,
11934                rp,
11935                ..
11936            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11937            _ => return self.matmul(w, x, m),
11938        };
11939        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
11940        // which does its own mirror pick).
11941        let (bytes, rp) = match w {
11942            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11943            _ => (bytes, rp),
11944        };
11945        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11946        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
11947        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
11948        // (token,row) by construction, which is exactly what this method exists to guarantee.
11949        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11950            return Ok(y);
11951        }
11952        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
11953        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
11954        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
11955        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
11956        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
11957        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
11958        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
11959        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
11960        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11961            && std::env::var("MEMRA_NO_BATCHED").is_err()
11962            && (m <= 4 || Self::b8_enabled())
11963            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
11964            // no mirror precondition, `rp` selects the layout only.
11965            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
11966                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
11967        {
11968            let mcols = Self::batched_mcols(m);
11969            return self.qmatvec_mmvq_batched(
11970                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11971            );
11972        }
11973        if self.mmvq_supports(qtype) {
11974            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
11975            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
11976            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11977        }
11978        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
11979        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
11980        self.matmul_pre(w, &aq, &ad, x, m)
11981    }
11982
11983    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
11984    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
11985    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
11986    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
11987    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
11988    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
11989    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
11990    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
11991    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
11992    pub fn matmul_decode_exact_pre(
11993        &self,
11994        w: &crate::model::GpuTensor,
11995        aq: &CudaSlice<i8>,
11996        ad: &CudaSlice<f32>,
11997        m: usize,
11998    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11999        use crate::model::GpuTensor;
12000        debug_assert!(
12001            self.uses_q8_1_fast(w),
12002            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12003        );
12004        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
12005        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12006            return Ok(y);
12007        }
12008        let in_f = w.in_features();
12009        let out_f = w.out_features();
12010        let (bytes, qtype, row_bytes, scale, rp) = match w {
12011            GpuTensor::Quant {
12012                bytes,
12013                qtype,
12014                row_bytes,
12015                scale,
12016                rp,
12017                ..
12018            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12019            _ => {
12020                return Err(
12021                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12022                );
12023            }
12024        };
12025        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
12026        let (bytes, rp) = match w {
12027            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12028            _ => (bytes, rp),
12029        };
12030        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
12031        if (2..=16).contains(&m)
12032            && self.batched_supports(qtype)
12033            && self.mmvq_supports(qtype)
12034            && std::env::var("MEMRA_NO_BATCHED").is_err()
12035            && (m <= 4 || Self::b8_enabled())
12036            && (m <= 8
12037                || qtype == QT_Q4_0
12038                || qtype == QT_Q6_K
12039                || qtype == QT_F8_E4M3
12040                || qtype == QT_NVFP4
12041                || qtype == QT_Q4_K
12042                || qtype == QT_Q5_K
12043                || qtype == QT_Q8_0)
12044        {
12045            let mcols = Self::batched_mcols(m);
12046            return self.qmatvec_mmvq_batched(
12047                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12048            );
12049        }
12050        if self.mmvq_supports(qtype) {
12051            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12052        }
12053        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
12054        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
12055        let x0 = self.zeros(0)?;
12056        self.matmul_pre(w, aq, ad, &x0, m)
12057    }
12058
12059    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
12060    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
12061    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
12062    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
12063    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
12064    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
12065    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
12066    /// per-tensor path.
12067    pub fn matmul_decode_exact_dual_pre(
12068        &self,
12069        w0: &crate::model::GpuTensor,
12070        w1: &crate::model::GpuTensor,
12071        aq: &CudaSlice<i8>,
12072        ad: &CudaSlice<f32>,
12073        m: usize,
12074    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12075    {
12076        use crate::model::GpuTensor;
12077        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12078        let on = *ON.get_or_init(|| {
12079            std::env::var("MEMRA_SPEC_DUAL_T")
12080                .map(|v| v != "0")
12081                .unwrap_or(true)
12082        });
12083        if !on
12084            || !(2..=7).contains(&m)
12085            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12086            || !self.uses_q8_1_fast(w0)
12087            || !self.uses_q8_1_fast(w1)
12088        {
12089            return Ok(None);
12090        }
12091        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
12092        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
12093        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
12094        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
12095        if !self.mmvq_supports(QT_NVFP4) {
12096            return Ok(None);
12097        }
12098        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12099        if w1.in_features() != in_f || w1.out_features() != out_f {
12100            return Ok(None);
12101        }
12102        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12103            (
12104                GpuTensor::Quant {
12105                    bytes: b0,
12106                    qtype: q0,
12107                    row_bytes: rb0,
12108                    scale: s0,
12109                    rp: rp0,
12110                    rp4: None,
12111                    ..
12112                },
12113                GpuTensor::Quant {
12114                    bytes: b1,
12115                    qtype: q1,
12116                    row_bytes: rb1,
12117                    scale: s1,
12118                    rp: rp1,
12119                    rp4: None,
12120                    ..
12121                },
12122            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12123                (b0, b1, *rb0, *s0, *s1, *rp0)
12124            }
12125            _ => return Ok(None),
12126        };
12127        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
12128        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
12129        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12130        {
12131            return Ok(None);
12132        }
12133        let (y0, y1) =
12134            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12135        Ok(Some(((y0, s0), (y1, s1))))
12136    }
12137
12138    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
12139    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
12140    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
12141    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
12142    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
12143    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
12144    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
12145    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
12146    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
12147    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
12148    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
12149    pub fn matmul_decode_exact_group4_pre(
12150        &self,
12151        ws: [&crate::model::GpuTensor; 4],
12152        aq: &CudaSlice<i8>,
12153        ad: &CudaSlice<f32>,
12154        m: usize,
12155    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12156        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12157        let on = *ON.get_or_init(|| {
12158            std::env::var("MEMRA_TK_GDN_GROUP")
12159                .map(|v| v != "0")
12160                .unwrap_or(true)
12161        });
12162        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12163    }
12164
12165    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
12166    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
12167    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
12168    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
12169    pub fn matmul_decode_exact_group3_pre(
12170        &self,
12171        ws: [&crate::model::GpuTensor; 3],
12172        aq: &CudaSlice<i8>,
12173        ad: &CudaSlice<f32>,
12174        m: usize,
12175    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12176        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12177        let on = *ON.get_or_init(|| {
12178            std::env::var("MEMRA_TK_FA_GROUP")
12179                .map(|v| v != "0")
12180                .unwrap_or(true)
12181        });
12182        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12183    }
12184
12185    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
12186    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
12187    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
12188    fn matmul_decode_exact_group_pre(
12189        &self,
12190        ws: &[&crate::model::GpuTensor],
12191        aq: &CudaSlice<i8>,
12192        ad: &CudaSlice<f32>,
12193        m: usize,
12194        on: bool,
12195        tag: &'static str,
12196    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12197        use crate::model::GpuTensor;
12198        if !on
12199            || !(2..=16).contains(&m)
12200            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12201            || (m > 4 && !Self::b8_enabled())
12202            || !self.mmvq_supports(QT_NVFP4)
12203            || !self.batched_supports(QT_NVFP4)
12204        {
12205            return Ok(None);
12206        }
12207        let in_f = ws[0].in_features();
12208        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12209        for w in ws {
12210            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12211                return Ok(None);
12212            }
12213            match w {
12214                GpuTensor::Quant {
12215                    bytes,
12216                    qtype,
12217                    scale,
12218                    rp: true,
12219                    rp4: None,
12220                    ..
12221                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12222                    parts.push((bytes, w.out_features(), *scale));
12223                }
12224                _ => return Ok(None),
12225            }
12226        }
12227        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12228        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12229        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12230        let mcols = if (5..=7).contains(&m) && b567 {
12231            m
12232        } else {
12233            Self::batched_mcols(m)
12234        };
12235        let kname: &'static str = match mcols {
12236            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12237            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12238            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12239            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12240            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12241            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12242            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12243            _ => return Ok(None),
12244        };
12245        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12246        // the second door's print on the slice-D battery — key the once-set by tag.
12247        if std::env::var("MEMRA_DEBUG").is_ok() {
12248            use std::sync::Mutex;
12249            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12250            let mut seen = SEEN.lock().unwrap();
12251            if !seen.contains(&tag) {
12252                seen.push(tag);
12253                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12254            }
12255        }
12256        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12257        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12258        let total: usize = parts.iter().map(|p| p.1).sum();
12259        let three = parts.len() == 3;
12260        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12261        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12262        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12263        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12264        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12265        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12266        let cfg = LaunchConfig {
12267            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12268            block_dim: (32, ROWS_PER_BLOCK, 1),
12269            shared_mem_bytes: 0,
12270        };
12271        let (inf, mi) = (in_f as i32, m as i32);
12272        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12273        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12274        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12275        let s3 = if three { 1.0f32 } else { parts[3].2 };
12276        let w3 = if three { parts[0].0 } else { parts[3].0 };
12277        let f = self.func(kname);
12278        let __s_b = self.gpu.stream();
12279        let mut b = __s_b.launch_builder(&f);
12280        b.arg(parts[0].0)
12281            .arg(parts[1].0)
12282            .arg(parts[2].0)
12283            .arg(w3)
12284            .arg(aq)
12285            .arg(ad)
12286            .arg(&mut y0)
12287            .arg(&mut y1)
12288            .arg(&mut y2)
12289            .arg(&mut y3)
12290            .arg(&inf)
12291            .arg(&n0)
12292            .arg(&n1)
12293            .arg(&n2)
12294            .arg(&n3)
12295            .arg(&mi)
12296            .arg(&s0)
12297            .arg(&s1)
12298            .arg(&s2)
12299            .arg(&s3);
12300        unsafe {
12301            b.launch(cfg)?;
12302        }
12303        Ok(Some(if three {
12304            vec![y0, y1, y2]
12305        } else {
12306            vec![y0, y1, y2, y3]
12307        }))
12308    }
12309
12310    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12311    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12312    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12313    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12314    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12315    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12316    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12317    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12318    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12319    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12320    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
12321    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
12322    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
12323    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
12324    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
12325    pub fn matmul_decode_exact_dual(
12326        &self,
12327        w0: &crate::model::GpuTensor,
12328        w1: &crate::model::GpuTensor,
12329        x: &CudaSlice<f32>,
12330        m: usize,
12331    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12332        use crate::model::GpuTensor;
12333        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12334        let on = *ON.get_or_init(|| {
12335            std::env::var("MEMRA_SPEC_DUAL_T")
12336                .map(|v| v != "0")
12337                .unwrap_or(true)
12338        });
12339        if !on
12340            || !(2..=4).contains(&m)
12341            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12342            || !self.uses_q8_1_fast(w0)
12343            || !self.uses_q8_1_fast(w1)
12344        {
12345            return Ok(None);
12346        }
12347        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
12348        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
12349        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
12350        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
12351        if !self.mmvq_supports(QT_NVFP4) {
12352            return Ok(None);
12353        }
12354        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12355        if w1.in_features() != in_f || w1.out_features() != out_f {
12356            return Ok(None);
12357        }
12358        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12359            (
12360                GpuTensor::Quant {
12361                    bytes: b0,
12362                    qtype: q0,
12363                    row_bytes: rb0,
12364                    scale: s0,
12365                    rp: rp0,
12366                    rp4: None,
12367                    ..
12368                },
12369                GpuTensor::Quant {
12370                    bytes: b1,
12371                    qtype: q1,
12372                    row_bytes: rb1,
12373                    scale: s1,
12374                    rp: rp1,
12375                    rp4: None,
12376                    ..
12377                },
12378            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12379                (b0, b1, *rb0, *s0, *s1, *rp0)
12380            }
12381            _ => return Ok(None),
12382        };
12383        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
12384        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
12385        if std::env::var("MEMRA_DEBUG").is_ok() {
12386            static ONCE: std::sync::Once = std::sync::Once::new();
12387            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12388        }
12389        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12390        let (y0, y1) =
12391            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12392        let mut y0 = y0;
12393        let mut y1 = y1;
12394        if s0 != 1.0 {
12395            self.scale_inplace(&mut y0, s0, m * out_f)?;
12396        }
12397        if s1 != 1.0 {
12398            self.scale_inplace(&mut y1, s1, m * out_f)?;
12399        }
12400        Ok(Some((y0, y1)))
12401    }
12402
12403    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
12404    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
12405    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
12406    /// twins (both buffers must be the repacked layout).
12407    #[allow(clippy::too_many_arguments)]
12408    pub fn qmatvec_batched_dual_raw(
12409        &self,
12410        b0: &CudaSlice<u8>,
12411        b1: &CudaSlice<u8>,
12412        aq: &CudaSlice<i8>,
12413        ad: &CudaSlice<f32>,
12414        m: usize,
12415        in_f: usize,
12416        out_f: usize,
12417        row_bytes: usize,
12418        rp: bool,
12419    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12420        const ROWS_PER_BLOCK: u32 = 4;
12421        let mcols = Self::batched_mcols(m);
12422        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
12423        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
12424        let tiny_rp1 = rp
12425            && mcols == 4
12426            && out_f <= 128
12427            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12428        let (name, rows_per_block) = if tiny_rp1 {
12429            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12430        } else {
12431            match (mcols, rp, m) {
12432                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12433                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12434                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12435                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12436                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12437                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12438                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12439                _ => {
12440                    return Err(
12441                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12442                    );
12443                }
12444            }
12445        };
12446        let f = self.func(name);
12447        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12448        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12449        let cfg = LaunchConfig {
12450            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12451            block_dim: (32, ROWS_PER_BLOCK, 1),
12452            shared_mem_bytes: 0,
12453        };
12454        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12455        let __s_b = self.gpu.stream();
12456        let mut b = __s_b.launch_builder(&f);
12457        b.arg(b0)
12458            .arg(b1)
12459            .arg(aq)
12460            .arg(ad)
12461            .arg(&mut y0)
12462            .arg(&mut y1)
12463            .arg(&inf)
12464            .arg(&outf)
12465            .arg(&mi)
12466            .arg(&rb);
12467        unsafe {
12468            b.launch(cfg)?;
12469        }
12470        Ok((y0, y1))
12471    }
12472
12473    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
12474    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
12475    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
12476    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
12477    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
12478    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
12479    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
12480    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
12481    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
12482    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
12483    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
12484    pub fn matmul_pre_dual_noscale(
12485        &self,
12486        w0: &crate::model::GpuTensor,
12487        w1: &crate::model::GpuTensor,
12488        aq: &CudaSlice<i8>,
12489        ad: &CudaSlice<f32>,
12490        m: usize,
12491    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12492    {
12493        use crate::model::GpuTensor;
12494        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12495            return Ok(None);
12496        }
12497        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
12498        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
12499        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
12500        // would mix dispatch families across the pair — the exact class `q8_fused_params`
12501        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
12502        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
12503        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
12504        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
12505        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
12506        if !self.mmvq_supports(QT_NVFP4) {
12507            return Ok(None);
12508        }
12509        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12510        if w1.in_features() != in_f || w1.out_features() != out_f {
12511            return Ok(None);
12512        }
12513        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
12514        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
12515        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
12516        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
12517        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
12518        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
12519        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
12520        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
12521        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
12522        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
12523        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
12524        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
12525        let no_mirror =
12526            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12527        if self.q8_ffn_fuse2_on()
12528            && no_mirror(w0)
12529            && no_mirror(w1)
12530            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12531        {
12532            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12533            return Ok(Some(((y0, 1.0), (y1, 1.0))));
12534        }
12535        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
12536        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
12537        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
12538        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
12539        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
12540        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
12541        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
12542        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
12543        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
12544        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12545            let (y0, y1) =
12546                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12547            return Ok(Some(((y0, p0.3), (y1, p1.3))));
12548        }
12549        let (b0, q0, rb0, s0, rp0) = match w0 {
12550            GpuTensor::Quant {
12551                bytes,
12552                qtype,
12553                row_bytes,
12554                scale,
12555                rp,
12556                ..
12557            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12558            _ => return Ok(None),
12559        };
12560        let (b1, q1, rb1, s1, rp1) = match w1 {
12561            GpuTensor::Quant {
12562                bytes,
12563                qtype,
12564                row_bytes,
12565                scale,
12566                rp,
12567                ..
12568            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12569            _ => return Ok(None),
12570        };
12571        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12572            return Ok(None);
12573        }
12574        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12575        const RPW: u32 = 2;
12576        let rows_per_block = ROWS_PER_BLOCK * RPW;
12577        let f = self.func(if rp0 {
12578            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12579        } else {
12580            "qmatvec_nvfp4_mmvq_dual_mr2"
12581        });
12582        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12583        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12584        let cfg = LaunchConfig {
12585            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12586            block_dim: (32, ROWS_PER_BLOCK, 1),
12587            shared_mem_bytes: 0,
12588        };
12589        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12590        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
12591        // yscale args stay 1.0 here (they exist for the single-tensor callers).
12592        let one = 1.0f32;
12593        let __s_b = self.gpu.stream();
12594        let mut b = __s_b.launch_builder(&f);
12595        b.arg(b0)
12596            .arg(b1)
12597            .arg(aq)
12598            .arg(ad)
12599            .arg(&mut y0)
12600            .arg(&mut y1)
12601            .arg(&inf)
12602            .arg(&outf)
12603            .arg(&mi)
12604            .arg(&rb)
12605            .arg(&one)
12606            .arg(&one);
12607        unsafe {
12608            b.launch(cfg)?;
12609        }
12610        Ok(Some(((y0, s0), (y1, s1))))
12611    }
12612
12613    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
12614    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
12615    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
12616    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
12617    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
12618    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
12619    /// back to the three singles.
12620    #[allow(clippy::too_many_arguments)]
12621    pub fn matmul_nvfp4_fused3(
12622        &self,
12623        w0: &crate::model::GpuTensor,
12624        w1: &crate::model::GpuTensor,
12625        w2: &crate::model::GpuTensor,
12626        aq: &CudaSlice<i8>,
12627        ad: &CudaSlice<f32>,
12628        m: usize,
12629    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12630    {
12631        use crate::model::GpuTensor;
12632        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
12633        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
12634        // verbatim, weight rows read once for all m columns, bit-identical per
12635        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
12636        // segments would re-read the weight per row" note described the grid.y=m lift,
12637        // which this twin deliberately is NOT.
12638        if !self.mmvq_supports(QT_NVFP4)
12639            || !self.uses_q8_1_fast(w0)
12640            || !self.uses_q8_1_fast(w1)
12641            || !self.uses_q8_1_fast(w2)
12642        {
12643            return Ok(None);
12644        }
12645        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
12646        // door — same family and bit-identity law as the fused4 delegate above.
12647        if (9..=16).contains(&m) {
12648            return Ok(
12649                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12650                    Some(mut ys) => {
12651                        let y2 = ys.pop().unwrap();
12652                        let y1 = ys.pop().unwrap();
12653                        let y0 = ys.pop().unwrap();
12654                        Some((y0, y1, y2))
12655                    }
12656                    None => None,
12657                },
12658            );
12659        }
12660        if !(1..=8).contains(&m) {
12661            return Ok(None);
12662        }
12663        if m > 1 {
12664            let in_f = w0.in_features();
12665            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12666                || !self.batched_supports(QT_NVFP4)
12667                || std::env::var("MEMRA_NO_BATCHED").is_ok()
12668                || (m > 4 && !Self::b8_enabled())
12669                || in_f % 512 != 0
12670                || in_f / 64 > 272
12671            {
12672                return Ok(None);
12673            }
12674        }
12675        let unpack = |w: &crate::model::GpuTensor| match w {
12676            GpuTensor::Quant {
12677                bytes,
12678                qtype,
12679                scale,
12680                rp,
12681                ..
12682            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12683            _ => None,
12684        };
12685        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12686            return Ok(None);
12687        };
12688        let in_f = w0.in_features();
12689        if w1.in_features() != in_f || w2.in_features() != in_f {
12690            return Ok(None);
12691        }
12692        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12693        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12694        const RPW: u32 = 2;
12695        let rows_pb = ROWS_PER_BLOCK * RPW;
12696        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12697        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12698        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12699        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12700        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12701        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12702        // only dereferenced for the launch-arg build inside this call.
12703        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12704        if m > 1 {
12705            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
12706            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12707                return Ok(None);
12708            }
12709            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12710            let cfg = LaunchConfig {
12711                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12712                block_dim: (32, ROWS_PER_BLOCK, 1),
12713                shared_mem_bytes: 0,
12714            };
12715            let __s_b = self.gpu.stream();
12716            let mut b = __s_b.launch_builder(&f);
12717            b.arg(b0)
12718                .arg(b1)
12719                .arg(b2)
12720                .arg(aq)
12721                .arg(ad)
12722                .arg(&mut y0)
12723                .arg(&mut y1)
12724                .arg(&mut y2)
12725                .arg(&inf)
12726                .arg(&oi0)
12727                .arg(&oi1)
12728                .arg(&oi2)
12729                .arg(&mi);
12730            unsafe {
12731                b.launch(cfg)?;
12732            }
12733            return Ok(Some((y0, y1, y2)));
12734        }
12735        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12736        let cfg = LaunchConfig {
12737            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12738            block_dim: (32, ROWS_PER_BLOCK, 1),
12739            shared_mem_bytes: 0,
12740        };
12741        let __s_b = self.gpu.stream();
12742        let mut b = __s_b.launch_builder(&f);
12743        b.arg(b0)
12744            .arg(b1)
12745            .arg(b2)
12746            .arg(aq)
12747            .arg(ad)
12748            .arg(&mut y0)
12749            .arg(&mut y1)
12750            .arg(&mut y2)
12751            .arg(&inf)
12752            .arg(&oi0)
12753            .arg(&oi1)
12754            .arg(&oi2)
12755            .arg(&mi)
12756            .arg(&p0.1)
12757            .arg(&p1.1)
12758            .arg(&p2.1);
12759        unsafe {
12760            b.launch(cfg)?;
12761        }
12762        Ok(Some((y0, y1, y2)))
12763    }
12764
12765    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
12766    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
12767    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
12768    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
12769    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
12770    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
12771    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
12772    /// same-binary interleaved A/B arm.
12773    pub fn matmul_nvfp4_fused2(
12774        &self,
12775        w0: &crate::model::GpuTensor,
12776        w1: &crate::model::GpuTensor,
12777        aq: &CudaSlice<i8>,
12778        ad: &CudaSlice<f32>,
12779        m: usize,
12780    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12781        use crate::model::GpuTensor;
12782        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12783        let off =
12784            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12785        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
12786        // read serves all m rows); the fused segments would re-read the weight per row.
12787        if off
12788            || m != 1
12789            || !self.mmvq_supports(QT_NVFP4)
12790            || !self.uses_q8_1_fast(w0)
12791            || !self.uses_q8_1_fast(w1)
12792        {
12793            return Ok(None);
12794        }
12795        let unpack = |w: &crate::model::GpuTensor| match w {
12796            GpuTensor::Quant {
12797                bytes,
12798                qtype,
12799                scale,
12800                rp,
12801                ..
12802            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12803            _ => None,
12804        };
12805        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12806            return Ok(None);
12807        };
12808        let in_f = w0.in_features();
12809        if w1.in_features() != in_f {
12810            return Ok(None);
12811        }
12812        let (o0, o1) = (w0.out_features(), w1.out_features());
12813        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12814        const RPW: u32 = 2;
12815        let rows_pb = ROWS_PER_BLOCK * RPW;
12816        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12817        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12818        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12819        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12820        let cfg = LaunchConfig {
12821            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12822            block_dim: (32, ROWS_PER_BLOCK, 1),
12823            shared_mem_bytes: 0,
12824        };
12825        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12826        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12827        // only dereferenced for the launch-arg build inside this call.
12828        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12829        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
12830        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
12831        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12832            {
12833                use cudarc::driver::{DevicePtr, DevicePtrMut};
12834                let s = &self.gpu.stream();
12835                let (pw0, _g0) = b0.device_ptr(s);
12836                let (pw1, _g1) = b1.device_ptr(s);
12837                let (paq, _g2) = aq.device_ptr(s);
12838                let (pad, _g3) = ad.device_ptr(s);
12839                let (py0, _g4) = y0.device_ptr_mut(s);
12840                let (py1, _g5) = y1.device_ptr_mut(s);
12841                let (s0, s1) = (p0.1, p1.1);
12842                let mut ps = [
12843                    &pw0 as *const _ as *mut std::ffi::c_void,
12844                    &pw1 as *const _ as *mut _,
12845                    &paq as *const _ as *mut _,
12846                    &pad as *const _ as *mut _,
12847                    &py0 as *const _ as *mut _,
12848                    &py1 as *const _ as *mut _,
12849                    &inf as *const _ as *mut _,
12850                    &oi0 as *const _ as *mut _,
12851                    &oi1 as *const _ as *mut _,
12852                    &mi as *const _ as *mut _,
12853                    &s0 as *const _ as *mut _,
12854                    &s1 as *const _ as *mut _,
12855                ];
12856                unsafe {
12857                    self.launch_pdl(
12858                        "qmatvec_nvfp4_mmvq_fused2_rp",
12859                        cfg.grid_dim,
12860                        cfg.block_dim,
12861                        &mut ps,
12862                    )?;
12863                }
12864            }
12865            return Ok(Some((y0, y1)));
12866        }
12867        let __s_b = self.gpu.stream();
12868        let mut b = __s_b.launch_builder(&f);
12869        b.arg(b0)
12870            .arg(b1)
12871            .arg(aq)
12872            .arg(ad)
12873            .arg(&mut y0)
12874            .arg(&mut y1)
12875            .arg(&inf)
12876            .arg(&oi0)
12877            .arg(&oi1)
12878            .arg(&mi)
12879            .arg(&p0.1)
12880            .arg(&p1.1);
12881        unsafe {
12882            b.launch(cfg)?;
12883        }
12884        Ok(Some((y0, y1)))
12885    }
12886
12887    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
12888    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
12889    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
12890    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
12891    pub fn matmul_nvfp4_fused2_into(
12892        &self,
12893        w0: &crate::model::GpuTensor,
12894        w1: &crate::model::GpuTensor,
12895        aq: &CudaSlice<i8>,
12896        ad: &CudaSlice<f32>,
12897        y0: &mut CudaSlice<f32>,
12898        y1: &mut CudaSlice<f32>,
12899    ) -> Result<bool, Box<dyn std::error::Error>> {
12900        use crate::model::GpuTensor;
12901        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12902        let off =
12903            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12904        if off
12905            || !self.mmvq_supports(QT_NVFP4)
12906            || !self.uses_q8_1_fast(w0)
12907            || !self.uses_q8_1_fast(w1)
12908        {
12909            return Ok(false);
12910        }
12911        let unpack = |w: &crate::model::GpuTensor| match w {
12912            GpuTensor::Quant {
12913                bytes,
12914                qtype,
12915                scale,
12916                rp,
12917                ..
12918            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12919            _ => None,
12920        };
12921        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12922            return Ok(false);
12923        };
12924        let in_f = w0.in_features();
12925        if w1.in_features() != in_f {
12926            return Ok(false);
12927        }
12928        let (o0, o1) = (w0.out_features(), w1.out_features());
12929        if y0.len() < o0 || y1.len() < o1 {
12930            return Ok(false);
12931        }
12932        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12933        const RPW: u32 = 2;
12934        let rows_pb = ROWS_PER_BLOCK * RPW;
12935        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12936        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12937        let cfg = LaunchConfig {
12938            grid_dim: (nb(o0) + nb(o1), 1, 1),
12939            block_dim: (32, ROWS_PER_BLOCK, 1),
12940            shared_mem_bytes: 0,
12941        };
12942        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
12943        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12944        // only dereferenced for the launch-arg build inside this call.
12945        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12946        let __s_b = self.gpu.stream();
12947        let mut b = __s_b.launch_builder(&f);
12948        b.arg(b0)
12949            .arg(b1)
12950            .arg(aq)
12951            .arg(ad)
12952            .arg(&mut *y0)
12953            .arg(&mut *y1)
12954            .arg(&inf)
12955            .arg(&oi0)
12956            .arg(&oi1)
12957            .arg(&mi)
12958            .arg(&p0.1)
12959            .arg(&p1.1);
12960        unsafe {
12961            b.launch(cfg)?;
12962        }
12963        Ok(true)
12964    }
12965
12966    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
12967    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
12968    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
12969    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
12970    #[allow(clippy::type_complexity)]
12971    pub fn matmul_nvfp4_fused4(
12972        &self,
12973        w0: &crate::model::GpuTensor,
12974        w1: &crate::model::GpuTensor,
12975        w2: &crate::model::GpuTensor,
12976        w3: &crate::model::GpuTensor,
12977        aq: &CudaSlice<i8>,
12978        ad: &CudaSlice<f32>,
12979        m: usize,
12980    ) -> Result<
12981        Option<(
12982            CudaSlice<f32>,
12983            CudaSlice<f32>,
12984            CudaSlice<f32>,
12985            CudaSlice<f32>,
12986        )>,
12987        Box<dyn std::error::Error>,
12988    > {
12989        use crate::model::GpuTensor;
12990        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
12991        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
12992        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
12993        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
12994        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
12995        // Admission mirrors the singles' batched gates below.
12996        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
12997            || !self.mmvq_supports(QT_NVFP4)
12998            || !self.uses_q8_1_fast(w0)
12999            || !self.uses_q8_1_fast(w1)
13000            || !self.uses_q8_1_fast(w2)
13001            || !self.uses_q8_1_fast(w3)
13002        {
13003            return Ok(None);
13004        }
13005        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
13006        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
13007        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
13008        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
13009        if (9..=16).contains(&m) {
13010            return Ok(
13011                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13012                    Some(mut ys) => {
13013                        let y3 = ys.pop().unwrap();
13014                        let y2 = ys.pop().unwrap();
13015                        let y1 = ys.pop().unwrap();
13016                        let y0 = ys.pop().unwrap();
13017                        Some((y0, y1, y2, y3))
13018                    }
13019                    None => None,
13020                },
13021            );
13022        }
13023        if !(1..=8).contains(&m) {
13024            return Ok(None);
13025        }
13026        if m > 1 {
13027            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
13028            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
13029            let in_f = w0.in_features();
13030            if !self.batched_supports(QT_NVFP4)
13031                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13032                || (m > 4 && !Self::b8_enabled())
13033                || in_f % 512 != 0
13034                || in_f / 64 > 272
13035            {
13036                return Ok(None);
13037            }
13038        }
13039        let unpack = |w: &crate::model::GpuTensor| match w {
13040            GpuTensor::Quant {
13041                bytes,
13042                qtype,
13043                scale,
13044                rp,
13045                ..
13046            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13047            _ => None,
13048        };
13049        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13050            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13051        else {
13052            return Ok(None);
13053        };
13054        let in_f = w0.in_features();
13055        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13056            return Ok(None);
13057        }
13058        let (o0, o1, o2, o3) = (
13059            w0.out_features(),
13060            w1.out_features(),
13061            w2.out_features(),
13062            w3.out_features(),
13063        );
13064        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13065        const RPW: u32 = 2;
13066        let rows_pb = ROWS_PER_BLOCK * RPW;
13067        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13068        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13069        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13070        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13071        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13072        let (inf, oi0, oi1, oi2, oi3, mi) = (
13073            in_f as i32,
13074            o0 as i32,
13075            o1 as i32,
13076            o2 as i32,
13077            o3 as i32,
13078            m as i32,
13079        );
13080        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13081        // only dereferenced for the launch-arg build inside this call.
13082        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13083        if m > 1 {
13084            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
13085            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
13086            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13087                return Ok(None);
13088            }
13089            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13090            let cfg = LaunchConfig {
13091                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13092                block_dim: (32, ROWS_PER_BLOCK, 1),
13093                shared_mem_bytes: 0,
13094            };
13095            let __s_b = self.gpu.stream();
13096            let mut b = __s_b.launch_builder(&f);
13097            b.arg(b0)
13098                .arg(b1)
13099                .arg(b2)
13100                .arg(b3)
13101                .arg(aq)
13102                .arg(ad)
13103                .arg(&mut y0)
13104                .arg(&mut y1)
13105                .arg(&mut y2)
13106                .arg(&mut y3)
13107                .arg(&inf)
13108                .arg(&oi0)
13109                .arg(&oi1)
13110                .arg(&oi2)
13111                .arg(&oi3)
13112                .arg(&mi);
13113            unsafe {
13114                b.launch(cfg)?;
13115            }
13116            return Ok(Some((y0, y1, y2, y3)));
13117        }
13118        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13119        let cfg = LaunchConfig {
13120            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13121            block_dim: (32, ROWS_PER_BLOCK, 1),
13122            shared_mem_bytes: 0,
13123        };
13124        let __s_b = self.gpu.stream();
13125        let mut b = __s_b.launch_builder(&f);
13126        b.arg(b0)
13127            .arg(b1)
13128            .arg(b2)
13129            .arg(b3)
13130            .arg(aq)
13131            .arg(ad)
13132            .arg(&mut y0)
13133            .arg(&mut y1)
13134            .arg(&mut y2)
13135            .arg(&mut y3)
13136            .arg(&inf)
13137            .arg(&oi0)
13138            .arg(&oi1)
13139            .arg(&oi2)
13140            .arg(&oi3)
13141            .arg(&mi)
13142            .arg(&p0.1)
13143            .arg(&p1.1)
13144            .arg(&p2.1)
13145            .arg(&p3.1);
13146        unsafe {
13147            b.launch(cfg)?;
13148        }
13149        Ok(Some((y0, y1, y2, y3)))
13150    }
13151
13152    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
13153    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
13154    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
13155    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
13156    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
13157    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
13158    /// back to the per-tensor path.
13159    pub fn matmul_q8_fused2(
13160        &self,
13161        w0: &crate::model::GpuTensor,
13162        w1: &crate::model::GpuTensor,
13163        aq: &CudaSlice<i8>,
13164        ad: &CudaSlice<f32>,
13165    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13166        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
13167        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
13168        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
13169        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
13170        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
13171        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13172            return Ok(Some(self.e4m3_fused2_core(
13173                p0.0,
13174                p1.0,
13175                aq,
13176                ad,
13177                w0.in_features(),
13178                p0.1,
13179                p1.1,
13180                p0.2,
13181                p0.3,
13182                p1.3,
13183            )?));
13184        }
13185        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13186            return Ok(None);
13187        };
13188        Ok(Some(self.q8_fused2_core(
13189            p0.0,
13190            p1.0,
13191            aq,
13192            ad,
13193            w0.in_features(),
13194            p0.1,
13195            p1.1,
13196            p0.2,
13197        )?))
13198    }
13199
13200    #[allow(clippy::too_many_arguments)]
13201    fn q8_fused2_core(
13202        &self,
13203        b0: &CudaSlice<u8>,
13204        b1: &CudaSlice<u8>,
13205        aq: &CudaSlice<i8>,
13206        ad: &CudaSlice<f32>,
13207        in_f: usize,
13208        out0: usize,
13209        out1: usize,
13210        row_bytes: usize,
13211    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13212        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13213        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13214        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13215        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13216        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13217        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13218        let cfg = LaunchConfig {
13219            grid_dim: (nb0 + nb1, 1, 1),
13220            block_dim: (32, ROWS_PER_BLOCK, 1),
13221            shared_mem_bytes: 0,
13222        };
13223        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13224        let __s_b = self.gpu.stream();
13225        let mut b = __s_b.launch_builder(&f);
13226        b.arg(b0)
13227            .arg(b1)
13228            .arg(aq)
13229            .arg(ad)
13230            .arg(&mut y0)
13231            .arg(&mut y1)
13232            .arg(&inf)
13233            .arg(&o0)
13234            .arg(&o1)
13235            .arg(&rbl);
13236        unsafe {
13237            b.launch(cfg)?;
13238        }
13239        Ok((y0, y1))
13240    }
13241
13242    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13243    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13244    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13245    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13246    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13247    pub fn matmul_q8_fused2_x(
13248        &self,
13249        w0: &crate::model::GpuTensor,
13250        w1: &crate::model::GpuTensor,
13251        x: &CudaSlice<f32>,
13252    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13253        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13254            return Ok(None);
13255        }
13256        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13257            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13258            return Ok(Some(self.e4m3_fused2_core(
13259                p0.0,
13260                p1.0,
13261                &aq,
13262                &ad,
13263                w0.in_features(),
13264                p0.1,
13265                p1.1,
13266                p0.2,
13267                p0.3,
13268                p1.3,
13269            )?));
13270        }
13271        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13272            return Ok(None);
13273        };
13274        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13275        Ok(Some(self.q8_fused2_core(
13276            p0.0,
13277            p1.0,
13278            &aq,
13279            &ad,
13280            w0.in_features(),
13281            p0.1,
13282            p1.1,
13283            p0.2,
13284        )?))
13285    }
13286
13287    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13288    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13289    #[allow(clippy::too_many_arguments)]
13290    pub fn qmatvec_q8_fused2_raw(
13291        &self,
13292        b0: &CudaSlice<u8>,
13293        b1: &CudaSlice<u8>,
13294        x: &CudaSlice<f32>,
13295        in_f: usize,
13296        out0: usize,
13297        out1: usize,
13298        row_bytes: usize,
13299    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13300        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13301        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13302    }
13303
13304    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13305    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13306    /// (tensor,row) to three separate m=1 MMVQ launches.
13307    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13308    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13309    pub fn matmul_q4_fused3(
13310        &self,
13311        w0: &crate::model::GpuTensor,
13312        w1: &crate::model::GpuTensor,
13313        w2: &crate::model::GpuTensor,
13314        aq: &CudaSlice<i8>,
13315        ad: &CudaSlice<f32>,
13316    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13317    {
13318        use crate::model::GpuTensor;
13319        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13320            match w {
13321                GpuTensor::Quant {
13322                    qtype, row_bytes, ..
13323                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13324                _ => None,
13325            }
13326        };
13327        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13328            return Ok(None);
13329        };
13330        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13331            return Ok(None);
13332        }
13333        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
13334        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
13335        // the separate matvecs (each routes its own rp).
13336        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13337            match w {
13338                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13339                    Some(m) => (m, true),
13340                    None => (bytes, *rp),
13341                },
13342                _ => unreachable!(),
13343            }
13344        }
13345        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13346        if rp0 != rp1 || rp1 != rp2 {
13347            return Ok(None);
13348        }
13349        let rp = rp0;
13350        let rpb: u32 = 4;
13351        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
13352        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
13353        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
13354        let mr1 = rp && Self::q40_mr1_on();
13355        let nb = |o: usize| {
13356            if mr1 {
13357                (o as u32).div_ceil(rpb)
13358            } else {
13359                (o as u32).div_ceil(2).div_ceil(rpb)
13360            }
13361        };
13362        let grid = nb(o0) + nb(o1) + nb(o2);
13363        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13364        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13365        let mut y2 = self.alloc_uninit::<f32>(o2)?;
13366        let f = self.func(if mr1 {
13367            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13368        } else if rp {
13369            "qmatvec_q4_0_mmvq_fused3_rp"
13370        } else {
13371            "qmatvec_q4_0_mmvq_fused3"
13372        });
13373        let cfg = LaunchConfig {
13374            grid_dim: (grid, 1, 1),
13375            block_dim: (32, rpb, 1),
13376            shared_mem_bytes: 0,
13377        };
13378        let inf = w0.in_features() as i32;
13379        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13380        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13381        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
13382        // variant may take the programmatic-serialization launch.
13383        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13384            {
13385                use cudarc::driver::{DevicePtr, DevicePtrMut};
13386                let s = &self.gpu.stream();
13387                let (p0, _g0) = b0.device_ptr(s);
13388                let (p1, _g1) = b1.device_ptr(s);
13389                let (p2, _g2) = b2.device_ptr(s);
13390                let (paq, _g3) = aq.device_ptr(s);
13391                let (pad, _g4) = ad.device_ptr(s);
13392                let (py0, _g5) = y0.device_ptr_mut(s);
13393                let (py1, _g6) = y1.device_ptr_mut(s);
13394                let (py2, _g7) = y2.device_ptr_mut(s);
13395                let mut ps = [
13396                    &p0 as *const _ as *mut std::ffi::c_void,
13397                    &p1 as *const _ as *mut _,
13398                    &p2 as *const _ as *mut _,
13399                    &paq as *const _ as *mut _,
13400                    &pad as *const _ as *mut _,
13401                    &py0 as *const _ as *mut _,
13402                    &py1 as *const _ as *mut _,
13403                    &py2 as *const _ as *mut _,
13404                    &inf as *const _ as *mut _,
13405                    &oo0 as *const _ as *mut _,
13406                    &oo1 as *const _ as *mut _,
13407                    &oo2 as *const _ as *mut _,
13408                    &r0 as *const _ as *mut _,
13409                    &r1 as *const _ as *mut _,
13410                    &r2 as *const _ as *mut _,
13411                ];
13412                unsafe {
13413                    self.launch_pdl(
13414                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13415                        (grid, 1, 1),
13416                        (32, rpb, 1),
13417                        &mut ps,
13418                    )?;
13419                }
13420            }
13421            return Ok(Some((y0, y1, y2)));
13422        }
13423        let __s_b = self.gpu.stream();
13424        let mut b = __s_b.launch_builder(&f);
13425        b.arg(b0)
13426            .arg(b1)
13427            .arg(b2)
13428            .arg(aq)
13429            .arg(ad)
13430            .arg(&mut y0)
13431            .arg(&mut y1)
13432            .arg(&mut y2)
13433            .arg(&inf)
13434            .arg(&oo0)
13435            .arg(&oo1)
13436            .arg(&oo2)
13437            .arg(&r0)
13438            .arg(&r1)
13439            .arg(&r2);
13440        unsafe {
13441            b.launch(cfg)?;
13442        }
13443        Ok(Some((y0, y1, y2)))
13444    }
13445
13446    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13447    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
13448    #[allow(clippy::too_many_arguments)]
13449    pub fn matmul_q4_fused3_into(
13450        &self,
13451        w0: &crate::model::GpuTensor,
13452        w1: &crate::model::GpuTensor,
13453        w2: &crate::model::GpuTensor,
13454        aq: &CudaSlice<i8>,
13455        ad: &CudaSlice<f32>,
13456        y0: &mut CudaSlice<f32>,
13457        y1: &mut CudaSlice<f32>,
13458        y2: &mut CudaSlice<f32>,
13459    ) -> Result<bool, Box<dyn std::error::Error>> {
13460        use crate::model::GpuTensor;
13461        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13462            match w {
13463                GpuTensor::Quant {
13464                    qtype, row_bytes, ..
13465                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13466                _ => None,
13467            }
13468        };
13469        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13470            return Ok(false);
13471        };
13472        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13473            return Ok(false);
13474        }
13475        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13476            match w {
13477                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13478                    Some(m) => (m, true),
13479                    None => (bytes, *rp),
13480                },
13481                _ => unreachable!(),
13482            }
13483        }
13484        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13485        if rp0 != rp1 || rp1 != rp2 {
13486            return Ok(false);
13487        }
13488        let rp = rp0;
13489        let rpb: u32 = 4;
13490        let mr1 = rp && Self::q40_mr1_on();
13491        let nb = |o: usize| {
13492            if mr1 {
13493                (o as u32).div_ceil(rpb)
13494            } else {
13495                (o as u32).div_ceil(2).div_ceil(rpb)
13496            }
13497        };
13498        let grid = nb(o0) + nb(o1) + nb(o2);
13499        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13500        let f = self.func(if mr1 {
13501            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13502        } else if rp {
13503            "qmatvec_q4_0_mmvq_fused3_rp"
13504        } else {
13505            "qmatvec_q4_0_mmvq_fused3"
13506        });
13507        let cfg = LaunchConfig {
13508            grid_dim: (grid, 1, 1),
13509            block_dim: (32, rpb, 1),
13510            shared_mem_bytes: 0,
13511        };
13512        let inf = w0.in_features() as i32;
13513        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13514        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13515        // PDL wave-A: identical to the owned twin (capture-lane parity).
13516        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13517            use cudarc::driver::{DevicePtr, DevicePtrMut};
13518            let s = &self.gpu.stream();
13519            let (p0, _g0) = b0.device_ptr(s);
13520            let (p1, _g1) = b1.device_ptr(s);
13521            let (p2, _g2) = b2.device_ptr(s);
13522            let (paq, _g3) = aq.device_ptr(s);
13523            let (pad, _g4) = ad.device_ptr(s);
13524            let (py0, _g5) = y0.device_ptr_mut(s);
13525            let (py1, _g6) = y1.device_ptr_mut(s);
13526            let (py2, _g7) = y2.device_ptr_mut(s);
13527            let mut ps = [
13528                &p0 as *const _ as *mut std::ffi::c_void,
13529                &p1 as *const _ as *mut _,
13530                &p2 as *const _ as *mut _,
13531                &paq as *const _ as *mut _,
13532                &pad as *const _ as *mut _,
13533                &py0 as *const _ as *mut _,
13534                &py1 as *const _ as *mut _,
13535                &py2 as *const _ as *mut _,
13536                &inf as *const _ as *mut _,
13537                &oo0 as *const _ as *mut _,
13538                &oo1 as *const _ as *mut _,
13539                &oo2 as *const _ as *mut _,
13540                &r0 as *const _ as *mut _,
13541                &r1 as *const _ as *mut _,
13542                &r2 as *const _ as *mut _,
13543            ];
13544            unsafe {
13545                self.launch_pdl(
13546                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13547                    (grid, 1, 1),
13548                    (32, rpb, 1),
13549                    &mut ps,
13550                )?;
13551            }
13552            return Ok(true);
13553        }
13554        let __s_b = self.gpu.stream();
13555        let mut b = __s_b.launch_builder(&f);
13556        b.arg(b0)
13557            .arg(b1)
13558            .arg(b2)
13559            .arg(aq)
13560            .arg(ad)
13561            .arg(&mut *y0)
13562            .arg(&mut *y1)
13563            .arg(&mut *y2)
13564            .arg(&inf)
13565            .arg(&oo0)
13566            .arg(&oo1)
13567            .arg(&oo2)
13568            .arg(&r0)
13569            .arg(&r1)
13570            .arg(&r2);
13571        unsafe {
13572            b.launch(cfg)?;
13573        }
13574        Ok(true)
13575    }
13576
13577    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
13578    pub fn matmul_q4_fused2(
13579        &self,
13580        w0: &crate::model::GpuTensor,
13581        w1: &crate::model::GpuTensor,
13582        aq: &CudaSlice<i8>,
13583        ad: &CudaSlice<f32>,
13584    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13585        use crate::model::GpuTensor;
13586        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13587            match w {
13588                GpuTensor::Quant {
13589                    qtype, row_bytes, ..
13590                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13591                _ => None,
13592            }
13593        };
13594        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13595            return Ok(None);
13596        };
13597        if w0.in_features() != w1.in_features() {
13598            return Ok(None);
13599        }
13600        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
13601        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13602            match w {
13603                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13604                    Some(m) => (m, true),
13605                    None => (bytes, *rp),
13606                },
13607                _ => unreachable!(),
13608            }
13609        }
13610        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13611        if rp0 != rp1 {
13612            return Ok(None);
13613        }
13614        let rp = rp0;
13615        let rpb: u32 = 4;
13616        // mr1 twin — see matmul_q4_fused3.
13617        let mr1 = rp && Self::q40_mr1_on();
13618        let nb = |o: usize| {
13619            if mr1 {
13620                (o as u32).div_ceil(rpb)
13621            } else {
13622                (o as u32).div_ceil(2).div_ceil(rpb)
13623            }
13624        };
13625        let grid = nb(o0) + nb(o1);
13626        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13627        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13628        let f = self.func(if mr1 {
13629            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13630        } else if rp {
13631            "qmatvec_q4_0_mmvq_fused2_rp"
13632        } else {
13633            "qmatvec_q4_0_mmvq_fused2"
13634        });
13635        let cfg = LaunchConfig {
13636            grid_dim: (grid, 1, 1),
13637            block_dim: (32, rpb, 1),
13638            shared_mem_bytes: 0,
13639        };
13640        let inf = w0.in_features() as i32;
13641        let (oo0, oo1) = (o0 as i32, o1 as i32);
13642        let (r0, r1) = (rb0 as i64, rb1 as i64);
13643        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
13644        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13645            {
13646                use cudarc::driver::{DevicePtr, DevicePtrMut};
13647                let s = &self.gpu.stream();
13648                let (p0, _g0) = b0.device_ptr(s);
13649                let (p1, _g1) = b1.device_ptr(s);
13650                let (paq, _g2) = aq.device_ptr(s);
13651                let (pad, _g3) = ad.device_ptr(s);
13652                let (py0, _g4) = y0.device_ptr_mut(s);
13653                let (py1, _g5) = y1.device_ptr_mut(s);
13654                let mut ps = [
13655                    &p0 as *const _ as *mut std::ffi::c_void,
13656                    &p1 as *const _ as *mut _,
13657                    &paq as *const _ as *mut _,
13658                    &pad as *const _ as *mut _,
13659                    &py0 as *const _ as *mut _,
13660                    &py1 as *const _ as *mut _,
13661                    &inf as *const _ as *mut _,
13662                    &oo0 as *const _ as *mut _,
13663                    &oo1 as *const _ as *mut _,
13664                    &r0 as *const _ as *mut _,
13665                    &r1 as *const _ as *mut _,
13666                ];
13667                unsafe {
13668                    self.launch_pdl(
13669                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13670                        (grid, 1, 1),
13671                        (32, rpb, 1),
13672                        &mut ps,
13673                    )?;
13674                }
13675            }
13676            return Ok(Some((y0, y1)));
13677        }
13678        let __s_b = self.gpu.stream();
13679        let mut b = __s_b.launch_builder(&f);
13680        b.arg(b0)
13681            .arg(b1)
13682            .arg(aq)
13683            .arg(ad)
13684            .arg(&mut y0)
13685            .arg(&mut y1)
13686            .arg(&inf)
13687            .arg(&oo0)
13688            .arg(&oo1)
13689            .arg(&r0)
13690            .arg(&r1);
13691        unsafe {
13692            b.launch(cfg)?;
13693        }
13694        Ok(Some((y0, y1)))
13695    }
13696
13697    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13698    pub fn matmul_q4_fused2_into(
13699        &self,
13700        w0: &crate::model::GpuTensor,
13701        w1: &crate::model::GpuTensor,
13702        aq: &CudaSlice<i8>,
13703        ad: &CudaSlice<f32>,
13704        y0: &mut CudaSlice<f32>,
13705        y1: &mut CudaSlice<f32>,
13706    ) -> Result<bool, Box<dyn std::error::Error>> {
13707        use crate::model::GpuTensor;
13708        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13709            match w {
13710                GpuTensor::Quant {
13711                    qtype, row_bytes, ..
13712                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13713                _ => None,
13714            }
13715        };
13716        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13717            return Ok(false);
13718        };
13719        if w0.in_features() != w1.in_features() {
13720            return Ok(false);
13721        }
13722        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13723            match w {
13724                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13725                    Some(m) => (m, true),
13726                    None => (bytes, *rp),
13727                },
13728                _ => unreachable!(),
13729            }
13730        }
13731        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13732        if rp0 != rp1 {
13733            return Ok(false);
13734        }
13735        let rp = rp0;
13736        let rpb: u32 = 4;
13737        let mr1 = rp && Self::q40_mr1_on();
13738        let nb = |o: usize| {
13739            if mr1 {
13740                (o as u32).div_ceil(rpb)
13741            } else {
13742                (o as u32).div_ceil(2).div_ceil(rpb)
13743            }
13744        };
13745        let grid = nb(o0) + nb(o1);
13746        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13747        let f = self.func(if mr1 {
13748            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13749        } else if rp {
13750            "qmatvec_q4_0_mmvq_fused2_rp"
13751        } else {
13752            "qmatvec_q4_0_mmvq_fused2"
13753        });
13754        let cfg = LaunchConfig {
13755            grid_dim: (grid, 1, 1),
13756            block_dim: (32, rpb, 1),
13757            shared_mem_bytes: 0,
13758        };
13759        let inf = w0.in_features() as i32;
13760        let (oo0, oo1) = (o0 as i32, o1 as i32);
13761        let (r0, r1) = (rb0 as i64, rb1 as i64);
13762        // PDL wave-A: identical to the owned twin (capture-lane parity).
13763        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13764            use cudarc::driver::{DevicePtr, DevicePtrMut};
13765            let s = &self.gpu.stream();
13766            let (p0, _g0) = b0.device_ptr(s);
13767            let (p1, _g1) = b1.device_ptr(s);
13768            let (paq, _g2) = aq.device_ptr(s);
13769            let (pad, _g3) = ad.device_ptr(s);
13770            let (py0, _g4) = y0.device_ptr_mut(s);
13771            let (py1, _g5) = y1.device_ptr_mut(s);
13772            let mut ps = [
13773                &p0 as *const _ as *mut std::ffi::c_void,
13774                &p1 as *const _ as *mut _,
13775                &paq as *const _ as *mut _,
13776                &pad as *const _ as *mut _,
13777                &py0 as *const _ as *mut _,
13778                &py1 as *const _ as *mut _,
13779                &inf as *const _ as *mut _,
13780                &oo0 as *const _ as *mut _,
13781                &oo1 as *const _ as *mut _,
13782                &r0 as *const _ as *mut _,
13783                &r1 as *const _ as *mut _,
13784            ];
13785            unsafe {
13786                self.launch_pdl(
13787                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13788                    (grid, 1, 1),
13789                    (32, rpb, 1),
13790                    &mut ps,
13791                )?;
13792            }
13793            return Ok(true);
13794        }
13795        let __s_b = self.gpu.stream();
13796        let mut b = __s_b.launch_builder(&f);
13797        b.arg(b0)
13798            .arg(b1)
13799            .arg(aq)
13800            .arg(ad)
13801            .arg(&mut *y0)
13802            .arg(&mut *y1)
13803            .arg(&inf)
13804            .arg(&oo0)
13805            .arg(&oo1)
13806            .arg(&r0)
13807            .arg(&r1);
13808        unsafe {
13809            b.launch(cfg)?;
13810        }
13811        Ok(true)
13812    }
13813
13814    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
13815    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
13816    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
13817    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
13818    pub fn matmul_q4_fused2_batched(
13819        &self,
13820        w0: &crate::model::GpuTensor,
13821        w1: &crate::model::GpuTensor,
13822        aq: &CudaSlice<i8>,
13823        ad: &CudaSlice<f32>,
13824        m: usize,
13825    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13826        use crate::model::GpuTensor;
13827        if m < 2 || m > 8 {
13828            return Ok(None);
13829        }
13830        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13831            match w {
13832                GpuTensor::Quant {
13833                    qtype, row_bytes, ..
13834                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13835                _ => None,
13836            }
13837        };
13838        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
13839            return Ok(None);
13840        };
13841        if w0.in_features() != w1.in_features() {
13842            return Ok(None);
13843        }
13844        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13845            match w {
13846                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13847                    Some(mr) => (mr, true),
13848                    None => (bytes, *rp),
13849                },
13850                _ => unreachable!(),
13851            }
13852        }
13853        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13854        if !rp0 || !rp1 {
13855            return Ok(None);
13856        }
13857        let mcols = Self::batched_mcols(m);
13858        let rpb: u32 = 4;
13859        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13860        let grid = nb(o0) + nb(o1);
13861        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13862        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13863        let f = self.func(match mcols {
13864            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
13865            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
13866            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
13867        });
13868        let cfg = LaunchConfig {
13869            grid_dim: (grid, 1, 1),
13870            block_dim: (32, rpb, 1),
13871            shared_mem_bytes: 0,
13872        };
13873        let inf = w0.in_features() as i32;
13874        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
13875        let rb = rb0 as i64;
13876        let __s_b = self.gpu.stream();
13877        let mut b = __s_b.launch_builder(&f);
13878        b.arg(b0)
13879            .arg(b1)
13880            .arg(aq)
13881            .arg(ad)
13882            .arg(&mut y0)
13883            .arg(&mut y1)
13884            .arg(&inf)
13885            .arg(&oo0)
13886            .arg(&oo1)
13887            .arg(&mi)
13888            .arg(&rb);
13889        unsafe {
13890            b.launch(cfg)?;
13891        }
13892        Ok(Some((y0, y1)))
13893    }
13894
13895    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
13896    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
13897    #[allow(clippy::too_many_arguments)]
13898    pub fn matmul_q4_fused3_batched(
13899        &self,
13900        w0: &crate::model::GpuTensor,
13901        w1: &crate::model::GpuTensor,
13902        w2: &crate::model::GpuTensor,
13903        aq: &CudaSlice<i8>,
13904        ad: &CudaSlice<f32>,
13905        m: usize,
13906    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13907    {
13908        use crate::model::GpuTensor;
13909        if m < 2 || m > 8 {
13910            return Ok(None);
13911        }
13912        let q4 = |w: &GpuTensor| -> Option<usize> {
13913            match w {
13914                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
13915                _ => None,
13916            }
13917        };
13918        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
13919            return Ok(None);
13920        };
13921        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13922            return Ok(None);
13923        }
13924        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13925            match w {
13926                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13927                    Some(mr) => (mr, true),
13928                    None => (bytes, *rp),
13929                },
13930                _ => unreachable!(),
13931            }
13932        }
13933        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13934        if !rp0 || !rp1 || !rp2 {
13935            return Ok(None);
13936        }
13937        let mcols = Self::batched_mcols(m);
13938        let rpb: u32 = 4;
13939        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13940        let grid = nb(o0) + nb(o1) + nb(o2);
13941        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13942        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13943        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13944        let f = self.func(match mcols {
13945            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
13946            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
13947            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
13948        });
13949        let cfg = LaunchConfig {
13950            grid_dim: (grid, 1, 1),
13951            block_dim: (32, rpb, 1),
13952            shared_mem_bytes: 0,
13953        };
13954        let inf = w0.in_features() as i32;
13955        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
13956        let rb = 0i64;
13957        let __s_b = self.gpu.stream();
13958        let mut b = __s_b.launch_builder(&f);
13959        b.arg(b0)
13960            .arg(b1)
13961            .arg(b2)
13962            .arg(aq)
13963            .arg(ad)
13964            .arg(&mut y0)
13965            .arg(&mut y1)
13966            .arg(&mut y2)
13967            .arg(&inf)
13968            .arg(&oo0)
13969            .arg(&oo1)
13970            .arg(&oo2)
13971            .arg(&mi)
13972            .arg(&rb);
13973        unsafe {
13974            b.launch(cfg)?;
13975        }
13976        Ok(Some((y0, y1, y2)))
13977    }
13978
13979    pub fn matmul_q8_fused3(
13980        &self,
13981        w0: &crate::model::GpuTensor,
13982        w1: &crate::model::GpuTensor,
13983        w2: &crate::model::GpuTensor,
13984        aq: &CudaSlice<i8>,
13985        ad: &CudaSlice<f32>,
13986    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13987    {
13988        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
13989        // are per-tensor FP8, so native residency without this arm meant three separate launches.
13990        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13991            return Ok(Some(self.e4m3_fused3_core(
13992                p0.0,
13993                p1.0,
13994                p2.0,
13995                aq,
13996                ad,
13997                w0.in_features(),
13998                p0.1,
13999                p1.1,
14000                p2.1,
14001                p0.2,
14002                p0.3,
14003                p1.3,
14004                p2.3,
14005            )?));
14006        }
14007        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14008            return Ok(None);
14009        };
14010        Ok(Some(self.q8_fused3_core(
14011            p0.0,
14012            p1.0,
14013            p2.0,
14014            aq,
14015            ad,
14016            w0.in_features(),
14017            p0.1,
14018            p1.1,
14019            p2.1,
14020            p0.2,
14021        )?))
14022    }
14023
14024    #[allow(clippy::too_many_arguments)]
14025    fn q8_fused3_core(
14026        &self,
14027        b0: &CudaSlice<u8>,
14028        b1: &CudaSlice<u8>,
14029        b2: &CudaSlice<u8>,
14030        aq: &CudaSlice<i8>,
14031        ad: &CudaSlice<f32>,
14032        in_f: usize,
14033        out0: usize,
14034        out1: usize,
14035        out2: usize,
14036        row_bytes: usize,
14037    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14038        const ROWS_PER_BLOCK: u32 = 4;
14039        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14040        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14041        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14042        let f = self.func("qmatvec_q8_0_mmvq_fused3");
14043        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14044        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14045        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14046        let cfg = LaunchConfig {
14047            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14048            block_dim: (32, ROWS_PER_BLOCK, 1),
14049            shared_mem_bytes: 0,
14050        };
14051        let (inf, o0, o1, o2, rbl) = (
14052            in_f as i32,
14053            out0 as i32,
14054            out1 as i32,
14055            out2 as i32,
14056            row_bytes as i64,
14057        );
14058        let __s_b = self.gpu.stream();
14059        let mut b = __s_b.launch_builder(&f);
14060        b.arg(b0)
14061            .arg(b1)
14062            .arg(b2)
14063            .arg(aq)
14064            .arg(ad)
14065            .arg(&mut y0)
14066            .arg(&mut y1)
14067            .arg(&mut y2)
14068            .arg(&inf)
14069            .arg(&o0)
14070            .arg(&o1)
14071            .arg(&o2)
14072            .arg(&rbl);
14073        unsafe {
14074            b.launch(cfg)?;
14075        }
14076        Ok((y0, y1, y2))
14077    }
14078
14079    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
14080    #[allow(clippy::too_many_arguments)]
14081    pub fn qmatvec_q8_fused3_raw(
14082        &self,
14083        b0: &CudaSlice<u8>,
14084        b1: &CudaSlice<u8>,
14085        b2: &CudaSlice<u8>,
14086        x: &CudaSlice<f32>,
14087        in_f: usize,
14088        out0: usize,
14089        out1: usize,
14090        out2: usize,
14091        row_bytes: usize,
14092    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14093        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14094        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14095    }
14096
14097    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
14098    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
14099    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
14100    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
14101    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
14102    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
14103    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
14104    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
14105    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
14106    /// twin must not introduce a batched program the reference path would not run).
14107    pub fn matmul_q8_fused2_t(
14108        &self,
14109        w0: &crate::model::GpuTensor,
14110        w1: &crate::model::GpuTensor,
14111        aq: &CudaSlice<i8>,
14112        ad: &CudaSlice<f32>,
14113        m: usize,
14114    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14115        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
14116        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
14117        // fuses too — same template body, still bit-identical to the two _b8 launches.
14118        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14119            return Ok(None);
14120        }
14121        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
14122        // so the fused b8 launch would introduce a batched program the reference path would not run.
14123        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14124            if m > 4 && !Self::b8_enabled() {
14125                return Ok(None);
14126            }
14127            return Ok(Some(self.e4m3_fused2_t_core(
14128                p0.0,
14129                p1.0,
14130                aq,
14131                ad,
14132                m,
14133                w0.in_features(),
14134                p0.1,
14135                p1.1,
14136                p0.2,
14137                p0.3,
14138                p1.3,
14139            )?));
14140        }
14141        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14142            return Ok(None);
14143        };
14144        Ok(Some(self.q8_fused2_t_core(
14145            p0.0,
14146            p1.0,
14147            aq,
14148            ad,
14149            m,
14150            w0.in_features(),
14151            p0.1,
14152            p1.1,
14153            p0.2,
14154        )?))
14155    }
14156
14157    #[allow(clippy::too_many_arguments)]
14158    fn q8_fused2_t_core(
14159        &self,
14160        b0: &CudaSlice<u8>,
14161        b1: &CudaSlice<u8>,
14162        aq: &CudaSlice<i8>,
14163        ad: &CudaSlice<f32>,
14164        m: usize,
14165        in_f: usize,
14166        out0: usize,
14167        out1: usize,
14168        row_bytes: usize,
14169    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14170        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14171        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14172        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14173        let f = self.func(match Self::batched_mcols(m) {
14174            2 => "qmatvec_q8_0_mmvq_fused2_b2",
14175            4 => "qmatvec_q8_0_mmvq_fused2_b4",
14176            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
14177            _ => "qmatvec_q8_0_mmvq_fused2_b8",
14178        });
14179        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14180        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14181        let cfg = LaunchConfig {
14182            grid_dim: (nb0 + nb1, 1, 1),
14183            block_dim: (32, ROWS_PER_BLOCK, 1),
14184            shared_mem_bytes: 0,
14185        };
14186        let (inf, o0, o1, mi, rbl) = (
14187            in_f as i32,
14188            out0 as i32,
14189            out1 as i32,
14190            m as i32,
14191            row_bytes as i64,
14192        );
14193        let __s_b = self.gpu.stream();
14194        let mut b = __s_b.launch_builder(&f);
14195        b.arg(b0)
14196            .arg(b1)
14197            .arg(aq)
14198            .arg(ad)
14199            .arg(&mut y0)
14200            .arg(&mut y1)
14201            .arg(&inf)
14202            .arg(&o0)
14203            .arg(&o1)
14204            .arg(&mi)
14205            .arg(&rbl);
14206        unsafe {
14207            b.launch(cfg)?;
14208        }
14209        Ok((y0, y1))
14210    }
14211
14212    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14213    /// q8_1 quant of the [m, in_f] activation), no env gating.
14214    #[allow(clippy::too_many_arguments)]
14215    pub fn qmatvec_q8_fused2_t_raw(
14216        &self,
14217        b0: &CudaSlice<u8>,
14218        b1: &CudaSlice<u8>,
14219        x: &CudaSlice<f32>,
14220        m: usize,
14221        in_f: usize,
14222        out0: usize,
14223        out1: usize,
14224        row_bytes: usize,
14225    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14226        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14227        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14228    }
14229
14230    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14231    /// `matmul_q8_fused2_t` with three ranges.
14232    #[allow(clippy::too_many_arguments)]
14233    pub fn matmul_q8_fused3_t(
14234        &self,
14235        w0: &crate::model::GpuTensor,
14236        w1: &crate::model::GpuTensor,
14237        w2: &crate::model::GpuTensor,
14238        aq: &CudaSlice<i8>,
14239        ad: &CudaSlice<f32>,
14240        m: usize,
14241    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14242    {
14243        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14244            return Ok(None);
14245        }
14246        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14247            return Ok(Some(self.e4m3_fused3_t_core(
14248                p0.0,
14249                p1.0,
14250                p2.0,
14251                aq,
14252                ad,
14253                m,
14254                w0.in_features(),
14255                p0.1,
14256                p1.1,
14257                p2.1,
14258                p0.2,
14259                p0.3,
14260                p1.3,
14261                p2.3,
14262            )?));
14263        }
14264        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14265            return Ok(None);
14266        };
14267        Ok(Some(self.q8_fused3_t_core(
14268            p0.0,
14269            p1.0,
14270            p2.0,
14271            aq,
14272            ad,
14273            m,
14274            w0.in_features(),
14275            p0.1,
14276            p1.1,
14277            p2.1,
14278            p0.2,
14279        )?))
14280    }
14281
14282    #[allow(clippy::too_many_arguments)]
14283    fn q8_fused3_t_core(
14284        &self,
14285        b0: &CudaSlice<u8>,
14286        b1: &CudaSlice<u8>,
14287        b2: &CudaSlice<u8>,
14288        aq: &CudaSlice<i8>,
14289        ad: &CudaSlice<f32>,
14290        m: usize,
14291        in_f: usize,
14292        out0: usize,
14293        out1: usize,
14294        out2: usize,
14295        row_bytes: usize,
14296    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14297        const ROWS_PER_BLOCK: u32 = 4;
14298        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14299        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14300        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14301        let f = self.func(if Self::batched_mcols(m) == 2 {
14302            "qmatvec_q8_0_mmvq_fused3_b2"
14303        } else {
14304            "qmatvec_q8_0_mmvq_fused3_b4"
14305        });
14306        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14307        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14308        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14309        let cfg = LaunchConfig {
14310            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14311            block_dim: (32, ROWS_PER_BLOCK, 1),
14312            shared_mem_bytes: 0,
14313        };
14314        let (inf, o0, o1, o2, mi, rbl) = (
14315            in_f as i32,
14316            out0 as i32,
14317            out1 as i32,
14318            out2 as i32,
14319            m as i32,
14320            row_bytes as i64,
14321        );
14322        let __s_b = self.gpu.stream();
14323        let mut b = __s_b.launch_builder(&f);
14324        b.arg(b0)
14325            .arg(b1)
14326            .arg(b2)
14327            .arg(aq)
14328            .arg(ad)
14329            .arg(&mut y0)
14330            .arg(&mut y1)
14331            .arg(&mut y2)
14332            .arg(&inf)
14333            .arg(&o0)
14334            .arg(&o1)
14335            .arg(&o2)
14336            .arg(&mi)
14337            .arg(&rbl);
14338        unsafe {
14339            b.launch(cfg)?;
14340        }
14341        Ok((y0, y1, y2))
14342    }
14343
14344    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
14345    #[allow(clippy::too_many_arguments)]
14346    pub fn qmatvec_q8_fused3_t_raw(
14347        &self,
14348        b0: &CudaSlice<u8>,
14349        b1: &CudaSlice<u8>,
14350        b2: &CudaSlice<u8>,
14351        x: &CudaSlice<f32>,
14352        m: usize,
14353        in_f: usize,
14354        out0: usize,
14355        out1: usize,
14356        out2: usize,
14357        row_bytes: usize,
14358    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14359        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14360        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14361    }
14362
14363    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
14364    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
14365    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
14366    pub fn q8_ffn_fuse2_on(&self) -> bool {
14367        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14368        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14369    }
14370
14371    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
14372    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
14373    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
14374    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
14375    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
14376    #[allow(clippy::type_complexity)]
14377    fn q8_fused_params<'w, const N: usize>(
14378        &self,
14379        ws: &[&'w crate::model::GpuTensor; N],
14380    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14381        use crate::model::GpuTensor;
14382        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14383            return None;
14384        }
14385        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14386            return None;
14387        }
14388        let in_f = ws[0].in_features();
14389        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14390        for (i, w) in ws.iter().enumerate() {
14391            match w {
14392                GpuTensor::Quant {
14393                    bytes,
14394                    qtype,
14395                    row_bytes,
14396                    scale,
14397                    ..
14398                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14399                    out[i] = Some((bytes, w.out_features(), *row_bytes))
14400                }
14401                _ => return None,
14402            }
14403        }
14404        Some(out.map(|o| o.unwrap()))
14405    }
14406
14407    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
14408    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
14409    pub fn e4m3_dual_on(&self) -> bool {
14410        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14411        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14412    }
14413
14414    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
14415    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
14416    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
14417    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
14418    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
14419    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
14420    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
14421    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
14422    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
14423    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
14424    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
14425    #[allow(clippy::type_complexity)]
14426    fn e4m3_fused_params<'w, const N: usize>(
14427        &self,
14428        ws: &[&'w crate::model::GpuTensor; N],
14429    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14430        use crate::model::GpuTensor;
14431        if !self.e4m3_dual_on() {
14432            return None;
14433        }
14434        let in_f = ws[0].in_features();
14435        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14436        for (i, w) in ws.iter().enumerate() {
14437            match w {
14438                GpuTensor::Quant {
14439                    bytes,
14440                    qtype,
14441                    row_bytes,
14442                    scale,
14443                    rp,
14444                    rp4,
14445                    ..
14446                } if *qtype == QT_F8_E4M3
14447                    && w.in_features() == in_f
14448                    && *row_bytes == in_f
14449                    && !*rp
14450                    && rp4.is_none() =>
14451                {
14452                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14453                }
14454                _ => return None,
14455            }
14456        }
14457        Some(out.map(|o| o.unwrap()))
14458    }
14459
14460    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
14461    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
14462    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
14463    #[allow(clippy::too_many_arguments)]
14464    fn e4m3_fused2_core(
14465        &self,
14466        b0: &CudaSlice<u8>,
14467        b1: &CudaSlice<u8>,
14468        aq: &CudaSlice<i8>,
14469        ad: &CudaSlice<f32>,
14470        in_f: usize,
14471        out0: usize,
14472        out1: usize,
14473        row_bytes: usize,
14474        ws0: f32,
14475        ws1: f32,
14476    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14477        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14478        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14479        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14480        let f = self.func("qmatvec_e4m3_mmvq_fused2");
14481        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14482        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14483        let cfg = LaunchConfig {
14484            grid_dim: (nb0 + nb1, 1, 1),
14485            block_dim: (32, ROWS_PER_BLOCK, 1),
14486            shared_mem_bytes: 0,
14487        };
14488        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14489        let __s_b = self.gpu.stream();
14490        let mut b = __s_b.launch_builder(&f);
14491        b.arg(b0)
14492            .arg(b1)
14493            .arg(aq)
14494            .arg(ad)
14495            .arg(&mut y0)
14496            .arg(&mut y1)
14497            .arg(&inf)
14498            .arg(&o0)
14499            .arg(&o1)
14500            .arg(&rbl)
14501            .arg(&ws0)
14502            .arg(&ws1);
14503        unsafe {
14504            b.launch(cfg)?;
14505        }
14506        Ok((y0, y1))
14507    }
14508
14509    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
14510    #[allow(clippy::too_many_arguments)]
14511    fn e4m3_fused3_core(
14512        &self,
14513        b0: &CudaSlice<u8>,
14514        b1: &CudaSlice<u8>,
14515        b2: &CudaSlice<u8>,
14516        aq: &CudaSlice<i8>,
14517        ad: &CudaSlice<f32>,
14518        in_f: usize,
14519        out0: usize,
14520        out1: usize,
14521        out2: usize,
14522        row_bytes: usize,
14523        ws0: f32,
14524        ws1: f32,
14525        ws2: f32,
14526    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14527        const ROWS_PER_BLOCK: u32 = 4;
14528        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14529        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14530        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14531        let f = self.func("qmatvec_e4m3_mmvq_fused3");
14532        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14533        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14534        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14535        let cfg = LaunchConfig {
14536            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14537            block_dim: (32, ROWS_PER_BLOCK, 1),
14538            shared_mem_bytes: 0,
14539        };
14540        let (inf, o0, o1, o2, rbl) = (
14541            in_f as i32,
14542            out0 as i32,
14543            out1 as i32,
14544            out2 as i32,
14545            row_bytes as i64,
14546        );
14547        let __s_b = self.gpu.stream();
14548        let mut b = __s_b.launch_builder(&f);
14549        b.arg(b0)
14550            .arg(b1)
14551            .arg(b2)
14552            .arg(aq)
14553            .arg(ad)
14554            .arg(&mut y0)
14555            .arg(&mut y1)
14556            .arg(&mut y2)
14557            .arg(&inf)
14558            .arg(&o0)
14559            .arg(&o1)
14560            .arg(&o2)
14561            .arg(&rbl)
14562            .arg(&ws0)
14563            .arg(&ws1)
14564            .arg(&ws2);
14565        unsafe {
14566            b.launch(cfg)?;
14567        }
14568        Ok((y0, y1, y2))
14569    }
14570
14571    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
14572    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
14573    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
14574    #[allow(clippy::too_many_arguments)]
14575    fn e4m3_fused2_t_core(
14576        &self,
14577        b0: &CudaSlice<u8>,
14578        b1: &CudaSlice<u8>,
14579        aq: &CudaSlice<i8>,
14580        ad: &CudaSlice<f32>,
14581        m: usize,
14582        in_f: usize,
14583        out0: usize,
14584        out1: usize,
14585        row_bytes: usize,
14586        ws0: f32,
14587        ws1: f32,
14588    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14589        const ROWS_PER_BLOCK: u32 = 4;
14590        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14591        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14592        let f = self.func(match Self::batched_mcols(m) {
14593            2 => "qmatvec_e4m3_mmvq_fused2_b2",
14594            4 => "qmatvec_e4m3_mmvq_fused2_b4",
14595            _ => "qmatvec_e4m3_mmvq_fused2_b8",
14596        });
14597        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14598        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14599        let cfg = LaunchConfig {
14600            grid_dim: (nb0 + nb1, 1, 1),
14601            block_dim: (32, ROWS_PER_BLOCK, 1),
14602            shared_mem_bytes: 0,
14603        };
14604        let (inf, o0, o1, mi, rbl) = (
14605            in_f as i32,
14606            out0 as i32,
14607            out1 as i32,
14608            m as i32,
14609            row_bytes as i64,
14610        );
14611        let __s_b = self.gpu.stream();
14612        let mut b = __s_b.launch_builder(&f);
14613        b.arg(b0)
14614            .arg(b1)
14615            .arg(aq)
14616            .arg(ad)
14617            .arg(&mut y0)
14618            .arg(&mut y1)
14619            .arg(&inf)
14620            .arg(&o0)
14621            .arg(&o1)
14622            .arg(&mi)
14623            .arg(&rbl);
14624        unsafe {
14625            b.launch(cfg)?;
14626        }
14627        if ws0 != 1.0 {
14628            self.scale_inplace(&mut y0, ws0, m * out0)?;
14629        }
14630        if ws1 != 1.0 {
14631            self.scale_inplace(&mut y1, ws1, m * out1)?;
14632        }
14633        Ok((y0, y1))
14634    }
14635
14636    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
14637    #[allow(clippy::too_many_arguments)]
14638    fn e4m3_fused3_t_core(
14639        &self,
14640        b0: &CudaSlice<u8>,
14641        b1: &CudaSlice<u8>,
14642        b2: &CudaSlice<u8>,
14643        aq: &CudaSlice<i8>,
14644        ad: &CudaSlice<f32>,
14645        m: usize,
14646        in_f: usize,
14647        out0: usize,
14648        out1: usize,
14649        out2: usize,
14650        row_bytes: usize,
14651        ws0: f32,
14652        ws1: f32,
14653        ws2: f32,
14654    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14655        const ROWS_PER_BLOCK: u32 = 4;
14656        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14657        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14658        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14659        let f = self.func(if Self::batched_mcols(m) == 2 {
14660            "qmatvec_e4m3_mmvq_fused3_b2"
14661        } else {
14662            "qmatvec_e4m3_mmvq_fused3_b4"
14663        });
14664        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14665        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14666        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14667        let cfg = LaunchConfig {
14668            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14669            block_dim: (32, ROWS_PER_BLOCK, 1),
14670            shared_mem_bytes: 0,
14671        };
14672        let (inf, o0, o1, o2, mi, rbl) = (
14673            in_f as i32,
14674            out0 as i32,
14675            out1 as i32,
14676            out2 as i32,
14677            m as i32,
14678            row_bytes as i64,
14679        );
14680        let __s_b = self.gpu.stream();
14681        let mut b = __s_b.launch_builder(&f);
14682        b.arg(b0)
14683            .arg(b1)
14684            .arg(b2)
14685            .arg(aq)
14686            .arg(ad)
14687            .arg(&mut y0)
14688            .arg(&mut y1)
14689            .arg(&mut y2)
14690            .arg(&inf)
14691            .arg(&o0)
14692            .arg(&o1)
14693            .arg(&o2)
14694            .arg(&mi)
14695            .arg(&rbl);
14696        unsafe {
14697            b.launch(cfg)?;
14698        }
14699        if ws0 != 1.0 {
14700            self.scale_inplace(&mut y0, ws0, m * out0)?;
14701        }
14702        if ws1 != 1.0 {
14703            self.scale_inplace(&mut y1, ws1, m * out1)?;
14704        }
14705        if ws2 != 1.0 {
14706            self.scale_inplace(&mut y2, ws2, m * out2)?;
14707        }
14708        Ok((y0, y1, y2))
14709    }
14710
14711    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
14712    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
14713    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
14714    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
14715    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
14716    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
14717    ///
14718    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
14719    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
14720    pub fn qmatvec_e4m3_blk_mmvq(
14721        &self,
14722        bytes: &CudaSlice<u8>,
14723        aq: &CudaSlice<i8>,
14724        ad: &CudaSlice<f32>,
14725        scales: &CudaSlice<f32>,
14726        m: usize,
14727        in_f: usize,
14728        out_f: usize,
14729        row_bytes: usize,
14730        scale_cols: usize,
14731    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14732        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
14733        self.qmatvec_e4m3_blk_mmvq_into(
14734            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14735        )?;
14736        Ok(y)
14737    }
14738
14739    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
14740    #[allow(clippy::too_many_arguments)]
14741    pub fn qmatvec_e4m3_blk_mmvq_into(
14742        &self,
14743        bytes: &CudaSlice<u8>,
14744        aq: &CudaSlice<i8>,
14745        ad: &CudaSlice<f32>,
14746        scales: &CudaSlice<f32>,
14747        m: usize,
14748        in_f: usize,
14749        out_f: usize,
14750        row_bytes: usize,
14751        scale_cols: usize,
14752        y: &mut CudaSlice<f32>,
14753    ) -> Result<(), Box<dyn std::error::Error>> {
14754        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14755        let f = self.func("qmatvec_e4m3_blk_mmvq");
14756        let cfg = LaunchConfig {
14757            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14758            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
14759            shared_mem_bytes: 0,                // warp-only reduce
14760        };
14761        let (inf, outf, mi, rb, sc) = (
14762            in_f as i32,
14763            out_f as i32,
14764            m as i32,
14765            row_bytes as i64,
14766            scale_cols as i32,
14767        );
14768        let __s_b = self.gpu.stream();
14769        let mut b = __s_b.launch_builder(&f);
14770        b.arg(bytes)
14771            .arg(aq)
14772            .arg(ad)
14773            .arg(scales)
14774            .arg(&mut *y)
14775            .arg(&inf)
14776            .arg(&outf)
14777            .arg(&mi)
14778            .arg(&rb)
14779            .arg(&sc);
14780        unsafe {
14781            b.launch(cfg)?;
14782        }
14783        Ok(())
14784    }
14785
14786    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
14787    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
14788    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
14789    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
14790    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
14791    #[allow(clippy::too_many_arguments)]
14792    pub fn qmatvec_e4m3_blk_mmvq_batched(
14793        &self,
14794        bytes: &CudaSlice<u8>,
14795        aq: &CudaSlice<i8>,
14796        ad: &CudaSlice<f32>,
14797        scales: &CudaSlice<f32>,
14798        m: usize,
14799        in_f: usize,
14800        out_f: usize,
14801        row_bytes: usize,
14802        scale_cols: usize,
14803        mcols: usize,
14804    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14805        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14806        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14807        let name = match mcols {
14808            2 => "qmatvec_e4m3_blk_mmvq_b2",
14809            4 => "qmatvec_e4m3_blk_mmvq_b4",
14810            8 => "qmatvec_e4m3_blk_mmvq_b8",
14811            16 => "qmatvec_e4m3_blk_mmvq_b16",
14812            _ => {
14813                return Err(
14814                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14815                );
14816            }
14817        };
14818        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14819        let f = self.func(name);
14820        let cfg = LaunchConfig {
14821            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14822            block_dim: (32, ROWS_PER_BLOCK, 1),
14823            shared_mem_bytes: 0,
14824        };
14825        let (inf, outf, mi, rb, sc) = (
14826            in_f as i32,
14827            out_f as i32,
14828            m as i32,
14829            row_bytes as i64,
14830            scale_cols as i32,
14831        );
14832        let __s_b = self.gpu.stream();
14833        let mut b = __s_b.launch_builder(&f);
14834        b.arg(bytes)
14835            .arg(aq)
14836            .arg(ad)
14837            .arg(scales)
14838            .arg(&mut y)
14839            .arg(&inf)
14840            .arg(&outf)
14841            .arg(&mi)
14842            .arg(&rb)
14843            .arg(&sc);
14844        unsafe {
14845            b.launch(cfg)?;
14846        }
14847        Ok(y)
14848    }
14849
14850    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
14851    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
14852    #[allow(clippy::too_many_arguments)]
14853    pub fn qmatvec_e4m3_blk_batched_raw(
14854        &self,
14855        bytes: &CudaSlice<u8>,
14856        x: &CudaSlice<f32>,
14857        scales: &CudaSlice<f32>,
14858        m: usize,
14859        in_f: usize,
14860        out_f: usize,
14861        row_bytes: usize,
14862        scale_cols: usize,
14863        mcols: usize,
14864    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14865        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14866        self.qmatvec_e4m3_blk_mmvq_batched(
14867            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
14868        )
14869    }
14870
14871    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
14872    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
14873    #[allow(clippy::too_many_arguments)]
14874    pub fn qmatvec_e4m3_blk_mmvq_raw(
14875        &self,
14876        bytes: &CudaSlice<u8>,
14877        x: &CudaSlice<f32>,
14878        scales: &CudaSlice<f32>,
14879        m: usize,
14880        in_f: usize,
14881        out_f: usize,
14882        row_bytes: usize,
14883        scale_cols: usize,
14884    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14885        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14886        self.qmatvec_e4m3_blk_mmvq(
14887            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
14888        )
14889    }
14890
14891    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
14892    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
14893    #[allow(clippy::too_many_arguments)]
14894    pub fn qmatvec_e4m3_fused2_raw(
14895        &self,
14896        b0: &CudaSlice<u8>,
14897        b1: &CudaSlice<u8>,
14898        x: &CudaSlice<f32>,
14899        in_f: usize,
14900        out0: usize,
14901        out1: usize,
14902        row_bytes: usize,
14903        ws0: f32,
14904        ws1: f32,
14905    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14906        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14907        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
14908    }
14909
14910    #[allow(clippy::too_many_arguments)]
14911    pub fn qmatvec_e4m3_fused3_raw(
14912        &self,
14913        b0: &CudaSlice<u8>,
14914        b1: &CudaSlice<u8>,
14915        b2: &CudaSlice<u8>,
14916        x: &CudaSlice<f32>,
14917        in_f: usize,
14918        out0: usize,
14919        out1: usize,
14920        out2: usize,
14921        row_bytes: usize,
14922        ws0: f32,
14923        ws1: f32,
14924        ws2: f32,
14925    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14926        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14927        self.e4m3_fused3_core(
14928            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14929        )
14930    }
14931
14932    #[allow(clippy::too_many_arguments)]
14933    pub fn qmatvec_e4m3_fused2_t_raw(
14934        &self,
14935        b0: &CudaSlice<u8>,
14936        b1: &CudaSlice<u8>,
14937        x: &CudaSlice<f32>,
14938        m: usize,
14939        in_f: usize,
14940        out0: usize,
14941        out1: usize,
14942        row_bytes: usize,
14943        ws0: f32,
14944        ws1: f32,
14945    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14946        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14947        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
14948    }
14949
14950    #[allow(clippy::too_many_arguments)]
14951    pub fn qmatvec_e4m3_fused3_t_raw(
14952        &self,
14953        b0: &CudaSlice<u8>,
14954        b1: &CudaSlice<u8>,
14955        b2: &CudaSlice<u8>,
14956        x: &CudaSlice<f32>,
14957        m: usize,
14958        in_f: usize,
14959        out0: usize,
14960        out1: usize,
14961        out2: usize,
14962        row_bytes: usize,
14963        ws0: f32,
14964        ws1: f32,
14965        ws2: f32,
14966    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14967        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14968        self.e4m3_fused3_t_core(
14969            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14970        )
14971    }
14972
14973    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
14974    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
14975    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
14976    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
14977    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
14978    ///
14979    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
14980    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
14981    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
14982    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
14983    fn try_e4m3_blk_pre(
14984        &self,
14985        w: &crate::model::GpuTensor,
14986        aq: &CudaSlice<i8>,
14987        ad: &CudaSlice<f32>,
14988        m: usize,
14989    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14990        use crate::model::GpuTensor;
14991        if let GpuTensor::Quant {
14992            bytes,
14993            qtype,
14994            row_bytes,
14995            blk: Some(g),
14996            ..
14997        } = w
14998        {
14999            if *qtype == QT_F8_E4M3_BLK {
15000                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
15001                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
15002                // below, so the decode-exactness contract is preserved at every width. Gated by
15003                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
15004                // one rollback door covers every dtype's batched tier.
15005                if (2..=16).contains(&m)
15006                    && std::env::var("MEMRA_NO_BATCHED").is_err()
15007                    && (m <= 4 || Self::b8_enabled())
15008                {
15009                    let mcols = Self::batched_mcols(m);
15010                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15011                        bytes,
15012                        aq,
15013                        ad,
15014                        &g.scales,
15015                        m,
15016                        w.in_features(),
15017                        w.out_features(),
15018                        *row_bytes,
15019                        g.cols,
15020                        mcols,
15021                    )?));
15022                }
15023                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15024                    bytes,
15025                    aq,
15026                    ad,
15027                    &g.scales,
15028                    m,
15029                    w.in_features(),
15030                    w.out_features(),
15031                    *row_bytes,
15032                    g.cols,
15033                )?));
15034            }
15035        }
15036        Ok(None)
15037    }
15038
15039    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
15040    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
15041    ///
15042    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
15043    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
15044    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
15045    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
15046    /// prefill keeps the floor's arithmetic and the floor's kernels.
15047    ///
15048    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
15049    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
15050    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
15051    /// (projection, prefill call) and frees immediately.
15052    ///
15053    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
15054    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
15055    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
15056    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
15057    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
15058    /// single-variable comparison instead of a two-variable one.
15059    ///
15060    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
15061    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
15062    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
15063    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
15064    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
15065    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
15066    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
15067    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
15068    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
15069    ///
15070    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
15071    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
15072    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
15073    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
15074    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
15075    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
15076    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
15077    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
15078    /// because v2's denominator had its slab already resident while this class's floor must build it
15079    /// every call; same tile, opposite sign, because the question changed.
15080    ///
15081    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
15082    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
15083    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
15084    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
15085    fn try_e4m3_blk_prefill(
15086        &self,
15087        w: &crate::model::GpuTensor,
15088        x: &CudaSlice<f32>,
15089        m: usize,
15090    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15091        use crate::model::GpuTensor;
15092        let GpuTensor::Quant {
15093            bytes,
15094            qtype,
15095            blk: Some(g),
15096            ..
15097        } = w
15098        else {
15099            return Ok(None);
15100        };
15101        if *qtype != QT_F8_E4M3_BLK {
15102            return Ok(None);
15103        }
15104        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
15105        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
15106        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
15107        // through to the dequant below when they do, never silently produce nothing.
15108        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15109            return Ok(Some(y));
15110        }
15111        let (in_f, out_f) = (w.in_features(), w.out_features());
15112        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15113        let tmp = GpuTensor::Quant {
15114            bytes: slab,
15115            qtype: QT_Q8_0,
15116            row_bytes: in_f / 32 * 34,
15117            ne: vec![in_f as u64, out_f as u64],
15118            scale: 1.0,
15119            rp: false,
15120            #[cfg(memra_cutlass)]
15121            cutlass: None,
15122            fp8: None,
15123            blk: None,
15124            f16: None,
15125            rp4: None,
15126        };
15127        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
15128        Ok(Some(self.matmul(&tmp, x, m)?))
15129    }
15130
15131    pub fn matmul_pre_noscale(
15132        &self,
15133        w: &crate::model::GpuTensor,
15134        aq: &CudaSlice<i8>,
15135        ad: &CudaSlice<f32>,
15136        m: usize,
15137    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15138        use crate::model::GpuTensor;
15139        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
15140        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
15141        // rather than let the tail below refuse and cost the caller a re-dispatch.
15142        if m == 1 {
15143            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15144                return Ok(Some((y, 1.0)));
15145            }
15146        }
15147        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
15148        if m != 1 || !self.uses_q8_1_fast(w) {
15149            return Ok(None);
15150        }
15151        let in_f = w.in_features();
15152        let out_f = w.out_features();
15153        let (bytes, qtype, row_bytes, scale, rp) = match w {
15154            GpuTensor::Quant {
15155                bytes,
15156                qtype,
15157                row_bytes,
15158                scale,
15159                rp,
15160                ..
15161            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15162            _ => return Ok(None),
15163        };
15164        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
15165        if self.mmvq_supports(qtype) {
15166            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
15167            let (mbytes, mrp) = match w {
15168                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15169                _ => (bytes, rp),
15170            };
15171            let y = self.qmatvec_mmvq(
15172                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
15173            )?;
15174            return Ok(Some((y, scale)));
15175        }
15176        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
15177        let name = match qtype {
15178            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15179            QT_Q4_K => "qmatvec_q4_K_dp4a",
15180            QT_Q6_K => "qmatvec_q6_K_dp4a",
15181            QT_Q5_K => "qmatvec_q5_K_dp4a",
15182            QT_Q3_K => "qmatvec_q3_K_dp4a",
15183            QT_NVFP4 => {
15184                if rp {
15185                    "qmatvec_nvfp4_dp4a_rp"
15186                } else {
15187                    "qmatvec_nvfp4_dp4a"
15188                }
15189            }
15190            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15191            _ => return Ok(None),
15192        };
15193        let f = self.func(name);
15194        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15195        let cfg = LaunchConfig {
15196            grid_dim: (out_f as u32, m as u32, 1),
15197            block_dim: (128, 1, 1),
15198            shared_mem_bytes: 0,
15199        };
15200        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15201        let __s_b = self.gpu.stream();
15202        let mut b = __s_b.launch_builder(&f);
15203        b.arg(bytes)
15204            .arg(aq)
15205            .arg(ad)
15206            .arg(&mut y)
15207            .arg(&inf)
15208            .arg(&outf)
15209            .arg(&mi)
15210            .arg(&rb);
15211        unsafe {
15212            b.launch(cfg)?;
15213        }
15214        Ok(Some((y, scale)))
15215    }
15216
15217    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15218    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15219    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15220        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15221        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15222        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15223        // is a pure function of the dtype — the decode-parity law holds under every env.
15224        if qtype == QT_F8_E4M3 {
15225            return true;
15226        }
15227        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15228            return false;
15229        }
15230        matches!(
15231            qtype,
15232            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15233        )
15234    }
15235
15236    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15237    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15238    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15239    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15240    pub fn qmatvec_mmvq(
15241        &self,
15242        bytes: &CudaSlice<u8>,
15243        aq: &CudaSlice<i8>,
15244        ad: &CudaSlice<f32>,
15245        m: usize,
15246        in_f: usize,
15247        out_f: usize,
15248        qtype: i32,
15249        row_bytes: usize,
15250        scale: f32,
15251        rp: bool,
15252    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15253        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15254        self.qmatvec_mmvq_into(
15255            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15256        )?;
15257        Ok(y)
15258    }
15259
15260    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15261    #[allow(clippy::too_many_arguments)]
15262    pub fn qmatvec_mmvq_into(
15263        &self,
15264        bytes: &CudaSlice<u8>,
15265        aq: &CudaSlice<i8>,
15266        ad: &CudaSlice<f32>,
15267        m: usize,
15268        in_f: usize,
15269        out_f: usize,
15270        qtype: i32,
15271        row_bytes: usize,
15272        scale: f32,
15273        rp: bool,
15274        y: &mut CudaSlice<f32>,
15275    ) -> Result<(), Box<dyn std::error::Error>> {
15276        debug_assert!(y.len() >= m * out_f);
15277        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15278        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15279        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15280        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15281        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15282        if qtype == QT_Q8_0
15283            && rp
15284            && m == 1
15285            && out_f >= 64
15286            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15287            && {
15288                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15289                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15290            }
15291        {
15292            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15293            let cfg = LaunchConfig {
15294                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15295                block_dim: (32, 2, 1),
15296                shared_mem_bytes: 0,
15297            };
15298            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15299            let __s_b = self.gpu.stream();
15300            let mut b = __s_b.launch_builder(&f);
15301            b.arg(bytes)
15302                .arg(aq)
15303                .arg(ad)
15304                .arg(&mut *y)
15305                .arg(&inf)
15306                .arg(&outf)
15307                .arg(&mi)
15308                .arg(&rb);
15309            unsafe {
15310                b.launch(cfg)?;
15311            }
15312            if scale != 1.0 {
15313                self.scale_inplace(y, scale, out_f)?;
15314            }
15315            return Ok(());
15316        }
15317        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15318        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15319        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15320        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
15321        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
15322        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
15323        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
15324        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
15325        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15326            2
15327        } else {
15328            1
15329        };
15330        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
15331        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
15332        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
15333        // valid-window interleaved, bit-identical per row — same dot program).
15334        if m == 1 && qtype == QT_Q4_0 {
15335            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15336            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
15337            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
15338            mr = *Q40MR.get_or_init(|| {
15339                std::env::var("MEMRA_Q40_MR")
15340                    .ok()
15341                    .and_then(|v| v.parse().ok())
15342                    .unwrap_or(1)
15343            });
15344        }
15345        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
15346        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
15347        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
15348        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
15349        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
15350        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
15351        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
15352        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
15353        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
15354        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
15355        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15356        let q5_force = q5_mode.as_deref() == Some("2");
15357        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
15358        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
15359        let q5_il = qtype == QT_Q5_K
15360            && m == 1
15361            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15362        if q5_il && !q5_force && out_f > 65536 {
15363            mr = 1;
15364        }
15365        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
15366        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
15367        if qtype == QT_Q4_0 && rp && mr != 1 {
15368            mr = 2;
15369        }
15370        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
15371        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
15372        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
15373        if qtype == QT_Q8_0 && rp {
15374            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15375            mr = *Q80MR.get_or_init(|| {
15376                std::env::var("MEMRA_Q80_MR")
15377                    .ok()
15378                    .and_then(|v| v.parse().ok())
15379                    .unwrap_or(1)
15380            });
15381        }
15382        let name = match (qtype, mr, rp) {
15383            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15384            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15385            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15386            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15387            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15388            (QT_Q5_K, 2, _) => {
15389                if q5_il {
15390                    "qmatvec_q5_K_mmvq_mr2_il"
15391                } else {
15392                    "qmatvec_q5_K_mmvq_mr2"
15393                }
15394            }
15395            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15396            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
15397            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
15398            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
15399            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
15400            (QT_Q8_0, _, true)
15401                if in_f % 1024 == 0 && {
15402                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15403                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15404                } =>
15405            {
15406                "qmatvec_q8_0_mmvq_rpca"
15407            }
15408            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15409            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15410            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
15411            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
15412            // reach a GGUF-layout kernel or vice versa.
15413            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15414            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15415            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15416            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15417            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15418            (QT_Q5_K, _, _) => {
15419                if q5_il {
15420                    "qmatvec_q5_K_mmvq_il"
15421                } else {
15422                    "qmatvec_q5_K_mmvq"
15423                }
15424            }
15425            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15426            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15427            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15428            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15429        };
15430        let f = self.func(name);
15431        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
15432        let rows_per_block = ROWS_PER_BLOCK * mr;
15433        let cfg = LaunchConfig {
15434            grid_dim: (
15435                (out_f as u32 + rows_per_block - 1) / rows_per_block,
15436                m as u32,
15437                1,
15438            ),
15439            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
15440            shared_mem_bytes: 0,                // warp-only reduce at m=1
15441        };
15442        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15443        let __s_b = self.gpu.stream();
15444        let mut b = __s_b.launch_builder(&f);
15445        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
15446        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
15447        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
15448        // weight_scale). Other mmvq kernels keep the 8-arg signature.
15449        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15450            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
15451            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
15452            if Self::pdl_on()
15453                && Self::pdl_mmvq_on()
15454                && Self::pdl_nvfp4q8_on()
15455                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15456            {
15457                use cudarc::driver::{DevicePtr, DevicePtrMut};
15458                let s = &self.gpu.stream();
15459                let (pw, _g0) = bytes.device_ptr(s);
15460                let (paq, _g1) = aq.device_ptr(s);
15461                let (pad, _g2) = ad.device_ptr(s);
15462                let (py, _g3) = y.device_ptr_mut(s);
15463                let mut ps = [
15464                    &pw as *const _ as *mut std::ffi::c_void,
15465                    &paq as *const _ as *mut _,
15466                    &pad as *const _ as *mut _,
15467                    &py as *const _ as *mut _,
15468                    &inf as *const _ as *mut _,
15469                    &outf as *const _ as *mut _,
15470                    &mi as *const _ as *mut _,
15471                    &rb as *const _ as *mut _,
15472                    &scale as *const _ as *mut _,
15473                ];
15474                unsafe {
15475                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15476                }
15477                return Ok(());
15478            }
15479            b.arg(bytes)
15480                .arg(aq)
15481                .arg(ad)
15482                .arg(&mut *y)
15483                .arg(&inf)
15484                .arg(&outf)
15485                .arg(&mi)
15486                .arg(&rb)
15487                .arg(&scale);
15488            unsafe {
15489                b.launch(cfg)?;
15490            }
15491        } else if Self::pdl_on()
15492            && Self::pdl_mmvq_on()
15493            && (matches!(
15494                name,
15495                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15496            ) || (Self::pdl_nvfp4q8_on()
15497                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15498        {
15499            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
15500            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
15501            // names may take this launch (unmarked kernels would read unordered).
15502            {
15503                use cudarc::driver::{DevicePtr, DevicePtrMut};
15504                let s = &self.gpu.stream();
15505                let (pw, _g0) = bytes.device_ptr(s);
15506                let (paq, _g1) = aq.device_ptr(s);
15507                let (pad, _g2) = ad.device_ptr(s);
15508                let (py, _g3) = y.device_ptr_mut(s);
15509                let mut ps = [
15510                    &pw as *const _ as *mut std::ffi::c_void,
15511                    &paq as *const _ as *mut _,
15512                    &pad as *const _ as *mut _,
15513                    &py as *const _ as *mut _,
15514                    &inf as *const _ as *mut _,
15515                    &outf as *const _ as *mut _,
15516                    &mi as *const _ as *mut _,
15517                    &rb as *const _ as *mut _,
15518                ];
15519                unsafe {
15520                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15521                }
15522            }
15523            if scale != 1.0 {
15524                self.scale_inplace(y, scale, m * out_f)?;
15525            }
15526        } else {
15527            b.arg(bytes)
15528                .arg(aq)
15529                .arg(ad)
15530                .arg(&mut *y)
15531                .arg(&inf)
15532                .arg(&outf)
15533                .arg(&mi)
15534                .arg(&rb);
15535            unsafe {
15536                b.launch(cfg)?;
15537            }
15538            if scale != 1.0 {
15539                self.scale_inplace(y, scale, m * out_f)?;
15540            }
15541        }
15542        Ok(())
15543    }
15544
15545    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
15546    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
15547    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
15548    pub fn qmatvec_mmvq_raw(
15549        &self,
15550        bytes: &CudaSlice<u8>,
15551        x: &CudaSlice<f32>,
15552        m: usize,
15553        in_f: usize,
15554        out_f: usize,
15555        qtype: i32,
15556        row_bytes: usize,
15557        rp: bool,
15558    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15559        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15560        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15561    }
15562
15563    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
15564    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
15565    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
15566    pub fn batched_supports(&self, qtype: i32) -> bool {
15567        matches!(
15568            qtype,
15569            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15570        )
15571    }
15572
15573    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
15574    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
15575    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
15576    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
15577    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
15578    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
15579    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
15580    pub fn iq_fast_enabled() -> bool {
15581        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15582        *ON.get_or_init(|| {
15583            std::env::var("MEMRA_IQ_FAST")
15584                .map(|v| v != "0")
15585                .unwrap_or(true)
15586        })
15587    }
15588
15589    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
15590    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
15591    pub fn b8_enabled() -> bool {
15592        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15593        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15594    }
15595
15596    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
15597    pub fn batched_mcols(m: usize) -> usize {
15598        if m == 2 {
15599            2
15600        } else if m <= 4 {
15601            4
15602        } else if m <= 8 {
15603            8
15604        } else {
15605            16
15606        }
15607    }
15608
15609    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
15610    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
15611    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
15612    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
15613    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15614        Some(match (qtype, mcols) {
15615            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15616            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15617            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15618            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
15619            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
15620            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
15621            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
15622            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
15623            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15624            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15625            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15626            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15627            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
15628            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
15629            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15630            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15631            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15632            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15633            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
15634            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
15635            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15636            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15637            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15638            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15639            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15640            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15641            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15642            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15643            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
15644            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
15645            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
15646            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15647            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15648            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15649            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15650            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
15651            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
15652            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
15653            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15654            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15655            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15656            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15657            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15658            _ => return None,
15659        })
15660    }
15661
15662    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
15663    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
15664    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
15665    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
15666    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
15667    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
15668    ///
15669    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
15670    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
15671    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
15672    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
15673    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
15674    /// msweep on all six 27B shapes (2026-07-03):
15675    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
15676    ///          it applies for b4 (-3..-14%), never loses;
15677    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
15678    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
15679    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
15680    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
15681    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
15682    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
15683    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
15684    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
15685    /// b2: in_f>=6144 -> r2, else base.
15686    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
15687    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
15688    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
15689    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
15690    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
15691    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
15692    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
15693    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
15694    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
15695    /// Device SM count (cached) — grid-fill policy input.
15696    pub fn sm_count(&self) -> i32 {
15697        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15698        *SMS.get_or_init(|| {
15699            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15700            self.gpu
15701                .ctx
15702                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15703                .unwrap_or(82)
15704        })
15705    }
15706
15707    pub fn batched_variant(
15708        &self,
15709        _m: usize,
15710        in_f: usize,
15711        out_f: usize,
15712        qtype: i32,
15713        row_bytes: usize,
15714        mcols: usize,
15715        rp: bool,
15716    ) -> &'static str {
15717        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
15718        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
15719        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
15720        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
15721        if qtype == QT_Q8_0 {
15722            return if rp { "rp" } else { "base" };
15723        }
15724        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15725        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15726            Ok("base") => "base",
15727            Ok("pf") => "pf",
15728            Ok("r2") => "r2",
15729            Ok("r2w8") => "r2w8",
15730            Ok("pfr2") => "pfr2",
15731            Ok("ca") => "ca",
15732            Ok("car2") => "car2",
15733            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
15734            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
15735            Ok("rp") => "rp",
15736            Ok("rpr2") => "rpr2",
15737            Ok("rpr2w8") => "rpr2w8",
15738            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
15739            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
15740            Ok("rpca") => "rpca",
15741            Ok("rpcar2") => "rpcar2",
15742            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
15743            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
15744            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
15745            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
15746            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
15747            // bit-identical to the decode path — measurement corpus ONLY, never auto).
15748            Ok("rpsc") => "rpsc",
15749            Ok("rpms") => "rpms",
15750            Ok("rpmsc") => "rpmsc",
15751            Ok("rpks") => "rpks",
15752            Ok("rpksc") => "rpksc",
15753            _ => "auto",
15754        });
15755        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
15756        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
15757        // shapes qualify; anything else falls back to the register variants.
15758        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15759        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
15760        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
15761        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
15762        // forced MEMRA_MMVQ_BV values still work).
15763        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15764        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15765        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15766        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15767        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15768        let sms = *SMS.get_or_init(|| {
15769            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15770            self.gpu
15771                .ctx
15772                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15773                .unwrap_or(82)
15774        });
15775        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
15776        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
15777        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
15778        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
15779        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
15780        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
15781        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
15782        // AUTO RULE = the measured winners table (differs from NVFP4's!):
15783        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
15784        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
15785        //     r2 1258us) — kernels kept behind the force seam for the corpus;
15786        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
15787        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
15788        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
15789        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
15790        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
15791        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
15792        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
15793        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
15794        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15795        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
15796        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
15797        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15798        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15799            Ok("base") => "base",
15800            Ok("r2") => "r2",
15801            Ok("r2w8") => "r2w8",
15802            _ => "auto",
15803        });
15804        let variant: &'static str = if qtype == QT_Q4_0 {
15805            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
15806            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
15807            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
15808            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15809            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15810                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
15811                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
15812                // + syncs cost more than the stalls, bank-pad made no difference);
15813                // register load-ahead flat (nvcc already reorders). The b-tier limiter
15814                // is still unidentified — see the jsonl row.
15815                Ok("base") => "base",
15816                Ok("r2") => "r2",
15817                Ok("ms") => "ms",
15818                Ok("sm") => "sm",
15819                Ok("la") => "la",
15820                _ => "auto",
15821            });
15822            let v = if q40 != "auto" {
15823                q40
15824            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15825                "r2"
15826            } else {
15827                "base"
15828            };
15829            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
15830            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
15831            // and the limiter is the per-column activation load chain (long_scoreboard
15832            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
15833            if rp {
15834                match v {
15835                    "ms" => "r2ms_rp",
15836                    "sm" => "r2sm_rp",
15837                    "la" => "r2la_rp",
15838                    "r2" => "r2_rp",
15839                    _ => "rp",
15840                }
15841            } else if matches!(v, "ms" | "sm" | "la") {
15842                "r2"
15843            } else {
15844                v
15845            }
15846        } else if qtype != QT_NVFP4 && !kq_r2 {
15847            "base"
15848        } else if kq_r2 && rp {
15849            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
15850            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
15851            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
15852            "rp"
15853        } else if kq_r2 {
15854            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
15855            // mcols != 4 forced r2w8 falls to unbounded r2.
15856            if kq_bv != "auto" {
15857                if kq_bv == "r2w8" && mcols != 4 {
15858                    "r2"
15859                } else {
15860                    kq_bv
15861                }
15862            } else if bv != "auto" {
15863                match bv {
15864                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
15865                    "r2w8" | "rpr2w8" => {
15866                        if mcols != 4 {
15867                            "r2"
15868                        } else {
15869                            "r2w8"
15870                        }
15871                    }
15872                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
15873                }
15874            } else {
15875                let blocks = (out_f + 7) / 8;
15876                let waves = blocks as f64 / (7 * sms as usize) as f64;
15877                let filled = blocks >= 4 * sms as usize;
15878                let use_r2 = if qtype == QT_Q4_K {
15879                    filled
15880                } else {
15881                    waves >= 2.0
15882                };
15883                if use_r2 { "r2" } else { "base" }
15884            }
15885        } else if bv != "auto" {
15886            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
15887            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
15888            // unsupported (shape, mcols) combos fall back to pf/r2.
15889            // On rp buffers, forced legacy names map to their rp twins (layout law).
15890            let v = if bv == "r2w8" && mcols == 2 {
15891                "r2"
15892            } else if bv == "ca" && (!ca_ok || mcols == 8) {
15893                "pf"
15894            } else if bv == "car2" && (!ca_ok || mcols == 8) {
15895                "r2"
15896            } else if bv == "pfr2" && mcols == 8 {
15897                "r2"
15898            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
15899                "rpr2"
15900            }
15901            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
15902            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
15903                if mcols == 8 { "rpr2w8" } else { "rpr2" }
15904            } else if bv == "rpcar2" && mcols == 2 {
15905                "rpca"
15906            }
15907            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
15908            // (rpms has no smem and no alignment need — always valid on rp buffers).
15909            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
15910                "rpr2"
15911            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
15912                "rpr2"
15913            } else {
15914                bv
15915            };
15916            if rp {
15917                match v {
15918                    "base" | "pf" | "ca" | "rp" => "rp",
15919                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
15920                    "r2w8" | "rpr2w8" => {
15921                        if mcols == 2 {
15922                            "rpr2"
15923                        } else {
15924                            "rpr2w8"
15925                        }
15926                    }
15927                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
15928                }
15929            } else {
15930                v
15931            }
15932        } else if mcols == 8 {
15933            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
15934            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
15935            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
15936            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
15937            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
15938            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
15939            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
15940            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
15941            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
15942            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
15943            if rp {
15944                if sc_ok { "rpsc" } else { "rpr2w8" }
15945            } else {
15946                "r2w8"
15947            }
15948        } else if mcols >= 4 {
15949            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
15950            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
15951            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
15952            let blocks = (out_f + 7) / 8;
15953            let r7 = 7 * sms as usize;
15954            let r8 = 8 * sms as usize;
15955            let waves = blocks as f64 / r7 as f64;
15956            let filled = blocks >= 4 * sms as usize;
15957            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
15958            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
15959            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
15960            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
15961                // the extra residency drops the INTEGER wave count -> the straggler wave a
15962                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
15963                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
15964                if rp { "rpr2w8" } else { "r2w8" }
15965            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
15966                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
15967                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
15968                if rp { "rpr2" } else { "r2" }
15969            } else {
15970                // fractional straggler-wave window with no crossing, or grid too small to fill
15971                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
15972                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
15973                if rp { "rp" } else { "pf" }
15974            }
15975        } else if in_f >= 6144 {
15976            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
15977            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
15978            // stays.
15979            if rp { "rpr2" } else { "r2" }
15980        } else if rp {
15981            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
15982            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
15983            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
15984            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
15985            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
15986            if sc_ok && waves >= 0.9 && waves <= 1.1 {
15987                "rpsc"
15988            } else {
15989                "rp"
15990            }
15991        } else {
15992            "base"
15993        };
15994        variant
15995    }
15996
15997    pub fn qmatvec_mmvq_batched(
15998        &self,
15999        bytes: &CudaSlice<u8>,
16000        aq: &CudaSlice<i8>,
16001        ad: &CudaSlice<f32>,
16002        m: usize,
16003        in_f: usize,
16004        out_f: usize,
16005        qtype: i32,
16006        row_bytes: usize,
16007        mcols: usize,
16008        scale: f32,
16009        rp: bool,
16010    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16011        const ROWS_PER_BLOCK: u32 = 4;
16012        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
16013        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
16014        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
16015        // weight keeps its rp-layout kernel family regardless of the override.
16016        let forced: Option<&'static str> = {
16017            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16018            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16019                .as_deref()
16020                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16021        };
16022        let variant = match forced {
16023            Some(v) if !rp || v.contains("rp") => v,
16024            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16025        };
16026        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16027            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16028        })?;
16029        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
16030        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
16031        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
16032        let variant = if mcols == 16 {
16033            if rp { "rp" } else { "base" }
16034        } else {
16035            variant
16036        };
16037        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
16038        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
16039        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
16040        // per-(token,row) chain (columns c >= m never execute in either form) ->
16041        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
16042        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
16043        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16044        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16045        if b567
16046            && qtype == QT_NVFP4
16047            && rp
16048            && mcols == 8
16049            && (5..=7).contains(&m)
16050            && matches!(variant, "rpsc" | "rpr2w8")
16051        {
16052            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16053            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
16054            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16055            let cfg = LaunchConfig {
16056                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16057                block_dim: (32, ROWS_PER_BLOCK, 1),
16058                shared_mem_bytes: 0,
16059            };
16060            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16061            let __s_b = self.gpu.stream();
16062            let mut b = __s_b.launch_builder(&f);
16063            b.arg(bytes)
16064                .arg(aq)
16065                .arg(ad)
16066                .arg(&mut y)
16067                .arg(&inf)
16068                .arg(&outf)
16069                .arg(&mi)
16070                .arg(&rb);
16071            unsafe {
16072                b.launch(cfg)?;
16073            }
16074            if scale != 1.0 {
16075                self.scale_inplace(&mut y, scale, m * out_f)?;
16076            }
16077            return Ok(y);
16078        }
16079        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16080            "base" => (base_name.into(), ROWS_PER_BLOCK),
16081            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16082            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16083            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16084            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
16085            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
16086            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
16087            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16088            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16089            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16090            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16091            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16092            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16093            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16094            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
16095        };
16096        debug_assert!(
16097            !rp || name.contains("_rp"),
16098            "rp weight dispatched to a GGUF-layout kernel"
16099        );
16100        let f = self.func(&name);
16101        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16102        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
16103        let smem = if name.contains("_r2sm_rp") {
16104            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16105        } else {
16106            0
16107        };
16108        let cfg = LaunchConfig {
16109            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16110            block_dim: (32, ROWS_PER_BLOCK, 1),
16111            shared_mem_bytes: smem,
16112        };
16113        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16114        let __s_b = self.gpu.stream();
16115        let mut b = __s_b.launch_builder(&f);
16116        b.arg(bytes)
16117            .arg(aq)
16118            .arg(ad)
16119            .arg(&mut y)
16120            .arg(&inf)
16121            .arg(&outf)
16122            .arg(&mi)
16123            .arg(&rb);
16124        unsafe {
16125            b.launch(cfg)?;
16126        }
16127        if scale != 1.0 {
16128            self.scale_inplace(&mut y, scale, m * out_f)?;
16129        }
16130        Ok(y)
16131    }
16132
16133    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
16134    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
16135    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
16136    pub fn qmatvec_batched_raw(
16137        &self,
16138        bytes: &CudaSlice<u8>,
16139        x: &CudaSlice<f32>,
16140        m: usize,
16141        in_f: usize,
16142        out_f: usize,
16143        qtype: i32,
16144        row_bytes: usize,
16145        mcols: usize,
16146        rp: bool,
16147    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16148        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16149        self.qmatvec_mmvq_batched(
16150            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16151        )
16152    }
16153
16154    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
16155    pub fn qmatvec_nvfp4_batched_raw(
16156        &self,
16157        bytes: &CudaSlice<u8>,
16158        x: &CudaSlice<f32>,
16159        m: usize,
16160        in_f: usize,
16161        out_f: usize,
16162        row_bytes: usize,
16163        mcols: usize,
16164        rp: bool,
16165    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16166        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16167    }
16168
16169    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
16170    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
16171    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
16172    fn try_fp4_gemm(
16173        &self,
16174        w: &crate::model::GpuTensor,
16175        x: &CudaSlice<f32>,
16176        m: usize,
16177        in_f: usize,
16178        out_f: usize,
16179    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16180        use crate::model::GpuTensor;
16181        if cfg!(memra_portable_cuda) {
16182            return Ok(None);
16183        }
16184        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which cu/qmatvec_gemm.cu:1234 omits on a
16185        // portable build (the mxf4 block-scale MMA is sm_120a-only). Refuse at the door.
16186        if std::env::var("MEMRA_FP4").is_ok() {
16187            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16188        }
16189        if std::env::var("MEMRA_FP4").is_err() {
16190            return Ok(None);
16191        }
16192        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
16193        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
16194        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
16195        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
16196        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
16197        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
16198        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
16199        // for the common no-macro-scale case.
16200        #[cfg(memra_cutlass)]
16201        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16202            if let GpuTensor::Quant {
16203                bytes,
16204                qtype,
16205                scale,
16206                row_bytes,
16207                cutlass,
16208                ..
16209            } = w
16210            {
16211                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16212                    if let Some(cw) = cutlass {
16213                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16214                        let y = self.cutlass_fp4_gemm(
16215                            &cw.b_packed,
16216                            &cw.sfb_swizzled,
16217                            x,
16218                            *scale,
16219                            m,
16220                            out_f,
16221                            in_f,
16222                        )?;
16223                        return Ok(Some(y));
16224                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16225                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16226                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16227                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16228                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16229                        let (b_packed, sfb_sw) =
16230                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16231                        let y =
16232                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16233                        return Ok(Some(y));
16234                    }
16235                }
16236            }
16237        }
16238        if let GpuTensor::Quant {
16239            bytes,
16240            qtype,
16241            row_bytes,
16242            scale,
16243            rp,
16244            ..
16245        } = w
16246        {
16247            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16248            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16249            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16250                let y =
16251                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16252                return Ok(Some(y));
16253            }
16254        }
16255        Ok(None)
16256    }
16257
16258    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16259    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16260    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16261    pub fn rms_norm_f16out(
16262        &self,
16263        x: &CudaSlice<f32>,
16264        w: &CudaSlice<f32>,
16265        dst: &mut CudaSlice<f32>,
16266        dst16: &mut CudaSlice<u8>,
16267        ncols: usize,
16268        nrows: usize,
16269        eps: f32,
16270    ) -> Result<(), Box<dyn std::error::Error>> {
16271        let f = self.func("rms_norm_f16out_f32");
16272        let cfg = LaunchConfig {
16273            grid_dim: (nrows as u32, 1, 1),
16274            block_dim: (rms_block(), 1, 1),
16275            shared_mem_bytes: 0,
16276        };
16277        let (nc, e) = (ncols as i32, eps);
16278        let __s_b = self.gpu.stream();
16279        let mut b = __s_b.launch_builder(&f);
16280        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16281        unsafe {
16282            b.launch(cfg)?;
16283        }
16284        Ok(())
16285    }
16286
16287    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16288    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16289    #[allow(clippy::too_many_arguments)]
16290    pub fn add_rms_norm_f16out(
16291        &self,
16292        a: &CudaSlice<f32>,
16293        b: &CudaSlice<f32>,
16294        w: &CudaSlice<f32>,
16295        res: &mut CudaSlice<f32>,
16296        dst: &mut CudaSlice<f32>,
16297        dst16: &mut CudaSlice<u8>,
16298        ncols: usize,
16299        nrows: usize,
16300        eps: f32,
16301    ) -> Result<(), Box<dyn std::error::Error>> {
16302        let f = self.func("add_rms_norm_f16out_f32");
16303        let cfg = LaunchConfig {
16304            grid_dim: (nrows as u32, 1, 1),
16305            block_dim: (rms_block(), 1, 1),
16306            shared_mem_bytes: 0,
16307        };
16308        let (nc, e) = (ncols as i32, eps);
16309        let __s_lb = self.gpu.stream();
16310        let mut lb = __s_lb.launch_builder(&f);
16311        lb.arg(a)
16312            .arg(b)
16313            .arg(w)
16314            .arg(res)
16315            .arg(dst)
16316            .arg(dst16)
16317            .arg(&nc)
16318            .arg(&e);
16319        unsafe {
16320            lb.launch(cfg)?;
16321        }
16322        Ok(())
16323    }
16324
16325    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
16326    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
16327    pub fn matmul_group_xh(
16328        &self,
16329        ws: &[&crate::model::GpuTensor],
16330        x: &CudaSlice<f32>,
16331        xh: &CudaSlice<u8>,
16332        m: usize,
16333    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16334        let mut out = Vec::with_capacity(ws.len());
16335        let in_f = ws[0].in_features();
16336        for w in ws {
16337            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16338                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16339                    out.push(y);
16340                    continue;
16341                }
16342            }
16343            out.push(self.matmul(w, x, m)?);
16344        }
16345        Ok(out)
16346    }
16347
16348    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
16349    /// GDN steps). Layouts [T, H].
16350    pub fn gdn_pad_mask(
16351        &self,
16352        beta: &mut CudaSlice<f32>,
16353        g_log: &mut CudaSlice<f32>,
16354        len_d: &CudaSlice<i32>,
16355        h: usize,
16356        t: usize,
16357    ) -> Result<(), Box<dyn std::error::Error>> {
16358        let f = self.func("gdn_pad_mask_f32");
16359        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16360        let (hi, ti) = (h as i32, t as i32);
16361        let __s_b = self.gpu.stream();
16362        let mut b = __s_b.launch_builder(&f);
16363        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16364        unsafe {
16365            b.launch(cfg)?;
16366        }
16367        Ok(())
16368    }
16369
16370    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
16371    /// gather for the padded prime graph's h_seed/hlast.
16372    pub fn row_gather_dev(
16373        &self,
16374        src: &CudaSlice<f32>,
16375        dst: &mut CudaSlice<f32>,
16376        len_d: &CudaSlice<i32>,
16377        ncols: usize,
16378    ) -> Result<(), Box<dyn std::error::Error>> {
16379        let f = self.func("row_gather_dev_f32");
16380        let cfg = LaunchConfig::for_num_elems(ncols as u32);
16381        let nc = ncols as i32;
16382        let __s_b = self.gpu.stream();
16383        let mut b = __s_b.launch_builder(&f);
16384        b.arg(src).arg(dst).arg(len_d).arg(&nc);
16385        unsafe {
16386            b.launch(cfg)?;
16387        }
16388        Ok(())
16389    }
16390
16391    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
16392    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
16393    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
16394    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
16395    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
16396    /// different in_f) falls back to its own `matmul` — behavior unchanged.
16397    pub fn matmul_group(
16398        &self,
16399        ws: &[&crate::model::GpuTensor],
16400        x: &CudaSlice<f32>,
16401        m: usize,
16402    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16403        use crate::model::GpuTensor;
16404        let mut out = Vec::with_capacity(ws.len());
16405        let any_mirror = ws
16406            .iter()
16407            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16408        if m >= 16 && any_mirror && !self.verify_exact_on() {
16409            let in_f = ws[0].in_features();
16410            let xh = self.f16_act(x, m * in_f, in_f)?;
16411            for w in ws {
16412                if w.in_features() == in_f {
16413                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16414                        out.push(y);
16415                        continue;
16416                    }
16417                }
16418                out.push(self.matmul(w, x, m)?);
16419            }
16420            return Ok(out);
16421        }
16422        for w in ws {
16423            out.push(self.matmul(w, x, m)?);
16424        }
16425        Ok(out)
16426    }
16427
16428    /// Cross-request grouped matmul (task #13): run ONE projection group over the
16429    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
16430    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
16431    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
16432    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
16433    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
16434    pub fn matmul_group_multi(
16435        &self,
16436        ws: &[&crate::model::GpuTensor],
16437        xs: &[&CudaSlice<f32>],
16438        ms: &[usize],
16439    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16440        assert_eq!(xs.len(), ms.len());
16441        let in_f = ws[0].in_features();
16442        let total: usize = ms.iter().sum();
16443        let mut xcat = self.uninit(total * in_f)?;
16444        let mut off = 0usize;
16445        for (x, &m) in xs.iter().zip(ms) {
16446            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16447            off += m;
16448        }
16449        let ys = self.matmul_group(ws, &xcat, total)?;
16450        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16451        for (w, y) in ws.iter().zip(ys) {
16452            let out_f = w.out_features();
16453            let mut off = 0usize;
16454            for (s, &m) in ms.iter().enumerate() {
16455                let mut ys_s = self.uninit(m * out_f)?;
16456                let src = y.slice(off * out_f..(off + m) * out_f);
16457                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16458                out[s].push(ys_s);
16459                off += m;
16460            }
16461        }
16462        Ok(out)
16463    }
16464
16465    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
16466    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
16467    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
16468    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
16469    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
16470    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
16471    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
16472    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
16473    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
16474    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16475        use crate::model::GpuTensor;
16476        if !legacy_quant_gemm_allowed(
16477            cfg!(memra_portable_cuda),
16478            cfg!(memra_hopper_mma),
16479            std::env::var_os("MEMRA_NO_GEMM").is_some(),
16480        ) {
16481            return false;
16482        }
16483        match w {
16484            GpuTensor::Quant { qtype, .. } => {
16485                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16486                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16487            }
16488            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16489        }
16490    }
16491
16492    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
16493    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
16494    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
16495    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
16496    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
16497    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
16498    pub fn qmatvec_gemm(
16499        &self,
16500        w: &crate::model::GpuTensor,
16501        aq: &CudaSlice<i8>,
16502        ad: &CudaSlice<f32>,
16503        m: usize,
16504    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16505        use crate::model::GpuTensor;
16506        let in_f = w.in_features();
16507        let out_f = w.out_features();
16508        let (bytes, qtype, row_bytes, scale, rp) = match w {
16509            GpuTensor::Quant {
16510                bytes,
16511                qtype,
16512                row_bytes,
16513                scale,
16514                rp,
16515                ..
16516            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16517            _ => unreachable!("gemm_supports guaranteed Quant"),
16518        };
16519        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
16520        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
16521        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
16522        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
16523        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
16524        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16525            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16526                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16527                if scale != 1.0 {
16528                    self.scale_inplace(&mut y, scale, m * out_f)?;
16529                }
16530                return Ok(y);
16531            }
16532        }
16533        let name = match qtype {
16534            QT_Q8_0 => "qmatvec_gemm_q8_0",
16535            QT_Q4_K => "qmatvec_gemm_q4_K",
16536            QT_Q4_0 => {
16537                if rp {
16538                    "qmatvec_gemm_q4_0_rp"
16539                } else {
16540                    "qmatvec_gemm_q4_0"
16541                }
16542            }
16543            QT_Q5_K => "qmatvec_gemm_q5_K",
16544            QT_Q6_K => "qmatvec_gemm_q6_K",
16545            QT_NVFP4 => {
16546                if rp {
16547                    "qmatvec_gemm_nvfp4_rp"
16548                } else {
16549                    "qmatvec_gemm_nvfp4"
16550                }
16551            }
16552            _ => unreachable!(),
16553        };
16554        let f = self.func(name);
16555        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16556        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
16557        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
16558        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
16559        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16560        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16561        let k1_tile = if is_k1 {
16562            k1_launch_override().unwrap_or((128, 128, 8))
16563        } else {
16564            (128, 128, 8)
16565        };
16566        let (bm, bn): (u32, u32) = if is_k1 {
16567            (k1_tile.0, k1_tile.1)
16568        } else {
16569            (64, 256)
16570        };
16571        let warps: u32 = if is_k1 {
16572            k1_tile.2
16573        } else {
16574            match qtype {
16575                QT_NVFP4 => 8,
16576                _ => 4,
16577            }
16578        };
16579        let cfg = LaunchConfig {
16580            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16581            block_dim: (32, warps, 1),
16582            shared_mem_bytes: 0,
16583        };
16584        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16585        let __s_b = self.gpu.stream();
16586        let mut b = __s_b.launch_builder(&f);
16587        b.arg(bytes)
16588            .arg(aq)
16589            .arg(ad)
16590            .arg(&mut y)
16591            .arg(&inf)
16592            .arg(&outf)
16593            .arg(&mi)
16594            .arg(&rb);
16595        unsafe {
16596            b.launch(cfg)?;
16597        }
16598        if scale != 1.0 {
16599            self.scale_inplace(&mut y, scale, m * out_f)?;
16600        }
16601        Ok(y)
16602    }
16603
16604    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
16605    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
16606    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
16607    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
16608    pub fn qmatvec_gemm_raw(
16609        &self,
16610        bytes: &CudaSlice<u8>,
16611        x: &CudaSlice<f32>,
16612        m: usize,
16613        in_f: usize,
16614        out_f: usize,
16615        qtype: i32,
16616        row_bytes: usize,
16617    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16618        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16619        let name = match qtype {
16620            QT_Q8_0 => "qmatvec_gemm_q8_0",
16621            QT_Q4_K => "qmatvec_gemm_q4_K",
16622            QT_Q4_0 => "qmatvec_gemm_q4_0",
16623            QT_Q5_K => "qmatvec_gemm_q5_K",
16624            QT_Q6_K => "qmatvec_gemm_q6_K",
16625            QT_NVFP4 => "qmatvec_gemm_nvfp4",
16626            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16627            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16628        };
16629        let f = self.func(name);
16630        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16631        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
16632        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
16633        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16634        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16635        let k1_tile = if is_k1 {
16636            k1_launch_override().unwrap_or((128, 128, 8))
16637        } else {
16638            (128, 128, 8)
16639        };
16640        let (bm, bn): (u32, u32) = if is_k1 {
16641            (k1_tile.0, k1_tile.1)
16642        } else {
16643            (64, 256)
16644        };
16645        let warps: u32 = if is_k1 {
16646            k1_tile.2
16647        } else {
16648            match qtype {
16649                QT_NVFP4 | QT_NVFP4_RP => 8,
16650                _ => 4,
16651            }
16652        };
16653        let cfg = LaunchConfig {
16654            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16655            block_dim: (32, warps, 1),
16656            shared_mem_bytes: 0,
16657        };
16658        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16659        let __s_b = self.gpu.stream();
16660        let mut b = __s_b.launch_builder(&f);
16661        b.arg(bytes)
16662            .arg(&aq)
16663            .arg(&ad)
16664            .arg(&mut y)
16665            .arg(&inf)
16666            .arg(&outf)
16667            .arg(&mi)
16668            .arg(&rb);
16669        unsafe {
16670            b.launch(cfg)?;
16671        }
16672        Ok(y)
16673    }
16674
16675    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
16676    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
16677    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
16678    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
16679    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
16680    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
16681    pub fn qmatvec_gemm_q8_0_wgmma_raw(
16682        &self,
16683        rp4: &CudaSlice<u8>,
16684        aq: &CudaSlice<i8>,
16685        ad: &CudaSlice<f32>,
16686        m: usize,
16687        in_f: usize,
16688        out_f: usize,
16689    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16690        assert!(
16691            out_f % 64 == 0 && in_f % 32 == 0,
16692            "wgmma GEMM needs out_f%64==0, in_f%32==0"
16693        );
16694        let f = self.func("qmatvec_gemm_q8_0_wgmma");
16695        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
16696        let cfg = LaunchConfig {
16697            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16698            block_dim: (128, 1, 1),
16699            shared_mem_bytes: 0,
16700        };
16701        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16702        let __s_b = self.gpu.stream();
16703        let mut b = __s_b.launch_builder(&f);
16704        b.arg(rp4)
16705            .arg(aq)
16706            .arg(ad)
16707            .arg(&mut y)
16708            .arg(&inf)
16709            .arg(&outf)
16710            .arg(&mi);
16711        unsafe {
16712            b.launch(cfg)?;
16713        }
16714        Ok(y)
16715    }
16716
16717    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
16718    pub fn scale_inplace(
16719        &self,
16720        y: &mut CudaSlice<f32>,
16721        s: f32,
16722        n: usize,
16723    ) -> Result<(), Box<dyn std::error::Error>> {
16724        let f = self.func("scale_f32");
16725        let cfg = LaunchConfig::for_num_elems(n as u32);
16726        let (sf, ni) = (s, n as i32);
16727        let __s_b = self.gpu.stream();
16728        let mut b = __s_b.launch_builder(&f);
16729        b.arg(y).arg(&sf).arg(&ni);
16730        unsafe {
16731            b.launch(cfg)?;
16732        }
16733        Ok(())
16734    }
16735
16736    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
16737    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
16738    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
16739    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
16740    pub fn bf16_to_f32(
16741        &self,
16742        data: &cudarc::driver::CudaView<'_, u8>,
16743        n: usize,
16744    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16745        let mut out = self.alloc_uninit::<f32>(n)?;
16746        let f = self.func("bf16_to_f32");
16747        let cfg = LaunchConfig::for_num_elems(n as u32);
16748        let ni = n as i32;
16749        let __s_b = self.gpu.stream();
16750        let mut b = __s_b.launch_builder(&f);
16751        b.arg(data).arg(&mut out).arg(&ni);
16752        unsafe {
16753            b.launch(cfg)?;
16754        }
16755        Ok(out)
16756    }
16757
16758    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
16759    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
16760    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
16761    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
16762    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
16763    /// calls, the spec-verify contract) vs plain linear.
16764    fn linear_bf16_chunked(
16765        &self,
16766        x: &CudaSlice<f32>,
16767        data: &CudaSlice<u8>,
16768        m: usize,
16769        in_f: usize,
16770        out_f: usize,
16771        exact: bool,
16772        canonical_chunk_rows: Option<usize>,
16773    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16774        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
16775        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
16776        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
16777        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16778        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16779        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16780        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16781        let started = timing.then(std::time::Instant::now);
16782        let result =
16783            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16784        if let Some(started) = started {
16785            use std::sync::atomic::Ordering;
16786            self.stream().synchronize()?;
16787            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16788                + started.elapsed().as_nanos() as u64;
16789            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16790                + (in_f * out_f * 2) as u64;
16791            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16792            if calls % 1024 == 0 {
16793                eprintln!(
16794                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16795                     weight_gb={:.2}",
16796                    ns as f64 / 1.0e6,
16797                    ns as f64 / calls as f64 / 1.0e3,
16798                    wb as f64 / 1.0e9,
16799                );
16800            }
16801        }
16802        result
16803    }
16804
16805    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
16806    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
16807    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
16808    /// numeric-class doors (DEV_ROUTES precedent).
16809    pub(crate) fn bf16_mmv_on() -> bool {
16810        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16811        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16812    }
16813
16814    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
16815    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
16816    fn matvec_bf16(
16817        &self,
16818        data: &CudaSlice<u8>,
16819        x: &CudaSlice<f32>,
16820        in_f: usize,
16821        out_f: usize,
16822    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16823        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16824            return Err(format!(
16825                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16826                data.len(),
16827                x.len()
16828            )
16829            .into());
16830        }
16831        let mut y = self.alloc_uninit::<f32>(out_f)?;
16832        let f = self.func("matvec_bf16_f32acc");
16833        let cfg = LaunchConfig {
16834            grid_dim: (out_f as u32, 1, 1),
16835            block_dim: (mmv_block(), 1, 1),
16836            shared_mem_bytes: 0,
16837        };
16838        let ini = in_f as i32;
16839        let __s_bld = self.gpu.stream();
16840        let mut bld = __s_bld.launch_builder(&f);
16841        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
16842        unsafe {
16843            bld.launch(cfg)?;
16844        }
16845        Ok(y)
16846    }
16847
16848    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
16849    /// launches, a position upload, and the rope launch; the position is read directly from
16850    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
16851    #[allow(clippy::too_many_arguments)]
16852    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
16853    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
16854    /// Bit-identical to the split kernels; requires head_dim == 128 and
16855    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
16856    #[allow(clippy::too_many_arguments)]
16857    pub fn qk_norm_rope_append_inc_dcw(
16858        &self,
16859        q_raw: &CudaSlice<f32>,
16860        k_raw: &CudaSlice<f32>,
16861        v_raw: &CudaSlice<f32>,
16862        qw: &CudaSlice<f32>,
16863        kw: &CudaSlice<f32>,
16864        q_out: &mut CudaSlice<f32>,
16865        k_out: &mut CudaSlice<f32>,
16866        pos: &CudaSlice<i32>,
16867        k_plane: &mut CudaSlice<u8>,
16868        v_plane: &mut CudaSlice<u8>,
16869        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
16870        // (single) writer, exactly like the split append+inc pair it replaces.
16871        len_dev: &CudaSlice<i32>,
16872        base_dev: Option<&CudaSlice<i32>>,
16873        done_ctr: &mut CudaSlice<u32>,
16874        kv_dim_k: usize,
16875        kv_dim_v: usize,
16876        k_tok_bytes: usize,
16877        v_tok_bytes: usize,
16878        head_dim: usize,
16879        n_dims: usize,
16880        nh_q: usize,
16881        nh_k: usize,
16882        eps: f32,
16883        freq_base: f32,
16884        freq_scale: f32,
16885        ff: Option<&CudaSlice<f32>>,
16886    ) -> Result<(), Box<dyn std::error::Error>> {
16887        if head_dim != 128
16888            || kv_dim_v != kv_dim_k
16889            || kv_dim_k != nh_k * head_dim
16890            || q_raw.len() < nh_q * head_dim
16891            || k_raw.len() < nh_k * head_dim
16892            || v_raw.len() < kv_dim_v
16893            || q_out.len() < nh_q * head_dim
16894            || k_out.len() < nh_k * head_dim
16895            || pos.is_empty()
16896            || done_ctr.is_empty()
16897        {
16898            return Err(format!(
16899                "qk_norm_rope_append_inc geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}                  kv_k={kv_dim_k} kv_v={kv_dim_v}"
16900            )
16901            .into());
16902        }
16903        let f = self.func("qk_norm_rope_append_inc_dcw");
16904        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16905        let cfg = LaunchConfig {
16906            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16907            block_dim: (128, 1, 1),
16908            shared_mem_bytes: 0,
16909        };
16910        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
16911        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16912        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16913        let null: u64 = 0;
16914        let __s_b = self.gpu.stream();
16915        let mut b = __s_b.launch_builder(&f);
16916        b.arg(q_raw)
16917            .arg(k_raw)
16918            .arg(v_raw)
16919            .arg(qw)
16920            .arg(kw)
16921            .arg(q_out)
16922            .arg(k_out)
16923            .arg(pos)
16924            .arg(&mut *k_plane)
16925            .arg(&mut *v_plane)
16926            .arg(len_dev);
16927        match base_dev {
16928            Some(base) => {
16929                b.arg(base);
16930            }
16931            None => {
16932                b.arg(&null);
16933            }
16934        }
16935        b.arg(&mut *done_ctr)
16936            .arg(&kvk)
16937            .arg(&kvv)
16938            .arg(&ktb)
16939            .arg(&vtb)
16940            .arg(&hd)
16941            .arg(&nd)
16942            .arg(&nq)
16943            .arg(&eps)
16944            .arg(&theta_scale)
16945            .arg(&freq_scale);
16946        match ff {
16947            Some(freqs) => {
16948                b.arg(freqs);
16949            }
16950            None => {
16951                b.arg(&null);
16952            }
16953        }
16954        unsafe {
16955            b.launch(cfg)?;
16956        }
16957        Ok(())
16958    }
16959
16960    pub fn qk_norm_rope_into(
16961        &self,
16962        q_raw: &CudaSlice<f32>,
16963        k_raw: &CudaSlice<f32>,
16964        qw: &CudaSlice<f32>,
16965        kw: &CudaSlice<f32>,
16966        q_out: &mut CudaSlice<f32>,
16967        k_out: &mut CudaSlice<f32>,
16968        pos: &CudaSlice<i32>,
16969        head_dim: usize,
16970        n_dims: usize,
16971        nh_q: usize,
16972        nh_k: usize,
16973        eps: f32,
16974        freq_base: f32,
16975        freq_scale: f32,
16976        ff: Option<&CudaSlice<f32>>,
16977    ) -> Result<(), Box<dyn std::error::Error>> {
16978        if head_dim > 512
16979            || q_raw.len() < nh_q * head_dim
16980            || k_raw.len() < nh_k * head_dim
16981            || q_out.len() < nh_q * head_dim
16982            || k_out.len() < nh_k * head_dim
16983            || qw.len() < head_dim
16984            || kw.len() < head_dim
16985            || pos.is_empty()
16986        {
16987            return Err(format!(
16988                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
16989            )
16990            .into());
16991        }
16992        let f = self.func("qk_norm_rope_f32");
16993        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16994        let cfg = LaunchConfig {
16995            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16996            block_dim: (128, 1, 1),
16997            shared_mem_bytes: 0,
16998        };
16999        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17000        let __s_b = self.gpu.stream();
17001        let mut b = __s_b.launch_builder(&f);
17002        b.arg(q_raw)
17003            .arg(k_raw)
17004            .arg(qw)
17005            .arg(kw)
17006            .arg(q_out)
17007            .arg(k_out)
17008            .arg(pos)
17009            .arg(&hd)
17010            .arg(&nd)
17011            .arg(&nq)
17012            .arg(&eps)
17013            .arg(&theta_scale)
17014            .arg(&freq_scale);
17015        match ff {
17016            Some(ffv) => {
17017                b.arg(ffv);
17018                unsafe {
17019                    b.launch(cfg)?;
17020                }
17021            }
17022            None => {
17023                let null: u64 = 0;
17024                b.arg(&null);
17025                unsafe {
17026                    b.launch(cfg)?;
17027                }
17028            }
17029        }
17030        Ok(())
17031    }
17032
17033    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
17034    /// launch computes a rank's whole O partial from its four canonical column blocks.
17035    #[allow(clippy::too_many_arguments)]
17036    pub fn matvec_f32_b4_into(
17037        &self,
17038        w: [&CudaSlice<f32>; 4],
17039        x: &CudaSlice<f32>,
17040        y: &mut CudaSlice<f32>,
17041        block_cols: usize,
17042        out_f: usize,
17043    ) -> Result<(), Box<dyn std::error::Error>> {
17044        if block_cols % 4 != 0
17045            || x.len() < 4 * block_cols
17046            || y.len() < out_f
17047            || w.iter().any(|w| w.len() != out_f * block_cols)
17048        {
17049            return Err(format!(
17050                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17051                x.len()
17052            )
17053            .into());
17054        }
17055        let f = self.func("matvec_f32_b4");
17056        let cfg = LaunchConfig {
17057            grid_dim: (out_f as u32, 1, 1),
17058            block_dim: (128, 1, 1),
17059            shared_mem_bytes: 0,
17060        };
17061        let (bc, of) = (block_cols as i32, out_f as i32);
17062        let __s_b = self.gpu.stream();
17063        let mut b = __s_b.launch_builder(&f);
17064        b.arg(w[0])
17065            .arg(w[1])
17066            .arg(w[2])
17067            .arg(w[3])
17068            .arg(x)
17069            .arg(y)
17070            .arg(&bc)
17071            .arg(&of);
17072        unsafe {
17073            b.launch(cfg)?;
17074        }
17075        Ok(())
17076    }
17077
17078    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
17079    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
17080    pub fn axpy_rows_seq_into(
17081        &self,
17082        x: &CudaSlice<f32>,
17083        w: &CudaSlice<f32>,
17084        y: &mut CudaSlice<f32>,
17085        width: usize,
17086        n_rows: usize,
17087    ) -> Result<(), Box<dyn std::error::Error>> {
17088        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17089            return Err(format!(
17090                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17091                x.len(),
17092                w.len(),
17093                y.len()
17094            )
17095            .into());
17096        }
17097        let f = self.func("axpy_rows_seq_f32");
17098        let cfg = LaunchConfig::for_num_elems(width as u32);
17099        let (wi, nr) = (width as i32, n_rows as i32);
17100        let __s_b = self.gpu.stream();
17101        let mut b = __s_b.launch_builder(&f);
17102        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17103        unsafe {
17104            b.launch(cfg)?;
17105        }
17106        Ok(())
17107    }
17108
17109    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
17110    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
17111    /// exact sequential FP chain of the base kernel over that window.
17112    #[allow(clippy::too_many_arguments)]
17113    pub fn axpy_rows_seq_md_off_into(
17114        &self,
17115        x: &CudaSlice<f32>,
17116        w_route: &CudaSlice<f32>,
17117        md: &CudaSlice<f32>,
17118        sel: &CudaSlice<i32>,
17119        y: &mut CudaSlice<f32>,
17120        width: usize,
17121        n_rows: usize,
17122        row0: usize,
17123    ) -> Result<(), Box<dyn std::error::Error>> {
17124        if x.len() < (row0 + n_rows) * width
17125            || w_route.len() < row0 + n_rows
17126            || sel.len() < row0 + n_rows
17127            || y.len() < width
17128        {
17129            return Err(format!(
17130                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17131                 rows={n_rows} row0={row0}",
17132                x.len(),
17133                w_route.len(),
17134                sel.len(),
17135                y.len()
17136            )
17137            .into());
17138        }
17139        let f = self.func("axpy_rows_seq_md_off_f32");
17140        let cfg = LaunchConfig::for_num_elems(width as u32);
17141        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17142        let __s_b = self.gpu.stream();
17143        let mut b = __s_b.launch_builder(&f);
17144        b.arg(x)
17145            .arg(w_route)
17146            .arg(md)
17147            .arg(sel)
17148            .arg(y)
17149            .arg(&wi)
17150            .arg(&nr)
17151            .arg(&r0);
17152        unsafe {
17153            b.launch(cfg)?;
17154        }
17155        Ok(())
17156    }
17157
17158    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
17159    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
17160    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
17161    /// outputs are bit-equal to its own t=1 launch.
17162    #[allow(clippy::too_many_arguments)]
17163    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17164        &self,
17165        gate_bank: &CudaSlice<u8>,
17166        up_bank: &CudaSlice<u8>,
17167        sel: &CudaSlice<i32>,
17168        aq: &CudaSlice<i8>,
17169        ad: &CudaSlice<f32>,
17170        yg: &mut CudaSlice<f32>,
17171        yu: &mut CudaSlice<f32>,
17172        n_sel: usize,
17173        n_sel_col: usize,
17174        in_f: usize,
17175        out_f: usize,
17176        row_bytes: usize,
17177        expert_stride: usize,
17178        act_row_stride: usize,
17179        ad_row_stride: usize,
17180    ) -> Result<(), Box<dyn std::error::Error>> {
17181        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17182        if yg.len() < n_sel * out_f
17183            || yu.len() < n_sel * out_f
17184            || sel.len() < n_sel
17185            || n_sel_col == 0
17186            || n_sel % n_sel_col != 0
17187        {
17188            return Err("NVFP4 gu tcol geometry".into());
17189        }
17190        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17191        let cfg = LaunchConfig {
17192            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17193            block_dim: (128, 1, 1),
17194            shared_mem_bytes: 0,
17195        };
17196        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17197        let (rb, es) = (row_bytes as i64, expert_stride as i64);
17198        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17199        let __s_b = self.gpu.stream();
17200        let mut b = __s_b.launch_builder(&f);
17201        b.arg(gate_bank)
17202            .arg(up_bank)
17203            .arg(sel)
17204            .arg(aq)
17205            .arg(ad)
17206            .arg(yg)
17207            .arg(yu)
17208            .arg(&inf)
17209            .arg(&outf)
17210            .arg(&ns)
17211            .arg(&rb)
17212            .arg(&es)
17213            .arg(&ars)
17214            .arg(&adrs)
17215            .arg(&nsc);
17216        unsafe {
17217            b.launch(cfg)?;
17218        }
17219        Ok(())
17220    }
17221
17222    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
17223    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
17224    #[allow(clippy::too_many_arguments)]
17225    pub fn axpy_rows_seq_md_into(
17226        &self,
17227        x: &CudaSlice<f32>,
17228        w_route: &CudaSlice<f32>,
17229        md: &CudaSlice<f32>,
17230        sel: &CudaSlice<i32>,
17231        y: &mut CudaSlice<f32>,
17232        width: usize,
17233        n_rows: usize,
17234    ) -> Result<(), Box<dyn std::error::Error>> {
17235        if x.len() < n_rows * width
17236            || w_route.len() < n_rows
17237            || sel.len() < n_rows
17238            || y.len() < width
17239        {
17240            return Err(format!(
17241                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17242                x.len(),
17243                w_route.len(),
17244                sel.len(),
17245                y.len()
17246            )
17247            .into());
17248        }
17249        let f = self.func("axpy_rows_seq_md_f32");
17250        let cfg = LaunchConfig::for_num_elems(width as u32);
17251        let (wi, nr) = (width as i32, n_rows as i32);
17252        let __s_b = self.gpu.stream();
17253        let mut b = __s_b.launch_builder(&f);
17254        b.arg(x)
17255            .arg(w_route)
17256            .arg(md)
17257            .arg(sel)
17258            .arg(y)
17259            .arg(&wi)
17260            .arg(&nr);
17261        unsafe {
17262            b.launch(cfg)?;
17263        }
17264        Ok(())
17265    }
17266
17267    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
17268    #[allow(clippy::too_many_arguments)]
17269    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
17270    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
17271    /// land column-major-of-rows: yq[c*out_q + row] etc.
17272    #[allow(clippy::too_many_arguments)]
17273    pub fn matvec_bf16_qkvg_tcol_into(
17274        &self,
17275        wq: &CudaSlice<u8>,
17276        wk: &CudaSlice<u8>,
17277        wv: &CudaSlice<u8>,
17278        wg: &CudaSlice<u8>,
17279        x_t: &CudaSlice<f32>,
17280        yq: &mut CudaSlice<f32>,
17281        yk: &mut CudaSlice<f32>,
17282        yv: &mut CudaSlice<f32>,
17283        yg: &mut CudaSlice<f32>,
17284        in_f: usize,
17285        out_q: usize,
17286        out_kv: usize,
17287        out_g: usize,
17288        t: usize,
17289    ) -> Result<(), Box<dyn std::error::Error>> {
17290        if t == 0
17291            || t > 8
17292            || in_f % 8 != 0
17293            || x_t.len() < t * in_f
17294            || yq.len() < t * out_q
17295            || yk.len() < t * out_kv
17296            || yv.len() < t * out_kv
17297            || (out_g > 0 && yg.len() < t * out_g)
17298        {
17299            return Err("matvec_bf16_qkvg_tcol geometry".into());
17300        }
17301        let f = self.func("matvec_bf16_qkvg_tcol");
17302        let grid = out_q + 2 * out_kv + out_g;
17303        let cfg = LaunchConfig {
17304            grid_dim: (grid as u32, 1, 1),
17305            block_dim: (mmv_block(), 1, 1),
17306            shared_mem_bytes: 0,
17307        };
17308        let (ini, oq, okv, og, ti) = (
17309            in_f as i32,
17310            out_q as i32,
17311            out_kv as i32,
17312            out_g as i32,
17313            t as i32,
17314        );
17315        let __s_b = self.gpu.stream();
17316        let mut b = __s_b.launch_builder(&f);
17317        b.arg(wq)
17318            .arg(wk)
17319            .arg(wv)
17320            .arg(wg)
17321            .arg(x_t)
17322            .arg(yq)
17323            .arg(yk)
17324            .arg(yv)
17325            .arg(yg)
17326            .arg(&ini)
17327            .arg(&oq)
17328            .arg(&okv)
17329            .arg(&og)
17330            .arg(&ti);
17331        unsafe {
17332            b.launch(cfg)?;
17333        }
17334        Ok(())
17335    }
17336
17337    pub fn matvec_bf16_qkvg_into(
17338        &self,
17339        wq: &CudaSlice<u8>,
17340        wk: &CudaSlice<u8>,
17341        wv: &CudaSlice<u8>,
17342        wg: &CudaSlice<u8>,
17343        x: &CudaSlice<f32>,
17344        yq: &mut CudaSlice<f32>,
17345        yk: &mut CudaSlice<f32>,
17346        yv: &mut CudaSlice<f32>,
17347        yg: &mut CudaSlice<f32>,
17348        in_f: usize,
17349        out_q: usize,
17350        out_kv: usize,
17351        out_g: usize,
17352    ) -> Result<(), Box<dyn std::error::Error>> {
17353        if in_f % 8 != 0
17354            || wq.len() != out_q * in_f * 2
17355            || wk.len() != out_kv * in_f * 2
17356            || wv.len() != out_kv * in_f * 2
17357            || wg.len() < out_g * in_f * 2
17358            || x.len() < in_f
17359            || yq.len() < out_q
17360            || yk.len() < out_kv
17361            || yv.len() < out_kv
17362            || (out_g > 0 && yg.len() < out_g)
17363        {
17364            return Err(format!(
17365                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17366            )
17367            .into());
17368        }
17369        let f = self.func("matvec_bf16_qkvg");
17370        let cfg = LaunchConfig {
17371            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17372            block_dim: (mmv_block(), 1, 1),
17373            shared_mem_bytes: 0,
17374        };
17375        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17376        let __s_b = self.gpu.stream();
17377        let mut b = __s_b.launch_builder(&f);
17378        b.arg(wq)
17379            .arg(wk)
17380            .arg(wv)
17381            .arg(wg)
17382            .arg(x)
17383            .arg(yq)
17384            .arg(yk)
17385            .arg(yv)
17386            .arg(yg)
17387            .arg(&inf)
17388            .arg(&oq)
17389            .arg(&okv)
17390            .arg(&og);
17391        unsafe {
17392            b.launch(cfg)?;
17393        }
17394        Ok(())
17395    }
17396
17397    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
17398    pub fn matvec_bf16_b4_into(
17399        &self,
17400        w: [&CudaSlice<u8>; 4],
17401        x: &CudaSlice<f32>,
17402        y: &mut CudaSlice<f32>,
17403        block_cols: usize,
17404        out_f: usize,
17405    ) -> Result<(), Box<dyn std::error::Error>> {
17406        if block_cols % 8 != 0
17407            || x.len() < 4 * block_cols
17408            || y.len() < out_f
17409            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17410        {
17411            return Err(format!(
17412                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17413                x.len()
17414            )
17415            .into());
17416        }
17417        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
17418        // bit-identical per row (the second row's stream hides the first's reduce tail).
17419        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17420        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17421        let f = self.func(if x2 {
17422            "matvec_bf16_b4_x2"
17423        } else {
17424            "matvec_bf16_b4"
17425        });
17426        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17427        let cfg = LaunchConfig {
17428            grid_dim: (grid as u32, 1, 1),
17429            block_dim: (mmv_block(), 1, 1),
17430            shared_mem_bytes: 0,
17431        };
17432        let (bc, of) = (block_cols as i32, out_f as i32);
17433        let __s_b = self.gpu.stream();
17434        let mut b = __s_b.launch_builder(&f);
17435        b.arg(w[0])
17436            .arg(w[1])
17437            .arg(w[2])
17438            .arg(w[3])
17439            .arg(x)
17440            .arg(y)
17441            .arg(&bc)
17442            .arg(&of);
17443        unsafe {
17444            b.launch(cfg)?;
17445        }
17446        Ok(())
17447    }
17448
17449    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
17450    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
17451    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
17452    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
17453    /// t=1 program).
17454    pub fn matvec_bf16_b4_tcol_into(
17455        &self,
17456        w: [&CudaSlice<u8>; 4],
17457        x_t: &CudaSlice<f32>,
17458        y_t: &mut CudaSlice<f32>,
17459        block_cols: usize,
17460        out_f: usize,
17461        t: usize,
17462    ) -> Result<(), Box<dyn std::error::Error>> {
17463        if block_cols % 8 != 0
17464            || t == 0
17465            || t > 8
17466            || x_t.len() < t * 4 * block_cols
17467            || y_t.len() < t * out_f
17468            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17469        {
17470            return Err(format!(
17471                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17472                x_t.len()
17473            )
17474            .into());
17475        }
17476        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17477            return Err(
17478                "b4 tcol verify is qualified against the plain b4 kernel only \
17479                        (MEMRA_B4_X2=1 is a different t=1 program)"
17480                    .into(),
17481            );
17482        }
17483        let f = self.func("matvec_bf16_b4_tcol");
17484        let cfg = LaunchConfig {
17485            grid_dim: (out_f as u32, 1, 1),
17486            block_dim: (mmv_block(), 1, 1),
17487            shared_mem_bytes: 0,
17488        };
17489        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17490        let __s_b = self.gpu.stream();
17491        let mut b = __s_b.launch_builder(&f);
17492        b.arg(w[0])
17493            .arg(w[1])
17494            .arg(w[2])
17495            .arg(w[3])
17496            .arg(x_t)
17497            .arg(y_t)
17498            .arg(&bc)
17499            .arg(&of)
17500            .arg(&ti);
17501        unsafe {
17502            b.launch(cfg)?;
17503        }
17504        Ok(())
17505    }
17506
17507    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
17508    pub fn matvec_bf16_into(
17509        &self,
17510        data: &CudaSlice<u8>,
17511        x: &CudaSlice<f32>,
17512        y: &mut CudaSlice<f32>,
17513        in_f: usize,
17514        out_f: usize,
17515    ) -> Result<(), Box<dyn std::error::Error>> {
17516        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17517            return Err(format!(
17518                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17519                data.len(),
17520                x.len(),
17521                y.len()
17522            )
17523            .into());
17524        }
17525        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
17526        // block, exact f32acc per-row program — cures the 1-iteration latency
17527        // starvation (shexp down measured 420GB/s at in_f=1280).
17528        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17529        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17530            && in_f <= 2048;
17531        if x4 {
17532            let f = self.func("matvec_bf16_f32acc_x4");
17533            let cfg = LaunchConfig {
17534                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17535                block_dim: (mmv_block(), 1, 1),
17536                shared_mem_bytes: 0,
17537            };
17538            let (ini, outi) = (in_f as i32, out_f as i32);
17539            let __s_b = self.gpu.stream();
17540            let mut b = __s_b.launch_builder(&f);
17541            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17542            unsafe {
17543                b.launch(cfg)?;
17544            }
17545            return Ok(());
17546        }
17547        let f = self.func("matvec_bf16_f32acc");
17548        let cfg = LaunchConfig {
17549            grid_dim: (out_f as u32, 1, 1),
17550            block_dim: (mmv_block(), 1, 1),
17551            shared_mem_bytes: 0,
17552        };
17553        let ini = in_f as i32;
17554        let __s_b = self.gpu.stream();
17555        let mut b = __s_b.launch_builder(&f);
17556        b.arg(data).arg(x).arg(y).arg(&ini);
17557        unsafe {
17558            b.launch(cfg)?;
17559        }
17560        Ok(())
17561    }
17562
17563    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
17564    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
17565    pub fn matvec_bf16_view_into(
17566        &self,
17567        data: &cudarc::driver::CudaView<'_, u8>,
17568        x: &CudaSlice<f32>,
17569        y: &mut CudaSlice<f32>,
17570        in_f: usize,
17571        out_f: usize,
17572    ) -> Result<(), Box<dyn std::error::Error>> {
17573        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17574            return Err(format!(
17575                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17576                data.len(),
17577                x.len(),
17578                y.len()
17579            )
17580            .into());
17581        }
17582        let f = self.func("matvec_bf16_f32acc");
17583        let cfg = LaunchConfig {
17584            grid_dim: (out_f as u32, 1, 1),
17585            block_dim: (mmv_block(), 1, 1),
17586            shared_mem_bytes: 0,
17587        };
17588        let ini = in_f as i32;
17589        let __s_b = self.gpu.stream();
17590        let mut b = __s_b.launch_builder(&f);
17591        b.arg(data).arg(x).arg(y).arg(&ini);
17592        unsafe {
17593            b.launch(cfg)?;
17594        }
17595        Ok(())
17596    }
17597
17598    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
17599    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
17600    pub fn matvec_bf16_raw_out(
17601        &self,
17602        w: &CudaSlice<u8>,
17603        x: &CudaSlice<f32>,
17604        y_raw: u64,
17605        in_f: usize,
17606        out_f: usize,
17607    ) -> Result<(), Box<dyn std::error::Error>> {
17608        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17609            return Err("matvec_bf16_raw_out geometry".into());
17610        }
17611        let f = self.func("matvec_bf16_f32acc");
17612        let cfg = LaunchConfig {
17613            grid_dim: (out_f as u32, 1, 1),
17614            block_dim: (mmv_block(), 1, 1),
17615            shared_mem_bytes: 0,
17616        };
17617        let ini = in_f as i32;
17618        let __s_b = self.gpu.stream();
17619        let mut b = __s_b.launch_builder(&f);
17620        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17621        unsafe {
17622            b.launch(cfg)?;
17623        }
17624        Ok(())
17625    }
17626
17627    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
17628    /// UVA pointers so the caller passes persistent-static rows without holding locks).
17629    /// Exact per-element sequence of the split add + add_scaled_rows pair.
17630    pub fn add3_raw(
17631        &self,
17632        a: &CudaSlice<f32>,
17633        b: &CudaSlice<f32>,
17634        sh_raw: u64,
17635        scale_raw: u64,
17636        dst: &mut CudaSlice<f32>,
17637        n: usize,
17638    ) -> Result<(), Box<dyn std::error::Error>> {
17639        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17640            return Err("add3_raw geometry".into());
17641        }
17642        let f = self.func("add3_f32");
17643        let cfg = LaunchConfig {
17644            grid_dim: ((n as u32).div_ceil(256), 1, 1),
17645            block_dim: (256, 1, 1),
17646            shared_mem_bytes: 0,
17647        };
17648        let ni = n as i32;
17649        let __s_b = self.gpu.stream();
17650        let mut bld = __s_b.launch_builder(&f);
17651        bld.arg(a)
17652            .arg(b)
17653            .arg(&sh_raw)
17654            .arg(&scale_raw)
17655            .arg(dst)
17656            .arg(&ni);
17657        unsafe {
17658            bld.launch(cfg)?;
17659        }
17660        Ok(())
17661    }
17662
17663    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
17664    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
17665    pub fn matvec_bf16_down_addscale_into(
17666        &self,
17667        w: &CudaSlice<u8>,
17668        x: &CudaSlice<f32>,
17669        scale: &CudaSlice<f32>,
17670        dst: &mut CudaSlice<f32>,
17671        in_f: usize,
17672        out_f: usize,
17673    ) -> Result<(), Box<dyn std::error::Error>> {
17674        if w.len() != in_f * out_f * 2
17675            || x.len() < in_f
17676            || in_f % 8 != 0
17677            || dst.len() < out_f
17678            || scale.is_empty()
17679        {
17680            return Err("matvec_bf16_down_addscale geometry".into());
17681        }
17682        let f = self.func("matvec_bf16_down_addscale");
17683        let cfg = LaunchConfig {
17684            grid_dim: (out_f as u32, 1, 1),
17685            block_dim: (mmv_block(), 1, 1),
17686            shared_mem_bytes: 0,
17687        };
17688        let ini = in_f as i32;
17689        let __s_b = self.gpu.stream();
17690        let mut b = __s_b.launch_builder(&f);
17691        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17692        unsafe {
17693            b.launch(cfg)?;
17694        }
17695        Ok(())
17696    }
17697
17698    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
17699    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
17700    pub fn matvec_bf16_dual_silu_into(
17701        &self,
17702        wg: &CudaSlice<u8>,
17703        wu: &CudaSlice<u8>,
17704        x: &CudaSlice<f32>,
17705        act: &mut CudaSlice<f32>,
17706        in_f: usize,
17707        out_f: usize,
17708        limit: Option<f32>,
17709    ) -> Result<(), Box<dyn std::error::Error>> {
17710        if wg.len() != in_f * out_f * 2
17711            || wu.len() != in_f * out_f * 2
17712            || x.len() < in_f
17713            || in_f % 8 != 0
17714            || act.len() < out_f
17715        {
17716            return Err("matvec_bf16_dual_silu geometry".into());
17717        }
17718        let f = self.func("matvec_bf16_dual_silu");
17719        let cfg = LaunchConfig {
17720            grid_dim: (out_f as u32, 1, 1),
17721            block_dim: (mmv_block(), 1, 1),
17722            shared_mem_bytes: 0,
17723        };
17724        let (ini, outi) = (in_f as i32, out_f as i32);
17725        let lim = limit.unwrap_or(0.0);
17726        let __s_b = self.gpu.stream();
17727        let mut b = __s_b.launch_builder(&f);
17728        b.arg(wg)
17729            .arg(wu)
17730            .arg(x)
17731            .arg(act)
17732            .arg(&ini)
17733            .arg(&outi)
17734            .arg(&lim);
17735        unsafe {
17736            b.launch(cfg)?;
17737        }
17738        Ok(())
17739    }
17740
17741    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
17742    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
17743    #[allow(clippy::too_many_arguments)]
17744    pub fn matvec_bf16_dual_view_into(
17745        &self,
17746        wg: &cudarc::driver::CudaView<'_, u8>,
17747        wu: &cudarc::driver::CudaView<'_, u8>,
17748        x: &CudaSlice<f32>,
17749        yg: &mut CudaSlice<f32>,
17750        yu: &mut CudaSlice<f32>,
17751        in_f: usize,
17752        out_f: usize,
17753    ) -> Result<(), Box<dyn std::error::Error>> {
17754        if wg.len() != in_f * out_f * 2
17755            || wu.len() != in_f * out_f * 2
17756            || x.len() < in_f
17757            || in_f % 8 != 0
17758            || yg.len() < out_f
17759            || yu.len() < out_f
17760        {
17761            return Err(format!(
17762                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17763                wg.len(),
17764                wu.len(),
17765                x.len()
17766            )
17767            .into());
17768        }
17769        let f = self.func("matvec_bf16_dual");
17770        let cfg = LaunchConfig {
17771            grid_dim: ((2 * out_f) as u32, 1, 1),
17772            block_dim: (mmv_block(), 1, 1),
17773            shared_mem_bytes: 0,
17774        };
17775        let (ini, outi) = (in_f as i32, out_f as i32);
17776        let __s_b = self.gpu.stream();
17777        let mut b = __s_b.launch_builder(&f);
17778        b.arg(wg)
17779            .arg(wu)
17780            .arg(x)
17781            .arg(yg)
17782            .arg(yu)
17783            .arg(&ini)
17784            .arg(&outi);
17785        unsafe {
17786            b.launch(cfg)?;
17787        }
17788        Ok(())
17789    }
17790
17791    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
17792    #[allow(clippy::too_many_arguments)]
17793    pub fn matvec_bf16_dual_into(
17794        &self,
17795        wg: &CudaSlice<u8>,
17796        wu: &CudaSlice<u8>,
17797        x: &CudaSlice<f32>,
17798        yg: &mut CudaSlice<f32>,
17799        yu: &mut CudaSlice<f32>,
17800        in_f: usize,
17801        out_f: usize,
17802    ) -> Result<(), Box<dyn std::error::Error>> {
17803        if wg.len() != in_f * out_f * 2
17804            || wu.len() != in_f * out_f * 2
17805            || x.len() < in_f
17806            || in_f % 8 != 0
17807            || yg.len() < out_f
17808            || yu.len() < out_f
17809        {
17810            return Err(format!(
17811                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17812                wg.len(),
17813                wu.len(),
17814                x.len()
17815            )
17816            .into());
17817        }
17818        let f = self.func("matvec_bf16_dual");
17819        let cfg = LaunchConfig {
17820            grid_dim: ((2 * out_f) as u32, 1, 1),
17821            block_dim: (mmv_block(), 1, 1),
17822            shared_mem_bytes: 0,
17823        };
17824        let (ini, outi) = (in_f as i32, out_f as i32);
17825        let __s_b = self.gpu.stream();
17826        let mut b = __s_b.launch_builder(&f);
17827        b.arg(wg)
17828            .arg(wu)
17829            .arg(x)
17830            .arg(yg)
17831            .arg(yu)
17832            .arg(&ini)
17833            .arg(&outi);
17834        unsafe {
17835            b.launch(cfg)?;
17836        }
17837        Ok(())
17838    }
17839
17840    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
17841    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
17842    pub(crate) fn matvec_bf16_dual(
17843        &self,
17844        wg: &CudaSlice<u8>,
17845        wu: &CudaSlice<u8>,
17846        x: &CudaSlice<f32>,
17847        in_f: usize,
17848        out_f: usize,
17849    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17850        if wg.len() != in_f * out_f * 2
17851            || wu.len() != in_f * out_f * 2
17852            || x.len() < in_f
17853            || in_f % 8 != 0
17854        {
17855            return Err(format!(
17856                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
17857                wg.len(),
17858                wu.len(),
17859                x.len()
17860            )
17861            .into());
17862        }
17863        let mut yg = self.alloc_uninit::<f32>(out_f)?;
17864        let mut yu = self.alloc_uninit::<f32>(out_f)?;
17865        let f = self.func("matvec_bf16_dual");
17866        let cfg = LaunchConfig {
17867            grid_dim: ((2 * out_f) as u32, 1, 1),
17868            block_dim: (mmv_block(), 1, 1),
17869            shared_mem_bytes: 0,
17870        };
17871        let (ini, outi) = (in_f as i32, out_f as i32);
17872        let __s_b = self.gpu.stream();
17873        let mut b = __s_b.launch_builder(&f);
17874        b.arg(wg)
17875            .arg(wu)
17876            .arg(x)
17877            .arg(&mut yg)
17878            .arg(&mut yu)
17879            .arg(&ini)
17880            .arg(&outi);
17881        unsafe {
17882            b.launch(cfg)?;
17883        }
17884        Ok((yg, yu))
17885    }
17886
17887    #[allow(clippy::too_many_arguments)]
17888    fn linear_bf16_chunked_inner(
17889        &self,
17890        x: &CudaSlice<f32>,
17891        data: &CudaSlice<u8>,
17892        m: usize,
17893        in_f: usize,
17894        out_f: usize,
17895        exact: bool,
17896        canonical_chunk_rows: Option<usize>,
17897    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17898        const CHUNK_BYTES: usize = 256 << 20;
17899        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
17900        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
17901        if m == 1
17902            && !exact
17903            && canonical_chunk_rows.is_none()
17904            && in_f % 8 == 0
17905            && Self::bf16_mmv_on()
17906        {
17907            return self.matvec_bf16(data, x, in_f, out_f);
17908        }
17909        let row_bytes = in_f
17910            .checked_mul(std::mem::size_of::<f32>())
17911            .ok_or("BF16 chunk row byte count overflow")?;
17912        if row_bytes == 0 || out_f == 0 {
17913            return Err("BF16 chunk dimensions must be nonzero".into());
17914        }
17915        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
17916        let chunk_rows = match canonical_chunk_rows {
17917            Some(rows) if rows == 0 => {
17918                return Err("canonical BF16 chunk rows must be nonzero".into());
17919            }
17920            Some(rows) if rows > max_chunk_rows => {
17921                return Err(format!(
17922                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
17923                )
17924                .into());
17925            }
17926            Some(rows) if out_f % rows != 0 => {
17927                return Err(format!(
17928                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
17929                )
17930                .into());
17931            }
17932            Some(rows) => rows,
17933            None => max_chunk_rows,
17934        };
17935        if chunk_rows >= out_f {
17936            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
17937            return if exact {
17938                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
17939            } else {
17940                self.linear(x, &wf32, m, in_f, out_f)
17941            };
17942        }
17943        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17944        let mut r0 = 0usize;
17945        while r0 < out_f {
17946            let rows = chunk_rows.min(out_f - r0);
17947            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
17948            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
17949            let yc = if exact {
17950                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
17951            } else {
17952                self.linear(x, &wf32, m, in_f, rows)?
17953            };
17954            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
17955            for mi in 0..m {
17956                let src = yc.slice(mi * rows..(mi + 1) * rows);
17957                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
17958                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
17959            }
17960            r0 += rows;
17961        }
17962        Ok(y)
17963    }
17964
17965    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
17966    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
17967    /// chunked BF16 numerical program instead of re-encoding the weight.
17968    pub fn linear_bf16_resident(
17969        &self,
17970        x: &CudaSlice<f32>,
17971        data: &CudaSlice<u8>,
17972        m: usize,
17973        in_f: usize,
17974        out_f: usize,
17975    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17976        if data.len() != in_f * out_f * 2 {
17977            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17978        }
17979        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
17980    }
17981
17982    /// Execute a resident BF16 projection as fixed-width output-row chunks.
17983    ///
17984    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
17985    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
17986    /// model topology rather than the active rank count.
17987    pub fn linear_bf16_resident_canonical_rows(
17988        &self,
17989        x: &CudaSlice<f32>,
17990        data: &CudaSlice<u8>,
17991        m: usize,
17992        in_f: usize,
17993        out_f: usize,
17994        canonical_chunk_rows: usize,
17995    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17996        if data.len() != in_f * out_f * 2 {
17997            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17998        }
17999        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18000    }
18001
18002    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
18003    ///
18004    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
18005    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
18006    pub fn linear_f32_resident_canonical_rows(
18007        &self,
18008        x: &CudaSlice<f32>,
18009        data: &CudaSlice<f32>,
18010        m: usize,
18011        in_f: usize,
18012        out_f: usize,
18013        canonical_chunk_rows: usize,
18014    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18015        self.linear_f32_resident_canonical_rows_inner(
18016            x,
18017            data,
18018            m,
18019            in_f,
18020            out_f,
18021            canonical_chunk_rows,
18022            false,
18023        )
18024    }
18025
18026    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
18027    ///
18028    /// The projection shapes and values are identical to
18029    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
18030    /// changes, replacing one device copy per token with one placement kernel per output chunk.
18031    pub fn linear_f32_resident_canonical_rows_strided(
18032        &self,
18033        x: &CudaSlice<f32>,
18034        data: &CudaSlice<f32>,
18035        m: usize,
18036        in_f: usize,
18037        out_f: usize,
18038        canonical_chunk_rows: usize,
18039    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18040        self.linear_f32_resident_canonical_rows_inner(
18041            x,
18042            data,
18043            m,
18044            in_f,
18045            out_f,
18046            canonical_chunk_rows,
18047            true,
18048        )
18049    }
18050
18051    fn linear_f32_resident_canonical_rows_inner(
18052        &self,
18053        x: &CudaSlice<f32>,
18054        data: &CudaSlice<f32>,
18055        m: usize,
18056        in_f: usize,
18057        out_f: usize,
18058        canonical_chunk_rows: usize,
18059        strided_output: bool,
18060    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18061        if data.len() != in_f * out_f {
18062            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18063        }
18064        if canonical_chunk_rows == 0
18065            || canonical_chunk_rows > out_f
18066            || out_f % canonical_chunk_rows != 0
18067        {
18068            return Err(format!(
18069                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18070            )
18071            .into());
18072        }
18073        if canonical_chunk_rows == out_f {
18074            return self.linear(x, data, m, in_f, out_f);
18075        }
18076
18077        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18078        let input = x.slice(0..x.len());
18079        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18080            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18081            if m == 1 {
18082                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18083                self.linear_device_into(
18084                    &input,
18085                    &weights,
18086                    &mut destination,
18087                    1,
18088                    in_f,
18089                    canonical_chunk_rows,
18090                )?;
18091                continue;
18092            }
18093            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18094            if strided_output {
18095                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18096            } else {
18097                for token in 0..m {
18098                    let source = chunk
18099                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18100                    let mut destination =
18101                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18102                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18103                }
18104            }
18105        }
18106        Ok(y)
18107    }
18108
18109    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
18110    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
18111    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
18112    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
18113    pub fn linear_f32_resident_canonical_rows_t1_into(
18114        &self,
18115        x: &CudaSlice<f32>,
18116        data: &CudaSlice<f32>,
18117        y: &mut CudaSlice<f32>,
18118        in_f: usize,
18119        out_f: usize,
18120        canonical_chunk_rows: usize,
18121    ) -> Result<(), Box<dyn std::error::Error>> {
18122        if data.len() != in_f * out_f {
18123            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18124        }
18125        if y.len() != out_f || x.len() != in_f {
18126            return Err(format!(
18127                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18128                x.len(),
18129                y.len()
18130            )
18131            .into());
18132        }
18133        if canonical_chunk_rows == 0
18134            || canonical_chunk_rows > out_f
18135            || out_f % canonical_chunk_rows != 0
18136        {
18137            return Err(format!(
18138                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18139            )
18140            .into());
18141        }
18142        let input = x.slice(0..x.len());
18143        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18144            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18145            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18146            self.linear_device_into(
18147                &input,
18148                &weights,
18149                &mut destination,
18150                1,
18151                in_f,
18152                canonical_chunk_rows,
18153            )?;
18154        }
18155        Ok(())
18156    }
18157
18158    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
18159    /// without the allocation, for workspace-resident operands.
18160    pub fn linear_t1_into(
18161        &self,
18162        x: &cudarc::driver::CudaView<'_, f32>,
18163        w: &cudarc::driver::CudaView<'_, f32>,
18164        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18165        in_f: usize,
18166        out_f: usize,
18167    ) -> Result<(), Box<dyn std::error::Error>> {
18168        self.linear_device_into(x, w, y, 1, in_f, out_f)
18169    }
18170
18171    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
18172    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
18173    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
18174    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
18175    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
18176    /// router/shexp sites and matmul_decode_exact's Float arm.
18177    pub fn linear_decode_exact(
18178        &self,
18179        x: &CudaSlice<f32>,
18180        w: &CudaSlice<f32>,
18181        m_tokens: usize,
18182        in_f: usize,
18183        out_f: usize,
18184    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18185        if m_tokens == 1 {
18186            return self.linear(x, w, 1, in_f, out_f);
18187        }
18188        let xv = self.view(x, m_tokens * in_f);
18189        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18190        for t in 0..m_tokens {
18191            let row = xv.slice(t * in_f..(t + 1) * in_f);
18192            let mut xr = self.alloc_uninit::<f32>(in_f)?;
18193            self.copy_view_into(&mut xr, 0, &row, in_f)?;
18194            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18195            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18196        }
18197        Ok(y)
18198    }
18199
18200    pub fn linear(
18201        &self,
18202        x: &CudaSlice<f32>,
18203        w: &CudaSlice<f32>,
18204        m_tokens: usize,
18205        in_f: usize,
18206        out_f: usize,
18207    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18208        self.linear_device(x, w, m_tokens, in_f, out_f)
18209    }
18210
18211    fn linear_device<I>(
18212        &self,
18213        x: &I,
18214        w: &I,
18215        m_tokens: usize,
18216        in_f: usize,
18217        out_f: usize,
18218    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18219    where
18220        I: cudarc::driver::DevicePtr<f32>,
18221    {
18222        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
18223        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
18224        Ok(c)
18225    }
18226
18227    fn linear_device_into<I, O>(
18228        &self,
18229        x: &I,
18230        w: &I,
18231        c: &mut O,
18232        m_tokens: usize,
18233        in_f: usize,
18234        out_f: usize,
18235    ) -> Result<(), Box<dyn std::error::Error>>
18236    where
18237        I: cudarc::driver::DevicePtr<f32>,
18238        O: cudarc::driver::DevicePtrMut<f32>,
18239    {
18240        use cudarc::cublaslt::{Matmul, MatmulConfig};
18241        let cfg = MatmulConfig {
18242            transa: true,
18243            transb: false,
18244            transc: false,
18245            m: out_f as u64,
18246            n: m_tokens as u64,
18247            k: in_f as u64,
18248            alpha: 1.0,
18249            lda: in_f as i64,
18250            ldb: in_f as i64,
18251            beta: 0.0,
18252            ldc: out_f as i64,
18253            stride_a: None,
18254            stride_b: None,
18255            stride_c: None,
18256            stride_bias: None,
18257            batch_size: None,
18258        };
18259        let blas = self.gpu.blas();
18260        unsafe {
18261            blas.matmul(cfg, w, x, c, None, None)?;
18262        }
18263        Ok(())
18264    }
18265
18266    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
18267    pub fn sdpa_naive(
18268        &self,
18269        q: &CudaSlice<f32>,
18270        k: &CudaSlice<f32>,
18271        v: &CudaSlice<f32>,
18272        o: &mut CudaSlice<f32>,
18273        head_dim: usize,
18274        n_head: usize,
18275        n_head_kv: usize,
18276        t: usize,
18277        t_kv: usize,
18278        scale: f32,
18279        causal: bool,
18280    ) -> Result<(), Box<dyn std::error::Error>> {
18281        let f = self.func("sdpa_naive_f32");
18282        let cfg = LaunchConfig {
18283            grid_dim: (n_head as u32, t as u32, 1),
18284            block_dim: (128, 1, 1),
18285            shared_mem_bytes: (t_kv * 4) as u32,
18286        };
18287        let (hd, nh, nhkv, ti, tkvi, cz) = (
18288            head_dim as i32,
18289            n_head as i32,
18290            n_head_kv as i32,
18291            t as i32,
18292            t_kv as i32,
18293            causal as i32,
18294        );
18295        let __s_b = self.gpu.stream();
18296        let mut b = __s_b.launch_builder(&f);
18297        b.arg(q)
18298            .arg(k)
18299            .arg(v)
18300            .arg(o)
18301            .arg(&hd)
18302            .arg(&nh)
18303            .arg(&nhkv)
18304            .arg(&ti)
18305            .arg(&tkvi)
18306            .arg(&scale)
18307            .arg(&cz);
18308        unsafe {
18309            b.launch(cfg)?;
18310        }
18311        Ok(())
18312    }
18313
18314    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
18315    /// bidirectional image islands. `span_id` labels each absolute kv position
18316    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
18317    /// reproducing the reference's non-causal image batch. window 0 = no window.
18318    #[allow(clippy::too_many_arguments)]
18319    pub fn sdpa_naive_island(
18320        &self,
18321        q: &CudaSlice<f32>,
18322        k: &CudaSlice<f32>,
18323        v: &CudaSlice<f32>,
18324        o: &mut CudaSlice<f32>,
18325        span_id: &CudaSlice<i32>,
18326        head_dim: usize,
18327        n_head: usize,
18328        n_head_kv: usize,
18329        t: usize,
18330        t_kv: usize,
18331        scale: f32,
18332        window: usize,
18333    ) -> Result<(), Box<dyn std::error::Error>> {
18334        let f = self.func("sdpa_naive_island_f32");
18335        let cfg = LaunchConfig {
18336            grid_dim: (n_head as u32, t as u32, 1),
18337            block_dim: (128, 1, 1),
18338            shared_mem_bytes: (t_kv * 4) as u32,
18339        };
18340        let (hd, nh, nhkv, ti, tkvi, wi) = (
18341            head_dim as i32,
18342            n_head as i32,
18343            n_head_kv as i32,
18344            t as i32,
18345            t_kv as i32,
18346            window as i32,
18347        );
18348        let __s_b = self.gpu.stream();
18349        let mut b = __s_b.launch_builder(&f);
18350        b.arg(q)
18351            .arg(k)
18352            .arg(v)
18353            .arg(o)
18354            .arg(span_id)
18355            .arg(&hd)
18356            .arg(&nh)
18357            .arg(&nhkv)
18358            .arg(&ti)
18359            .arg(&tkvi)
18360            .arg(&scale)
18361            .arg(&wi);
18362        unsafe {
18363            b.launch(cfg)?;
18364        }
18365        Ok(())
18366    }
18367
18368    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
18369    #[allow(clippy::too_many_arguments)]
18370    pub fn sdpa_naive_w(
18371        &self,
18372        q: &CudaSlice<f32>,
18373        k: &CudaSlice<f32>,
18374        v: &CudaSlice<f32>,
18375        o: &mut CudaSlice<f32>,
18376        head_dim: usize,
18377        n_head: usize,
18378        n_head_kv: usize,
18379        t: usize,
18380        t_kv: usize,
18381        scale: f32,
18382        causal: bool,
18383        window: usize,
18384    ) -> Result<(), Box<dyn std::error::Error>> {
18385        let f = self.func("sdpa_naive_w_f32");
18386        let cfg = LaunchConfig {
18387            grid_dim: (n_head as u32, t as u32, 1),
18388            block_dim: (128, 1, 1),
18389            shared_mem_bytes: (t_kv * 4) as u32,
18390        };
18391        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18392            head_dim as i32,
18393            n_head as i32,
18394            n_head_kv as i32,
18395            t as i32,
18396            t_kv as i32,
18397            causal as i32,
18398            window as i32,
18399        );
18400        let __s_b = self.gpu.stream();
18401        let mut b = __s_b.launch_builder(&f);
18402        b.arg(q)
18403            .arg(k)
18404            .arg(v)
18405            .arg(o)
18406            .arg(&hd)
18407            .arg(&nh)
18408            .arg(&nhkv)
18409            .arg(&ti)
18410            .arg(&tkvi)
18411            .arg(&scale)
18412            .arg(&cz)
18413            .arg(&wi);
18414        unsafe {
18415            b.launch(cfg)?;
18416        }
18417        Ok(())
18418    }
18419
18420    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
18421    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
18422    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
18423    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
18424    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
18425    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
18426    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
18427    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
18428    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
18429    #[allow(clippy::too_many_arguments)]
18430    pub fn sdpa_naive_w_lo(
18431        &self,
18432        q: &CudaSlice<f32>,
18433        k: &CudaSlice<f32>,
18434        v: &CudaSlice<f32>,
18435        o: &mut CudaSlice<f32>,
18436        head_dim: usize,
18437        n_head: usize,
18438        n_head_kv: usize,
18439        t: usize,
18440        t_kv: usize,
18441        scale: f32,
18442        causal: bool,
18443        window: usize,
18444    ) -> Result<(), Box<dyn std::error::Error>> {
18445        let kv_lo = if window > 0 {
18446            (t_kv - t + 1).saturating_sub(window)
18447        } else {
18448            0
18449        };
18450        let smem = (t_kv - kv_lo) * 4;
18451        if smem > 48 * 1024 {
18452            return Err(format!(
18453                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18454                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18455                 a window this wide needs the multi-pass long-ctx kernel"
18456            )
18457            .into());
18458        }
18459        let f = self.func("sdpa_naive_w_lo_f32");
18460        let cfg = LaunchConfig {
18461            grid_dim: (n_head as u32, t as u32, 1),
18462            block_dim: (128, 1, 1),
18463            shared_mem_bytes: smem as u32,
18464        };
18465        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18466            head_dim as i32,
18467            n_head as i32,
18468            n_head_kv as i32,
18469            t as i32,
18470            t_kv as i32,
18471            causal as i32,
18472            window as i32,
18473            kv_lo as i32,
18474        );
18475        let __s_b = self.gpu.stream();
18476        let mut b = __s_b.launch_builder(&f);
18477        b.arg(q)
18478            .arg(k)
18479            .arg(v)
18480            .arg(o)
18481            .arg(&hd)
18482            .arg(&nh)
18483            .arg(&nhkv)
18484            .arg(&ti)
18485            .arg(&tkvi)
18486            .arg(&scale)
18487            .arg(&cz)
18488            .arg(&wi)
18489            .arg(&lo);
18490        unsafe {
18491            b.launch(cfg)?;
18492        }
18493        Ok(())
18494    }
18495
18496    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
18497    pub fn sdpa_naive_view(
18498        &self,
18499        q: &CudaSlice<f32>,
18500        k: &cudarc::driver::CudaView<f32>,
18501        v: &cudarc::driver::CudaView<f32>,
18502        o: &mut CudaSlice<f32>,
18503        head_dim: usize,
18504        n_head: usize,
18505        n_head_kv: usize,
18506        t: usize,
18507        t_kv: usize,
18508        scale: f32,
18509        causal: bool,
18510    ) -> Result<(), Box<dyn std::error::Error>> {
18511        let f = self.func("sdpa_naive_f32");
18512        let cfg = LaunchConfig {
18513            grid_dim: (n_head as u32, t as u32, 1),
18514            block_dim: (128, 1, 1),
18515            shared_mem_bytes: (t_kv * 4) as u32,
18516        };
18517        let (hd, nh, nhkv, ti, tkvi, cz) = (
18518            head_dim as i32,
18519            n_head as i32,
18520            n_head_kv as i32,
18521            t as i32,
18522            t_kv as i32,
18523            causal as i32,
18524        );
18525        let __s_b = self.gpu.stream();
18526        let mut b = __s_b.launch_builder(&f);
18527        b.arg(q)
18528            .arg(k)
18529            .arg(v)
18530            .arg(o)
18531            .arg(&hd)
18532            .arg(&nh)
18533            .arg(&nhkv)
18534            .arg(&ti)
18535            .arg(&tkvi)
18536            .arg(&scale)
18537            .arg(&cz);
18538        unsafe {
18539            b.launch(cfg)?;
18540        }
18541        Ok(())
18542    }
18543
18544    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
18545    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
18546    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
18547    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
18548    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
18549    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
18550    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
18551    #[allow(clippy::too_many_arguments)]
18552    pub fn fa_dequant_kv_view_f32(
18553        &self,
18554        k: &cudarc::driver::CudaView<u8>,
18555        v: &cudarc::driver::CudaView<u8>,
18556        kf: &mut CudaSlice<f32>,
18557        vf: &mut CudaSlice<f32>,
18558        kv_dim_k: usize,
18559        kv_dim_v: usize,
18560        t_kv: usize,
18561        k_tok_bytes: usize,
18562        v_tok_bytes: usize,
18563        g: bool,
18564    ) -> Result<(), Box<dyn std::error::Error>> {
18565        let f = if g {
18566            self.func_g("fa_dequant_kv_ws_f32")
18567        } else {
18568            self.func("fa_dequant_kv_ws_f32")
18569        };
18570        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18571        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18572        let cfg = LaunchConfig {
18573            grid_dim: (nblk.max(1), 1, 1),
18574            block_dim: (256, 1, 1),
18575            shared_mem_bytes: 0,
18576        };
18577        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18578        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18579        let __s_b = self.gpu.stream();
18580        let mut b = __s_b.launch_builder(&f);
18581        b.arg(k)
18582            .arg(v)
18583            .arg(&mut *kf)
18584            .arg(&mut *vf)
18585            .arg(&kdk)
18586            .arg(&kdv)
18587            .arg(&tkvi)
18588            .arg(&ktb)
18589            .arg(&vtb);
18590        unsafe {
18591            b.launch(cfg)?;
18592        }
18593        Ok(())
18594    }
18595
18596    #[allow(clippy::too_many_arguments)]
18597    pub fn sdpa_naive_quantized_view(
18598        &self,
18599        q: &CudaSlice<f32>,
18600        k: &cudarc::driver::CudaView<u8>,
18601        v: &cudarc::driver::CudaView<u8>,
18602        o: &mut CudaSlice<f32>,
18603        head_dim: usize,
18604        n_head: usize,
18605        n_head_kv: usize,
18606        t: usize,
18607        t_kv: usize,
18608        scale: f32,
18609        causal: bool,
18610        k_tok_bytes: usize,
18611        v_tok_bytes: usize,
18612    ) -> Result<(), Box<dyn std::error::Error>> {
18613        let kv_dim = n_head_kv * head_dim;
18614        let mut kf = self.uninit(t_kv * kv_dim)?;
18615        let mut vf = self.uninit(t_kv * kv_dim)?;
18616        let f = self.func("fa_dequant_kv_ws_f32");
18617        let total = (2 * t_kv * kv_dim) as u64;
18618        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18619        let cfg = LaunchConfig {
18620            grid_dim: (nblk.max(1), 1, 1),
18621            block_dim: (256, 1, 1),
18622            shared_mem_bytes: 0,
18623        };
18624        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18625        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18626        let __s_b = self.gpu.stream();
18627        let mut b = __s_b.launch_builder(&f);
18628        b.arg(k)
18629            .arg(v)
18630            .arg(&mut kf)
18631            .arg(&mut vf)
18632            .arg(&kv_dim_i)
18633            .arg(&kv_dim_i)
18634            .arg(&t_kv_i)
18635            .arg(&k_tok_bytes_i)
18636            .arg(&v_tok_bytes_i);
18637        unsafe { b.launch(cfg)? };
18638        self.sdpa_naive(
18639            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18640        )
18641    }
18642
18643    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
18644    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
18645    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
18646    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
18647    /// unwindowed function above and produces bit-identical output at window == 0.
18648    ///
18649    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
18650    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
18651    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
18652    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
18653    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
18654    #[allow(clippy::too_many_arguments)]
18655    pub fn sdpa_naive_w_quantized_view(
18656        &self,
18657        q: &CudaSlice<f32>,
18658        k: &cudarc::driver::CudaView<u8>,
18659        v: &cudarc::driver::CudaView<u8>,
18660        o: &mut CudaSlice<f32>,
18661        head_dim: usize,
18662        n_head: usize,
18663        n_head_kv: usize,
18664        t: usize,
18665        t_kv: usize,
18666        scale: f32,
18667        causal: bool,
18668        window: usize,
18669        k_tok_bytes: usize,
18670        v_tok_bytes: usize,
18671    ) -> Result<(), Box<dyn std::error::Error>> {
18672        let kv_dim = n_head_kv * head_dim;
18673        let mut kf = self.uninit(t_kv * kv_dim)?;
18674        let mut vf = self.uninit(t_kv * kv_dim)?;
18675        let f = self.func("fa_dequant_kv_ws_f32");
18676        let total = (2 * t_kv * kv_dim) as u64;
18677        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18678        let cfg = LaunchConfig {
18679            grid_dim: (nblk.max(1), 1, 1),
18680            block_dim: (256, 1, 1),
18681            shared_mem_bytes: 0,
18682        };
18683        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18684        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18685        let __s_b = self.gpu.stream();
18686        let mut b = __s_b.launch_builder(&f);
18687        b.arg(k)
18688            .arg(v)
18689            .arg(&mut kf)
18690            .arg(&mut vf)
18691            .arg(&kv_dim_i)
18692            .arg(&kv_dim_i)
18693            .arg(&t_kv_i)
18694            .arg(&k_tok_bytes_i)
18695            .arg(&v_tok_bytes_i);
18696        unsafe { b.launch(cfg)? };
18697        self.sdpa_naive_w(
18698            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18699        )
18700    }
18701
18702    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
18703    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
18704    /// Q/K/V/O [head_dim, n_head(_kv), T].
18705    pub fn fa_prefill(
18706        &self,
18707        q: &CudaSlice<f32>,
18708        k: &CudaSlice<f32>,
18709        v: &CudaSlice<f32>,
18710        o: &mut CudaSlice<f32>,
18711        head_dim: usize,
18712        n_head: usize,
18713        n_head_kv: usize,
18714        t: usize,
18715        t_kv: usize,
18716        scale: f32,
18717        causal: bool,
18718    ) -> Result<(), Box<dyn std::error::Error>> {
18719        if portable_mma_gated() {
18720            return self.sdpa_naive(
18721                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18722            );
18723        }
18724        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
18725        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
18726        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
18727        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
18728        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
18729        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
18730        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
18731        let fa3_on = head_dim == 256
18732            && causal
18733            && t == t_kv
18734            && match std::env::var("MEMRA_FA3").as_deref() {
18735                Ok("0") => false,
18736                // The force arm consults the arch now: the bf16 stage below calls
18737                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
18738                // a portable build. Refuse at the switch, not at the lookup.
18739                Ok("1") => {
18740                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
18741                    true
18742                }
18743                _ => cfg!(memra_hopper_mma),
18744            };
18745        if fa3_on {
18746            let n = t * n_head * head_dim;
18747            let nkv = t * n_head_kv * head_dim;
18748            let mut q16 = self.alloc_u8_uninit(n * 2)?;
18749            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18750            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18751            self.f32_to_bf16_into(q, &mut q16, n)?;
18752            self.f32_to_bf16_into(k, &mut k16, nkv)?;
18753            self.f32_to_bf16_into(v, &mut v16, nkv)?;
18754            let rc = {
18755                use cudarc::driver::{DevicePtr, DevicePtrMut};
18756                let stream = self.gpu.stream();
18757                let (qp, _g1) = q16.device_ptr(&stream);
18758                let (kp, _g2) = k16.device_ptr(&stream);
18759                let (vp, _g3) = v16.device_ptr(&stream);
18760                let (op, _g4) = o.device_ptr_mut(&stream);
18761                unsafe {
18762                    memra_fa3_prefill(
18763                        qp as *const core::ffi::c_void,
18764                        kp as *const core::ffi::c_void,
18765                        vp as *const core::ffi::c_void,
18766                        op as *mut f32,
18767                        t as i32,
18768                        n_head as i32,
18769                        n_head_kv as i32,
18770                        head_dim as i32,
18771                        scale,
18772                        stream.cu_stream() as *mut core::ffi::c_void,
18773                    )
18774                }
18775            };
18776            if rc != 0 {
18777                return Err(format!("memra_fa3_prefill rc={rc}").into());
18778            }
18779            return Ok(());
18780        }
18781        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
18782        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
18783        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
18784        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
18785        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18786        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18787        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18788            const BLOCK_Q: usize = 64;
18789            const BKX: usize = 32;
18790            let f = self.func("fa_prefill_bf16_p1");
18791            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18792                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18793            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18794            f.set_attribute(
18795                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18796                shmem as i32,
18797            )?;
18798            let cfg = LaunchConfig {
18799                grid_dim: (
18800                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18801                    n_head as u32,
18802                    1,
18803                ),
18804                block_dim: (32, 4, 1),
18805                shared_mem_bytes: shmem,
18806            };
18807            let (hd, nh, nhkv, ti, tkvi, cz) = (
18808                head_dim as i32,
18809                n_head as i32,
18810                n_head_kv as i32,
18811                t as i32,
18812                t_kv as i32,
18813                causal as i32,
18814            );
18815            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18816            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18817            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18818            let __s_b = self.gpu.stream();
18819            let mut b = __s_b.launch_builder(&f);
18820            b.arg(&qb)
18821                .arg(&kb)
18822                .arg(&vb)
18823                .arg(o)
18824                .arg(&hd)
18825                .arg(&nh)
18826                .arg(&nhkv)
18827                .arg(&ti)
18828                .arg(&tkvi)
18829                .arg(&scale)
18830                .arg(&cz);
18831            unsafe {
18832                b.launch(cfg)?;
18833            }
18834            return Ok(());
18835        }
18836        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
18837        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
18838        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
18839        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
18840        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
18841        const BK: usize = 32;
18842        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
18843        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
18844        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
18845        let (block_q, warps, w2_sfx): (usize, u32, &str) =
18846            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
18847        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
18848        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
18849        // other head_dims to sdpa_naive before reaching here.
18850        let hd_sfx = fa_hd_suffix(head_dim)?;
18851        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18852        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
18853        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
18854        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
18855        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
18856        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
18857        let (kb16, vb16) = if bf16kv {
18858            let n = t_kv * n_head_kv * head_dim;
18859            let mut kb = self.alloc_u8_uninit(n * 2)?;
18860            let mut vb = self.alloc_u8_uninit(n * 2)?;
18861            let fcv = self.func("f32_to_bf16_bulk");
18862            let ni = n as i64;
18863            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
18864            let __s_b = self.gpu.stream();
18865            let mut b = __s_b.launch_builder(&fcv);
18866            b.arg(k).arg(&mut kb).arg(&ni);
18867            unsafe {
18868                b.launch(cfgc)?;
18869            }
18870            let __s_b = self.gpu.stream();
18871            let mut b = __s_b.launch_builder(&fcv);
18872            b.arg(v).arg(&mut vb).arg(&ni);
18873            unsafe {
18874                b.launch(cfgc)?;
18875            }
18876            (Some(kb), Some(vb))
18877        } else {
18878            (None, None)
18879        };
18880        let f = self.func(&if bf16kv {
18881            format!("fa_prefill_bf16kv_pp{hd_sfx}")
18882        } else {
18883            format!(
18884                "fa_prefill_f32{}{}{hd_sfx}",
18885                if floor { "" } else { "_pp" },
18886                if floor { "" } else { w2_sfx }
18887            )
18888        });
18889        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
18890        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
18891        let kv_stages = if bf16kv { 2 } else { 1 };
18892        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
18893            + 4 * (block_q * BK + 2 * block_q)) as u32;
18894        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18895        f.set_attribute(
18896            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18897            shmem as i32,
18898        )?;
18899        let cfg = LaunchConfig {
18900            grid_dim: (
18901                (t as u32 + block_q as u32 - 1) / block_q as u32,
18902                n_head as u32,
18903                1,
18904            ),
18905            block_dim: (32, warps, 1),
18906            shared_mem_bytes: shmem,
18907        };
18908        let (hd, nh, nhkv, ti, tkvi, cz) = (
18909            head_dim as i32,
18910            n_head as i32,
18911            n_head_kv as i32,
18912            t as i32,
18913            t_kv as i32,
18914            causal as i32,
18915        );
18916        let __s_b = self.gpu.stream();
18917        let mut b = __s_b.launch_builder(&f);
18918        b.arg(q);
18919        match (&kb16, &vb16) {
18920            (Some(kb), Some(vb)) => {
18921                b.arg(kb).arg(vb);
18922            }
18923            _ => {
18924                b.arg(k).arg(v);
18925            }
18926        }
18927        b.arg(o)
18928            .arg(&hd)
18929            .arg(&nh)
18930            .arg(&nhkv)
18931            .arg(&ti)
18932            .arg(&tkvi)
18933            .arg(&scale)
18934            .arg(&cz);
18935        unsafe {
18936            b.launch(cfg)?;
18937        }
18938        Ok(())
18939    }
18940
18941    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
18942    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
18943    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
18944    #[allow(clippy::too_many_arguments)]
18945    pub fn fa_prefill_w(
18946        &self,
18947        q: &CudaSlice<f32>,
18948        k: &CudaSlice<f32>,
18949        v: &CudaSlice<f32>,
18950        o: &mut CudaSlice<f32>,
18951        head_dim: usize,
18952        n_head: usize,
18953        n_head_kv: usize,
18954        t: usize,
18955        t_kv: usize,
18956        scale: f32,
18957        causal: bool,
18958        window: usize,
18959    ) -> Result<(), Box<dyn std::error::Error>> {
18960        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
18961        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
18962        if portable_mma_gated() {
18963            return self.sdpa_naive_w(
18964                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18965            );
18966        }
18967        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
18968        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
18969        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
18970        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18971        let faw_f32 =
18972            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
18973        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18974        self.fa_prefill_w_arm(
18975            q,
18976            k,
18977            v,
18978            o,
18979            head_dim,
18980            n_head,
18981            n_head_kv,
18982            t,
18983            t_kv,
18984            scale,
18985            causal,
18986            window,
18987            floor || faw_f32,
18988            floor,
18989        )
18990    }
18991
18992    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
18993    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
18994    #[allow(clippy::too_many_arguments)]
18995    pub fn fa_prefill_w_pre(
18996        &self,
18997        qb: &CudaSlice<u8>,
18998        kb: &CudaSlice<u8>,
18999        vb: &CudaSlice<u8>,
19000        o: &mut CudaSlice<f32>,
19001        head_dim: usize,
19002        n_head: usize,
19003        n_head_kv: usize,
19004        t: usize,
19005        t_kv: usize,
19006        scale: f32,
19007        causal: bool,
19008        window: usize,
19009        v_f16: bool,
19010    ) -> Result<(), Box<dyn std::error::Error>> {
19011        const BLOCK_Q: usize = 64;
19012        const BK: usize = 32;
19013        debug_assert_eq!(head_dim, 256);
19014        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19015        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19016        if hp {
19017            const BLOCK_QH: usize = 32;
19018            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
19019            // else re-encode through the pooled scratch (stream-ordered reuse).
19020            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19021            let vh: &CudaSlice<u8> = if v_f16 {
19022                vb
19023            } else {
19024                let n = t_kv * n_head_kv * head_dim;
19025                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19026                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19027                }
19028                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19029                vguard.as_ref().unwrap()
19030            };
19031            let f = self.func("fa_prefill_w_bf16_p1h2");
19032            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19033            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19034            f.set_attribute(
19035                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19036                shmem as i32,
19037            )?;
19038            let cfg = LaunchConfig {
19039                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19040                block_dim: (32, 4, 1),
19041                shared_mem_bytes: shmem,
19042            };
19043            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19044                head_dim as i32,
19045                n_head as i32,
19046                n_head_kv as i32,
19047                t as i32,
19048                t_kv as i32,
19049                causal as i32,
19050                window as i32,
19051            );
19052            let __s_b = self.gpu.stream();
19053            let mut b = __s_b.launch_builder(&f);
19054            b.arg(qb)
19055                .arg(kb)
19056                .arg(vh)
19057                .arg(o)
19058                .arg(&hd)
19059                .arg(&nh)
19060                .arg(&nhkv)
19061                .arg(&ti)
19062                .arg(&tkvi)
19063                .arg(&scale)
19064                .arg(&cz)
19065                .arg(&wi);
19066            unsafe {
19067                b.launch(cfg)?;
19068            }
19069            return Ok(());
19070        }
19071        let f = self.func("fa_prefill_w_bf16_p1");
19072        let shmem =
19073            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19074        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19075        f.set_attribute(
19076            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19077            shmem as i32,
19078        )?;
19079        let cfg = LaunchConfig {
19080            grid_dim: (
19081                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19082                n_head as u32,
19083                1,
19084            ),
19085            block_dim: (32, 4, 1),
19086            shared_mem_bytes: shmem,
19087        };
19088        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19089            head_dim as i32,
19090            n_head as i32,
19091            n_head_kv as i32,
19092            t as i32,
19093            t_kv as i32,
19094            causal as i32,
19095            window as i32,
19096        );
19097        let __s_b = self.gpu.stream();
19098        let mut b = __s_b.launch_builder(&f);
19099        b.arg(qb)
19100            .arg(kb)
19101            .arg(vb)
19102            .arg(o)
19103            .arg(&hd)
19104            .arg(&nh)
19105            .arg(&nhkv)
19106            .arg(&ti)
19107            .arg(&tkvi)
19108            .arg(&scale)
19109            .arg(&cz)
19110            .arg(&wi);
19111        unsafe {
19112            b.launch(cfg)?;
19113        }
19114        Ok(())
19115    }
19116
19117    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
19118    #[allow(clippy::too_many_arguments)]
19119    pub fn fa_prefill_w_arm(
19120        &self,
19121        q: &CudaSlice<f32>,
19122        k: &CudaSlice<f32>,
19123        v: &CudaSlice<f32>,
19124        o: &mut CudaSlice<f32>,
19125        head_dim: usize,
19126        n_head: usize,
19127        n_head_kv: usize,
19128        t: usize,
19129        t_kv: usize,
19130        scale: f32,
19131        causal: bool,
19132        window: usize,
19133        f32_stage: bool,
19134        floor: bool,
19135    ) -> Result<(), Box<dyn std::error::Error>> {
19136        const BLOCK_Q: usize = 64;
19137        const BK: usize = 32;
19138        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19139        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
19140        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
19141        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
19142        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19143        let p1 = !floor
19144            && !f32_stage
19145            && *P1_ON.get_or_init(|| {
19146                std::env::var("MEMRA_FAW_P1")
19147                    .map(|v| v != "0")
19148                    .unwrap_or(true)
19149            });
19150        let hp =
19151            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19152        if hp {
19153            const BLOCK_QH: usize = 32;
19154            let f = self.func("fa_prefill_w_bf16_p1h2");
19155            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19156            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19157            f.set_attribute(
19158                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19159                shmem as i32,
19160            )?;
19161            let cfg = LaunchConfig {
19162                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19163                block_dim: (32, 4, 1),
19164                shared_mem_bytes: shmem,
19165            };
19166            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19167                head_dim as i32,
19168                n_head as i32,
19169                n_head_kv as i32,
19170                t as i32,
19171                t_kv as i32,
19172                causal as i32,
19173                window as i32,
19174            );
19175            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19176            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19177            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19178            let __s_b = self.gpu.stream();
19179            let mut b = __s_b.launch_builder(&f);
19180            b.arg(&qb)
19181                .arg(&kb)
19182                .arg(&vh)
19183                .arg(o)
19184                .arg(&hd)
19185                .arg(&nh)
19186                .arg(&nhkv)
19187                .arg(&ti)
19188                .arg(&tkvi)
19189                .arg(&scale)
19190                .arg(&cz)
19191                .arg(&wi);
19192            unsafe {
19193                b.launch(cfg)?;
19194            }
19195            return Ok(());
19196        }
19197        if p1 {
19198            let f = self.func("fa_prefill_w_bf16_p1");
19199            let shmem =
19200                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19201            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19202            f.set_attribute(
19203                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19204                shmem as i32,
19205            )?;
19206            let cfg = LaunchConfig {
19207                grid_dim: (
19208                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19209                    n_head as u32,
19210                    1,
19211                ),
19212                block_dim: (32, 4, 1),
19213                shared_mem_bytes: shmem,
19214            };
19215            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19216                head_dim as i32,
19217                n_head as i32,
19218                n_head_kv as i32,
19219                t as i32,
19220                t_kv as i32,
19221                causal as i32,
19222                window as i32,
19223            );
19224            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19225            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19226            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19227            let __s_b = self.gpu.stream();
19228            let mut b = __s_b.launch_builder(&f);
19229            b.arg(&qb)
19230                .arg(&kb)
19231                .arg(&vb)
19232                .arg(o)
19233                .arg(&hd)
19234                .arg(&nh)
19235                .arg(&nhkv)
19236                .arg(&ti)
19237                .arg(&tkvi)
19238                .arg(&scale)
19239                .arg(&cz)
19240                .arg(&wi);
19241            unsafe {
19242                b.launch(cfg)?;
19243            }
19244            return Ok(());
19245        }
19246        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
19247        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
19248        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19249        let g4 = !floor
19250            && !f32_stage
19251            && n_head_kv == 1
19252            && n_head % 4 == 0
19253            && *G4_ON.get_or_init(|| {
19254                std::env::var("MEMRA_FAW_G4")
19255                    .map(|v| v != "0")
19256                    .unwrap_or(true)
19257            });
19258        if g4 {
19259            const SP_M: usize = 16;
19260            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
19261            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
19262            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19263            let o2 = *O2_ON.get_or_init(|| {
19264                std::env::var("MEMRA_FAW_O2")
19265                    .map(|v| v != "0")
19266                    .unwrap_or(true)
19267            });
19268            let f = self.func(if o2 {
19269                "fa_prefill_w_bf16_g4o2"
19270            } else {
19271                "fa_prefill_w_bf16_g4"
19272            });
19273            let shmem = if o2 {
19274                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19275            } else {
19276                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19277                    as u32
19278            };
19279            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19280            f.set_attribute(
19281                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19282                shmem as i32,
19283            )?;
19284            let cfg = LaunchConfig {
19285                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19286                block_dim: (32, 4, 1),
19287                shared_mem_bytes: shmem,
19288            };
19289            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19290                head_dim as i32,
19291                n_head as i32,
19292                n_head_kv as i32,
19293                t as i32,
19294                t_kv as i32,
19295                causal as i32,
19296                window as i32,
19297            );
19298            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19299            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19300            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19301            let __s_b = self.gpu.stream();
19302            let mut b = __s_b.launch_builder(&f);
19303            b.arg(&qb)
19304                .arg(&kb)
19305                .arg(&vb)
19306                .arg(o)
19307                .arg(&hd)
19308                .arg(&nh)
19309                .arg(&nhkv)
19310                .arg(&ti)
19311                .arg(&tkvi)
19312                .arg(&scale)
19313                .arg(&cz)
19314                .arg(&wi);
19315            unsafe {
19316                b.launch(cfg)?;
19317            }
19318            return Ok(());
19319        }
19320        let f = self.func(if floor {
19321            "fa_prefill_w_f32"
19322        } else if f32_stage {
19323            "fa_prefill_w_f32_pp"
19324        } else {
19325            "fa_prefill_w_bf16_pp"
19326        });
19327        let shmem =
19328            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19329        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19330        f.set_attribute(
19331            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19332            shmem as i32,
19333        )?;
19334        let cfg = LaunchConfig {
19335            grid_dim: (
19336                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19337                n_head as u32,
19338                1,
19339            ),
19340            block_dim: (32, 4, 1),
19341            shared_mem_bytes: shmem,
19342        };
19343        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19344            head_dim as i32,
19345            n_head as i32,
19346            n_head_kv as i32,
19347            t as i32,
19348            t_kv as i32,
19349            causal as i32,
19350            window as i32,
19351        );
19352        if f32_stage {
19353            let __s_b = self.gpu.stream();
19354            let mut b = __s_b.launch_builder(&f);
19355            b.arg(q)
19356                .arg(k)
19357                .arg(v)
19358                .arg(o)
19359                .arg(&hd)
19360                .arg(&nh)
19361                .arg(&nhkv)
19362                .arg(&ti)
19363                .arg(&tkvi)
19364                .arg(&scale)
19365                .arg(&cz)
19366                .arg(&wi);
19367            unsafe {
19368                b.launch(cfg)?;
19369            }
19370        } else {
19371            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19372            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19373            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19374            let __s_b = self.gpu.stream();
19375            let mut b = __s_b.launch_builder(&f);
19376            b.arg(&qb)
19377                .arg(&kb)
19378                .arg(&vb)
19379                .arg(o)
19380                .arg(&hd)
19381                .arg(&nh)
19382                .arg(&nhkv)
19383                .arg(&ti)
19384                .arg(&tkvi)
19385                .arg(&scale)
19386                .arg(&cz)
19387                .arg(&wi);
19388            unsafe {
19389                b.launch(cfg)?;
19390            }
19391        }
19392        Ok(())
19393    }
19394
19395    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
19396    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
19397    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
19398    #[allow(clippy::too_many_arguments)]
19399    pub fn fa_prefill_hd512(
19400        &self,
19401        q: &CudaSlice<f32>,
19402        k: &CudaSlice<f32>,
19403        v: &CudaSlice<f32>,
19404        o: &mut CudaSlice<f32>,
19405        head_dim: usize,
19406        n_head: usize,
19407        n_head_kv: usize,
19408        t: usize,
19409        t_kv: usize,
19410        scale: f32,
19411        causal: bool,
19412    ) -> Result<(), Box<dyn std::error::Error>> {
19413        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
19414        if portable_mma_gated() {
19415            return self.sdpa_naive(
19416                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19417            );
19418        }
19419        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
19420        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
19421        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
19422        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
19423        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
19424        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19425        let f32_stage =
19426            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19427        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
19428        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
19429        // Own numeric config (partial-sum order) — battery-gated.
19430        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19431        let sp = !f32_stage
19432            && *SP_ON.get_or_init(|| {
19433                std::env::var("MEMRA_FA512_SP")
19434                    .map(|v| v != "0")
19435                    .unwrap_or(true)
19436            });
19437        self.fa_prefill_hd512_arm(
19438            q,
19439            k,
19440            v,
19441            o,
19442            head_dim,
19443            n_head,
19444            n_head_kv,
19445            t,
19446            t_kv,
19447            scale,
19448            causal,
19449            f32_stage,
19450            sp,
19451            sp && fa_f16pv_on(),
19452        )
19453    }
19454
19455    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
19456    #[allow(clippy::too_many_arguments)]
19457    pub fn fa_prefill_hd512_pre(
19458        &self,
19459        qb: &CudaSlice<u8>,
19460        kb: &CudaSlice<u8>,
19461        vb: &CudaSlice<u8>,
19462        o: &mut CudaSlice<f32>,
19463        head_dim: usize,
19464        n_head: usize,
19465        n_head_kv: usize,
19466        t: usize,
19467        t_kv: usize,
19468        scale: f32,
19469        causal: bool,
19470        v_f16: bool,
19471    ) -> Result<(), Box<dyn std::error::Error>> {
19472        debug_assert_eq!(head_dim, 512);
19473        const SP_M: usize = 16;
19474        const BKS: usize = 32;
19475        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
19476        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
19477        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
19478        let f16pv = fa_f16pv_on();
19479        let nw = if f16pv { fa512_wide_warps() } else { 2 };
19480        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19481        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19482        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19483        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19484            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
19485            let n = t_kv * n_head_kv * head_dim;
19486            let need = n * 2;
19487            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19488                *vguard = Some(self.alloc_uninit::<u8>(need)?);
19489            }
19490            let dst = vguard.as_mut().unwrap();
19491            self.bf16_to_f16_into(vb, n, dst)?;
19492            vguard.as_ref().unwrap()
19493        } else {
19494            vb
19495        };
19496        let f = self.func(if hp {
19497            "fa_prefill_bf16_hd512_sp16h2"
19498        } else {
19499            match (f16pv, nw) {
19500                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19501                (true, _) => "fa_prefill_bf16_hd512_sp16",
19502                _ => "fa_prefill_bf16_hd512_sp",
19503            }
19504        });
19505        let (nwarp, npart) = if hp {
19506            (4usize, 4usize)
19507        } else if nw > 2 {
19508            (nw, nw)
19509        } else {
19510            (2, 1)
19511        };
19512        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
19513        let shmem = if hp {
19514            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19515                as u32
19516        } else {
19517            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19518                + 4 * (npart * SP_M * BKS + SP_M)) as u32
19519        };
19520        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19521        f.set_attribute(
19522            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19523            shmem as i32,
19524        )?;
19525        let grid_y = if hp {
19526            (n_head / 2) as u32
19527        } else {
19528            n_head as u32
19529        };
19530        let cfg = LaunchConfig {
19531            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19532            block_dim: (32, nwarp as u32, 1),
19533            shared_mem_bytes: shmem,
19534        };
19535        let (hd, nh, nhkv, ti, tkvi, cz) = (
19536            head_dim as i32,
19537            n_head as i32,
19538            n_head_kv as i32,
19539            t as i32,
19540            t_kv as i32,
19541            causal as i32,
19542        );
19543        let __s_b = self.gpu.stream();
19544        let mut b = __s_b.launch_builder(&f);
19545        b.arg(qb)
19546            .arg(kb)
19547            .arg(vref)
19548            .arg(o)
19549            .arg(&hd)
19550            .arg(&nh)
19551            .arg(&nhkv)
19552            .arg(&ti)
19553            .arg(&tkvi)
19554            .arg(&scale)
19555            .arg(&cz);
19556        unsafe {
19557            b.launch(cfg)?;
19558        }
19559        Ok(())
19560    }
19561
19562    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
19563    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
19564    #[allow(clippy::too_many_arguments)]
19565    pub fn fa_prefill_hd512_arm(
19566        &self,
19567        q: &CudaSlice<f32>,
19568        k: &CudaSlice<f32>,
19569        v: &CudaSlice<f32>,
19570        o: &mut CudaSlice<f32>,
19571        head_dim: usize,
19572        n_head: usize,
19573        n_head_kv: usize,
19574        t: usize,
19575        t_kv: usize,
19576        scale: f32,
19577        causal: bool,
19578        f32_stage: bool,
19579        sp: bool,
19580        f16pv: bool,
19581    ) -> Result<(), Box<dyn std::error::Error>> {
19582        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19583        if sp && !f32_stage {
19584            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
19585            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
19586            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
19587            const SP_M: usize = 16;
19588            const BKS: usize = 32;
19589            let nw = if f16pv { fa512_wide_warps() } else { 2 };
19590            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19591            let f = self.func(if hp {
19592                "fa_prefill_bf16_hd512_sp16h2"
19593            } else {
19594                match (f16pv, nw) {
19595                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19596                    (true, _) => "fa_prefill_bf16_hd512_sp16",
19597                    _ => "fa_prefill_bf16_hd512_sp",
19598                }
19599            });
19600            let (nwarp, npart) = if hp {
19601                (4usize, 4usize)
19602            } else if nw > 2 {
19603                (nw, nw)
19604            } else {
19605                (2, 1)
19606            };
19607            let shmem = if hp {
19608                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19609                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19610            } else {
19611                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19612                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
19613            };
19614            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19615            f.set_attribute(
19616                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19617                shmem as i32,
19618            )?;
19619            let grid_y = if hp {
19620                (n_head / 2) as u32
19621            } else {
19622                n_head as u32
19623            };
19624            let cfg = LaunchConfig {
19625                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19626                block_dim: (32, nwarp as u32, 1),
19627                shared_mem_bytes: shmem,
19628            };
19629            let (hd, nh, nhkv, ti, tkvi, cz) = (
19630                head_dim as i32,
19631                n_head as i32,
19632                n_head_kv as i32,
19633                t as i32,
19634                t_kv as i32,
19635                causal as i32,
19636            );
19637            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19638            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19639            let vb = if f16pv {
19640                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19641            } else {
19642                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19643            };
19644            let __s_b = self.gpu.stream();
19645            let mut b = __s_b.launch_builder(&f);
19646            b.arg(&qb)
19647                .arg(&kb)
19648                .arg(&vb)
19649                .arg(o)
19650                .arg(&hd)
19651                .arg(&nh)
19652                .arg(&nhkv)
19653                .arg(&ti)
19654                .arg(&tkvi)
19655                .arg(&scale)
19656                .arg(&cz);
19657            unsafe {
19658                b.launch(cfg)?;
19659            }
19660            return Ok(());
19661        }
19662        const BLOCK_Q: usize = 32;
19663        const BK: usize = 32;
19664        const HALF: usize = 256;
19665        let f = self.func(if f32_stage {
19666            "fa_prefill_f32_hd512"
19667        } else {
19668            "fa_prefill_bf16_hd512"
19669        });
19670        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
19671        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19672            + 4 * BLOCK_Q) as u32;
19673        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19674        f.set_attribute(
19675            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19676            shmem as i32,
19677        )?;
19678        let cfg = LaunchConfig {
19679            grid_dim: (
19680                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19681                n_head as u32,
19682                2,
19683            ),
19684            block_dim: (32, 2, 1),
19685            shared_mem_bytes: shmem,
19686        };
19687        let (hd, nh, nhkv, ti, tkvi, cz) = (
19688            head_dim as i32,
19689            n_head as i32,
19690            n_head_kv as i32,
19691            t as i32,
19692            t_kv as i32,
19693            causal as i32,
19694        );
19695        if f32_stage {
19696            let __s_b = self.gpu.stream();
19697            let mut b = __s_b.launch_builder(&f);
19698            b.arg(q)
19699                .arg(k)
19700                .arg(v)
19701                .arg(o)
19702                .arg(&hd)
19703                .arg(&nh)
19704                .arg(&nhkv)
19705                .arg(&ti)
19706                .arg(&tkvi)
19707                .arg(&scale)
19708                .arg(&cz);
19709            unsafe {
19710                b.launch(cfg)?;
19711            }
19712        } else {
19713            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19714            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19715            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19716            let __s_b = self.gpu.stream();
19717            let mut b = __s_b.launch_builder(&f);
19718            b.arg(&qb)
19719                .arg(&kb)
19720                .arg(&vb)
19721                .arg(o)
19722                .arg(&hd)
19723                .arg(&nh)
19724                .arg(&nhkv)
19725                .arg(&ti)
19726                .arg(&tkvi)
19727                .arg(&scale)
19728                .arg(&cz);
19729            unsafe {
19730                b.launch(cfg)?;
19731            }
19732        }
19733        Ok(())
19734    }
19735
19736    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
19737    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
19738    /// separate f32_to_bf16 the FA entries would run).
19739    #[allow(clippy::too_many_arguments)]
19740    pub fn rope_neox2_bf16e(
19741        &self,
19742        q: &mut CudaSlice<f32>,
19743        k: &mut CudaSlice<f32>,
19744        qb: &mut CudaSlice<u8>,
19745        kb: &mut CudaSlice<u8>,
19746        pos: &CudaSlice<i32>,
19747        head_dim: usize,
19748        n_dims: usize,
19749        nh_q: usize,
19750        nh_k: usize,
19751        n_tokens: usize,
19752        base: f32,
19753        freq_scale: f32,
19754        ff: Option<&CudaSlice<f32>>,
19755    ) -> Result<(), Box<dyn std::error::Error>> {
19756        let f = self.func("rope_neox2_bf16e_f32");
19757        let rows = ((nh_q + nh_k) * n_tokens) as u32;
19758        let cfg = LaunchConfig {
19759            grid_dim: (rows, 1, 1),
19760            block_dim: ((head_dim / 2) as u32, 1, 1),
19761            shared_mem_bytes: 0,
19762        };
19763        let theta_scale = base.powf(-2.0 / n_dims as f32);
19764        let (hd, nd, nhq, nhk, nt) = (
19765            head_dim as i32,
19766            n_dims as i32,
19767            nh_q as i32,
19768            nh_k as i32,
19769            n_tokens as i32,
19770        );
19771        let __s_b = self.gpu.stream();
19772        let mut b = __s_b.launch_builder(&f);
19773        match ff {
19774            Some(t) => {
19775                b.arg(&mut *q)
19776                    .arg(&mut *k)
19777                    .arg(&mut *qb)
19778                    .arg(&mut *kb)
19779                    .arg(pos)
19780                    .arg(&hd)
19781                    .arg(&nd)
19782                    .arg(&nhq)
19783                    .arg(&nhk)
19784                    .arg(&nt)
19785                    .arg(&theta_scale)
19786                    .arg(&freq_scale)
19787                    .arg(t);
19788                unsafe {
19789                    b.launch(cfg)?;
19790                }
19791            }
19792            None => {
19793                let null: u64 = 0;
19794                b.arg(&mut *q)
19795                    .arg(&mut *k)
19796                    .arg(&mut *qb)
19797                    .arg(&mut *kb)
19798                    .arg(pos)
19799                    .arg(&hd)
19800                    .arg(&nd)
19801                    .arg(&nhq)
19802                    .arg(&nhk)
19803                    .arg(&nt)
19804                    .arg(&theta_scale)
19805                    .arg(&freq_scale)
19806                    .arg(&null);
19807                unsafe {
19808                    b.launch(cfg)?;
19809                }
19810            }
19811        }
19812        Ok(())
19813    }
19814
19815    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
19816    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
19817    pub fn f32_to_bf16(
19818        &self,
19819        x: &CudaSlice<f32>,
19820        n: usize,
19821    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19822        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19823        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19824        let f = self.func("f32_to_bf16_flat");
19825        let n_i = n as i64;
19826        let cfg = LaunchConfig {
19827            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19828            block_dim: (256, 1, 1),
19829            shared_mem_bytes: 0,
19830        };
19831        let __s_b = self.gpu.stream();
19832        let mut b = __s_b.launch_builder(&f);
19833        b.arg(x).arg(&mut y).arg(&n_i);
19834        unsafe {
19835            b.launch(cfg)?;
19836        }
19837        Ok(y)
19838    }
19839
19840    pub fn f32_to_f16(
19841        &self,
19842        x: &CudaSlice<f32>,
19843        n: usize,
19844    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19845        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
19846        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19847        let f = self.func("f32_to_f16_flat");
19848        let n_i = n as i64;
19849        let cfg = LaunchConfig {
19850            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19851            block_dim: (256, 1, 1),
19852            shared_mem_bytes: 0,
19853        };
19854        let __s_b = self.gpu.stream();
19855        let mut b = __s_b.launch_builder(&f);
19856        b.arg(x).arg(&mut y).arg(&n_i);
19857        unsafe {
19858            b.launch(cfg)?;
19859        }
19860        Ok(y)
19861    }
19862
19863    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
19864    pub fn bf16_to_f16(
19865        &self,
19866        xb: &CudaSlice<u8>,
19867        n: usize,
19868    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19869        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19870        self.bf16_to_f16_into(xb, n, &mut y)?;
19871        Ok(y)
19872    }
19873
19874    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
19875    pub fn bf16_to_f16_into(
19876        &self,
19877        xb: &CudaSlice<u8>,
19878        n: usize,
19879        y: &mut CudaSlice<u8>,
19880    ) -> Result<(), Box<dyn std::error::Error>> {
19881        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
19882        assert!(y.len() >= n * 2);
19883        let f = self.func("bf16_to_f16_flat");
19884        let n2 = (n / 2) as i64;
19885        let cfg = LaunchConfig {
19886            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
19887            block_dim: (256, 1, 1),
19888            shared_mem_bytes: 0,
19889        };
19890        let __s_b = self.gpu.stream();
19891        let mut b = __s_b.launch_builder(&f);
19892        b.arg(xb).arg(y).arg(&n2);
19893        unsafe {
19894            b.launch(cfg)?;
19895        }
19896        Ok(())
19897    }
19898
19899    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
19900    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
19901    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
19902    /// head_dim in {256, 128}, bf16kv lane on.
19903    #[allow(clippy::too_many_arguments)]
19904    pub fn fa_prefill_vl8(
19905        &self,
19906        seqs: &[FaSeqVl],
19907        head_dim: usize,
19908        n_head: usize,
19909        n_head_kv: usize,
19910        scale: f32,
19911    ) -> Result<(), Box<dyn std::error::Error>> {
19912        const BK: usize = 32;
19913        let b = seqs.len();
19914        assert!(b >= 1 && b <= 8);
19915        let mut packed = [FaSeqVl::default(); 8];
19916        packed[..b].copy_from_slice(seqs);
19917        let v = FaVl8(packed);
19918        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19919        let ept = (n_head_kv * head_dim) as i32;
19920        {
19921            let f = self.func("fa_mirror_vl");
19922            let max_n = (max_t as i64) * ept as i64;
19923            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
19924            for which in 0..2i32 {
19925                let cfg = LaunchConfig {
19926                    grid_dim: (blocks, 1, b as u32),
19927                    block_dim: (256, 1, 1),
19928                    shared_mem_bytes: 0,
19929                };
19930                let __s_lb = self.gpu.stream();
19931                let mut lb = __s_lb.launch_builder(&f);
19932                lb.arg(&v).arg(&ept).arg(&which);
19933                unsafe {
19934                    lb.launch(cfg)?;
19935                }
19936            }
19937        }
19938        let hd_sfx = fa_hd_suffix(head_dim)?;
19939        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
19940        let block_q = 64usize;
19941        let kv_stages = 2usize;
19942        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19943            + 4 * (block_q * BK + 2 * block_q)) as u32;
19944        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19945        f.set_attribute(
19946            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19947            shmem as i32,
19948        )?;
19949        let cfg = LaunchConfig {
19950            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
19951            block_dim: (32, 4, 1),
19952            shared_mem_bytes: shmem,
19953        };
19954        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19955        let __s_lb = self.gpu.stream();
19956        let mut lb = __s_lb.launch_builder(&f);
19957        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
19958        unsafe {
19959            lb.launch(cfg)?;
19960        }
19961        Ok(())
19962    }
19963
19964    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
19965    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
19966    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
19967    #[allow(clippy::too_many_arguments)]
19968    pub fn attn_pre_vl8(
19969        &self,
19970        seqs: &[AttnPreVl],
19971        wq: &CudaSlice<f32>,
19972        wk: &CudaSlice<f32>,
19973        head_dim: usize,
19974        rope_dims: usize,
19975        n_head: usize,
19976        n_head_kv: usize,
19977        eps: f32,
19978        freq_base: f32,
19979        freq_scale: f32,
19980        kv_dim_k: usize,
19981        kv_dim_v: usize,
19982        k_tok_bytes: usize,
19983        v_tok_bytes: usize,
19984    ) -> Result<(), Box<dyn std::error::Error>> {
19985        let b = seqs.len();
19986        assert!(b >= 1 && b <= 8);
19987        let mut packed = [AttnPreVl::default(); 8];
19988        packed[..b].copy_from_slice(seqs);
19989        let v = AttnPreVl8(packed);
19990        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19991        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19992        {
19993            let f = self.func("q_gate_split_vl");
19994            let n = max_t * (n_head * head_dim) as u32;
19995            let cfg = LaunchConfig {
19996                grid_dim: (n.div_ceil(256), 1, b as u32),
19997                block_dim: (256, 1, 1),
19998                shared_mem_bytes: 0,
19999            };
20000            let __s_lb = self.gpu.stream();
20001            let mut lb = __s_lb.launch_builder(&f);
20002            lb.arg(&v).arg(&hd).arg(&nh);
20003            unsafe {
20004                lb.launch(cfg)?;
20005            }
20006        }
20007        {
20008            let f = self.func("attn_rms_vl");
20009            let cfg = LaunchConfig {
20010                grid_dim: (max_t * n_head as u32, 2, b as u32),
20011                block_dim: (rms_block(), 1, 1),
20012                shared_mem_bytes: 0,
20013            };
20014            let __s_lb = self.gpu.stream();
20015            let mut lb = __s_lb.launch_builder(&f);
20016            lb.arg(&v)
20017                .arg(wq)
20018                .arg(wk)
20019                .arg(&hd)
20020                .arg(&nh)
20021                .arg(&nhkv)
20022                .arg(&eps);
20023            unsafe {
20024                lb.launch(cfg)?;
20025            }
20026        }
20027        {
20028            let f = self.func("attn_rope_vl");
20029            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20030            let nd = rope_dims as i32;
20031            let cfg = LaunchConfig {
20032                grid_dim: (max_t * n_head as u32, 2, b as u32),
20033                block_dim: ((head_dim / 2) as u32, 1, 1),
20034                shared_mem_bytes: 0,
20035            };
20036            let __s_lb = self.gpu.stream();
20037            let mut lb = __s_lb.launch_builder(&f);
20038            lb.arg(&v)
20039                .arg(&hd)
20040                .arg(&nd)
20041                .arg(&nh)
20042                .arg(&nhkv)
20043                .arg(&theta_scale)
20044                .arg(&freq_scale);
20045            unsafe {
20046                lb.launch(cfg)?;
20047            }
20048        }
20049        {
20050            let f = self.func("append_kv_vl");
20051            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20052            let cfg = LaunchConfig {
20053                grid_dim: (nblk, max_t, b as u32),
20054                block_dim: (32, 1, 1),
20055                shared_mem_bytes: 0,
20056            };
20057            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20058            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20059            let __s_lb = self.gpu.stream();
20060            let mut lb = __s_lb.launch_builder(&f);
20061            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20062            unsafe {
20063                lb.launch(cfg)?;
20064            }
20065        }
20066        Ok(())
20067    }
20068
20069    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
20070    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
20071    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
20072    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
20073    pub fn fa_prefill_view(
20074        &self,
20075        q: &CudaSlice<f32>,
20076        k: &cudarc::driver::CudaView<u8>,
20077        v: &cudarc::driver::CudaView<u8>,
20078        o: &mut CudaSlice<f32>,
20079        head_dim: usize,
20080        n_head: usize,
20081        n_head_kv: usize,
20082        t: usize,
20083        t_kv: usize,
20084        scale: f32,
20085        causal: bool,
20086        k_tok_bytes: usize,
20087        v_tok_bytes: usize,
20088        g: bool,
20089    ) -> Result<(), Box<dyn std::error::Error>> {
20090        if portable_mma_gated() {
20091            return self.sdpa_naive_quantized_view(
20092                q,
20093                k,
20094                v,
20095                o,
20096                head_dim,
20097                n_head,
20098                n_head_kv,
20099                t,
20100                t_kv,
20101                scale,
20102                causal,
20103                k_tok_bytes,
20104                v_tok_bytes,
20105            );
20106        }
20107        const BLOCK_Q: usize = 64;
20108        const BK: usize = 32;
20109        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
20110        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
20111        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20112        let f = if g {
20113            self.func_g(&name)
20114        } else {
20115            self.func(&name)
20116        };
20117        let shmem =
20118            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20119        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20120        f.set_attribute(
20121            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20122            shmem as i32,
20123        )?;
20124        let cfg = LaunchConfig {
20125            grid_dim: (
20126                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20127                n_head as u32,
20128                1,
20129            ),
20130            block_dim: (32, 4, 1),
20131            shared_mem_bytes: shmem,
20132        };
20133        let (hd, nh, nhkv, ti, tkvi, cz) = (
20134            head_dim as i32,
20135            n_head as i32,
20136            n_head_kv as i32,
20137            t as i32,
20138            t_kv as i32,
20139            causal as i32,
20140        );
20141        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20142        let __s_b = self.gpu.stream();
20143        let mut b = __s_b.launch_builder(&f);
20144        b.arg(q)
20145            .arg(k)
20146            .arg(v)
20147            .arg(o)
20148            .arg(&hd)
20149            .arg(&nh)
20150            .arg(&nhkv)
20151            .arg(&ti)
20152            .arg(&tkvi)
20153            .arg(&scale)
20154            .arg(&cz)
20155            .arg(&ktb)
20156            .arg(&vtb);
20157        unsafe {
20158            b.launch(cfg)?;
20159        }
20160        Ok(())
20161    }
20162
20163    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
20164    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
20165    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
20166    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
20167    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
20168    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
20169    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
20170    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
20171    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
20172    #[allow(clippy::too_many_arguments)]
20173    pub fn fa_prefill_view_ws(
20174        &self,
20175        q: &CudaSlice<f32>,
20176        k: &cudarc::driver::CudaView<u8>,
20177        v: &cudarc::driver::CudaView<u8>,
20178        o: &mut CudaSlice<f32>,
20179        head_dim: usize,
20180        n_head: usize,
20181        n_head_kv: usize,
20182        t: usize,
20183        t_kv: usize,
20184        scale: f32,
20185        causal: bool,
20186        k_tok_bytes: usize,
20187        v_tok_bytes: usize,
20188        g: bool,
20189    ) -> Result<(), Box<dyn std::error::Error>> {
20190        if portable_mma_gated() {
20191            return self.sdpa_naive_quantized_view(
20192                q,
20193                k,
20194                v,
20195                o,
20196                head_dim,
20197                n_head,
20198                n_head_kv,
20199                t,
20200                t_kv,
20201                scale,
20202                causal,
20203                k_tok_bytes,
20204                v_tok_bytes,
20205            );
20206        }
20207        const BLOCK_Q: usize = 64;
20208        const BK: usize = 32;
20209        let kv_dim_k = n_head_kv * head_dim;
20210        let kv_dim_v = n_head_kv * head_dim;
20211        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20212        let v_ws_bytes = t_kv * kv_dim_v * 2;
20213        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
20214        let mut guard = self.prime_deqw_ws.lock().unwrap();
20215        let need_grow = match guard.as_ref() {
20216            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20217            None => true,
20218        };
20219        if need_grow {
20220            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20221            let (ck, cv) = guard
20222                .as_ref()
20223                .map(|(a, b)| (a.len(), b.len()))
20224                .unwrap_or((0, 0));
20225            *guard = Some((
20226                self.alloc_u8(grow(ck, k_ws_bytes))?,
20227                self.alloc_u8(grow(cv, v_ws_bytes))?,
20228            ));
20229        }
20230        let (kw, vw) = guard.as_mut().unwrap();
20231        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
20232        {
20233            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
20234            let f = if g {
20235                self.func_g("fa_dequant_kv_ws_bf16")
20236            } else {
20237                self.func("fa_dequant_kv_ws_bf16")
20238            };
20239            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20240            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20241            let cfg = LaunchConfig {
20242                grid_dim: (nblk.max(1), 1, 1),
20243                block_dim: (256, 1, 1),
20244                shared_mem_bytes: 0,
20245            };
20246            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20247            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20248            let __s_b = self.gpu.stream();
20249            let mut b = __s_b.launch_builder(&f);
20250            b.arg(k)
20251                .arg(v)
20252                .arg(&mut *kw)
20253                .arg(&mut *vw)
20254                .arg(&kdk)
20255                .arg(&kdv)
20256                .arg(&tkvi)
20257                .arg(&ktb)
20258                .arg(&vtb);
20259            unsafe {
20260                b.launch(cfg)?;
20261            }
20262        }
20263        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
20264        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
20265        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
20266        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
20267        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
20268        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
20269        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
20270        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20271            .map(|v| v != "0")
20272            .unwrap_or(true);
20273        {
20274            let hd_sfx = fa_hd_suffix(head_dim)?;
20275            let f = self.func(&format!(
20276                "fa_prefill_qw{}{hd_sfx}",
20277                if db { "_db" } else { "" }
20278            ));
20279            let shmem = if db {
20280                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
20281                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20282            } else {
20283                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20284            };
20285            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20286            f.set_attribute(
20287                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20288                shmem as i32,
20289            )?;
20290            let cfg = LaunchConfig {
20291                grid_dim: (
20292                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20293                    n_head as u32,
20294                    1,
20295                ),
20296                block_dim: (32, 4, 1),
20297                shared_mem_bytes: shmem,
20298            };
20299            let (hd, nh, nhkv, ti, tkvi, cz) = (
20300                head_dim as i32,
20301                n_head as i32,
20302                n_head_kv as i32,
20303                t as i32,
20304                t_kv as i32,
20305                causal as i32,
20306            );
20307            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20308            let __s_b = self.gpu.stream();
20309            let mut b = __s_b.launch_builder(&f);
20310            b.arg(q)
20311                .arg(&*kw)
20312                .arg(&*vw)
20313                .arg(o)
20314                .arg(&hd)
20315                .arg(&nh)
20316                .arg(&nhkv)
20317                .arg(&ti)
20318                .arg(&tkvi)
20319                .arg(&scale)
20320                .arg(&cz)
20321                .arg(&kdk)
20322                .arg(&kdv);
20323            unsafe {
20324                b.launch(cfg)?;
20325            }
20326        }
20327        Ok(())
20328    }
20329
20330    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
20331    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
20332    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
20333    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
20334    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
20335    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
20336    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
20337    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
20338    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
20339    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
20340    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
20341    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
20342    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
20343    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
20344    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
20345    #[allow(clippy::too_many_arguments)]
20346    pub fn fa_prefill_view_ws_w_hd128(
20347        &self,
20348        q: &CudaSlice<f32>,
20349        k: &cudarc::driver::CudaView<u8>,
20350        v: &cudarc::driver::CudaView<u8>,
20351        o: &mut CudaSlice<f32>,
20352        head_dim: usize,
20353        n_head: usize,
20354        n_head_kv: usize,
20355        t: usize,
20356        t_kv: usize,
20357        scale: f32,
20358        causal: bool,
20359        window: usize,
20360        k_tok_bytes: usize,
20361        v_tok_bytes: usize,
20362    ) -> Result<(), Box<dyn std::error::Error>> {
20363        assert_eq!(
20364            head_dim, 128,
20365            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20366        );
20367        if portable_mma_gated() {
20368            return self.sdpa_naive_w_quantized_view(
20369                q,
20370                k,
20371                v,
20372                o,
20373                head_dim,
20374                n_head,
20375                n_head_kv,
20376                t,
20377                t_kv,
20378                scale,
20379                causal,
20380                window,
20381                k_tok_bytes,
20382                v_tok_bytes,
20383            );
20384        }
20385        const BLOCK_Q: usize = 64;
20386        const BK: usize = 32;
20387        let kv_dim_k = n_head_kv * head_dim;
20388        let kv_dim_v = n_head_kv * head_dim;
20389        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20390        let v_ws_bytes = t_kv * kv_dim_v * 2;
20391        let mut guard = self.prime_deqw_ws.lock().unwrap();
20392        let need_grow = match guard.as_ref() {
20393            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20394            None => true,
20395        };
20396        if need_grow {
20397            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20398            let (ck, cv) = guard
20399                .as_ref()
20400                .map(|(a, b)| (a.len(), b.len()))
20401                .unwrap_or((0, 0));
20402            *guard = Some((
20403                self.alloc_u8(grow(ck, k_ws_bytes))?,
20404                self.alloc_u8(grow(cv, v_ws_bytes))?,
20405            ));
20406        }
20407        let (kw, vw) = guard.as_mut().unwrap();
20408        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
20409        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
20410        {
20411            let f = self.func("fa_dequant_kv_ws_bf16");
20412            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20413            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20414            let cfg = LaunchConfig {
20415                grid_dim: (nblk.max(1), 1, 1),
20416                block_dim: (256, 1, 1),
20417                shared_mem_bytes: 0,
20418            };
20419            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20420            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20421            let __s_b = self.gpu.stream();
20422            let mut b = __s_b.launch_builder(&f);
20423            b.arg(k)
20424                .arg(v)
20425                .arg(&mut *kw)
20426                .arg(&mut *vw)
20427                .arg(&kdk)
20428                .arg(&kdv)
20429                .arg(&tkvi)
20430                .arg(&ktb)
20431                .arg(&vtb);
20432            unsafe {
20433                b.launch(cfg)?;
20434            }
20435        }
20436        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
20437        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20438            .map(|v| v != "0")
20439            .unwrap_or(true);
20440        {
20441            let f = self.func(if db {
20442                "fa_prefill_qw_db_w_hd128"
20443            } else {
20444                "fa_prefill_qw_w_hd128"
20445            });
20446            let shmem = if db {
20447                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20448            } else {
20449                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20450            };
20451            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20452            f.set_attribute(
20453                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20454                shmem as i32,
20455            )?;
20456            let cfg = LaunchConfig {
20457                grid_dim: (
20458                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20459                    n_head as u32,
20460                    1,
20461                ),
20462                block_dim: (32, 4, 1),
20463                shared_mem_bytes: shmem,
20464            };
20465            let (hd, nh, nhkv, ti, tkvi, cz) = (
20466                head_dim as i32,
20467                n_head as i32,
20468                n_head_kv as i32,
20469                t as i32,
20470                t_kv as i32,
20471                causal as i32,
20472            );
20473            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20474            let __s_b = self.gpu.stream();
20475            let mut b = __s_b.launch_builder(&f);
20476            b.arg(q)
20477                .arg(&*kw)
20478                .arg(&*vw)
20479                .arg(o)
20480                .arg(&hd)
20481                .arg(&nh)
20482                .arg(&nhkv)
20483                .arg(&ti)
20484                .arg(&tkvi)
20485                .arg(&scale)
20486                .arg(&cz)
20487                .arg(&kdk)
20488                .arg(&kdv)
20489                .arg(&wnd);
20490            unsafe {
20491                b.launch(cfg)?;
20492            }
20493        }
20494        Ok(())
20495    }
20496
20497    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
20498    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
20499    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
20500    pub fn fa_decode(
20501        &self,
20502        q: &CudaSlice<f32>,
20503        k: &cudarc::driver::CudaView<u8>,
20504        v: &cudarc::driver::CudaView<u8>,
20505        o: &mut CudaSlice<f32>,
20506        head_dim: usize,
20507        n_head: usize,
20508        n_head_kv: usize,
20509        t_kv: usize,
20510        scale: f32,
20511        k_tok_bytes: usize,
20512        v_tok_bytes: usize,
20513    ) -> Result<(), Box<dyn std::error::Error>> {
20514        self.fa_decode_kvmod(
20515            q,
20516            k,
20517            v,
20518            o,
20519            head_dim,
20520            n_head,
20521            n_head_kv,
20522            t_kv,
20523            scale,
20524            k_tok_bytes,
20525            v_tok_bytes,
20526            false,
20527        )
20528    }
20529
20530    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
20531    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
20532    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
20533    #[allow(clippy::too_many_arguments)]
20534    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
20535    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
20536    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
20537    #[allow(clippy::too_many_arguments)]
20538    #[allow(clippy::too_many_arguments)]
20539    fn fa_decode_scalar_unified(
20540        &self,
20541        q: &cudarc::driver::CudaView<f32>,
20542        k: &cudarc::driver::CudaView<u8>,
20543        v: &cudarc::driver::CudaView<u8>,
20544        o: &mut cudarc::driver::CudaViewMut<f32>,
20545        head_dim: usize,
20546        n_head: usize,
20547        n_head_kv: usize,
20548        t_kv_host: usize,
20549        t_kv_dev: Option<&CudaSlice<i32>>,
20550        scale: f32,
20551        n_splits: usize,
20552        split_keys: usize,
20553        k_tok_bytes: usize,
20554        v_tok_bytes: usize,
20555        g: bool,
20556        part_o: &mut CudaSlice<f32>,
20557        part_m: &mut CudaSlice<f32>,
20558        part_l: &mut CudaSlice<f32>,
20559        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20560    ) -> Result<(), Box<dyn std::error::Error>> {
20561        let f = if g {
20562            self.func_g("fa_decode_f32")
20563        } else {
20564            self.fa_func("fa_decode_f32", head_dim)
20565        };
20566        let cfg = LaunchConfig {
20567            grid_dim: (n_head as u32, n_splits as u32, 1),
20568            block_dim: (head_dim as u32, 1, 1),
20569            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20570        };
20571        let (hd, nh, nhkv, nsp) = (
20572            head_dim as i32,
20573            n_head as i32,
20574            n_head_kv as i32,
20575            n_splits as i32,
20576        );
20577        let (ktb, vtb, tkvi, ski) = (
20578            k_tok_bytes as i64,
20579            v_tok_bytes as i64,
20580            t_kv_host as i32,
20581            split_keys as i32,
20582        );
20583        let __s_b = self.gpu.stream();
20584        let mut b = __s_b.launch_builder(&f);
20585        match t_kv_dev {
20586            Some(d) => {
20587                b.arg(q)
20588                    .arg(k)
20589                    .arg(v)
20590                    .arg(&mut *part_o)
20591                    .arg(&mut *part_m)
20592                    .arg(&mut *part_l)
20593                    .arg(&hd)
20594                    .arg(&nh)
20595                    .arg(&nhkv)
20596                    .arg(&tkvi)
20597                    .arg(d)
20598                    .arg(&scale)
20599                    .arg(&nsp)
20600                    .arg(&ski)
20601                    .arg(&ktb)
20602                    .arg(&vtb);
20603                unsafe {
20604                    b.launch(cfg)?;
20605                }
20606            }
20607            None => {
20608                let null: u64 = 0;
20609                b.arg(q)
20610                    .arg(k)
20611                    .arg(v)
20612                    .arg(&mut *part_o)
20613                    .arg(&mut *part_m)
20614                    .arg(&mut *part_l)
20615                    .arg(&hd)
20616                    .arg(&nh)
20617                    .arg(&nhkv)
20618                    .arg(&tkvi)
20619                    .arg(&null)
20620                    .arg(&scale)
20621                    .arg(&nsp)
20622                    .arg(&ski)
20623                    .arg(&ktb)
20624                    .arg(&vtb);
20625                unsafe {
20626                    b.launch(cfg)?;
20627                }
20628            }
20629        }
20630        let cfg2 = LaunchConfig {
20631            grid_dim: (n_head as u32, 1, 1),
20632            block_dim: (head_dim as u32, 1, 1),
20633            shared_mem_bytes: 0,
20634        };
20635        if let Some((oq, od)) = q8_out {
20636            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
20637            let fc = if g {
20638                self.func_g("fa_decode_combine_q8_1")
20639            } else {
20640                self.fa_func("fa_decode_combine_q8_1", head_dim)
20641            };
20642            let __s_b2 = self.gpu.stream();
20643            let mut b2 = __s_b2.launch_builder(&fc);
20644            b2.arg(&*part_o)
20645                .arg(&*part_m)
20646                .arg(&*part_l)
20647                .arg(oq)
20648                .arg(od)
20649                .arg(&hd)
20650                .arg(&nh)
20651                .arg(&nsp);
20652            unsafe {
20653                b2.launch(cfg2)?;
20654            }
20655            return Ok(());
20656        }
20657        let fc = if g {
20658            self.func_g("fa_decode_combine_f32")
20659        } else {
20660            self.fa_func("fa_decode_combine_f32", head_dim)
20661        };
20662        let __s_b2 = self.gpu.stream();
20663        let mut b2 = __s_b2.launch_builder(&fc);
20664        b2.arg(&*part_o)
20665            .arg(&*part_m)
20666            .arg(&*part_l)
20667            .arg(o)
20668            .arg(&hd)
20669            .arg(&nh)
20670            .arg(&nsp);
20671        unsafe {
20672            b2.launch(cfg2)?;
20673        }
20674        Ok(())
20675    }
20676
20677    pub fn fa_decode_kvmod(
20678        &self,
20679        q: &CudaSlice<f32>,
20680        k: &cudarc::driver::CudaView<u8>,
20681        v: &cudarc::driver::CudaView<u8>,
20682        o: &mut CudaSlice<f32>,
20683        head_dim: usize,
20684        n_head: usize,
20685        n_head_kv: usize,
20686        t_kv: usize,
20687        scale: f32,
20688        k_tok_bytes: usize,
20689        v_tok_bytes: usize,
20690        g: bool,
20691    ) -> Result<(), Box<dyn std::error::Error>> {
20692        let q_view = q.as_view();
20693        let mut o_view = o.as_view_mut();
20694        self.fa_decode_kvmod_view(
20695            &q_view,
20696            k,
20697            v,
20698            &mut o_view,
20699            head_dim,
20700            n_head,
20701            n_head_kv,
20702            t_kv,
20703            scale,
20704            k_tok_bytes,
20705            v_tok_bytes,
20706            g,
20707        )
20708    }
20709
20710    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
20711    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
20712    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
20713    /// per-session KV view and FA launch.
20714    #[allow(clippy::too_many_arguments)]
20715    pub fn fa_decode_kvmod_view(
20716        &self,
20717        q: &cudarc::driver::CudaView<f32>,
20718        k: &cudarc::driver::CudaView<u8>,
20719        v: &cudarc::driver::CudaView<u8>,
20720        o: &mut cudarc::driver::CudaViewMut<f32>,
20721        head_dim: usize,
20722        n_head: usize,
20723        n_head_kv: usize,
20724        t_kv: usize,
20725        scale: f32,
20726        k_tok_bytes: usize,
20727        v_tok_bytes: usize,
20728        g: bool,
20729    ) -> Result<(), Box<dyn std::error::Error>> {
20730        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
20731        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
20732        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
20733        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
20734        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
20735        //
20736        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
20737        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
20738        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
20739        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
20740        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
20741        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
20742        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
20743        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
20744        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
20745        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
20746        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
20747        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
20748        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
20749        // fall to the exact scalar there instead of the broken register arm.
20750        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20751        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
20752        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
20753        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
20754        if g && head_dim == 256 && !fa_v4_at(t_kv) {
20755            fa_vec = false;
20756        }
20757        let sp = fa_split_keys(t_kv, n_head_kv);
20758        let n_splits = if fa_vec {
20759            ((t_kv + sp - 1) / sp).max(1)
20760        } else {
20761            ((t_kv + 255) / 256).max(1)
20762        };
20763        let o_len = n_head * n_splits * head_dim;
20764        let ml_len = n_head * n_splits;
20765        let mut part_guard = self.fa_part_pool.lock().unwrap();
20766        if part_guard
20767            .as_ref()
20768            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20769            .unwrap_or(true)
20770        {
20771            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
20772            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
20773            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
20774            // later live allocations land at those addresses, and the next graph REPLAY writes
20775            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
20776            // output corruption began the burst after the trunk's t_kv growth first realloc'd
20777            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
20778            // the baked addresses alive (single-stream: eager writes the new buffers, replays
20779            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
20780            // (total retired < final size).
20781            let old = part_guard.take();
20782            let (co, cm) = old
20783                .as_ref()
20784                .map(|pp| (pp.0.len(), pp.1.len()))
20785                .unwrap_or((0, 0));
20786            if let Some(old) = old {
20787                self.fa_part_retired.lock().unwrap().push(old);
20788            }
20789            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20790                eprintln!(
20791                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20792                    co, o_len, cm, ml_len
20793                );
20794            }
20795            *part_guard = Some((
20796                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20797                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20798                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20799            ));
20800        }
20801        let pg = part_guard.as_mut().unwrap();
20802        self.gpu
20803            .stream()
20804            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20805        self.gpu
20806            .stream()
20807            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20808        self.gpu
20809            .stream()
20810            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20811        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20812        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20813        let (hd, nh, nhkv, tkvi, nsp) = (
20814            head_dim as i32,
20815            n_head as i32,
20816            n_head_kv as i32,
20817            t_kv as i32,
20818            n_splits as i32,
20819        );
20820        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20821        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
20822        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
20823        // silently truncating the accumulator.
20824        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20825        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
20826        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
20827        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
20828        // 178.4 -> 173.7 when 512 rode vec unconditionally).
20829        let fa512_min = fa512_min_tkv();
20830        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
20831        // g-module keeps the v4 pick (its class is not the depth-decay class).
20832        let deep = fa_vec
20833            && head_dim == 256
20834            && fa_v4_at(t_kv)
20835            && !g
20836            && fa_deep_at(t_kv)
20837            && !matches!(fa_v4_mode(), "noB3" | "stage");
20838        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
20839            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
20840            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
20841            let gqa = (n_head / n_head_kv).max(1) as u32;
20842            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
20843            (
20844                fv,
20845                LaunchConfig {
20846                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20847                    block_dim: (32, gqa, 1),
20848                    shared_mem_bytes: 0,
20849                },
20850            )
20851        } else if fa_vec && head_dim <= 256 {
20852            let gqa = (n_head / n_head_kv).max(1) as u32;
20853            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
20854            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
20855            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
20856            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
20857            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
20858            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
20859            // dequant each tile ONCE per block.
20860            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
20861            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
20862            // there by 12x — latency, not bandwidth, rules small KV).
20863            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20864            let smem_tkv = *SMEM_TKV.get_or_init(|| {
20865                std::env::var("MEMRA_FA_SMEM_TKV")
20866                    .ok()
20867                    .and_then(|v| v.parse().ok())
20868                    .unwrap_or_else(|| {
20869                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
20870                    })
20871            });
20872            if fa_v4_at(t_kv) && head_dim == 256 {
20873                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
20874                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
20875                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
20876                let v4name = match fa_v4_mode() {
20877                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
20878                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
20879                    _ if deep => "fa_decode_vec_q_v4_deep",
20880                    _ => "fa_decode_vec_q_v4",
20881                };
20882                let fv = if g {
20883                    self.func_g(v4name)
20884                } else {
20885                    self.func(v4name)
20886                };
20887                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
20888                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
20889                let shmem = (if deep { 12160 } else { 11520 }
20890                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20891                use cudarc::driver::sys::CUfunction_attribute_enum as A;
20892                fv.set_attribute(
20893                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20894                    shmem as i32,
20895                )?;
20896                (
20897                    fv,
20898                    LaunchConfig {
20899                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20900                        block_dim: (32, gqa, 1),
20901                        shared_mem_bytes: shmem,
20902                    },
20903                )
20904            } else if fa_v3_active(head_dim) {
20905                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
20906                // smem = sV only (half of v2's).
20907                let fv = if g {
20908                    self.func_g("fa_decode_vec_q_v3")
20909                } else {
20910                    self.func("fa_decode_vec_q_v3")
20911                };
20912                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
20913                (
20914                    fv,
20915                    LaunchConfig {
20916                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20917                        block_dim: (32, gqa, 1),
20918                        shared_mem_bytes: shmem,
20919                    },
20920                )
20921            } else if fa_v2_on() {
20922                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
20923                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
20924                // partials; same 32KB sK+sV tile as the smem twin.
20925                let fv = if g {
20926                    self.func_g("fa_decode_vec_q_v2")
20927                } else {
20928                    self.func("fa_decode_vec_q_v2")
20929                };
20930                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20931                (
20932                    fv,
20933                    LaunchConfig {
20934                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20935                        block_dim: (32, gqa, 1),
20936                        shared_mem_bytes: shmem,
20937                    },
20938                )
20939            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
20940            {
20941                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
20942                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
20943                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
20944                let fv = if g {
20945                    self.func_g("fa_decode_vec_q_smem")
20946                } else {
20947                    self.func("fa_decode_vec_q_smem")
20948                };
20949                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20950                use cudarc::driver::sys::CUfunction_attribute_enum as A;
20951                fv.set_attribute(
20952                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20953                    shmem as i32,
20954                )?;
20955                (
20956                    fv,
20957                    LaunchConfig {
20958                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20959                        block_dim: (32, gqa, 1),
20960                        shared_mem_bytes: shmem,
20961                    },
20962                )
20963            } else {
20964                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
20965                // dequant, zero dynamic shared memory.
20966                let fv = if g {
20967                    self.func_g("fa_decode_vec_q")
20968                } else {
20969                    self.func("fa_decode_vec_q")
20970                };
20971                (
20972                    fv,
20973                    LaunchConfig {
20974                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20975                        block_dim: (32, gqa, 1),
20976                        shared_mem_bytes: 0,
20977                    },
20978                )
20979            }
20980        } else {
20981            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
20982            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
20983            return self.fa_decode_scalar_unified(
20984                q,
20985                k,
20986                v,
20987                o,
20988                head_dim,
20989                n_head,
20990                n_head_kv,
20991                t_kv,
20992                None,
20993                scale,
20994                n_splits,
20995                if fa_vec { sp } else { 256 },
20996                k_tok_bytes,
20997                v_tok_bytes,
20998                g,
20999                part_o,
21000                part_m,
21001                part_l,
21002                None,
21003            );
21004        };
21005        let __s_b = self.gpu.stream();
21006        let mut b = __s_b.launch_builder(&f);
21007        b.arg(q)
21008            .arg(k)
21009            .arg(v)
21010            .arg(&mut *part_o)
21011            .arg(&mut *part_m)
21012            .arg(&mut *part_l)
21013            .arg(&hd)
21014            .arg(&nh)
21015            .arg(&nhkv)
21016            .arg(&tkvi)
21017            .arg(&scale)
21018            .arg(&nsp)
21019            .arg(&ktb)
21020            .arg(&vtb);
21021        unsafe {
21022            b.launch(cfg)?;
21023        }
21024        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
21025        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
21026        let (fc, cfg2) = (
21027            if g {
21028                self.func_g("fa_decode_combine_f32")
21029            } else {
21030                self.fa_func("fa_decode_combine_f32", head_dim)
21031            },
21032            LaunchConfig {
21033                grid_dim: (n_head as u32, 1, 1),
21034                block_dim: (head_dim as u32, 1, 1),
21035                shared_mem_bytes: 0,
21036            },
21037        );
21038        let __s_b2 = self.gpu.stream();
21039        let mut b2 = __s_b2.launch_builder(&fc);
21040        b2.arg(&*part_o)
21041            .arg(&*part_m)
21042            .arg(&*part_l)
21043            .arg(o)
21044            .arg(&hd)
21045            .arg(&nh)
21046            .arg(&nsp);
21047        unsafe {
21048            b2.launch(cfg2)?;
21049        }
21050        Ok(())
21051    }
21052
21053    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
21054    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
21055    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
21056    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
21057    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
21058    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
21059    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
21060    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
21061    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
21062    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
21063    #[allow(clippy::too_many_arguments)]
21064    pub fn fa_decode_batch_seqs_v4(
21065        &self,
21066        q: &CudaSlice<f32>,
21067        kv_ptrs: &cudarc::driver::CudaView<u64>,
21068        pos_seq: &CudaSlice<i32>,
21069        o: &mut CudaSlice<f32>,
21070        head_dim: usize,
21071        n_head: usize,
21072        n_head_kv: usize,
21073        b_n: usize,
21074        t_kv_max: usize,
21075        scale: f32,
21076        split_keys: usize,
21077        k_tok_bytes: usize,
21078        v_tok_bytes: usize,
21079    ) -> Result<(), Box<dyn std::error::Error>> {
21080        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21081        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21082        let o_len = b_n * n_head * n_splits_max * head_dim;
21083        let ml_len = b_n * n_head * n_splits_max;
21084        let mut part_guard = self.fa_part_pool.lock().unwrap();
21085        if part_guard
21086            .as_ref()
21087            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21088            .unwrap_or(true)
21089        {
21090            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21091            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21092            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21093            // later live allocations land at those addresses, and the next graph REPLAY writes
21094            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21095            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21096            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21097            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21098            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21099            // (total retired < final size).
21100            let old = part_guard.take();
21101            let (co, cm) = old
21102                .as_ref()
21103                .map(|pp| (pp.0.len(), pp.1.len()))
21104                .unwrap_or((0, 0));
21105            if let Some(old) = old {
21106                self.fa_part_retired.lock().unwrap().push(old);
21107            }
21108            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21109                eprintln!(
21110                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21111                    co, o_len, cm, ml_len
21112                );
21113            }
21114            *part_guard = Some((
21115                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21116                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21117                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21118            ));
21119        }
21120        let pg = part_guard.as_mut().unwrap();
21121        self.gpu
21122            .stream()
21123            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21124        self.gpu
21125            .stream()
21126            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21127        self.gpu
21128            .stream()
21129            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21130        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21131        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21132        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21133        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21134        let gqa = (n_head / n_head_kv).max(1) as u32;
21135        let f = self.func("fa_decode_vec_q_seqs_v4");
21136        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
21137        let shmem = (11520 + 32 * head_dim * 2) as u32;
21138        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21139        f.set_attribute(
21140            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21141            shmem as i32,
21142        )?;
21143        let cfg = LaunchConfig {
21144            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21145            block_dim: (32, gqa, 1),
21146            shared_mem_bytes: shmem,
21147        };
21148        {
21149            let __s_b = self.gpu.stream();
21150            let mut b = __s_b.launch_builder(&f);
21151            b.arg(q)
21152                .arg(kv_ptrs)
21153                .arg(pos_seq)
21154                .arg(&mut *part_o)
21155                .arg(&mut *part_m)
21156                .arg(&mut *part_l)
21157                .arg(&hd)
21158                .arg(&nh)
21159                .arg(&nhkv)
21160                .arg(&scale)
21161                .arg(&nspm)
21162                .arg(&spk)
21163                .arg(&ktb)
21164                .arg(&vtb);
21165            unsafe {
21166                b.launch(cfg)?;
21167            }
21168        }
21169        let fc = self.func("fa_decode_combine_seqs");
21170        let cfg2 = LaunchConfig {
21171            grid_dim: (n_head as u32, b_n as u32, 1),
21172            block_dim: (head_dim as u32, 1, 1),
21173            shared_mem_bytes: 0,
21174        };
21175        let __s_b2 = self.gpu.stream();
21176        let mut b2 = __s_b2.launch_builder(&fc);
21177        b2.arg(&*part_o)
21178            .arg(&*part_m)
21179            .arg(&*part_l)
21180            .arg(o)
21181            .arg(&hd)
21182            .arg(&nh)
21183            .arg(pos_seq)
21184            .arg(&nspm)
21185            .arg(&spk);
21186        unsafe {
21187            b2.launch(cfg2)?;
21188        }
21189        Ok(())
21190    }
21191
21192    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
21193    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
21194    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
21195    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
21196    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
21197    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
21198    #[allow(clippy::too_many_arguments)]
21199    pub fn append_kv_quantized_seqs(
21200        &self,
21201        k_rows: &CudaSlice<f32>,
21202        v_rows: &CudaSlice<f32>,
21203        kv_ptrs: &cudarc::driver::CudaView<u64>,
21204        pos_seq: &CudaSlice<i32>,
21205        b_n: usize,
21206        kv_dim_k: usize,
21207        kv_dim_v: usize,
21208        k_tok_bytes: usize,
21209        v_tok_bytes: usize,
21210    ) -> Result<(), Box<dyn std::error::Error>> {
21211        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21212        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21213        let cfg = LaunchConfig {
21214            grid_dim: (nblk, b_n as u32, 1),
21215            block_dim: (32, 1, 1),
21216            shared_mem_bytes: 0,
21217        };
21218        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21219        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21220        let __s_b = self.gpu.stream();
21221        let mut b = __s_b.launch_builder(&f);
21222        b.arg(k_rows)
21223            .arg(v_rows)
21224            .arg(kv_ptrs)
21225            .arg(pos_seq)
21226            .arg(&kdk)
21227            .arg(&kdv)
21228            .arg(&ktb)
21229            .arg(&vtb);
21230        unsafe {
21231            b.launch(cfg)?;
21232        }
21233        Ok(())
21234    }
21235
21236    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
21237    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
21238    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
21239    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
21240    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
21241    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21242        std::env::var("MEMRA_NO_FA_VEC").is_err()
21243            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21244            && base_len + 1 >= fa_vec_min_tkv()
21245            && head_dim <= 256
21246            && head_dim % 32 == 0
21247    }
21248
21249    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
21250    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
21251    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
21252    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
21253    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
21254    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
21255    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
21256    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
21257    #[allow(clippy::too_many_arguments)]
21258    pub fn fa_decode_rows(
21259        &self,
21260        q: &CudaSlice<f32>,
21261        k: &cudarc::driver::CudaView<u8>,
21262        v: &cudarc::driver::CudaView<u8>,
21263        o: &mut CudaSlice<f32>,
21264        head_dim: usize,
21265        n_head: usize,
21266        n_head_kv: usize,
21267        base_len: usize,
21268        t: usize,
21269        scale: f32,
21270        k_tok_bytes: usize,
21271        v_tok_bytes: usize,
21272        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
21273        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
21274        // keep the host arg. None is a bug for hd512 (asserted below).
21275        base_dev: Option<(&CudaSlice<i32>, i32)>,
21276        // K and V planes hold the same values (gemma globals, wv:=wk): pick
21277        // the _kv twin — V plane never read, value rides the q8_0 key dq.
21278        kv_shared: bool,
21279        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
21280        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
21281        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
21282        g: bool,
21283        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
21284        // (hd512 path) — the standalone quantize launch folds away.
21285        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21286    ) -> Result<(), Box<dyn std::error::Error>> {
21287        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21288        let t_kv_max = base_len + t; // LAST row's key bound
21289        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
21290        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
21291        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
21292        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
21293        // (parity law), so the partition is freely tunable — verify and decode move together.
21294        if head_dim == 512 {
21295            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21296            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
21297            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
21298            let v = *SP512.get_or_init(|| {
21299                std::env::var("MEMRA_FA_SP512")
21300                    .ok()
21301                    .and_then(|x| x.parse().ok())
21302                    .unwrap_or(0)
21303            });
21304            sp = if v >= 8 {
21305                v
21306            } else {
21307                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21308            };
21309        }
21310        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21311        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21312        let gqa = (n_head / n_head_kv).max(1) as u32;
21313        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
21314        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
21315        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
21316        // the different partition changes the combine's FP order (greedy tie flips at depth;
21317        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
21318        // consecutive rows by their OWN ladder value and launch once per group — each row then
21319        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
21320        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
21321        // sp override is t_kv-independent by construction).
21322        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
21323        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21324            groups.push((0, t, sp));
21325        } else {
21326            let mut r0 = 0usize;
21327            while r0 < t {
21328                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21329                let mut r1 = r0 + 1;
21330                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21331                    r1 += 1;
21332                }
21333                groups.push((r0, r1 - r0, sp_g));
21334                r0 = r1;
21335            }
21336        }
21337        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
21338        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
21339        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
21340        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21341        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21342            std::env::var("MEMRA_FA_SMEM_TKV")
21343                .ok()
21344                .and_then(|v| v.parse().ok())
21345                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21346        });
21347        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21348        let v3 = fa_v3_active(head_dim);
21349        let smem_rows =
21350            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21351        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
21352        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
21353        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
21354        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
21355        let _ = kv_shared;
21356        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
21357        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
21358        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21359        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
21360        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
21361        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
21362        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
21363        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
21364        // (kv_head, split) stages its tile once and loops the rows over it — kills the
21365        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
21366        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
21367        // shared by every hd512 caller through this wrapper (decode+verify flip together;
21368        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
21369        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
21370        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
21371        // not unpack-bound; jsonl 2026-07-14.
21372        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21373        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
21374        let tb512 = head_dim == 512
21375            && sp <= 32
21376            && n_head / n_head_kv.max(1) <= 16
21377            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21378        let fname = if tb512 {
21379            "fa_decode_vec_q_rows_v4_512_tb"
21380        } else if i2 {
21381            "fa_decode_vec_q_rows_dpl16_i2"
21382        } else if head_dim == 512 {
21383            "fa_decode_vec_q_rows_dpl16"
21384        }
21385        // gemma globals (parity law)
21386        else if v4 {
21387            "fa_decode_vec_q_rows_v4"
21388        } else if v3 {
21389            "fa_decode_vec_q_rows_v3"
21390        } else if fa_v2_on() {
21391            "fa_decode_vec_q_rows_v2"
21392        } else if smem_rows {
21393            "fa_decode_vec_q_rows_smem"
21394        } else {
21395            "fa_decode_vec_q_rows"
21396        };
21397        let f = if head_dim == 512 {
21398            self.fa_func(fname, head_dim)
21399        } else if g {
21400            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
21401            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
21402            // g-module rows against decode's g-module v4 — different programs, short-VG
21403            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
21404            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
21405            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
21406            // dq macros are format-aware.
21407            self.func_g(if smem_rows {
21408                "fa_decode_vec_q_rows"
21409            } else {
21410                fname
21411            })
21412        } else {
21413            self.func(fname)
21414        };
21415        let shmem = if tb512 {
21416            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
21417            let gk = Self::gkv_on();
21418            let sh =
21419                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21420            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21421            f.set_attribute(
21422                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21423                sh as i32,
21424            )?;
21425            sh
21426        } else if v4 || v3 || smem_rows || fa_v2_on() {
21427            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
21428            let sh = (if v4 {
21429                11520 + 32 * head_dim * if g { 1 } else { 2 }
21430            } else if v3 {
21431                32 * head_dim * 2
21432            } else {
21433                2 * 32 * head_dim * 2
21434            }) as u32;
21435            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21436            f.set_attribute(
21437                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21438                sh as i32,
21439            )?;
21440            sh
21441        } else {
21442            0
21443        };
21444        // Per-GROUP launches (single group in the common case — identical to the pre-fix
21445        // single launch there): each group gets its own partials (the rows kernel indexes
21446        // partials by its LOCAL grid.z row) and q/o row-offset views.
21447        for &(r0, t_g, sp_g) in &groups {
21448            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21449            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21450            let base_i = (base_len + r0) as i32;
21451            let o_len = t_g * n_head * n_splits_g * head_dim;
21452            let ml_len = t_g * n_head * n_splits_g;
21453            let mut part_guard = self.fa_part_pool.lock().unwrap();
21454            if part_guard
21455                .as_ref()
21456                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21457                .unwrap_or(true)
21458            {
21459                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21460                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21461                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21462                // later live allocations land at those addresses, and the next graph REPLAY writes
21463                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21464                // output corruption began the burst after the trunk's t_kv growth first realloc'd
21465                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21466                // the baked addresses alive (single-stream: eager writes the new buffers, replays
21467                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21468                // (total retired < final size).
21469                let old = part_guard.take();
21470                let (co, cm) = old
21471                    .as_ref()
21472                    .map(|pp| (pp.0.len(), pp.1.len()))
21473                    .unwrap_or((0, 0));
21474                if let Some(old) = old {
21475                    self.fa_part_retired.lock().unwrap().push(old);
21476                }
21477                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21478                    eprintln!(
21479                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21480                        co, o_len, cm, ml_len
21481                    );
21482                }
21483                *part_guard = Some((
21484                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21485                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21486                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21487                ));
21488            }
21489            let pg = part_guard.as_mut().unwrap();
21490            self.gpu
21491                .stream()
21492                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21493            self.gpu
21494                .stream()
21495                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21496            self.gpu
21497                .stream()
21498                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21499            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21500            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21501            let qv = self.view(q, t * n_head * head_dim);
21502            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21503            let cfg = LaunchConfig {
21504                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21505                block_dim: (32, gqa, 1),
21506                shared_mem_bytes: shmem,
21507            };
21508            {
21509                let __s_b = self.gpu.stream();
21510                let mut b = __s_b.launch_builder(&f);
21511                if tb512 {
21512                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
21513                    let (bd, plus) =
21514                        base_dev.expect("hd512 rows twin requires a device base counter");
21515                    let plus_g = plus + r0 as i32;
21516                    let nr = t_g as i32;
21517                    if Self::pdl_on() && Self::pdl_wb_on() {
21518                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
21519                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21520                        let s = &self.gpu.stream();
21521                        let (pq, _b0) = q_g.device_ptr(s);
21522                        let (pk, _b1) = k.device_ptr(s);
21523                        let (pv, _b2) = v.device_ptr(s);
21524                        let (po, _b3) = part_o.device_ptr_mut(s);
21525                        let (pm, _b4) = part_m.device_ptr_mut(s);
21526                        let (pl, _b5) = part_l.device_ptr_mut(s);
21527                        let (pb, _b6) = bd.device_ptr(s);
21528                        let mut ps = [
21529                            &pq as *const _ as *mut std::ffi::c_void,
21530                            &pk as *const _ as *mut _,
21531                            &pv as *const _ as *mut _,
21532                            &po as *const _ as *mut _,
21533                            &pm as *const _ as *mut _,
21534                            &pl as *const _ as *mut _,
21535                            &hd as *const _ as *mut _,
21536                            &nh as *const _ as *mut _,
21537                            &nhkv as *const _ as *mut _,
21538                            &pb as *const _ as *mut _,
21539                            &plus_g as *const _ as *mut _,
21540                            &scale as *const _ as *mut _,
21541                            &nspm as *const _ as *mut _,
21542                            &spk as *const _ as *mut _,
21543                            &ktb as *const _ as *mut _,
21544                            &vtb as *const _ as *mut _,
21545                            &nr as *const _ as *mut _,
21546                        ];
21547                        unsafe {
21548                            self.launch_pdl_flash(
21549                                Self::gkv_on(),
21550                                "fa_decode_vec_q_rows_v4_512_tb",
21551                                (n_head_kv as u32, n_splits_g as u32, 1),
21552                                (32, gqa, 1),
21553                                shmem,
21554                                &mut ps,
21555                            )?;
21556                        }
21557                    } else {
21558                        let cfg_tb = LaunchConfig {
21559                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21560                            block_dim: (32, gqa, 1),
21561                            shared_mem_bytes: shmem,
21562                        };
21563                        b.arg(&q_g)
21564                            .arg(k)
21565                            .arg(v)
21566                            .arg(&mut *part_o)
21567                            .arg(&mut *part_m)
21568                            .arg(&mut *part_l)
21569                            .arg(&hd)
21570                            .arg(&nh)
21571                            .arg(&nhkv)
21572                            .arg(bd)
21573                            .arg(&plus_g)
21574                            .arg(&scale)
21575                            .arg(&nspm)
21576                            .arg(&spk)
21577                            .arg(&ktb)
21578                            .arg(&vtb)
21579                            .arg(&nr);
21580                        unsafe {
21581                            b.launch(cfg_tb)?;
21582                        }
21583                    }
21584                } else if head_dim == 512 {
21585                    let (bd, plus) =
21586                        base_dev.expect("hd512 rows twin requires a device base counter");
21587                    let plus_g = plus + r0 as i32;
21588                    b.arg(&q_g)
21589                        .arg(k)
21590                        .arg(v)
21591                        .arg(&mut *part_o)
21592                        .arg(&mut *part_m)
21593                        .arg(&mut *part_l)
21594                        .arg(&hd)
21595                        .arg(&nh)
21596                        .arg(&nhkv)
21597                        .arg(bd)
21598                        .arg(&plus_g)
21599                        .arg(&scale)
21600                        .arg(&nspm)
21601                        .arg(&spk)
21602                        .arg(&ktb)
21603                        .arg(&vtb);
21604                    unsafe {
21605                        b.launch(cfg)?;
21606                    }
21607                } else {
21608                    b.arg(&q_g)
21609                        .arg(k)
21610                        .arg(v)
21611                        .arg(&mut *part_o)
21612                        .arg(&mut *part_m)
21613                        .arg(&mut *part_l)
21614                        .arg(&hd)
21615                        .arg(&nh)
21616                        .arg(&nhkv)
21617                        .arg(&base_i)
21618                        .arg(&scale)
21619                        .arg(&nspm)
21620                        .arg(&spk)
21621                        .arg(&ktb)
21622                        .arg(&vtb);
21623                    unsafe {
21624                        b.launch(cfg)?;
21625                    }
21626                }
21627            }
21628            let cfg2 = LaunchConfig {
21629                grid_dim: (n_head as u32, t_g as u32, 1),
21630                block_dim: (head_dim as u32, 1, 1),
21631                shared_mem_bytes: 0,
21632            };
21633            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21634            if head_dim == 512 {
21635                // device-len combine (shared by verify/eager/graph — parity by symbol): the
21636                // per-row n_splits derives from the SAME counter the rows kernel read.
21637                let (bd, plus) = base_dev.unwrap();
21638                let plus_g = plus + r0 as i32;
21639                if let Some((oq, od)) = q8_out.as_mut() {
21640                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
21641                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21642                    if Self::pdl_on() && Self::pdl_wb_on() {
21643                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
21644                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21645                        let s = &self.gpu.stream();
21646                        let (po, _g0) = part_o.device_ptr(s);
21647                        let (pm, _g1) = part_m.device_ptr(s);
21648                        let (pl, _g2) = part_l.device_ptr(s);
21649                        let (pq, _g3) = oq.device_ptr_mut(s);
21650                        let (pd, _g4) = od.device_ptr_mut(s);
21651                        let (pb, _g5) = bd.device_ptr(s);
21652                        let mut ps = [
21653                            &po as *const _ as *mut std::ffi::c_void,
21654                            &pm as *const _ as *mut _,
21655                            &pl as *const _ as *mut _,
21656                            &pq as *const _ as *mut _,
21657                            &pd as *const _ as *mut _,
21658                            &hd as *const _ as *mut _,
21659                            &nh as *const _ as *mut _,
21660                            &pb as *const _ as *mut _,
21661                            &plus_g as *const _ as *mut _,
21662                            &nspm as *const _ as *mut _,
21663                            &spk as *const _ as *mut _,
21664                        ];
21665                        unsafe {
21666                            self.launch_pdl_flash(
21667                                Self::gkv_on(),
21668                                "fa_decode_combine_rows_dc_q8_1",
21669                                cfg2.grid_dim,
21670                                cfg2.block_dim,
21671                                0,
21672                                &mut ps,
21673                            )?;
21674                        }
21675                        continue;
21676                    }
21677                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21678                    let __s_b2 = self.gpu.stream();
21679                    let mut b2 = __s_b2.launch_builder(&fc);
21680                    b2.arg(&*part_o)
21681                        .arg(&*part_m)
21682                        .arg(&*part_l)
21683                        .arg(&mut **oq)
21684                        .arg(&mut **od)
21685                        .arg(&hd)
21686                        .arg(&nh)
21687                        .arg(bd)
21688                        .arg(&plus_g)
21689                        .arg(&nspm)
21690                        .arg(&spk);
21691                    unsafe {
21692                        b2.launch(cfg2)?;
21693                    }
21694                    continue;
21695                }
21696                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21697                let __s_b2 = self.gpu.stream();
21698                let mut b2 = __s_b2.launch_builder(&fc);
21699                b2.arg(&*part_o)
21700                    .arg(&*part_m)
21701                    .arg(&*part_l)
21702                    .arg(&mut o_g)
21703                    .arg(&hd)
21704                    .arg(&nh)
21705                    .arg(bd)
21706                    .arg(&plus_g)
21707                    .arg(&nspm)
21708                    .arg(&spk);
21709                unsafe {
21710                    b2.launch(cfg2)?;
21711                }
21712            } else {
21713                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
21714                // leave the caller's pair unwritten (consumer would read garbage).
21715                assert!(
21716                    q8_out.is_none(),
21717                    "rows q8 emit requires the hd512 dc combine"
21718                );
21719                let fc = self.func("fa_decode_combine_rows");
21720                let __s_b2 = self.gpu.stream();
21721                let mut b2 = __s_b2.launch_builder(&fc);
21722                b2.arg(&*part_o)
21723                    .arg(&*part_m)
21724                    .arg(&*part_l)
21725                    .arg(&mut o_g)
21726                    .arg(&hd)
21727                    .arg(&nh)
21728                    .arg(&base_i)
21729                    .arg(&nspm)
21730                    .arg(&spk);
21731                unsafe {
21732                    b2.launch(cfg2)?;
21733                }
21734            }
21735        }
21736        Ok(())
21737    }
21738
21739    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
21740    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
21741    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
21742    #[allow(clippy::too_many_arguments)]
21743    pub fn fa_decode_rows_w(
21744        &self,
21745        q: &CudaSlice<f32>,
21746        k: &cudarc::driver::CudaView<u8>,
21747        v: &cudarc::driver::CudaView<u8>,
21748        o: &mut CudaSlice<f32>,
21749        head_dim: usize,
21750        n_head: usize,
21751        n_head_kv: usize,
21752        base_dev: &CudaSlice<i32>,
21753        base_plus: i32,
21754        t: usize,
21755        scale: f32,
21756        window: usize,
21757        k_tok_bytes: usize,
21758        v_tok_bytes: usize,
21759        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21760    ) -> Result<(), Box<dyn std::error::Error>> {
21761        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
21762        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
21763        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
21764        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
21765        debug_assert!(head_dim == 256);
21766        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
21767        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
21768        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
21769        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
21770        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
21771        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
21772        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
21773        let sp = {
21774            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21775            let v = *SPW.get_or_init(|| {
21776                std::env::var("MEMRA_FA_SPW")
21777                    .ok()
21778                    .and_then(|x| x.parse().ok())
21779                    .unwrap_or(0)
21780            });
21781            if v >= 8 {
21782                v
21783            } else {
21784                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21785            }
21786        };
21787        let n_splits_max = (window + sp - 1) / sp;
21788        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21789        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21790        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21791        let gqa = (n_head / n_head_kv).max(1) as u32;
21792        let o_len = t * n_head * n_splits_max * head_dim;
21793        let ml_len = t * n_head * n_splits_max;
21794        let mut part_guard = self.fa_part_pool.lock().unwrap();
21795        if part_guard
21796            .as_ref()
21797            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21798            .unwrap_or(true)
21799        {
21800            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21801            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21802            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21803            // later live allocations land at those addresses, and the next graph REPLAY writes
21804            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21805            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21806            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21807            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21808            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21809            // (total retired < final size).
21810            let old = part_guard.take();
21811            let (co, cm) = old
21812                .as_ref()
21813                .map(|pp| (pp.0.len(), pp.1.len()))
21814                .unwrap_or((0, 0));
21815            if let Some(old) = old {
21816                self.fa_part_retired.lock().unwrap().push(old);
21817            }
21818            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21819                eprintln!(
21820                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21821                    co, o_len, cm, ml_len
21822                );
21823            }
21824            *part_guard = Some((
21825                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21826                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21827                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21828            ));
21829        }
21830        let pg = part_guard.as_mut().unwrap();
21831        self.gpu
21832            .stream()
21833            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21834        self.gpu
21835            .stream()
21836            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21837        self.gpu
21838            .stream()
21839            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21840        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21841        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
21842        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
21843        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
21844        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
21845        // floor (deep-ctx broadcast win); register twin between.
21846        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21847        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
21848            std::env::var("MEMRA_FA_SMEM_TKV")
21849                .ok()
21850                .and_then(|v| v.parse().ok())
21851                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21852        });
21853        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
21854        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
21855        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
21856        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
21857        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
21858        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21859        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
21860        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
21861        // per (lane, format-module) keeps parity structural; the old register-i2 detour
21862        // (-33%) is retired.
21863        let wg = Self::wkv_on();
21864        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
21865        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
21866        let sp2 =
21867            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
21868        if sp2 {
21869            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21870            if Self::pdl_on() && Self::pdl_wb_on() {
21871                // wave-B2b: flavor mirrors wg.
21872                use cudarc::driver::{DevicePtr, DevicePtrMut};
21873                let s = &self.gpu.stream();
21874                let (pq, _b0) = q.device_ptr(s);
21875                let (pk, _b1) = k.device_ptr(s);
21876                let (pv, _b2) = v.device_ptr(s);
21877                let (po, _b3) = part_o.device_ptr_mut(s);
21878                let (pm, _b4) = part_m.device_ptr_mut(s);
21879                let (pl, _b5) = part_l.device_ptr_mut(s);
21880                let (pb, _b6) = base_dev.device_ptr(s);
21881                let mut ps = [
21882                    &pq as *const _ as *mut std::ffi::c_void,
21883                    &pk as *const _ as *mut _,
21884                    &pv as *const _ as *mut _,
21885                    &po as *const _ as *mut _,
21886                    &pm as *const _ as *mut _,
21887                    &pl as *const _ as *mut _,
21888                    &hd as *const _ as *mut _,
21889                    &nh as *const _ as *mut _,
21890                    &nhkv as *const _ as *mut _,
21891                    &pb as *const _ as *mut _,
21892                    &base_plus as *const _ as *mut _,
21893                    &scale as *const _ as *mut _,
21894                    &nspm as *const _ as *mut _,
21895                    &spk as *const _ as *mut _,
21896                    &ktb as *const _ as *mut _,
21897                    &vtb as *const _ as *mut _,
21898                    &wini as *const _ as *mut _,
21899                ];
21900                unsafe {
21901                    self.launch_pdl_flash(
21902                        wg,
21903                        "fa_decode_vec_q_rows_v4_w_sp",
21904                        (n_head_kv as u32, n_splits_max as u32, t as u32),
21905                        (32, gqa + 1, 1),
21906                        sh,
21907                        &mut ps,
21908                    )?;
21909                }
21910            } else {
21911                let f = if wg {
21912                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
21913                } else {
21914                    self.func("fa_decode_vec_q_rows_v4_w_sp")
21915                };
21916                f.set_attribute(
21917                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21918                    sh as i32,
21919                )?;
21920                let cfg = LaunchConfig {
21921                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21922                    block_dim: (32, gqa + 1, 1),
21923                    shared_mem_bytes: sh,
21924                };
21925                let __s_b = self.gpu.stream();
21926                let mut b = __s_b.launch_builder(&f);
21927                b.arg(q)
21928                    .arg(k)
21929                    .arg(v)
21930                    .arg(&mut *part_o)
21931                    .arg(&mut *part_m)
21932                    .arg(&mut *part_l)
21933                    .arg(&hd)
21934                    .arg(&nh)
21935                    .arg(&nhkv)
21936                    .arg(base_dev)
21937                    .arg(&base_plus)
21938                    .arg(&scale)
21939                    .arg(&nspm)
21940                    .arg(&spk)
21941                    .arg(&ktb)
21942                    .arg(&vtb)
21943                    .arg(&wini);
21944                unsafe {
21945                    b.launch(cfg)?;
21946                }
21947            }
21948        } else {
21949            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
21950                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
21951                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21952                use cudarc::driver::{DevicePtr, DevicePtrMut};
21953                let s = &self.gpu.stream();
21954                let (pq, _b0) = q.device_ptr(s);
21955                let (pk, _b1) = k.device_ptr(s);
21956                let (pv, _b2) = v.device_ptr(s);
21957                let (po, _b3) = part_o.device_ptr_mut(s);
21958                let (pm, _b4) = part_m.device_ptr_mut(s);
21959                let (pl, _b5) = part_l.device_ptr_mut(s);
21960                let (pb, _b6) = base_dev.device_ptr(s);
21961                let mut ps = [
21962                    &pq as *const _ as *mut std::ffi::c_void,
21963                    &pk as *const _ as *mut _,
21964                    &pv as *const _ as *mut _,
21965                    &po as *const _ as *mut _,
21966                    &pm as *const _ as *mut _,
21967                    &pl as *const _ as *mut _,
21968                    &hd as *const _ as *mut _,
21969                    &nh as *const _ as *mut _,
21970                    &nhkv as *const _ as *mut _,
21971                    &pb as *const _ as *mut _,
21972                    &base_plus as *const _ as *mut _,
21973                    &scale as *const _ as *mut _,
21974                    &nspm as *const _ as *mut _,
21975                    &spk as *const _ as *mut _,
21976                    &ktb as *const _ as *mut _,
21977                    &vtb as *const _ as *mut _,
21978                    &wini as *const _ as *mut _,
21979                ];
21980                unsafe {
21981                    self.launch_pdl_flash(
21982                        wg,
21983                        "fa_decode_vec_q_rows_v4_w",
21984                        (n_head_kv as u32, n_splits_max as u32, t as u32),
21985                        (32, gqa, 1),
21986                        sh,
21987                        &mut ps,
21988                    )?;
21989                }
21990            } else {
21991                let pick = |name: &str| {
21992                    if wg {
21993                        self.func_g(name)
21994                    } else {
21995                        self.func(name)
21996                    }
21997                };
21998                let (f, sh) = if fa_v4_at(window) {
21999                    let f = pick("fa_decode_vec_q_rows_v4_w");
22000                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22001                } else if smem_tkv > 0 && window >= smem_tkv {
22002                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
22003                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
22004                    (
22005                        pick("fa_decode_vec_q_rows_smem_w"),
22006                        (2 * 32 * head_dim * 2) as u32,
22007                    )
22008                } else {
22009                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22010                };
22011                f.set_attribute(
22012                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22013                    sh as i32,
22014                )?;
22015                let cfg = LaunchConfig {
22016                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22017                    block_dim: (32, gqa, 1),
22018                    shared_mem_bytes: sh,
22019                };
22020                let __s_b = self.gpu.stream();
22021                let mut b = __s_b.launch_builder(&f);
22022                b.arg(q)
22023                    .arg(k)
22024                    .arg(v)
22025                    .arg(&mut *part_o)
22026                    .arg(&mut *part_m)
22027                    .arg(&mut *part_l)
22028                    .arg(&hd)
22029                    .arg(&nh)
22030                    .arg(&nhkv)
22031                    .arg(base_dev)
22032                    .arg(&base_plus)
22033                    .arg(&scale)
22034                    .arg(&nspm)
22035                    .arg(&spk)
22036                    .arg(&ktb)
22037                    .arg(&vtb)
22038                    .arg(&wini);
22039                unsafe {
22040                    b.launch(cfg)?;
22041                }
22042            }
22043        }
22044        let cfg2 = LaunchConfig {
22045            grid_dim: (n_head as u32, t as u32, 1),
22046            block_dim: (head_dim as u32, 1, 1),
22047            shared_mem_bytes: 0,
22048        };
22049        if let Some((oq, od)) = q8_out {
22050            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
22051            // consumes the pair directly; the standalone quantize launch folds away.
22052            if Self::pdl_on() && Self::pdl_wb_on() {
22053                // wave-B2: flavor mirrors the builder's wg choice.
22054                use cudarc::driver::{DevicePtr, DevicePtrMut};
22055                let s = &self.gpu.stream();
22056                let (po, _g0) = part_o.device_ptr(s);
22057                let (pm, _g1) = part_m.device_ptr(s);
22058                let (pl, _g2) = part_l.device_ptr(s);
22059                let (pq, _g3) = oq.device_ptr_mut(s);
22060                let (pd, _g4) = od.device_ptr_mut(s);
22061                let mut ps = [
22062                    &po as *const _ as *mut std::ffi::c_void,
22063                    &pm as *const _ as *mut _,
22064                    &pl as *const _ as *mut _,
22065                    &pq as *const _ as *mut _,
22066                    &pd as *const _ as *mut _,
22067                    &hd as *const _ as *mut _,
22068                    &nh as *const _ as *mut _,
22069                    &nspm as *const _ as *mut _,
22070                    &spk as *const _ as *mut _,
22071                    &wini as *const _ as *mut _,
22072                ];
22073                unsafe {
22074                    self.launch_pdl_flash(
22075                        wg,
22076                        "fa_decode_combine_rows_w_q8_1",
22077                        cfg2.grid_dim,
22078                        cfg2.block_dim,
22079                        0,
22080                        &mut ps,
22081                    )?;
22082                }
22083                return Ok(());
22084            }
22085            let fc = if wg {
22086                self.func_g("fa_decode_combine_rows_w_q8_1")
22087            } else {
22088                self.func("fa_decode_combine_rows_w_q8_1")
22089            };
22090            let __s_b2 = self.gpu.stream();
22091            let mut b2 = __s_b2.launch_builder(&fc);
22092            b2.arg(&*part_o)
22093                .arg(&*part_m)
22094                .arg(&*part_l)
22095                .arg(oq)
22096                .arg(od)
22097                .arg(&hd)
22098                .arg(&nh)
22099                .arg(&nspm)
22100                .arg(&spk)
22101                .arg(&wini);
22102            unsafe {
22103                b2.launch(cfg2)?;
22104            }
22105            return Ok(());
22106        }
22107        let fc = if wg {
22108            self.func_g("fa_decode_combine_rows_w")
22109        } else {
22110            self.func("fa_decode_combine_rows_w")
22111        };
22112        let __s_b2 = self.gpu.stream();
22113        let mut b2 = __s_b2.launch_builder(&fc);
22114        b2.arg(&*part_o)
22115            .arg(&*part_m)
22116            .arg(&*part_l)
22117            .arg(o)
22118            .arg(&hd)
22119            .arg(&nh)
22120            .arg(&nspm)
22121            .arg(&spk)
22122            .arg(&wini);
22123        unsafe {
22124            b2.launch(cfg2)?;
22125        }
22126        Ok(())
22127    }
22128
22129    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
22130    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
22131    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
22132    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
22133    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
22134    #[allow(clippy::too_many_arguments)]
22135    pub fn fa_decode_rows_dc(
22136        &self,
22137        q: &CudaSlice<f32>,
22138        k: &cudarc::driver::CudaView<u8>,
22139        v: &cudarc::driver::CudaView<u8>,
22140        o: &mut CudaSlice<f32>,
22141        head_dim: usize,
22142        n_head: usize,
22143        n_head_kv: usize,
22144        base_dev: &CudaSlice<i32>,
22145        t_kv_upper: usize,
22146        t: usize,
22147        scale: f32,
22148        k_tok_bytes: usize,
22149        v_tok_bytes: usize,
22150        base_plus: i32,
22151        g: bool,
22152    ) -> Result<(), Box<dyn std::error::Error>> {
22153        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22154        assert!(
22155            v4 || fa_v3_active(head_dim),
22156            "stream fa rows requires the v3 or v4 lane"
22157        );
22158        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22159        if v4 {
22160            let sp = fa_split_keys(t_kv_upper, n_head_kv);
22161            let n_splits_max = (t_kv_upper + sp - 1) / sp;
22162            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22163            let (nspm, spk) = (n_splits_max as i32, sp as i32);
22164            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22165            let gqa = (n_head / n_head_kv).max(1) as u32;
22166            let o_len = t * n_head * n_splits_max * head_dim;
22167            let ml_len = t * n_head * n_splits_max;
22168            let mut part_guard = self.fa_part_pool.lock().unwrap();
22169            if part_guard
22170                .as_ref()
22171                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22172                .unwrap_or(true)
22173            {
22174                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22175                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22176                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22177                // later live allocations land at those addresses, and the next graph REPLAY writes
22178                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22179                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22180                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22181                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22182                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22183                // (total retired < final size).
22184                let old = part_guard.take();
22185                let (co, cm) = old
22186                    .as_ref()
22187                    .map(|pp| (pp.0.len(), pp.1.len()))
22188                    .unwrap_or((0, 0));
22189                if let Some(old) = old {
22190                    self.fa_part_retired.lock().unwrap().push(old);
22191                }
22192                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22193                    eprintln!(
22194                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22195                        co, o_len, cm, ml_len
22196                    );
22197                }
22198                *part_guard = Some((
22199                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22200                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22201                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22202                ));
22203            }
22204            let pg = part_guard.as_mut().unwrap();
22205            self.gpu
22206                .stream()
22207                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22208            self.gpu
22209                .stream()
22210                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22211            self.gpu
22212                .stream()
22213                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22214            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22215            let f = if g {
22216                self.func_g("fa_decode_vec_q_rows_v4_dc")
22217            } else {
22218                self.func("fa_decode_vec_q_rows_v4_dc")
22219            };
22220            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22221            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22222            f.set_attribute(
22223                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22224                sh as i32,
22225            )?;
22226            let cfg = LaunchConfig {
22227                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22228                block_dim: (32, gqa, 1),
22229                shared_mem_bytes: sh,
22230            };
22231            let __s_b = self.gpu.stream();
22232            let mut b = __s_b.launch_builder(&f);
22233            b.arg(q)
22234                .arg(k)
22235                .arg(v)
22236                .arg(&mut *part_o)
22237                .arg(&mut *part_m)
22238                .arg(&mut *part_l)
22239                .arg(&hd)
22240                .arg(&nh)
22241                .arg(&nhkv)
22242                .arg(base_dev)
22243                .arg(&base_plus)
22244                .arg(&scale)
22245                .arg(&nspm)
22246                .arg(&spk)
22247                .arg(&ktb)
22248                .arg(&vtb);
22249            unsafe {
22250                b.launch(cfg)?;
22251            }
22252            let fc = self.func("fa_decode_combine_rows_dc");
22253            let cfg2 = LaunchConfig {
22254                grid_dim: (n_head as u32, t as u32, 1),
22255                block_dim: (head_dim as u32, 1, 1),
22256                shared_mem_bytes: 0,
22257            };
22258            let __s_b2 = self.gpu.stream();
22259            let mut b2 = __s_b2.launch_builder(&fc);
22260            b2.arg(&*part_o)
22261                .arg(&*part_m)
22262                .arg(&*part_l)
22263                .arg(o)
22264                .arg(&hd)
22265                .arg(&nh)
22266                .arg(base_dev)
22267                .arg(&base_plus)
22268                .arg(&nspm)
22269                .arg(&spk);
22270            unsafe {
22271                b2.launch(cfg2)?;
22272            }
22273            return Ok(());
22274        }
22275        let sp = fa_split_keys(t_kv_upper, n_head_kv);
22276        let n_splits_max = (t_kv_upper + sp - 1) / sp;
22277        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22278        let (nspm, spk) = (n_splits_max as i32, sp as i32);
22279        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22280        let gqa = (n_head / n_head_kv).max(1) as u32;
22281        let o_len = t * n_head * n_splits_max * head_dim;
22282        let ml_len = t * n_head * n_splits_max;
22283        let mut part_guard = self.fa_part_pool.lock().unwrap();
22284        if part_guard
22285            .as_ref()
22286            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22287            .unwrap_or(true)
22288        {
22289            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22290            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22291            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22292            // later live allocations land at those addresses, and the next graph REPLAY writes
22293            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22294            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22295            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22296            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22297            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22298            // (total retired < final size).
22299            let old = part_guard.take();
22300            let (co, cm) = old
22301                .as_ref()
22302                .map(|pp| (pp.0.len(), pp.1.len()))
22303                .unwrap_or((0, 0));
22304            if let Some(old) = old {
22305                self.fa_part_retired.lock().unwrap().push(old);
22306            }
22307            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22308                eprintln!(
22309                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22310                    co, o_len, cm, ml_len
22311                );
22312            }
22313            *part_guard = Some((
22314                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22315                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22316                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22317            ));
22318        }
22319        let pg = part_guard.as_mut().unwrap();
22320        self.gpu
22321            .stream()
22322            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22323        self.gpu
22324            .stream()
22325            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22326        self.gpu
22327            .stream()
22328            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22329        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22330        let f = self.func("fa_decode_vec_q_rows_v3_dc");
22331        let sh = (32 * head_dim * 2) as u32;
22332        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22333        f.set_attribute(
22334            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22335            sh as i32,
22336        )?;
22337        let cfg = LaunchConfig {
22338            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22339            block_dim: (32, gqa, 1),
22340            shared_mem_bytes: sh,
22341        };
22342        let __s_b = self.gpu.stream();
22343        let mut b = __s_b.launch_builder(&f);
22344        b.arg(q)
22345            .arg(k)
22346            .arg(v)
22347            .arg(&mut *part_o)
22348            .arg(&mut *part_m)
22349            .arg(&mut *part_l)
22350            .arg(&hd)
22351            .arg(&nh)
22352            .arg(&nhkv)
22353            .arg(base_dev)
22354            .arg(&scale)
22355            .arg(&nspm)
22356            .arg(&spk)
22357            .arg(&ktb)
22358            .arg(&vtb);
22359        unsafe {
22360            b.launch(cfg)?;
22361        }
22362        let fc = self.func("fa_decode_combine_rows_dc");
22363        let cfg2 = LaunchConfig {
22364            grid_dim: (n_head as u32, t as u32, 1),
22365            block_dim: (head_dim as u32, 1, 1),
22366            shared_mem_bytes: 0,
22367        };
22368        let plus0 = 0i32;
22369        let __s_b2 = self.gpu.stream();
22370        let mut b2 = __s_b2.launch_builder(&fc);
22371        b2.arg(&*part_o)
22372            .arg(&*part_m)
22373            .arg(&*part_l)
22374            .arg(o)
22375            .arg(&hd)
22376            .arg(&nh)
22377            .arg(base_dev)
22378            .arg(&plus0)
22379            .arg(&nspm)
22380            .arg(&spk);
22381        unsafe {
22382            b2.launch(cfg2)?;
22383        }
22384        Ok(())
22385    }
22386
22387    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
22388    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
22389    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
22390    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
22391    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
22392    ///
22393    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
22394    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
22395    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
22396    /// grouping (different but mathematically-equal log-sum-exp merge).
22397    pub fn fa_decode_dc(
22398        &self,
22399        q: &CudaSlice<f32>,
22400        k: &cudarc::driver::CudaView<u8>,
22401        v: &cudarc::driver::CudaView<u8>,
22402        o: &mut CudaSlice<f32>,
22403        head_dim: usize,
22404        n_head: usize,
22405        n_head_kv: usize,
22406        t_kv_dev: &CudaSlice<i32>,
22407        bucket_max: usize,
22408        scale: f32,
22409        k_tok_bytes: usize,
22410        v_tok_bytes: usize,
22411        g: bool,
22412    ) -> Result<(), Box<dyn std::error::Error>> {
22413        self.fa_decode_dc_q8(
22414            q,
22415            k,
22416            v,
22417            o,
22418            head_dim,
22419            n_head,
22420            n_head_kv,
22421            t_kv_dev,
22422            bucket_max,
22423            scale,
22424            k_tok_bytes,
22425            v_tok_bytes,
22426            g,
22427            None,
22428        )
22429    }
22430
22431    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
22432    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
22433    #[allow(clippy::too_many_arguments)]
22434    pub fn fa_decode_dc_q8(
22435        &self,
22436        q: &CudaSlice<f32>,
22437        k: &cudarc::driver::CudaView<u8>,
22438        v: &cudarc::driver::CudaView<u8>,
22439        o: &mut CudaSlice<f32>,
22440        head_dim: usize,
22441        n_head: usize,
22442        n_head_kv: usize,
22443        t_kv_dev: &CudaSlice<i32>,
22444        bucket_max: usize,
22445        scale: f32,
22446        k_tok_bytes: usize,
22447        v_tok_bytes: usize,
22448        g: bool,
22449        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22450    ) -> Result<(), Box<dyn std::error::Error>> {
22451        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
22452        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
22453        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
22454        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
22455        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
22456        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
22457        // 2026-07-12).
22458        let mut fa_vec =
22459            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22460        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22461            fa_vec = false;
22462        } // mirror kvmod/geom
22463        let sp = fa_split_keys(bucket_max, n_head_kv);
22464        let n_splits = if fa_vec {
22465            ((bucket_max + sp - 1) / sp).max(1)
22466        } else {
22467            ((bucket_max + 255) / 256).max(1)
22468        };
22469        let o_len = n_head * n_splits * head_dim;
22470        let ml_len = n_head * n_splits;
22471        let mut part_guard = self.fa_part_pool.lock().unwrap();
22472        if part_guard
22473            .as_ref()
22474            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22475            .unwrap_or(true)
22476        {
22477            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22478            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22479            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22480            // later live allocations land at those addresses, and the next graph REPLAY writes
22481            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22482            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22483            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22484            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22485            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22486            // (total retired < final size).
22487            let old = part_guard.take();
22488            let (co, cm) = old
22489                .as_ref()
22490                .map(|pp| (pp.0.len(), pp.1.len()))
22491                .unwrap_or((0, 0));
22492            if let Some(old) = old {
22493                self.fa_part_retired.lock().unwrap().push(old);
22494            }
22495            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22496                eprintln!(
22497                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22498                    co, o_len, cm, ml_len
22499                );
22500            }
22501            *part_guard = Some((
22502                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22503                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22504                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22505            ));
22506        }
22507        let pg = part_guard.as_mut().unwrap();
22508        self.gpu
22509            .stream()
22510            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22511        self.gpu
22512            .stream()
22513            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22514        self.gpu
22515            .stream()
22516            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22517        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22518        let (hd, nh, nhkv, nsp) = (
22519            head_dim as i32,
22520            n_head as i32,
22521            n_head_kv as i32,
22522            n_splits as i32,
22523        );
22524        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22525        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22526        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
22527        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
22528        let deep = fa_vec
22529            && head_dim == 256
22530            && fa_v4_at(bucket_max)
22531            && !g
22532            && fa_deep_at(bucket_max)
22533            && !matches!(fa_v4_mode(), "noB3" | "stage");
22534        let (f, cfg) = if fa_vec
22535            && head_dim == 512
22536            && bucket_max >= {
22537                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22538                *FA512_MIN_DC.get_or_init(|| {
22539                    std::env::var("MEMRA_FA512_MIN")
22540                        .ok()
22541                        .and_then(|v| v.parse().ok())
22542                        .unwrap_or(512)
22543                })
22544            } {
22545            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
22546            let gqa = (n_head / n_head_kv).max(1) as u32;
22547            (
22548                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22549                LaunchConfig {
22550                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22551                    block_dim: (32, gqa, 1),
22552                    shared_mem_bytes: 0,
22553                },
22554            )
22555        } else if fa_vec && head_dim == 512 {
22556            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
22557            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
22558            let q_view = q.as_view();
22559            let mut o_view = o.as_view_mut();
22560            return self.fa_decode_scalar_unified(
22561                &q_view,
22562                k,
22563                v,
22564                &mut o_view,
22565                head_dim,
22566                n_head,
22567                n_head_kv,
22568                0,
22569                Some(t_kv_dev),
22570                scale,
22571                n_splits,
22572                sp,
22573                k_tok_bytes,
22574                v_tok_bytes,
22575                g,
22576                &mut *part_o,
22577                &mut *part_m,
22578                &mut *part_l,
22579                q8_out,
22580            );
22581        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22582            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
22583            // incl the g-module route + raw-e4m3 sV sizing.
22584            let gqa = (n_head / n_head_kv).max(1) as u32;
22585            let fv = if g {
22586                self.func_g("fa_decode_vec_q_v4_dc")
22587            } else if deep {
22588                self.func("fa_decode_vec_q_v4_deep_dc")
22589            } else {
22590                self.func("fa_decode_vec_q_v4_dc")
22591            };
22592            let shmem =
22593                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22594            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22595            fv.set_attribute(
22596                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22597                shmem as i32,
22598            )?;
22599            (
22600                fv,
22601                LaunchConfig {
22602                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22603                    block_dim: (32, gqa, 1),
22604                    shared_mem_bytes: shmem,
22605                },
22606            )
22607        } else if fa_vec && fa_v3_active(head_dim) {
22608            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
22609            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
22610            let gqa = (n_head / n_head_kv).max(1) as u32;
22611            let fv = if g {
22612                self.func_g("fa_decode_vec_q_v3_dc")
22613            } else {
22614                self.func("fa_decode_vec_q_v3_dc")
22615            };
22616            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
22617            (
22618                fv,
22619                LaunchConfig {
22620                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22621                    block_dim: (32, gqa, 1),
22622                    shared_mem_bytes: shmem,
22623                },
22624            )
22625        } else if fa_vec && fa_v2_on() {
22626            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
22627            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
22628            // a numeric config; eager, rows-verify and graph all switch together).
22629            let gqa = (n_head / n_head_kv).max(1) as u32;
22630            let fv = if g {
22631                self.func_g("fa_decode_vec_q_v2_dc")
22632            } else {
22633                self.func("fa_decode_vec_q_v2_dc")
22634            };
22635            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22636            (
22637                fv,
22638                LaunchConfig {
22639                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22640                    block_dim: (32, gqa, 1),
22641                    shared_mem_bytes: shmem,
22642                },
22643            )
22644        } else if fa_vec {
22645            let gqa = (n_head / n_head_kv).max(1) as u32;
22646            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
22647            let fv = if g {
22648                self.func_g("fa_decode_vec_q_dc")
22649            } else {
22650                self.func("fa_decode_vec_q_dc")
22651            };
22652            (
22653                fv,
22654                LaunchConfig {
22655                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22656                    block_dim: (32, gqa, 1),
22657                    shared_mem_bytes: 0,
22658                },
22659            )
22660        } else {
22661            let q_view = q.as_view();
22662            let mut o_view = o.as_view_mut();
22663            return self.fa_decode_scalar_unified(
22664                &q_view,
22665                k,
22666                v,
22667                &mut o_view,
22668                head_dim,
22669                n_head,
22670                n_head_kv,
22671                0,
22672                Some(t_kv_dev),
22673                scale,
22674                n_splits,
22675                if fa_vec { sp } else { 256 },
22676                k_tok_bytes,
22677                v_tok_bytes,
22678                g,
22679                &mut *part_o,
22680                &mut *part_m,
22681                &mut *part_l,
22682                q8_out,
22683            );
22684        };
22685        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
22686        let __s_b = self.gpu.stream();
22687        let mut b = __s_b.launch_builder(&f);
22688        b.arg(q)
22689            .arg(k)
22690            .arg(v)
22691            .arg(&mut *part_o)
22692            .arg(&mut *part_m)
22693            .arg(&mut *part_l)
22694            .arg(&hd)
22695            .arg(&nh)
22696            .arg(&nhkv)
22697            .arg(t_kv_dev)
22698            .arg(&scale)
22699            .arg(&nsp)
22700            .arg(&ski)
22701            .arg(&ktb)
22702            .arg(&vtb);
22703        unsafe {
22704            b.launch(cfg)?;
22705        }
22706        let cfg2 = LaunchConfig {
22707            grid_dim: (n_head as u32, 1, 1),
22708            block_dim: (head_dim as u32, 1, 1),
22709            shared_mem_bytes: 0,
22710        };
22711        if let Some((oq, od)) = q8_out {
22712            let fc = if g {
22713                self.func_g("fa_decode_combine_q8_1")
22714            } else {
22715                self.fa_func("fa_decode_combine_q8_1", head_dim)
22716            };
22717            let __s_b2 = self.gpu.stream();
22718            let mut b2 = __s_b2.launch_builder(&fc);
22719            b2.arg(&*part_o)
22720                .arg(&*part_m)
22721                .arg(&*part_l)
22722                .arg(oq)
22723                .arg(od)
22724                .arg(&hd)
22725                .arg(&nh)
22726                .arg(&nsp);
22727            unsafe {
22728                b2.launch(cfg2)?;
22729            }
22730            return Ok(());
22731        }
22732        let fc = if g {
22733            self.func_g("fa_decode_combine_f32")
22734        } else {
22735            self.fa_func("fa_decode_combine_f32", head_dim)
22736        };
22737        let __s_b2 = self.gpu.stream();
22738        let mut b2 = __s_b2.launch_builder(&fc);
22739        b2.arg(&*part_o)
22740            .arg(&*part_m)
22741            .arg(&*part_l)
22742            .arg(o)
22743            .arg(&hd)
22744            .arg(&nh)
22745            .arg(&nsp);
22746        unsafe {
22747            b2.launch(cfg2)?;
22748        }
22749        Ok(())
22750    }
22751
22752    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
22753    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
22754    /// at equal rows.
22755    #[allow(clippy::too_many_arguments)]
22756    pub fn append_kv_quantized_dcw(
22757        &self,
22758        k_row: &CudaSlice<f32>,
22759        v_row: &CudaSlice<f32>,
22760        kc: &mut CudaSlice<u8>,
22761        vc: &mut CudaSlice<u8>,
22762        len_dev: &CudaSlice<i32>,
22763        base_dev: Option<&CudaSlice<i32>>,
22764        kv_dim_k: usize,
22765        kv_dim_v: usize,
22766        k_tok_bytes: usize,
22767        v_tok_bytes: usize,
22768    ) -> Result<(), Box<dyn std::error::Error>> {
22769        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22770        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22771        let cfg = LaunchConfig {
22772            grid_dim: (nblk, 1, 1),
22773            block_dim: (32, 1, 1),
22774            shared_mem_bytes: 0,
22775        };
22776        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22777        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22778        let null: u64 = 0;
22779        let __s_b = self.gpu.stream();
22780        let mut b = __s_b.launch_builder(&f);
22781        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22782        match base_dev {
22783            Some(base) => {
22784                b.arg(base);
22785            }
22786            None => {
22787                b.arg(&null);
22788            }
22789        }
22790        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22791        unsafe {
22792            b.launch(cfg)?;
22793        }
22794        Ok(())
22795    }
22796
22797    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
22798    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22799        let f = self.func("inc_i32");
22800        let cfg = LaunchConfig {
22801            grid_dim: (1, 1, 1),
22802            block_dim: (1, 1, 1),
22803            shared_mem_bytes: 0,
22804        };
22805        let __s_b = self.gpu.stream();
22806        let mut b = __s_b.launch_builder(&f);
22807        b.arg(counter);
22808        unsafe {
22809            b.launch(cfg)?;
22810        }
22811        Ok(())
22812    }
22813
22814    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
22815    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
22816    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
22817    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
22818    /// kernel class on this lane); callers keep eager below the vec floor and for any other
22819    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
22820    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
22821    /// alive across bucket growth.
22822    #[allow(clippy::too_many_arguments)]
22823    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
22824    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
22825    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
22826    fn fa_part_pool_grow(
22827        &self,
22828        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22829        o_len: usize,
22830        ml_len: usize,
22831    ) -> Result<(), Box<dyn std::error::Error>> {
22832        if part_guard
22833            .as_ref()
22834            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22835            .unwrap_or(true)
22836        {
22837            let old = part_guard.take();
22838            let (co, cm) = old
22839                .as_ref()
22840                .map(|pp| (pp.0.len(), pp.1.len()))
22841                .unwrap_or((0, 0));
22842            if let Some(old) = old {
22843                self.fa_part_retired.lock().unwrap().push(old);
22844            }
22845            *part_guard = Some((
22846                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22847                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22848                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22849            ));
22850        }
22851        Ok(())
22852    }
22853
22854    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
22855    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
22856    pub fn fa_dcw_pool_ensure(
22857        &self,
22858        head_dim: usize,
22859        n_head: usize,
22860        n_head_kv: usize,
22861        bucket_max: usize,
22862    ) -> Result<(), Box<dyn std::error::Error>> {
22863        let sp = fa_split_keys(bucket_max, n_head_kv);
22864        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22865        let o_len = n_head * n_splits * head_dim;
22866        let ml_len = n_head * n_splits;
22867        let mut part_guard = self.fa_part_pool.lock().unwrap();
22868        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
22869    }
22870
22871    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
22872    /// appended; one launch walks the KV stream once with two query rows (per-row causal
22873    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
22874    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
22875    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
22876    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
22877    /// outputs (the head gate fuses into the combine as in the t=1 path).
22878    #[allow(clippy::too_many_arguments)]
22879    pub fn fa_decode_dcw2(
22880        &self,
22881        q2: &CudaSlice<f32>,
22882        k_ring: &cudarc::driver::CudaView<u8>,
22883        v_ring: &cudarc::driver::CudaView<u8>,
22884        o2: &mut CudaSlice<f32>,
22885        head_dim: usize,
22886        n_head: usize,
22887        n_head_kv: usize,
22888        len_dev: &CudaSlice<i32>,
22889        base_dev: Option<&CudaSlice<i32>>,
22890        window: usize,
22891        bucket_max: usize,
22892        scale: f32,
22893        k_tok_bytes: usize,
22894        v_tok_bytes: usize,
22895        gate2: &CudaSlice<f32>,
22896    ) -> Result<(), Box<dyn std::error::Error>> {
22897        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22898        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
22899            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
22900        }
22901        let sp = fa_split_keys(bucket_max, n_head_kv);
22902        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22903        // Partials for BOTH rows: row-major halves.
22904        let o_len = 2 * n_head * n_splits * head_dim;
22905        let ml_len = 2 * n_head * n_splits;
22906        let mut part_guard = self.fa_part_pool.lock().unwrap();
22907        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
22908        let pg = part_guard.as_mut().unwrap();
22909        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22910        let (hd, nh, nhkv, nsp) = (
22911            head_dim as i32,
22912            n_head as i32,
22913            n_head_kv as i32,
22914            n_splits as i32,
22915        );
22916        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22917        let (ski, win) = (sp as i32, window as i32);
22918        let gqa = (n_head / n_head_kv).max(1) as u32;
22919        let smem = (32 * head_dim * 2) as u32;
22920        let f = self.func("fa_decode_vec_q_v3_dcw2");
22921        let cfg = LaunchConfig {
22922            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22923            block_dim: (32, gqa, 1),
22924            shared_mem_bytes: smem,
22925        };
22926        let null: u64 = 0;
22927        {
22928            let __s_b = self.gpu.stream();
22929            let mut b = __s_b.launch_builder(&f);
22930            b.arg(q2)
22931                .arg(k_ring)
22932                .arg(v_ring)
22933                .arg(&mut *part_o)
22934                .arg(&mut *part_m)
22935                .arg(&mut *part_l)
22936                .arg(&hd)
22937                .arg(&nh)
22938                .arg(&nhkv)
22939                .arg(len_dev);
22940            match base_dev {
22941                Some(base) => {
22942                    b.arg(base);
22943                }
22944                None => {
22945                    b.arg(&null);
22946                }
22947            }
22948            b.arg(&win)
22949                .arg(&scale)
22950                .arg(&nsp)
22951                .arg(&ski)
22952                .arg(&ktb)
22953                .arg(&vtb);
22954            unsafe {
22955                b.launch(cfg)?;
22956            }
22957        }
22958        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
22959        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
22960        // one launch covers both rows with the exact t=1 program per (row, head).
22961        let fc = {
22962            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22963            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
22964                self.func("fa_decode_combine_gate_f32_s")
22965            } else {
22966                self.func("fa_decode_combine_gate_f32")
22967            }
22968        };
22969        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
22970        let nh2 = (2 * n_head) as i32;
22971        let cfg2 = LaunchConfig {
22972            grid_dim: ((2 * n_head) as u32, 1, 1),
22973            block_dim: (head_dim as u32, 1, 1),
22974            shared_mem_bytes: if combine_shared {
22975                (2 * n_splits * 4) as u32
22976            } else {
22977                0
22978            },
22979        };
22980        let __s_b2 = self.gpu.stream();
22981        let mut b2 = __s_b2.launch_builder(&fc);
22982        b2.arg(&*part_o)
22983            .arg(&*part_m)
22984            .arg(&*part_l)
22985            .arg(gate2)
22986            .arg(o2)
22987            .arg(&hd)
22988            .arg(&nh2)
22989            .arg(&nsp);
22990        unsafe {
22991            b2.launch(cfg2)?;
22992        }
22993        Ok(())
22994    }
22995
22996    pub fn fa_decode_dcw(
22997        &self,
22998        q: &CudaSlice<f32>,
22999        k_ring: &cudarc::driver::CudaView<u8>,
23000        v_ring: &cudarc::driver::CudaView<u8>,
23001        o: &mut CudaSlice<f32>,
23002        head_dim: usize,
23003        n_head: usize,
23004        n_head_kv: usize,
23005        len_dev: &CudaSlice<i32>,
23006        base_dev: Option<&CudaSlice<i32>>,
23007        window: usize,
23008        bucket_max: usize,
23009        scale: f32,
23010        k_tok_bytes: usize,
23011        v_tok_bytes: usize,
23012        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
23013        // one launch saved); `o` then receives the GATED output and the caller skips its
23014        // attn_head_gate call.
23015        fused_gate: Option<&CudaSlice<f32>>,
23016    ) -> Result<(), Box<dyn std::error::Error>> {
23017        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23018        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23019            return Err("fa_decode_dcw supports the default v3-vec class only                         (bucket >= vec floor, head_dim <= 256, MEMRA_FA_V3 on);                         keep eager outside it"
23020                .into());
23021        }
23022        let sp = fa_split_keys(bucket_max, n_head_kv);
23023        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23024        let o_len = n_head * n_splits * head_dim;
23025        let ml_len = n_head * n_splits;
23026        let mut part_guard = self.fa_part_pool.lock().unwrap();
23027        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23028        let pg = part_guard.as_mut().unwrap();
23029        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
23030        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
23031        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
23032        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
23033        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23034        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
23035        // finds the attention children BY their three-memset signature and updates the
23036        // memset widths per bucket — capturing without them silently kills retargeting
23037        // (battery-v8 token drift, 2026-08-21).
23038        let memset_on = *MEMSET_ON
23039            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23040            || crate::tp::token_graph_building();
23041        if memset_on {
23042            self.gpu
23043                .stream()
23044                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23045            self.gpu
23046                .stream()
23047                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23048            self.gpu
23049                .stream()
23050                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23051        }
23052        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23053        let (hd, nh, nhkv, nsp) = (
23054            head_dim as i32,
23055            n_head as i32,
23056            n_head_kv as i32,
23057            n_splits as i32,
23058        );
23059        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23060        let (ski, win) = (sp as i32, window as i32);
23061        let gqa = (n_head / n_head_kv).max(1) as u32;
23062        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
23063        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
23064        // see fa_dec_v3_walk_u). Same launch geometry.
23065        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23066        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23067        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23068            Ok("2") => 2,
23069            Ok("1") => 1,
23070            _ => 0,
23071        });
23072        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
23073        // permission-blocked in this container and the module params are not exposed, so this
23074        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
23075        // prints cumulative cycle shares every 430 launches.
23076        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23077        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23078        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23079            std::sync::Mutex::new(None);
23080        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
23081        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
23082        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
23083        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23084        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23085            && (n_head / n_head_kv) % 2 == 0
23086            && (n_head / n_head_kv) >= 2;
23087        let f = if fprof {
23088            self.func("fa_decode_vec_q_v3_dcw_prof")
23089        } else if hs2 {
23090            self.func("fa_decode_vec_q_v3_dcw_hs2")
23091        } else if hoist == 2 {
23092            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
23093            self.func("fa_decode_vec_q_v3_dcw_hc")
23094        } else if hoist == 1 {
23095            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
23096            self.func("fa_decode_vec_q_v3_dcw_h")
23097        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23098            self.func("fa_decode_vec_q_v3_dcw_u8")
23099        } else {
23100            self.func("fa_decode_vec_q_v3_dcw")
23101        };
23102        let cfg = LaunchConfig {
23103            grid_dim: if hs2 {
23104                ((2 * n_head_kv) as u32, n_splits as u32, 1)
23105            } else {
23106                (n_head_kv as u32, n_splits as u32, 1)
23107            },
23108            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23109            shared_mem_bytes: smem,
23110        };
23111        let null: u64 = 0;
23112        let __s_b = self.gpu.stream();
23113        let mut b = __s_b.launch_builder(&f);
23114        b.arg(q)
23115            .arg(k_ring)
23116            .arg(v_ring)
23117            .arg(&mut *part_o)
23118            .arg(&mut *part_m)
23119            .arg(&mut *part_l)
23120            .arg(&hd)
23121            .arg(&nh)
23122            .arg(&nhkv)
23123            .arg(len_dev);
23124        match base_dev {
23125            Some(base) => {
23126                b.arg(base);
23127            }
23128            None => {
23129                b.arg(&null);
23130            }
23131        }
23132        b.arg(&win)
23133            .arg(&scale)
23134            .arg(&nsp)
23135            .arg(&ski)
23136            .arg(&ktb)
23137            .arg(&vtb);
23138        if fprof {
23139            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
23140            if guard
23141                .as_ref()
23142                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
23143            {
23144                *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
23145            }
23146            let (_, buf) = guard.as_mut().expect("armed above");
23147            b.arg(&*buf);
23148            unsafe {
23149                b.launch(cfg)?;
23150            }
23151            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23152            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
23153            if n % 430 == 0 {
23154                self.stream().synchronize()?;
23155                let h = self.dtoh_u64(buf)?;
23156                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
23157                let tot: u64 = h[..6].iter().sum();
23158                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
23159                for (i, name) in phases.iter().enumerate() {
23160                    let pct = if tot > 0 {
23161                        h[i] as f64 / tot as f64 * 100.0
23162                    } else {
23163                        0.0
23164                    };
23165                    line.push_str(&format!(" {name}={pct:.1}%"));
23166                }
23167                if h[6] > 0 {
23168                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
23169                }
23170                eprintln!("{line}");
23171            }
23172        } else {
23173            unsafe {
23174                b.launch(cfg)?;
23175            }
23176        }
23177        let mut combine_shared = false;
23178        let fc = if fused_gate.is_some() {
23179            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
23180            // n_splits-deep dependent global load chain every thread used to walk twice).
23181            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23182            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23183                combine_shared = true;
23184                self.func("fa_decode_combine_gate_f32_s")
23185            } else {
23186                self.func("fa_decode_combine_gate_f32")
23187            }
23188        } else {
23189            self.fa_func("fa_decode_combine_f32", head_dim)
23190        };
23191        let cfg2 = LaunchConfig {
23192            grid_dim: (n_head as u32, 1, 1),
23193            block_dim: (head_dim as u32, 1, 1),
23194            shared_mem_bytes: if combine_shared {
23195                (2 * n_splits * 4) as u32
23196            } else {
23197                0
23198            },
23199        };
23200        let __s_b2 = self.gpu.stream();
23201        let mut b2 = __s_b2.launch_builder(&fc);
23202        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
23203        if let Some(gate_row) = fused_gate {
23204            b2.arg(gate_row);
23205        }
23206        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
23207        unsafe {
23208            b2.launch(cfg2)?;
23209        }
23210        Ok(())
23211    }
23212
23213    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
23214    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
23215    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
23216    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
23217    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
23218    pub fn fa_geom_eager(
23219        &self,
23220        t_kv: usize,
23221        head_dim: usize,
23222        n_head_kv: usize,
23223        g: bool,
23224    ) -> (bool, usize) {
23225        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
23226        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
23227        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
23228        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
23229        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
23230        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
23231        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
23232        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
23233        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
23234        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
23235        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
23236        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
23237        // family; everything else falls to the g-module scalar.
23238        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
23239        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
23240        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
23241        if g && head_dim == 256 && !fa_v4_at(t_kv) {
23242            fa_vec = false;
23243        }
23244        let sp = fa_split_keys(t_kv, n_head_kv);
23245        let n_splits = if fa_vec {
23246            ((t_kv + sp - 1) / sp).max(1)
23247        } else {
23248            ((t_kv + 255) / 256).max(1)
23249        };
23250        (fa_vec, n_splits)
23251    }
23252
23253    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
23254    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
23255    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
23256    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
23257    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
23258    pub fn fa_bucket_key(
23259        &self,
23260        t_kv: usize,
23261        head_dim: usize,
23262        n_head_kv: usize,
23263        g: bool,
23264    ) -> (bool, usize) {
23265        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
23266    }
23267
23268    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
23269    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
23270    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
23271    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
23272    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
23273    /// device data) — every per-step varying scalar must come from a device counter. Returns the
23274    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
23275    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
23276    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
23277    /// replays (transients returning to the pool get reused by unrelated work and corrupt
23278    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
23279    pub fn capture_graph_retained<F>(
23280        &self,
23281        step: F,
23282    ) -> Result<
23283        (
23284            cudarc::driver::CudaGraph,
23285            Vec<Box<dyn std::any::Any + Send>>,
23286        ),
23287        Box<dyn std::error::Error>,
23288    >
23289    where
23290        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23291    {
23292        use cudarc::driver::sys::CUgraphInstantiate_flags;
23293        self.capture_graph_retained_flags(
23294            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23295            step,
23296        )
23297    }
23298
23299    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
23300    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
23301    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
23302    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
23303    pub fn capture_graph_retained_flags<F>(
23304        &self,
23305        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
23306        mut step: F,
23307    ) -> Result<
23308        (
23309            cudarc::driver::CudaGraph,
23310            Vec<Box<dyn std::any::Any + Send>>,
23311        ),
23312        Box<dyn std::error::Error>,
23313    >
23314    where
23315        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23316    {
23317        use cudarc::driver::sys::CUstreamCaptureMode;
23318        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
23319        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
23320        // while the capture region is open become dead copy NODES replayed every launch
23321        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
23322        // warmup runs allocate the same transient sequence at the same pool addresses, so
23323        // retaining the warmup clones preserves the draft-graph fix without polluting the
23324        // captured graph.
23325        self.capture_keep.lock().unwrap().clear();
23326        let was_tracking = self.gpu.ctx.is_event_tracking();
23327        if was_tracking {
23328            unsafe {
23329                self.gpu.ctx.disable_event_tracking();
23330            }
23331        }
23332        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23333            self.capture_keep_on
23334                .store(true, std::sync::atomic::Ordering::Relaxed);
23335            let w = (|| {
23336                step(self)?;
23337                step(self)
23338            })();
23339            self.capture_keep_on
23340                .store(false, std::sync::atomic::Ordering::Relaxed);
23341            w?;
23342            self.gpu.stream().synchronize()?;
23343            self.gpu
23344                .stream()
23345                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23346            let r = step(self);
23347            let g = self.gpu.stream().end_capture(flags);
23348            r?;
23349            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23350            graph.upload()?;
23351            Ok(graph)
23352        };
23353        let result = run();
23354        self.capture_keep_on
23355            .store(false, std::sync::atomic::Ordering::Relaxed);
23356        if was_tracking {
23357            unsafe {
23358                self.gpu.ctx.enable_event_tracking();
23359            }
23360        }
23361        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
23362        Ok((result?, keeper))
23363    }
23364
23365    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
23366    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
23367    /// alloc-free with persistent operands, and their bodies carry device side effects
23368    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
23369    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
23370    pub fn capture_graph_retained_nowarm<F>(
23371        &self,
23372        mut step: F,
23373    ) -> Result<
23374        (
23375            cudarc::driver::CudaGraph,
23376            Vec<Box<dyn std::any::Any + Send>>,
23377        ),
23378        Box<dyn std::error::Error>,
23379    >
23380    where
23381        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23382    {
23383        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23384        let was_tracking = self.gpu.ctx.is_event_tracking();
23385        if was_tracking {
23386            unsafe {
23387                self.gpu.ctx.disable_event_tracking();
23388            }
23389        }
23390        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23391            self.gpu.stream().synchronize()?;
23392            self.gpu
23393                .stream()
23394                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23395            let r = step(self);
23396            let g = self.gpu.stream().end_capture(
23397                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23398            );
23399            r?;
23400            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23401            graph.upload()?;
23402            Ok(graph)
23403        };
23404        let result = run();
23405        if was_tracking {
23406            unsafe {
23407                self.gpu.ctx.enable_event_tracking();
23408            }
23409        }
23410        Ok((result?, Vec::new()))
23411    }
23412
23413    pub fn capture_graph<F>(
23414        &self,
23415        mut step: F,
23416    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
23417    where
23418        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23419    {
23420        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23421        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
23422        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
23423        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
23424        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
23425        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
23426        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
23427        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
23428        let was_tracking = self.gpu.ctx.is_event_tracking();
23429        if was_tracking {
23430            unsafe {
23431                self.gpu.ctx.disable_event_tracking();
23432            }
23433        }
23434        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
23435        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
23436        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
23437        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
23438        // measure that scan's real cost on the generic path. Diagnostic door only; the
23439        // default stays AUTO_FREE until a measured A/B justifies moving it.
23440        let iflag = {
23441            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23442            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23443                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
23444                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
23445                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23446                Ok("priority") => {
23447                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23448                }
23449                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23450            })
23451        };
23452        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
23453        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
23454        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
23455        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
23456        // eager step executions and are node-count-invariant. Printing the split bounds the
23457        // refactor's ceiling instead of assuming it.
23458        let ct = {
23459            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23460            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23461        };
23462        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
23463        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
23464        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
23465        // chased, and node-count-invariant, so no capture-body refactor could touch it.
23466        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
23467        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
23468        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
23469        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
23470        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
23471        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
23472        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
23473        // grow and never frees, resident counters/scratch, cache set in place), and the
23474        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
23475        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
23476        // settling and pool mapping. Arbitrated adversarially, not by taste:
23477        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
23478        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
23479        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
23480        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
23481        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
23482        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
23483        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
23484        let warmups = {
23485            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23486            *W.get_or_init(|| {
23487                std::env::var("MEMRA_GRAPH_WARMUPS")
23488                    .ok()
23489                    .and_then(|v| v.parse().ok())
23490                    .filter(|n| *n >= 1)
23491                    .unwrap_or(1)
23492            })
23493        };
23494        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23495            let t_w = std::time::Instant::now();
23496            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
23497            for _ in 0..warmups {
23498                step(self)?;
23499            }
23500            self.gpu.stream().synchronize()?;
23501            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23502            // capture the third run.
23503            let t_c = std::time::Instant::now();
23504            self.gpu
23505                .stream()
23506                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23507            // If the body errors mid-capture, end the capture before propagating so the stream isn't
23508            // left in a capturing state.
23509            let r = step(self);
23510            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23511            let t_i = std::time::Instant::now();
23512            let g = self.gpu.stream().end_capture(iflag);
23513            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23514            r?;
23515            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23516            let t_u = std::time::Instant::now();
23517            graph.upload()?;
23518            if ct {
23519                println!(
23520                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
23521                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
23522                    t_u.elapsed().as_secs_f64() * 1e3
23523                );
23524            }
23525            Ok(graph)
23526        };
23527        let result = run();
23528        if was_tracking {
23529            unsafe {
23530                self.gpu.ctx.enable_event_tracking();
23531            }
23532        }
23533        result
23534    }
23535
23536    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
23537    pub fn gdn_scan_s128_view(
23538        &self,
23539        q: &CudaSlice<f32>,
23540        k: &CudaSlice<f32>,
23541        v: &CudaSlice<f32>,
23542        g: &CudaSlice<f32>,
23543        beta: &CudaSlice<f32>,
23544        state_in: &cudarc::driver::CudaView<f32>,
23545        state_out: &mut cudarc::driver::CudaViewMut<f32>,
23546        o: &mut CudaSlice<f32>,
23547        n_head: usize,
23548        t: usize,
23549        scale: f32,
23550    ) -> Result<(), Box<dyn std::error::Error>> {
23551        let f = self.func("gdn_scan_s128");
23552        const S_V: u32 = 128;
23553        const WARP: u32 = 32;
23554        const COLS: u32 = 4;
23555        let cfg = LaunchConfig {
23556            grid_dim: (n_head as u32, 1, S_V / COLS),
23557            block_dim: (WARP, COLS, 1),
23558            shared_mem_bytes: 0,
23559        };
23560        let (h, ti) = (n_head as i32, t as i32);
23561        let __s_b = self.gpu.stream();
23562        let mut b = __s_b.launch_builder(&f);
23563        b.arg(q)
23564            .arg(k)
23565            .arg(v)
23566            .arg(g)
23567            .arg(beta)
23568            .arg(state_in)
23569            .arg(state_out)
23570            .arg(o)
23571            .arg(&h)
23572            .arg(&ti)
23573            .arg(&scale);
23574        unsafe {
23575            b.launch(cfg)?;
23576        }
23577        Ok(())
23578    }
23579
23580    /// conv1d where the input is a CudaView (resident conv state assembled in place).
23581    pub fn ssm_conv1d_view(
23582        &self,
23583        x: &cudarc::driver::CudaView<f32>,
23584        w: &CudaSlice<f32>,
23585        y: &mut CudaSlice<f32>,
23586        conv_dim: usize,
23587        t: usize,
23588        d_conv: usize,
23589        silu: bool,
23590    ) -> Result<(), Box<dyn std::error::Error>> {
23591        let f = self.func("ssm_conv1d_silu_f32");
23592        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
23593        let cfg = LaunchConfig {
23594            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23595            block_dim: (256, 1, 1),
23596            shared_mem_bytes: 0,
23597        };
23598        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23599        let __s_b = self.gpu.stream();
23600        let mut b = __s_b.launch_builder(&f);
23601        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23602        unsafe {
23603            b.launch(cfg)?;
23604        }
23605        Ok(())
23606    }
23607
23608    /// Depthwise causal conv1d + optional SiLU.
23609    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
23610    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
23611    /// FUSED prefill conv (token-major input, zero left-state): replaces
23612    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
23613    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
23614    pub fn ssm_conv1d_tm(
23615        &self,
23616        qkv_tm: &CudaSlice<f32>,
23617        w: &CudaSlice<f32>,
23618        y: &mut CudaSlice<f32>,
23619        conv_dim: usize,
23620        t: usize,
23621        d_conv: usize,
23622    ) -> Result<(), Box<dyn std::error::Error>> {
23623        let f = self.func("ssm_conv1d_tm_f32");
23624        let cfg = LaunchConfig {
23625            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23626            block_dim: (256, 1, 1),
23627            shared_mem_bytes: 0,
23628        };
23629        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23630        let __s_b = self.gpu.stream();
23631        let mut b = __s_b.launch_builder(&f);
23632        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23633        unsafe {
23634            b.launch(cfg)?;
23635        }
23636        Ok(())
23637    }
23638
23639    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
23640    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
23641    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
23642    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
23643    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
23644    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
23645    /// columns; the final ring == what T sequential decode ring rolls leave).
23646    pub fn ssm_conv1d_tm_state(
23647        &self,
23648        qkv_tm: &CudaSlice<f32>,
23649        conv_state: &mut CudaSlice<f32>,
23650        w: &CudaSlice<f32>,
23651        y: &mut CudaSlice<f32>,
23652        conv_dim: usize,
23653        t: usize,
23654        d_conv: usize,
23655    ) -> Result<(), Box<dyn std::error::Error>> {
23656        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23657    }
23658
23659    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
23660    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
23661    #[allow(clippy::too_many_arguments)]
23662    pub fn ssm_conv1d_tm_state_pad(
23663        &self,
23664        qkv_tm: &CudaSlice<f32>,
23665        conv_state: &mut CudaSlice<f32>,
23666        w: &CudaSlice<f32>,
23667        y: &mut CudaSlice<f32>,
23668        conv_dim: usize,
23669        t: usize,
23670        d_conv: usize,
23671        pad_len: Option<&CudaSlice<i32>>,
23672    ) -> Result<(), Box<dyn std::error::Error>> {
23673        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23674        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23675        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23676        // cloning first keeps the ordering trivially correct under any future stream split.
23677        let ring_old = if t < d_conv - 1 {
23678            Some(self.clone_dtod(conv_state)?)
23679        } else {
23680            None
23681        };
23682        {
23683            let f = self.func("ssm_conv1d_tm_state_f32");
23684            let cfg = LaunchConfig {
23685                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23686                block_dim: (256, 1, 1),
23687                shared_mem_bytes: 0,
23688            };
23689            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23690            let __s_b = self.gpu.stream();
23691            let mut b = __s_b.launch_builder(&f);
23692            b.arg(qkv_tm)
23693                .arg(&*conv_state)
23694                .arg(w)
23695                .arg(y)
23696                .arg(&cd)
23697                .arg(&ti)
23698                .arg(&dc);
23699            unsafe {
23700                b.launch(cfg)?;
23701            }
23702        }
23703        match (ring_old, pad_len) {
23704            (None, Some(len_d)) => {
23705                let f = self.func("ssm_conv_ring_update_dev_f32");
23706                let n = conv_dim * (d_conv - 1);
23707                let cfg = LaunchConfig::for_num_elems(n as u32);
23708                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23709                let __s_b = self.gpu.stream();
23710                let mut b = __s_b.launch_builder(&f);
23711                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23712                unsafe {
23713                    b.launch(cfg)?;
23714                }
23715            }
23716            (None, None) => {
23717                let f = self.func("ssm_conv_ring_update_f32");
23718                let n = conv_dim * (d_conv - 1);
23719                let cfg = LaunchConfig::for_num_elems(n as u32);
23720                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23721                let __s_b = self.gpu.stream();
23722                let mut b = __s_b.launch_builder(&f);
23723                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23724                unsafe {
23725                    b.launch(cfg)?;
23726                }
23727            }
23728            (Some(old), _) => {
23729                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23730            }
23731        }
23732        Ok(())
23733    }
23734
23735    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
23736    pub fn ssm_conv1d_tm_state_pad_v(
23737        &self,
23738        qkv_tm: &cudarc::driver::CudaView<f32>,
23739        conv_state: &mut CudaSlice<f32>,
23740        w: &CudaSlice<f32>,
23741        y: &mut CudaSlice<f32>,
23742        conv_dim: usize,
23743        t: usize,
23744        d_conv: usize,
23745        pad_len: Option<&CudaSlice<i32>>,
23746    ) -> Result<(), Box<dyn std::error::Error>> {
23747        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23748        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23749        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23750        // cloning first keeps the ordering trivially correct under any future stream split.
23751        let ring_old = if t < d_conv - 1 {
23752            Some(self.clone_dtod(conv_state)?)
23753        } else {
23754            None
23755        };
23756        {
23757            let f = self.func("ssm_conv1d_tm_state_f32");
23758            let cfg = LaunchConfig {
23759                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23760                block_dim: (256, 1, 1),
23761                shared_mem_bytes: 0,
23762            };
23763            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23764            let __s_b = self.gpu.stream();
23765            let mut b = __s_b.launch_builder(&f);
23766            b.arg(qkv_tm)
23767                .arg(&*conv_state)
23768                .arg(w)
23769                .arg(y)
23770                .arg(&cd)
23771                .arg(&ti)
23772                .arg(&dc);
23773            unsafe {
23774                b.launch(cfg)?;
23775            }
23776        }
23777        match (ring_old, pad_len) {
23778            (None, Some(len_d)) => {
23779                let f = self.func("ssm_conv_ring_update_dev_f32");
23780                let n = conv_dim * (d_conv - 1);
23781                let cfg = LaunchConfig::for_num_elems(n as u32);
23782                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23783                let __s_b = self.gpu.stream();
23784                let mut b = __s_b.launch_builder(&f);
23785                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23786                unsafe {
23787                    b.launch(cfg)?;
23788                }
23789            }
23790            (None, None) => {
23791                let f = self.func("ssm_conv_ring_update_f32");
23792                let n = conv_dim * (d_conv - 1);
23793                let cfg = LaunchConfig::for_num_elems(n as u32);
23794                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23795                let __s_b = self.gpu.stream();
23796                let mut b = __s_b.launch_builder(&f);
23797                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23798                unsafe {
23799                    b.launch(cfg)?;
23800                }
23801            }
23802            (Some(_), _) => unreachable!(
23803                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23804            ),
23805        }
23806        Ok(())
23807    }
23808
23809    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
23810    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
23811    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
23812    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
23813    pub fn ssm_conv_ring_rebuild(
23814        &self,
23815        qkv_tm: &CudaSlice<f32>,
23816        ring_old: &CudaSlice<f32>,
23817        conv_state: &mut CudaSlice<f32>,
23818        conv_dim: usize,
23819        tc: usize,
23820        d_conv: usize,
23821    ) -> Result<(), Box<dyn std::error::Error>> {
23822        let f = self.func("ssm_conv_ring_rebuild_f32");
23823        let n = conv_dim * (d_conv - 1);
23824        let cfg = LaunchConfig::for_num_elems(n as u32);
23825        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23826        let __s_b = self.gpu.stream();
23827        let mut b = __s_b.launch_builder(&f);
23828        b.arg(qkv_tm)
23829            .arg(ring_old)
23830            .arg(conv_state)
23831            .arg(&cd)
23832            .arg(&ti)
23833            .arg(&dc);
23834        unsafe {
23835            b.launch(cfg)?;
23836        }
23837        Ok(())
23838    }
23839
23840    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
23841    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
23842    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
23843    /// the argmax + run-spec gates are the authority.
23844    #[allow(clippy::too_many_arguments)]
23845    pub fn gdn_prep_decode(
23846        &self,
23847        conv_out: &CudaSlice<f32>,
23848        beta_raw: &CudaSlice<f32>,
23849        alpha: &CudaSlice<f32>,
23850        dt_bias: &CudaSlice<f32>,
23851        a: &CudaSlice<f32>,
23852        q_l2: &mut CudaSlice<f32>,
23853        k_l2: &mut CudaSlice<f32>,
23854        v_g: &mut CudaSlice<f32>,
23855        beta: &mut CudaSlice<f32>,
23856        g_log: &mut CudaSlice<f32>,
23857        d_state: usize,
23858        num_v: usize,
23859        num_k: usize,
23860        key_dim: usize,
23861        eps: f32,
23862    ) -> Result<(), Box<dyn std::error::Error>> {
23863        let f = self.func("gdn_prep_decode_f32");
23864        let cfg = LaunchConfig {
23865            grid_dim: (num_v as u32, 1, 1),
23866            block_dim: (32, 4, 1),
23867            shared_mem_bytes: 0,
23868        };
23869        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23870        let __s_b = self.gpu.stream();
23871        let mut b = __s_b.launch_builder(&f);
23872        b.arg(conv_out)
23873            .arg(beta_raw)
23874            .arg(alpha)
23875            .arg(dt_bias)
23876            .arg(a)
23877            .arg(q_l2)
23878            .arg(k_l2)
23879            .arg(v_g)
23880            .arg(beta)
23881            .arg(g_log)
23882            .arg(&ds)
23883            .arg(&nv)
23884            .arg(&nk)
23885            .arg(&kd)
23886            .arg(&eps);
23887        unsafe {
23888            b.launch(cfg)?;
23889        }
23890        Ok(())
23891    }
23892
23893    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
23894    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
23895    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
23896    #[allow(clippy::too_many_arguments)]
23897    pub fn ssm_conv1d_gdn(
23898        &self,
23899        qkv_tm: &CudaSlice<f32>,
23900        w: &CudaSlice<f32>,
23901        q_g: &mut CudaSlice<f32>,
23902        k_g: &mut CudaSlice<f32>,
23903        v_g: &mut CudaSlice<f32>,
23904        conv_dim: usize,
23905        t: usize,
23906        d_conv: usize,
23907        d_state: usize,
23908        num_v: usize,
23909        num_k: usize,
23910        key_dim: usize,
23911    ) -> Result<(), Box<dyn std::error::Error>> {
23912        let f = self.func("ssm_conv1d_gdn_f32");
23913        let cfg = LaunchConfig {
23914            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23915            block_dim: (256, 1, 1),
23916            shared_mem_bytes: 0,
23917        };
23918        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23919        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23920        let __s_b = self.gpu.stream();
23921        let mut b = __s_b.launch_builder(&f);
23922        b.arg(qkv_tm)
23923            .arg(w)
23924            .arg(q_g)
23925            .arg(k_g)
23926            .arg(v_g)
23927            .arg(&cd)
23928            .arg(&ti)
23929            .arg(&dc)
23930            .arg(&ds)
23931            .arg(&nv)
23932            .arg(&nk)
23933            .arg(&kd);
23934        unsafe {
23935            b.launch(cfg)?;
23936        }
23937        Ok(())
23938    }
23939
23940    pub fn ssm_conv1d(
23941        &self,
23942        x: &CudaSlice<f32>,
23943        w: &CudaSlice<f32>,
23944        y: &mut CudaSlice<f32>,
23945        conv_dim: usize,
23946        t: usize,
23947        d_conv: usize,
23948        silu: bool,
23949    ) -> Result<(), Box<dyn std::error::Error>> {
23950        let f = self.func("ssm_conv1d_silu_f32");
23951        let cfg = LaunchConfig {
23952            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23953            block_dim: (256, 1, 1),
23954            shared_mem_bytes: 0,
23955        };
23956        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23957        let __s_b = self.gpu.stream();
23958        let mut b = __s_b.launch_builder(&f);
23959        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23960        unsafe {
23961            b.launch(cfg)?;
23962        }
23963        Ok(())
23964    }
23965
23966    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
23967    /// o:[128,H,T]. Single sequence.
23968    pub fn gdn_scan_s128(
23969        &self,
23970        q: &CudaSlice<f32>,
23971        k: &CudaSlice<f32>,
23972        v: &CudaSlice<f32>,
23973        g: &CudaSlice<f32>,
23974        beta: &CudaSlice<f32>,
23975        state_in: &CudaSlice<f32>,
23976        state_out: &mut CudaSlice<f32>,
23977        o: &mut CudaSlice<f32>,
23978        n_head: usize,
23979        t: usize,
23980        scale: f32,
23981    ) -> Result<(), Box<dyn std::error::Error>> {
23982        let f = self.func("gdn_scan_s128");
23983        const S_V: u32 = 128;
23984        const WARP: u32 = 32;
23985        const COLS_PER_BLOCK: u32 = 4;
23986        let cfg = LaunchConfig {
23987            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
23988            block_dim: (WARP, COLS_PER_BLOCK, 1),
23989            shared_mem_bytes: 0,
23990        };
23991        let (h, ti) = (n_head as i32, t as i32);
23992        let __s_b = self.gpu.stream();
23993        let mut b = __s_b.launch_builder(&f);
23994        b.arg(q)
23995            .arg(k)
23996            .arg(v)
23997            .arg(g)
23998            .arg(beta)
23999            .arg(state_in)
24000            .arg(state_out)
24001            .arg(o)
24002            .arg(&h)
24003            .arg(&ti)
24004            .arg(&scale);
24005        unsafe {
24006            b.launch(cfg)?;
24007        }
24008        Ok(())
24009    }
24010
24011    // ==== B2' batched decode state ops (decode_batch.rs) ====
24012    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
24013    // Bodies are the single-seq kernels per sequence — bit-identical per row.
24014
24015    #[allow(clippy::too_many_arguments)]
24016    pub fn ssm_conv1d_fused_decode_b(
24017        &self,
24018        qkv_cols: &CudaSlice<f32>,
24019        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24020        w: &CudaSlice<f32>,
24021        conv_outs: &mut CudaSlice<f32>,
24022        conv_dim: usize,
24023        d_conv: usize,
24024        b_n: usize,
24025    ) -> Result<(), Box<dyn std::error::Error>> {
24026        let f = self.func("ssm_conv1d_fused_decode_b_f32");
24027        let cfg = LaunchConfig {
24028            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24029            block_dim: (256, 1, 1),
24030            shared_mem_bytes: 0,
24031        };
24032        let (cd, dc) = (conv_dim as i32, d_conv as i32);
24033        let __s_b = self.gpu.stream();
24034        let mut b = __s_b.launch_builder(&f);
24035        b.arg(qkv_cols)
24036            .arg(conv_state_ptrs)
24037            .arg(w)
24038            .arg(conv_outs)
24039            .arg(&cd)
24040            .arg(&dc);
24041        unsafe {
24042            b.launch(cfg)?;
24043        }
24044        Ok(())
24045    }
24046
24047    #[allow(clippy::too_many_arguments)]
24048    pub fn gdn_prep_decode_b(
24049        &self,
24050        conv_outs: &CudaSlice<f32>,
24051        beta_raws: &CudaSlice<f32>,
24052        alphas: &CudaSlice<f32>,
24053        dt_bias: &CudaSlice<f32>,
24054        a: &CudaSlice<f32>,
24055        q_l2: &mut CudaSlice<f32>,
24056        k_l2: &mut CudaSlice<f32>,
24057        v_g: &mut CudaSlice<f32>,
24058        beta: &mut CudaSlice<f32>,
24059        g_log: &mut CudaSlice<f32>,
24060        d_state: usize,
24061        num_v: usize,
24062        num_k: usize,
24063        key_dim: usize,
24064        eps: f32,
24065        conv_dim: usize,
24066        b_n: usize,
24067    ) -> Result<(), Box<dyn std::error::Error>> {
24068        let f = self.func("gdn_prep_decode_b_f32");
24069        let cfg = LaunchConfig {
24070            grid_dim: (num_v as u32, 1, b_n as u32),
24071            block_dim: (32, 4, 1),
24072            shared_mem_bytes: 0,
24073        };
24074        let (ds, nv, nk, kd, cd) = (
24075            d_state as i32,
24076            num_v as i32,
24077            num_k as i32,
24078            key_dim as i32,
24079            conv_dim as i32,
24080        );
24081        let __s_b = self.gpu.stream();
24082        let mut b = __s_b.launch_builder(&f);
24083        b.arg(conv_outs)
24084            .arg(beta_raws)
24085            .arg(alphas)
24086            .arg(dt_bias)
24087            .arg(a)
24088            .arg(q_l2)
24089            .arg(k_l2)
24090            .arg(v_g)
24091            .arg(beta)
24092            .arg(g_log)
24093            .arg(&ds)
24094            .arg(&nv)
24095            .arg(&nk)
24096            .arg(&kd)
24097            .arg(&eps)
24098            .arg(&cd);
24099        unsafe {
24100            b.launch(cfg)?;
24101        }
24102        Ok(())
24103    }
24104
24105    #[allow(clippy::too_many_arguments)]
24106    pub fn gdn_scan_s128_batched(
24107        &self,
24108        q: &CudaSlice<f32>,
24109        k: &CudaSlice<f32>,
24110        v: &CudaSlice<f32>,
24111        g: &CudaSlice<f32>,
24112        beta: &CudaSlice<f32>,
24113        state_in_ptrs: &cudarc::driver::CudaView<u64>,
24114        state_out_ptrs: &cudarc::driver::CudaView<u64>,
24115        o: &mut CudaSlice<f32>,
24116        n_head: usize,
24117        b_n: usize,
24118        scale: f32,
24119    ) -> Result<(), Box<dyn std::error::Error>> {
24120        let f = self.func("gdn_scan_s128_b");
24121        const S_V: u32 = 128;
24122        const WARP: u32 = 32;
24123        const COLS_PER_BLOCK: u32 = 4;
24124        let cfg = LaunchConfig {
24125            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24126            block_dim: (WARP, COLS_PER_BLOCK, 1),
24127            shared_mem_bytes: 0,
24128        };
24129        let h = n_head as i32;
24130        let __s_b = self.gpu.stream();
24131        let mut b = __s_b.launch_builder(&f);
24132        b.arg(q)
24133            .arg(k)
24134            .arg(v)
24135            .arg(g)
24136            .arg(beta)
24137            .arg(state_in_ptrs)
24138            .arg(state_out_ptrs)
24139            .arg(o)
24140            .arg(&h)
24141            .arg(&scale);
24142        unsafe {
24143            b.launch(cfg)?;
24144        }
24145        Ok(())
24146    }
24147
24148    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
24149    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
24150    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
24151    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
24152    /// numeric class; only the pointer arithmetic moved host-side.
24153    #[allow(clippy::too_many_arguments)]
24154    pub fn ssm_conv1d_fused_decode_b_view(
24155        &self,
24156        qkv_cols: &cudarc::driver::CudaView<f32>,
24157        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24158        w: &CudaSlice<f32>,
24159        conv_outs: &mut CudaSlice<f32>,
24160        conv_dim: usize,
24161        d_conv: usize,
24162        b_n: usize,
24163    ) -> Result<(), Box<dyn std::error::Error>> {
24164        let f = self.func("ssm_conv1d_fused_decode_b_f32");
24165        let cfg = LaunchConfig {
24166            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24167            block_dim: (256, 1, 1),
24168            shared_mem_bytes: 0,
24169        };
24170        let (cd, dc) = (conv_dim as i32, d_conv as i32);
24171        let __s_b = self.gpu.stream();
24172        let mut b = __s_b.launch_builder(&f);
24173        b.arg(qkv_cols)
24174            .arg(conv_state_ptrs)
24175            .arg(w)
24176            .arg(conv_outs)
24177            .arg(&cd)
24178            .arg(&dc);
24179        unsafe {
24180            b.launch(cfg)?;
24181        }
24182        Ok(())
24183    }
24184
24185    #[allow(clippy::too_many_arguments)]
24186    pub fn gdn_prep_decode_b_view(
24187        &self,
24188        conv_outs: &CudaSlice<f32>,
24189        beta_raws: &cudarc::driver::CudaView<f32>,
24190        alphas: &cudarc::driver::CudaView<f32>,
24191        dt_bias: &CudaSlice<f32>,
24192        a: &CudaSlice<f32>,
24193        q_l2: &mut CudaSlice<f32>,
24194        k_l2: &mut CudaSlice<f32>,
24195        v_g: &mut CudaSlice<f32>,
24196        beta: &mut CudaSlice<f32>,
24197        g_log: &mut CudaSlice<f32>,
24198        d_state: usize,
24199        num_v: usize,
24200        num_k: usize,
24201        key_dim: usize,
24202        eps: f32,
24203        conv_dim: usize,
24204        b_n: usize,
24205    ) -> Result<(), Box<dyn std::error::Error>> {
24206        let f = self.func("gdn_prep_decode_b_f32");
24207        let cfg = LaunchConfig {
24208            grid_dim: (num_v as u32, 1, b_n as u32),
24209            block_dim: (32, 4, 1),
24210            shared_mem_bytes: 0,
24211        };
24212        let (ds, nv, nk, kd, cd) = (
24213            d_state as i32,
24214            num_v as i32,
24215            num_k as i32,
24216            key_dim as i32,
24217            conv_dim as i32,
24218        );
24219        let __s_b = self.gpu.stream();
24220        let mut b = __s_b.launch_builder(&f);
24221        b.arg(conv_outs)
24222            .arg(beta_raws)
24223            .arg(alphas)
24224            .arg(dt_bias)
24225            .arg(a)
24226            .arg(q_l2)
24227            .arg(k_l2)
24228            .arg(v_g)
24229            .arg(beta)
24230            .arg(g_log)
24231            .arg(&ds)
24232            .arg(&nv)
24233            .arg(&nk)
24234            .arg(&kd)
24235            .arg(&eps)
24236            .arg(&cd);
24237        unsafe {
24238            b.launch(cfg)?;
24239        }
24240        Ok(())
24241    }
24242
24243    #[allow(clippy::too_many_arguments)]
24244    pub fn gdn_scan_s128_batched_view(
24245        &self,
24246        q: &CudaSlice<f32>,
24247        k: &CudaSlice<f32>,
24248        v: &CudaSlice<f32>,
24249        g: &CudaSlice<f32>,
24250        beta: &CudaSlice<f32>,
24251        state_in_ptrs: &cudarc::driver::CudaView<u64>,
24252        state_out_ptrs: &cudarc::driver::CudaView<u64>,
24253        o: &mut cudarc::driver::CudaViewMut<f32>,
24254        n_head: usize,
24255        b_n: usize,
24256        scale: f32,
24257    ) -> Result<(), Box<dyn std::error::Error>> {
24258        let f = self.func("gdn_scan_s128_b");
24259        const S_V: u32 = 128;
24260        const WARP: u32 = 32;
24261        const COLS_PER_BLOCK: u32 = 4;
24262        let cfg = LaunchConfig {
24263            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24264            block_dim: (WARP, COLS_PER_BLOCK, 1),
24265            shared_mem_bytes: 0,
24266        };
24267        let h = n_head as i32;
24268        let __s_b = self.gpu.stream();
24269        let mut b = __s_b.launch_builder(&f);
24270        b.arg(q)
24271            .arg(k)
24272            .arg(v)
24273            .arg(g)
24274            .arg(beta)
24275            .arg(state_in_ptrs)
24276            .arg(state_out_ptrs)
24277            .arg(o)
24278            .arg(&h)
24279            .arg(&scale);
24280        unsafe {
24281            b.launch(cfg)?;
24282        }
24283        Ok(())
24284    }
24285
24286    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
24287    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
24288    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
24289    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
24290    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
24291    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
24292    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
24293    /// identity law); prime_cache/forward/forward_last are the only callers.
24294    pub fn gdn_chunked_enabled() -> bool {
24295        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24296        *E.get_or_init(|| {
24297            std::env::var("MEMRA_GDN_CHUNKED")
24298                .map(|v| v != "0")
24299                .unwrap_or(true)
24300        })
24301    }
24302
24303    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
24304    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
24305    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
24306    /// of 32 in [32, 128] (kernel row mappings require it).
24307    pub fn gdn_chunk_size() -> usize {
24308        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24309        *C.get_or_init(|| {
24310            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
24311                .ok()
24312                .and_then(|v| v.parse().ok())
24313                .unwrap_or(32);
24314            c.clamp(32, 128) / 32 * 32
24315        })
24316    }
24317
24318    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
24319    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
24320    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
24321    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
24322    #[allow(clippy::too_many_arguments)]
24323    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
24324    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
24325    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
24326    #[allow(clippy::too_many_arguments)]
24327    pub fn gdn_chunk_k123(
24328        &self,
24329        q: &CudaSlice<f32>,
24330        k: &CudaSlice<f32>,
24331        v: &CudaSlice<f32>,
24332        g: &CudaSlice<f32>,
24333        beta: &CudaSlice<f32>,
24334        wb16: Option<&mut CudaSlice<u8>>,
24335        n_head: usize,
24336        t: usize,
24337        c: usize,
24338        hk: usize,
24339        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
24340    ) -> Result<
24341        (
24342            CudaSlice<f32>,
24343            CudaSlice<f32>,
24344            CudaSlice<f32>,
24345            CudaSlice<f32>,
24346        ),
24347        Box<dyn std::error::Error>,
24348    > {
24349        const D: usize = 128;
24350        let h = n_head;
24351        let nc = (t + c - 1) / c;
24352        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24353        let mut gcum = self.uninit(t * h)?;
24354        let mut a = self.uninit(nc * h * c * c)?;
24355        let mut p = self.uninit(nc * h * c * c)?;
24356        let mut u = self.uninit(nc * h * c * D)?;
24357        let mut w = self.uninit(nc * h * c * D)?;
24358        {
24359            // K1
24360            let f = self.func("gdn_chunk_cumgate_f32");
24361            let cfg = LaunchConfig {
24362                grid_dim: (nc as u32, h as u32, 1),
24363                block_dim: (32, 1, 1),
24364                shared_mem_bytes: 0,
24365            };
24366            let __s_b = self.gpu.stream();
24367            let mut b = __s_b.launch_builder(&f);
24368            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
24369            unsafe {
24370                b.launch(cfg)?;
24371            }
24372        }
24373        if let Some((qb, kb, pb)) = k2w {
24374            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
24375            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
24376            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
24377            let f = self.func("gdn_k2_wgmma");
24378            let cfg = LaunchConfig {
24379                grid_dim: (nc as u32, h as u32, 1),
24380                block_dim: (128, 1, 1),
24381                shared_mem_bytes: 0,
24382            };
24383            let hki = hk as i32;
24384            let __s_b = self.gpu.stream();
24385            let mut b = __s_b.launch_builder(&f);
24386            b.arg(qb)
24387                .arg(kb)
24388                .arg(&gcum)
24389                .arg(beta)
24390                .arg(&mut a)
24391                .arg(&mut *pb)
24392                .arg(&hi)
24393                .arg(&ti)
24394                .arg(&ci)
24395                .arg(&hki);
24396            unsafe {
24397                b.launch(cfg)?;
24398            }
24399        } else if c <= 64 && !portable_mma_gated() {
24400            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
24401            let f = self.func("gdn_chunk_attn_f32");
24402            f.set_attribute(
24403                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24404                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24405            )?;
24406            let jt = ((c + 31) / 32) as u32;
24407            let cfg = LaunchConfig {
24408                grid_dim: (nc as u32, h as u32, jt),
24409                block_dim: (256, 1, 1),
24410                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24411            };
24412            let hki = hk as i32;
24413            let __s_b = self.gpu.stream();
24414            let mut b = __s_b.launch_builder(&f);
24415            b.arg(q)
24416                .arg(k)
24417                .arg(&gcum)
24418                .arg(beta)
24419                .arg(&mut a)
24420                .arg(&mut p)
24421                .arg(&hi)
24422                .arg(&ti)
24423                .arg(&ci)
24424                .arg(&hki);
24425            unsafe {
24426                b.launch(cfg)?;
24427            }
24428        } else {
24429            // K2 generic (C = 128, or the portable target's low-smem fallback)
24430            assert!(
24431                hk == h,
24432                "generic K2 is broadcast-only (de-broadcast rides C==32)"
24433            );
24434            let f = self.func("gdn_chunk_attn_g_f32");
24435            let cfg = LaunchConfig {
24436                grid_dim: (nc as u32, h as u32, 1),
24437                block_dim: (32, 8, 1),
24438                shared_mem_bytes: 0,
24439            };
24440            let __s_b = self.gpu.stream();
24441            let mut b = __s_b.launch_builder(&f);
24442            b.arg(q)
24443                .arg(k)
24444                .arg(&gcum)
24445                .arg(beta)
24446                .arg(&mut a)
24447                .arg(&mut p)
24448                .arg(&hi)
24449                .arg(&ti)
24450                .arg(&ci);
24451            unsafe {
24452                b.launch(cfg)?;
24453            }
24454        }
24455        {
24456            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
24457            let cfg = LaunchConfig {
24458                grid_dim: (nc as u32, h as u32, 1),
24459                block_dim: (256, 1, 1),
24460                shared_mem_bytes: 0,
24461            };
24462            match c {
24463                32 | 64 => {
24464                    let f = self.func(if c == 32 {
24465                        "gdn_chunk_solve32_f32"
24466                    } else {
24467                        "gdn_chunk_solve64_f32"
24468                    });
24469                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
24470                    let wb: u64 = match wb16 {
24471                        Some(d) => self.addr_u8(d),
24472                        None => 0,
24473                    };
24474                    let hki = hk as i32;
24475                    let __s_b = self.gpu.stream();
24476                    let mut b = __s_b.launch_builder(&f);
24477                    b.arg(v)
24478                        .arg(k)
24479                        .arg(&a)
24480                        .arg(&gcum)
24481                        .arg(&mut u)
24482                        .arg(&mut w)
24483                        .arg(&wb)
24484                        .arg(&hi)
24485                        .arg(&ti)
24486                        .arg(&hki);
24487                    unsafe {
24488                        b.launch(cfg)?;
24489                    }
24490                }
24491                _ => {
24492                    assert!(hk == h, "generic K3 is broadcast-only");
24493                    let f = self.func("gdn_chunk_solve_f32");
24494                    let __s_b = self.gpu.stream();
24495                    let mut b = __s_b.launch_builder(&f);
24496                    b.arg(v)
24497                        .arg(k)
24498                        .arg(&a)
24499                        .arg(&gcum)
24500                        .arg(&mut u)
24501                        .arg(&mut w)
24502                        .arg(&hi)
24503                        .arg(&ti)
24504                        .arg(&ci);
24505                    unsafe {
24506                        b.launch(cfg)?;
24507                    }
24508                }
24509            }
24510        }
24511        Ok((gcum, p, u, w))
24512    }
24513
24514    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
24515    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
24516    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
24517    pub fn gdn_db_on() -> bool {
24518        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24519    }
24520
24521    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
24522    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
24523    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
24524    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
24525    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
24526    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
24527    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
24528    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
24529    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24530        !portable_mma_gated()
24531            && c == 32
24532            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24533                Ok("1") => true,
24534                Ok("0") => false,
24535                _ => gdn_mma_default_on(),
24536            }
24537    }
24538
24539    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
24540    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
24541    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
24542    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
24543    /// force would silently produce garbage. Required since the sm_120a mma default
24544    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
24545    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24546        cfg!(memra_hopper_mma)
24547            && self.gdn_mma_enabled(c)
24548            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24549    }
24550
24551    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
24552    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
24553    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
24554    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
24555    #[allow(clippy::too_many_arguments)]
24556    pub fn ssm_conv1d_gdn_state_pad(
24557        &self,
24558        qkv_tm: &cudarc::driver::CudaView<f32>,
24559        conv_state: &mut CudaSlice<f32>,
24560        w: &CudaSlice<f32>,
24561        q_g: &mut CudaSlice<f32>,
24562        k_g: &mut CudaSlice<f32>,
24563        v_g: &mut CudaSlice<f32>,
24564        conv_dim: usize,
24565        t: usize,
24566        d_conv: usize,
24567        d_state: usize,
24568        num_v: usize,
24569        num_k: usize,
24570        key_dim: usize,
24571        hk: usize,
24572        pad_len: Option<&CudaSlice<i32>>,
24573    ) -> Result<(), Box<dyn std::error::Error>> {
24574        assert!(
24575            t >= d_conv - 1,
24576            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24577        );
24578        {
24579            let f = self.func("ssm_conv1d_gdn_state_f32");
24580            let cfg = LaunchConfig {
24581                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24582                block_dim: (256, 1, 1),
24583                shared_mem_bytes: 0,
24584            };
24585            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24586            let (ds, nv, nk, kd, hki) = (
24587                d_state as i32,
24588                num_v as i32,
24589                num_k as i32,
24590                key_dim as i32,
24591                hk as i32,
24592            );
24593            let __s_b = self.gpu.stream();
24594            let mut b = __s_b.launch_builder(&f);
24595            b.arg(qkv_tm)
24596                .arg(&*conv_state)
24597                .arg(w)
24598                .arg(q_g)
24599                .arg(k_g)
24600                .arg(v_g)
24601                .arg(&cd)
24602                .arg(&ti)
24603                .arg(&dc)
24604                .arg(&ds)
24605                .arg(&nv)
24606                .arg(&nk)
24607                .arg(&kd)
24608                .arg(&hki);
24609            unsafe {
24610                b.launch(cfg)?;
24611            }
24612        }
24613        match pad_len {
24614            Some(len_d) => {
24615                let f = self.func("ssm_conv_ring_update_dev_f32");
24616                let n = conv_dim * (d_conv - 1);
24617                let cfg = LaunchConfig::for_num_elems(n as u32);
24618                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24619                let __s_b = self.gpu.stream();
24620                let mut b = __s_b.launch_builder(&f);
24621                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24622                unsafe {
24623                    b.launch(cfg)?;
24624                }
24625            }
24626            None => {
24627                let f = self.func("ssm_conv_ring_update_f32");
24628                let n = conv_dim * (d_conv - 1);
24629                let cfg = LaunchConfig::for_num_elems(n as u32);
24630                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24631                let __s_b = self.gpu.stream();
24632                let mut b = __s_b.launch_builder(&f);
24633                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24634                unsafe {
24635                    b.launch(cfg)?;
24636                }
24637            }
24638        }
24639        Ok(())
24640    }
24641
24642    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
24643    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
24644    /// K2/K3 can write them.
24645    pub fn gdn_chunk_alloc(
24646        &self,
24647        n_head: usize,
24648        t: usize,
24649        c: usize,
24650        hk: usize,
24651    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24652        const D: usize = 128;
24653        assert!(
24654            c == 32,
24655            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24656        );
24657        let h = n_head;
24658        let nc = (t + c - 1) / c;
24659        Ok(GdnChunkBufs {
24660            gcum: self.uninit(t * h)?,
24661            a: self.uninit(nc * h * c * c)?,
24662            p: self.uninit(nc * h * c * c)?,
24663            u: self.uninit(nc * h * c * D)?,
24664            w: self.uninit(nc * h * c * D)?,
24665            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24666            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24667            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24668            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24669            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24670            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24671            o: self.uninit(D * h * t)?,
24672            t,
24673            nc,
24674        })
24675    }
24676
24677    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
24678    pub fn f32_to_bf16_v(
24679        &self,
24680        x: &cudarc::driver::CudaView<f32>,
24681        dst: &mut CudaSlice<u8>,
24682        n: usize,
24683    ) -> Result<(), Box<dyn std::error::Error>> {
24684        let f = self.func("f32_to_bf16_bulk");
24685        let ni = n as i64;
24686        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24687        let __s_b = self.gpu.stream();
24688        let mut b = __s_b.launch_builder(&f);
24689        b.arg(x).arg(dst).arg(&ni);
24690        unsafe {
24691            b.launch(cfg)?;
24692        }
24693        Ok(())
24694    }
24695
24696    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
24697    pub fn f32_to_bf16_into(
24698        &self,
24699        x: &CudaSlice<f32>,
24700        dst: &mut CudaSlice<u8>,
24701        n: usize,
24702    ) -> Result<(), Box<dyn std::error::Error>> {
24703        let f = self.func("f32_to_bf16_bulk");
24704        let ni = n as i64;
24705        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24706        let __s_b = self.gpu.stream();
24707        let mut b = __s_b.launch_builder(&f);
24708        b.arg(x).arg(dst).arg(&ni);
24709        unsafe {
24710            b.launch(cfg)?;
24711        }
24712        Ok(())
24713    }
24714
24715    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
24716    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
24717    pub fn gdn_chunk_k123_vl8(
24718        &self,
24719        seqs: &[GdnSeqVl],
24720        n_head: usize,
24721        hk: usize,
24722        wq: Option<&GdnWVl8>,
24723    ) -> Result<(), Box<dyn std::error::Error>> {
24724        let b = seqs.len();
24725        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24726        let mut packed = [GdnSeqVl::default(); 8];
24727        packed[..b].copy_from_slice(seqs);
24728        let v = GdnVl8(packed);
24729        let (hi, ci) = (n_head as i32, 32i32);
24730        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24731        {
24732            let f = self.func("gdn_chunk_cumgate_vl");
24733            let cfg = LaunchConfig {
24734                grid_dim: (max_nc, n_head as u32, b as u32),
24735                block_dim: (32, 1, 1),
24736                shared_mem_bytes: 0,
24737            };
24738            let __s_lb = self.gpu.stream();
24739            let mut lb = __s_lb.launch_builder(&f);
24740            lb.arg(&v).arg(&hi).arg(&ci);
24741            unsafe {
24742                lb.launch(cfg)?;
24743            }
24744        }
24745        let hki = hk as i32;
24746        if let Some(w) = wq {
24747            // K2-wgmma vl twin (writes A + pre-masked Pb16)
24748            let f = self.func("gdn_k2_wgmma_vl");
24749            let cfg = LaunchConfig {
24750                grid_dim: (max_nc, n_head as u32, b as u32),
24751                block_dim: (128, 1, 1),
24752                shared_mem_bytes: 0,
24753            };
24754            let __s_lb = self.gpu.stream();
24755            let mut lb = __s_lb.launch_builder(&f);
24756            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24757            unsafe {
24758                lb.launch(cfg)?;
24759            }
24760        } else {
24761            let f = self.func("gdn_chunk_attn_vl");
24762            f.set_attribute(
24763                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24764                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24765            )?;
24766            let cfg = LaunchConfig {
24767                grid_dim: (max_nc, n_head as u32, b as u32),
24768                block_dim: (256, 1, 1),
24769                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24770            };
24771            let __s_lb = self.gpu.stream();
24772            let mut lb = __s_lb.launch_builder(&f);
24773            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24774            unsafe {
24775                lb.launch(cfg)?;
24776            }
24777        }
24778        {
24779            let f = self.func("gdn_chunk_solve32_vl");
24780            let cfg = LaunchConfig {
24781                grid_dim: (max_nc, n_head as u32, b as u32),
24782                block_dim: (256, 1, 1),
24783                shared_mem_bytes: 0,
24784            };
24785            let __s_lb = self.gpu.stream();
24786            let mut lb = __s_lb.launch_builder(&f);
24787            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24788            unsafe {
24789                lb.launch(cfg)?;
24790            }
24791        }
24792        Ok(())
24793    }
24794
24795    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
24796    /// fused gate-prep, 5 launches for every sequence (per-element math identical
24797    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
24798    #[allow(clippy::too_many_arguments)]
24799    pub fn gdn_prep_vl8(
24800        &self,
24801        seqs: &[GdnPrepVl],
24802        conv_w: &CudaSlice<f32>,
24803        dt_bias: &CudaSlice<f32>,
24804        a: &CudaSlice<f32>,
24805        conv_dim: usize,
24806        d_conv: usize,
24807        d_state: usize,
24808        num_v: usize,
24809        num_k: usize,
24810        key_dim: usize,
24811        hk: usize,
24812        eps: f32,
24813    ) -> Result<(), Box<dyn std::error::Error>> {
24814        let b = seqs.len();
24815        assert!(b >= 1 && b <= 8);
24816        let mut packed = [GdnPrepVl::default(); 8];
24817        packed[..b].copy_from_slice(seqs);
24818        let v = GdnPrepVl8(packed);
24819        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24820        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24821        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24822        assert!(
24823            conv_fuse || hk == num_v,
24824            "de-broadcast requires the fused conv"
24825        );
24826        if conv_fuse {
24827            let f = self.func("ssm_conv1d_gdn_state_vl");
24828            let cfg = LaunchConfig {
24829                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24830                block_dim: (256, 1, 1),
24831                shared_mem_bytes: 0,
24832            };
24833            let (dsi, nvi, nki, kdi, hki) = (
24834                d_state as i32,
24835                num_v as i32,
24836                num_k as i32,
24837                key_dim as i32,
24838                hk as i32,
24839            );
24840            let __s_lb = self.gpu.stream();
24841            let mut lb = __s_lb.launch_builder(&f);
24842            lb.arg(&v)
24843                .arg(conv_w)
24844                .arg(&cdi)
24845                .arg(&dci)
24846                .arg(&dsi)
24847                .arg(&nvi)
24848                .arg(&nki)
24849                .arg(&kdi)
24850                .arg(&hki);
24851            unsafe {
24852                lb.launch(cfg)?;
24853            }
24854        } else {
24855            let f = self.func("ssm_conv1d_tm_state_vl");
24856            let cfg = LaunchConfig {
24857                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24858                block_dim: (256, 1, 1),
24859                shared_mem_bytes: 0,
24860            };
24861            let __s_lb = self.gpu.stream();
24862            let mut lb = __s_lb.launch_builder(&f);
24863            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
24864            unsafe {
24865                lb.launch(cfg)?;
24866            }
24867        }
24868        {
24869            let f = self.func("ssm_conv_ring_update_vl");
24870            let n = (conv_dim * (d_conv - 1)) as u32;
24871            let cfg = LaunchConfig {
24872                grid_dim: (n.div_ceil(256), 1, b as u32),
24873                block_dim: (256, 1, 1),
24874                shared_mem_bytes: 0,
24875            };
24876            let __s_lb = self.gpu.stream();
24877            let mut lb = __s_lb.launch_builder(&f);
24878            lb.arg(&v).arg(&cdi).arg(&dci);
24879            unsafe {
24880                lb.launch(cfg)?;
24881            }
24882        }
24883        if !conv_fuse {
24884            let f = self.func("qkv_to_gdn_repack_vl");
24885            let n = max_t * (num_v * d_state) as u32;
24886            let cfg = LaunchConfig {
24887                grid_dim: (n.div_ceil(256), 1, b as u32),
24888                block_dim: (256, 1, 1),
24889                shared_mem_bytes: 0,
24890            };
24891            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24892            let __s_lb = self.gpu.stream();
24893            let mut lb = __s_lb.launch_builder(&f);
24894            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
24895            unsafe {
24896                lb.launch(cfg)?;
24897            }
24898        }
24899        if Self::l2_v2_on(d_state) {
24900            let f = self.func("gdn_l2_v2_vl");
24901            let cfg = LaunchConfig {
24902                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
24903                block_dim: (256, 1, 1),
24904                shared_mem_bytes: 0,
24905            };
24906            let (dsi, nvi) = (d_state as i32, hk as i32);
24907            let __s_lb = self.gpu.stream();
24908            let mut lb = __s_lb.launch_builder(&f);
24909            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24910            unsafe {
24911                lb.launch(cfg)?;
24912            }
24913        } else {
24914            let f = self.func("gdn_l2_vl");
24915            let cfg = LaunchConfig {
24916                grid_dim: (max_t * hk as u32, 2, b as u32),
24917                block_dim: (256, 1, 1),
24918                shared_mem_bytes: 0,
24919            };
24920            let (dsi, nvi) = (d_state as i32, hk as i32);
24921            let __s_lb = self.gpu.stream();
24922            let mut lb = __s_lb.launch_builder(&f);
24923            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24924            unsafe {
24925                lb.launch(cfg)?;
24926            }
24927        }
24928        {
24929            let f = self.func("gdn_gate_prep_vl");
24930            let n = max_t * num_v as u32;
24931            let cfg = LaunchConfig {
24932                grid_dim: (n.div_ceil(256), 1, b as u32),
24933                block_dim: (256, 1, 1),
24934                shared_mem_bytes: 0,
24935            };
24936            let nvi = num_v as i32;
24937            let __s_lb = self.gpu.stream();
24938            let mut lb = __s_lb.launch_builder(&f);
24939            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
24940            unsafe {
24941                lb.launch(cfg)?;
24942            }
24943        }
24944        Ok(())
24945    }
24946
24947    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
24948    pub fn gdn_mirror_vl8(
24949        &self,
24950        seqs: &[GdnSeqVl],
24951        n_head: usize,
24952        which: i32,
24953        hk: usize,
24954    ) -> Result<(), Box<dyn std::error::Error>> {
24955        let b = seqs.len();
24956        assert!(b >= 1 && b <= 8);
24957        let mut packed = [GdnSeqVl::default(); 8];
24958        packed[..b].copy_from_slice(seqs);
24959        let v = GdnVl8(packed);
24960        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
24961        let max_n = seqs
24962            .iter()
24963            .map(|s| {
24964                if which == 0 {
24965                    s.t as i64 * ept as i64
24966                } else {
24967                    s.nc as i64 * ept as i64 * 32
24968                }
24969            })
24970            .max()
24971            .unwrap();
24972        let f = self.func("gdn_mirror_vl");
24973        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
24974        let cfg = LaunchConfig {
24975            grid_dim: (blocks, 1, b as u32),
24976            block_dim: (256, 1, 1),
24977            shared_mem_bytes: 0,
24978        };
24979        let __s_lb = self.gpu.stream();
24980        let mut lb = __s_lb.launch_builder(&f);
24981        lb.arg(&v).arg(&ept).arg(&which);
24982        unsafe {
24983            lb.launch(cfg)?;
24984        }
24985        Ok(())
24986    }
24987
24988    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
24989    pub fn gdn_tail_vl8(
24990        &self,
24991        seqs: &[GdnPrepVl],
24992        norm_w: &CudaSlice<f32>,
24993        d_state: usize,
24994        num_v: usize,
24995        eps: f32,
24996    ) -> Result<(), Box<dyn std::error::Error>> {
24997        let b = seqs.len();
24998        assert!(b >= 1 && b <= 8);
24999        let mut packed = [GdnPrepVl::default(); 8];
25000        packed[..b].copy_from_slice(seqs);
25001        let v = GdnPrepVl8(packed);
25002        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25003        let f = self.func("gated_rmsnorm_f16out_vl");
25004        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25005        let cfg = LaunchConfig {
25006            grid_dim: (max_t * num_v as u32, 1, b as u32),
25007            block_dim: (128, 1, 1),
25008            shared_mem_bytes: 0,
25009        };
25010        let (dsi, nvi) = (d_state as i32, num_v as i32);
25011        let __s_lb = self.gpu.stream();
25012        let mut lb = __s_lb.launch_builder(&f);
25013        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25014        unsafe {
25015            lb.launch(cfg)?;
25016        }
25017        Ok(())
25018    }
25019
25020    /// Raw device address helpers for the varlen by-value arg struct (single-stream
25021    /// launches; every buffer outlives the call — the f16 FFI discipline).
25022    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25023        use cudarc::driver::DevicePtr;
25024        let s = self.gpu.stream();
25025        let (p, _g) = x.device_ptr(&s);
25026        p as u64
25027    }
25028    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25029        use cudarc::driver::DevicePtrMut;
25030        let s = self.gpu.stream();
25031        let (p, _g) = x.device_ptr_mut(&s);
25032        p as u64
25033    }
25034    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25035        use cudarc::driver::DevicePtr;
25036        let s = self.gpu.stream();
25037        let (p, _g) = x.device_ptr(&s);
25038        p as u64
25039    }
25040    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25041        use cudarc::driver::DevicePtr;
25042        let s = self.gpu.stream();
25043        let (p, _g) = x.device_ptr(&s);
25044        p as u64
25045    }
25046
25047    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
25048    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
25049    /// launches, so this is strictly bit-gateable against them).
25050    pub fn gdn_chunk_vl8(
25051        &self,
25052        seqs: &[GdnSeqVl],
25053        n_head: usize,
25054        scale: f32,
25055        hk: usize,
25056        wq: Option<&GdnWVl8>,
25057    ) -> Result<(), Box<dyn std::error::Error>> {
25058        const NSPLIT: u32 = 4;
25059        let b = seqs.len();
25060        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25061        let mut packed = [GdnSeqVl::default(); 8];
25062        packed[..b].copy_from_slice(seqs);
25063        let v = GdnVl8(packed);
25064        let (hi, ci) = (n_head as i32, 32i32);
25065        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25066        let hki = hk as i32;
25067        if let Some(w) = wq {
25068            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
25069            let f = self.func("gdn_k45_wgmma_vl");
25070            let cfg = LaunchConfig {
25071                grid_dim: (n_head as u32, NSPLIT, b as u32),
25072                block_dim: (256, 1, 1),
25073                shared_mem_bytes: 0,
25074            };
25075            let __s_lb = self.gpu.stream();
25076            let mut lb = __s_lb.launch_builder(&f);
25077            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25078            unsafe {
25079                lb.launch(cfg)?;
25080            }
25081            let _ = max_nc;
25082            return Ok(());
25083        }
25084        {
25085            let f = self.func("gdn_chunk_state_mma_vl");
25086            let cfg = LaunchConfig {
25087                grid_dim: (n_head as u32, NSPLIT, b as u32),
25088                block_dim: (256, 1, 1),
25089                shared_mem_bytes: 0,
25090            };
25091            let __s_lb = self.gpu.stream();
25092            let mut lb = __s_lb.launch_builder(&f);
25093            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25094            unsafe {
25095                lb.launch(cfg)?;
25096            }
25097        }
25098        {
25099            let f = self.func("gdn_chunk_output_mma_vl");
25100            let cfg = LaunchConfig {
25101                grid_dim: (max_nc, n_head as u32, b as u32),
25102                block_dim: (256, 1, 1),
25103                shared_mem_bytes: 0,
25104            };
25105            let __s_lb = self.gpu.stream();
25106            let mut lb = __s_lb.launch_builder(&f);
25107            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25108            unsafe {
25109                lb.launch(cfg)?;
25110            }
25111        }
25112        Ok(())
25113    }
25114    pub fn gdn_scan_chunked(
25115        &self,
25116        q: &CudaSlice<f32>,
25117        k: &CudaSlice<f32>,
25118        v: &CudaSlice<f32>,
25119        g: &CudaSlice<f32>,
25120        beta: &CudaSlice<f32>,
25121        kb16_pre: Option<&CudaSlice<u8>>,
25122        qb16_pre: Option<&CudaSlice<u8>>,
25123        state_in: &CudaSlice<f32>,
25124        state_out: &mut CudaSlice<f32>,
25125        o: &mut CudaSlice<f32>,
25126        n_head: usize,
25127        t: usize,
25128        scale: f32,
25129        c: usize,
25130        hk: usize,
25131    ) -> Result<(), Box<dyn std::error::Error>> {
25132        const D: usize = 128;
25133        const NSPLIT: u32 = 4;
25134        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25135        let h = n_head;
25136        let nc = (t + c - 1) / c;
25137        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25138        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
25139        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
25140        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
25141        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
25142        let gdn_mma_pre = !portable_mma_gated()
25143            && c == 32
25144            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25145                Ok("1") => true,
25146                Ok("0") => false,
25147                _ => gdn_mma_default_on(),
25148            };
25149        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
25150            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
25151        } else {
25152            None
25153        };
25154        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
25155        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
25156        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
25157        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
25158        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
25159            && gdn_mma_pre
25160            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
25161        let nk = t * hk * D;
25162        let mut kb16_local: Option<CudaSlice<u8>> = None;
25163        if gdn_mma_pre && kb16_pre.is_none() {
25164            let mut kb = self.alloc_u8_uninit(nk * 2)?;
25165            let f = self.func("f32_to_bf16_bulk");
25166            let n2 = nk as i64;
25167            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25168            let __s_b = self.gpu.stream();
25169            let mut b = __s_b.launch_builder(&f);
25170            b.arg(k).arg(&mut kb).arg(&n2);
25171            unsafe {
25172                b.launch(cfg2)?;
25173            }
25174            kb16_local = Some(kb);
25175        }
25176        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
25177        if let Some(kb) = kb16_pre {
25178            assert!(kb.len() >= nk * 2, "kb16_pre too small");
25179        }
25180        let mut qb16: Option<CudaSlice<u8>> = None;
25181        let mut pb16: Option<CudaSlice<u8>> = None;
25182        if gdn_wgmma_pre {
25183            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
25184            // the standalone bulk cvt only serves callers without the prep mirror.
25185            if qb16_pre.is_none() {
25186                let mut qb = self.alloc_u8_uninit(nk * 2)?;
25187                let f = self.func("f32_to_bf16_bulk");
25188                let n2 = nk as i64;
25189                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25190                let __s_b = self.gpu.stream();
25191                let mut b = __s_b.launch_builder(&f);
25192                b.arg(q).arg(&mut qb).arg(&n2);
25193                unsafe {
25194                    b.launch(cfg2)?;
25195                }
25196                qb16 = Some(qb);
25197            } else if let Some(qb) = qb16_pre {
25198                assert!(qb.len() >= nk * 2, "qb16_pre too small");
25199            }
25200            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
25201        }
25202        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
25203        let k2w = if gdn_wgmma_pre {
25204            Some((
25205                *qb16_ref0.as_ref().unwrap(),
25206                *kb16_ref0.as_ref().unwrap(),
25207                pb16.as_mut().unwrap(),
25208            ))
25209        } else {
25210            None
25211        };
25212        let (gcum, p, u, w) =
25213            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
25214        let _ = &w;
25215        let mut y = self.uninit(nc * h * c * D)?;
25216        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
25217        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
25218        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
25219        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
25220        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
25221        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
25222        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
25223        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
25224        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
25225        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
25226        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
25227        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
25228        // sites must agree or the pre-work arms while the scan takes the scalar route.
25229        let gdn_mma = !portable_mma_gated()
25230            && c == 32
25231            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25232                Ok("1") => true,
25233                Ok("0") => false,
25234                _ => gdn_mma_default_on(),
25235            };
25236        if gdn_mma {
25237            let wb16 = wb16_pre
25238                .take()
25239                .expect("mma path pre-allocates wb16 (K3 store fold)");
25240            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
25241            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
25242            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
25243            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
25244            // pass runs inside the persistent-M kernel; Y and Ssnap are never
25245            // materialized. New numeric class (gk folds into k^T instead of ys) —
25246            // explicit opt-in until the state-carry battery promotes it. Env read per
25247            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
25248            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
25249            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
25250            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
25251            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
25252            if gdn_wgmma_pre {
25253                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
25254                let qb16 = qb16_ref0.unwrap();
25255                let pb16 = pb16.as_ref().unwrap();
25256                {
25257                    let f = self.func("gdn_k45_wgmma");
25258                    let cfg = LaunchConfig {
25259                        grid_dim: (h as u32, 4, 1),
25260                        block_dim: (256, 1, 1),
25261                        shared_mem_bytes: 0,
25262                    };
25263                    let hki = hk as i32;
25264                    let __s_b = self.gpu.stream();
25265                    let mut b = __s_b.launch_builder(&f);
25266                    b.arg(kb16_ref)
25267                        .arg(&gcum)
25268                        .arg(beta)
25269                        .arg(&u)
25270                        .arg(&wb16)
25271                        .arg(qb16)
25272                        .arg(pb16)
25273                        .arg(o)
25274                        .arg(&scale)
25275                        .arg(state_in)
25276                        .arg(&mut *state_out)
25277                        .arg(&hi)
25278                        .arg(&ti)
25279                        .arg(&ci)
25280                        .arg(&hki);
25281                    unsafe {
25282                        b.launch(cfg)?;
25283                    }
25284                }
25285                return Ok(());
25286            }
25287            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
25288            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
25289            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
25290            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
25291            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
25292            {
25293                let f = self.func("gdn_chunk_state_mma");
25294                let cfg = LaunchConfig {
25295                    grid_dim: (h as u32, NSPLIT, 1),
25296                    block_dim: (256, 1, 1),
25297                    shared_mem_bytes: 0,
25298                };
25299                let hki = hk as i32;
25300                let __s_b = self.gpu.stream();
25301                let mut b = __s_b.launch_builder(&f);
25302                b.arg(kb16_ref)
25303                    .arg(&gcum)
25304                    .arg(beta)
25305                    .arg(&u)
25306                    .arg(&wb16)
25307                    .arg(&mut y16)
25308                    .arg(&mut ssnap16)
25309                    .arg(state_in)
25310                    .arg(&mut *state_out)
25311                    .arg(&hi)
25312                    .arg(&ti)
25313                    .arg(&ci)
25314                    .arg(&hki);
25315                unsafe {
25316                    b.launch(cfg)?;
25317                }
25318            }
25319            {
25320                // K5-mma (bf16 St/Y consumers)
25321                let f = self.func("gdn_chunk_output_mma");
25322                let jt = ((c + 31) / 32) as u32;
25323                let cfg = LaunchConfig {
25324                    grid_dim: (nc as u32, h as u32, jt),
25325                    block_dim: (256, 1, 1),
25326                    shared_mem_bytes: 0,
25327                };
25328                let hki = hk as i32;
25329                let __s_b = self.gpu.stream();
25330                let mut b = __s_b.launch_builder(&f);
25331                b.arg(q)
25332                    .arg(&gcum)
25333                    .arg(&p)
25334                    .arg(&y16)
25335                    .arg(&ssnap16)
25336                    .arg(o)
25337                    .arg(&hi)
25338                    .arg(&ti)
25339                    .arg(&ci)
25340                    .arg(&scale)
25341                    .arg(&hki);
25342                unsafe {
25343                    b.launch(cfg)?;
25344                }
25345            }
25346            return Ok(());
25347        }
25348        {
25349            // K4 (sequential over chunks inside; blocks col-partition the state)
25350            let f = self.func("gdn_chunk_state_f32");
25351            let cfg = LaunchConfig {
25352                grid_dim: (h as u32, NSPLIT, 1),
25353                block_dim: (256, 1, 1),
25354                shared_mem_bytes: 0,
25355            };
25356            let __s_b = self.gpu.stream();
25357            let mut b = __s_b.launch_builder(&f);
25358            b.arg(k)
25359                .arg(&gcum)
25360                .arg(beta)
25361                .arg(&u)
25362                .arg(&w)
25363                .arg(&mut y)
25364                .arg(&mut ssnap)
25365                .arg(state_in)
25366                .arg(&mut *state_out)
25367                .arg(&hi)
25368                .arg(&ti)
25369                .arg(&ci);
25370            unsafe {
25371                b.launch(cfg)?;
25372            }
25373        }
25374        {
25375            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
25376            let f = self.func("gdn_chunk_output_f32");
25377            let jt = ((c + 31) / 32) as u32;
25378            let cfg = LaunchConfig {
25379                grid_dim: (nc as u32, h as u32, jt),
25380                block_dim: (256, 1, 1),
25381                shared_mem_bytes: 0,
25382            };
25383            let __s_b = self.gpu.stream();
25384            let mut b = __s_b.launch_builder(&f);
25385            b.arg(q)
25386                .arg(&gcum)
25387                .arg(&p)
25388                .arg(&y)
25389                .arg(&ssnap)
25390                .arg(o)
25391                .arg(&hi)
25392                .arg(&ti)
25393                .arg(&ci)
25394                .arg(&scale);
25395            unsafe {
25396                b.launch(cfg)?;
25397            }
25398        }
25399        Ok(())
25400    }
25401
25402    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
25403    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
25404    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
25405    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
25406    ///
25407    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
25408    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
25409    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
25410    #[allow(clippy::too_many_arguments)]
25411    #[allow(clippy::too_many_arguments)]
25412    pub fn gdn_scan_prefill(
25413        &self,
25414        q: &CudaSlice<f32>,
25415        k: &CudaSlice<f32>,
25416        v: &CudaSlice<f32>,
25417        g: &CudaSlice<f32>,
25418        beta: &CudaSlice<f32>,
25419        kb16_pre: Option<&CudaSlice<u8>>,
25420        qb16_pre: Option<&CudaSlice<u8>>,
25421        state_in: &CudaSlice<f32>,
25422        state_out: &mut CudaSlice<f32>,
25423        o: &mut CudaSlice<f32>,
25424        n_head: usize,
25425        t: usize,
25426        scale: f32,
25427        hk: usize,
25428    ) -> Result<(), Box<dyn std::error::Error>> {
25429        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
25430            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
25431            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25432        }
25433        if Self::gdn_chunked_enabled() && t >= 16 {
25434            self.gdn_scan_chunked(
25435                q,
25436                k,
25437                v,
25438                g,
25439                beta,
25440                kb16_pre,
25441                qb16_pre,
25442                state_in,
25443                state_out,
25444                o,
25445                n_head,
25446                t,
25447                scale,
25448                Self::gdn_chunk_size(),
25449                hk,
25450            )
25451        } else {
25452            assert!(
25453                hk == n_head,
25454                "s128 scan is broadcast-only (prep guarantees by predicate)"
25455            );
25456            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25457        }
25458    }
25459
25460    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
25461    #[allow(clippy::too_many_arguments)]
25462    fn gdn_scan_diff(
25463        &self,
25464        q: &CudaSlice<f32>,
25465        k: &CudaSlice<f32>,
25466        v: &CudaSlice<f32>,
25467        g: &CudaSlice<f32>,
25468        beta: &CudaSlice<f32>,
25469        state_in: &CudaSlice<f32>,
25470        state_out: &mut CudaSlice<f32>,
25471        o: &mut CudaSlice<f32>,
25472        n_head: usize,
25473        t: usize,
25474        scale: f32,
25475    ) -> Result<(), Box<dyn std::error::Error>> {
25476        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25477        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25478        let mut o_c = self.uninit(o.len())?;
25479        let mut st_c = self.uninit(state_out.len())?;
25480        self.gdn_scan_chunked(
25481            q,
25482            k,
25483            v,
25484            g,
25485            beta,
25486            None,
25487            None,
25488            state_in,
25489            &mut st_c,
25490            &mut o_c,
25491            n_head,
25492            t,
25493            scale,
25494            Self::gdn_chunk_size(),
25495            n_head,
25496        )?;
25497        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25498        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25499        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25500        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25501            let mut max_abs = 0f32;
25502            let mut max_rel = 0f32;
25503            let mut sum_rel = 0f64;
25504            for (x, y) in a.iter().zip(b) {
25505                let ad = (x - y).abs();
25506                let rel = ad / x.abs().max(y.abs()).max(1e-3);
25507                if ad > max_abs {
25508                    max_abs = ad;
25509                }
25510                if rel > max_rel {
25511                    max_rel = rel;
25512                }
25513                sum_rel += rel as f64;
25514            }
25515            (max_abs, max_rel, sum_rel / a.len() as f64)
25516        };
25517        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25518        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25519        println!(
25520            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25521                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25522            Self::gdn_chunk_size()
25523        );
25524        Ok(())
25525    }
25526
25527    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
25528    pub fn gdn_glog(
25529        &self,
25530        alpha: &CudaSlice<f32>,
25531        dt_bias: &CudaSlice<f32>,
25532        a: &CudaSlice<f32>,
25533        g_log: &mut CudaSlice<f32>,
25534        n_head: usize,
25535        t: usize,
25536    ) -> Result<(), Box<dyn std::error::Error>> {
25537        let f = self.func("gdn_glog_f32");
25538        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25539        let (h, ti) = (n_head as i32, t as i32);
25540        let __s_b = self.gpu.stream();
25541        let mut b = __s_b.launch_builder(&f);
25542        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25543        unsafe {
25544            b.launch(cfg)?;
25545        }
25546        Ok(())
25547    }
25548
25549    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
25550    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
25551    pub fn sigmoid_v(
25552        &self,
25553        x: &cudarc::driver::CudaView<f32>,
25554        y: &mut CudaSlice<f32>,
25555        n: usize,
25556    ) -> Result<(), Box<dyn std::error::Error>> {
25557        let f = self.func("sigmoid_f32");
25558        let cfg = LaunchConfig::for_num_elems(n as u32);
25559        let ni = n as i32;
25560        let __s_b = self.gpu.stream();
25561        let mut b = __s_b.launch_builder(&f);
25562        b.arg(x).arg(y).arg(&ni);
25563        unsafe {
25564            b.launch(cfg)?;
25565        }
25566        Ok(())
25567    }
25568
25569    pub fn gdn_glog_v(
25570        &self,
25571        alpha: &cudarc::driver::CudaView<f32>,
25572        dt_bias: &CudaSlice<f32>,
25573        a: &CudaSlice<f32>,
25574        g_log: &mut CudaSlice<f32>,
25575        n_head: usize,
25576        t: usize,
25577    ) -> Result<(), Box<dyn std::error::Error>> {
25578        let f = self.func("gdn_glog_f32");
25579        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25580        let (h, ti) = (n_head as i32, t as i32);
25581        let __s_b = self.gpu.stream();
25582        let mut b = __s_b.launch_builder(&f);
25583        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25584        unsafe {
25585            b.launch(cfg)?;
25586        }
25587        Ok(())
25588    }
25589
25590    pub fn sigmoid(
25591        &self,
25592        x: &CudaSlice<f32>,
25593        y: &mut CudaSlice<f32>,
25594        n: usize,
25595    ) -> Result<(), Box<dyn std::error::Error>> {
25596        let f = self.func("sigmoid_f32");
25597        let cfg = LaunchConfig::for_num_elems(n as u32);
25598        let ni = n as i32;
25599        let __s_b = self.gpu.stream();
25600        let mut b = __s_b.launch_builder(&f);
25601        b.arg(x).arg(y).arg(&ni);
25602        unsafe {
25603            b.launch(cfg)?;
25604        }
25605        Ok(())
25606    }
25607
25608    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
25609    /// (replaces sigmoid + mul + convert). Bit-identical class.
25610    pub fn sig_mul_f16out(
25611        &self,
25612        a: &CudaSlice<f32>,
25613        g: &CudaSlice<f32>,
25614        dst: &mut CudaSlice<f32>,
25615        dst16: &mut CudaSlice<u8>,
25616        n: usize,
25617    ) -> Result<(), Box<dyn std::error::Error>> {
25618        let f = self.func("sig_mul_f16out_f32");
25619        let cfg = LaunchConfig::for_num_elems(n as u32);
25620        let ni = n as i32;
25621        let __s_b = self.gpu.stream();
25622        let mut b = __s_b.launch_builder(&f);
25623        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25624        unsafe {
25625            b.launch(cfg)?;
25626        }
25627        Ok(())
25628    }
25629
25630    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
25631    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
25632    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
25633    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
25634    ///
25635    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
25636    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
25637    /// applies the wrong number of distinct gate values.
25638    #[allow(clippy::too_many_arguments)]
25639    pub fn attn_head_gate(
25640        &self,
25641        a: &CudaSlice<f32>,
25642        g: &CudaSlice<f32>,
25643        dst: &mut CudaSlice<f32>,
25644        dst16: Option<&mut CudaSlice<u8>>,
25645        head_dim: usize,
25646        n_head: usize,
25647        t: usize,
25648    ) -> Result<(), Box<dyn std::error::Error>> {
25649        let f = self.func("attn_head_gate_f32");
25650        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25651        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25652        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
25653        let d16: u64 = match dst16 {
25654            Some(d) => self.addr_u8(d),
25655            None => 0,
25656        };
25657        let __s_b = self.gpu.stream();
25658        let mut b = __s_b.launch_builder(&f);
25659        b.arg(a)
25660            .arg(g)
25661            .arg(dst)
25662            .arg(&d16)
25663            .arg(&hd)
25664            .arg(&nh)
25665            .arg(&ti);
25666        unsafe {
25667            b.launch(cfg)?;
25668        }
25669        Ok(())
25670    }
25671
25672    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
25673    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
25674    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
25675    ///
25676    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
25677    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
25678    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
25679    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
25680    #[allow(clippy::too_many_arguments)]
25681    pub fn swiglu_clamped_mul_scaled(
25682        &self,
25683        gate: &CudaSlice<f32>,
25684        up: &CudaSlice<f32>,
25685        gs: f32,
25686        us: f32,
25687        limit: f32,
25688        dst: &mut CudaSlice<f32>,
25689        n: usize,
25690    ) -> Result<(), Box<dyn std::error::Error>> {
25691        debug_assert!(
25692            limit > 1e-6,
25693            "swiglu_clamped needs a live limit; use silu_mul_scaled"
25694        );
25695        let f = self.func("swiglu_clamped_mul_scaled_f32");
25696        let cfg = LaunchConfig::for_num_elems(n as u32);
25697        let ni = n as i32;
25698        let __s_b = self.gpu.stream();
25699        let mut b = __s_b.launch_builder(&f);
25700        b.arg(gate)
25701            .arg(up)
25702            .arg(&gs)
25703            .arg(&us)
25704            .arg(&limit)
25705            .arg(dst)
25706            .arg(&ni);
25707        unsafe {
25708            b.launch(cfg)?;
25709        }
25710        Ok(())
25711    }
25712
25713    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
25714    pub fn gated_rmsnorm(
25715        &self,
25716        o: &CudaSlice<f32>,
25717        w: &CudaSlice<f32>,
25718        z: &CudaSlice<f32>,
25719        dst: &mut CudaSlice<f32>,
25720        ncols: usize,
25721        nrows: usize,
25722        eps: f32,
25723    ) -> Result<(), Box<dyn std::error::Error>> {
25724        let f = self.func("gated_rmsnorm_f32");
25725        let cfg = LaunchConfig {
25726            grid_dim: (nrows as u32, 1, 1),
25727            block_dim: (128, 1, 1),
25728            shared_mem_bytes: 0,
25729        };
25730        let (nc, e) = (ncols as i32, eps);
25731        let __s_b = self.gpu.stream();
25732        let mut b = __s_b.launch_builder(&f);
25733        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25734        unsafe {
25735            b.launch(cfg)?;
25736        }
25737        Ok(())
25738    }
25739
25740    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
25741    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
25742    pub fn gated_rmsnorm_f16out(
25743        &self,
25744        o: &CudaSlice<f32>,
25745        w: &CudaSlice<f32>,
25746        z: &CudaSlice<f32>,
25747        dst: &mut CudaSlice<f32>,
25748        dst16: &mut CudaSlice<u8>,
25749        ncols: usize,
25750        nrows: usize,
25751        eps: f32,
25752    ) -> Result<(), Box<dyn std::error::Error>> {
25753        let f = self.func("gated_rmsnorm_f16out_f32");
25754        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25755        let cfg = LaunchConfig {
25756            grid_dim: (nrows as u32, 1, 1),
25757            block_dim: (128, 1, 1),
25758            shared_mem_bytes: 0,
25759        };
25760        let (nc, e) = (ncols as i32, eps);
25761        let __s_b = self.gpu.stream();
25762        let mut b = __s_b.launch_builder(&f);
25763        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25764        unsafe {
25765            b.launch(cfg)?;
25766        }
25767        Ok(())
25768    }
25769
25770    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
25771    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
25772    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
25773    #[allow(clippy::too_many_arguments)]
25774    pub fn add_rms_norm_zq8(
25775        &self,
25776        a: &CudaSlice<f32>,
25777        b_in: &CudaSlice<f32>,
25778        w: &CudaSlice<f32>,
25779        res: &mut CudaSlice<f32>,
25780        z: &mut CudaSlice<f32>,
25781        ncols: usize,
25782        nrows: usize,
25783        eps: f32,
25784    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25785        assert!(ncols % 32 == 0);
25786        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25787        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25788        let f = self.func("add_rms_norm_zq8");
25789        let cfg = LaunchConfig {
25790            grid_dim: (nrows as u32, 1, 1),
25791            block_dim: (1024, 1, 1),
25792            shared_mem_bytes: 0,
25793        };
25794        let (nc, ep) = (ncols as i32, eps);
25795        let __s_b = self.gpu.stream();
25796        let mut b = __s_b.launch_builder(&f);
25797        b.arg(a)
25798            .arg(b_in)
25799            .arg(w)
25800            .arg(res)
25801            .arg(z)
25802            .arg(&mut q)
25803            .arg(&mut d)
25804            .arg(&nc)
25805            .arg(&ep);
25806        unsafe {
25807            b.launch(cfg)?;
25808        }
25809        Ok((q, d))
25810    }
25811
25812    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
25813    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
25814    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
25815    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
25816    pub fn gated_rmsnorm_zv(
25817        &self,
25818        o: &CudaSlice<f32>,
25819        w: &CudaSlice<f32>,
25820        z: &cudarc::driver::CudaView<f32>,
25821        dst: &mut CudaSlice<f32>,
25822        ncols: usize,
25823        nrows: usize,
25824        eps: f32,
25825    ) -> Result<(), Box<dyn std::error::Error>> {
25826        let f = self.func("gated_rmsnorm_f32");
25827        let cfg = LaunchConfig {
25828            grid_dim: (nrows as u32, 1, 1),
25829            block_dim: (128, 1, 1),
25830            shared_mem_bytes: 0,
25831        };
25832        let (nc, e) = (ncols as i32, eps);
25833        let __s_b = self.gpu.stream();
25834        let mut b = __s_b.launch_builder(&f);
25835        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25836        unsafe {
25837            b.launch(cfg)?;
25838        }
25839        Ok(())
25840    }
25841
25842    pub fn gated_rmsnorm_f16out_zv(
25843        &self,
25844        o: &CudaSlice<f32>,
25845        w: &CudaSlice<f32>,
25846        z: &cudarc::driver::CudaView<f32>,
25847        dst: &mut CudaSlice<f32>,
25848        dst16: &mut CudaSlice<u8>,
25849        ncols: usize,
25850        nrows: usize,
25851        eps: f32,
25852    ) -> Result<(), Box<dyn std::error::Error>> {
25853        let f = self.func("gated_rmsnorm_f16out_f32");
25854        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25855        let cfg = LaunchConfig {
25856            grid_dim: (nrows as u32, 1, 1),
25857            block_dim: (128, 1, 1),
25858            shared_mem_bytes: 0,
25859        };
25860        let (nc, e) = (ncols as i32, eps);
25861        let __s_b = self.gpu.stream();
25862        let mut b = __s_b.launch_builder(&f);
25863        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25864        unsafe {
25865            b.launch(cfg)?;
25866        }
25867        Ok(())
25868    }
25869
25870    pub fn gated_rmsnorm_q8_1(
25871        &self,
25872        o: &CudaSlice<f32>,
25873        w: &CudaSlice<f32>,
25874        z: &CudaSlice<f32>,
25875        ncols: usize,
25876        nrows: usize,
25877        eps: f32,
25878    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25879        assert!(ncols % 32 == 0);
25880        let f = self.func("gated_rmsnorm_q8_1");
25881        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
25882        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25883        let cfg = LaunchConfig {
25884            grid_dim: (nrows as u32, 1, 1),
25885            block_dim: (128, 1, 1),
25886            shared_mem_bytes: 0,
25887        };
25888        let (nc, ep) = (ncols as i32, eps);
25889        let __s_b = self.gpu.stream();
25890        let mut b = __s_b.launch_builder(&f);
25891        b.arg(o)
25892            .arg(w)
25893            .arg(z)
25894            .arg(&mut out_q)
25895            .arg(&mut out_d)
25896            .arg(&nc)
25897            .arg(&ep);
25898        unsafe {
25899            b.launch(cfg)?;
25900        }
25901        Ok((out_q, out_d))
25902    }
25903
25904    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
25905    pub fn transpose(
25906        &self,
25907        inp: &CudaSlice<f32>,
25908        rows: usize,
25909        cols: usize,
25910    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25911        let f = self.func("transpose_f32");
25912        let mut out = self.zeros(rows * cols)?;
25913        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
25914        let (r, c) = (rows as i32, cols as i32);
25915        let __s_b = self.gpu.stream();
25916        let mut b = __s_b.launch_builder(&f);
25917        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
25918        unsafe {
25919            b.launch(cfg)?;
25920        }
25921        Ok(out)
25922    }
25923
25924    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
25925    pub fn repeat_heads(
25926        &self,
25927        inp: &CudaSlice<f32>,
25928        out: &mut CudaSlice<f32>,
25929        head_dim: usize,
25930        n_in: usize,
25931        n_out: usize,
25932        t: usize,
25933    ) -> Result<(), Box<dyn std::error::Error>> {
25934        let f = self.func("repeat_heads_f32");
25935        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
25936        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
25937        let __s_b = self.gpu.stream();
25938        let mut b = __s_b.launch_builder(&f);
25939        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
25940        unsafe {
25941            b.launch(cfg)?;
25942        }
25943        Ok(())
25944    }
25945
25946    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
25947    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
25948    ///
25949    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
25950    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
25951    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
25952    pub fn q_gate_split(
25953        &self,
25954        qf: &CudaSlice<f32>,
25955        q_out: &mut CudaSlice<f32>,
25956        gate_out: &mut CudaSlice<f32>,
25957        head_dim: usize,
25958        n_head: usize,
25959        t: usize,
25960    ) -> Result<(), Box<dyn std::error::Error>> {
25961        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
25962        let out_need = head_dim * n_head * t;
25963        if q_out.len() < out_need || gate_out.len() < out_need {
25964            return Err(format!(
25965                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
25966                q_out.len(),
25967                gate_out.len()
25968            )
25969            .into());
25970        }
25971        let f = self.func("q_gate_split_f32");
25972        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25973        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25974        let __s_b = self.gpu.stream();
25975        let mut b = __s_b.launch_builder(&f);
25976        b.arg(qf)
25977            .arg(q_out)
25978            .arg(gate_out)
25979            .arg(&hd)
25980            .arg(&nh)
25981            .arg(&ti);
25982        unsafe {
25983            b.launch(cfg)?;
25984        }
25985        Ok(())
25986    }
25987
25988    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
25989    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
25990    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
25991    pub fn qkv_to_gdn_repack(
25992        &self,
25993        conv_out: &CudaSlice<f32>,
25994        q_g: &mut CudaSlice<f32>,
25995        k_g: &mut CudaSlice<f32>,
25996        v_g: &mut CudaSlice<f32>,
25997        d_state: usize,
25998        num_v: usize,
25999        num_k: usize,
26000        key_dim: usize,
26001        t: usize,
26002    ) -> Result<(), Box<dyn std::error::Error>> {
26003        let f = self.func("qkv_to_gdn_repack_f32");
26004        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26005        let (ds, nv, nk, kd, ti) = (
26006            d_state as i32,
26007            num_v as i32,
26008            num_k as i32,
26009            key_dim as i32,
26010            t as i32,
26011        );
26012        let __s_b = self.gpu.stream();
26013        let mut b = __s_b.launch_builder(&f);
26014        b.arg(conv_out)
26015            .arg(q_g)
26016            .arg(k_g)
26017            .arg(v_g)
26018            .arg(&ds)
26019            .arg(&nv)
26020            .arg(&nk)
26021            .arg(&kd)
26022            .arg(&ti);
26023        unsafe {
26024            b.launch(cfg)?;
26025        }
26026        Ok(())
26027    }
26028
26029    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
26030    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
26031    pub fn conv_left_pad(
26032        &self,
26033        src: &CudaSlice<f32>,
26034        dst: &mut CudaSlice<f32>,
26035        conv_dim: usize,
26036        t: usize,
26037        pad: usize,
26038    ) -> Result<(), Box<dyn std::error::Error>> {
26039        let f = self.func("conv_left_pad_f32");
26040        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26041        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26042        let __s_b = self.gpu.stream();
26043        let mut b = __s_b.launch_builder(&f);
26044        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26045        unsafe {
26046            b.launch(cfg)?;
26047        }
26048        Ok(())
26049    }
26050
26051    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
26052    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
26053    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
26054    pub fn conv_assemble_and_roll(
26055        &self,
26056        qkv_col: &CudaSlice<f32>,
26057        conv_state: &mut CudaSlice<f32>,
26058        conv_in: &mut CudaSlice<f32>,
26059        conv_dim: usize,
26060        pad: usize,
26061    ) -> Result<(), Box<dyn std::error::Error>> {
26062        let f = self.func("conv_assemble_and_roll_f32");
26063        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26064        let (cd, p) = (conv_dim as i32, pad as i32);
26065        let __s_b = self.gpu.stream();
26066        let mut b = __s_b.launch_builder(&f);
26067        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26068        unsafe {
26069            b.launch(cfg)?;
26070        }
26071        Ok(())
26072    }
26073
26074    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
26075    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
26076    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
26077    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
26078    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
26079    pub fn ssm_conv1d_fused_decode(
26080        &self,
26081        qkv_col: &CudaSlice<f32>,
26082        conv_state: &mut CudaSlice<f32>,
26083        w: &CudaSlice<f32>,
26084        conv_out: &mut CudaSlice<f32>,
26085        conv_dim: usize,
26086        d_conv: usize,
26087    ) -> Result<(), Box<dyn std::error::Error>> {
26088        let f = self.func("ssm_conv1d_fused_decode_f32");
26089        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26090        let (cd, dc) = (conv_dim as i32, d_conv as i32);
26091        let __s_b = self.gpu.stream();
26092        let mut b = __s_b.launch_builder(&f);
26093        b.arg(qkv_col)
26094            .arg(conv_state)
26095            .arg(w)
26096            .arg(conv_out)
26097            .arg(&cd)
26098            .arg(&dc);
26099        unsafe {
26100            b.launch(cfg)?;
26101        }
26102        Ok(())
26103    }
26104
26105    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
26106    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
26107    pub fn slice_range(
26108        &self,
26109        src: &CudaSlice<f32>,
26110        start: usize,
26111        len: usize,
26112    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26113        let host = self.gpu.stream().clone_dtoh(src)?;
26114        self.gpu.stream().synchronize()?;
26115        Ok(self.htod(&host[start..start + len])?)
26116    }
26117}
26118
26119#[cfg(test)]
26120mod target_dispatch_tests {
26121    use super::legacy_quant_gemm_allowed;
26122
26123    #[test]
26124    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26125        // sm_120a native lane
26126        assert!(legacy_quant_gemm_allowed(false, false, false));
26127        assert!(!legacy_quant_gemm_allowed(false, false, true));
26128        // pure portable lane (sm_89): gated
26129        assert!(!legacy_quant_gemm_allowed(true, false, false));
26130        assert!(!legacy_quant_gemm_allowed(true, false, true));
26131        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
26132        assert!(legacy_quant_gemm_allowed(true, true, false));
26133        assert!(!legacy_quant_gemm_allowed(true, true, true));
26134    }
26135
26136    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26137    #[test]
26138    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26139        assert!(!legacy_quant_gemm_allowed(
26140            cfg!(memra_portable_cuda),
26141            cfg!(memra_hopper_mma),
26142            false
26143        ));
26144    }
26145
26146    #[cfg(memra_hopper_mma)]
26147    #[test]
26148    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
26149        assert!(legacy_quant_gemm_allowed(
26150            cfg!(memra_portable_cuda),
26151            cfg!(memra_hopper_mma),
26152            false
26153        ));
26154        assert!(super::portable_mma_gated() == false);
26155    }
26156}
26157
26158/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
26159/// inherent methods (inherent methods win name resolution, so no recursion).
26160impl memra_kv::KvDev for Engine {
26161    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26162        Engine::zeros(self, n)
26163    }
26164    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26165        Engine::uninit(self, n)
26166    }
26167    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
26168        Engine::alloc_u8(self, n)
26169    }
26170    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
26171        Engine::htod_i32(self, v)
26172    }
26173    fn clone_dtod(
26174        &self,
26175        src: &CudaSlice<f32>,
26176    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26177        Engine::clone_dtod(self, src)
26178    }
26179    fn copy_into(
26180        &self,
26181        dst: &mut CudaSlice<f32>,
26182        off: usize,
26183        src: &CudaSlice<f32>,
26184        len: usize,
26185    ) -> Result<(), Box<dyn std::error::Error>> {
26186        Engine::copy_into(self, dst, off, src, len)
26187    }
26188    fn set_i32_one(
26189        &self,
26190        d: &mut CudaSlice<i32>,
26191        v: i32,
26192    ) -> Result<(), Box<dyn std::error::Error>> {
26193        Engine::set_i32_one(self, d, v)
26194    }
26195}
26196
26197#[cfg(test)]
26198mod fused_gate_bounds_tests {
26199    use super::*;
26200
26201    /// The fused `[q|gate]` split's read-site guard, on the device.
26202    ///
26203    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
26204    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
26205    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
26206    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
26207    /// `FusedQGateExtent` before the launch.
26208    ///
26209    /// Catch demonstration for this test (guard temporarily removed, then restored):
26210    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
26211    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
26212    /// the call returns `Err`. Receipt in the lane report.
26213    #[test]
26214    #[ignore = "requires a CUDA GPU"]
26215    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
26216        let e = Engine::new(0).unwrap();
26217        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
26218        let fused = 2 * head_dim * n_head * t;
26219        let out_n = head_dim * n_head * t;
26220
26221        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
26222        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
26223        let mut q = e.uninit(out_n).unwrap();
26224        let mut gate = e.uninit(out_n).unwrap();
26225        let err = e
26226            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
26227            .expect_err("half-width wq must be refused, not read past")
26228            .to_string();
26229        assert!(err.contains("NO fused gate"), "{err}");
26230        assert!(err.contains(&format!("{fused}")), "{err}");
26231
26232        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
26233        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
26234        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
26235        let wide = e.htod(&host).unwrap();
26236        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
26237            .expect("full-width wq splits");
26238        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
26239        for tok in 0..t {
26240            for hh in 0..n_head {
26241                for d in 0..head_dim {
26242                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
26243                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
26244                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
26245                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
26246                }
26247            }
26248        }
26249
26250        // undersized destinations are refused too (the other half of the extent contract)
26251        let mut small = e.uninit(out_n - 1).unwrap();
26252        assert!(
26253            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
26254                .is_err()
26255        );
26256    }
26257}
26258
26259/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
26260/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
26261/// any launch, so the refusal is testable without a device.
26262#[cfg(test)]
26263mod fused_rope_width_tests {
26264    use super::Engine;
26265
26266    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
26267    /// safetensors route derives the same), which is why the fusion is legal there today.
26268    #[test]
26269    fn full_width_is_accepted() {
26270        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
26271        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
26272        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
26273    }
26274
26275    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
26276    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
26277    ///
26278    /// ```text
26279    /// attention.key_length     512   rope.dimension_count     512   (global class)
26280    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
26281    /// ```
26282    ///
26283    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
26284    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
26285    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
26286    /// instead of a silently over-rotated head.
26287    #[test]
26288    fn gemma4_official_artifact_widths_pass() {
26289        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
26290        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
26291    }
26292
26293    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
26294    /// with no `n_dims`, silently rotating the pass-through band.
26295    #[test]
26296    fn partial_rotary_is_refused_with_the_geometry_named() {
26297        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
26298        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
26299            .expect_err("partial rotary must refuse");
26300        let msg = err.to_string();
26301        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
26302        assert!(msg.contains("n_rot 64"), "{msg}");
26303        assert!(msg.contains("head_dim 256"), "{msg}");
26304        assert!(
26305            msg.contains("64..256"),
26306            "names the band it would corrupt: {msg}"
26307        );
26308        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
26309        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
26310        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
26311        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
26312    }
26313}