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    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
7754    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
7755    #[allow(clippy::too_many_arguments)]
7756    pub fn qmatvec_nvfp4_sel_gu_ep_into(
7757        &self,
7758        gate_bank: &CudaSlice<u8>,
7759        up_bank: &CudaSlice<u8>,
7760        sel: &CudaSlice<i32>,
7761        aq: &CudaSlice<i8>,
7762        ad: &CudaSlice<f32>,
7763        yg: &mut CudaSlice<f32>,
7764        yu: &mut CudaSlice<f32>,
7765        n_sel: usize,
7766        in_f: usize,
7767        out_f: usize,
7768        row_bytes: usize,
7769        expert_stride: usize,
7770        owner: usize,
7771    ) -> Result<(), Box<dyn std::error::Error>> {
7772        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7773        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7774            return Err("NVFP4 gu ep geometry".into());
7775        }
7776        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
7777        let cfg = LaunchConfig {
7778            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
7779            block_dim: (128, 1, 1),
7780            shared_mem_bytes: 0,
7781        };
7782        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
7783        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7784        let (ars, adrs) = (0i64, 0i64);
7785        let __s_b = self.gpu.stream();
7786        let mut b = __s_b.launch_builder(&f);
7787        b.arg(gate_bank)
7788            .arg(up_bank)
7789            .arg(sel)
7790            .arg(aq)
7791            .arg(ad)
7792            .arg(yg)
7793            .arg(yu)
7794            .arg(&inf)
7795            .arg(&outf)
7796            .arg(&ns)
7797            .arg(&rb)
7798            .arg(&es)
7799            .arg(&ars)
7800            .arg(&adrs)
7801            .arg(&own);
7802        unsafe {
7803            b.launch(cfg)?;
7804        }
7805        Ok(())
7806    }
7807
7808    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
7809    #[allow(clippy::too_many_arguments)]
7810    pub fn silu_mul_scaled_q8_1_sel_ep_into(
7811        &self,
7812        gate: &CudaSlice<f32>,
7813        up: &CudaSlice<f32>,
7814        gmac: &CudaSlice<f32>,
7815        umac: &CudaSlice<f32>,
7816        sel: &CudaSlice<i32>,
7817        limit: Option<f32>,
7818        out_q: &mut CudaSlice<i8>,
7819        out_d: &mut CudaSlice<f32>,
7820        n_per: usize,
7821        n_sel: usize,
7822        owner: usize,
7823    ) -> Result<(), Box<dyn std::error::Error>> {
7824        if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
7825            return Err("NVFP4 silu ep geometry".into());
7826        }
7827        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
7828        let warps = n_sel * n_per / 32;
7829        let cfg = LaunchConfig {
7830            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
7831            block_dim: (128, 1, 1),
7832            shared_mem_bytes: 0,
7833        };
7834        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
7835        let (lim, has) = match limit {
7836            Some(l) => (l, 1i32),
7837            None => (0.0f32, 0i32),
7838        };
7839        let __s_b = self.gpu.stream();
7840        let mut b = __s_b.launch_builder(&f);
7841        b.arg(gate)
7842            .arg(up)
7843            .arg(gmac)
7844            .arg(umac)
7845            .arg(sel)
7846            .arg(&lim)
7847            .arg(&has)
7848            .arg(out_q)
7849            .arg(out_d)
7850            .arg(&np)
7851            .arg(&ns)
7852            .arg(&own);
7853        unsafe {
7854            b.launch(cfg)?;
7855        }
7856        Ok(())
7857    }
7858
7859    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
7860    #[allow(clippy::too_many_arguments)]
7861    pub fn qmatvec_nvfp4_sel_down8_ep_into(
7862        &self,
7863        bank: &CudaSlice<u8>,
7864        sel: &CudaSlice<i32>,
7865        aq: &CudaSlice<i8>,
7866        ad: &CudaSlice<f32>,
7867        route_w: &CudaSlice<f32>,
7868        md: &CudaSlice<f32>,
7869        dst: &mut CudaSlice<f32>,
7870        n_sel: usize,
7871        in_f: usize,
7872        out_f: usize,
7873        row_bytes: usize,
7874        expert_stride: usize,
7875        act_row_stride: usize,
7876        ad_row_stride: usize,
7877        owner: usize,
7878    ) -> Result<(), Box<dyn std::error::Error>> {
7879        if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
7880            return Err("NVFP4 down8 ep geometry".into());
7881        }
7882        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
7883        let cfg = LaunchConfig {
7884            grid_dim: (out_f as u32, 1, 1),
7885            block_dim: (32, n_sel as u32, 1),
7886            shared_mem_bytes: 0,
7887        };
7888        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
7889        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7890        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7891        let __s_b = self.gpu.stream();
7892        let mut b = __s_b.launch_builder(&f);
7893        b.arg(bank)
7894            .arg(sel)
7895            .arg(aq)
7896            .arg(ad)
7897            .arg(route_w)
7898            .arg(md)
7899            .arg(dst)
7900            .arg(&inf)
7901            .arg(&outf)
7902            .arg(&ns)
7903            .arg(&rb)
7904            .arg(&es)
7905            .arg(&ars)
7906            .arg(&adrs)
7907            .arg(&own);
7908        unsafe {
7909            b.launch(cfg)?;
7910        }
7911        Ok(())
7912    }
7913
7914    pub fn qmatvec_nvfp4_sel_into(
7915        &self,
7916        bank: &CudaSlice<u8>,
7917        sel: &CudaSlice<i32>,
7918        aq: &CudaSlice<i8>,
7919        ad: &CudaSlice<f32>,
7920        y: &mut CudaSlice<f32>,
7921        n_sel: usize,
7922        in_f: usize,
7923        out_f: usize,
7924        row_bytes: usize,
7925        expert_stride: usize,
7926        act_row_stride: usize,
7927        ad_row_stride: usize,
7928    ) -> Result<(), Box<dyn std::error::Error>> {
7929        assert!(
7930            in_f % 64 == 0,
7931            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7932        );
7933        if y.len() < n_sel * out_f || sel.len() < n_sel {
7934            return Err(format!(
7935                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7936                y.len(),
7937                sel.len()
7938            )
7939            .into());
7940        }
7941        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
7942        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
7943        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
7944        // sequential-rows variant was flat). Default stays the single-row form.
7945        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
7946        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
7947        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7948        let mode = *MR.get_or_init(|| {
7949            if crate::tp::nvfp4_bank_v2_on() {
7950                3
7951            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7952                2
7953            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7954                1
7955            } else {
7956                0
7957            }
7958        });
7959        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7960        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
7961        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
7962        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
7963        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7964        let v2s = mode == 3
7965            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7966            && row_bytes % 16 == 0
7967            && in_f <= 4096;
7968        let f = match (mode, v2s) {
7969            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7970            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7971            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7972            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7973            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7974        };
7975        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
7976        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
7977        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
7978        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
7979        let nsb = in_f >> 5;
7980        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7981            32
7982        } else if mode == 1 {
7983            512
7984        } else {
7985            128
7986        };
7987        let cfg = LaunchConfig {
7988            grid_dim: (
7989                if v2s {
7990                    (out_f as u32).div_ceil(8)
7991                } else {
7992                    match mode {
7993                        2 => (out_f as u32).div_ceil(16),
7994                        1 => (out_f as u32).div_ceil(4),
7995                        _ => out_f as u32,
7996                    }
7997                },
7998                n_sel as u32,
7999                1,
8000            ),
8001            block_dim: (fit_block, 1, 1),
8002            shared_mem_bytes: 0,
8003        };
8004        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8005        let (rb, es, ars, adrs) = (
8006            row_bytes as i64,
8007            expert_stride as i64,
8008            act_row_stride as i64,
8009            ad_row_stride as i64,
8010        );
8011        let __s_b = self.gpu.stream();
8012        let mut b = __s_b.launch_builder(&f);
8013        b.arg(bank)
8014            .arg(sel)
8015            .arg(aq)
8016            .arg(ad)
8017            .arg(y)
8018            .arg(&inf)
8019            .arg(&outf)
8020            .arg(&ns)
8021            .arg(&rb)
8022            .arg(&es)
8023            .arg(&ars)
8024            .arg(&adrs);
8025        unsafe {
8026            b.launch(cfg)?;
8027        }
8028        Ok(())
8029    }
8030
8031    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
8032    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
8033    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
8034    /// takes the plain SiLU kernel.
8035    #[allow(clippy::too_many_arguments)]
8036    pub fn silu_mul_scaled_q8_1_sel_into(
8037        &self,
8038        gate: &CudaSlice<f32>,
8039        up: &CudaSlice<f32>,
8040        gmac: &CudaSlice<f32>,
8041        umac: &CudaSlice<f32>,
8042        sel: &CudaSlice<i32>,
8043        limit: Option<f32>,
8044        out_q: &mut CudaSlice<i8>,
8045        out_d: &mut CudaSlice<f32>,
8046        n_per: usize,
8047        n_sel: usize,
8048    ) -> Result<(), Box<dyn std::error::Error>> {
8049        let n = n_per * n_sel;
8050        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8051            return Err(format!(
8052                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8053                out_q.len(),
8054                out_d.len()
8055            )
8056            .into());
8057        }
8058        if let Some(limit) = limit {
8059            if limit <= 1e-6 {
8060                return Err(format!(
8061                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8062                )
8063                .into());
8064            }
8065            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8066            let cfg = LaunchConfig::for_num_elems(n as u32);
8067            let (np, ns) = (n_per as i32, n_sel as i32);
8068            let __s_b = self.gpu.stream();
8069            let mut b = __s_b.launch_builder(&f);
8070            b.arg(gate)
8071                .arg(up)
8072                .arg(gmac)
8073                .arg(umac)
8074                .arg(sel)
8075                .arg(&limit)
8076                .arg(out_q)
8077                .arg(out_d)
8078                .arg(&np)
8079                .arg(&ns);
8080            unsafe {
8081                b.launch(cfg)?;
8082            }
8083            return Ok(());
8084        }
8085        let f = self.func("silu_mul_scaled_q8_1_sel");
8086        let cfg = LaunchConfig::for_num_elems(n as u32);
8087        let (np, ns) = (n_per as i32, n_sel as i32);
8088        let __s_b = self.gpu.stream();
8089        let mut b = __s_b.launch_builder(&f);
8090        b.arg(gate)
8091            .arg(up)
8092            .arg(gmac)
8093            .arg(umac)
8094            .arg(sel)
8095            .arg(out_q)
8096            .arg(out_d)
8097            .arg(&np)
8098            .arg(&ns);
8099        unsafe {
8100            b.launch(cfg)?;
8101        }
8102        Ok(())
8103    }
8104
8105    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8106        Ok(self.gpu.stream().clone_htod(v)?)
8107    }
8108    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8109        Ok(self.gpu.stream().clone_htod(v)?)
8110    }
8111    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
8112    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8113        Ok(self.gpu.stream().clone_htod(v)?)
8114    }
8115    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8116        Ok(self.gpu.stream().clone_htod(v)?)
8117    }
8118    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
8119    pub fn dtoh_view(
8120        &self,
8121        d: &cudarc::driver::CudaView<f32>,
8122    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8123        let v = self.gpu.stream().clone_dtoh(d)?;
8124        self.gpu.stream().synchronize()?;
8125        Ok(v)
8126    }
8127    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8128        let v = self.gpu.stream().clone_dtoh(d)?;
8129        self.gpu.stream().synchronize()?;
8130        Ok(v)
8131    }
8132    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
8133    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
8134    /// issuing them together avoids a second stream synchronization in every trunk layer.
8135    pub fn dtoh_pair(
8136        &self,
8137        a: &CudaSlice<f32>,
8138        b: &CudaSlice<f32>,
8139    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8140        let av = self.gpu.stream().clone_dtoh(a)?;
8141        let bv = self.gpu.stream().clone_dtoh(b)?;
8142        self.gpu.stream().synchronize()?;
8143        Ok((av, bv))
8144    }
8145    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
8146    /// cross a shape-sensitive host boundary.
8147    pub fn dtoh_pair_views(
8148        &self,
8149        a: &cudarc::driver::CudaView<f32>,
8150        b: &cudarc::driver::CudaView<f32>,
8151    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8152        let av = self.gpu.stream().clone_dtoh(a)?;
8153        let bv = self.gpu.stream().clone_dtoh(b)?;
8154        self.gpu.stream().synchronize()?;
8155        Ok((av, bv))
8156    }
8157    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
8158    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8159        let v = self.gpu.stream().clone_dtoh(d)?;
8160        self.gpu.stream().synchronize()?;
8161        Ok(v)
8162    }
8163    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
8164    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8165        let v = self.gpu.stream().clone_dtoh(d)?;
8166        self.gpu.stream().synchronize()?;
8167        Ok(v)
8168    }
8169    pub fn dtoh_u8_view(
8170        &self,
8171        d: &cudarc::driver::CudaView<u8>,
8172    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8173        let v = self.gpu.stream().clone_dtoh(d)?;
8174        self.gpu.stream().synchronize()?;
8175        Ok(v)
8176    }
8177    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8178        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8179        self.keep_if_capturing(&s);
8180        Ok(s)
8181    }
8182
8183    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
8184    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
8185    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
8186    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
8187    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
8188    /// back (or kept resident for graph replay). Returns the device token buffer.
8189    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
8190    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
8191    pub fn prob_of_token_device(
8192        &self,
8193        logits: &CudaSlice<f32>,
8194        tok: &CudaSlice<u32>,
8195        n_vocab: usize,
8196    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8197        let nb = ARGMAX_NB;
8198        let mut part = self.alloc_uninit::<f32>(nb)?;
8199        let mut p = self.alloc_uninit::<f32>(1)?;
8200        let f1 = self.func("prob_of_token_partial_f32");
8201        let cfg1 = LaunchConfig {
8202            grid_dim: (nb as u32, 1, 1),
8203            block_dim: (256, 1, 1),
8204            shared_mem_bytes: 0,
8205        };
8206        let nv = n_vocab as i32;
8207        let __s_b1 = self.gpu.stream();
8208        let mut b1 = __s_b1.launch_builder(&f1);
8209        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8210        unsafe {
8211            b1.launch(cfg1)?;
8212        }
8213        let f2 = self.func("prob_of_token_final_f32");
8214        let cfg2 = LaunchConfig {
8215            grid_dim: (1, 1, 1),
8216            block_dim: (256, 1, 1),
8217            shared_mem_bytes: 0,
8218        };
8219        let nbi = nb as i32;
8220        let __s_b2 = self.gpu.stream();
8221        let mut b2 = __s_b2.launch_builder(&f2);
8222        b2.arg(&part).arg(&mut p).arg(&nbi);
8223        unsafe {
8224            b2.launch(cfg2)?;
8225        }
8226        Ok(p)
8227    }
8228
8229    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
8230    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
8231    /// where the host reads the p-min confidence between replays. Same kernels, same math.
8232    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
8233    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
8234    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
8235    pub fn prob_of_token_device_col(
8236        &self,
8237        logits: &CudaSlice<f32>,
8238        tok_all: &CudaSlice<u32>,
8239        tok_idx: usize,
8240        p_out: &mut CudaSlice<f32>,
8241        p_idx: usize,
8242        n_vocab: usize,
8243    ) -> Result<(), Box<dyn std::error::Error>> {
8244        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8245        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8246        let nb = ARGMAX_NB;
8247        let mut part = self.alloc_uninit::<f32>(nb)?;
8248        let f1 = self.func("prob_of_token_partial_f32");
8249        let cfg1 = LaunchConfig {
8250            grid_dim: (nb as u32, 1, 1),
8251            block_dim: (256, 1, 1),
8252            shared_mem_bytes: 0,
8253        };
8254        let nv = n_vocab as i32;
8255        let __s_b1 = self.gpu.stream();
8256        let mut b1 = __s_b1.launch_builder(&f1);
8257        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8258        unsafe {
8259            b1.launch(cfg1)?;
8260        }
8261        let f2 = self.func("prob_of_token_final_f32");
8262        let cfg2 = LaunchConfig {
8263            grid_dim: (1, 1, 1),
8264            block_dim: (256, 1, 1),
8265            shared_mem_bytes: 0,
8266        };
8267        let nbi = nb as i32;
8268        let __s_b2 = self.gpu.stream();
8269        let mut b2 = __s_b2.launch_builder(&f2);
8270        b2.arg(&part).arg(&mut p_v).arg(&nbi);
8271        unsafe {
8272            b2.launch(cfg2)?;
8273        }
8274        Ok(())
8275    }
8276
8277    pub fn prob_of_token_device_into(
8278        &self,
8279        logits: &CudaSlice<f32>,
8280        tok: &CudaSlice<u32>,
8281        p_out: &mut CudaSlice<f32>,
8282        n_vocab: usize,
8283    ) -> Result<(), Box<dyn std::error::Error>> {
8284        let nb = ARGMAX_NB;
8285        let mut part = self.alloc_uninit::<f32>(nb)?;
8286        let f1 = self.func("prob_of_token_partial_f32");
8287        let cfg1 = LaunchConfig {
8288            grid_dim: (nb as u32, 1, 1),
8289            block_dim: (256, 1, 1),
8290            shared_mem_bytes: 0,
8291        };
8292        let nv = n_vocab as i32;
8293        let __s_b1 = self.gpu.stream();
8294        let mut b1 = __s_b1.launch_builder(&f1);
8295        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8296        unsafe {
8297            b1.launch(cfg1)?;
8298        }
8299        let f2 = self.func("prob_of_token_final_f32");
8300        let cfg2 = LaunchConfig {
8301            grid_dim: (1, 1, 1),
8302            block_dim: (256, 1, 1),
8303            shared_mem_bytes: 0,
8304        };
8305        let nbi = nb as i32;
8306        let __s_b2 = self.gpu.stream();
8307        let mut b2 = __s_b2.launch_builder(&f2);
8308        b2.arg(&part).arg(p_out).arg(&nbi);
8309        unsafe {
8310            b2.launch(cfg2)?;
8311        }
8312        Ok(())
8313    }
8314
8315    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
8316    /// (graph-constant params, device-varying index). Capture-safe.
8317    pub fn u32_hist_append(
8318        &self,
8319        tok: &CudaSlice<u32>,
8320        hist: &mut CudaSlice<u32>,
8321        idx: &mut CudaSlice<i32>,
8322    ) -> Result<(), Box<dyn std::error::Error>> {
8323        let f = self.func("u32_hist_append");
8324        let cfg = LaunchConfig {
8325            grid_dim: (1, 1, 1),
8326            block_dim: (32, 1, 1),
8327            shared_mem_bytes: 0,
8328        };
8329        let __s_b = self.gpu.stream();
8330        let mut b = __s_b.launch_builder(&f);
8331        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8332        unsafe {
8333            b.launch(cfg)?;
8334        }
8335        Ok(())
8336    }
8337
8338    pub fn argmax_token_device(
8339        &self,
8340        logits: &CudaSlice<f32>,
8341        n_vocab: usize,
8342    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8343        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8344        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8345        Ok(tok)
8346    }
8347    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
8348    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
8349    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
8350    /// pointer is baked once and the token id never round-trips to host inside steady state. The
8351    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
8352    /// captured passes bake fixed addresses.
8353    pub fn argmax_token_device_into(
8354        &self,
8355        logits: &CudaSlice<f32>,
8356        tok: &mut CudaSlice<u32>,
8357        n_vocab: usize,
8358    ) -> Result<(), Box<dyn std::error::Error>> {
8359        let nb = ARGMAX_NB;
8360        let f1 = self.func("argmax_partial_f32");
8361        let f2 = self.func("argmax_final_f32");
8362        let mut guard = self.argmax_partials.lock().unwrap();
8363        if guard.is_none() {
8364            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8365            // buffers carry no cudarc events (illegal inside capture).
8366            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8367            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8368            *guard = Some((pv, pi));
8369        }
8370        let (part_v, part_i) = guard.as_mut().unwrap();
8371        let nv = n_vocab as i32;
8372        let nbi = nb as i32;
8373        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8374        let cfg1 = LaunchConfig {
8375            grid_dim: (nb as u32, 1, 1),
8376            block_dim: (256, 1, 1),
8377            shared_mem_bytes: 0,
8378        };
8379        let __s_b1 = self.gpu.stream();
8380        let mut b1 = __s_b1.launch_builder(&f1);
8381        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8382        unsafe {
8383            b1.launch(cfg1)?;
8384        }
8385        // pass 2: one block reduces NB partials -> token_out[0].
8386        let cfg2 = LaunchConfig {
8387            grid_dim: (1, 1, 1),
8388            block_dim: (256, 1, 1),
8389            shared_mem_bytes: 0,
8390        };
8391        let __s_b2 = self.gpu.stream();
8392        let mut b2 = __s_b2.launch_builder(&f2);
8393        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8394        unsafe {
8395            b2.launch(cfg2)?;
8396        }
8397        Ok(())
8398    }
8399    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8400    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8401    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8402    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8403    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8404    pub fn argmax_token_device_col(
8405        &self,
8406        logits: &CudaSlice<f32>,
8407        col: usize,
8408        n_vocab: usize,
8409        toks: &mut CudaSlice<u32>,
8410        out_idx: usize,
8411    ) -> Result<(), Box<dyn std::error::Error>> {
8412        let nb = ARGMAX_NB;
8413        let f1 = self.func("argmax_partial_f32");
8414        let f2 = self.func("argmax_final_f32");
8415        let mut guard = self.argmax_partials.lock().unwrap();
8416        if guard.is_none() {
8417            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8418            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8419            *guard = Some((pv, pi));
8420        }
8421        let (part_v, part_i) = guard.as_mut().unwrap();
8422        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8423        let nv = n_vocab as i32;
8424        let nbi = nb as i32;
8425        let cfg1 = LaunchConfig {
8426            grid_dim: (nb as u32, 1, 1),
8427            block_dim: (256, 1, 1),
8428            shared_mem_bytes: 0,
8429        };
8430        let __s_b1 = self.gpu.stream();
8431        let mut b1 = __s_b1.launch_builder(&f1);
8432        b1.arg(&col_view)
8433            .arg(&mut *part_v)
8434            .arg(&mut *part_i)
8435            .arg(&nv);
8436        unsafe {
8437            b1.launch(cfg1)?;
8438        }
8439        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8440        let cfg2 = LaunchConfig {
8441            grid_dim: (1, 1, 1),
8442            block_dim: (256, 1, 1),
8443            shared_mem_bytes: 0,
8444        };
8445        let __s_b2 = self.gpu.stream();
8446        let mut b2 = __s_b2.launch_builder(&f2);
8447        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8448        unsafe {
8449            b2.launch(cfg2)?;
8450        }
8451        Ok(())
8452    }
8453    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8454    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8455        Ok(self.gpu.stream().clone_htod(v)?)
8456    }
8457    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8458        let v = self.gpu.stream().clone_dtoh(d)?;
8459        self.gpu.stream().synchronize()?;
8460        Ok(v)
8461    }
8462
8463    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8464        let v = self.gpu.stream().clone_dtoh(d)?;
8465        self.gpu.stream().synchronize()?;
8466        Ok(v)
8467    }
8468    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8469    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8470    /// contents change every step, the address must not, so a captured graph can read it).
8471    pub fn htod_u32_into(
8472        &self,
8473        dst: &mut CudaSlice<u32>,
8474        src: &[u32],
8475    ) -> Result<(), Box<dyn std::error::Error>> {
8476        let mut view = dst.slice_mut(0..src.len());
8477        self.gpu.stream().memcpy_htod(src, &mut view)?;
8478        Ok(())
8479    }
8480
8481    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8482    /// table without changing the device address its reconcile kernel consumes.
8483    pub fn htod_i32_into(
8484        &self,
8485        dst: &mut CudaSlice<i32>,
8486        src: &[i32],
8487    ) -> Result<(), Box<dyn std::error::Error>> {
8488        let mut view = dst.slice_mut(0..src.len());
8489        self.gpu.stream().memcpy_htod(src, &mut view)?;
8490        Ok(())
8491    }
8492
8493    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8494        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8495        self.keep_if_capturing(&s);
8496        Ok(s)
8497    }
8498    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8499    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8500    pub fn embed_gather_device_into(
8501        &self,
8502        embd: &CudaSlice<u8>,
8503        token_d: &CudaSlice<u32>,
8504        x_out: &mut CudaSlice<f32>,
8505        n_embd: usize,
8506        qtype: i32,
8507        row_bytes: usize,
8508    ) -> Result<(), Box<dyn std::error::Error>> {
8509        let f = self.func("embed_gather_u32");
8510        let cfg = LaunchConfig {
8511            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8512            block_dim: (256, 1, 1),
8513            shared_mem_bytes: 0,
8514        };
8515        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8516        let __s_b = self.gpu.stream();
8517        let mut b = __s_b.launch_builder(&f);
8518        b.arg(embd)
8519            .arg(token_d)
8520            .arg(x_out)
8521            .arg(&ne)
8522            .arg(&qt)
8523            .arg(&rb);
8524        unsafe {
8525            b.launch(cfg)?;
8526        }
8527        Ok(())
8528    }
8529    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
8530    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8531        let v = self.gpu.stream().clone_dtoh(d)?;
8532        self.gpu.stream().synchronize()?;
8533        Ok(v[0])
8534    }
8535    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
8536    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
8537    /// the counter value after the throwaway capture warmups corrupt it.
8538    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
8539    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
8540    /// copy (fine at stream-idle boundaries, poison mid-round).
8541    pub fn i32_set_k(
8542        &self,
8543        dst: &mut CudaSlice<i32>,
8544        v: i32,
8545    ) -> Result<(), Box<dyn std::error::Error>> {
8546        let f = self.func("i32_set_k");
8547        let cfg = LaunchConfig {
8548            grid_dim: (1, 1, 1),
8549            block_dim: (1, 1, 1),
8550            shared_mem_bytes: 0,
8551        };
8552        let idx = 0i32;
8553        let __s_b = self.gpu.stream();
8554        let mut b = __s_b.launch_builder(&f);
8555        b.arg(dst).arg(&v).arg(&idx);
8556        unsafe {
8557            b.launch(cfg)?;
8558        }
8559        Ok(())
8560    }
8561
8562    pub fn set_i32_one(
8563        &self,
8564        d: &mut CudaSlice<i32>,
8565        v: i32,
8566    ) -> Result<(), Box<dyn std::error::Error>> {
8567        self.gpu.stream().memcpy_htod(&[v], d)?;
8568        Ok(())
8569    }
8570    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
8571    /// during priming / capture-state restore.
8572    pub fn set_u32_one(
8573        &self,
8574        d: &mut CudaSlice<u32>,
8575        v: u32,
8576    ) -> Result<(), Box<dyn std::error::Error>> {
8577        self.gpu.stream().memcpy_htod(&[v], d)?;
8578        Ok(())
8579    }
8580    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
8581    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8582        let v = self.gpu.stream().clone_dtoh(d)?;
8583        self.gpu.stream().synchronize()?;
8584        Ok(v[0])
8585    }
8586    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
8587    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8588        Ok(self.gpu.stream().clone_htod(bytes)?)
8589    }
8590    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
8591    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
8592    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
8593    pub fn embed_gather_device(
8594        &self,
8595        embd: &CudaSlice<u8>,
8596        token_d: &CudaSlice<u32>,
8597        n_embd: usize,
8598        qtype: i32,
8599        row_bytes: usize,
8600    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8601        let f = self.func("embed_gather_u32");
8602        let mut x = self.alloc_uninit::<f32>(n_embd)?;
8603        let cfg = LaunchConfig {
8604            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8605            block_dim: (256, 1, 1),
8606            shared_mem_bytes: 0,
8607        };
8608        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8609        let __s_b = self.gpu.stream();
8610        let mut b = __s_b.launch_builder(&f);
8611        b.arg(embd)
8612            .arg(token_d)
8613            .arg(&mut x)
8614            .arg(&ne)
8615            .arg(&qt)
8616            .arg(&rb);
8617        unsafe {
8618            b.launch(cfg)?;
8619        }
8620        Ok(x)
8621    }
8622
8623    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
8624    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
8625    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
8626    pub fn embed_gather_device_t(
8627        &self,
8628        embd: &CudaSlice<u8>,
8629        tokens: &[u32],
8630        n_embd: usize,
8631        qtype: i32,
8632        row_bytes: usize,
8633    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8634        let t = tokens.len();
8635        let tok_d = self.gpu.stream().clone_htod(tokens)?;
8636        let f = self.func("embed_gather_u32_t");
8637        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8638        let cfg = LaunchConfig {
8639            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8640            block_dim: (256, 1, 1),
8641            shared_mem_bytes: 0,
8642        };
8643        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8644        let __s_b = self.gpu.stream();
8645        let mut b = __s_b.launch_builder(&f);
8646        b.arg(embd)
8647            .arg(&tok_d)
8648            .arg(&mut x)
8649            .arg(&ne)
8650            .arg(&qt)
8651            .arg(&rb)
8652            .arg(&ti);
8653        unsafe {
8654            b.launch(cfg)?;
8655        }
8656        Ok(x)
8657    }
8658
8659    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
8660    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
8661    /// as embed_gather_device_t — bit-identical rows.
8662    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
8663    pub fn embed_gather_device_tv(
8664        &self,
8665        embd: &CudaSlice<u8>,
8666        tok_v: &cudarc::driver::CudaView<u32>,
8667        t: usize,
8668        n_embd: usize,
8669        qtype: i32,
8670        row_bytes: usize,
8671    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8672        let f = self.func("embed_gather_u32_t");
8673        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8674        let cfg = LaunchConfig {
8675            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8676            block_dim: (256, 1, 1),
8677            shared_mem_bytes: 0,
8678        };
8679        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8680        let __s_b = self.gpu.stream();
8681        let mut b = __s_b.launch_builder(&f);
8682        b.arg(embd)
8683            .arg(tok_v)
8684            .arg(&mut x)
8685            .arg(&ne)
8686            .arg(&qt)
8687            .arg(&rb)
8688            .arg(&ti);
8689        unsafe {
8690            b.launch(cfg)?;
8691        }
8692        Ok(x)
8693    }
8694
8695    pub fn embed_gather_device_td(
8696        &self,
8697        embd: &CudaSlice<u8>,
8698        tok_d: &CudaSlice<u32>,
8699        t: usize,
8700        n_embd: usize,
8701        qtype: i32,
8702        row_bytes: usize,
8703    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8704        let f = self.func("embed_gather_u32_t");
8705        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8706        let cfg = LaunchConfig {
8707            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8708            block_dim: (256, 1, 1),
8709            shared_mem_bytes: 0,
8710        };
8711        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8712        let __s_b = self.gpu.stream();
8713        let mut b = __s_b.launch_builder(&f);
8714        b.arg(embd)
8715            .arg(tok_d)
8716            .arg(&mut x)
8717            .arg(&ne)
8718            .arg(&qt)
8719            .arg(&rb)
8720            .arg(&ti);
8721        unsafe {
8722            b.launch(cfg)?;
8723        }
8724        Ok(x)
8725    }
8726
8727    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
8728    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
8729    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
8730    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
8731    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
8732    #[inline]
8733    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
8734    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8735        if self
8736            .capture_keep_on
8737            .load(std::sync::atomic::Ordering::Relaxed)
8738        {
8739            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8740        }
8741    }
8742
8743    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8744        &self,
8745        n: usize,
8746    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8747        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8748        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
8749        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
8750        // not cover engine-internal buffers). Debug-only: massive launch overhead.
8751        {
8752            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8753            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8754                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
8755                use cudarc::driver::DevicePtrMut;
8756                let n_bytes = s.len() * std::mem::size_of::<T>();
8757                let stream = self.gpu.stream();
8758                let (p_, _g) = s.device_ptr_mut(&stream);
8759                unsafe {
8760                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8761                        .result()?;
8762                }
8763            }
8764        }
8765        self.keep_if_capturing(&s);
8766        Ok(s)
8767    }
8768
8769    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
8770    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
8771    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
8772    /// consumers alloc through this (m=1 decode arms).
8773    pub fn uninit_q8_pair(
8774        &self,
8775        n: usize,
8776    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8777        Ok((
8778            self.alloc_uninit::<i8>(n)?,
8779            self.alloc_uninit::<f32>(n / 32)?,
8780        ))
8781    }
8782
8783    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8784        self.alloc_uninit::<f32>(n)
8785    }
8786
8787    /// i8 uninitialized scratch (same contract as `uninit`).
8788    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8789        self.alloc_uninit::<i8>(n)
8790    }
8791
8792    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
8793    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
8794    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
8795    #[allow(clippy::too_many_arguments)]
8796    pub fn rms_norm3(
8797        &self,
8798        x: &CudaSlice<f32>,
8799        w0: &CudaSlice<f32>,
8800        w1: &CudaSlice<f32>,
8801        w2: &CudaSlice<f32>,
8802        d0: &mut CudaSlice<f32>,
8803        d1: &mut CudaSlice<f32>,
8804        d2: &mut CudaSlice<f32>,
8805        ncols: usize,
8806        nrows: usize,
8807        eps: f32,
8808    ) -> Result<(), Box<dyn std::error::Error>> {
8809        let f = self.func("rms_norm3_f32");
8810        let cfg = LaunchConfig {
8811            grid_dim: (nrows as u32, 1, 1),
8812            block_dim: (rms_block(), 1, 1),
8813            shared_mem_bytes: 0,
8814        };
8815        let (nc, e) = (ncols as i32, eps);
8816        let __s_b = self.gpu.stream();
8817        let mut b = __s_b.launch_builder(&f);
8818        b.arg(x)
8819            .arg(w0)
8820            .arg(w1)
8821            .arg(w2)
8822            .arg(d0)
8823            .arg(d1)
8824            .arg(d2)
8825            .arg(&nc)
8826            .arg(&e);
8827        unsafe {
8828            b.launch(cfg)?;
8829        }
8830        Ok(())
8831    }
8832
8833    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
8834    #[allow(clippy::too_many_arguments)]
8835    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
8836    /// piggybacks on the same conditions.
8837    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8838        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8839        *WARP_ON.get_or_init(|| {
8840            std::env::var("MEMRA_QKVNORM_W")
8841                .map(|v| v != "0")
8842                .unwrap_or(true)
8843        }) && ncols % 4 == 0
8844            && rows >= 64
8845    }
8846
8847    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
8848    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
8849    #[allow(clippy::too_many_arguments)]
8850    pub fn rms_norm_qkv_w4b(
8851        &self,
8852        q: &CudaSlice<f32>,
8853        k: &CudaSlice<f32>,
8854        v: &CudaSlice<f32>,
8855        wq: &CudaSlice<f32>,
8856        wk: &CudaSlice<f32>,
8857        wv: &CudaSlice<f32>,
8858        dq: &mut CudaSlice<f32>,
8859        dk: &mut CudaSlice<f32>,
8860        dv: &mut CudaSlice<f32>,
8861        dvb: &mut CudaSlice<u8>,
8862        ncols: usize,
8863        rq: usize,
8864        rk: usize,
8865        eps: f32,
8866        vf16: bool,
8867    ) -> Result<(), Box<dyn std::error::Error>> {
8868        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8869        let f = self.func("rms_norm_qkv_w4b_f32");
8870        let rows = (rq + 2 * rk) as u32;
8871        let cfg = LaunchConfig {
8872            grid_dim: (rows.div_ceil(8), 1, 1),
8873            block_dim: (256, 1, 1),
8874            shared_mem_bytes: 0,
8875        };
8876        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8877        let vf = vf16 as i32;
8878        let __s_b = self.gpu.stream();
8879        let mut b = __s_b.launch_builder(&f);
8880        b.arg(q)
8881            .arg(k)
8882            .arg(v)
8883            .arg(wq)
8884            .arg(wk)
8885            .arg(wv)
8886            .arg(dq)
8887            .arg(dk)
8888            .arg(dv)
8889            .arg(&mut *dvb)
8890            .arg(&nc)
8891            .arg(&rqi)
8892            .arg(&rki)
8893            .arg(&rvi)
8894            .arg(&e)
8895            .arg(&vf);
8896        unsafe {
8897            b.launch(cfg)?;
8898        }
8899        Ok(())
8900    }
8901
8902    pub fn rms_norm_qkv(
8903        &self,
8904        q: &CudaSlice<f32>,
8905        k: &CudaSlice<f32>,
8906        v: &CudaSlice<f32>,
8907        wq: &CudaSlice<f32>,
8908        wk: &CudaSlice<f32>,
8909        wv: &CudaSlice<f32>,
8910        dq: &mut CudaSlice<f32>,
8911        dk: &mut CudaSlice<f32>,
8912        dv: &mut CudaSlice<f32>,
8913        ncols: usize,
8914        rq: usize,
8915        rk: usize,
8916        eps: f32,
8917    ) -> Result<(), Box<dyn std::error::Error>> {
8918        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
8919        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
8920        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
8921        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8922        let warp_on = *WARP_ON.get_or_init(|| {
8923            std::env::var("MEMRA_QKVNORM_W")
8924                .map(|v| v != "0")
8925                .unwrap_or(true)
8926        });
8927        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
8928        // replay numerics are untouched on every model; only prefill depth takes the new config.
8929        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8930            let f = self.func("rms_norm_qkv_w4_f32");
8931            let rows = (rq + 2 * rk) as u32;
8932            let cfg = LaunchConfig {
8933                grid_dim: (rows.div_ceil(8), 1, 1),
8934                block_dim: (256, 1, 1),
8935                shared_mem_bytes: 0,
8936            };
8937            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8938            let __s_b = self.gpu.stream();
8939            let mut b = __s_b.launch_builder(&f);
8940            b.arg(q)
8941                .arg(k)
8942                .arg(v)
8943                .arg(wq)
8944                .arg(wk)
8945                .arg(wv)
8946                .arg(dq)
8947                .arg(dk)
8948                .arg(dv)
8949                .arg(&nc)
8950                .arg(&rqi)
8951                .arg(&rki)
8952                .arg(&rvi)
8953                .arg(&e);
8954            unsafe {
8955                b.launch(cfg)?;
8956            }
8957            return Ok(());
8958        }
8959        let f = self.func("rms_norm_qkv_f32");
8960        let grid = (rq + 2 * rk) as u32;
8961        let cfg = LaunchConfig {
8962            grid_dim: (grid, 1, 1),
8963            block_dim: (rms_block(), 1, 1),
8964            shared_mem_bytes: 0,
8965        };
8966        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8967        let __s_b = self.gpu.stream();
8968        let mut b = __s_b.launch_builder(&f);
8969        b.arg(q)
8970            .arg(k)
8971            .arg(v)
8972            .arg(wq)
8973            .arg(wk)
8974            .arg(wv)
8975            .arg(dq)
8976            .arg(dk)
8977            .arg(dv)
8978            .arg(&nc)
8979            .arg(&rqi)
8980            .arg(&rki)
8981            .arg(&e);
8982        unsafe {
8983            b.launch(cfg)?;
8984        }
8985        Ok(())
8986    }
8987
8988    /// gemma4 fused pair of rms_norms over two different inputs (same width).
8989    #[allow(clippy::too_many_arguments)]
8990    pub fn rms_norm2x(
8991        &self,
8992        a: &CudaSlice<f32>,
8993        bb: &CudaSlice<f32>,
8994        wa: &CudaSlice<f32>,
8995        wb: &CudaSlice<f32>,
8996        da: &mut CudaSlice<f32>,
8997        db: &mut CudaSlice<f32>,
8998        ncols: usize,
8999        nrows: usize,
9000        eps: f32,
9001    ) -> Result<(), Box<dyn std::error::Error>> {
9002        let f = self.func("rms_norm2x_f32");
9003        let cfg = LaunchConfig {
9004            grid_dim: (2 * nrows as u32, 1, 1),
9005            block_dim: (rms_block(), 1, 1),
9006            shared_mem_bytes: 0,
9007        };
9008        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9009        let __s_b = self.gpu.stream();
9010        let mut b = __s_b.launch_builder(&f);
9011        b.arg(a)
9012            .arg(bb)
9013            .arg(wa)
9014            .arg(wb)
9015            .arg(da)
9016            .arg(db)
9017            .arg(&nc)
9018            .arg(&nr)
9019            .arg(&e);
9020        unsafe {
9021            b.launch(cfg)?;
9022        }
9023        Ok(())
9024    }
9025
9026    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
9027    pub fn softcap(
9028        &self,
9029        y: &mut CudaSlice<f32>,
9030        cap: f32,
9031        n: usize,
9032    ) -> Result<(), Box<dyn std::error::Error>> {
9033        let f = self.func("softcap_f32");
9034        let cfg = LaunchConfig::for_num_elems(n as u32);
9035        let ni = n as i32;
9036        let __s_b = self.gpu.stream();
9037        let mut b = __s_b.launch_builder(&f);
9038        b.arg(y).arg(&cap).arg(&ni);
9039        unsafe {
9040            b.launch(cfg)?;
9041        }
9042        Ok(())
9043    }
9044
9045    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
9046    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
9047    pub fn mask_ids_rows(
9048        &self,
9049        y: &mut CudaSlice<f32>,
9050        ids: &CudaSlice<i32>,
9051        n_ids: usize,
9052        n_vocab: usize,
9053        t: usize,
9054    ) -> Result<(), Box<dyn std::error::Error>> {
9055        let f = self.func("mask_ids_rows_f32");
9056        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9057        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9058        let __s_b = self.gpu.stream();
9059        let mut b = __s_b.launch_builder(&f);
9060        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9061        unsafe {
9062            b.launch(cfg)?;
9063        }
9064        Ok(())
9065    }
9066
9067    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
9068    #[allow(clippy::too_many_arguments)]
9069    pub fn add_scale_rms_norm(
9070        &self,
9071        a: &CudaSlice<f32>,
9072        b_in: &CudaSlice<f32>,
9073        c: f32,
9074        w: &CudaSlice<f32>,
9075        res: &mut CudaSlice<f32>,
9076        dst: &mut CudaSlice<f32>,
9077        ncols: usize,
9078        nrows: usize,
9079        eps: f32,
9080    ) -> Result<(), Box<dyn std::error::Error>> {
9081        let f = self.func("add_scale_rms_norm_f32");
9082        let cfg = LaunchConfig {
9083            grid_dim: (nrows as u32, 1, 1),
9084            block_dim: (rms_block(), 1, 1),
9085            shared_mem_bytes: 0,
9086        };
9087        let (nc, e2) = (ncols as i32, eps);
9088        let __s_b = self.gpu.stream();
9089        let mut b = __s_b.launch_builder(&f);
9090        b.arg(a)
9091            .arg(b_in)
9092            .arg(&c)
9093            .arg(w)
9094            .arg(res)
9095            .arg(dst)
9096            .arg(&nc)
9097            .arg(&e2);
9098        unsafe {
9099            b.launch(cfg)?;
9100        }
9101        Ok(())
9102    }
9103
9104    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
9105    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
9106    #[allow(clippy::too_many_arguments)]
9107    pub fn add_scale_rms_norm_q8_1(
9108        &self,
9109        a: &CudaSlice<f32>,
9110        b_in: &CudaSlice<f32>,
9111        c: f32,
9112        w: &CudaSlice<f32>,
9113        res: &mut CudaSlice<f32>,
9114        ncols: usize,
9115        nrows: usize,
9116        eps: f32,
9117    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9118        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9119        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9120        let (nc, e2) = (ncols as i32, eps);
9121        if Self::pdl_on() && Self::pdl_wb_on() {
9122            {
9123                use cudarc::driver::{DevicePtr, DevicePtrMut};
9124                let s = &self.gpu.stream();
9125                let (pa, _g0) = a.device_ptr(s);
9126                let (pb, _g1) = b_in.device_ptr(s);
9127                let (pw, _g2) = w.device_ptr(s);
9128                let (pr, _g3) = res.device_ptr_mut(s);
9129                let (pq, _g4) = out_q.device_ptr_mut(s);
9130                let (pd, _g5) = out_d.device_ptr_mut(s);
9131                let mut ps = [
9132                    &pa as *const _ as *mut std::ffi::c_void,
9133                    &pb as *const _ as *mut _,
9134                    &c as *const _ as *mut _,
9135                    &pw as *const _ as *mut _,
9136                    &pr as *const _ as *mut _,
9137                    &pq as *const _ as *mut _,
9138                    &pd as *const _ as *mut _,
9139                    &nc as *const _ as *mut _,
9140                    &e2 as *const _ as *mut _,
9141                ];
9142                unsafe {
9143                    self.launch_pdl(
9144                        "add_scale_rms_norm_q8_1",
9145                        (nrows as u32, 1, 1),
9146                        (rms_block(), 1, 1),
9147                        &mut ps,
9148                    )?;
9149                }
9150            }
9151            return Ok((out_q, out_d));
9152        }
9153        let f = self.func("add_scale_rms_norm_q8_1");
9154        let cfg = LaunchConfig {
9155            grid_dim: (nrows as u32, 1, 1),
9156            block_dim: (rms_block(), 1, 1),
9157            shared_mem_bytes: 0,
9158        };
9159        let __s_b = self.gpu.stream();
9160        let mut b = __s_b.launch_builder(&f);
9161        b.arg(a)
9162            .arg(b_in)
9163            .arg(&c)
9164            .arg(w)
9165            .arg(res)
9166            .arg(&mut out_q)
9167            .arg(&mut out_d)
9168            .arg(&nc)
9169            .arg(&e2);
9170        unsafe {
9171            b.launch(cfg)?;
9172        }
9173        Ok((out_q, out_d))
9174    }
9175
9176    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
9177    #[allow(clippy::too_many_arguments)]
9178    pub fn add_scale_rms_norm_q8_1_into(
9179        &self,
9180        a: &CudaSlice<f32>,
9181        b_in: &CudaSlice<f32>,
9182        c: f32,
9183        w: &CudaSlice<f32>,
9184        res: &mut CudaSlice<f32>,
9185        ncols: usize,
9186        nrows: usize,
9187        eps: f32,
9188        out_q: &mut CudaSlice<i8>,
9189        out_d: &mut CudaSlice<f32>,
9190    ) -> Result<(), Box<dyn std::error::Error>> {
9191        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9192        let (nc, e2) = (ncols as i32, eps);
9193        if Self::pdl_on() && Self::pdl_wb_on() {
9194            use cudarc::driver::{DevicePtr, DevicePtrMut};
9195            let s = &self.gpu.stream();
9196            let (pa, _g0) = a.device_ptr(s);
9197            let (pb, _g1) = b_in.device_ptr(s);
9198            let (pw, _g2) = w.device_ptr(s);
9199            let (pr, _g3) = res.device_ptr_mut(s);
9200            let (pq, _g4) = out_q.device_ptr_mut(s);
9201            let (pd, _g5) = out_d.device_ptr_mut(s);
9202            let mut ps = [
9203                &pa as *const _ as *mut std::ffi::c_void,
9204                &pb as *const _ as *mut _,
9205                &c as *const _ as *mut _,
9206                &pw as *const _ as *mut _,
9207                &pr as *const _ as *mut _,
9208                &pq as *const _ as *mut _,
9209                &pd as *const _ as *mut _,
9210                &nc as *const _ as *mut _,
9211                &e2 as *const _ as *mut _,
9212            ];
9213            unsafe {
9214                self.launch_pdl(
9215                    "add_scale_rms_norm_q8_1",
9216                    (nrows as u32, 1, 1),
9217                    (rms_block(), 1, 1),
9218                    &mut ps,
9219                )?;
9220            }
9221            return Ok(());
9222        }
9223        let f = self.func("add_scale_rms_norm_q8_1");
9224        let cfg = LaunchConfig {
9225            grid_dim: (nrows as u32, 1, 1),
9226            block_dim: (rms_block(), 1, 1),
9227            shared_mem_bytes: 0,
9228        };
9229        let __s_b = self.gpu.stream();
9230        let mut b = __s_b.launch_builder(&f);
9231        b.arg(a)
9232            .arg(b_in)
9233            .arg(&c)
9234            .arg(w)
9235            .arg(res)
9236            .arg(&mut *out_q)
9237            .arg(&mut *out_d)
9238            .arg(&nc)
9239            .arg(&e2);
9240        unsafe {
9241            b.launch(cfg)?;
9242        }
9243        Ok(())
9244    }
9245
9246    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
9247    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
9248    #[allow(clippy::too_many_arguments)]
9249    pub fn rms_pre_add_scale_rms_norm_q8_1(
9250        &self,
9251        a: &CudaSlice<f32>,
9252        wa: &CudaSlice<f32>,
9253        b_in: &CudaSlice<f32>,
9254        c: f32,
9255        w: &CudaSlice<f32>,
9256        res: &mut CudaSlice<f32>,
9257        ncols: usize,
9258        nrows: usize,
9259        eps: f32,
9260    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9261        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9262        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9263        let (nc, e2) = (ncols as i32, eps);
9264        if Self::pdl_on() {
9265            {
9266                use cudarc::driver::{DevicePtr, DevicePtrMut};
9267                let s = &self.gpu.stream();
9268                let (pa, _g0) = a.device_ptr(s);
9269                let (pwa, _g1) = wa.device_ptr(s);
9270                let (pb, _g2) = b_in.device_ptr(s);
9271                let (pw, _g3) = w.device_ptr(s);
9272                let (pr, _g4) = res.device_ptr_mut(s);
9273                let (pq, _g5) = out_q.device_ptr_mut(s);
9274                let (pd, _g6) = out_d.device_ptr_mut(s);
9275                let mut ps = [
9276                    &pa as *const _ as *mut std::ffi::c_void,
9277                    &pwa as *const _ as *mut _,
9278                    &pb as *const _ as *mut _,
9279                    &c as *const _ as *mut _,
9280                    &pw as *const _ as *mut _,
9281                    &pr as *const _ as *mut _,
9282                    &pq as *const _ as *mut _,
9283                    &pd as *const _ as *mut _,
9284                    &nc as *const _ as *mut _,
9285                    &e2 as *const _ as *mut _,
9286                ];
9287                unsafe {
9288                    self.launch_pdl(
9289                        "rms_pre_add_scale_rms_norm_q8_1",
9290                        (nrows as u32, 1, 1),
9291                        (rms_block(), 1, 1),
9292                        &mut ps,
9293                    )?;
9294                }
9295            }
9296            return Ok((out_q, out_d));
9297        }
9298        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9299        let cfg = LaunchConfig {
9300            grid_dim: (nrows as u32, 1, 1),
9301            block_dim: (rms_block(), 1, 1),
9302            shared_mem_bytes: 0,
9303        };
9304        let __s_b = self.gpu.stream();
9305        let mut b = __s_b.launch_builder(&f);
9306        b.arg(a)
9307            .arg(wa)
9308            .arg(b_in)
9309            .arg(&c)
9310            .arg(w)
9311            .arg(res)
9312            .arg(&mut out_q)
9313            .arg(&mut out_d)
9314            .arg(&nc)
9315            .arg(&e2);
9316        unsafe {
9317            b.launch(cfg)?;
9318        }
9319        Ok((out_q, out_d))
9320    }
9321
9322    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
9323    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
9324    pub fn gelu_tanh_mul_q8_1(
9325        &self,
9326        gate: &CudaSlice<f32>,
9327        up: &cudarc::driver::CudaView<f32>,
9328        act: &mut CudaSlice<f32>,
9329        ncols: usize,
9330        nrows: usize,
9331    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9332        debug_assert!(ncols % 128 == 0);
9333        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9334        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9335        let nc = ncols as i32;
9336        if Self::pdl_on() {
9337            {
9338                use cudarc::driver::{DevicePtr, DevicePtrMut};
9339                let s = &self.gpu.stream();
9340                let (pg, _g0) = gate.device_ptr(s);
9341                let (pu, _g1) = up.device_ptr(s);
9342                let (pact, _g2) = act.device_ptr_mut(s);
9343                let (pq, _g3) = out_q.device_ptr_mut(s);
9344                let (pd, _g4) = out_d.device_ptr_mut(s);
9345                let mut ps = [
9346                    &pg as *const _ as *mut std::ffi::c_void,
9347                    &pu as *const _ as *mut _,
9348                    &pact as *const _ as *mut _,
9349                    &pq as *const _ as *mut _,
9350                    &pd as *const _ as *mut _,
9351                    &nc as *const _ as *mut _,
9352                ];
9353                unsafe {
9354                    self.launch_pdl(
9355                        "gelu_tanh_mul_q8_1",
9356                        (nrows as u32, 1, 1),
9357                        (rms_block(), 1, 1),
9358                        &mut ps,
9359                    )?;
9360                }
9361            }
9362            return Ok((out_q, out_d));
9363        }
9364        let f = self.func("gelu_tanh_mul_q8_1");
9365        let cfg = LaunchConfig {
9366            grid_dim: (nrows as u32, 1, 1),
9367            block_dim: (rms_block(), 1, 1),
9368            shared_mem_bytes: 0,
9369        };
9370        let __s_b = self.gpu.stream();
9371        let mut b = __s_b.launch_builder(&f);
9372        b.arg(gate)
9373            .arg(up)
9374            .arg(act)
9375            .arg(&mut out_q)
9376            .arg(&mut out_d)
9377            .arg(&nc);
9378        unsafe {
9379            b.launch(cfg)?;
9380        }
9381        Ok((out_q, out_d))
9382    }
9383
9384    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9385    #[allow(clippy::too_many_arguments)]
9386    pub fn gelu_tanh_mul_q8_1_into(
9387        &self,
9388        gate: &CudaSlice<f32>,
9389        up: &cudarc::driver::CudaView<f32>,
9390        act: &mut CudaSlice<f32>,
9391        ncols: usize,
9392        nrows: usize,
9393        out_q: &mut CudaSlice<i8>,
9394        out_d: &mut CudaSlice<f32>,
9395    ) -> Result<(), Box<dyn std::error::Error>> {
9396        debug_assert!(ncols % 128 == 0);
9397        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9398        let nc = ncols as i32;
9399        if Self::pdl_on() {
9400            use cudarc::driver::{DevicePtr, DevicePtrMut};
9401            let s = &self.gpu.stream();
9402            let (pg, _g0) = gate.device_ptr(s);
9403            let (pu, _g1) = up.device_ptr(s);
9404            let (pact, _g2) = act.device_ptr_mut(s);
9405            let (pq, _g3) = out_q.device_ptr_mut(s);
9406            let (pd, _g4) = out_d.device_ptr_mut(s);
9407            let mut ps = [
9408                &pg as *const _ as *mut std::ffi::c_void,
9409                &pu as *const _ as *mut _,
9410                &pact as *const _ as *mut _,
9411                &pq as *const _ as *mut _,
9412                &pd as *const _ as *mut _,
9413                &nc as *const _ as *mut _,
9414            ];
9415            unsafe {
9416                self.launch_pdl(
9417                    "gelu_tanh_mul_q8_1",
9418                    (nrows as u32, 1, 1),
9419                    (rms_block(), 1, 1),
9420                    &mut ps,
9421                )?;
9422            }
9423            return Ok(());
9424        }
9425        let f = self.func("gelu_tanh_mul_q8_1");
9426        let cfg = LaunchConfig {
9427            grid_dim: (nrows as u32, 1, 1),
9428            block_dim: (rms_block(), 1, 1),
9429            shared_mem_bytes: 0,
9430        };
9431        let __s_b = self.gpu.stream();
9432        let mut b = __s_b.launch_builder(&f);
9433        b.arg(gate)
9434            .arg(up)
9435            .arg(&mut *act)
9436            .arg(&mut *out_q)
9437            .arg(&mut *out_d)
9438            .arg(&nc);
9439        unsafe {
9440            b.launch(cfg)?;
9441        }
9442        Ok(())
9443    }
9444
9445    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9446    #[allow(clippy::too_many_arguments)]
9447    pub fn add_rms_norm3_q8z(
9448        &self,
9449        a: &CudaSlice<f32>,
9450        b_in: &CudaSlice<f32>,
9451        w0: &CudaSlice<f32>,
9452        w1: &CudaSlice<f32>,
9453        w2: &CudaSlice<f32>,
9454        res: &mut CudaSlice<f32>,
9455        out1: &mut CudaSlice<f32>,
9456        ncols: usize,
9457        nrows: usize,
9458        eps: f32,
9459    ) -> Result<
9460        (
9461            (CudaSlice<i8>, CudaSlice<f32>),
9462            (CudaSlice<i8>, CudaSlice<f32>),
9463        ),
9464        Box<dyn std::error::Error>,
9465    > {
9466        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9467        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9468        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9469        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9470        let f = self.func("add_rms_norm3_q8z_f32");
9471        let cfg = LaunchConfig {
9472            grid_dim: (nrows as u32, 1, 1),
9473            block_dim: (rms_block(), 1, 1),
9474            shared_mem_bytes: 0,
9475        };
9476        let (nc, e2) = (ncols as i32, eps);
9477        let __s_b = self.gpu.stream();
9478        let mut b = __s_b.launch_builder(&f);
9479        b.arg(a)
9480            .arg(b_in)
9481            .arg(w0)
9482            .arg(w1)
9483            .arg(w2)
9484            .arg(res)
9485            .arg(&mut q0)
9486            .arg(&mut d0)
9487            .arg(out1)
9488            .arg(&mut q2)
9489            .arg(&mut d2)
9490            .arg(&nc)
9491            .arg(&e2);
9492        unsafe {
9493            b.launch(cfg)?;
9494        }
9495        Ok(((q0, d0), (q2, d2)))
9496    }
9497
9498    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9499    #[allow(clippy::too_many_arguments)]
9500    pub fn add_rms_norm3(
9501        &self,
9502        a: &CudaSlice<f32>,
9503        b_in: &CudaSlice<f32>,
9504        w0: &CudaSlice<f32>,
9505        w1: &CudaSlice<f32>,
9506        w2: &CudaSlice<f32>,
9507        res: &mut CudaSlice<f32>,
9508        d0: &mut CudaSlice<f32>,
9509        d1: &mut CudaSlice<f32>,
9510        d2: &mut CudaSlice<f32>,
9511        ncols: usize,
9512        nrows: usize,
9513        eps: f32,
9514    ) -> Result<(), Box<dyn std::error::Error>> {
9515        let f = self.func("add_rms_norm3_f32");
9516        let cfg = LaunchConfig {
9517            grid_dim: (nrows as u32, 1, 1),
9518            block_dim: (rms_block(), 1, 1),
9519            shared_mem_bytes: 0,
9520        };
9521        let (nc, e2) = (ncols as i32, eps);
9522        let __s_b = self.gpu.stream();
9523        let mut b = __s_b.launch_builder(&f);
9524        b.arg(a)
9525            .arg(b_in)
9526            .arg(w0)
9527            .arg(w1)
9528            .arg(w2)
9529            .arg(res)
9530            .arg(d0)
9531            .arg(d1)
9532            .arg(d2)
9533            .arg(&nc)
9534            .arg(&e2);
9535        unsafe {
9536            b.launch(cfg)?;
9537        }
9538        Ok(())
9539    }
9540
9541    /// dst = (a + b) * c (residual add + layer scale, one launch).
9542    pub fn add_scale(
9543        &self,
9544        a: &CudaSlice<f32>,
9545        b_in: &CudaSlice<f32>,
9546        c: f32,
9547        dst: &mut CudaSlice<f32>,
9548        n: usize,
9549    ) -> Result<(), Box<dyn std::error::Error>> {
9550        let f = self.func("add_scale_f32");
9551        let cfg = LaunchConfig::for_num_elems(n as u32);
9552        let ni = n as i32;
9553        let __s_b = self.gpu.stream();
9554        let mut b = __s_b.launch_builder(&f);
9555        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9556        unsafe {
9557            b.launch(cfg)?;
9558        }
9559        Ok(())
9560    }
9561
9562    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
9563    pub fn layer_norm_bias(
9564        &self,
9565        x: &CudaSlice<f32>,
9566        w: &CudaSlice<f32>,
9567        b: &CudaSlice<f32>,
9568        dst: &mut CudaSlice<f32>,
9569        ncols: usize,
9570        nrows: usize,
9571        eps: f32,
9572    ) -> Result<(), Box<dyn std::error::Error>> {
9573        let f = self.func("layer_norm_bias_f32");
9574        let (nc, e) = (ncols as i32, eps);
9575        let cfg = LaunchConfig {
9576            grid_dim: (nrows as u32, 1, 1),
9577            block_dim: (256, 1, 1),
9578            shared_mem_bytes: 0,
9579        };
9580        let __s_b = self.gpu.stream();
9581        let mut lb = __s_b.launch_builder(&f);
9582        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9583        unsafe {
9584            lb.launch(cfg)?;
9585        }
9586        Ok(())
9587    }
9588
9589    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
9590    pub fn gelu_tanh(
9591        &self,
9592        x: &CudaSlice<f32>,
9593        dst: &mut CudaSlice<f32>,
9594        n: usize,
9595    ) -> Result<(), Box<dyn std::error::Error>> {
9596        let f = self.func("gelu_tanh_f32");
9597        let ni = n as i64;
9598        let cfg = LaunchConfig {
9599            grid_dim: (n.div_ceil(256) as u32, 1, 1),
9600            block_dim: (256, 1, 1),
9601            shared_mem_bytes: 0,
9602        };
9603        let __s_b = self.gpu.stream();
9604        let mut lb = __s_b.launch_builder(&f);
9605        lb.arg(x).arg(&mut *dst).arg(&ni);
9606        unsafe {
9607            lb.launch(cfg)?;
9608        }
9609        Ok(())
9610    }
9611
9612    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
9613    pub fn row_softmax(
9614        &self,
9615        x: &mut CudaSlice<f32>,
9616        ncols: usize,
9617        nrows: usize,
9618    ) -> Result<(), Box<dyn std::error::Error>> {
9619        let f = self.func("row_softmax_f32");
9620        let nc = ncols as i32;
9621        let cfg = LaunchConfig {
9622            grid_dim: (nrows as u32, 1, 1),
9623            block_dim: (256, 1, 1),
9624            shared_mem_bytes: 0,
9625        };
9626        let __s_b = self.gpu.stream();
9627        let mut lb = __s_b.launch_builder(&f);
9628        lb.arg(&mut *x).arg(&nc);
9629        unsafe {
9630            lb.launch(cfg)?;
9631        }
9632        Ok(())
9633    }
9634
9635    pub fn rms_norm(
9636        &self,
9637        x: &CudaSlice<f32>,
9638        w: &CudaSlice<f32>,
9639        dst: &mut CudaSlice<f32>,
9640        ncols: usize,
9641        nrows: usize,
9642        eps: f32,
9643    ) -> Result<(), Box<dyn std::error::Error>> {
9644        let (nc, e) = (ncols as i32, eps);
9645        let kname = if Self::norm_ilp_on() {
9646            "rms_norm_f32_v2"
9647        } else {
9648            "rms_norm_f32"
9649        };
9650        if Self::pdl_on() && Self::pdl_wb_on() {
9651            use cudarc::driver::{DevicePtr, DevicePtrMut};
9652            let s = &self.gpu.stream();
9653            let (px, _g0) = x.device_ptr(s);
9654            let (pw, _g1) = w.device_ptr(s);
9655            let (pd, _g2) = dst.device_ptr_mut(s);
9656            let mut ps = [
9657                &px as *const _ as *mut std::ffi::c_void,
9658                &pw as *const _ as *mut _,
9659                &pd as *const _ as *mut _,
9660                &nc as *const _ as *mut _,
9661                &e as *const _ as *mut _,
9662            ];
9663            unsafe {
9664                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9665            }
9666            return Ok(());
9667        }
9668        let f = self.func(kname);
9669        let cfg = LaunchConfig {
9670            grid_dim: (nrows as u32, 1, 1),
9671            block_dim: (rms_block(), 1, 1),
9672            shared_mem_bytes: 0,
9673        };
9674        let __s_b = self.gpu.stream();
9675        let mut b = __s_b.launch_builder(&f);
9676        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9677        unsafe {
9678            b.launch(cfg)?;
9679        }
9680        Ok(())
9681    }
9682
9683    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
9684    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
9685    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
9686    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
9687    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
9688    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
9689    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
9690    pub fn rms_norm_decode(
9691        &self,
9692        x: &CudaSlice<f32>,
9693        w: &CudaSlice<f32>,
9694        dst: &mut CudaSlice<f32>,
9695        ncols: usize,
9696        nrows: usize,
9697        eps: f32,
9698    ) -> Result<(), Box<dyn std::error::Error>> {
9699        let f = self.func(if Self::norm_ilp_on() {
9700            "rms_norm_f32_v2"
9701        } else {
9702            "rms_norm_f32"
9703        });
9704        let cfg = LaunchConfig {
9705            grid_dim: (nrows as u32, 1, 1),
9706            block_dim: (1024, 1, 1),
9707            shared_mem_bytes: 0,
9708        };
9709        let (nc, e) = (ncols as i32, eps);
9710        let __s_b = self.gpu.stream();
9711        let mut b = __s_b.launch_builder(&f);
9712        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9713        unsafe {
9714            b.launch(cfg)?;
9715        }
9716        Ok(())
9717    }
9718
9719    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
9720    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
9721    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
9722    pub fn rms_norm_q8_1(
9723        &self,
9724        x: &CudaSlice<f32>,
9725        w: &CudaSlice<f32>,
9726        ncols: usize,
9727        nrows: usize,
9728        eps: f32,
9729    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9730        let nblk = ncols / 32;
9731        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9732        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9733        let (nc, e) = (ncols as i32, eps);
9734        if Self::pdl_on() {
9735            {
9736                use cudarc::driver::{DevicePtr, DevicePtrMut};
9737                let s = &self.gpu.stream();
9738                let (px, _g0) = x.device_ptr(s);
9739                let (pw, _g1) = w.device_ptr(s);
9740                let (pq, _g2) = q.device_ptr_mut(s);
9741                let (pd, _g3) = d.device_ptr_mut(s);
9742                let mut ps = [
9743                    &px as *const _ as *mut std::ffi::c_void,
9744                    &pw as *const _ as *mut _,
9745                    &pq as *const _ as *mut _,
9746                    &pd as *const _ as *mut _,
9747                    &nc as *const _ as *mut _,
9748                    &e as *const _ as *mut _,
9749                ];
9750                unsafe {
9751                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9752                }
9753            }
9754            return Ok((q, d));
9755        }
9756        let f = self.func("rms_norm_q8_1");
9757        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
9758        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
9759        let cfg = LaunchConfig {
9760            grid_dim: (nrows as u32, 1, 1),
9761            block_dim: (1024, 1, 1),
9762            shared_mem_bytes: 0,
9763        };
9764        let __s_b = self.gpu.stream();
9765        let mut b = __s_b.launch_builder(&f);
9766        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9767        unsafe {
9768            b.launch(cfg)?;
9769        }
9770        Ok((q, d))
9771    }
9772
9773    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
9774    /// PDL arm), caller-owned outputs.
9775    pub fn rms_norm_q8_1_into(
9776        &self,
9777        x: &CudaSlice<f32>,
9778        w: &CudaSlice<f32>,
9779        ncols: usize,
9780        nrows: usize,
9781        eps: f32,
9782        q: &mut CudaSlice<i8>,
9783        d: &mut CudaSlice<f32>,
9784    ) -> Result<(), Box<dyn std::error::Error>> {
9785        let nblk = ncols / 32;
9786        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9787        let (nc, e) = (ncols as i32, eps);
9788        if Self::pdl_on() {
9789            use cudarc::driver::{DevicePtr, DevicePtrMut};
9790            let s = &self.gpu.stream();
9791            let (px, _g0) = x.device_ptr(s);
9792            let (pw, _g1) = w.device_ptr(s);
9793            let (pq, _g2) = q.device_ptr_mut(s);
9794            let (pd, _g3) = d.device_ptr_mut(s);
9795            let mut ps = [
9796                &px as *const _ as *mut std::ffi::c_void,
9797                &pw as *const _ as *mut _,
9798                &pq as *const _ as *mut _,
9799                &pd as *const _ as *mut _,
9800                &nc as *const _ as *mut _,
9801                &e as *const _ as *mut _,
9802            ];
9803            unsafe {
9804                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9805            }
9806            return Ok(());
9807        }
9808        let f = self.func("rms_norm_q8_1");
9809        let cfg = LaunchConfig {
9810            grid_dim: (nrows as u32, 1, 1),
9811            block_dim: (1024, 1, 1),
9812            shared_mem_bytes: 0,
9813        };
9814        let __s_b = self.gpu.stream();
9815        let mut b = __s_b.launch_builder(&f);
9816        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9817        unsafe {
9818            b.launch(cfg)?;
9819        }
9820        Ok(())
9821    }
9822
9823    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
9824    pub fn quantize_q8_1_into(
9825        &self,
9826        x: &CudaSlice<f32>,
9827        m: usize,
9828        in_f: usize,
9829        q: &mut CudaSlice<i8>,
9830        d: &mut CudaSlice<f32>,
9831    ) -> Result<(), Box<dyn std::error::Error>> {
9832        let nblk = in_f / 32;
9833        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9834        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9835        let (inf, mi) = (in_f as i32, m as i32);
9836        if Self::pdl_on() && Self::pdl_wb_on() {
9837            use cudarc::driver::{DevicePtr, DevicePtrMut};
9838            let s = &self.gpu.stream();
9839            let (px, _g0) = x.device_ptr(s);
9840            let (pq, _g1) = q.device_ptr_mut(s);
9841            let (pd, _g2) = d.device_ptr_mut(s);
9842            let mut ps = [
9843                &px as *const _ as *mut std::ffi::c_void,
9844                &pq as *const _ as *mut _,
9845                &pd as *const _ as *mut _,
9846                &inf as *const _ as *mut _,
9847                &mi as *const _ as *mut _,
9848            ];
9849            unsafe {
9850                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9851            }
9852            return Ok(());
9853        }
9854        let f = self.func("quantize_q8_1");
9855        let __s_b = self.gpu.stream();
9856        let mut b = __s_b.launch_builder(&f);
9857        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9858        unsafe {
9859            b.launch(cfg)?;
9860        }
9861        Ok(())
9862    }
9863
9864    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
9865    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
9866    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
9867    pub fn add_rms_norm_q8_1(
9868        &self,
9869        a: &CudaSlice<f32>,
9870        b_in: &CudaSlice<f32>,
9871        w: &CudaSlice<f32>,
9872        res: &mut CudaSlice<f32>,
9873        ncols: usize,
9874        nrows: usize,
9875        eps: f32,
9876    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9877        let nblk = ncols / 32;
9878        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9879        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9880        let f = self.func("add_rms_norm_q8_1");
9881        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
9882        let cfg = LaunchConfig {
9883            grid_dim: (nrows as u32, 1, 1),
9884            block_dim: (1024, 1, 1),
9885            shared_mem_bytes: 0,
9886        };
9887        let (nc, e) = (ncols as i32, eps);
9888        let __s_bld = self.gpu.stream();
9889        let mut bld = __s_bld.launch_builder(&f);
9890        bld.arg(a)
9891            .arg(b_in)
9892            .arg(w)
9893            .arg(res)
9894            .arg(&mut q)
9895            .arg(&mut d)
9896            .arg(&nc)
9897            .arg(&e);
9898        unsafe {
9899            bld.launch(cfg)?;
9900        }
9901        Ok((q, d))
9902    }
9903
9904    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
9905    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
9906    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
9907    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
9908    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
9909    #[allow(clippy::too_many_arguments)]
9910    pub fn join_add_rms_norm_raw(
9911        &self,
9912        a0_raw: u64,
9913        a1_raw: u64,
9914        x: &CudaSlice<f32>,
9915        w: &CudaSlice<f32>,
9916        res: &mut CudaSlice<f32>,
9917        dst: &mut CudaSlice<f32>,
9918        ncols: usize,
9919        eps: f32,
9920    ) -> Result<(), Box<dyn std::error::Error>> {
9921        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9922            return Err("join_add_rms_norm geometry".into());
9923        }
9924        let f = self.func("join_add_rms_norm_f32");
9925        let cfg = LaunchConfig {
9926            grid_dim: (1, 1, 1),
9927            block_dim: (rms_block(), 1, 1),
9928            shared_mem_bytes: 0,
9929        };
9930        let (nc, e) = (ncols as i32, eps);
9931        let __s_b = self.gpu.stream();
9932        let mut b = __s_b.launch_builder(&f);
9933        b.arg(&a0_raw)
9934            .arg(&a1_raw)
9935            .arg(x)
9936            .arg(w)
9937            .arg(&mut *res)
9938            .arg(&mut *dst)
9939            .arg(&nc)
9940            .arg(&e);
9941        unsafe {
9942            b.launch(cfg)?;
9943        }
9944        Ok(())
9945    }
9946
9947    pub fn add_rms_norm(
9948        &self,
9949        a: &CudaSlice<f32>,
9950        b: &CudaSlice<f32>,
9951        w: &CudaSlice<f32>,
9952        res: &mut CudaSlice<f32>,
9953        dst: &mut CudaSlice<f32>,
9954        ncols: usize,
9955        nrows: usize,
9956        eps: f32,
9957    ) -> Result<(), Box<dyn std::error::Error>> {
9958        let (nc, e) = (ncols as i32, eps);
9959        let kname = if Self::norm_ilp_on() {
9960            "add_rms_norm_f32_v2"
9961        } else {
9962            "add_rms_norm_f32"
9963        };
9964        if Self::pdl_on() && Self::pdl_wb_on() {
9965            use cudarc::driver::{DevicePtr, DevicePtrMut};
9966            let s = &self.gpu.stream();
9967            let (pa, _g0) = a.device_ptr(s);
9968            let (pb, _g1) = b.device_ptr(s);
9969            let (pw, _g2) = w.device_ptr(s);
9970            let (pr, _g3) = res.device_ptr_mut(s);
9971            let (pd, _g4) = dst.device_ptr_mut(s);
9972            let mut ps = [
9973                &pa as *const _ as *mut std::ffi::c_void,
9974                &pb as *const _ as *mut _,
9975                &pw as *const _ as *mut _,
9976                &pr as *const _ as *mut _,
9977                &pd as *const _ as *mut _,
9978                &nc as *const _ as *mut _,
9979                &e as *const _ as *mut _,
9980            ];
9981            unsafe {
9982                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9983            }
9984            return Ok(());
9985        }
9986        let f = self.func(kname);
9987        let cfg = LaunchConfig {
9988            grid_dim: (nrows as u32, 1, 1),
9989            block_dim: (rms_block(), 1, 1),
9990            shared_mem_bytes: 0,
9991        };
9992        let __s_b2 = self.gpu.stream();
9993        let mut b2 = __s_b2.launch_builder(&f);
9994        b2.arg(a)
9995            .arg(b)
9996            .arg(w)
9997            .arg(&mut *res)
9998            .arg(&mut *dst)
9999            .arg(&nc)
10000            .arg(&e);
10001        unsafe {
10002            b2.launch(cfg)?;
10003        }
10004        Ok(())
10005    }
10006
10007    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
10008    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
10009    #[allow(clippy::too_many_arguments)]
10010    pub fn rms_pre_add_rms_norm(
10011        &self,
10012        a: &CudaSlice<f32>,
10013        wa: &CudaSlice<f32>,
10014        b: &CudaSlice<f32>,
10015        w: &CudaSlice<f32>,
10016        res: &mut CudaSlice<f32>,
10017        dst: &mut CudaSlice<f32>,
10018        ncols: usize,
10019        nrows: usize,
10020        eps: f32,
10021    ) -> Result<(), Box<dyn std::error::Error>> {
10022        let f = self.func("rms_pre_add_rms_norm_f32");
10023        let cfg = LaunchConfig {
10024            grid_dim: (nrows as u32, 1, 1),
10025            block_dim: (rms_block(), 1, 1),
10026            shared_mem_bytes: 0,
10027        };
10028        let (nc, e) = (ncols as i32, eps);
10029        let __s_b2 = self.gpu.stream();
10030        let mut b2 = __s_b2.launch_builder(&f);
10031        b2.arg(a)
10032            .arg(wa)
10033            .arg(b)
10034            .arg(w)
10035            .arg(&mut *res)
10036            .arg(&mut *dst)
10037            .arg(&nc)
10038            .arg(&e);
10039        unsafe {
10040            b2.launch(cfg)?;
10041        }
10042        Ok(())
10043    }
10044
10045    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
10046    #[allow(clippy::too_many_arguments)]
10047    pub fn rms_pre_add_rms_norm_q8z(
10048        &self,
10049        a: &CudaSlice<f32>,
10050        wa: &CudaSlice<f32>,
10051        b: &CudaSlice<f32>,
10052        w: &CudaSlice<f32>,
10053        res: &mut CudaSlice<f32>,
10054        dst: &mut CudaSlice<f32>,
10055        ncols: usize,
10056        nrows: usize,
10057        eps: f32,
10058    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10059        debug_assert!(ncols % 128 == 0);
10060        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10061        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10062        let (nc, e) = (ncols as i32, eps);
10063        if Self::pdl_on() {
10064            {
10065                use cudarc::driver::{DevicePtr, DevicePtrMut};
10066                let s = &self.gpu.stream();
10067                let (pa, _g0) = a.device_ptr(s);
10068                let (pwa, _g1) = wa.device_ptr(s);
10069                let (pb, _g2) = b.device_ptr(s);
10070                let (pw, _g3) = w.device_ptr(s);
10071                let (pr, _g4) = res.device_ptr_mut(s);
10072                let (pdst, _g5) = dst.device_ptr_mut(s);
10073                let (pq, _g6) = out_q.device_ptr_mut(s);
10074                let (pd, _g7) = out_d.device_ptr_mut(s);
10075                let mut ps = [
10076                    &pa as *const _ as *mut std::ffi::c_void,
10077                    &pwa as *const _ as *mut _,
10078                    &pb as *const _ as *mut _,
10079                    &pw as *const _ as *mut _,
10080                    &pr as *const _ as *mut _,
10081                    &pdst as *const _ as *mut _,
10082                    &pq as *const _ as *mut _,
10083                    &pd as *const _ as *mut _,
10084                    &nc as *const _ as *mut _,
10085                    &e as *const _ as *mut _,
10086                ];
10087                unsafe {
10088                    self.launch_pdl(
10089                        "rms_pre_add_rms_norm_q8z_f32",
10090                        (nrows as u32, 1, 1),
10091                        (rms_block(), 1, 1),
10092                        &mut ps,
10093                    )?;
10094                }
10095            }
10096            return Ok((out_q, out_d));
10097        }
10098        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10099        let cfg = LaunchConfig {
10100            grid_dim: (nrows as u32, 1, 1),
10101            block_dim: (rms_block(), 1, 1),
10102            shared_mem_bytes: 0,
10103        };
10104        let __s_b2 = self.gpu.stream();
10105        let mut b2 = __s_b2.launch_builder(&f);
10106        b2.arg(a)
10107            .arg(wa)
10108            .arg(b)
10109            .arg(w)
10110            .arg(&mut *res)
10111            .arg(&mut *dst)
10112            .arg(&mut out_q)
10113            .arg(&mut out_d)
10114            .arg(&nc)
10115            .arg(&e);
10116        unsafe {
10117            b2.launch(cfg)?;
10118        }
10119        Ok((out_q, out_d))
10120    }
10121
10122    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
10123    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
10124    /// body must stay attribute-free (the fused2_into precedent).
10125    #[allow(clippy::too_many_arguments)]
10126    pub fn rms_pre_add_rms_norm_q8z_into(
10127        &self,
10128        a: &CudaSlice<f32>,
10129        wa: &CudaSlice<f32>,
10130        b: &CudaSlice<f32>,
10131        w: &CudaSlice<f32>,
10132        res: &mut CudaSlice<f32>,
10133        dst: &mut CudaSlice<f32>,
10134        ncols: usize,
10135        nrows: usize,
10136        eps: f32,
10137        out_q: &mut CudaSlice<i8>,
10138        out_d: &mut CudaSlice<f32>,
10139    ) -> Result<(), Box<dyn std::error::Error>> {
10140        debug_assert!(ncols % 128 == 0);
10141        let (nc, e) = (ncols as i32, eps);
10142        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10143        let cfg = LaunchConfig {
10144            grid_dim: (nrows as u32, 1, 1),
10145            block_dim: (rms_block(), 1, 1),
10146            shared_mem_bytes: 0,
10147        };
10148        let __s_b = self.gpu.stream();
10149        let mut b2 = __s_b.launch_builder(&f);
10150        b2.arg(a)
10151            .arg(wa)
10152            .arg(b)
10153            .arg(w)
10154            .arg(&mut *res)
10155            .arg(&mut *dst)
10156            .arg(&mut *out_q)
10157            .arg(&mut *out_d)
10158            .arg(&nc)
10159            .arg(&e);
10160        unsafe {
10161            b2.launch(cfg)?;
10162        }
10163        Ok(())
10164    }
10165
10166    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
10167    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
10168    #[allow(clippy::too_many_arguments)]
10169    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10170        &self,
10171        a: &CudaSlice<f32>,
10172        wa: &CudaSlice<f32>,
10173        b_in: &CudaSlice<f32>,
10174        c: f32,
10175        w: &CudaSlice<f32>,
10176        res: &mut CudaSlice<f32>,
10177        ncols: usize,
10178        nrows: usize,
10179        eps: f32,
10180        out_q: &mut CudaSlice<i8>,
10181        out_d: &mut CudaSlice<f32>,
10182    ) -> Result<(), Box<dyn std::error::Error>> {
10183        debug_assert!(ncols % 128 == 0);
10184        let (nc, e2) = (ncols as i32, eps);
10185        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10186        let cfg = LaunchConfig {
10187            grid_dim: (nrows as u32, 1, 1),
10188            block_dim: (rms_block(), 1, 1),
10189            shared_mem_bytes: 0,
10190        };
10191        let __s_b = self.gpu.stream();
10192        let mut b2 = __s_b.launch_builder(&f);
10193        b2.arg(a)
10194            .arg(wa)
10195            .arg(b_in)
10196            .arg(&c)
10197            .arg(w)
10198            .arg(&mut *res)
10199            .arg(&mut *out_q)
10200            .arg(&mut *out_d)
10201            .arg(&nc)
10202            .arg(&e2);
10203        unsafe {
10204            b2.launch(cfg)?;
10205        }
10206        Ok(())
10207    }
10208
10209    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
10210    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
10211    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
10212    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
10213    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
10214    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
10215    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
10216    pub fn g4_pnfold_on() -> bool {
10217        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10218        *ON.get_or_init(|| {
10219            std::env::var("MEMRA_G4_PNFOLD")
10220                .map(|v| v != "0")
10221                .unwrap_or(true)
10222        })
10223    }
10224
10225    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
10226    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
10227    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
10228    pub fn build_q4_out_concat3(
10229        &self,
10230        w0: &crate::model::GpuTensor,
10231        w1: &crate::model::GpuTensor,
10232        w2: &crate::model::GpuTensor,
10233    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10234        use crate::model::GpuTensor;
10235        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10236            match w {
10237                GpuTensor::Quant {
10238                    qtype,
10239                    row_bytes,
10240                    rp,
10241                    ..
10242                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10243                _ => None,
10244            }
10245        };
10246        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10247        else {
10248            return Ok(None);
10249        };
10250        if rb0 != rb1
10251            || rb0 != rb2
10252            || w0.in_features() != w1.in_features()
10253            || w0.in_features() != w2.in_features()
10254        {
10255            return Ok(None);
10256        }
10257        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10258            match w {
10259                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10260                _ => unreachable!(),
10261            }
10262        }
10263        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10264        let total = rb0 * (o0 + o1 + o2);
10265        let mut cat = self.alloc_u8(total)?;
10266        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10267        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10268        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10269        Ok(Some(GpuTensor::Quant {
10270            bytes: cat,
10271            qtype: QT_Q4_0,
10272            row_bytes: rb0,
10273            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10274            scale: 1.0,
10275            rp: false,
10276            #[cfg(memra_cutlass)]
10277            cutlass: None,
10278            fp8: None,
10279            blk: None,
10280            rp4: None,
10281            f16: None,
10282        }))
10283    }
10284
10285    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
10286    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
10287    ///
10288    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
10289    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
10290    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
10291    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
10292    ///
10293    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
10294    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
10295    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
10296    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
10297    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
10298    ///
10299    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
10300    /// width. A future partial-rotary caller fails at its first launch with the geometry named
10301    /// instead of serving quietly wrong logits.
10302    fn full_width_rope_only(
10303        kernel: &str,
10304        n_rot: usize,
10305        head_dim: usize,
10306    ) -> Result<(), Box<dyn std::error::Error>> {
10307        if n_rot == head_dim {
10308            return Ok(());
10309        }
10310        Err(format!(
10311            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10312             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10313             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10314             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10315             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10316        )
10317        .into())
10318    }
10319
10320    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
10321    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10322    /// ([`Engine::full_width_rope_only`]).
10323    #[allow(clippy::too_many_arguments)]
10324    pub fn rms_norm_qkv_rope_cat(
10325        &self,
10326        qkv: &CudaSlice<f32>,
10327        wq: &CudaSlice<f32>,
10328        wk: &CudaSlice<f32>,
10329        wv: &CudaSlice<f32>,
10330        q: &mut CudaSlice<f32>,
10331        k: &mut CudaSlice<f32>,
10332        v: &mut CudaSlice<f32>,
10333        head_dim: usize,
10334        n_rot: usize,
10335        rq: usize,
10336        rk: usize,
10337        pos: &CudaSlice<i32>,
10338        nh_q: usize,
10339        nh_k: usize,
10340        base: f32,
10341        freq_scale: f32,
10342        ff: Option<&CudaSlice<f32>>,
10343        eps: f32,
10344    ) -> Result<(), Box<dyn std::error::Error>> {
10345        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10346        let rows = rq + rk + rk;
10347        let theta_scale = base.powf(-2.0 / head_dim as f32);
10348        let (nc, rqi, rki, nhq, nhk) = (
10349            head_dim as i32,
10350            rq as i32,
10351            rk as i32,
10352            nh_q as i32,
10353            nh_k as i32,
10354        );
10355        if Self::pdl_on() {
10356            use cudarc::driver::{DevicePtr, DevicePtrMut};
10357            let s = &self.gpu.stream();
10358            let (pqkv, _g0) = qkv.device_ptr(s);
10359            let (pwq, _g1) = wq.device_ptr(s);
10360            let (pwk, _g2) = wk.device_ptr(s);
10361            let (pwv, _g3) = wv.device_ptr(s);
10362            let (pq, _g4) = q.device_ptr_mut(s);
10363            let (pk, _g5) = k.device_ptr_mut(s);
10364            let (pv, _g6) = v.device_ptr_mut(s);
10365            let (ppos, _g7) = pos.device_ptr(s);
10366            let (pff, _g8) = match ff {
10367                Some(t) => {
10368                    let (p, g) = t.device_ptr(s);
10369                    (p, Some(g))
10370                }
10371                None => (0, None),
10372            };
10373            let mut ps = [
10374                &pqkv as *const _ as *mut std::ffi::c_void,
10375                &pwq as *const _ as *mut _,
10376                &pwk as *const _ as *mut _,
10377                &pwv as *const _ as *mut _,
10378                &pq as *const _ as *mut _,
10379                &pk as *const _ as *mut _,
10380                &pv as *const _ as *mut _,
10381                &nc as *const _ as *mut _,
10382                &rqi as *const _ as *mut _,
10383                &rki as *const _ as *mut _,
10384                &ppos as *const _ as *mut _,
10385                &nhq as *const _ as *mut _,
10386                &nhk as *const _ as *mut _,
10387                &theta_scale as *const _ as *mut _,
10388                &freq_scale as *const _ as *mut _,
10389                &pff as *const _ as *mut _,
10390                &eps as *const _ as *mut _,
10391            ];
10392            unsafe {
10393                self.launch_pdl(
10394                    "rms_norm_qkv_rope_cat_f32",
10395                    (rows as u32, 1, 1),
10396                    (rms_block(), 1, 1),
10397                    &mut ps,
10398                )?;
10399            }
10400            return Ok(());
10401        }
10402        let f = self.func("rms_norm_qkv_rope_cat_f32");
10403        let cfg = LaunchConfig {
10404            grid_dim: (rows as u32, 1, 1),
10405            block_dim: (rms_block(), 1, 1),
10406            shared_mem_bytes: 0,
10407        };
10408        let __s_b = self.gpu.stream();
10409        let mut b = __s_b.launch_builder(&f);
10410        match ff {
10411            Some(t) => {
10412                b.arg(qkv)
10413                    .arg(wq)
10414                    .arg(wk)
10415                    .arg(wv)
10416                    .arg(&mut *q)
10417                    .arg(&mut *k)
10418                    .arg(&mut *v)
10419                    .arg(&nc)
10420                    .arg(&rqi)
10421                    .arg(&rki)
10422                    .arg(pos)
10423                    .arg(&nhq)
10424                    .arg(&nhk)
10425                    .arg(&theta_scale)
10426                    .arg(&freq_scale)
10427                    .arg(t)
10428                    .arg(&eps);
10429                unsafe {
10430                    b.launch(cfg)?;
10431                }
10432            }
10433            None => {
10434                let null: u64 = 0;
10435                b.arg(qkv)
10436                    .arg(wq)
10437                    .arg(wk)
10438                    .arg(wv)
10439                    .arg(&mut *q)
10440                    .arg(&mut *k)
10441                    .arg(&mut *v)
10442                    .arg(&nc)
10443                    .arg(&rqi)
10444                    .arg(&rki)
10445                    .arg(pos)
10446                    .arg(&nhq)
10447                    .arg(&nhk)
10448                    .arg(&theta_scale)
10449                    .arg(&freq_scale)
10450                    .arg(&null)
10451                    .arg(&eps);
10452                unsafe {
10453                    b.launch(cfg)?;
10454                }
10455            }
10456        }
10457        Ok(())
10458    }
10459
10460    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10461    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10462    /// ([`Engine::full_width_rope_only`]).
10463    #[allow(clippy::too_many_arguments)]
10464    pub fn rms_norm_qkv_rope(
10465        &self,
10466        q0: &CudaSlice<f32>,
10467        k0: &CudaSlice<f32>,
10468        v0: &CudaSlice<f32>,
10469        wq: &CudaSlice<f32>,
10470        wk: &CudaSlice<f32>,
10471        wv: &CudaSlice<f32>,
10472        q: &mut CudaSlice<f32>,
10473        k: &mut CudaSlice<f32>,
10474        v: &mut CudaSlice<f32>,
10475        head_dim: usize,
10476        n_rot: usize,
10477        rq: usize,
10478        rk: usize,
10479        pos: &CudaSlice<i32>,
10480        nh_q: usize,
10481        nh_k: usize,
10482        base: f32,
10483        freq_scale: f32,
10484        ff: Option<&CudaSlice<f32>>,
10485        eps: f32,
10486    ) -> Result<(), Box<dyn std::error::Error>> {
10487        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10488        let f = self.func("rms_norm_qkv_rope_f32");
10489        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10490        let cfg = LaunchConfig {
10491            grid_dim: (rows as u32, 1, 1),
10492            block_dim: (rms_block(), 1, 1),
10493            shared_mem_bytes: 0,
10494        };
10495        let theta_scale = base.powf(-2.0 / head_dim as f32);
10496        let (nc, rqi, rki, nhq, nhk) = (
10497            head_dim as i32,
10498            rq as i32,
10499            rk as i32,
10500            nh_q as i32,
10501            nh_k as i32,
10502        );
10503        let __s_b = self.gpu.stream();
10504        let mut b = __s_b.launch_builder(&f);
10505        match ff {
10506            Some(t) => {
10507                b.arg(q0)
10508                    .arg(k0)
10509                    .arg(v0)
10510                    .arg(wq)
10511                    .arg(wk)
10512                    .arg(wv)
10513                    .arg(&mut *q)
10514                    .arg(&mut *k)
10515                    .arg(&mut *v)
10516                    .arg(&nc)
10517                    .arg(&rqi)
10518                    .arg(&rki)
10519                    .arg(pos)
10520                    .arg(&nhq)
10521                    .arg(&nhk)
10522                    .arg(&theta_scale)
10523                    .arg(&freq_scale)
10524                    .arg(t)
10525                    .arg(&eps);
10526                unsafe {
10527                    b.launch(cfg)?;
10528                }
10529            }
10530            None => {
10531                let null: u64 = 0;
10532                b.arg(q0)
10533                    .arg(k0)
10534                    .arg(v0)
10535                    .arg(wq)
10536                    .arg(wk)
10537                    .arg(wv)
10538                    .arg(&mut *q)
10539                    .arg(&mut *k)
10540                    .arg(&mut *v)
10541                    .arg(&nc)
10542                    .arg(&rqi)
10543                    .arg(&rki)
10544                    .arg(pos)
10545                    .arg(&nhq)
10546                    .arg(&nhk)
10547                    .arg(&theta_scale)
10548                    .arg(&freq_scale)
10549                    .arg(&null)
10550                    .arg(&eps);
10551                unsafe {
10552                    b.launch(cfg)?;
10553                }
10554            }
10555        }
10556        Ok(())
10557    }
10558
10559    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
10560    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
10561    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
10562    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10563    /// ([`Engine::full_width_rope_only`]).
10564    #[allow(clippy::too_many_arguments)]
10565    pub fn rms_norm_qkv_rope_append_dc(
10566        &self,
10567        q0: &CudaSlice<f32>,
10568        k0: &CudaSlice<f32>,
10569        v0: &CudaSlice<f32>,
10570        wq: &CudaSlice<f32>,
10571        wk: &CudaSlice<f32>,
10572        wv: &CudaSlice<f32>,
10573        q: &mut CudaSlice<f32>,
10574        k: &mut CudaSlice<f32>,
10575        v: &mut CudaSlice<f32>,
10576        head_dim: usize,
10577        n_rot: usize,
10578        rq: usize,
10579        rk: usize,
10580        pos: &CudaSlice<i32>,
10581        nh_q: usize,
10582        nh_k: usize,
10583        base: f32,
10584        freq_scale: f32,
10585        ff: Option<&CudaSlice<f32>>,
10586        eps: f32,
10587        kc: &mut CudaSlice<u8>,
10588        vc: &mut CudaSlice<u8>,
10589        t_dev: &CudaSlice<i32>,
10590        k_tok_bytes: usize,
10591        v_tok_bytes: usize,
10592        g: bool,
10593    ) -> Result<(), Box<dyn std::error::Error>> {
10594        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10595        let rows = rq + rk + rk;
10596        let theta_scale = base.powf(-2.0 / head_dim as f32);
10597        let (nc, rqi, rki, nhq, nhk) = (
10598            head_dim as i32,
10599            rq as i32,
10600            rk as i32,
10601            nh_q as i32,
10602            nh_k as i32,
10603        );
10604        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10605        if Self::pdl_on() && Self::pdl_wb_on() {
10606            use cudarc::driver::{DevicePtr, DevicePtrMut};
10607            let s = &self.gpu.stream();
10608            let (p0, _a0) = q0.device_ptr(s);
10609            let (p1, _a1) = k0.device_ptr(s);
10610            let (p2, _a2) = v0.device_ptr(s);
10611            let (pwq, _a3) = wq.device_ptr(s);
10612            let (pwk, _a4) = wk.device_ptr(s);
10613            let (pwv, _a5) = wv.device_ptr(s);
10614            let (pq, _a6) = q.device_ptr_mut(s);
10615            let (pk, _a7) = k.device_ptr_mut(s);
10616            let (pv, _a8) = v.device_ptr_mut(s);
10617            let (pp, _a9) = pos.device_ptr(s);
10618            let pff: u64 = match ff {
10619                Some(t) => {
10620                    let (p, _gg) = t.device_ptr(s);
10621                    p as u64
10622                }
10623                None => 0,
10624            };
10625            let (pkc, _a10) = kc.device_ptr_mut(s);
10626            let (pvc, _a11) = vc.device_ptr_mut(s);
10627            let (pt, _a12) = t_dev.device_ptr(s);
10628            let mut ps = [
10629                &p0 as *const _ as *mut std::ffi::c_void,
10630                &p1 as *const _ as *mut _,
10631                &p2 as *const _ as *mut _,
10632                &pwq as *const _ as *mut _,
10633                &pwk as *const _ as *mut _,
10634                &pwv as *const _ as *mut _,
10635                &pq as *const _ as *mut _,
10636                &pk as *const _ as *mut _,
10637                &pv as *const _ as *mut _,
10638                &nc as *const _ as *mut _,
10639                &rqi as *const _ as *mut _,
10640                &rki as *const _ as *mut _,
10641                &pp as *const _ as *mut _,
10642                &nhq as *const _ as *mut _,
10643                &nhk as *const _ as *mut _,
10644                &theta_scale as *const _ as *mut _,
10645                &freq_scale as *const _ as *mut _,
10646                &pff as *const _ as *mut _,
10647                &eps as *const _ as *mut _,
10648                &pkc as *const _ as *mut _,
10649                &pvc as *const _ as *mut _,
10650                &pt as *const _ as *mut _,
10651                &ktb as *const _ as *mut _,
10652                &vtb as *const _ as *mut _,
10653            ];
10654            unsafe {
10655                self.launch_pdl_flash(
10656                    g,
10657                    "rms_norm_qkv_rope_append_dc_f32",
10658                    (rows as u32, 1, 1),
10659                    (rms_block(), 1, 1),
10660                    0,
10661                    &mut ps,
10662                )?;
10663            }
10664            return Ok(());
10665        }
10666        let f = if g {
10667            self.func_g("rms_norm_qkv_rope_append_dc_f32")
10668        } else {
10669            self.func("rms_norm_qkv_rope_append_dc_f32")
10670        };
10671        let cfg = LaunchConfig {
10672            grid_dim: (rows as u32, 1, 1),
10673            block_dim: (rms_block(), 1, 1),
10674            shared_mem_bytes: 0,
10675        };
10676        let __s_b = self.gpu.stream();
10677        let mut b = __s_b.launch_builder(&f);
10678        match ff {
10679            Some(t) => {
10680                b.arg(q0)
10681                    .arg(k0)
10682                    .arg(v0)
10683                    .arg(wq)
10684                    .arg(wk)
10685                    .arg(wv)
10686                    .arg(&mut *q)
10687                    .arg(&mut *k)
10688                    .arg(&mut *v)
10689                    .arg(&nc)
10690                    .arg(&rqi)
10691                    .arg(&rki)
10692                    .arg(pos)
10693                    .arg(&nhq)
10694                    .arg(&nhk)
10695                    .arg(&theta_scale)
10696                    .arg(&freq_scale)
10697                    .arg(t)
10698                    .arg(&eps)
10699                    .arg(&mut *kc)
10700                    .arg(&mut *vc)
10701                    .arg(t_dev)
10702                    .arg(&ktb)
10703                    .arg(&vtb);
10704                unsafe {
10705                    b.launch(cfg)?;
10706                }
10707            }
10708            None => {
10709                let null: u64 = 0;
10710                b.arg(q0)
10711                    .arg(k0)
10712                    .arg(v0)
10713                    .arg(wq)
10714                    .arg(wk)
10715                    .arg(wv)
10716                    .arg(&mut *q)
10717                    .arg(&mut *k)
10718                    .arg(&mut *v)
10719                    .arg(&nc)
10720                    .arg(&rqi)
10721                    .arg(&rki)
10722                    .arg(pos)
10723                    .arg(&nhq)
10724                    .arg(&nhk)
10725                    .arg(&theta_scale)
10726                    .arg(&freq_scale)
10727                    .arg(&null)
10728                    .arg(&eps)
10729                    .arg(&mut *kc)
10730                    .arg(&mut *vc)
10731                    .arg(t_dev)
10732                    .arg(&ktb)
10733                    .arg(&vtb);
10734                unsafe {
10735                    b.launch(cfg)?;
10736                }
10737            }
10738        }
10739        Ok(())
10740    }
10741
10742    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
10743    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
10744    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
10745    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
10746    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
10747    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
10748    /// `head_dim` ([`Engine::full_width_rope_only`]).
10749    #[allow(clippy::too_many_arguments)]
10750    pub fn rms_norm_qkv_rope_append(
10751        &self,
10752        q0: &CudaSlice<f32>,
10753        k0: &CudaSlice<f32>,
10754        v0: &CudaSlice<f32>,
10755        wq: &CudaSlice<f32>,
10756        wk: &CudaSlice<f32>,
10757        wv: &CudaSlice<f32>,
10758        q: &mut CudaSlice<f32>,
10759        k: &mut CudaSlice<f32>,
10760        v: &mut CudaSlice<f32>,
10761        head_dim: usize,
10762        n_rot: usize,
10763        rq: usize,
10764        rk: usize,
10765        pos: &CudaSlice<i32>,
10766        nh_q: usize,
10767        nh_k: usize,
10768        base: f32,
10769        freq_scale: f32,
10770        ff: Option<&CudaSlice<f32>>,
10771        eps: f32,
10772        kc: &mut CudaSlice<u8>,
10773        vc: &mut CudaSlice<u8>,
10774        t: usize,
10775        k_tok_bytes: usize,
10776        v_tok_bytes: usize,
10777        g: bool,
10778    ) -> Result<(), Box<dyn std::error::Error>> {
10779        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10780        let rows = rq + rk + rk;
10781        let theta_scale = base.powf(-2.0 / head_dim as f32);
10782        let (nc, rqi, rki, nhq, nhk) = (
10783            head_dim as i32,
10784            rq as i32,
10785            rk as i32,
10786            nh_q as i32,
10787            nh_k as i32,
10788        );
10789        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10790        let ti = t as i32;
10791        if Self::pdl_on() && Self::pdl_wb_on() {
10792            use cudarc::driver::{DevicePtr, DevicePtrMut};
10793            let s = &self.gpu.stream();
10794            let (p0, _a0) = q0.device_ptr(s);
10795            let (p1, _a1) = k0.device_ptr(s);
10796            let (p2, _a2) = v0.device_ptr(s);
10797            let (pwq, _a3) = wq.device_ptr(s);
10798            let (pwk, _a4) = wk.device_ptr(s);
10799            let (pwv, _a5) = wv.device_ptr(s);
10800            let (pq, _a6) = q.device_ptr_mut(s);
10801            let (pk, _a7) = k.device_ptr_mut(s);
10802            let (pv, _a8) = v.device_ptr_mut(s);
10803            let (pp, _a9) = pos.device_ptr(s);
10804            let pff: u64 = match ff {
10805                Some(t) => {
10806                    let (p, _gg) = t.device_ptr(s);
10807                    p as u64
10808                }
10809                None => 0,
10810            };
10811            let (pkc, _a10) = kc.device_ptr_mut(s);
10812            let (pvc, _a11) = vc.device_ptr_mut(s);
10813            let mut ps = [
10814                &p0 as *const _ as *mut std::ffi::c_void,
10815                &p1 as *const _ as *mut _,
10816                &p2 as *const _ as *mut _,
10817                &pwq as *const _ as *mut _,
10818                &pwk as *const _ as *mut _,
10819                &pwv as *const _ as *mut _,
10820                &pq as *const _ as *mut _,
10821                &pk as *const _ as *mut _,
10822                &pv as *const _ as *mut _,
10823                &nc as *const _ as *mut _,
10824                &rqi as *const _ as *mut _,
10825                &rki as *const _ as *mut _,
10826                &pp as *const _ as *mut _,
10827                &nhq as *const _ as *mut _,
10828                &nhk as *const _ as *mut _,
10829                &theta_scale as *const _ as *mut _,
10830                &freq_scale as *const _ as *mut _,
10831                &pff as *const _ as *mut _,
10832                &eps as *const _ as *mut _,
10833                &pkc as *const _ as *mut _,
10834                &pvc as *const _ as *mut _,
10835                &ti as *const _ as *mut _,
10836                &ktb as *const _ as *mut _,
10837                &vtb as *const _ as *mut _,
10838            ];
10839            unsafe {
10840                self.launch_pdl_flash(
10841                    g,
10842                    "rms_norm_qkv_rope_append_f32",
10843                    (rows as u32, 1, 1),
10844                    (rms_block(), 1, 1),
10845                    0,
10846                    &mut ps,
10847                )?;
10848            }
10849            return Ok(());
10850        }
10851        let f = if g {
10852            self.func_g("rms_norm_qkv_rope_append_f32")
10853        } else {
10854            self.func("rms_norm_qkv_rope_append_f32")
10855        };
10856        let cfg = LaunchConfig {
10857            grid_dim: (rows as u32, 1, 1),
10858            block_dim: (rms_block(), 1, 1),
10859            shared_mem_bytes: 0,
10860        };
10861        let __s_b = self.gpu.stream();
10862        let mut b = __s_b.launch_builder(&f);
10863        let null: u64 = 0;
10864        b.arg(q0)
10865            .arg(k0)
10866            .arg(v0)
10867            .arg(wq)
10868            .arg(wk)
10869            .arg(wv)
10870            .arg(&mut *q)
10871            .arg(&mut *k)
10872            .arg(&mut *v)
10873            .arg(&nc)
10874            .arg(&rqi)
10875            .arg(&rki)
10876            .arg(pos)
10877            .arg(&nhq)
10878            .arg(&nhk)
10879            .arg(&theta_scale)
10880            .arg(&freq_scale);
10881        match ff {
10882            Some(t) => {
10883                b.arg(t);
10884            }
10885            None => {
10886                b.arg(&null);
10887            }
10888        }
10889        b.arg(&eps)
10890            .arg(&mut *kc)
10891            .arg(&mut *vc)
10892            .arg(&ti)
10893            .arg(&ktb)
10894            .arg(&vtb);
10895        unsafe {
10896            b.launch(cfg)?;
10897        }
10898        Ok(())
10899    }
10900
10901    pub fn add_q8_1(
10902        &self,
10903        a: &CudaSlice<f32>,
10904        b: &CudaSlice<f32>,
10905        res: &mut CudaSlice<f32>,
10906        ncols: usize,
10907        nrows: usize,
10908    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10909        debug_assert!(ncols % 128 == 0);
10910        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10911        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10912        let f = self.func("add_q8_1_f32");
10913        let cfg = LaunchConfig {
10914            grid_dim: (nrows as u32, 1, 1),
10915            block_dim: (rms_block(), 1, 1),
10916            shared_mem_bytes: 0,
10917        };
10918        let nc = ncols as i32;
10919        let __s_b2 = self.gpu.stream();
10920        let mut b2 = __s_b2.launch_builder(&f);
10921        b2.arg(a)
10922            .arg(b)
10923            .arg(&mut *res)
10924            .arg(&mut out_q)
10925            .arg(&mut out_d)
10926            .arg(&nc);
10927        unsafe {
10928            b2.launch(cfg)?;
10929        }
10930        Ok((out_q, out_d))
10931    }
10932
10933    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
10934    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
10935    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
10936    pub fn rms_pre_add_q8_1(
10937        &self,
10938        a: &CudaSlice<f32>,
10939        wa: &CudaSlice<f32>,
10940        b: &CudaSlice<f32>,
10941        res: &mut CudaSlice<f32>,
10942        ncols: usize,
10943        nrows: usize,
10944        eps: f32,
10945    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10946        debug_assert!(ncols % 128 == 0);
10947        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10948        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10949        let f = self.func("rms_pre_add_q8_1_f32");
10950        let cfg = LaunchConfig {
10951            grid_dim: (nrows as u32, 1, 1),
10952            block_dim: (rms_block(), 1, 1),
10953            shared_mem_bytes: 0,
10954        };
10955        let (nc, ep) = (ncols as i32, eps);
10956        let __s_b2 = self.gpu.stream();
10957        let mut b2 = __s_b2.launch_builder(&f);
10958        b2.arg(a)
10959            .arg(wa)
10960            .arg(b)
10961            .arg(&mut *res)
10962            .arg(&mut out_q)
10963            .arg(&mut out_d)
10964            .arg(&nc)
10965            .arg(&ep);
10966        unsafe {
10967            b2.launch(cfg)?;
10968        }
10969        Ok((out_q, out_d))
10970    }
10971
10972    /// L2 norm per row (head_dim), no weight.
10973    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
10974    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
10975    pub fn l2_v2_on(ncols: usize) -> bool {
10976        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10977    }
10978
10979    pub fn l2_norm_pp(
10980        &self,
10981        x: &CudaSlice<f32>,
10982        dst: &mut CudaSlice<f32>,
10983        dst16: Option<&mut CudaSlice<u8>>,
10984        ncols: usize,
10985        nrows: usize,
10986        eps: f32,
10987    ) -> Result<(), Box<dyn std::error::Error>> {
10988        if Self::l2_v2_on(ncols) {
10989            let f = self.func("l2_norm_pp_v2_f32");
10990            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
10991            let cfg = LaunchConfig {
10992                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10993                block_dim: (256, 1, 1),
10994                shared_mem_bytes: 0,
10995            };
10996            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10997            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
10998            let d16: u64 = match dst16 {
10999                Some(d) => self.addr_u8(d),
11000                None => 0,
11001            };
11002            let __s_b = self.gpu.stream();
11003            let mut b = __s_b.launch_builder(&f);
11004            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11005            unsafe {
11006                b.launch(cfg)?;
11007            }
11008            return Ok(());
11009        }
11010        self.l2_norm(x, dst, ncols, nrows, eps)
11011    }
11012
11013    pub fn l2_norm(
11014        &self,
11015        x: &CudaSlice<f32>,
11016        dst: &mut CudaSlice<f32>,
11017        ncols: usize,
11018        nrows: usize,
11019        eps: f32,
11020    ) -> Result<(), Box<dyn std::error::Error>> {
11021        let f = self.func("l2_norm_f32");
11022        let cfg = LaunchConfig {
11023            grid_dim: (nrows as u32, 1, 1),
11024            block_dim: (256, 1, 1),
11025            shared_mem_bytes: 0,
11026        };
11027        let (nc, e) = (ncols as i32, eps);
11028        let __s_b = self.gpu.stream();
11029        let mut b = __s_b.launch_builder(&f);
11030        b.arg(x).arg(dst).arg(&nc).arg(&e);
11031        unsafe {
11032            b.launch(cfg)?;
11033        }
11034        Ok(())
11035    }
11036
11037    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
11038    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
11039    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
11040    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
11041    /// propagate through gdn_scan and flip argmax on marginal logits.
11042    pub fn l2_norm_decode(
11043        &self,
11044        x: &CudaSlice<f32>,
11045        dst: &mut CudaSlice<f32>,
11046        ncols: usize,
11047        nrows: usize,
11048        eps: f32,
11049    ) -> Result<(), Box<dyn std::error::Error>> {
11050        let f = self.func("l2_norm_f32");
11051        let cfg = LaunchConfig {
11052            grid_dim: (nrows as u32, 1, 1),
11053            block_dim: (32, 1, 1),
11054            shared_mem_bytes: 0,
11055        };
11056        let (nc, e) = (ncols as i32, eps);
11057        let __s_b = self.gpu.stream();
11058        let mut b = __s_b.launch_builder(&f);
11059        b.arg(x).arg(dst).arg(&nc).arg(&e);
11060        unsafe {
11061            b.launch(cfg)?;
11062        }
11063        Ok(())
11064    }
11065
11066    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
11067    pub fn rope_neox(
11068        &self,
11069        x: &mut CudaSlice<f32>,
11070        pos: &CudaSlice<i32>,
11071        head_dim: usize,
11072        n_dims: usize,
11073        n_heads: usize,
11074        n_tokens: usize,
11075        freq_base: f32,
11076        freq_scale: f32,
11077    ) -> Result<(), Box<dyn std::error::Error>> {
11078        let f = self.func("rope_neox_f32");
11079        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11080        let grid = (n_heads * n_tokens) as u32;
11081        let cfg = LaunchConfig {
11082            grid_dim: (grid, 1, 1),
11083            block_dim: ((head_dim / 2) as u32, 1, 1),
11084            shared_mem_bytes: 0,
11085        };
11086        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11087        let __s_b = self.gpu.stream();
11088        let mut b = __s_b.launch_builder(&f);
11089        b.arg(x)
11090            .arg(pos)
11091            .arg(&hd)
11092            .arg(&nd)
11093            .arg(&nh)
11094            .arg(&theta_scale)
11095            .arg(&freq_scale);
11096        unsafe {
11097            b.launch(cfg)?;
11098        }
11099        Ok(())
11100    }
11101
11102    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
11103    pub fn rope_neox_ff(
11104        &self,
11105        x: &mut CudaSlice<f32>,
11106        pos: &CudaSlice<i32>,
11107        head_dim: usize,
11108        n_dims: usize,
11109        n_heads: usize,
11110        n_tokens: usize,
11111        freq_base: f32,
11112        freq_scale: f32,
11113        ff: &CudaSlice<f32>,
11114    ) -> Result<(), Box<dyn std::error::Error>> {
11115        let f = self.func("rope_neox_ff_f32");
11116        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11117        let grid = (n_heads * n_tokens) as u32;
11118        let cfg = LaunchConfig {
11119            grid_dim: (grid, 1, 1),
11120            block_dim: ((head_dim / 2) as u32, 1, 1),
11121            shared_mem_bytes: 0,
11122        };
11123        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11124        let __s_b = self.gpu.stream();
11125        let mut b = __s_b.launch_builder(&f);
11126        b.arg(x)
11127            .arg(pos)
11128            .arg(&hd)
11129            .arg(&nd)
11130            .arg(&nh)
11131            .arg(&theta_scale)
11132            .arg(&freq_scale)
11133            .arg(ff);
11134        unsafe {
11135            b.launch(cfg)?;
11136        }
11137        Ok(())
11138    }
11139
11140    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
11141    #[allow(clippy::too_many_arguments)]
11142    pub fn rope_neox2(
11143        &self,
11144        q: &mut CudaSlice<f32>,
11145        k: &mut CudaSlice<f32>,
11146        pos: &CudaSlice<i32>,
11147        head_dim: usize,
11148        n_dims: usize,
11149        nh_q: usize,
11150        nh_k: usize,
11151        n_tokens: usize,
11152        freq_base: f32,
11153        freq_scale: f32,
11154        ff: Option<&CudaSlice<f32>>,
11155    ) -> Result<(), Box<dyn std::error::Error>> {
11156        let f = self.func("rope_neox2_f32");
11157        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11158        let grid = ((nh_q + nh_k) * n_tokens) as u32;
11159        let cfg = LaunchConfig {
11160            grid_dim: (grid, 1, 1),
11161            block_dim: ((head_dim / 2) as u32, 1, 1),
11162            shared_mem_bytes: 0,
11163        };
11164        let (hd, nd, nq, nk, nt) = (
11165            head_dim as i32,
11166            n_dims as i32,
11167            nh_q as i32,
11168            nh_k as i32,
11169            n_tokens as i32,
11170        );
11171        let __s_b = self.gpu.stream();
11172        let mut b = __s_b.launch_builder(&f);
11173        b.arg(q)
11174            .arg(k)
11175            .arg(pos)
11176            .arg(&hd)
11177            .arg(&nd)
11178            .arg(&nq)
11179            .arg(&nk)
11180            .arg(&nt)
11181            .arg(&theta_scale)
11182            .arg(&freq_scale);
11183        match ff {
11184            Some(ffv) => {
11185                b.arg(ffv);
11186                unsafe {
11187                    b.launch(cfg)?;
11188                }
11189            }
11190            None => {
11191                let null: u64 = 0;
11192                b.arg(&null);
11193                unsafe {
11194                    b.launch(cfg)?;
11195                }
11196            }
11197        }
11198        Ok(())
11199    }
11200
11201    /// gemma4 R1: dst = GELU_tanh(gate) * up.
11202    pub fn gelu_tanh_mul(
11203        &self,
11204        gate: &CudaSlice<f32>,
11205        up: &CudaSlice<f32>,
11206        dst: &mut CudaSlice<f32>,
11207        n: usize,
11208    ) -> Result<(), Box<dyn std::error::Error>> {
11209        let f = self.func("gelu_tanh_mul_f32");
11210        let cfg = LaunchConfig::for_num_elems(n as u32);
11211        let ni = n as i32;
11212        let __s_b = self.gpu.stream();
11213        let mut b = __s_b.launch_builder(&f);
11214        b.arg(gate).arg(up).arg(dst).arg(&ni);
11215        unsafe {
11216            b.launch(cfg)?;
11217        }
11218        Ok(())
11219    }
11220
11221    pub fn silu_mul(
11222        &self,
11223        gate: &CudaSlice<f32>,
11224        up: &CudaSlice<f32>,
11225        dst: &mut CudaSlice<f32>,
11226        n: usize,
11227    ) -> Result<(), Box<dyn std::error::Error>> {
11228        let f = self.func("silu_mul_f32");
11229        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11230        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11231        let ni = n as i32;
11232        let __s_b = self.gpu.stream();
11233        let mut b = __s_b.launch_builder(&f);
11234        b.arg(gate).arg(up).arg(dst).arg(&ni);
11235        unsafe {
11236            b.launch(cfg)?;
11237        }
11238        Ok(())
11239    }
11240
11241    /// SwiGLU twin using Memra's host-matching expf transcription.
11242    pub fn silu_mul_host_expf(
11243        &self,
11244        gate: &CudaSlice<f32>,
11245        up: &CudaSlice<f32>,
11246        dst: &mut CudaSlice<f32>,
11247        n: usize,
11248    ) -> Result<(), Box<dyn std::error::Error>> {
11249        let f = self.func("silu_mul_host_expf_f32");
11250        let cfg = LaunchConfig::for_num_elems(n as u32);
11251        let ni = n as i32;
11252        let __s_b = self.gpu.stream();
11253        let mut b = __s_b.launch_builder(&f);
11254        b.arg(gate).arg(up).arg(dst).arg(&ni);
11255        unsafe {
11256            b.launch(cfg)?;
11257        }
11258        Ok(())
11259    }
11260
11261    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
11262    pub fn silu_clamped_mul_host_expf(
11263        &self,
11264        gate: &CudaSlice<f32>,
11265        up: &CudaSlice<f32>,
11266        limit: f32,
11267        dst: &mut CudaSlice<f32>,
11268        n: usize,
11269    ) -> Result<(), Box<dyn std::error::Error>> {
11270        if !limit.is_finite() || limit <= 0.0 {
11271            return Err(
11272                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11273            );
11274        }
11275        let f = self.func("silu_clamped_mul_host_expf_f32");
11276        let cfg = LaunchConfig::for_num_elems(n as u32);
11277        let ni = n as i32;
11278        let __s_b = self.gpu.stream();
11279        let mut b = __s_b.launch_builder(&f);
11280        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11281        unsafe {
11282            b.launch(cfg)?;
11283        }
11284        Ok(())
11285    }
11286
11287    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
11288    /// for the down projection — kills the standalone convert pass. Bit-identical class.
11289    pub fn silu_mul_f16out(
11290        &self,
11291        gate: &CudaSlice<f32>,
11292        up: &CudaSlice<f32>,
11293        dst: &mut CudaSlice<f32>,
11294        dst16: &mut CudaSlice<u8>,
11295        n: usize,
11296    ) -> Result<(), Box<dyn std::error::Error>> {
11297        let f = self.func("silu_mul_f16out_f32");
11298        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11299        let ni = n as i32;
11300        let __s_b = self.gpu.stream();
11301        let mut b = __s_b.launch_builder(&f);
11302        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11303        unsafe {
11304            b.launch(cfg)?;
11305        }
11306        Ok(())
11307    }
11308
11309    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
11310    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
11311    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
11312    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
11313    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
11314    /// launches per dense FFN layer (the gate+up post-matmul scales).
11315    pub fn silu_mul_scaled(
11316        &self,
11317        gate: &CudaSlice<f32>,
11318        up: &CudaSlice<f32>,
11319        gs: f32,
11320        us: f32,
11321        dst: &mut CudaSlice<f32>,
11322        n: usize,
11323    ) -> Result<(), Box<dyn std::error::Error>> {
11324        let f = self.func("silu_mul_scaled_f32");
11325        let cfg = LaunchConfig::for_num_elems(n as u32);
11326        let ni = n as i32;
11327        let (gsf, usf) = (gs, us);
11328        let __s_b = self.gpu.stream();
11329        let mut b = __s_b.launch_builder(&f);
11330        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11331        unsafe {
11332            b.launch(cfg)?;
11333        }
11334        Ok(())
11335    }
11336
11337    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
11338    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
11339    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
11340    #[allow(clippy::too_many_arguments)]
11341    pub fn swigluoai_mul_scaled(
11342        &self,
11343        gate: &CudaSlice<f32>,
11344        up: &CudaSlice<f32>,
11345        gs: f32,
11346        us: f32,
11347        alpha: f32,
11348        limit: f32,
11349        dst: &mut CudaSlice<f32>,
11350        n: usize,
11351    ) -> Result<(), Box<dyn std::error::Error>> {
11352        let f = self.func("swigluoai_mul_scaled_f32");
11353        let cfg = LaunchConfig::for_num_elems(n as u32);
11354        let ni = n as i32;
11355        let __s_b = self.gpu.stream();
11356        let mut b = __s_b.launch_builder(&f);
11357        b.arg(gate)
11358            .arg(up)
11359            .arg(&gs)
11360            .arg(&us)
11361            .arg(&alpha)
11362            .arg(&limit)
11363            .arg(dst)
11364            .arg(&ni);
11365        unsafe {
11366            b.launch(cfg)?;
11367        }
11368        Ok(())
11369    }
11370
11371    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11372    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11373    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11374    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11375    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11376    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11377    /// n must be a multiple of 32 (n_ff always is).
11378    pub fn silu_mul_scaled_q8_1(
11379        &self,
11380        gate: &CudaSlice<f32>,
11381        up: &CudaSlice<f32>,
11382        gs: f32,
11383        us: f32,
11384        n: usize,
11385    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11386        let f = self.func("silu_mul_scaled_q8_1");
11387        let nblk = n / 32;
11388        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11389        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11390        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11391        let cfg = LaunchConfig::for_num_elems(n as u32);
11392        let (gsf, usf, ni) = (gs, us, n as i32);
11393        let __s_b = self.gpu.stream();
11394        let mut b = __s_b.launch_builder(&f);
11395        b.arg(gate)
11396            .arg(up)
11397            .arg(&gsf)
11398            .arg(&usf)
11399            .arg(&mut aq)
11400            .arg(&mut ad)
11401            .arg(&ni);
11402        unsafe {
11403            b.launch(cfg)?;
11404        }
11405        Ok((aq, ad))
11406    }
11407
11408    pub fn add(
11409        &self,
11410        a: &CudaSlice<f32>,
11411        b_in: &CudaSlice<f32>,
11412        dst: &mut CudaSlice<f32>,
11413        n: usize,
11414    ) -> Result<(), Box<dyn std::error::Error>> {
11415        let f = self.func("add_f32");
11416        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11417        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11418        let ni = n as i32;
11419        let __s_bld = self.gpu.stream();
11420        let mut bld = __s_bld.launch_builder(&f);
11421        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11422        unsafe {
11423            bld.launch(cfg)?;
11424        }
11425        Ok(())
11426    }
11427
11428    pub fn mul(
11429        &self,
11430        a: &CudaSlice<f32>,
11431        b_in: &CudaSlice<f32>,
11432        dst: &mut CudaSlice<f32>,
11433        n: usize,
11434    ) -> Result<(), Box<dyn std::error::Error>> {
11435        let f = self.func("mul_f32");
11436        let cfg = LaunchConfig::for_num_elems(n as u32);
11437        let ni = n as i32;
11438        let __s_bld = self.gpu.stream();
11439        let mut bld = __s_bld.launch_builder(&f);
11440        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11441        unsafe {
11442            bld.launch(cfg)?;
11443        }
11444        Ok(())
11445    }
11446
11447    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11448    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11449    pub fn matmul(
11450        &self,
11451        w: &crate::model::GpuTensor,
11452        x: &CudaSlice<f32>,
11453        m: usize,
11454    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11455        use crate::model::GpuTensor;
11456        let in_f = w.in_features();
11457        let out_f = w.out_features();
11458        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11459        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11460        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11461        // gives nothing). Quantize the activation once here then call the GEMM.
11462        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11463        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11464        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11465        #[allow(non_snake_case)]
11466        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11467        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11468        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11469            usize::MAX
11470        } else {
11471            16usize
11472        };
11473
11474        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11475        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11476        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11477        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11478        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11479        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11480        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11481        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11482        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11483        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11484        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11485        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11486        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11487        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11488        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11489        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11490        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11491        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11492        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11493        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11494        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11495        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11496        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11497        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11498        if m >= GEMM_M_THRESHOLD {
11499            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11500                return Ok(y);
11501            }
11502            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
11503            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
11504            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
11505            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
11506            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
11507            // tile defaults differently by operand source.
11508            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11509                return Ok(y);
11510            }
11511            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
11512            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
11513            if let Some(y) = self.try_f16_gemm(w, x, m)? {
11514                return Ok(y);
11515            }
11516        }
11517        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
11518        // m threshold the rest of this method uses:
11519        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
11520        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
11521        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
11522        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
11523        //     across every tier by construction with no batched twin needed.
11524        //
11525        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
11526        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
11527        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
11528        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
11529        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
11530        // arms is what makes sure it never gets there.
11531        if let GpuTensor::Quant { qtype, .. } = w {
11532            if *qtype == QT_F8_E4M3_BLK {
11533                if m >= GEMM_M_THRESHOLD {
11534                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11535                        return Ok(y);
11536                    }
11537                }
11538                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11539                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11540                    return Ok(y);
11541                }
11542            }
11543        }
11544        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11545            return self.qmatvec_mmq(w, x, m);
11546        }
11547        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11548            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11549            return self.qmatvec_gemm(w, &aq, &ad, m);
11550        }
11551        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
11552        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
11553        if m >= GEMM_M_THRESHOLD {
11554            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11555                return Ok(y);
11556            }
11557        }
11558        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
11559        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
11560        // to Stage-A f32-dequant (the correctness oracle path).
11561        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11562        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
11563        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
11564        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
11565        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
11566        if m == 1 && fast {
11567            if let GpuTensor::Quant {
11568                bytes,
11569                qtype,
11570                row_bytes,
11571                rp,
11572                rp4,
11573                scale,
11574                ..
11575            } = w
11576            {
11577                if self.mmvq_supports(*qtype) {
11578                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
11579                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
11580                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
11581                    let (bytes, rp) = match rp4 {
11582                        Some(m4) => (m4, true),
11583                        None => (bytes, *rp),
11584                    };
11585                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11586                    return self.qmatvec_mmvq(
11587                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11588                    );
11589                }
11590            }
11591        }
11592        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
11593        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
11594        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
11595        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
11596        // block below. MEMRA_NO_BATCHED -> per-m path.
11597        //
11598        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
11599        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
11600        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
11601        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
11602        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
11603        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
11604        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
11605        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
11606        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
11607        if (2..=16).contains(&m)
11608            && fast
11609            && std::env::var("MEMRA_NO_BATCHED").is_err()
11610            && (m <= 4 || Self::b8_enabled())
11611        {
11612            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
11613            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
11614            // is present (rp4) — the mirror pick below then routes to the _rp family.
11615            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
11616            // because the native e4m3 row layout is already aligned and needs no mirror.
11617            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
11618            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
11619            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
11620            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
11621            let m_ok = m <= 8
11622                || matches!(w, GpuTensor::Quant { qtype, .. }
11623                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11624                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11625            if m_ok {
11626                if let GpuTensor::Quant {
11627                    bytes,
11628                    qtype,
11629                    row_bytes,
11630                    rp,
11631                    rp4,
11632                    ..
11633                } = w
11634                {
11635                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11636                        let (bytes, rp) = match rp4 {
11637                            Some(m4) => (m4, true),
11638                            None => (bytes, *rp),
11639                        };
11640                        let mcols = Self::batched_mcols(m);
11641                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11642                        let mut y = self.qmatvec_mmvq_batched(
11643                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11644                        )?;
11645                        if let GpuTensor::Quant { scale, .. } = w {
11646                            if *scale != 1.0 {
11647                                self.scale_inplace(&mut y, *scale, m * out_f)?;
11648                            }
11649                        }
11650                        return Ok(y);
11651                    }
11652                }
11653            }
11654        }
11655        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
11656        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
11657        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
11658        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
11659        // for this dtype, so the generic match below must never see it under `fast`.
11660        if fast {
11661            if let GpuTensor::Quant {
11662                bytes,
11663                qtype,
11664                row_bytes,
11665                scale,
11666                ..
11667            } = w
11668            {
11669                if *qtype == QT_F8_E4M3 {
11670                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11671                    return self.qmatvec_mmvq(
11672                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11673                    );
11674                }
11675            }
11676        }
11677        let mut y = match w {
11678            GpuTensor::Quant {
11679                bytes,
11680                qtype,
11681                row_bytes,
11682                ..
11683            } if fast && *qtype == QT_Q8_0 => {
11684                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11685            }
11686            GpuTensor::Quant {
11687                bytes,
11688                qtype,
11689                row_bytes,
11690                ..
11691            } if fast && *qtype == QT_Q4_K => {
11692                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11693            }
11694            GpuTensor::Quant {
11695                bytes,
11696                qtype,
11697                row_bytes,
11698                ..
11699            } if fast && *qtype == QT_Q6_K => {
11700                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11701            }
11702            GpuTensor::Quant {
11703                bytes,
11704                qtype,
11705                row_bytes,
11706                ..
11707            } if fast && *qtype == QT_Q5_K => {
11708                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11709            }
11710            GpuTensor::Quant {
11711                bytes,
11712                qtype,
11713                row_bytes,
11714                ..
11715            } if fast && *qtype == QT_Q3_K => {
11716                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11717            }
11718            GpuTensor::Quant {
11719                bytes,
11720                qtype,
11721                row_bytes,
11722                rp,
11723                ..
11724            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11725                if *rp {
11726                    "qmatvec_nvfp4_dp4a_rp"
11727                } else {
11728                    "qmatvec_nvfp4_dp4a"
11729                },
11730                &bytes.slice(0..bytes.len()),
11731                x,
11732                m,
11733                in_f,
11734                out_f,
11735                *row_bytes,
11736            )?,
11737            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
11738            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
11739            // anomaly (research/kat-anomaly-20260802/).
11740            GpuTensor::Quant {
11741                bytes,
11742                qtype,
11743                row_bytes,
11744                ..
11745            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11746                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11747            }
11748            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
11749            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
11750            // without first writing the matching kernel, or func() will panic
11751            // "kernel ... not in any fatbin".
11752            GpuTensor::Quant {
11753                bytes,
11754                qtype,
11755                row_bytes,
11756                rp,
11757                ..
11758            } =>
11759            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
11760            // deq(row,j) form cannot address the planes; same value/product order).
11761            {
11762                self.qmatvec(
11763                    bytes,
11764                    x,
11765                    m,
11766                    in_f,
11767                    out_f,
11768                    if *rp && *qtype == QT_NVFP4 {
11769                        QT_NVFP4_RP
11770                    } else {
11771                        *qtype
11772                    },
11773                    *row_bytes,
11774                )?
11775            }
11776            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11777            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
11778            // cuBLASLt f32 GEMV as the Float arm.
11779            GpuTensor::FloatBf16 { data, .. } => {
11780                self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11781            }
11782        };
11783        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
11784        if let GpuTensor::Quant { scale, .. } = w {
11785            if *scale != 1.0 {
11786                self.scale_inplace(&mut y, *scale, m * out_f)?;
11787            }
11788        }
11789        Ok(y)
11790    }
11791
11792    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
11793    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
11794    ///
11795    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
11796    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
11797    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
11798    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
11799    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
11800    /// path must not pay an env lookup for a flag that is off.
11801    pub fn stage_a_raw_needed() -> bool {
11802        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11803        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11804    }
11805
11806    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
11807    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
11808    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11809        use crate::model::GpuTensor;
11810        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11811            return false;
11812        }
11813        match w {
11814            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
11815            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
11816            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
11817            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
11818            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
11819            // block class has no fused twin yet, so each of its projections takes its own launch.
11820            GpuTensor::Quant { qtype, .. } => {
11821                matches!(
11822                    *qtype,
11823                    QT_Q8_0
11824                        | QT_Q4_K
11825                        | QT_Q6_K
11826                        | QT_Q5_K
11827                        | QT_Q3_K
11828                        | QT_NVFP4
11829                        | QT_F8_E4M3
11830                        | QT_F8_E4M3_BLK
11831                        | QT_Q4_0
11832                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11833            }
11834            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11835        }
11836    }
11837
11838    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
11839    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
11840    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
11841    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
11842    pub fn matmul_pre(
11843        &self,
11844        w: &crate::model::GpuTensor,
11845        aq: &CudaSlice<i8>,
11846        ad: &CudaSlice<f32>,
11847        x_fallback: &CudaSlice<f32>,
11848        m: usize,
11849    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11850        use crate::model::GpuTensor;
11851        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
11852        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
11853        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
11854        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
11855        // rc=30013 dig, 2026-07-31).
11856        let x_raw_ok = x_fallback.len() >= m * w.in_features();
11857        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
11858        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
11859        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11860            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11861                return Ok(y);
11862            }
11863            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
11864            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
11865            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11866                return Ok(y);
11867            }
11868            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
11869            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11870                return Ok(y);
11871            }
11872        }
11873        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
11874        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
11875        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
11876        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
11877        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
11878        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11879            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11880                return Ok(y);
11881            }
11882        }
11883        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11884            return Ok(y);
11885        }
11886        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
11887        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
11888        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
11889        // aq/ad.
11890        if m >= 16
11891            && w.out_features() >= 128
11892            && self.mmq_supports(w)
11893            && !self.verify_exact_on()
11894            && x_raw_ok
11895        {
11896            return self.qmatvec_mmq(w, x_fallback, m);
11897        }
11898        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
11899        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
11900        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11901            if let Some(y) =
11902                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11903            {
11904                return Ok(y);
11905            }
11906        }
11907        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
11908        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
11909        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11910            return self.qmatvec_gemm(w, aq, ad, m);
11911        }
11912        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
11913        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
11914        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
11915        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
11916        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
11917        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
11918        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
11919        // which reads `m * in_f` floats out of a 0-byte allocation ->
11920        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
11921        // it poisons the context, so every LATER request in that process fails with an unrelated
11922        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
11923        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
11924        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
11925        // dense artifact and left the arm with no working truth instrument.
11926        //
11927        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
11928        // strictly better than an illegal address surfacing later at an unrelated sync point, and
11929        // an oracle that cannot run must say so rather than corrupt the context it runs in.
11930        if !self.uses_q8_1_fast(w) {
11931            if !x_raw_ok {
11932                return Err(format!(
11933                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11934                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11935                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
11936                     activation (see Engine::rms_norm_decode, which is bit-identical to \
11937                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11938                    x_fallback.len(),
11939                    m,
11940                    w.in_features(),
11941                    m * w.in_features()
11942                )
11943                .into());
11944            }
11945            return self.matmul(w, x_fallback, m);
11946        }
11947        let in_f = w.in_features();
11948        let out_f = w.out_features();
11949        let (bytes, qtype, row_bytes, scale, rp) = match w {
11950            GpuTensor::Quant {
11951                bytes,
11952                qtype,
11953                row_bytes,
11954                scale,
11955                rp,
11956                ..
11957            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11958            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11959        };
11960        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
11961        // the dp4a/oracle tails below keep the raw GGUF bytes.
11962        let (mbytes, mrp) = match w {
11963            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11964            _ => (bytes, rp),
11965        };
11966        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
11967        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
11968        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
11969        if m == 1 && self.mmvq_supports(qtype) {
11970            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11971        }
11972        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
11973        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
11974        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
11975        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
11976        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
11977        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
11978        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
11979        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
11980        // m=5..8 on the old per-m path (b8-tier-only seam).
11981        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
11982        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
11983        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
11984        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11985            && std::env::var("MEMRA_NO_BATCHED").is_err()
11986            && (m <= 4 || Self::b8_enabled())
11987            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
11988            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
11989            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
11990            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11991                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11992        {
11993            let mcols = Self::batched_mcols(m);
11994            return self.qmatvec_mmvq_batched(
11995                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11996            );
11997        }
11998        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
11999        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
12000        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
12001        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
12002        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
12003        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12004            let (b2, r2) = if qtype == QT_Q4_0 {
12005                (mbytes, mrp)
12006            } else {
12007                (bytes, rp)
12008            };
12009            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12010        }
12011        let name = match qtype {
12012            QT_Q8_0 => "qmatvec_q8_0_dp4a",
12013            QT_Q4_K => "qmatvec_q4_K_dp4a",
12014            QT_Q6_K => "qmatvec_q6_K_dp4a",
12015            QT_Q5_K => "qmatvec_q5_K_dp4a",
12016            QT_Q3_K => "qmatvec_q3_K_dp4a",
12017            QT_NVFP4 => {
12018                if rp {
12019                    "qmatvec_nvfp4_dp4a_rp"
12020                } else {
12021                    "qmatvec_nvfp4_dp4a"
12022                }
12023            }
12024            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12025            _ => unreachable!(),
12026        };
12027        let f = self.func(name);
12028        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
12029        let cfg = LaunchConfig {
12030            grid_dim: (out_f as u32, m as u32, 1),
12031            block_dim: (128, 1, 1),
12032            shared_mem_bytes: 0,
12033        };
12034        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12035        let __s_b = self.gpu.stream();
12036        let mut b = __s_b.launch_builder(&f);
12037        b.arg(bytes)
12038            .arg(aq)
12039            .arg(ad)
12040            .arg(&mut y)
12041            .arg(&inf)
12042            .arg(&outf)
12043            .arg(&mi)
12044            .arg(&rb);
12045        unsafe {
12046            b.launch(cfg)?;
12047        }
12048        if scale != 1.0 {
12049            self.scale_inplace(&mut y, scale, m * out_f)?;
12050        }
12051        Ok(y)
12052    }
12053
12054    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
12055    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
12056    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
12057    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
12058    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
12059    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
12060    /// reduce as m=1); this method just forces that path unconditionally.
12061    pub fn matmul_decode_exact(
12062        &self,
12063        w: &crate::model::GpuTensor,
12064        x: &CudaSlice<f32>,
12065        m: usize,
12066    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12067        use crate::model::GpuTensor;
12068        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
12069        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
12070        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
12071        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
12072        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
12073        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
12074        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
12075        if let GpuTensor::Float { data, .. } = w {
12076            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12077        }
12078        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
12079        // float linear (same n-independent reduction contract as the Float arm above).
12080        if let GpuTensor::FloatBf16 { data, .. } = w {
12081            let (in_f, out_f) = (w.in_features(), w.out_features());
12082            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12083        }
12084        if !self.uses_q8_1_fast(w) {
12085            return self.matmul(w, x, m);
12086        }
12087        let in_f = w.in_features();
12088        let out_f = w.out_features();
12089        let (bytes, qtype, row_bytes, scale, rp) = match w {
12090            GpuTensor::Quant {
12091                bytes,
12092                qtype,
12093                row_bytes,
12094                scale,
12095                rp,
12096                ..
12097            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12098            _ => return self.matmul(w, x, m),
12099        };
12100        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
12101        // which does its own mirror pick).
12102        let (bytes, rp) = match w {
12103            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12104            _ => (bytes, rp),
12105        };
12106        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12107        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
12108        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
12109        // (token,row) by construction, which is exactly what this method exists to guarantee.
12110        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12111            return Ok(y);
12112        }
12113        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
12114        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
12115        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
12116        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
12117        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
12118        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
12119        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
12120        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
12121        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12122            && std::env::var("MEMRA_NO_BATCHED").is_err()
12123            && (m <= 4 || Self::b8_enabled())
12124            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
12125            // no mirror precondition, `rp` selects the layout only.
12126            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12127                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12128        {
12129            let mcols = Self::batched_mcols(m);
12130            return self.qmatvec_mmvq_batched(
12131                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12132            );
12133        }
12134        if self.mmvq_supports(qtype) {
12135            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
12136            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
12137            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12138        }
12139        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
12140        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
12141        self.matmul_pre(w, &aq, &ad, x, m)
12142    }
12143
12144    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
12145    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
12146    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
12147    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
12148    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
12149    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
12150    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
12151    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
12152    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
12153    pub fn matmul_decode_exact_pre(
12154        &self,
12155        w: &crate::model::GpuTensor,
12156        aq: &CudaSlice<i8>,
12157        ad: &CudaSlice<f32>,
12158        m: usize,
12159    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12160        use crate::model::GpuTensor;
12161        debug_assert!(
12162            self.uses_q8_1_fast(w),
12163            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12164        );
12165        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
12166        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12167            return Ok(y);
12168        }
12169        let in_f = w.in_features();
12170        let out_f = w.out_features();
12171        let (bytes, qtype, row_bytes, scale, rp) = match w {
12172            GpuTensor::Quant {
12173                bytes,
12174                qtype,
12175                row_bytes,
12176                scale,
12177                rp,
12178                ..
12179            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12180            _ => {
12181                return Err(
12182                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12183                );
12184            }
12185        };
12186        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
12187        let (bytes, rp) = match w {
12188            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12189            _ => (bytes, rp),
12190        };
12191        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
12192        if (2..=16).contains(&m)
12193            && self.batched_supports(qtype)
12194            && self.mmvq_supports(qtype)
12195            && std::env::var("MEMRA_NO_BATCHED").is_err()
12196            && (m <= 4 || Self::b8_enabled())
12197            && (m <= 8
12198                || qtype == QT_Q4_0
12199                || qtype == QT_Q6_K
12200                || qtype == QT_F8_E4M3
12201                || qtype == QT_NVFP4
12202                || qtype == QT_Q4_K
12203                || qtype == QT_Q5_K
12204                || qtype == QT_Q8_0)
12205        {
12206            let mcols = Self::batched_mcols(m);
12207            return self.qmatvec_mmvq_batched(
12208                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12209            );
12210        }
12211        if self.mmvq_supports(qtype) {
12212            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12213        }
12214        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
12215        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
12216        let x0 = self.zeros(0)?;
12217        self.matmul_pre(w, aq, ad, &x0, m)
12218    }
12219
12220    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
12221    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
12222    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
12223    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
12224    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
12225    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
12226    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
12227    /// per-tensor path.
12228    pub fn matmul_decode_exact_dual_pre(
12229        &self,
12230        w0: &crate::model::GpuTensor,
12231        w1: &crate::model::GpuTensor,
12232        aq: &CudaSlice<i8>,
12233        ad: &CudaSlice<f32>,
12234        m: usize,
12235    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12236    {
12237        use crate::model::GpuTensor;
12238        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12239        let on = *ON.get_or_init(|| {
12240            std::env::var("MEMRA_SPEC_DUAL_T")
12241                .map(|v| v != "0")
12242                .unwrap_or(true)
12243        });
12244        if !on
12245            || !(2..=7).contains(&m)
12246            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12247            || !self.uses_q8_1_fast(w0)
12248            || !self.uses_q8_1_fast(w1)
12249        {
12250            return Ok(None);
12251        }
12252        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
12253        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
12254        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
12255        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
12256        if !self.mmvq_supports(QT_NVFP4) {
12257            return Ok(None);
12258        }
12259        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12260        if w1.in_features() != in_f || w1.out_features() != out_f {
12261            return Ok(None);
12262        }
12263        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12264            (
12265                GpuTensor::Quant {
12266                    bytes: b0,
12267                    qtype: q0,
12268                    row_bytes: rb0,
12269                    scale: s0,
12270                    rp: rp0,
12271                    rp4: None,
12272                    ..
12273                },
12274                GpuTensor::Quant {
12275                    bytes: b1,
12276                    qtype: q1,
12277                    row_bytes: rb1,
12278                    scale: s1,
12279                    rp: rp1,
12280                    rp4: None,
12281                    ..
12282                },
12283            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12284                (b0, b1, *rb0, *s0, *s1, *rp0)
12285            }
12286            _ => return Ok(None),
12287        };
12288        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
12289        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
12290        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12291        {
12292            return Ok(None);
12293        }
12294        let (y0, y1) =
12295            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12296        Ok(Some(((y0, s0), (y1, s1))))
12297    }
12298
12299    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
12300    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
12301    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
12302    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
12303    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
12304    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
12305    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
12306    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
12307    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
12308    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
12309    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
12310    pub fn matmul_decode_exact_group4_pre(
12311        &self,
12312        ws: [&crate::model::GpuTensor; 4],
12313        aq: &CudaSlice<i8>,
12314        ad: &CudaSlice<f32>,
12315        m: usize,
12316    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12317        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12318        let on = *ON.get_or_init(|| {
12319            std::env::var("MEMRA_TK_GDN_GROUP")
12320                .map(|v| v != "0")
12321                .unwrap_or(true)
12322        });
12323        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12324    }
12325
12326    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
12327    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
12328    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
12329    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
12330    pub fn matmul_decode_exact_group3_pre(
12331        &self,
12332        ws: [&crate::model::GpuTensor; 3],
12333        aq: &CudaSlice<i8>,
12334        ad: &CudaSlice<f32>,
12335        m: usize,
12336    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12337        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12338        let on = *ON.get_or_init(|| {
12339            std::env::var("MEMRA_TK_FA_GROUP")
12340                .map(|v| v != "0")
12341                .unwrap_or(true)
12342        });
12343        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12344    }
12345
12346    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
12347    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
12348    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
12349    fn matmul_decode_exact_group_pre(
12350        &self,
12351        ws: &[&crate::model::GpuTensor],
12352        aq: &CudaSlice<i8>,
12353        ad: &CudaSlice<f32>,
12354        m: usize,
12355        on: bool,
12356        tag: &'static str,
12357    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12358        use crate::model::GpuTensor;
12359        if !on
12360            || !(2..=16).contains(&m)
12361            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12362            || (m > 4 && !Self::b8_enabled())
12363            || !self.mmvq_supports(QT_NVFP4)
12364            || !self.batched_supports(QT_NVFP4)
12365        {
12366            return Ok(None);
12367        }
12368        let in_f = ws[0].in_features();
12369        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12370        for w in ws {
12371            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12372                return Ok(None);
12373            }
12374            match w {
12375                GpuTensor::Quant {
12376                    bytes,
12377                    qtype,
12378                    scale,
12379                    rp: true,
12380                    rp4: None,
12381                    ..
12382                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12383                    parts.push((bytes, w.out_features(), *scale));
12384                }
12385                _ => return Ok(None),
12386            }
12387        }
12388        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12389        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12390        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12391        let mcols = if (5..=7).contains(&m) && b567 {
12392            m
12393        } else {
12394            Self::batched_mcols(m)
12395        };
12396        let kname: &'static str = match mcols {
12397            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12398            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12399            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12400            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12401            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12402            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12403            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12404            _ => return Ok(None),
12405        };
12406        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12407        // the second door's print on the slice-D battery — key the once-set by tag.
12408        if std::env::var("MEMRA_DEBUG").is_ok() {
12409            use std::sync::Mutex;
12410            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12411            let mut seen = SEEN.lock().unwrap();
12412            if !seen.contains(&tag) {
12413                seen.push(tag);
12414                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12415            }
12416        }
12417        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12418        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12419        let total: usize = parts.iter().map(|p| p.1).sum();
12420        let three = parts.len() == 3;
12421        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12422        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12423        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12424        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12425        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12426        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12427        let cfg = LaunchConfig {
12428            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12429            block_dim: (32, ROWS_PER_BLOCK, 1),
12430            shared_mem_bytes: 0,
12431        };
12432        let (inf, mi) = (in_f as i32, m as i32);
12433        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12434        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12435        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12436        let s3 = if three { 1.0f32 } else { parts[3].2 };
12437        let w3 = if three { parts[0].0 } else { parts[3].0 };
12438        let f = self.func(kname);
12439        let __s_b = self.gpu.stream();
12440        let mut b = __s_b.launch_builder(&f);
12441        b.arg(parts[0].0)
12442            .arg(parts[1].0)
12443            .arg(parts[2].0)
12444            .arg(w3)
12445            .arg(aq)
12446            .arg(ad)
12447            .arg(&mut y0)
12448            .arg(&mut y1)
12449            .arg(&mut y2)
12450            .arg(&mut y3)
12451            .arg(&inf)
12452            .arg(&n0)
12453            .arg(&n1)
12454            .arg(&n2)
12455            .arg(&n3)
12456            .arg(&mi)
12457            .arg(&s0)
12458            .arg(&s1)
12459            .arg(&s2)
12460            .arg(&s3);
12461        unsafe {
12462            b.launch(cfg)?;
12463        }
12464        Ok(Some(if three {
12465            vec![y0, y1, y2]
12466        } else {
12467            vec![y0, y1, y2, y3]
12468        }))
12469    }
12470
12471    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12472    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12473    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12474    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12475    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12476    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12477    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12478    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12479    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12480    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12481    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
12482    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
12483    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
12484    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
12485    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
12486    pub fn matmul_decode_exact_dual(
12487        &self,
12488        w0: &crate::model::GpuTensor,
12489        w1: &crate::model::GpuTensor,
12490        x: &CudaSlice<f32>,
12491        m: usize,
12492    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12493        use crate::model::GpuTensor;
12494        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12495        let on = *ON.get_or_init(|| {
12496            std::env::var("MEMRA_SPEC_DUAL_T")
12497                .map(|v| v != "0")
12498                .unwrap_or(true)
12499        });
12500        if !on
12501            || !(2..=4).contains(&m)
12502            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12503            || !self.uses_q8_1_fast(w0)
12504            || !self.uses_q8_1_fast(w1)
12505        {
12506            return Ok(None);
12507        }
12508        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
12509        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
12510        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
12511        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
12512        if !self.mmvq_supports(QT_NVFP4) {
12513            return Ok(None);
12514        }
12515        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12516        if w1.in_features() != in_f || w1.out_features() != out_f {
12517            return Ok(None);
12518        }
12519        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12520            (
12521                GpuTensor::Quant {
12522                    bytes: b0,
12523                    qtype: q0,
12524                    row_bytes: rb0,
12525                    scale: s0,
12526                    rp: rp0,
12527                    rp4: None,
12528                    ..
12529                },
12530                GpuTensor::Quant {
12531                    bytes: b1,
12532                    qtype: q1,
12533                    row_bytes: rb1,
12534                    scale: s1,
12535                    rp: rp1,
12536                    rp4: None,
12537                    ..
12538                },
12539            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12540                (b0, b1, *rb0, *s0, *s1, *rp0)
12541            }
12542            _ => return Ok(None),
12543        };
12544        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
12545        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
12546        if std::env::var("MEMRA_DEBUG").is_ok() {
12547            static ONCE: std::sync::Once = std::sync::Once::new();
12548            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12549        }
12550        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12551        let (y0, y1) =
12552            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12553        let mut y0 = y0;
12554        let mut y1 = y1;
12555        if s0 != 1.0 {
12556            self.scale_inplace(&mut y0, s0, m * out_f)?;
12557        }
12558        if s1 != 1.0 {
12559            self.scale_inplace(&mut y1, s1, m * out_f)?;
12560        }
12561        Ok(Some((y0, y1)))
12562    }
12563
12564    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
12565    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
12566    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
12567    /// twins (both buffers must be the repacked layout).
12568    #[allow(clippy::too_many_arguments)]
12569    pub fn qmatvec_batched_dual_raw(
12570        &self,
12571        b0: &CudaSlice<u8>,
12572        b1: &CudaSlice<u8>,
12573        aq: &CudaSlice<i8>,
12574        ad: &CudaSlice<f32>,
12575        m: usize,
12576        in_f: usize,
12577        out_f: usize,
12578        row_bytes: usize,
12579        rp: bool,
12580    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12581        const ROWS_PER_BLOCK: u32 = 4;
12582        let mcols = Self::batched_mcols(m);
12583        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
12584        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
12585        let tiny_rp1 = rp
12586            && mcols == 4
12587            && out_f <= 128
12588            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12589        let (name, rows_per_block) = if tiny_rp1 {
12590            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12591        } else {
12592            match (mcols, rp, m) {
12593                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12594                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12595                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12596                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12597                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12598                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12599                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12600                _ => {
12601                    return Err(
12602                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12603                    );
12604                }
12605            }
12606        };
12607        let f = self.func(name);
12608        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12609        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12610        let cfg = LaunchConfig {
12611            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12612            block_dim: (32, ROWS_PER_BLOCK, 1),
12613            shared_mem_bytes: 0,
12614        };
12615        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12616        let __s_b = self.gpu.stream();
12617        let mut b = __s_b.launch_builder(&f);
12618        b.arg(b0)
12619            .arg(b1)
12620            .arg(aq)
12621            .arg(ad)
12622            .arg(&mut y0)
12623            .arg(&mut y1)
12624            .arg(&inf)
12625            .arg(&outf)
12626            .arg(&mi)
12627            .arg(&rb);
12628        unsafe {
12629            b.launch(cfg)?;
12630        }
12631        Ok((y0, y1))
12632    }
12633
12634    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
12635    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
12636    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
12637    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
12638    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
12639    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
12640    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
12641    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
12642    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
12643    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
12644    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
12645    pub fn matmul_pre_dual_noscale(
12646        &self,
12647        w0: &crate::model::GpuTensor,
12648        w1: &crate::model::GpuTensor,
12649        aq: &CudaSlice<i8>,
12650        ad: &CudaSlice<f32>,
12651        m: usize,
12652    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12653    {
12654        use crate::model::GpuTensor;
12655        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12656            return Ok(None);
12657        }
12658        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
12659        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
12660        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
12661        // would mix dispatch families across the pair — the exact class `q8_fused_params`
12662        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
12663        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
12664        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
12665        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
12666        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
12667        if !self.mmvq_supports(QT_NVFP4) {
12668            return Ok(None);
12669        }
12670        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12671        if w1.in_features() != in_f || w1.out_features() != out_f {
12672            return Ok(None);
12673        }
12674        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
12675        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
12676        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
12677        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
12678        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
12679        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
12680        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
12681        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
12682        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
12683        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
12684        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
12685        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
12686        let no_mirror =
12687            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12688        if self.q8_ffn_fuse2_on()
12689            && no_mirror(w0)
12690            && no_mirror(w1)
12691            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12692        {
12693            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12694            return Ok(Some(((y0, 1.0), (y1, 1.0))));
12695        }
12696        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
12697        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
12698        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
12699        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
12700        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
12701        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
12702        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
12703        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
12704        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
12705        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12706            let (y0, y1) =
12707                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12708            return Ok(Some(((y0, p0.3), (y1, p1.3))));
12709        }
12710        let (b0, q0, rb0, s0, rp0) = match w0 {
12711            GpuTensor::Quant {
12712                bytes,
12713                qtype,
12714                row_bytes,
12715                scale,
12716                rp,
12717                ..
12718            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12719            _ => return Ok(None),
12720        };
12721        let (b1, q1, rb1, s1, rp1) = match w1 {
12722            GpuTensor::Quant {
12723                bytes,
12724                qtype,
12725                row_bytes,
12726                scale,
12727                rp,
12728                ..
12729            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12730            _ => return Ok(None),
12731        };
12732        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12733            return Ok(None);
12734        }
12735        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12736        const RPW: u32 = 2;
12737        let rows_per_block = ROWS_PER_BLOCK * RPW;
12738        let f = self.func(if rp0 {
12739            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12740        } else {
12741            "qmatvec_nvfp4_mmvq_dual_mr2"
12742        });
12743        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12744        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12745        let cfg = LaunchConfig {
12746            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12747            block_dim: (32, ROWS_PER_BLOCK, 1),
12748            shared_mem_bytes: 0,
12749        };
12750        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12751        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
12752        // yscale args stay 1.0 here (they exist for the single-tensor callers).
12753        let one = 1.0f32;
12754        let __s_b = self.gpu.stream();
12755        let mut b = __s_b.launch_builder(&f);
12756        b.arg(b0)
12757            .arg(b1)
12758            .arg(aq)
12759            .arg(ad)
12760            .arg(&mut y0)
12761            .arg(&mut y1)
12762            .arg(&inf)
12763            .arg(&outf)
12764            .arg(&mi)
12765            .arg(&rb)
12766            .arg(&one)
12767            .arg(&one);
12768        unsafe {
12769            b.launch(cfg)?;
12770        }
12771        Ok(Some(((y0, s0), (y1, s1))))
12772    }
12773
12774    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
12775    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
12776    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
12777    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
12778    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
12779    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
12780    /// back to the three singles.
12781    #[allow(clippy::too_many_arguments)]
12782    pub fn matmul_nvfp4_fused3(
12783        &self,
12784        w0: &crate::model::GpuTensor,
12785        w1: &crate::model::GpuTensor,
12786        w2: &crate::model::GpuTensor,
12787        aq: &CudaSlice<i8>,
12788        ad: &CudaSlice<f32>,
12789        m: usize,
12790    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12791    {
12792        use crate::model::GpuTensor;
12793        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
12794        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
12795        // verbatim, weight rows read once for all m columns, bit-identical per
12796        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
12797        // segments would re-read the weight per row" note described the grid.y=m lift,
12798        // which this twin deliberately is NOT.
12799        if !self.mmvq_supports(QT_NVFP4)
12800            || !self.uses_q8_1_fast(w0)
12801            || !self.uses_q8_1_fast(w1)
12802            || !self.uses_q8_1_fast(w2)
12803        {
12804            return Ok(None);
12805        }
12806        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
12807        // door — same family and bit-identity law as the fused4 delegate above.
12808        if (9..=16).contains(&m) {
12809            return Ok(
12810                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12811                    Some(mut ys) => {
12812                        let y2 = ys.pop().unwrap();
12813                        let y1 = ys.pop().unwrap();
12814                        let y0 = ys.pop().unwrap();
12815                        Some((y0, y1, y2))
12816                    }
12817                    None => None,
12818                },
12819            );
12820        }
12821        if !(1..=8).contains(&m) {
12822            return Ok(None);
12823        }
12824        if m > 1 {
12825            let in_f = w0.in_features();
12826            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12827                || !self.batched_supports(QT_NVFP4)
12828                || std::env::var("MEMRA_NO_BATCHED").is_ok()
12829                || (m > 4 && !Self::b8_enabled())
12830                || in_f % 512 != 0
12831                || in_f / 64 > 272
12832            {
12833                return Ok(None);
12834            }
12835        }
12836        let unpack = |w: &crate::model::GpuTensor| match w {
12837            GpuTensor::Quant {
12838                bytes,
12839                qtype,
12840                scale,
12841                rp,
12842                ..
12843            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12844            _ => None,
12845        };
12846        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12847            return Ok(None);
12848        };
12849        let in_f = w0.in_features();
12850        if w1.in_features() != in_f || w2.in_features() != in_f {
12851            return Ok(None);
12852        }
12853        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12854        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12855        const RPW: u32 = 2;
12856        let rows_pb = ROWS_PER_BLOCK * RPW;
12857        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12858        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12859        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12860        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12861        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12862        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12863        // only dereferenced for the launch-arg build inside this call.
12864        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12865        if m > 1 {
12866            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
12867            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12868                return Ok(None);
12869            }
12870            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12871            let cfg = LaunchConfig {
12872                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12873                block_dim: (32, ROWS_PER_BLOCK, 1),
12874                shared_mem_bytes: 0,
12875            };
12876            let __s_b = self.gpu.stream();
12877            let mut b = __s_b.launch_builder(&f);
12878            b.arg(b0)
12879                .arg(b1)
12880                .arg(b2)
12881                .arg(aq)
12882                .arg(ad)
12883                .arg(&mut y0)
12884                .arg(&mut y1)
12885                .arg(&mut y2)
12886                .arg(&inf)
12887                .arg(&oi0)
12888                .arg(&oi1)
12889                .arg(&oi2)
12890                .arg(&mi);
12891            unsafe {
12892                b.launch(cfg)?;
12893            }
12894            return Ok(Some((y0, y1, y2)));
12895        }
12896        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12897        let cfg = LaunchConfig {
12898            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12899            block_dim: (32, ROWS_PER_BLOCK, 1),
12900            shared_mem_bytes: 0,
12901        };
12902        let __s_b = self.gpu.stream();
12903        let mut b = __s_b.launch_builder(&f);
12904        b.arg(b0)
12905            .arg(b1)
12906            .arg(b2)
12907            .arg(aq)
12908            .arg(ad)
12909            .arg(&mut y0)
12910            .arg(&mut y1)
12911            .arg(&mut y2)
12912            .arg(&inf)
12913            .arg(&oi0)
12914            .arg(&oi1)
12915            .arg(&oi2)
12916            .arg(&mi)
12917            .arg(&p0.1)
12918            .arg(&p1.1)
12919            .arg(&p2.1);
12920        unsafe {
12921            b.launch(cfg)?;
12922        }
12923        Ok(Some((y0, y1, y2)))
12924    }
12925
12926    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
12927    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
12928    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
12929    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
12930    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
12931    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
12932    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
12933    /// same-binary interleaved A/B arm.
12934    pub fn matmul_nvfp4_fused2(
12935        &self,
12936        w0: &crate::model::GpuTensor,
12937        w1: &crate::model::GpuTensor,
12938        aq: &CudaSlice<i8>,
12939        ad: &CudaSlice<f32>,
12940        m: usize,
12941    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12942        use crate::model::GpuTensor;
12943        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12944        let off =
12945            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12946        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
12947        // read serves all m rows); the fused segments would re-read the weight per row.
12948        if off
12949            || m != 1
12950            || !self.mmvq_supports(QT_NVFP4)
12951            || !self.uses_q8_1_fast(w0)
12952            || !self.uses_q8_1_fast(w1)
12953        {
12954            return Ok(None);
12955        }
12956        let unpack = |w: &crate::model::GpuTensor| match w {
12957            GpuTensor::Quant {
12958                bytes,
12959                qtype,
12960                scale,
12961                rp,
12962                ..
12963            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12964            _ => None,
12965        };
12966        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12967            return Ok(None);
12968        };
12969        let in_f = w0.in_features();
12970        if w1.in_features() != in_f {
12971            return Ok(None);
12972        }
12973        let (o0, o1) = (w0.out_features(), w1.out_features());
12974        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12975        const RPW: u32 = 2;
12976        let rows_pb = ROWS_PER_BLOCK * RPW;
12977        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12978        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12979        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12980        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12981        let cfg = LaunchConfig {
12982            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12983            block_dim: (32, ROWS_PER_BLOCK, 1),
12984            shared_mem_bytes: 0,
12985        };
12986        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12987        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12988        // only dereferenced for the launch-arg build inside this call.
12989        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12990        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
12991        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
12992        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12993            {
12994                use cudarc::driver::{DevicePtr, DevicePtrMut};
12995                let s = &self.gpu.stream();
12996                let (pw0, _g0) = b0.device_ptr(s);
12997                let (pw1, _g1) = b1.device_ptr(s);
12998                let (paq, _g2) = aq.device_ptr(s);
12999                let (pad, _g3) = ad.device_ptr(s);
13000                let (py0, _g4) = y0.device_ptr_mut(s);
13001                let (py1, _g5) = y1.device_ptr_mut(s);
13002                let (s0, s1) = (p0.1, p1.1);
13003                let mut ps = [
13004                    &pw0 as *const _ as *mut std::ffi::c_void,
13005                    &pw1 as *const _ as *mut _,
13006                    &paq as *const _ as *mut _,
13007                    &pad as *const _ as *mut _,
13008                    &py0 as *const _ as *mut _,
13009                    &py1 as *const _ as *mut _,
13010                    &inf as *const _ as *mut _,
13011                    &oi0 as *const _ as *mut _,
13012                    &oi1 as *const _ as *mut _,
13013                    &mi as *const _ as *mut _,
13014                    &s0 as *const _ as *mut _,
13015                    &s1 as *const _ as *mut _,
13016                ];
13017                unsafe {
13018                    self.launch_pdl(
13019                        "qmatvec_nvfp4_mmvq_fused2_rp",
13020                        cfg.grid_dim,
13021                        cfg.block_dim,
13022                        &mut ps,
13023                    )?;
13024                }
13025            }
13026            return Ok(Some((y0, y1)));
13027        }
13028        let __s_b = self.gpu.stream();
13029        let mut b = __s_b.launch_builder(&f);
13030        b.arg(b0)
13031            .arg(b1)
13032            .arg(aq)
13033            .arg(ad)
13034            .arg(&mut y0)
13035            .arg(&mut y1)
13036            .arg(&inf)
13037            .arg(&oi0)
13038            .arg(&oi1)
13039            .arg(&mi)
13040            .arg(&p0.1)
13041            .arg(&p1.1);
13042        unsafe {
13043            b.launch(cfg)?;
13044        }
13045        Ok(Some((y0, y1)))
13046    }
13047
13048    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
13049    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
13050    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
13051    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
13052    pub fn matmul_nvfp4_fused2_into(
13053        &self,
13054        w0: &crate::model::GpuTensor,
13055        w1: &crate::model::GpuTensor,
13056        aq: &CudaSlice<i8>,
13057        ad: &CudaSlice<f32>,
13058        y0: &mut CudaSlice<f32>,
13059        y1: &mut CudaSlice<f32>,
13060    ) -> Result<bool, Box<dyn std::error::Error>> {
13061        use crate::model::GpuTensor;
13062        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13063        let off =
13064            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13065        if off
13066            || !self.mmvq_supports(QT_NVFP4)
13067            || !self.uses_q8_1_fast(w0)
13068            || !self.uses_q8_1_fast(w1)
13069        {
13070            return Ok(false);
13071        }
13072        let unpack = |w: &crate::model::GpuTensor| match w {
13073            GpuTensor::Quant {
13074                bytes,
13075                qtype,
13076                scale,
13077                rp,
13078                ..
13079            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13080            _ => None,
13081        };
13082        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13083            return Ok(false);
13084        };
13085        let in_f = w0.in_features();
13086        if w1.in_features() != in_f {
13087            return Ok(false);
13088        }
13089        let (o0, o1) = (w0.out_features(), w1.out_features());
13090        if y0.len() < o0 || y1.len() < o1 {
13091            return Ok(false);
13092        }
13093        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13094        const RPW: u32 = 2;
13095        let rows_pb = ROWS_PER_BLOCK * RPW;
13096        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13097        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13098        let cfg = LaunchConfig {
13099            grid_dim: (nb(o0) + nb(o1), 1, 1),
13100            block_dim: (32, ROWS_PER_BLOCK, 1),
13101            shared_mem_bytes: 0,
13102        };
13103        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13104        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13105        // only dereferenced for the launch-arg build inside this call.
13106        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13107        let __s_b = self.gpu.stream();
13108        let mut b = __s_b.launch_builder(&f);
13109        b.arg(b0)
13110            .arg(b1)
13111            .arg(aq)
13112            .arg(ad)
13113            .arg(&mut *y0)
13114            .arg(&mut *y1)
13115            .arg(&inf)
13116            .arg(&oi0)
13117            .arg(&oi1)
13118            .arg(&mi)
13119            .arg(&p0.1)
13120            .arg(&p1.1);
13121        unsafe {
13122            b.launch(cfg)?;
13123        }
13124        Ok(true)
13125    }
13126
13127    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
13128    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
13129    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
13130    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
13131    #[allow(clippy::type_complexity)]
13132    pub fn matmul_nvfp4_fused4(
13133        &self,
13134        w0: &crate::model::GpuTensor,
13135        w1: &crate::model::GpuTensor,
13136        w2: &crate::model::GpuTensor,
13137        w3: &crate::model::GpuTensor,
13138        aq: &CudaSlice<i8>,
13139        ad: &CudaSlice<f32>,
13140        m: usize,
13141    ) -> Result<
13142        Option<(
13143            CudaSlice<f32>,
13144            CudaSlice<f32>,
13145            CudaSlice<f32>,
13146            CudaSlice<f32>,
13147        )>,
13148        Box<dyn std::error::Error>,
13149    > {
13150        use crate::model::GpuTensor;
13151        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
13152        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13153        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
13154        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
13155        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
13156        // Admission mirrors the singles' batched gates below.
13157        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13158            || !self.mmvq_supports(QT_NVFP4)
13159            || !self.uses_q8_1_fast(w0)
13160            || !self.uses_q8_1_fast(w1)
13161            || !self.uses_q8_1_fast(w2)
13162            || !self.uses_q8_1_fast(w3)
13163        {
13164            return Ok(None);
13165        }
13166        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
13167        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
13168        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
13169        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
13170        if (9..=16).contains(&m) {
13171            return Ok(
13172                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13173                    Some(mut ys) => {
13174                        let y3 = ys.pop().unwrap();
13175                        let y2 = ys.pop().unwrap();
13176                        let y1 = ys.pop().unwrap();
13177                        let y0 = ys.pop().unwrap();
13178                        Some((y0, y1, y2, y3))
13179                    }
13180                    None => None,
13181                },
13182            );
13183        }
13184        if !(1..=8).contains(&m) {
13185            return Ok(None);
13186        }
13187        if m > 1 {
13188            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
13189            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
13190            let in_f = w0.in_features();
13191            if !self.batched_supports(QT_NVFP4)
13192                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13193                || (m > 4 && !Self::b8_enabled())
13194                || in_f % 512 != 0
13195                || in_f / 64 > 272
13196            {
13197                return Ok(None);
13198            }
13199        }
13200        let unpack = |w: &crate::model::GpuTensor| match w {
13201            GpuTensor::Quant {
13202                bytes,
13203                qtype,
13204                scale,
13205                rp,
13206                ..
13207            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13208            _ => None,
13209        };
13210        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13211            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13212        else {
13213            return Ok(None);
13214        };
13215        let in_f = w0.in_features();
13216        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13217            return Ok(None);
13218        }
13219        let (o0, o1, o2, o3) = (
13220            w0.out_features(),
13221            w1.out_features(),
13222            w2.out_features(),
13223            w3.out_features(),
13224        );
13225        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13226        const RPW: u32 = 2;
13227        let rows_pb = ROWS_PER_BLOCK * RPW;
13228        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13229        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13230        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13231        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13232        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13233        let (inf, oi0, oi1, oi2, oi3, mi) = (
13234            in_f as i32,
13235            o0 as i32,
13236            o1 as i32,
13237            o2 as i32,
13238            o3 as i32,
13239            m as i32,
13240        );
13241        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13242        // only dereferenced for the launch-arg build inside this call.
13243        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13244        if m > 1 {
13245            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
13246            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
13247            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13248                return Ok(None);
13249            }
13250            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13251            let cfg = LaunchConfig {
13252                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13253                block_dim: (32, ROWS_PER_BLOCK, 1),
13254                shared_mem_bytes: 0,
13255            };
13256            let __s_b = self.gpu.stream();
13257            let mut b = __s_b.launch_builder(&f);
13258            b.arg(b0)
13259                .arg(b1)
13260                .arg(b2)
13261                .arg(b3)
13262                .arg(aq)
13263                .arg(ad)
13264                .arg(&mut y0)
13265                .arg(&mut y1)
13266                .arg(&mut y2)
13267                .arg(&mut y3)
13268                .arg(&inf)
13269                .arg(&oi0)
13270                .arg(&oi1)
13271                .arg(&oi2)
13272                .arg(&oi3)
13273                .arg(&mi);
13274            unsafe {
13275                b.launch(cfg)?;
13276            }
13277            return Ok(Some((y0, y1, y2, y3)));
13278        }
13279        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13280        let cfg = LaunchConfig {
13281            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13282            block_dim: (32, ROWS_PER_BLOCK, 1),
13283            shared_mem_bytes: 0,
13284        };
13285        let __s_b = self.gpu.stream();
13286        let mut b = __s_b.launch_builder(&f);
13287        b.arg(b0)
13288            .arg(b1)
13289            .arg(b2)
13290            .arg(b3)
13291            .arg(aq)
13292            .arg(ad)
13293            .arg(&mut y0)
13294            .arg(&mut y1)
13295            .arg(&mut y2)
13296            .arg(&mut y3)
13297            .arg(&inf)
13298            .arg(&oi0)
13299            .arg(&oi1)
13300            .arg(&oi2)
13301            .arg(&oi3)
13302            .arg(&mi)
13303            .arg(&p0.1)
13304            .arg(&p1.1)
13305            .arg(&p2.1)
13306            .arg(&p3.1);
13307        unsafe {
13308            b.launch(cfg)?;
13309        }
13310        Ok(Some((y0, y1, y2, y3)))
13311    }
13312
13313    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
13314    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
13315    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
13316    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
13317    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
13318    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
13319    /// back to the per-tensor path.
13320    pub fn matmul_q8_fused2(
13321        &self,
13322        w0: &crate::model::GpuTensor,
13323        w1: &crate::model::GpuTensor,
13324        aq: &CudaSlice<i8>,
13325        ad: &CudaSlice<f32>,
13326    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13327        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
13328        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
13329        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
13330        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
13331        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
13332        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13333            return Ok(Some(self.e4m3_fused2_core(
13334                p0.0,
13335                p1.0,
13336                aq,
13337                ad,
13338                w0.in_features(),
13339                p0.1,
13340                p1.1,
13341                p0.2,
13342                p0.3,
13343                p1.3,
13344            )?));
13345        }
13346        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13347            return Ok(None);
13348        };
13349        Ok(Some(self.q8_fused2_core(
13350            p0.0,
13351            p1.0,
13352            aq,
13353            ad,
13354            w0.in_features(),
13355            p0.1,
13356            p1.1,
13357            p0.2,
13358        )?))
13359    }
13360
13361    #[allow(clippy::too_many_arguments)]
13362    fn q8_fused2_core(
13363        &self,
13364        b0: &CudaSlice<u8>,
13365        b1: &CudaSlice<u8>,
13366        aq: &CudaSlice<i8>,
13367        ad: &CudaSlice<f32>,
13368        in_f: usize,
13369        out0: usize,
13370        out1: usize,
13371        row_bytes: usize,
13372    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13373        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13374        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13375        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13376        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13377        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13378        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13379        let cfg = LaunchConfig {
13380            grid_dim: (nb0 + nb1, 1, 1),
13381            block_dim: (32, ROWS_PER_BLOCK, 1),
13382            shared_mem_bytes: 0,
13383        };
13384        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13385        let __s_b = self.gpu.stream();
13386        let mut b = __s_b.launch_builder(&f);
13387        b.arg(b0)
13388            .arg(b1)
13389            .arg(aq)
13390            .arg(ad)
13391            .arg(&mut y0)
13392            .arg(&mut y1)
13393            .arg(&inf)
13394            .arg(&o0)
13395            .arg(&o1)
13396            .arg(&rbl);
13397        unsafe {
13398            b.launch(cfg)?;
13399        }
13400        Ok((y0, y1))
13401    }
13402
13403    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13404    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13405    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13406    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13407    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13408    pub fn matmul_q8_fused2_x(
13409        &self,
13410        w0: &crate::model::GpuTensor,
13411        w1: &crate::model::GpuTensor,
13412        x: &CudaSlice<f32>,
13413    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13414        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13415            return Ok(None);
13416        }
13417        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13418            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13419            return Ok(Some(self.e4m3_fused2_core(
13420                p0.0,
13421                p1.0,
13422                &aq,
13423                &ad,
13424                w0.in_features(),
13425                p0.1,
13426                p1.1,
13427                p0.2,
13428                p0.3,
13429                p1.3,
13430            )?));
13431        }
13432        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13433            return Ok(None);
13434        };
13435        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13436        Ok(Some(self.q8_fused2_core(
13437            p0.0,
13438            p1.0,
13439            &aq,
13440            &ad,
13441            w0.in_features(),
13442            p0.1,
13443            p1.1,
13444            p0.2,
13445        )?))
13446    }
13447
13448    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13449    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13450    #[allow(clippy::too_many_arguments)]
13451    pub fn qmatvec_q8_fused2_raw(
13452        &self,
13453        b0: &CudaSlice<u8>,
13454        b1: &CudaSlice<u8>,
13455        x: &CudaSlice<f32>,
13456        in_f: usize,
13457        out0: usize,
13458        out1: usize,
13459        row_bytes: usize,
13460    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13461        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13462        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13463    }
13464
13465    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13466    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13467    /// (tensor,row) to three separate m=1 MMVQ launches.
13468    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13469    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13470    pub fn matmul_q4_fused3(
13471        &self,
13472        w0: &crate::model::GpuTensor,
13473        w1: &crate::model::GpuTensor,
13474        w2: &crate::model::GpuTensor,
13475        aq: &CudaSlice<i8>,
13476        ad: &CudaSlice<f32>,
13477    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13478    {
13479        use crate::model::GpuTensor;
13480        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13481            match w {
13482                GpuTensor::Quant {
13483                    qtype, row_bytes, ..
13484                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13485                _ => None,
13486            }
13487        };
13488        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13489            return Ok(None);
13490        };
13491        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13492            return Ok(None);
13493        }
13494        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
13495        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
13496        // the separate matvecs (each routes its own rp).
13497        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13498            match w {
13499                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13500                    Some(m) => (m, true),
13501                    None => (bytes, *rp),
13502                },
13503                _ => unreachable!(),
13504            }
13505        }
13506        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13507        if rp0 != rp1 || rp1 != rp2 {
13508            return Ok(None);
13509        }
13510        let rp = rp0;
13511        let rpb: u32 = 4;
13512        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
13513        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
13514        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
13515        let mr1 = rp && Self::q40_mr1_on();
13516        let nb = |o: usize| {
13517            if mr1 {
13518                (o as u32).div_ceil(rpb)
13519            } else {
13520                (o as u32).div_ceil(2).div_ceil(rpb)
13521            }
13522        };
13523        let grid = nb(o0) + nb(o1) + nb(o2);
13524        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13525        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13526        let mut y2 = self.alloc_uninit::<f32>(o2)?;
13527        let f = self.func(if mr1 {
13528            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13529        } else if rp {
13530            "qmatvec_q4_0_mmvq_fused3_rp"
13531        } else {
13532            "qmatvec_q4_0_mmvq_fused3"
13533        });
13534        let cfg = LaunchConfig {
13535            grid_dim: (grid, 1, 1),
13536            block_dim: (32, rpb, 1),
13537            shared_mem_bytes: 0,
13538        };
13539        let inf = w0.in_features() as i32;
13540        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13541        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13542        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
13543        // variant may take the programmatic-serialization launch.
13544        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13545            {
13546                use cudarc::driver::{DevicePtr, DevicePtrMut};
13547                let s = &self.gpu.stream();
13548                let (p0, _g0) = b0.device_ptr(s);
13549                let (p1, _g1) = b1.device_ptr(s);
13550                let (p2, _g2) = b2.device_ptr(s);
13551                let (paq, _g3) = aq.device_ptr(s);
13552                let (pad, _g4) = ad.device_ptr(s);
13553                let (py0, _g5) = y0.device_ptr_mut(s);
13554                let (py1, _g6) = y1.device_ptr_mut(s);
13555                let (py2, _g7) = y2.device_ptr_mut(s);
13556                let mut ps = [
13557                    &p0 as *const _ as *mut std::ffi::c_void,
13558                    &p1 as *const _ as *mut _,
13559                    &p2 as *const _ as *mut _,
13560                    &paq as *const _ as *mut _,
13561                    &pad as *const _ as *mut _,
13562                    &py0 as *const _ as *mut _,
13563                    &py1 as *const _ as *mut _,
13564                    &py2 as *const _ as *mut _,
13565                    &inf as *const _ as *mut _,
13566                    &oo0 as *const _ as *mut _,
13567                    &oo1 as *const _ as *mut _,
13568                    &oo2 as *const _ as *mut _,
13569                    &r0 as *const _ as *mut _,
13570                    &r1 as *const _ as *mut _,
13571                    &r2 as *const _ as *mut _,
13572                ];
13573                unsafe {
13574                    self.launch_pdl(
13575                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13576                        (grid, 1, 1),
13577                        (32, rpb, 1),
13578                        &mut ps,
13579                    )?;
13580                }
13581            }
13582            return Ok(Some((y0, y1, y2)));
13583        }
13584        let __s_b = self.gpu.stream();
13585        let mut b = __s_b.launch_builder(&f);
13586        b.arg(b0)
13587            .arg(b1)
13588            .arg(b2)
13589            .arg(aq)
13590            .arg(ad)
13591            .arg(&mut y0)
13592            .arg(&mut y1)
13593            .arg(&mut y2)
13594            .arg(&inf)
13595            .arg(&oo0)
13596            .arg(&oo1)
13597            .arg(&oo2)
13598            .arg(&r0)
13599            .arg(&r1)
13600            .arg(&r2);
13601        unsafe {
13602            b.launch(cfg)?;
13603        }
13604        Ok(Some((y0, y1, y2)))
13605    }
13606
13607    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13608    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
13609    #[allow(clippy::too_many_arguments)]
13610    pub fn matmul_q4_fused3_into(
13611        &self,
13612        w0: &crate::model::GpuTensor,
13613        w1: &crate::model::GpuTensor,
13614        w2: &crate::model::GpuTensor,
13615        aq: &CudaSlice<i8>,
13616        ad: &CudaSlice<f32>,
13617        y0: &mut CudaSlice<f32>,
13618        y1: &mut CudaSlice<f32>,
13619        y2: &mut CudaSlice<f32>,
13620    ) -> Result<bool, Box<dyn std::error::Error>> {
13621        use crate::model::GpuTensor;
13622        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13623            match w {
13624                GpuTensor::Quant {
13625                    qtype, row_bytes, ..
13626                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13627                _ => None,
13628            }
13629        };
13630        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13631            return Ok(false);
13632        };
13633        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13634            return Ok(false);
13635        }
13636        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13637            match w {
13638                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13639                    Some(m) => (m, true),
13640                    None => (bytes, *rp),
13641                },
13642                _ => unreachable!(),
13643            }
13644        }
13645        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13646        if rp0 != rp1 || rp1 != rp2 {
13647            return Ok(false);
13648        }
13649        let rp = rp0;
13650        let rpb: u32 = 4;
13651        let mr1 = rp && Self::q40_mr1_on();
13652        let nb = |o: usize| {
13653            if mr1 {
13654                (o as u32).div_ceil(rpb)
13655            } else {
13656                (o as u32).div_ceil(2).div_ceil(rpb)
13657            }
13658        };
13659        let grid = nb(o0) + nb(o1) + nb(o2);
13660        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13661        let f = self.func(if mr1 {
13662            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13663        } else if rp {
13664            "qmatvec_q4_0_mmvq_fused3_rp"
13665        } else {
13666            "qmatvec_q4_0_mmvq_fused3"
13667        });
13668        let cfg = LaunchConfig {
13669            grid_dim: (grid, 1, 1),
13670            block_dim: (32, rpb, 1),
13671            shared_mem_bytes: 0,
13672        };
13673        let inf = w0.in_features() as i32;
13674        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13675        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13676        // PDL wave-A: identical to the owned twin (capture-lane parity).
13677        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13678            use cudarc::driver::{DevicePtr, DevicePtrMut};
13679            let s = &self.gpu.stream();
13680            let (p0, _g0) = b0.device_ptr(s);
13681            let (p1, _g1) = b1.device_ptr(s);
13682            let (p2, _g2) = b2.device_ptr(s);
13683            let (paq, _g3) = aq.device_ptr(s);
13684            let (pad, _g4) = ad.device_ptr(s);
13685            let (py0, _g5) = y0.device_ptr_mut(s);
13686            let (py1, _g6) = y1.device_ptr_mut(s);
13687            let (py2, _g7) = y2.device_ptr_mut(s);
13688            let mut ps = [
13689                &p0 as *const _ as *mut std::ffi::c_void,
13690                &p1 as *const _ as *mut _,
13691                &p2 as *const _ as *mut _,
13692                &paq as *const _ as *mut _,
13693                &pad as *const _ as *mut _,
13694                &py0 as *const _ as *mut _,
13695                &py1 as *const _ as *mut _,
13696                &py2 as *const _ as *mut _,
13697                &inf as *const _ as *mut _,
13698                &oo0 as *const _ as *mut _,
13699                &oo1 as *const _ as *mut _,
13700                &oo2 as *const _ as *mut _,
13701                &r0 as *const _ as *mut _,
13702                &r1 as *const _ as *mut _,
13703                &r2 as *const _ as *mut _,
13704            ];
13705            unsafe {
13706                self.launch_pdl(
13707                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13708                    (grid, 1, 1),
13709                    (32, rpb, 1),
13710                    &mut ps,
13711                )?;
13712            }
13713            return Ok(true);
13714        }
13715        let __s_b = self.gpu.stream();
13716        let mut b = __s_b.launch_builder(&f);
13717        b.arg(b0)
13718            .arg(b1)
13719            .arg(b2)
13720            .arg(aq)
13721            .arg(ad)
13722            .arg(&mut *y0)
13723            .arg(&mut *y1)
13724            .arg(&mut *y2)
13725            .arg(&inf)
13726            .arg(&oo0)
13727            .arg(&oo1)
13728            .arg(&oo2)
13729            .arg(&r0)
13730            .arg(&r1)
13731            .arg(&r2);
13732        unsafe {
13733            b.launch(cfg)?;
13734        }
13735        Ok(true)
13736    }
13737
13738    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
13739    pub fn matmul_q4_fused2(
13740        &self,
13741        w0: &crate::model::GpuTensor,
13742        w1: &crate::model::GpuTensor,
13743        aq: &CudaSlice<i8>,
13744        ad: &CudaSlice<f32>,
13745    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13746        use crate::model::GpuTensor;
13747        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13748            match w {
13749                GpuTensor::Quant {
13750                    qtype, row_bytes, ..
13751                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13752                _ => None,
13753            }
13754        };
13755        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13756            return Ok(None);
13757        };
13758        if w0.in_features() != w1.in_features() {
13759            return Ok(None);
13760        }
13761        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
13762        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13763            match w {
13764                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13765                    Some(m) => (m, true),
13766                    None => (bytes, *rp),
13767                },
13768                _ => unreachable!(),
13769            }
13770        }
13771        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13772        if rp0 != rp1 {
13773            return Ok(None);
13774        }
13775        let rp = rp0;
13776        let rpb: u32 = 4;
13777        // mr1 twin — see matmul_q4_fused3.
13778        let mr1 = rp && Self::q40_mr1_on();
13779        let nb = |o: usize| {
13780            if mr1 {
13781                (o as u32).div_ceil(rpb)
13782            } else {
13783                (o as u32).div_ceil(2).div_ceil(rpb)
13784            }
13785        };
13786        let grid = nb(o0) + nb(o1);
13787        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13788        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13789        let f = self.func(if mr1 {
13790            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13791        } else if rp {
13792            "qmatvec_q4_0_mmvq_fused2_rp"
13793        } else {
13794            "qmatvec_q4_0_mmvq_fused2"
13795        });
13796        let cfg = LaunchConfig {
13797            grid_dim: (grid, 1, 1),
13798            block_dim: (32, rpb, 1),
13799            shared_mem_bytes: 0,
13800        };
13801        let inf = w0.in_features() as i32;
13802        let (oo0, oo1) = (o0 as i32, o1 as i32);
13803        let (r0, r1) = (rb0 as i64, rb1 as i64);
13804        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
13805        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13806            {
13807                use cudarc::driver::{DevicePtr, DevicePtrMut};
13808                let s = &self.gpu.stream();
13809                let (p0, _g0) = b0.device_ptr(s);
13810                let (p1, _g1) = b1.device_ptr(s);
13811                let (paq, _g2) = aq.device_ptr(s);
13812                let (pad, _g3) = ad.device_ptr(s);
13813                let (py0, _g4) = y0.device_ptr_mut(s);
13814                let (py1, _g5) = y1.device_ptr_mut(s);
13815                let mut ps = [
13816                    &p0 as *const _ as *mut std::ffi::c_void,
13817                    &p1 as *const _ as *mut _,
13818                    &paq as *const _ as *mut _,
13819                    &pad as *const _ as *mut _,
13820                    &py0 as *const _ as *mut _,
13821                    &py1 as *const _ as *mut _,
13822                    &inf as *const _ as *mut _,
13823                    &oo0 as *const _ as *mut _,
13824                    &oo1 as *const _ as *mut _,
13825                    &r0 as *const _ as *mut _,
13826                    &r1 as *const _ as *mut _,
13827                ];
13828                unsafe {
13829                    self.launch_pdl(
13830                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13831                        (grid, 1, 1),
13832                        (32, rpb, 1),
13833                        &mut ps,
13834                    )?;
13835                }
13836            }
13837            return Ok(Some((y0, y1)));
13838        }
13839        let __s_b = self.gpu.stream();
13840        let mut b = __s_b.launch_builder(&f);
13841        b.arg(b0)
13842            .arg(b1)
13843            .arg(aq)
13844            .arg(ad)
13845            .arg(&mut y0)
13846            .arg(&mut y1)
13847            .arg(&inf)
13848            .arg(&oo0)
13849            .arg(&oo1)
13850            .arg(&r0)
13851            .arg(&r1);
13852        unsafe {
13853            b.launch(cfg)?;
13854        }
13855        Ok(Some((y0, y1)))
13856    }
13857
13858    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13859    pub fn matmul_q4_fused2_into(
13860        &self,
13861        w0: &crate::model::GpuTensor,
13862        w1: &crate::model::GpuTensor,
13863        aq: &CudaSlice<i8>,
13864        ad: &CudaSlice<f32>,
13865        y0: &mut CudaSlice<f32>,
13866        y1: &mut CudaSlice<f32>,
13867    ) -> Result<bool, Box<dyn std::error::Error>> {
13868        use crate::model::GpuTensor;
13869        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13870            match w {
13871                GpuTensor::Quant {
13872                    qtype, row_bytes, ..
13873                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13874                _ => None,
13875            }
13876        };
13877        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13878            return Ok(false);
13879        };
13880        if w0.in_features() != w1.in_features() {
13881            return Ok(false);
13882        }
13883        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13884            match w {
13885                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13886                    Some(m) => (m, true),
13887                    None => (bytes, *rp),
13888                },
13889                _ => unreachable!(),
13890            }
13891        }
13892        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13893        if rp0 != rp1 {
13894            return Ok(false);
13895        }
13896        let rp = rp0;
13897        let rpb: u32 = 4;
13898        let mr1 = rp && Self::q40_mr1_on();
13899        let nb = |o: usize| {
13900            if mr1 {
13901                (o as u32).div_ceil(rpb)
13902            } else {
13903                (o as u32).div_ceil(2).div_ceil(rpb)
13904            }
13905        };
13906        let grid = nb(o0) + nb(o1);
13907        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13908        let f = self.func(if mr1 {
13909            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13910        } else if rp {
13911            "qmatvec_q4_0_mmvq_fused2_rp"
13912        } else {
13913            "qmatvec_q4_0_mmvq_fused2"
13914        });
13915        let cfg = LaunchConfig {
13916            grid_dim: (grid, 1, 1),
13917            block_dim: (32, rpb, 1),
13918            shared_mem_bytes: 0,
13919        };
13920        let inf = w0.in_features() as i32;
13921        let (oo0, oo1) = (o0 as i32, o1 as i32);
13922        let (r0, r1) = (rb0 as i64, rb1 as i64);
13923        // PDL wave-A: identical to the owned twin (capture-lane parity).
13924        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13925            use cudarc::driver::{DevicePtr, DevicePtrMut};
13926            let s = &self.gpu.stream();
13927            let (p0, _g0) = b0.device_ptr(s);
13928            let (p1, _g1) = b1.device_ptr(s);
13929            let (paq, _g2) = aq.device_ptr(s);
13930            let (pad, _g3) = ad.device_ptr(s);
13931            let (py0, _g4) = y0.device_ptr_mut(s);
13932            let (py1, _g5) = y1.device_ptr_mut(s);
13933            let mut ps = [
13934                &p0 as *const _ as *mut std::ffi::c_void,
13935                &p1 as *const _ as *mut _,
13936                &paq as *const _ as *mut _,
13937                &pad as *const _ as *mut _,
13938                &py0 as *const _ as *mut _,
13939                &py1 as *const _ as *mut _,
13940                &inf as *const _ as *mut _,
13941                &oo0 as *const _ as *mut _,
13942                &oo1 as *const _ as *mut _,
13943                &r0 as *const _ as *mut _,
13944                &r1 as *const _ as *mut _,
13945            ];
13946            unsafe {
13947                self.launch_pdl(
13948                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13949                    (grid, 1, 1),
13950                    (32, rpb, 1),
13951                    &mut ps,
13952                )?;
13953            }
13954            return Ok(true);
13955        }
13956        let __s_b = self.gpu.stream();
13957        let mut b = __s_b.launch_builder(&f);
13958        b.arg(b0)
13959            .arg(b1)
13960            .arg(aq)
13961            .arg(ad)
13962            .arg(&mut *y0)
13963            .arg(&mut *y1)
13964            .arg(&inf)
13965            .arg(&oo0)
13966            .arg(&oo1)
13967            .arg(&r0)
13968            .arg(&r1);
13969        unsafe {
13970            b.launch(cfg)?;
13971        }
13972        Ok(true)
13973    }
13974
13975    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
13976    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
13977    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
13978    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
13979    pub fn matmul_q4_fused2_batched(
13980        &self,
13981        w0: &crate::model::GpuTensor,
13982        w1: &crate::model::GpuTensor,
13983        aq: &CudaSlice<i8>,
13984        ad: &CudaSlice<f32>,
13985        m: usize,
13986    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13987        use crate::model::GpuTensor;
13988        if m < 2 || m > 8 {
13989            return Ok(None);
13990        }
13991        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13992            match w {
13993                GpuTensor::Quant {
13994                    qtype, row_bytes, ..
13995                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13996                _ => None,
13997            }
13998        };
13999        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14000            return Ok(None);
14001        };
14002        if w0.in_features() != w1.in_features() {
14003            return Ok(None);
14004        }
14005        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14006            match w {
14007                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14008                    Some(mr) => (mr, true),
14009                    None => (bytes, *rp),
14010                },
14011                _ => unreachable!(),
14012            }
14013        }
14014        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14015        if !rp0 || !rp1 {
14016            return Ok(None);
14017        }
14018        let mcols = Self::batched_mcols(m);
14019        let rpb: u32 = 4;
14020        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14021        let grid = nb(o0) + nb(o1);
14022        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14023        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14024        let f = self.func(match mcols {
14025            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14026            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14027            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14028        });
14029        let cfg = LaunchConfig {
14030            grid_dim: (grid, 1, 1),
14031            block_dim: (32, rpb, 1),
14032            shared_mem_bytes: 0,
14033        };
14034        let inf = w0.in_features() as i32;
14035        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14036        let rb = rb0 as i64;
14037        let __s_b = self.gpu.stream();
14038        let mut b = __s_b.launch_builder(&f);
14039        b.arg(b0)
14040            .arg(b1)
14041            .arg(aq)
14042            .arg(ad)
14043            .arg(&mut y0)
14044            .arg(&mut y1)
14045            .arg(&inf)
14046            .arg(&oo0)
14047            .arg(&oo1)
14048            .arg(&mi)
14049            .arg(&rb);
14050        unsafe {
14051            b.launch(cfg)?;
14052        }
14053        Ok(Some((y0, y1)))
14054    }
14055
14056    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
14057    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
14058    #[allow(clippy::too_many_arguments)]
14059    pub fn matmul_q4_fused3_batched(
14060        &self,
14061        w0: &crate::model::GpuTensor,
14062        w1: &crate::model::GpuTensor,
14063        w2: &crate::model::GpuTensor,
14064        aq: &CudaSlice<i8>,
14065        ad: &CudaSlice<f32>,
14066        m: usize,
14067    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14068    {
14069        use crate::model::GpuTensor;
14070        if m < 2 || m > 8 {
14071            return Ok(None);
14072        }
14073        let q4 = |w: &GpuTensor| -> Option<usize> {
14074            match w {
14075                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14076                _ => None,
14077            }
14078        };
14079        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14080            return Ok(None);
14081        };
14082        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14083            return Ok(None);
14084        }
14085        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14086            match w {
14087                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14088                    Some(mr) => (mr, true),
14089                    None => (bytes, *rp),
14090                },
14091                _ => unreachable!(),
14092            }
14093        }
14094        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14095        if !rp0 || !rp1 || !rp2 {
14096            return Ok(None);
14097        }
14098        let mcols = Self::batched_mcols(m);
14099        let rpb: u32 = 4;
14100        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14101        let grid = nb(o0) + nb(o1) + nb(o2);
14102        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14103        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14104        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14105        let f = self.func(match mcols {
14106            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14107            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14108            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14109        });
14110        let cfg = LaunchConfig {
14111            grid_dim: (grid, 1, 1),
14112            block_dim: (32, rpb, 1),
14113            shared_mem_bytes: 0,
14114        };
14115        let inf = w0.in_features() as i32;
14116        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14117        let rb = 0i64;
14118        let __s_b = self.gpu.stream();
14119        let mut b = __s_b.launch_builder(&f);
14120        b.arg(b0)
14121            .arg(b1)
14122            .arg(b2)
14123            .arg(aq)
14124            .arg(ad)
14125            .arg(&mut y0)
14126            .arg(&mut y1)
14127            .arg(&mut y2)
14128            .arg(&inf)
14129            .arg(&oo0)
14130            .arg(&oo1)
14131            .arg(&oo2)
14132            .arg(&mi)
14133            .arg(&rb);
14134        unsafe {
14135            b.launch(cfg)?;
14136        }
14137        Ok(Some((y0, y1, y2)))
14138    }
14139
14140    pub fn matmul_q8_fused3(
14141        &self,
14142        w0: &crate::model::GpuTensor,
14143        w1: &crate::model::GpuTensor,
14144        w2: &crate::model::GpuTensor,
14145        aq: &CudaSlice<i8>,
14146        ad: &CudaSlice<f32>,
14147    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14148    {
14149        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
14150        // are per-tensor FP8, so native residency without this arm meant three separate launches.
14151        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14152            return Ok(Some(self.e4m3_fused3_core(
14153                p0.0,
14154                p1.0,
14155                p2.0,
14156                aq,
14157                ad,
14158                w0.in_features(),
14159                p0.1,
14160                p1.1,
14161                p2.1,
14162                p0.2,
14163                p0.3,
14164                p1.3,
14165                p2.3,
14166            )?));
14167        }
14168        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14169            return Ok(None);
14170        };
14171        Ok(Some(self.q8_fused3_core(
14172            p0.0,
14173            p1.0,
14174            p2.0,
14175            aq,
14176            ad,
14177            w0.in_features(),
14178            p0.1,
14179            p1.1,
14180            p2.1,
14181            p0.2,
14182        )?))
14183    }
14184
14185    #[allow(clippy::too_many_arguments)]
14186    fn q8_fused3_core(
14187        &self,
14188        b0: &CudaSlice<u8>,
14189        b1: &CudaSlice<u8>,
14190        b2: &CudaSlice<u8>,
14191        aq: &CudaSlice<i8>,
14192        ad: &CudaSlice<f32>,
14193        in_f: usize,
14194        out0: usize,
14195        out1: usize,
14196        out2: usize,
14197        row_bytes: usize,
14198    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14199        const ROWS_PER_BLOCK: u32 = 4;
14200        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14201        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14202        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14203        let f = self.func("qmatvec_q8_0_mmvq_fused3");
14204        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14205        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14206        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14207        let cfg = LaunchConfig {
14208            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14209            block_dim: (32, ROWS_PER_BLOCK, 1),
14210            shared_mem_bytes: 0,
14211        };
14212        let (inf, o0, o1, o2, rbl) = (
14213            in_f as i32,
14214            out0 as i32,
14215            out1 as i32,
14216            out2 as i32,
14217            row_bytes as i64,
14218        );
14219        let __s_b = self.gpu.stream();
14220        let mut b = __s_b.launch_builder(&f);
14221        b.arg(b0)
14222            .arg(b1)
14223            .arg(b2)
14224            .arg(aq)
14225            .arg(ad)
14226            .arg(&mut y0)
14227            .arg(&mut y1)
14228            .arg(&mut y2)
14229            .arg(&inf)
14230            .arg(&o0)
14231            .arg(&o1)
14232            .arg(&o2)
14233            .arg(&rbl);
14234        unsafe {
14235            b.launch(cfg)?;
14236        }
14237        Ok((y0, y1, y2))
14238    }
14239
14240    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
14241    #[allow(clippy::too_many_arguments)]
14242    pub fn qmatvec_q8_fused3_raw(
14243        &self,
14244        b0: &CudaSlice<u8>,
14245        b1: &CudaSlice<u8>,
14246        b2: &CudaSlice<u8>,
14247        x: &CudaSlice<f32>,
14248        in_f: usize,
14249        out0: usize,
14250        out1: usize,
14251        out2: usize,
14252        row_bytes: usize,
14253    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14254        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14255        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14256    }
14257
14258    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
14259    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
14260    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
14261    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
14262    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
14263    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
14264    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
14265    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
14266    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
14267    /// twin must not introduce a batched program the reference path would not run).
14268    pub fn matmul_q8_fused2_t(
14269        &self,
14270        w0: &crate::model::GpuTensor,
14271        w1: &crate::model::GpuTensor,
14272        aq: &CudaSlice<i8>,
14273        ad: &CudaSlice<f32>,
14274        m: usize,
14275    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14276        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
14277        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
14278        // fuses too — same template body, still bit-identical to the two _b8 launches.
14279        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14280            return Ok(None);
14281        }
14282        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
14283        // so the fused b8 launch would introduce a batched program the reference path would not run.
14284        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14285            if m > 4 && !Self::b8_enabled() {
14286                return Ok(None);
14287            }
14288            return Ok(Some(self.e4m3_fused2_t_core(
14289                p0.0,
14290                p1.0,
14291                aq,
14292                ad,
14293                m,
14294                w0.in_features(),
14295                p0.1,
14296                p1.1,
14297                p0.2,
14298                p0.3,
14299                p1.3,
14300            )?));
14301        }
14302        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14303            return Ok(None);
14304        };
14305        Ok(Some(self.q8_fused2_t_core(
14306            p0.0,
14307            p1.0,
14308            aq,
14309            ad,
14310            m,
14311            w0.in_features(),
14312            p0.1,
14313            p1.1,
14314            p0.2,
14315        )?))
14316    }
14317
14318    #[allow(clippy::too_many_arguments)]
14319    fn q8_fused2_t_core(
14320        &self,
14321        b0: &CudaSlice<u8>,
14322        b1: &CudaSlice<u8>,
14323        aq: &CudaSlice<i8>,
14324        ad: &CudaSlice<f32>,
14325        m: usize,
14326        in_f: usize,
14327        out0: usize,
14328        out1: usize,
14329        row_bytes: usize,
14330    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14331        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14332        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14333        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14334        let f = self.func(match Self::batched_mcols(m) {
14335            2 => "qmatvec_q8_0_mmvq_fused2_b2",
14336            4 => "qmatvec_q8_0_mmvq_fused2_b4",
14337            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
14338            _ => "qmatvec_q8_0_mmvq_fused2_b8",
14339        });
14340        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14341        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14342        let cfg = LaunchConfig {
14343            grid_dim: (nb0 + nb1, 1, 1),
14344            block_dim: (32, ROWS_PER_BLOCK, 1),
14345            shared_mem_bytes: 0,
14346        };
14347        let (inf, o0, o1, mi, rbl) = (
14348            in_f as i32,
14349            out0 as i32,
14350            out1 as i32,
14351            m as i32,
14352            row_bytes as i64,
14353        );
14354        let __s_b = self.gpu.stream();
14355        let mut b = __s_b.launch_builder(&f);
14356        b.arg(b0)
14357            .arg(b1)
14358            .arg(aq)
14359            .arg(ad)
14360            .arg(&mut y0)
14361            .arg(&mut y1)
14362            .arg(&inf)
14363            .arg(&o0)
14364            .arg(&o1)
14365            .arg(&mi)
14366            .arg(&rbl);
14367        unsafe {
14368            b.launch(cfg)?;
14369        }
14370        Ok((y0, y1))
14371    }
14372
14373    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14374    /// q8_1 quant of the [m, in_f] activation), no env gating.
14375    #[allow(clippy::too_many_arguments)]
14376    pub fn qmatvec_q8_fused2_t_raw(
14377        &self,
14378        b0: &CudaSlice<u8>,
14379        b1: &CudaSlice<u8>,
14380        x: &CudaSlice<f32>,
14381        m: usize,
14382        in_f: usize,
14383        out0: usize,
14384        out1: usize,
14385        row_bytes: usize,
14386    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14387        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14388        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14389    }
14390
14391    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14392    /// `matmul_q8_fused2_t` with three ranges.
14393    #[allow(clippy::too_many_arguments)]
14394    pub fn matmul_q8_fused3_t(
14395        &self,
14396        w0: &crate::model::GpuTensor,
14397        w1: &crate::model::GpuTensor,
14398        w2: &crate::model::GpuTensor,
14399        aq: &CudaSlice<i8>,
14400        ad: &CudaSlice<f32>,
14401        m: usize,
14402    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14403    {
14404        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14405            return Ok(None);
14406        }
14407        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14408            return Ok(Some(self.e4m3_fused3_t_core(
14409                p0.0,
14410                p1.0,
14411                p2.0,
14412                aq,
14413                ad,
14414                m,
14415                w0.in_features(),
14416                p0.1,
14417                p1.1,
14418                p2.1,
14419                p0.2,
14420                p0.3,
14421                p1.3,
14422                p2.3,
14423            )?));
14424        }
14425        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14426            return Ok(None);
14427        };
14428        Ok(Some(self.q8_fused3_t_core(
14429            p0.0,
14430            p1.0,
14431            p2.0,
14432            aq,
14433            ad,
14434            m,
14435            w0.in_features(),
14436            p0.1,
14437            p1.1,
14438            p2.1,
14439            p0.2,
14440        )?))
14441    }
14442
14443    #[allow(clippy::too_many_arguments)]
14444    fn q8_fused3_t_core(
14445        &self,
14446        b0: &CudaSlice<u8>,
14447        b1: &CudaSlice<u8>,
14448        b2: &CudaSlice<u8>,
14449        aq: &CudaSlice<i8>,
14450        ad: &CudaSlice<f32>,
14451        m: usize,
14452        in_f: usize,
14453        out0: usize,
14454        out1: usize,
14455        out2: usize,
14456        row_bytes: usize,
14457    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14458        const ROWS_PER_BLOCK: u32 = 4;
14459        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14460        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14461        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14462        let f = self.func(if Self::batched_mcols(m) == 2 {
14463            "qmatvec_q8_0_mmvq_fused3_b2"
14464        } else {
14465            "qmatvec_q8_0_mmvq_fused3_b4"
14466        });
14467        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14468        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14469        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14470        let cfg = LaunchConfig {
14471            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14472            block_dim: (32, ROWS_PER_BLOCK, 1),
14473            shared_mem_bytes: 0,
14474        };
14475        let (inf, o0, o1, o2, mi, rbl) = (
14476            in_f as i32,
14477            out0 as i32,
14478            out1 as i32,
14479            out2 as i32,
14480            m as i32,
14481            row_bytes as i64,
14482        );
14483        let __s_b = self.gpu.stream();
14484        let mut b = __s_b.launch_builder(&f);
14485        b.arg(b0)
14486            .arg(b1)
14487            .arg(b2)
14488            .arg(aq)
14489            .arg(ad)
14490            .arg(&mut y0)
14491            .arg(&mut y1)
14492            .arg(&mut y2)
14493            .arg(&inf)
14494            .arg(&o0)
14495            .arg(&o1)
14496            .arg(&o2)
14497            .arg(&mi)
14498            .arg(&rbl);
14499        unsafe {
14500            b.launch(cfg)?;
14501        }
14502        Ok((y0, y1, y2))
14503    }
14504
14505    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
14506    #[allow(clippy::too_many_arguments)]
14507    pub fn qmatvec_q8_fused3_t_raw(
14508        &self,
14509        b0: &CudaSlice<u8>,
14510        b1: &CudaSlice<u8>,
14511        b2: &CudaSlice<u8>,
14512        x: &CudaSlice<f32>,
14513        m: usize,
14514        in_f: usize,
14515        out0: usize,
14516        out1: usize,
14517        out2: usize,
14518        row_bytes: usize,
14519    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14520        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14521        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14522    }
14523
14524    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
14525    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
14526    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
14527    pub fn q8_ffn_fuse2_on(&self) -> bool {
14528        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14529        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14530    }
14531
14532    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
14533    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
14534    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
14535    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
14536    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
14537    #[allow(clippy::type_complexity)]
14538    fn q8_fused_params<'w, const N: usize>(
14539        &self,
14540        ws: &[&'w crate::model::GpuTensor; N],
14541    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14542        use crate::model::GpuTensor;
14543        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14544            return None;
14545        }
14546        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14547            return None;
14548        }
14549        let in_f = ws[0].in_features();
14550        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14551        for (i, w) in ws.iter().enumerate() {
14552            match w {
14553                GpuTensor::Quant {
14554                    bytes,
14555                    qtype,
14556                    row_bytes,
14557                    scale,
14558                    ..
14559                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14560                    out[i] = Some((bytes, w.out_features(), *row_bytes))
14561                }
14562                _ => return None,
14563            }
14564        }
14565        Some(out.map(|o| o.unwrap()))
14566    }
14567
14568    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
14569    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
14570    pub fn e4m3_dual_on(&self) -> bool {
14571        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14572        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14573    }
14574
14575    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
14576    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
14577    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
14578    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
14579    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
14580    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
14581    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
14582    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
14583    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
14584    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
14585    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
14586    #[allow(clippy::type_complexity)]
14587    fn e4m3_fused_params<'w, const N: usize>(
14588        &self,
14589        ws: &[&'w crate::model::GpuTensor; N],
14590    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14591        use crate::model::GpuTensor;
14592        if !self.e4m3_dual_on() {
14593            return None;
14594        }
14595        let in_f = ws[0].in_features();
14596        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14597        for (i, w) in ws.iter().enumerate() {
14598            match w {
14599                GpuTensor::Quant {
14600                    bytes,
14601                    qtype,
14602                    row_bytes,
14603                    scale,
14604                    rp,
14605                    rp4,
14606                    ..
14607                } if *qtype == QT_F8_E4M3
14608                    && w.in_features() == in_f
14609                    && *row_bytes == in_f
14610                    && !*rp
14611                    && rp4.is_none() =>
14612                {
14613                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14614                }
14615                _ => return None,
14616            }
14617        }
14618        Some(out.map(|o| o.unwrap()))
14619    }
14620
14621    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
14622    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
14623    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
14624    #[allow(clippy::too_many_arguments)]
14625    fn e4m3_fused2_core(
14626        &self,
14627        b0: &CudaSlice<u8>,
14628        b1: &CudaSlice<u8>,
14629        aq: &CudaSlice<i8>,
14630        ad: &CudaSlice<f32>,
14631        in_f: usize,
14632        out0: usize,
14633        out1: usize,
14634        row_bytes: usize,
14635        ws0: f32,
14636        ws1: f32,
14637    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14638        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14639        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14640        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14641        let f = self.func("qmatvec_e4m3_mmvq_fused2");
14642        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14643        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14644        let cfg = LaunchConfig {
14645            grid_dim: (nb0 + nb1, 1, 1),
14646            block_dim: (32, ROWS_PER_BLOCK, 1),
14647            shared_mem_bytes: 0,
14648        };
14649        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14650        let __s_b = self.gpu.stream();
14651        let mut b = __s_b.launch_builder(&f);
14652        b.arg(b0)
14653            .arg(b1)
14654            .arg(aq)
14655            .arg(ad)
14656            .arg(&mut y0)
14657            .arg(&mut y1)
14658            .arg(&inf)
14659            .arg(&o0)
14660            .arg(&o1)
14661            .arg(&rbl)
14662            .arg(&ws0)
14663            .arg(&ws1);
14664        unsafe {
14665            b.launch(cfg)?;
14666        }
14667        Ok((y0, y1))
14668    }
14669
14670    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
14671    #[allow(clippy::too_many_arguments)]
14672    fn e4m3_fused3_core(
14673        &self,
14674        b0: &CudaSlice<u8>,
14675        b1: &CudaSlice<u8>,
14676        b2: &CudaSlice<u8>,
14677        aq: &CudaSlice<i8>,
14678        ad: &CudaSlice<f32>,
14679        in_f: usize,
14680        out0: usize,
14681        out1: usize,
14682        out2: usize,
14683        row_bytes: usize,
14684        ws0: f32,
14685        ws1: f32,
14686        ws2: f32,
14687    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14688        const ROWS_PER_BLOCK: u32 = 4;
14689        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14690        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14691        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14692        let f = self.func("qmatvec_e4m3_mmvq_fused3");
14693        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14694        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14695        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14696        let cfg = LaunchConfig {
14697            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14698            block_dim: (32, ROWS_PER_BLOCK, 1),
14699            shared_mem_bytes: 0,
14700        };
14701        let (inf, o0, o1, o2, rbl) = (
14702            in_f as i32,
14703            out0 as i32,
14704            out1 as i32,
14705            out2 as i32,
14706            row_bytes as i64,
14707        );
14708        let __s_b = self.gpu.stream();
14709        let mut b = __s_b.launch_builder(&f);
14710        b.arg(b0)
14711            .arg(b1)
14712            .arg(b2)
14713            .arg(aq)
14714            .arg(ad)
14715            .arg(&mut y0)
14716            .arg(&mut y1)
14717            .arg(&mut y2)
14718            .arg(&inf)
14719            .arg(&o0)
14720            .arg(&o1)
14721            .arg(&o2)
14722            .arg(&rbl)
14723            .arg(&ws0)
14724            .arg(&ws1)
14725            .arg(&ws2);
14726        unsafe {
14727            b.launch(cfg)?;
14728        }
14729        Ok((y0, y1, y2))
14730    }
14731
14732    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
14733    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
14734    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
14735    #[allow(clippy::too_many_arguments)]
14736    fn e4m3_fused2_t_core(
14737        &self,
14738        b0: &CudaSlice<u8>,
14739        b1: &CudaSlice<u8>,
14740        aq: &CudaSlice<i8>,
14741        ad: &CudaSlice<f32>,
14742        m: usize,
14743        in_f: usize,
14744        out0: usize,
14745        out1: usize,
14746        row_bytes: usize,
14747        ws0: f32,
14748        ws1: f32,
14749    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14750        const ROWS_PER_BLOCK: u32 = 4;
14751        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14752        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14753        let f = self.func(match Self::batched_mcols(m) {
14754            2 => "qmatvec_e4m3_mmvq_fused2_b2",
14755            4 => "qmatvec_e4m3_mmvq_fused2_b4",
14756            _ => "qmatvec_e4m3_mmvq_fused2_b8",
14757        });
14758        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14759        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14760        let cfg = LaunchConfig {
14761            grid_dim: (nb0 + nb1, 1, 1),
14762            block_dim: (32, ROWS_PER_BLOCK, 1),
14763            shared_mem_bytes: 0,
14764        };
14765        let (inf, o0, o1, mi, rbl) = (
14766            in_f as i32,
14767            out0 as i32,
14768            out1 as i32,
14769            m as i32,
14770            row_bytes as i64,
14771        );
14772        let __s_b = self.gpu.stream();
14773        let mut b = __s_b.launch_builder(&f);
14774        b.arg(b0)
14775            .arg(b1)
14776            .arg(aq)
14777            .arg(ad)
14778            .arg(&mut y0)
14779            .arg(&mut y1)
14780            .arg(&inf)
14781            .arg(&o0)
14782            .arg(&o1)
14783            .arg(&mi)
14784            .arg(&rbl);
14785        unsafe {
14786            b.launch(cfg)?;
14787        }
14788        if ws0 != 1.0 {
14789            self.scale_inplace(&mut y0, ws0, m * out0)?;
14790        }
14791        if ws1 != 1.0 {
14792            self.scale_inplace(&mut y1, ws1, m * out1)?;
14793        }
14794        Ok((y0, y1))
14795    }
14796
14797    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
14798    #[allow(clippy::too_many_arguments)]
14799    fn e4m3_fused3_t_core(
14800        &self,
14801        b0: &CudaSlice<u8>,
14802        b1: &CudaSlice<u8>,
14803        b2: &CudaSlice<u8>,
14804        aq: &CudaSlice<i8>,
14805        ad: &CudaSlice<f32>,
14806        m: usize,
14807        in_f: usize,
14808        out0: usize,
14809        out1: usize,
14810        out2: usize,
14811        row_bytes: usize,
14812        ws0: f32,
14813        ws1: f32,
14814        ws2: f32,
14815    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14816        const ROWS_PER_BLOCK: u32 = 4;
14817        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14818        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14819        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14820        let f = self.func(if Self::batched_mcols(m) == 2 {
14821            "qmatvec_e4m3_mmvq_fused3_b2"
14822        } else {
14823            "qmatvec_e4m3_mmvq_fused3_b4"
14824        });
14825        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14826        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14827        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14828        let cfg = LaunchConfig {
14829            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14830            block_dim: (32, ROWS_PER_BLOCK, 1),
14831            shared_mem_bytes: 0,
14832        };
14833        let (inf, o0, o1, o2, mi, rbl) = (
14834            in_f as i32,
14835            out0 as i32,
14836            out1 as i32,
14837            out2 as i32,
14838            m as i32,
14839            row_bytes as i64,
14840        );
14841        let __s_b = self.gpu.stream();
14842        let mut b = __s_b.launch_builder(&f);
14843        b.arg(b0)
14844            .arg(b1)
14845            .arg(b2)
14846            .arg(aq)
14847            .arg(ad)
14848            .arg(&mut y0)
14849            .arg(&mut y1)
14850            .arg(&mut y2)
14851            .arg(&inf)
14852            .arg(&o0)
14853            .arg(&o1)
14854            .arg(&o2)
14855            .arg(&mi)
14856            .arg(&rbl);
14857        unsafe {
14858            b.launch(cfg)?;
14859        }
14860        if ws0 != 1.0 {
14861            self.scale_inplace(&mut y0, ws0, m * out0)?;
14862        }
14863        if ws1 != 1.0 {
14864            self.scale_inplace(&mut y1, ws1, m * out1)?;
14865        }
14866        if ws2 != 1.0 {
14867            self.scale_inplace(&mut y2, ws2, m * out2)?;
14868        }
14869        Ok((y0, y1, y2))
14870    }
14871
14872    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
14873    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
14874    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
14875    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
14876    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
14877    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
14878    ///
14879    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
14880    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
14881    pub fn qmatvec_e4m3_blk_mmvq(
14882        &self,
14883        bytes: &CudaSlice<u8>,
14884        aq: &CudaSlice<i8>,
14885        ad: &CudaSlice<f32>,
14886        scales: &CudaSlice<f32>,
14887        m: usize,
14888        in_f: usize,
14889        out_f: usize,
14890        row_bytes: usize,
14891        scale_cols: usize,
14892    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14893        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
14894        self.qmatvec_e4m3_blk_mmvq_into(
14895            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14896        )?;
14897        Ok(y)
14898    }
14899
14900    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
14901    #[allow(clippy::too_many_arguments)]
14902    pub fn qmatvec_e4m3_blk_mmvq_into(
14903        &self,
14904        bytes: &CudaSlice<u8>,
14905        aq: &CudaSlice<i8>,
14906        ad: &CudaSlice<f32>,
14907        scales: &CudaSlice<f32>,
14908        m: usize,
14909        in_f: usize,
14910        out_f: usize,
14911        row_bytes: usize,
14912        scale_cols: usize,
14913        y: &mut CudaSlice<f32>,
14914    ) -> Result<(), Box<dyn std::error::Error>> {
14915        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14916        let f = self.func("qmatvec_e4m3_blk_mmvq");
14917        let cfg = LaunchConfig {
14918            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14919            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
14920            shared_mem_bytes: 0,                // warp-only reduce
14921        };
14922        let (inf, outf, mi, rb, sc) = (
14923            in_f as i32,
14924            out_f as i32,
14925            m as i32,
14926            row_bytes as i64,
14927            scale_cols as i32,
14928        );
14929        let __s_b = self.gpu.stream();
14930        let mut b = __s_b.launch_builder(&f);
14931        b.arg(bytes)
14932            .arg(aq)
14933            .arg(ad)
14934            .arg(scales)
14935            .arg(&mut *y)
14936            .arg(&inf)
14937            .arg(&outf)
14938            .arg(&mi)
14939            .arg(&rb)
14940            .arg(&sc);
14941        unsafe {
14942            b.launch(cfg)?;
14943        }
14944        Ok(())
14945    }
14946
14947    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
14948    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
14949    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
14950    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
14951    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
14952    #[allow(clippy::too_many_arguments)]
14953    pub fn qmatvec_e4m3_blk_mmvq_batched(
14954        &self,
14955        bytes: &CudaSlice<u8>,
14956        aq: &CudaSlice<i8>,
14957        ad: &CudaSlice<f32>,
14958        scales: &CudaSlice<f32>,
14959        m: usize,
14960        in_f: usize,
14961        out_f: usize,
14962        row_bytes: usize,
14963        scale_cols: usize,
14964        mcols: usize,
14965    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14966        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14967        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14968        let name = match mcols {
14969            2 => "qmatvec_e4m3_blk_mmvq_b2",
14970            4 => "qmatvec_e4m3_blk_mmvq_b4",
14971            8 => "qmatvec_e4m3_blk_mmvq_b8",
14972            16 => "qmatvec_e4m3_blk_mmvq_b16",
14973            _ => {
14974                return Err(
14975                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14976                );
14977            }
14978        };
14979        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14980        let f = self.func(name);
14981        let cfg = LaunchConfig {
14982            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14983            block_dim: (32, ROWS_PER_BLOCK, 1),
14984            shared_mem_bytes: 0,
14985        };
14986        let (inf, outf, mi, rb, sc) = (
14987            in_f as i32,
14988            out_f as i32,
14989            m as i32,
14990            row_bytes as i64,
14991            scale_cols as i32,
14992        );
14993        let __s_b = self.gpu.stream();
14994        let mut b = __s_b.launch_builder(&f);
14995        b.arg(bytes)
14996            .arg(aq)
14997            .arg(ad)
14998            .arg(scales)
14999            .arg(&mut y)
15000            .arg(&inf)
15001            .arg(&outf)
15002            .arg(&mi)
15003            .arg(&rb)
15004            .arg(&sc);
15005        unsafe {
15006            b.launch(cfg)?;
15007        }
15008        Ok(y)
15009    }
15010
15011    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
15012    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
15013    #[allow(clippy::too_many_arguments)]
15014    pub fn qmatvec_e4m3_blk_batched_raw(
15015        &self,
15016        bytes: &CudaSlice<u8>,
15017        x: &CudaSlice<f32>,
15018        scales: &CudaSlice<f32>,
15019        m: usize,
15020        in_f: usize,
15021        out_f: usize,
15022        row_bytes: usize,
15023        scale_cols: usize,
15024        mcols: usize,
15025    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15026        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15027        self.qmatvec_e4m3_blk_mmvq_batched(
15028            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15029        )
15030    }
15031
15032    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
15033    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
15034    #[allow(clippy::too_many_arguments)]
15035    pub fn qmatvec_e4m3_blk_mmvq_raw(
15036        &self,
15037        bytes: &CudaSlice<u8>,
15038        x: &CudaSlice<f32>,
15039        scales: &CudaSlice<f32>,
15040        m: usize,
15041        in_f: usize,
15042        out_f: usize,
15043        row_bytes: usize,
15044        scale_cols: usize,
15045    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15046        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15047        self.qmatvec_e4m3_blk_mmvq(
15048            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15049        )
15050    }
15051
15052    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
15053    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
15054    #[allow(clippy::too_many_arguments)]
15055    pub fn qmatvec_e4m3_fused2_raw(
15056        &self,
15057        b0: &CudaSlice<u8>,
15058        b1: &CudaSlice<u8>,
15059        x: &CudaSlice<f32>,
15060        in_f: usize,
15061        out0: usize,
15062        out1: usize,
15063        row_bytes: usize,
15064        ws0: f32,
15065        ws1: f32,
15066    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15067        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15068        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15069    }
15070
15071    #[allow(clippy::too_many_arguments)]
15072    pub fn qmatvec_e4m3_fused3_raw(
15073        &self,
15074        b0: &CudaSlice<u8>,
15075        b1: &CudaSlice<u8>,
15076        b2: &CudaSlice<u8>,
15077        x: &CudaSlice<f32>,
15078        in_f: usize,
15079        out0: usize,
15080        out1: usize,
15081        out2: usize,
15082        row_bytes: usize,
15083        ws0: f32,
15084        ws1: f32,
15085        ws2: f32,
15086    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15087        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15088        self.e4m3_fused3_core(
15089            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15090        )
15091    }
15092
15093    #[allow(clippy::too_many_arguments)]
15094    pub fn qmatvec_e4m3_fused2_t_raw(
15095        &self,
15096        b0: &CudaSlice<u8>,
15097        b1: &CudaSlice<u8>,
15098        x: &CudaSlice<f32>,
15099        m: usize,
15100        in_f: usize,
15101        out0: usize,
15102        out1: usize,
15103        row_bytes: usize,
15104        ws0: f32,
15105        ws1: f32,
15106    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15107        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15108        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15109    }
15110
15111    #[allow(clippy::too_many_arguments)]
15112    pub fn qmatvec_e4m3_fused3_t_raw(
15113        &self,
15114        b0: &CudaSlice<u8>,
15115        b1: &CudaSlice<u8>,
15116        b2: &CudaSlice<u8>,
15117        x: &CudaSlice<f32>,
15118        m: usize,
15119        in_f: usize,
15120        out0: usize,
15121        out1: usize,
15122        out2: usize,
15123        row_bytes: usize,
15124        ws0: f32,
15125        ws1: f32,
15126        ws2: f32,
15127    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15128        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15129        self.e4m3_fused3_t_core(
15130            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15131        )
15132    }
15133
15134    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
15135    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
15136    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
15137    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
15138    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
15139    ///
15140    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
15141    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
15142    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
15143    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
15144    fn try_e4m3_blk_pre(
15145        &self,
15146        w: &crate::model::GpuTensor,
15147        aq: &CudaSlice<i8>,
15148        ad: &CudaSlice<f32>,
15149        m: usize,
15150    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15151        use crate::model::GpuTensor;
15152        if let GpuTensor::Quant {
15153            bytes,
15154            qtype,
15155            row_bytes,
15156            blk: Some(g),
15157            ..
15158        } = w
15159        {
15160            if *qtype == QT_F8_E4M3_BLK {
15161                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
15162                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
15163                // below, so the decode-exactness contract is preserved at every width. Gated by
15164                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
15165                // one rollback door covers every dtype's batched tier.
15166                if (2..=16).contains(&m)
15167                    && std::env::var("MEMRA_NO_BATCHED").is_err()
15168                    && (m <= 4 || Self::b8_enabled())
15169                {
15170                    let mcols = Self::batched_mcols(m);
15171                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15172                        bytes,
15173                        aq,
15174                        ad,
15175                        &g.scales,
15176                        m,
15177                        w.in_features(),
15178                        w.out_features(),
15179                        *row_bytes,
15180                        g.cols,
15181                        mcols,
15182                    )?));
15183                }
15184                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15185                    bytes,
15186                    aq,
15187                    ad,
15188                    &g.scales,
15189                    m,
15190                    w.in_features(),
15191                    w.out_features(),
15192                    *row_bytes,
15193                    g.cols,
15194                )?));
15195            }
15196        }
15197        Ok(None)
15198    }
15199
15200    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
15201    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
15202    ///
15203    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
15204    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
15205    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
15206    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
15207    /// prefill keeps the floor's arithmetic and the floor's kernels.
15208    ///
15209    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
15210    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
15211    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
15212    /// (projection, prefill call) and frees immediately.
15213    ///
15214    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
15215    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
15216    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
15217    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
15218    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
15219    /// single-variable comparison instead of a two-variable one.
15220    ///
15221    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
15222    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
15223    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
15224    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
15225    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
15226    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
15227    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
15228    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
15229    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
15230    ///
15231    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
15232    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
15233    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
15234    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
15235    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
15236    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
15237    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
15238    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
15239    /// because v2's denominator had its slab already resident while this class's floor must build it
15240    /// every call; same tile, opposite sign, because the question changed.
15241    ///
15242    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
15243    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
15244    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
15245    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
15246    fn try_e4m3_blk_prefill(
15247        &self,
15248        w: &crate::model::GpuTensor,
15249        x: &CudaSlice<f32>,
15250        m: usize,
15251    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15252        use crate::model::GpuTensor;
15253        let GpuTensor::Quant {
15254            bytes,
15255            qtype,
15256            blk: Some(g),
15257            ..
15258        } = w
15259        else {
15260            return Ok(None);
15261        };
15262        if *qtype != QT_F8_E4M3_BLK {
15263            return Ok(None);
15264        }
15265        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
15266        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
15267        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
15268        // through to the dequant below when they do, never silently produce nothing.
15269        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15270            return Ok(Some(y));
15271        }
15272        let (in_f, out_f) = (w.in_features(), w.out_features());
15273        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15274        let tmp = GpuTensor::Quant {
15275            bytes: slab,
15276            qtype: QT_Q8_0,
15277            row_bytes: in_f / 32 * 34,
15278            ne: vec![in_f as u64, out_f as u64],
15279            scale: 1.0,
15280            rp: false,
15281            #[cfg(memra_cutlass)]
15282            cutlass: None,
15283            fp8: None,
15284            blk: None,
15285            f16: None,
15286            rp4: None,
15287        };
15288        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
15289        Ok(Some(self.matmul(&tmp, x, m)?))
15290    }
15291
15292    pub fn matmul_pre_noscale(
15293        &self,
15294        w: &crate::model::GpuTensor,
15295        aq: &CudaSlice<i8>,
15296        ad: &CudaSlice<f32>,
15297        m: usize,
15298    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15299        use crate::model::GpuTensor;
15300        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
15301        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
15302        // rather than let the tail below refuse and cost the caller a re-dispatch.
15303        if m == 1 {
15304            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15305                return Ok(Some((y, 1.0)));
15306            }
15307        }
15308        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
15309        if m != 1 || !self.uses_q8_1_fast(w) {
15310            return Ok(None);
15311        }
15312        let in_f = w.in_features();
15313        let out_f = w.out_features();
15314        let (bytes, qtype, row_bytes, scale, rp) = match w {
15315            GpuTensor::Quant {
15316                bytes,
15317                qtype,
15318                row_bytes,
15319                scale,
15320                rp,
15321                ..
15322            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15323            _ => return Ok(None),
15324        };
15325        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
15326        if self.mmvq_supports(qtype) {
15327            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
15328            let (mbytes, mrp) = match w {
15329                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15330                _ => (bytes, rp),
15331            };
15332            let y = self.qmatvec_mmvq(
15333                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
15334            )?;
15335            return Ok(Some((y, scale)));
15336        }
15337        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
15338        let name = match qtype {
15339            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15340            QT_Q4_K => "qmatvec_q4_K_dp4a",
15341            QT_Q6_K => "qmatvec_q6_K_dp4a",
15342            QT_Q5_K => "qmatvec_q5_K_dp4a",
15343            QT_Q3_K => "qmatvec_q3_K_dp4a",
15344            QT_NVFP4 => {
15345                if rp {
15346                    "qmatvec_nvfp4_dp4a_rp"
15347                } else {
15348                    "qmatvec_nvfp4_dp4a"
15349                }
15350            }
15351            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15352            _ => return Ok(None),
15353        };
15354        let f = self.func(name);
15355        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15356        let cfg = LaunchConfig {
15357            grid_dim: (out_f as u32, m as u32, 1),
15358            block_dim: (128, 1, 1),
15359            shared_mem_bytes: 0,
15360        };
15361        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15362        let __s_b = self.gpu.stream();
15363        let mut b = __s_b.launch_builder(&f);
15364        b.arg(bytes)
15365            .arg(aq)
15366            .arg(ad)
15367            .arg(&mut y)
15368            .arg(&inf)
15369            .arg(&outf)
15370            .arg(&mi)
15371            .arg(&rb);
15372        unsafe {
15373            b.launch(cfg)?;
15374        }
15375        Ok(Some((y, scale)))
15376    }
15377
15378    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15379    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15380    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15381        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15382        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15383        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15384        // is a pure function of the dtype — the decode-parity law holds under every env.
15385        if qtype == QT_F8_E4M3 {
15386            return true;
15387        }
15388        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15389            return false;
15390        }
15391        matches!(
15392            qtype,
15393            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15394        )
15395    }
15396
15397    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15398    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15399    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15400    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15401    pub fn qmatvec_mmvq(
15402        &self,
15403        bytes: &CudaSlice<u8>,
15404        aq: &CudaSlice<i8>,
15405        ad: &CudaSlice<f32>,
15406        m: usize,
15407        in_f: usize,
15408        out_f: usize,
15409        qtype: i32,
15410        row_bytes: usize,
15411        scale: f32,
15412        rp: bool,
15413    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15414        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15415        self.qmatvec_mmvq_into(
15416            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15417        )?;
15418        Ok(y)
15419    }
15420
15421    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15422    #[allow(clippy::too_many_arguments)]
15423    pub fn qmatvec_mmvq_into(
15424        &self,
15425        bytes: &CudaSlice<u8>,
15426        aq: &CudaSlice<i8>,
15427        ad: &CudaSlice<f32>,
15428        m: usize,
15429        in_f: usize,
15430        out_f: usize,
15431        qtype: i32,
15432        row_bytes: usize,
15433        scale: f32,
15434        rp: bool,
15435        y: &mut CudaSlice<f32>,
15436    ) -> Result<(), Box<dyn std::error::Error>> {
15437        debug_assert!(y.len() >= m * out_f);
15438        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15439        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15440        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15441        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15442        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15443        if qtype == QT_Q8_0
15444            && rp
15445            && m == 1
15446            && out_f >= 64
15447            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15448            && {
15449                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15450                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15451            }
15452        {
15453            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15454            let cfg = LaunchConfig {
15455                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15456                block_dim: (32, 2, 1),
15457                shared_mem_bytes: 0,
15458            };
15459            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15460            let __s_b = self.gpu.stream();
15461            let mut b = __s_b.launch_builder(&f);
15462            b.arg(bytes)
15463                .arg(aq)
15464                .arg(ad)
15465                .arg(&mut *y)
15466                .arg(&inf)
15467                .arg(&outf)
15468                .arg(&mi)
15469                .arg(&rb);
15470            unsafe {
15471                b.launch(cfg)?;
15472            }
15473            if scale != 1.0 {
15474                self.scale_inplace(y, scale, out_f)?;
15475            }
15476            return Ok(());
15477        }
15478        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15479        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15480        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15481        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
15482        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
15483        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
15484        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
15485        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
15486        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15487            2
15488        } else {
15489            1
15490        };
15491        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
15492        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
15493        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
15494        // valid-window interleaved, bit-identical per row — same dot program).
15495        if m == 1 && qtype == QT_Q4_0 {
15496            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15497            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
15498            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
15499            mr = *Q40MR.get_or_init(|| {
15500                std::env::var("MEMRA_Q40_MR")
15501                    .ok()
15502                    .and_then(|v| v.parse().ok())
15503                    .unwrap_or(1)
15504            });
15505        }
15506        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
15507        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
15508        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
15509        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
15510        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
15511        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
15512        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
15513        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
15514        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
15515        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
15516        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15517        let q5_force = q5_mode.as_deref() == Some("2");
15518        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
15519        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
15520        let q5_il = qtype == QT_Q5_K
15521            && m == 1
15522            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15523        if q5_il && !q5_force && out_f > 65536 {
15524            mr = 1;
15525        }
15526        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
15527        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
15528        if qtype == QT_Q4_0 && rp && mr != 1 {
15529            mr = 2;
15530        }
15531        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
15532        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
15533        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
15534        if qtype == QT_Q8_0 && rp {
15535            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15536            mr = *Q80MR.get_or_init(|| {
15537                std::env::var("MEMRA_Q80_MR")
15538                    .ok()
15539                    .and_then(|v| v.parse().ok())
15540                    .unwrap_or(1)
15541            });
15542        }
15543        let name = match (qtype, mr, rp) {
15544            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15545            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15546            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15547            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15548            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15549            (QT_Q5_K, 2, _) => {
15550                if q5_il {
15551                    "qmatvec_q5_K_mmvq_mr2_il"
15552                } else {
15553                    "qmatvec_q5_K_mmvq_mr2"
15554                }
15555            }
15556            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15557            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
15558            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
15559            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
15560            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
15561            (QT_Q8_0, _, true)
15562                if in_f % 1024 == 0 && {
15563                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15564                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15565                } =>
15566            {
15567                "qmatvec_q8_0_mmvq_rpca"
15568            }
15569            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15570            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15571            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
15572            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
15573            // reach a GGUF-layout kernel or vice versa.
15574            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15575            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15576            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15577            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15578            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15579            (QT_Q5_K, _, _) => {
15580                if q5_il {
15581                    "qmatvec_q5_K_mmvq_il"
15582                } else {
15583                    "qmatvec_q5_K_mmvq"
15584                }
15585            }
15586            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15587            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15588            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15589            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15590        };
15591        let f = self.func(name);
15592        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
15593        let rows_per_block = ROWS_PER_BLOCK * mr;
15594        let cfg = LaunchConfig {
15595            grid_dim: (
15596                (out_f as u32 + rows_per_block - 1) / rows_per_block,
15597                m as u32,
15598                1,
15599            ),
15600            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
15601            shared_mem_bytes: 0,                // warp-only reduce at m=1
15602        };
15603        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15604        let __s_b = self.gpu.stream();
15605        let mut b = __s_b.launch_builder(&f);
15606        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
15607        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
15608        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
15609        // weight_scale). Other mmvq kernels keep the 8-arg signature.
15610        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15611            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
15612            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
15613            if Self::pdl_on()
15614                && Self::pdl_mmvq_on()
15615                && Self::pdl_nvfp4q8_on()
15616                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15617            {
15618                use cudarc::driver::{DevicePtr, DevicePtrMut};
15619                let s = &self.gpu.stream();
15620                let (pw, _g0) = bytes.device_ptr(s);
15621                let (paq, _g1) = aq.device_ptr(s);
15622                let (pad, _g2) = ad.device_ptr(s);
15623                let (py, _g3) = y.device_ptr_mut(s);
15624                let mut ps = [
15625                    &pw as *const _ as *mut std::ffi::c_void,
15626                    &paq as *const _ as *mut _,
15627                    &pad as *const _ as *mut _,
15628                    &py as *const _ as *mut _,
15629                    &inf as *const _ as *mut _,
15630                    &outf as *const _ as *mut _,
15631                    &mi as *const _ as *mut _,
15632                    &rb as *const _ as *mut _,
15633                    &scale as *const _ as *mut _,
15634                ];
15635                unsafe {
15636                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15637                }
15638                return Ok(());
15639            }
15640            b.arg(bytes)
15641                .arg(aq)
15642                .arg(ad)
15643                .arg(&mut *y)
15644                .arg(&inf)
15645                .arg(&outf)
15646                .arg(&mi)
15647                .arg(&rb)
15648                .arg(&scale);
15649            unsafe {
15650                b.launch(cfg)?;
15651            }
15652        } else if Self::pdl_on()
15653            && Self::pdl_mmvq_on()
15654            && (matches!(
15655                name,
15656                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15657            ) || (Self::pdl_nvfp4q8_on()
15658                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15659        {
15660            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
15661            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
15662            // names may take this launch (unmarked kernels would read unordered).
15663            {
15664                use cudarc::driver::{DevicePtr, DevicePtrMut};
15665                let s = &self.gpu.stream();
15666                let (pw, _g0) = bytes.device_ptr(s);
15667                let (paq, _g1) = aq.device_ptr(s);
15668                let (pad, _g2) = ad.device_ptr(s);
15669                let (py, _g3) = y.device_ptr_mut(s);
15670                let mut ps = [
15671                    &pw as *const _ as *mut std::ffi::c_void,
15672                    &paq as *const _ as *mut _,
15673                    &pad as *const _ as *mut _,
15674                    &py as *const _ as *mut _,
15675                    &inf as *const _ as *mut _,
15676                    &outf as *const _ as *mut _,
15677                    &mi as *const _ as *mut _,
15678                    &rb as *const _ as *mut _,
15679                ];
15680                unsafe {
15681                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15682                }
15683            }
15684            if scale != 1.0 {
15685                self.scale_inplace(y, scale, m * out_f)?;
15686            }
15687        } else {
15688            b.arg(bytes)
15689                .arg(aq)
15690                .arg(ad)
15691                .arg(&mut *y)
15692                .arg(&inf)
15693                .arg(&outf)
15694                .arg(&mi)
15695                .arg(&rb);
15696            unsafe {
15697                b.launch(cfg)?;
15698            }
15699            if scale != 1.0 {
15700                self.scale_inplace(y, scale, m * out_f)?;
15701            }
15702        }
15703        Ok(())
15704    }
15705
15706    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
15707    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
15708    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
15709    pub fn qmatvec_mmvq_raw(
15710        &self,
15711        bytes: &CudaSlice<u8>,
15712        x: &CudaSlice<f32>,
15713        m: usize,
15714        in_f: usize,
15715        out_f: usize,
15716        qtype: i32,
15717        row_bytes: usize,
15718        rp: bool,
15719    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15720        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15721        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15722    }
15723
15724    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
15725    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
15726    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
15727    pub fn batched_supports(&self, qtype: i32) -> bool {
15728        matches!(
15729            qtype,
15730            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15731        )
15732    }
15733
15734    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
15735    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
15736    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
15737    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
15738    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
15739    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
15740    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
15741    pub fn iq_fast_enabled() -> bool {
15742        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15743        *ON.get_or_init(|| {
15744            std::env::var("MEMRA_IQ_FAST")
15745                .map(|v| v != "0")
15746                .unwrap_or(true)
15747        })
15748    }
15749
15750    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
15751    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
15752    pub fn b8_enabled() -> bool {
15753        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15754        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15755    }
15756
15757    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
15758    pub fn batched_mcols(m: usize) -> usize {
15759        if m == 2 {
15760            2
15761        } else if m <= 4 {
15762            4
15763        } else if m <= 8 {
15764            8
15765        } else {
15766            16
15767        }
15768    }
15769
15770    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
15771    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
15772    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
15773    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
15774    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15775        Some(match (qtype, mcols) {
15776            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15777            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15778            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15779            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
15780            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
15781            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
15782            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
15783            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
15784            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15785            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15786            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15787            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15788            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
15789            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
15790            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15791            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15792            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15793            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15794            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
15795            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
15796            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15797            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15798            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15799            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15800            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15801            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15802            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15803            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15804            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
15805            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
15806            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
15807            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15808            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15809            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15810            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15811            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
15812            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
15813            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
15814            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15815            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15816            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15817            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15818            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15819            _ => return None,
15820        })
15821    }
15822
15823    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
15824    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
15825    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
15826    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
15827    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
15828    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
15829    ///
15830    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
15831    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
15832    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
15833    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
15834    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
15835    /// msweep on all six 27B shapes (2026-07-03):
15836    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
15837    ///          it applies for b4 (-3..-14%), never loses;
15838    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
15839    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
15840    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
15841    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
15842    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
15843    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
15844    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
15845    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
15846    /// b2: in_f>=6144 -> r2, else base.
15847    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
15848    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
15849    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
15850    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
15851    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
15852    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
15853    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
15854    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
15855    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
15856    /// Device SM count (cached) — grid-fill policy input.
15857    pub fn sm_count(&self) -> i32 {
15858        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15859        *SMS.get_or_init(|| {
15860            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15861            self.gpu
15862                .ctx
15863                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15864                .unwrap_or(82)
15865        })
15866    }
15867
15868    pub fn batched_variant(
15869        &self,
15870        _m: usize,
15871        in_f: usize,
15872        out_f: usize,
15873        qtype: i32,
15874        row_bytes: usize,
15875        mcols: usize,
15876        rp: bool,
15877    ) -> &'static str {
15878        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
15879        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
15880        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
15881        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
15882        if qtype == QT_Q8_0 {
15883            return if rp { "rp" } else { "base" };
15884        }
15885        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15886        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15887            Ok("base") => "base",
15888            Ok("pf") => "pf",
15889            Ok("r2") => "r2",
15890            Ok("r2w8") => "r2w8",
15891            Ok("pfr2") => "pfr2",
15892            Ok("ca") => "ca",
15893            Ok("car2") => "car2",
15894            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
15895            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
15896            Ok("rp") => "rp",
15897            Ok("rpr2") => "rpr2",
15898            Ok("rpr2w8") => "rpr2w8",
15899            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
15900            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
15901            Ok("rpca") => "rpca",
15902            Ok("rpcar2") => "rpcar2",
15903            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
15904            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
15905            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
15906            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
15907            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
15908            // bit-identical to the decode path — measurement corpus ONLY, never auto).
15909            Ok("rpsc") => "rpsc",
15910            Ok("rpms") => "rpms",
15911            Ok("rpmsc") => "rpmsc",
15912            Ok("rpks") => "rpks",
15913            Ok("rpksc") => "rpksc",
15914            _ => "auto",
15915        });
15916        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
15917        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
15918        // shapes qualify; anything else falls back to the register variants.
15919        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15920        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
15921        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
15922        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
15923        // forced MEMRA_MMVQ_BV values still work).
15924        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15925        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15926        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15927        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15928        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15929        let sms = *SMS.get_or_init(|| {
15930            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15931            self.gpu
15932                .ctx
15933                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15934                .unwrap_or(82)
15935        });
15936        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
15937        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
15938        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
15939        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
15940        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
15941        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
15942        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
15943        // AUTO RULE = the measured winners table (differs from NVFP4's!):
15944        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
15945        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
15946        //     r2 1258us) — kernels kept behind the force seam for the corpus;
15947        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
15948        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
15949        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
15950        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
15951        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
15952        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
15953        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
15954        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
15955        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15956        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
15957        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
15958        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15959        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15960            Ok("base") => "base",
15961            Ok("r2") => "r2",
15962            Ok("r2w8") => "r2w8",
15963            _ => "auto",
15964        });
15965        let variant: &'static str = if qtype == QT_Q4_0 {
15966            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
15967            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
15968            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
15969            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15970            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15971                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
15972                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
15973                // + syncs cost more than the stalls, bank-pad made no difference);
15974                // register load-ahead flat (nvcc already reorders). The b-tier limiter
15975                // is still unidentified — see the jsonl row.
15976                Ok("base") => "base",
15977                Ok("r2") => "r2",
15978                Ok("ms") => "ms",
15979                Ok("sm") => "sm",
15980                Ok("la") => "la",
15981                _ => "auto",
15982            });
15983            let v = if q40 != "auto" {
15984                q40
15985            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15986                "r2"
15987            } else {
15988                "base"
15989            };
15990            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
15991            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
15992            // and the limiter is the per-column activation load chain (long_scoreboard
15993            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
15994            if rp {
15995                match v {
15996                    "ms" => "r2ms_rp",
15997                    "sm" => "r2sm_rp",
15998                    "la" => "r2la_rp",
15999                    "r2" => "r2_rp",
16000                    _ => "rp",
16001                }
16002            } else if matches!(v, "ms" | "sm" | "la") {
16003                "r2"
16004            } else {
16005                v
16006            }
16007        } else if qtype != QT_NVFP4 && !kq_r2 {
16008            "base"
16009        } else if kq_r2 && rp {
16010            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
16011            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
16012            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
16013            "rp"
16014        } else if kq_r2 {
16015            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
16016            // mcols != 4 forced r2w8 falls to unbounded r2.
16017            if kq_bv != "auto" {
16018                if kq_bv == "r2w8" && mcols != 4 {
16019                    "r2"
16020                } else {
16021                    kq_bv
16022                }
16023            } else if bv != "auto" {
16024                match bv {
16025                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16026                    "r2w8" | "rpr2w8" => {
16027                        if mcols != 4 {
16028                            "r2"
16029                        } else {
16030                            "r2w8"
16031                        }
16032                    }
16033                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
16034                }
16035            } else {
16036                let blocks = (out_f + 7) / 8;
16037                let waves = blocks as f64 / (7 * sms as usize) as f64;
16038                let filled = blocks >= 4 * sms as usize;
16039                let use_r2 = if qtype == QT_Q4_K {
16040                    filled
16041                } else {
16042                    waves >= 2.0
16043                };
16044                if use_r2 { "r2" } else { "base" }
16045            }
16046        } else if bv != "auto" {
16047            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
16048            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
16049            // unsupported (shape, mcols) combos fall back to pf/r2.
16050            // On rp buffers, forced legacy names map to their rp twins (layout law).
16051            let v = if bv == "r2w8" && mcols == 2 {
16052                "r2"
16053            } else if bv == "ca" && (!ca_ok || mcols == 8) {
16054                "pf"
16055            } else if bv == "car2" && (!ca_ok || mcols == 8) {
16056                "r2"
16057            } else if bv == "pfr2" && mcols == 8 {
16058                "r2"
16059            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16060                "rpr2"
16061            }
16062            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
16063            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16064                if mcols == 8 { "rpr2w8" } else { "rpr2" }
16065            } else if bv == "rpcar2" && mcols == 2 {
16066                "rpca"
16067            }
16068            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
16069            // (rpms has no smem and no alignment need — always valid on rp buffers).
16070            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16071                "rpr2"
16072            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16073                "rpr2"
16074            } else {
16075                bv
16076            };
16077            if rp {
16078                match v {
16079                    "base" | "pf" | "ca" | "rp" => "rp",
16080                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16081                    "r2w8" | "rpr2w8" => {
16082                        if mcols == 2 {
16083                            "rpr2"
16084                        } else {
16085                            "rpr2w8"
16086                        }
16087                    }
16088                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
16089                }
16090            } else {
16091                v
16092            }
16093        } else if mcols == 8 {
16094            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
16095            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
16096            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
16097            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
16098            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
16099            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
16100            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
16101            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
16102            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
16103            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
16104            if rp {
16105                if sc_ok { "rpsc" } else { "rpr2w8" }
16106            } else {
16107                "r2w8"
16108            }
16109        } else if mcols >= 4 {
16110            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
16111            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
16112            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
16113            let blocks = (out_f + 7) / 8;
16114            let r7 = 7 * sms as usize;
16115            let r8 = 8 * sms as usize;
16116            let waves = blocks as f64 / r7 as f64;
16117            let filled = blocks >= 4 * sms as usize;
16118            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
16119            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
16120            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
16121            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16122                // the extra residency drops the INTEGER wave count -> the straggler wave a
16123                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
16124                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
16125                if rp { "rpr2w8" } else { "r2w8" }
16126            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16127                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
16128                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
16129                if rp { "rpr2" } else { "r2" }
16130            } else {
16131                // fractional straggler-wave window with no crossing, or grid too small to fill
16132                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
16133                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
16134                if rp { "rp" } else { "pf" }
16135            }
16136        } else if in_f >= 6144 {
16137            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
16138            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
16139            // stays.
16140            if rp { "rpr2" } else { "r2" }
16141        } else if rp {
16142            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
16143            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
16144            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
16145            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
16146            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16147            if sc_ok && waves >= 0.9 && waves <= 1.1 {
16148                "rpsc"
16149            } else {
16150                "rp"
16151            }
16152        } else {
16153            "base"
16154        };
16155        variant
16156    }
16157
16158    pub fn qmatvec_mmvq_batched(
16159        &self,
16160        bytes: &CudaSlice<u8>,
16161        aq: &CudaSlice<i8>,
16162        ad: &CudaSlice<f32>,
16163        m: usize,
16164        in_f: usize,
16165        out_f: usize,
16166        qtype: i32,
16167        row_bytes: usize,
16168        mcols: usize,
16169        scale: f32,
16170        rp: bool,
16171    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16172        const ROWS_PER_BLOCK: u32 = 4;
16173        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
16174        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
16175        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
16176        // weight keeps its rp-layout kernel family regardless of the override.
16177        let forced: Option<&'static str> = {
16178            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16179            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16180                .as_deref()
16181                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16182        };
16183        let variant = match forced {
16184            Some(v) if !rp || v.contains("rp") => v,
16185            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16186        };
16187        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16188            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16189        })?;
16190        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
16191        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
16192        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
16193        let variant = if mcols == 16 {
16194            if rp { "rp" } else { "base" }
16195        } else {
16196            variant
16197        };
16198        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
16199        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
16200        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
16201        // per-(token,row) chain (columns c >= m never execute in either form) ->
16202        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
16203        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
16204        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16205        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16206        if b567
16207            && qtype == QT_NVFP4
16208            && rp
16209            && mcols == 8
16210            && (5..=7).contains(&m)
16211            && matches!(variant, "rpsc" | "rpr2w8")
16212        {
16213            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16214            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
16215            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16216            let cfg = LaunchConfig {
16217                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16218                block_dim: (32, ROWS_PER_BLOCK, 1),
16219                shared_mem_bytes: 0,
16220            };
16221            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16222            let __s_b = self.gpu.stream();
16223            let mut b = __s_b.launch_builder(&f);
16224            b.arg(bytes)
16225                .arg(aq)
16226                .arg(ad)
16227                .arg(&mut y)
16228                .arg(&inf)
16229                .arg(&outf)
16230                .arg(&mi)
16231                .arg(&rb);
16232            unsafe {
16233                b.launch(cfg)?;
16234            }
16235            if scale != 1.0 {
16236                self.scale_inplace(&mut y, scale, m * out_f)?;
16237            }
16238            return Ok(y);
16239        }
16240        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16241            "base" => (base_name.into(), ROWS_PER_BLOCK),
16242            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16243            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16244            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16245            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
16246            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
16247            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
16248            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16249            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16250            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16251            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16252            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16253            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16254            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16255            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
16256        };
16257        debug_assert!(
16258            !rp || name.contains("_rp"),
16259            "rp weight dispatched to a GGUF-layout kernel"
16260        );
16261        let f = self.func(&name);
16262        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16263        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
16264        let smem = if name.contains("_r2sm_rp") {
16265            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16266        } else {
16267            0
16268        };
16269        let cfg = LaunchConfig {
16270            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16271            block_dim: (32, ROWS_PER_BLOCK, 1),
16272            shared_mem_bytes: smem,
16273        };
16274        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16275        let __s_b = self.gpu.stream();
16276        let mut b = __s_b.launch_builder(&f);
16277        b.arg(bytes)
16278            .arg(aq)
16279            .arg(ad)
16280            .arg(&mut y)
16281            .arg(&inf)
16282            .arg(&outf)
16283            .arg(&mi)
16284            .arg(&rb);
16285        unsafe {
16286            b.launch(cfg)?;
16287        }
16288        if scale != 1.0 {
16289            self.scale_inplace(&mut y, scale, m * out_f)?;
16290        }
16291        Ok(y)
16292    }
16293
16294    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
16295    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
16296    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
16297    pub fn qmatvec_batched_raw(
16298        &self,
16299        bytes: &CudaSlice<u8>,
16300        x: &CudaSlice<f32>,
16301        m: usize,
16302        in_f: usize,
16303        out_f: usize,
16304        qtype: i32,
16305        row_bytes: usize,
16306        mcols: usize,
16307        rp: bool,
16308    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16309        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16310        self.qmatvec_mmvq_batched(
16311            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16312        )
16313    }
16314
16315    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
16316    pub fn qmatvec_nvfp4_batched_raw(
16317        &self,
16318        bytes: &CudaSlice<u8>,
16319        x: &CudaSlice<f32>,
16320        m: usize,
16321        in_f: usize,
16322        out_f: usize,
16323        row_bytes: usize,
16324        mcols: usize,
16325        rp: bool,
16326    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16327        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16328    }
16329
16330    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
16331    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
16332    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
16333    fn try_fp4_gemm(
16334        &self,
16335        w: &crate::model::GpuTensor,
16336        x: &CudaSlice<f32>,
16337        m: usize,
16338        in_f: usize,
16339        out_f: usize,
16340    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16341        use crate::model::GpuTensor;
16342        if cfg!(memra_portable_cuda) {
16343            return Ok(None);
16344        }
16345        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which cu/qmatvec_gemm.cu:1234 omits on a
16346        // portable build (the mxf4 block-scale MMA is sm_120a-only). Refuse at the door.
16347        if std::env::var("MEMRA_FP4").is_ok() {
16348            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16349        }
16350        if std::env::var("MEMRA_FP4").is_err() {
16351            return Ok(None);
16352        }
16353        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
16354        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
16355        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
16356        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
16357        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
16358        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
16359        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
16360        // for the common no-macro-scale case.
16361        #[cfg(memra_cutlass)]
16362        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16363            if let GpuTensor::Quant {
16364                bytes,
16365                qtype,
16366                scale,
16367                row_bytes,
16368                cutlass,
16369                ..
16370            } = w
16371            {
16372                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16373                    if let Some(cw) = cutlass {
16374                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16375                        let y = self.cutlass_fp4_gemm(
16376                            &cw.b_packed,
16377                            &cw.sfb_swizzled,
16378                            x,
16379                            *scale,
16380                            m,
16381                            out_f,
16382                            in_f,
16383                        )?;
16384                        return Ok(Some(y));
16385                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16386                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16387                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16388                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16389                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16390                        let (b_packed, sfb_sw) =
16391                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16392                        let y =
16393                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16394                        return Ok(Some(y));
16395                    }
16396                }
16397            }
16398        }
16399        if let GpuTensor::Quant {
16400            bytes,
16401            qtype,
16402            row_bytes,
16403            scale,
16404            rp,
16405            ..
16406        } = w
16407        {
16408            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16409            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16410            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16411                let y =
16412                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16413                return Ok(Some(y));
16414            }
16415        }
16416        Ok(None)
16417    }
16418
16419    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16420    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16421    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16422    pub fn rms_norm_f16out(
16423        &self,
16424        x: &CudaSlice<f32>,
16425        w: &CudaSlice<f32>,
16426        dst: &mut CudaSlice<f32>,
16427        dst16: &mut CudaSlice<u8>,
16428        ncols: usize,
16429        nrows: usize,
16430        eps: f32,
16431    ) -> Result<(), Box<dyn std::error::Error>> {
16432        let f = self.func("rms_norm_f16out_f32");
16433        let cfg = LaunchConfig {
16434            grid_dim: (nrows as u32, 1, 1),
16435            block_dim: (rms_block(), 1, 1),
16436            shared_mem_bytes: 0,
16437        };
16438        let (nc, e) = (ncols as i32, eps);
16439        let __s_b = self.gpu.stream();
16440        let mut b = __s_b.launch_builder(&f);
16441        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16442        unsafe {
16443            b.launch(cfg)?;
16444        }
16445        Ok(())
16446    }
16447
16448    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16449    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16450    #[allow(clippy::too_many_arguments)]
16451    pub fn add_rms_norm_f16out(
16452        &self,
16453        a: &CudaSlice<f32>,
16454        b: &CudaSlice<f32>,
16455        w: &CudaSlice<f32>,
16456        res: &mut CudaSlice<f32>,
16457        dst: &mut CudaSlice<f32>,
16458        dst16: &mut CudaSlice<u8>,
16459        ncols: usize,
16460        nrows: usize,
16461        eps: f32,
16462    ) -> Result<(), Box<dyn std::error::Error>> {
16463        let f = self.func("add_rms_norm_f16out_f32");
16464        let cfg = LaunchConfig {
16465            grid_dim: (nrows as u32, 1, 1),
16466            block_dim: (rms_block(), 1, 1),
16467            shared_mem_bytes: 0,
16468        };
16469        let (nc, e) = (ncols as i32, eps);
16470        let __s_lb = self.gpu.stream();
16471        let mut lb = __s_lb.launch_builder(&f);
16472        lb.arg(a)
16473            .arg(b)
16474            .arg(w)
16475            .arg(res)
16476            .arg(dst)
16477            .arg(dst16)
16478            .arg(&nc)
16479            .arg(&e);
16480        unsafe {
16481            lb.launch(cfg)?;
16482        }
16483        Ok(())
16484    }
16485
16486    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
16487    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
16488    pub fn matmul_group_xh(
16489        &self,
16490        ws: &[&crate::model::GpuTensor],
16491        x: &CudaSlice<f32>,
16492        xh: &CudaSlice<u8>,
16493        m: usize,
16494    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16495        let mut out = Vec::with_capacity(ws.len());
16496        let in_f = ws[0].in_features();
16497        for w in ws {
16498            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16499                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16500                    out.push(y);
16501                    continue;
16502                }
16503            }
16504            out.push(self.matmul(w, x, m)?);
16505        }
16506        Ok(out)
16507    }
16508
16509    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
16510    /// GDN steps). Layouts [T, H].
16511    pub fn gdn_pad_mask(
16512        &self,
16513        beta: &mut CudaSlice<f32>,
16514        g_log: &mut CudaSlice<f32>,
16515        len_d: &CudaSlice<i32>,
16516        h: usize,
16517        t: usize,
16518    ) -> Result<(), Box<dyn std::error::Error>> {
16519        let f = self.func("gdn_pad_mask_f32");
16520        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16521        let (hi, ti) = (h as i32, t as i32);
16522        let __s_b = self.gpu.stream();
16523        let mut b = __s_b.launch_builder(&f);
16524        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16525        unsafe {
16526            b.launch(cfg)?;
16527        }
16528        Ok(())
16529    }
16530
16531    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
16532    /// gather for the padded prime graph's h_seed/hlast.
16533    pub fn row_gather_dev(
16534        &self,
16535        src: &CudaSlice<f32>,
16536        dst: &mut CudaSlice<f32>,
16537        len_d: &CudaSlice<i32>,
16538        ncols: usize,
16539    ) -> Result<(), Box<dyn std::error::Error>> {
16540        let f = self.func("row_gather_dev_f32");
16541        let cfg = LaunchConfig::for_num_elems(ncols as u32);
16542        let nc = ncols as i32;
16543        let __s_b = self.gpu.stream();
16544        let mut b = __s_b.launch_builder(&f);
16545        b.arg(src).arg(dst).arg(len_d).arg(&nc);
16546        unsafe {
16547            b.launch(cfg)?;
16548        }
16549        Ok(())
16550    }
16551
16552    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
16553    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
16554    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
16555    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
16556    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
16557    /// different in_f) falls back to its own `matmul` — behavior unchanged.
16558    pub fn matmul_group(
16559        &self,
16560        ws: &[&crate::model::GpuTensor],
16561        x: &CudaSlice<f32>,
16562        m: usize,
16563    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16564        use crate::model::GpuTensor;
16565        let mut out = Vec::with_capacity(ws.len());
16566        let any_mirror = ws
16567            .iter()
16568            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16569        if m >= 16 && any_mirror && !self.verify_exact_on() {
16570            let in_f = ws[0].in_features();
16571            let xh = self.f16_act(x, m * in_f, in_f)?;
16572            for w in ws {
16573                if w.in_features() == in_f {
16574                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16575                        out.push(y);
16576                        continue;
16577                    }
16578                }
16579                out.push(self.matmul(w, x, m)?);
16580            }
16581            return Ok(out);
16582        }
16583        for w in ws {
16584            out.push(self.matmul(w, x, m)?);
16585        }
16586        Ok(out)
16587    }
16588
16589    /// Cross-request grouped matmul (task #13): run ONE projection group over the
16590    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
16591    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
16592    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
16593    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
16594    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
16595    pub fn matmul_group_multi(
16596        &self,
16597        ws: &[&crate::model::GpuTensor],
16598        xs: &[&CudaSlice<f32>],
16599        ms: &[usize],
16600    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16601        assert_eq!(xs.len(), ms.len());
16602        let in_f = ws[0].in_features();
16603        let total: usize = ms.iter().sum();
16604        let mut xcat = self.uninit(total * in_f)?;
16605        let mut off = 0usize;
16606        for (x, &m) in xs.iter().zip(ms) {
16607            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16608            off += m;
16609        }
16610        let ys = self.matmul_group(ws, &xcat, total)?;
16611        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16612        for (w, y) in ws.iter().zip(ys) {
16613            let out_f = w.out_features();
16614            let mut off = 0usize;
16615            for (s, &m) in ms.iter().enumerate() {
16616                let mut ys_s = self.uninit(m * out_f)?;
16617                let src = y.slice(off * out_f..(off + m) * out_f);
16618                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16619                out[s].push(ys_s);
16620                off += m;
16621            }
16622        }
16623        Ok(out)
16624    }
16625
16626    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
16627    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
16628    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
16629    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
16630    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
16631    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
16632    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
16633    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
16634    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
16635    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16636        use crate::model::GpuTensor;
16637        if !legacy_quant_gemm_allowed(
16638            cfg!(memra_portable_cuda),
16639            cfg!(memra_hopper_mma),
16640            std::env::var_os("MEMRA_NO_GEMM").is_some(),
16641        ) {
16642            return false;
16643        }
16644        match w {
16645            GpuTensor::Quant { qtype, .. } => {
16646                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16647                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16648            }
16649            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16650        }
16651    }
16652
16653    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
16654    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
16655    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
16656    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
16657    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
16658    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
16659    pub fn qmatvec_gemm(
16660        &self,
16661        w: &crate::model::GpuTensor,
16662        aq: &CudaSlice<i8>,
16663        ad: &CudaSlice<f32>,
16664        m: usize,
16665    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16666        use crate::model::GpuTensor;
16667        let in_f = w.in_features();
16668        let out_f = w.out_features();
16669        let (bytes, qtype, row_bytes, scale, rp) = match w {
16670            GpuTensor::Quant {
16671                bytes,
16672                qtype,
16673                row_bytes,
16674                scale,
16675                rp,
16676                ..
16677            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16678            _ => unreachable!("gemm_supports guaranteed Quant"),
16679        };
16680        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
16681        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
16682        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
16683        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
16684        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
16685        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16686            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16687                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16688                if scale != 1.0 {
16689                    self.scale_inplace(&mut y, scale, m * out_f)?;
16690                }
16691                return Ok(y);
16692            }
16693        }
16694        let name = match qtype {
16695            QT_Q8_0 => "qmatvec_gemm_q8_0",
16696            QT_Q4_K => "qmatvec_gemm_q4_K",
16697            QT_Q4_0 => {
16698                if rp {
16699                    "qmatvec_gemm_q4_0_rp"
16700                } else {
16701                    "qmatvec_gemm_q4_0"
16702                }
16703            }
16704            QT_Q5_K => "qmatvec_gemm_q5_K",
16705            QT_Q6_K => "qmatvec_gemm_q6_K",
16706            QT_NVFP4 => {
16707                if rp {
16708                    "qmatvec_gemm_nvfp4_rp"
16709                } else {
16710                    "qmatvec_gemm_nvfp4"
16711                }
16712            }
16713            _ => unreachable!(),
16714        };
16715        let f = self.func(name);
16716        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16717        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
16718        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
16719        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
16720        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16721        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16722        let k1_tile = if is_k1 {
16723            k1_launch_override().unwrap_or((128, 128, 8))
16724        } else {
16725            (128, 128, 8)
16726        };
16727        let (bm, bn): (u32, u32) = if is_k1 {
16728            (k1_tile.0, k1_tile.1)
16729        } else {
16730            (64, 256)
16731        };
16732        let warps: u32 = if is_k1 {
16733            k1_tile.2
16734        } else {
16735            match qtype {
16736                QT_NVFP4 => 8,
16737                _ => 4,
16738            }
16739        };
16740        let cfg = LaunchConfig {
16741            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16742            block_dim: (32, warps, 1),
16743            shared_mem_bytes: 0,
16744        };
16745        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16746        let __s_b = self.gpu.stream();
16747        let mut b = __s_b.launch_builder(&f);
16748        b.arg(bytes)
16749            .arg(aq)
16750            .arg(ad)
16751            .arg(&mut y)
16752            .arg(&inf)
16753            .arg(&outf)
16754            .arg(&mi)
16755            .arg(&rb);
16756        unsafe {
16757            b.launch(cfg)?;
16758        }
16759        if scale != 1.0 {
16760            self.scale_inplace(&mut y, scale, m * out_f)?;
16761        }
16762        Ok(y)
16763    }
16764
16765    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
16766    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
16767    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
16768    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
16769    pub fn qmatvec_gemm_raw(
16770        &self,
16771        bytes: &CudaSlice<u8>,
16772        x: &CudaSlice<f32>,
16773        m: usize,
16774        in_f: usize,
16775        out_f: usize,
16776        qtype: i32,
16777        row_bytes: usize,
16778    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16779        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16780        let name = match qtype {
16781            QT_Q8_0 => "qmatvec_gemm_q8_0",
16782            QT_Q4_K => "qmatvec_gemm_q4_K",
16783            QT_Q4_0 => "qmatvec_gemm_q4_0",
16784            QT_Q5_K => "qmatvec_gemm_q5_K",
16785            QT_Q6_K => "qmatvec_gemm_q6_K",
16786            QT_NVFP4 => "qmatvec_gemm_nvfp4",
16787            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16788            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16789        };
16790        let f = self.func(name);
16791        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16792        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
16793        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
16794        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16795        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16796        let k1_tile = if is_k1 {
16797            k1_launch_override().unwrap_or((128, 128, 8))
16798        } else {
16799            (128, 128, 8)
16800        };
16801        let (bm, bn): (u32, u32) = if is_k1 {
16802            (k1_tile.0, k1_tile.1)
16803        } else {
16804            (64, 256)
16805        };
16806        let warps: u32 = if is_k1 {
16807            k1_tile.2
16808        } else {
16809            match qtype {
16810                QT_NVFP4 | QT_NVFP4_RP => 8,
16811                _ => 4,
16812            }
16813        };
16814        let cfg = LaunchConfig {
16815            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16816            block_dim: (32, warps, 1),
16817            shared_mem_bytes: 0,
16818        };
16819        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16820        let __s_b = self.gpu.stream();
16821        let mut b = __s_b.launch_builder(&f);
16822        b.arg(bytes)
16823            .arg(&aq)
16824            .arg(&ad)
16825            .arg(&mut y)
16826            .arg(&inf)
16827            .arg(&outf)
16828            .arg(&mi)
16829            .arg(&rb);
16830        unsafe {
16831            b.launch(cfg)?;
16832        }
16833        Ok(y)
16834    }
16835
16836    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
16837    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
16838    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
16839    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
16840    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
16841    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
16842    pub fn qmatvec_gemm_q8_0_wgmma_raw(
16843        &self,
16844        rp4: &CudaSlice<u8>,
16845        aq: &CudaSlice<i8>,
16846        ad: &CudaSlice<f32>,
16847        m: usize,
16848        in_f: usize,
16849        out_f: usize,
16850    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16851        assert!(
16852            out_f % 64 == 0 && in_f % 32 == 0,
16853            "wgmma GEMM needs out_f%64==0, in_f%32==0"
16854        );
16855        let f = self.func("qmatvec_gemm_q8_0_wgmma");
16856        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
16857        let cfg = LaunchConfig {
16858            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16859            block_dim: (128, 1, 1),
16860            shared_mem_bytes: 0,
16861        };
16862        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16863        let __s_b = self.gpu.stream();
16864        let mut b = __s_b.launch_builder(&f);
16865        b.arg(rp4)
16866            .arg(aq)
16867            .arg(ad)
16868            .arg(&mut y)
16869            .arg(&inf)
16870            .arg(&outf)
16871            .arg(&mi);
16872        unsafe {
16873            b.launch(cfg)?;
16874        }
16875        Ok(y)
16876    }
16877
16878    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
16879    pub fn scale_inplace(
16880        &self,
16881        y: &mut CudaSlice<f32>,
16882        s: f32,
16883        n: usize,
16884    ) -> Result<(), Box<dyn std::error::Error>> {
16885        let f = self.func("scale_f32");
16886        let cfg = LaunchConfig::for_num_elems(n as u32);
16887        let (sf, ni) = (s, n as i32);
16888        let __s_b = self.gpu.stream();
16889        let mut b = __s_b.launch_builder(&f);
16890        b.arg(y).arg(&sf).arg(&ni);
16891        unsafe {
16892            b.launch(cfg)?;
16893        }
16894        Ok(())
16895    }
16896
16897    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
16898    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
16899    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
16900    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
16901    pub fn bf16_to_f32(
16902        &self,
16903        data: &cudarc::driver::CudaView<'_, u8>,
16904        n: usize,
16905    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16906        let mut out = self.alloc_uninit::<f32>(n)?;
16907        let f = self.func("bf16_to_f32");
16908        let cfg = LaunchConfig::for_num_elems(n as u32);
16909        let ni = n as i32;
16910        let __s_b = self.gpu.stream();
16911        let mut b = __s_b.launch_builder(&f);
16912        b.arg(data).arg(&mut out).arg(&ni);
16913        unsafe {
16914            b.launch(cfg)?;
16915        }
16916        Ok(out)
16917    }
16918
16919    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
16920    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
16921    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
16922    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
16923    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
16924    /// calls, the spec-verify contract) vs plain linear.
16925    fn linear_bf16_chunked(
16926        &self,
16927        x: &CudaSlice<f32>,
16928        data: &CudaSlice<u8>,
16929        m: usize,
16930        in_f: usize,
16931        out_f: usize,
16932        exact: bool,
16933        canonical_chunk_rows: Option<usize>,
16934    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16935        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
16936        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
16937        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
16938        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16939        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16940        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16941        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16942        let started = timing.then(std::time::Instant::now);
16943        let result =
16944            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16945        if let Some(started) = started {
16946            use std::sync::atomic::Ordering;
16947            self.stream().synchronize()?;
16948            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16949                + started.elapsed().as_nanos() as u64;
16950            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16951                + (in_f * out_f * 2) as u64;
16952            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16953            if calls % 1024 == 0 {
16954                eprintln!(
16955                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16956                     weight_gb={:.2}",
16957                    ns as f64 / 1.0e6,
16958                    ns as f64 / calls as f64 / 1.0e3,
16959                    wb as f64 / 1.0e9,
16960                );
16961            }
16962        }
16963        result
16964    }
16965
16966    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
16967    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
16968    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
16969    /// numeric-class doors (DEV_ROUTES precedent).
16970    pub(crate) fn bf16_mmv_on() -> bool {
16971        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16972        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16973    }
16974
16975    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
16976    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
16977    fn matvec_bf16(
16978        &self,
16979        data: &CudaSlice<u8>,
16980        x: &CudaSlice<f32>,
16981        in_f: usize,
16982        out_f: usize,
16983    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16984        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16985            return Err(format!(
16986                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16987                data.len(),
16988                x.len()
16989            )
16990            .into());
16991        }
16992        let mut y = self.alloc_uninit::<f32>(out_f)?;
16993        let f = self.func("matvec_bf16_f32acc");
16994        let cfg = LaunchConfig {
16995            grid_dim: (out_f as u32, 1, 1),
16996            block_dim: (mmv_block(), 1, 1),
16997            shared_mem_bytes: 0,
16998        };
16999        let ini = in_f as i32;
17000        let __s_bld = self.gpu.stream();
17001        let mut bld = __s_bld.launch_builder(&f);
17002        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17003        unsafe {
17004            bld.launch(cfg)?;
17005        }
17006        Ok(y)
17007    }
17008
17009    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
17010    /// launches, a position upload, and the rope launch; the position is read directly from
17011    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
17012    #[allow(clippy::too_many_arguments)]
17013    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
17014    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
17015    /// Bit-identical to the split kernels; requires head_dim == 128 and
17016    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
17017    #[allow(clippy::too_many_arguments)]
17018    pub fn qk_norm_rope_append_inc_dcw(
17019        &self,
17020        q_raw: &CudaSlice<f32>,
17021        k_raw: &CudaSlice<f32>,
17022        v_raw: &CudaSlice<f32>,
17023        qw: &CudaSlice<f32>,
17024        kw: &CudaSlice<f32>,
17025        q_out: &mut CudaSlice<f32>,
17026        k_out: &mut CudaSlice<f32>,
17027        pos: &CudaSlice<i32>,
17028        k_plane: &mut CudaSlice<u8>,
17029        v_plane: &mut CudaSlice<u8>,
17030        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
17031        // (single) writer, exactly like the split append+inc pair it replaces.
17032        len_dev: &CudaSlice<i32>,
17033        base_dev: Option<&CudaSlice<i32>>,
17034        done_ctr: &mut CudaSlice<u32>,
17035        kv_dim_k: usize,
17036        kv_dim_v: usize,
17037        k_tok_bytes: usize,
17038        v_tok_bytes: usize,
17039        head_dim: usize,
17040        n_dims: usize,
17041        nh_q: usize,
17042        nh_k: usize,
17043        eps: f32,
17044        freq_base: f32,
17045        freq_scale: f32,
17046        ff: Option<&CudaSlice<f32>>,
17047    ) -> Result<(), Box<dyn std::error::Error>> {
17048        if head_dim != 128
17049            || kv_dim_v != kv_dim_k
17050            || kv_dim_k != nh_k * head_dim
17051            || q_raw.len() < nh_q * head_dim
17052            || k_raw.len() < nh_k * head_dim
17053            || v_raw.len() < kv_dim_v
17054            || q_out.len() < nh_q * head_dim
17055            || k_out.len() < nh_k * head_dim
17056            || pos.is_empty()
17057            || done_ctr.is_empty()
17058        {
17059            return Err(format!(
17060                "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}"
17061            )
17062            .into());
17063        }
17064        let f = self.func("qk_norm_rope_append_inc_dcw");
17065        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17066        let cfg = LaunchConfig {
17067            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17068            block_dim: (128, 1, 1),
17069            shared_mem_bytes: 0,
17070        };
17071        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17072        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17073        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17074        let null: u64 = 0;
17075        let __s_b = self.gpu.stream();
17076        let mut b = __s_b.launch_builder(&f);
17077        b.arg(q_raw)
17078            .arg(k_raw)
17079            .arg(v_raw)
17080            .arg(qw)
17081            .arg(kw)
17082            .arg(q_out)
17083            .arg(k_out)
17084            .arg(pos)
17085            .arg(&mut *k_plane)
17086            .arg(&mut *v_plane)
17087            .arg(len_dev);
17088        match base_dev {
17089            Some(base) => {
17090                b.arg(base);
17091            }
17092            None => {
17093                b.arg(&null);
17094            }
17095        }
17096        b.arg(&mut *done_ctr)
17097            .arg(&kvk)
17098            .arg(&kvv)
17099            .arg(&ktb)
17100            .arg(&vtb)
17101            .arg(&hd)
17102            .arg(&nd)
17103            .arg(&nq)
17104            .arg(&eps)
17105            .arg(&theta_scale)
17106            .arg(&freq_scale);
17107        match ff {
17108            Some(freqs) => {
17109                b.arg(freqs);
17110            }
17111            None => {
17112                b.arg(&null);
17113            }
17114        }
17115        unsafe {
17116            b.launch(cfg)?;
17117        }
17118        Ok(())
17119    }
17120
17121    pub fn qk_norm_rope_into(
17122        &self,
17123        q_raw: &CudaSlice<f32>,
17124        k_raw: &CudaSlice<f32>,
17125        qw: &CudaSlice<f32>,
17126        kw: &CudaSlice<f32>,
17127        q_out: &mut CudaSlice<f32>,
17128        k_out: &mut CudaSlice<f32>,
17129        pos: &CudaSlice<i32>,
17130        head_dim: usize,
17131        n_dims: usize,
17132        nh_q: usize,
17133        nh_k: usize,
17134        eps: f32,
17135        freq_base: f32,
17136        freq_scale: f32,
17137        ff: Option<&CudaSlice<f32>>,
17138    ) -> Result<(), Box<dyn std::error::Error>> {
17139        if head_dim > 512
17140            || q_raw.len() < nh_q * head_dim
17141            || k_raw.len() < nh_k * head_dim
17142            || q_out.len() < nh_q * head_dim
17143            || k_out.len() < nh_k * head_dim
17144            || qw.len() < head_dim
17145            || kw.len() < head_dim
17146            || pos.is_empty()
17147        {
17148            return Err(format!(
17149                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17150            )
17151            .into());
17152        }
17153        let f = self.func("qk_norm_rope_f32");
17154        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17155        let cfg = LaunchConfig {
17156            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17157            block_dim: (128, 1, 1),
17158            shared_mem_bytes: 0,
17159        };
17160        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17161        let __s_b = self.gpu.stream();
17162        let mut b = __s_b.launch_builder(&f);
17163        b.arg(q_raw)
17164            .arg(k_raw)
17165            .arg(qw)
17166            .arg(kw)
17167            .arg(q_out)
17168            .arg(k_out)
17169            .arg(pos)
17170            .arg(&hd)
17171            .arg(&nd)
17172            .arg(&nq)
17173            .arg(&eps)
17174            .arg(&theta_scale)
17175            .arg(&freq_scale);
17176        match ff {
17177            Some(ffv) => {
17178                b.arg(ffv);
17179                unsafe {
17180                    b.launch(cfg)?;
17181                }
17182            }
17183            None => {
17184                let null: u64 = 0;
17185                b.arg(&null);
17186                unsafe {
17187                    b.launch(cfg)?;
17188                }
17189            }
17190        }
17191        Ok(())
17192    }
17193
17194    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
17195    /// launch computes a rank's whole O partial from its four canonical column blocks.
17196    #[allow(clippy::too_many_arguments)]
17197    pub fn matvec_f32_b4_into(
17198        &self,
17199        w: [&CudaSlice<f32>; 4],
17200        x: &CudaSlice<f32>,
17201        y: &mut CudaSlice<f32>,
17202        block_cols: usize,
17203        out_f: usize,
17204    ) -> Result<(), Box<dyn std::error::Error>> {
17205        if block_cols % 4 != 0
17206            || x.len() < 4 * block_cols
17207            || y.len() < out_f
17208            || w.iter().any(|w| w.len() != out_f * block_cols)
17209        {
17210            return Err(format!(
17211                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17212                x.len()
17213            )
17214            .into());
17215        }
17216        let f = self.func("matvec_f32_b4");
17217        let cfg = LaunchConfig {
17218            grid_dim: (out_f as u32, 1, 1),
17219            block_dim: (128, 1, 1),
17220            shared_mem_bytes: 0,
17221        };
17222        let (bc, of) = (block_cols as i32, out_f as i32);
17223        let __s_b = self.gpu.stream();
17224        let mut b = __s_b.launch_builder(&f);
17225        b.arg(w[0])
17226            .arg(w[1])
17227            .arg(w[2])
17228            .arg(w[3])
17229            .arg(x)
17230            .arg(y)
17231            .arg(&bc)
17232            .arg(&of);
17233        unsafe {
17234            b.launch(cfg)?;
17235        }
17236        Ok(())
17237    }
17238
17239    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
17240    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
17241    pub fn axpy_rows_seq_into(
17242        &self,
17243        x: &CudaSlice<f32>,
17244        w: &CudaSlice<f32>,
17245        y: &mut CudaSlice<f32>,
17246        width: usize,
17247        n_rows: usize,
17248    ) -> Result<(), Box<dyn std::error::Error>> {
17249        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17250            return Err(format!(
17251                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17252                x.len(),
17253                w.len(),
17254                y.len()
17255            )
17256            .into());
17257        }
17258        let f = self.func("axpy_rows_seq_f32");
17259        let cfg = LaunchConfig::for_num_elems(width as u32);
17260        let (wi, nr) = (width as i32, n_rows as i32);
17261        let __s_b = self.gpu.stream();
17262        let mut b = __s_b.launch_builder(&f);
17263        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17264        unsafe {
17265            b.launch(cfg)?;
17266        }
17267        Ok(())
17268    }
17269
17270    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
17271    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
17272    /// exact sequential FP chain of the base kernel over that window.
17273    #[allow(clippy::too_many_arguments)]
17274    pub fn axpy_rows_seq_md_off_into(
17275        &self,
17276        x: &CudaSlice<f32>,
17277        w_route: &CudaSlice<f32>,
17278        md: &CudaSlice<f32>,
17279        sel: &CudaSlice<i32>,
17280        y: &mut CudaSlice<f32>,
17281        width: usize,
17282        n_rows: usize,
17283        row0: usize,
17284    ) -> Result<(), Box<dyn std::error::Error>> {
17285        if x.len() < (row0 + n_rows) * width
17286            || w_route.len() < row0 + n_rows
17287            || sel.len() < row0 + n_rows
17288            || y.len() < width
17289        {
17290            return Err(format!(
17291                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17292                 rows={n_rows} row0={row0}",
17293                x.len(),
17294                w_route.len(),
17295                sel.len(),
17296                y.len()
17297            )
17298            .into());
17299        }
17300        let f = self.func("axpy_rows_seq_md_off_f32");
17301        let cfg = LaunchConfig::for_num_elems(width as u32);
17302        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17303        let __s_b = self.gpu.stream();
17304        let mut b = __s_b.launch_builder(&f);
17305        b.arg(x)
17306            .arg(w_route)
17307            .arg(md)
17308            .arg(sel)
17309            .arg(y)
17310            .arg(&wi)
17311            .arg(&nr)
17312            .arg(&r0);
17313        unsafe {
17314            b.launch(cfg)?;
17315        }
17316        Ok(())
17317    }
17318
17319    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
17320    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
17321    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
17322    /// outputs are bit-equal to its own t=1 launch.
17323    #[allow(clippy::too_many_arguments)]
17324    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17325        &self,
17326        gate_bank: &CudaSlice<u8>,
17327        up_bank: &CudaSlice<u8>,
17328        sel: &CudaSlice<i32>,
17329        aq: &CudaSlice<i8>,
17330        ad: &CudaSlice<f32>,
17331        yg: &mut CudaSlice<f32>,
17332        yu: &mut CudaSlice<f32>,
17333        n_sel: usize,
17334        n_sel_col: usize,
17335        in_f: usize,
17336        out_f: usize,
17337        row_bytes: usize,
17338        expert_stride: usize,
17339        act_row_stride: usize,
17340        ad_row_stride: usize,
17341    ) -> Result<(), Box<dyn std::error::Error>> {
17342        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17343        if yg.len() < n_sel * out_f
17344            || yu.len() < n_sel * out_f
17345            || sel.len() < n_sel
17346            || n_sel_col == 0
17347            || n_sel % n_sel_col != 0
17348        {
17349            return Err("NVFP4 gu tcol geometry".into());
17350        }
17351        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17352        let cfg = LaunchConfig {
17353            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17354            block_dim: (128, 1, 1),
17355            shared_mem_bytes: 0,
17356        };
17357        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17358        let (rb, es) = (row_bytes as i64, expert_stride as i64);
17359        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17360        let __s_b = self.gpu.stream();
17361        let mut b = __s_b.launch_builder(&f);
17362        b.arg(gate_bank)
17363            .arg(up_bank)
17364            .arg(sel)
17365            .arg(aq)
17366            .arg(ad)
17367            .arg(yg)
17368            .arg(yu)
17369            .arg(&inf)
17370            .arg(&outf)
17371            .arg(&ns)
17372            .arg(&rb)
17373            .arg(&es)
17374            .arg(&ars)
17375            .arg(&adrs)
17376            .arg(&nsc);
17377        unsafe {
17378            b.launch(cfg)?;
17379        }
17380        Ok(())
17381    }
17382
17383    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
17384    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
17385    #[allow(clippy::too_many_arguments)]
17386    pub fn axpy_rows_seq_md_into(
17387        &self,
17388        x: &CudaSlice<f32>,
17389        w_route: &CudaSlice<f32>,
17390        md: &CudaSlice<f32>,
17391        sel: &CudaSlice<i32>,
17392        y: &mut CudaSlice<f32>,
17393        width: usize,
17394        n_rows: usize,
17395    ) -> Result<(), Box<dyn std::error::Error>> {
17396        if x.len() < n_rows * width
17397            || w_route.len() < n_rows
17398            || sel.len() < n_rows
17399            || y.len() < width
17400        {
17401            return Err(format!(
17402                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17403                x.len(),
17404                w_route.len(),
17405                sel.len(),
17406                y.len()
17407            )
17408            .into());
17409        }
17410        let f = self.func("axpy_rows_seq_md_f32");
17411        let cfg = LaunchConfig::for_num_elems(width as u32);
17412        let (wi, nr) = (width as i32, n_rows as i32);
17413        let __s_b = self.gpu.stream();
17414        let mut b = __s_b.launch_builder(&f);
17415        b.arg(x)
17416            .arg(w_route)
17417            .arg(md)
17418            .arg(sel)
17419            .arg(y)
17420            .arg(&wi)
17421            .arg(&nr);
17422        unsafe {
17423            b.launch(cfg)?;
17424        }
17425        Ok(())
17426    }
17427
17428    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
17429    #[allow(clippy::too_many_arguments)]
17430    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
17431    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
17432    /// land column-major-of-rows: yq[c*out_q + row] etc.
17433    #[allow(clippy::too_many_arguments)]
17434    pub fn matvec_bf16_qkvg_tcol_into(
17435        &self,
17436        wq: &CudaSlice<u8>,
17437        wk: &CudaSlice<u8>,
17438        wv: &CudaSlice<u8>,
17439        wg: &CudaSlice<u8>,
17440        x_t: &CudaSlice<f32>,
17441        yq: &mut CudaSlice<f32>,
17442        yk: &mut CudaSlice<f32>,
17443        yv: &mut CudaSlice<f32>,
17444        yg: &mut CudaSlice<f32>,
17445        in_f: usize,
17446        out_q: usize,
17447        out_kv: usize,
17448        out_g: usize,
17449        t: usize,
17450    ) -> Result<(), Box<dyn std::error::Error>> {
17451        if t == 0
17452            || t > 8
17453            || in_f % 8 != 0
17454            || x_t.len() < t * in_f
17455            || yq.len() < t * out_q
17456            || yk.len() < t * out_kv
17457            || yv.len() < t * out_kv
17458            || (out_g > 0 && yg.len() < t * out_g)
17459        {
17460            return Err("matvec_bf16_qkvg_tcol geometry".into());
17461        }
17462        let f = self.func("matvec_bf16_qkvg_tcol");
17463        let grid = out_q + 2 * out_kv + out_g;
17464        let cfg = LaunchConfig {
17465            grid_dim: (grid as u32, 1, 1),
17466            block_dim: (mmv_block(), 1, 1),
17467            shared_mem_bytes: 0,
17468        };
17469        let (ini, oq, okv, og, ti) = (
17470            in_f as i32,
17471            out_q as i32,
17472            out_kv as i32,
17473            out_g as i32,
17474            t as i32,
17475        );
17476        let __s_b = self.gpu.stream();
17477        let mut b = __s_b.launch_builder(&f);
17478        b.arg(wq)
17479            .arg(wk)
17480            .arg(wv)
17481            .arg(wg)
17482            .arg(x_t)
17483            .arg(yq)
17484            .arg(yk)
17485            .arg(yv)
17486            .arg(yg)
17487            .arg(&ini)
17488            .arg(&oq)
17489            .arg(&okv)
17490            .arg(&og)
17491            .arg(&ti);
17492        unsafe {
17493            b.launch(cfg)?;
17494        }
17495        Ok(())
17496    }
17497
17498    pub fn matvec_bf16_qkvg_into(
17499        &self,
17500        wq: &CudaSlice<u8>,
17501        wk: &CudaSlice<u8>,
17502        wv: &CudaSlice<u8>,
17503        wg: &CudaSlice<u8>,
17504        x: &CudaSlice<f32>,
17505        yq: &mut CudaSlice<f32>,
17506        yk: &mut CudaSlice<f32>,
17507        yv: &mut CudaSlice<f32>,
17508        yg: &mut CudaSlice<f32>,
17509        in_f: usize,
17510        out_q: usize,
17511        out_kv: usize,
17512        out_g: usize,
17513    ) -> Result<(), Box<dyn std::error::Error>> {
17514        if in_f % 8 != 0
17515            || wq.len() != out_q * in_f * 2
17516            || wk.len() != out_kv * in_f * 2
17517            || wv.len() != out_kv * in_f * 2
17518            || wg.len() < out_g * in_f * 2
17519            || x.len() < in_f
17520            || yq.len() < out_q
17521            || yk.len() < out_kv
17522            || yv.len() < out_kv
17523            || (out_g > 0 && yg.len() < out_g)
17524        {
17525            return Err(format!(
17526                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17527            )
17528            .into());
17529        }
17530        let f = self.func("matvec_bf16_qkvg");
17531        let cfg = LaunchConfig {
17532            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17533            block_dim: (mmv_block(), 1, 1),
17534            shared_mem_bytes: 0,
17535        };
17536        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17537        let __s_b = self.gpu.stream();
17538        let mut b = __s_b.launch_builder(&f);
17539        b.arg(wq)
17540            .arg(wk)
17541            .arg(wv)
17542            .arg(wg)
17543            .arg(x)
17544            .arg(yq)
17545            .arg(yk)
17546            .arg(yv)
17547            .arg(yg)
17548            .arg(&inf)
17549            .arg(&oq)
17550            .arg(&okv)
17551            .arg(&og);
17552        unsafe {
17553            b.launch(cfg)?;
17554        }
17555        Ok(())
17556    }
17557
17558    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
17559    pub fn matvec_bf16_b4_into(
17560        &self,
17561        w: [&CudaSlice<u8>; 4],
17562        x: &CudaSlice<f32>,
17563        y: &mut CudaSlice<f32>,
17564        block_cols: usize,
17565        out_f: usize,
17566    ) -> Result<(), Box<dyn std::error::Error>> {
17567        if block_cols % 8 != 0
17568            || x.len() < 4 * block_cols
17569            || y.len() < out_f
17570            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17571        {
17572            return Err(format!(
17573                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17574                x.len()
17575            )
17576            .into());
17577        }
17578        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
17579        // bit-identical per row (the second row's stream hides the first's reduce tail).
17580        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17581        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17582        let f = self.func(if x2 {
17583            "matvec_bf16_b4_x2"
17584        } else {
17585            "matvec_bf16_b4"
17586        });
17587        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17588        let cfg = LaunchConfig {
17589            grid_dim: (grid as u32, 1, 1),
17590            block_dim: (mmv_block(), 1, 1),
17591            shared_mem_bytes: 0,
17592        };
17593        let (bc, of) = (block_cols as i32, out_f as i32);
17594        let __s_b = self.gpu.stream();
17595        let mut b = __s_b.launch_builder(&f);
17596        b.arg(w[0])
17597            .arg(w[1])
17598            .arg(w[2])
17599            .arg(w[3])
17600            .arg(x)
17601            .arg(y)
17602            .arg(&bc)
17603            .arg(&of);
17604        unsafe {
17605            b.launch(cfg)?;
17606        }
17607        Ok(())
17608    }
17609
17610    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
17611    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
17612    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
17613    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
17614    /// t=1 program).
17615    pub fn matvec_bf16_b4_tcol_into(
17616        &self,
17617        w: [&CudaSlice<u8>; 4],
17618        x_t: &CudaSlice<f32>,
17619        y_t: &mut CudaSlice<f32>,
17620        block_cols: usize,
17621        out_f: usize,
17622        t: usize,
17623    ) -> Result<(), Box<dyn std::error::Error>> {
17624        if block_cols % 8 != 0
17625            || t == 0
17626            || t > 8
17627            || x_t.len() < t * 4 * block_cols
17628            || y_t.len() < t * out_f
17629            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17630        {
17631            return Err(format!(
17632                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17633                x_t.len()
17634            )
17635            .into());
17636        }
17637        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17638            return Err(
17639                "b4 tcol verify is qualified against the plain b4 kernel only \
17640                        (MEMRA_B4_X2=1 is a different t=1 program)"
17641                    .into(),
17642            );
17643        }
17644        let f = self.func("matvec_bf16_b4_tcol");
17645        let cfg = LaunchConfig {
17646            grid_dim: (out_f as u32, 1, 1),
17647            block_dim: (mmv_block(), 1, 1),
17648            shared_mem_bytes: 0,
17649        };
17650        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17651        let __s_b = self.gpu.stream();
17652        let mut b = __s_b.launch_builder(&f);
17653        b.arg(w[0])
17654            .arg(w[1])
17655            .arg(w[2])
17656            .arg(w[3])
17657            .arg(x_t)
17658            .arg(y_t)
17659            .arg(&bc)
17660            .arg(&of)
17661            .arg(&ti);
17662        unsafe {
17663            b.launch(cfg)?;
17664        }
17665        Ok(())
17666    }
17667
17668    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
17669    pub fn matvec_bf16_into(
17670        &self,
17671        data: &CudaSlice<u8>,
17672        x: &CudaSlice<f32>,
17673        y: &mut CudaSlice<f32>,
17674        in_f: usize,
17675        out_f: usize,
17676    ) -> Result<(), Box<dyn std::error::Error>> {
17677        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17678            return Err(format!(
17679                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17680                data.len(),
17681                x.len(),
17682                y.len()
17683            )
17684            .into());
17685        }
17686        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
17687        // block, exact f32acc per-row program — cures the 1-iteration latency
17688        // starvation (shexp down measured 420GB/s at in_f=1280).
17689        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17690        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17691            && in_f <= 2048;
17692        if x4 {
17693            let f = self.func("matvec_bf16_f32acc_x4");
17694            let cfg = LaunchConfig {
17695                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17696                block_dim: (mmv_block(), 1, 1),
17697                shared_mem_bytes: 0,
17698            };
17699            let (ini, outi) = (in_f as i32, out_f as i32);
17700            let __s_b = self.gpu.stream();
17701            let mut b = __s_b.launch_builder(&f);
17702            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17703            unsafe {
17704                b.launch(cfg)?;
17705            }
17706            return Ok(());
17707        }
17708        let f = self.func("matvec_bf16_f32acc");
17709        let cfg = LaunchConfig {
17710            grid_dim: (out_f as u32, 1, 1),
17711            block_dim: (mmv_block(), 1, 1),
17712            shared_mem_bytes: 0,
17713        };
17714        let ini = in_f as i32;
17715        let __s_b = self.gpu.stream();
17716        let mut b = __s_b.launch_builder(&f);
17717        b.arg(data).arg(x).arg(y).arg(&ini);
17718        unsafe {
17719            b.launch(cfg)?;
17720        }
17721        Ok(())
17722    }
17723
17724    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
17725    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
17726    pub fn matvec_bf16_view_into(
17727        &self,
17728        data: &cudarc::driver::CudaView<'_, u8>,
17729        x: &CudaSlice<f32>,
17730        y: &mut CudaSlice<f32>,
17731        in_f: usize,
17732        out_f: usize,
17733    ) -> Result<(), Box<dyn std::error::Error>> {
17734        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17735            return Err(format!(
17736                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17737                data.len(),
17738                x.len(),
17739                y.len()
17740            )
17741            .into());
17742        }
17743        let f = self.func("matvec_bf16_f32acc");
17744        let cfg = LaunchConfig {
17745            grid_dim: (out_f as u32, 1, 1),
17746            block_dim: (mmv_block(), 1, 1),
17747            shared_mem_bytes: 0,
17748        };
17749        let ini = in_f as i32;
17750        let __s_b = self.gpu.stream();
17751        let mut b = __s_b.launch_builder(&f);
17752        b.arg(data).arg(x).arg(y).arg(&ini);
17753        unsafe {
17754            b.launch(cfg)?;
17755        }
17756        Ok(())
17757    }
17758
17759    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
17760    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
17761    pub fn matvec_bf16_raw_out(
17762        &self,
17763        w: &CudaSlice<u8>,
17764        x: &CudaSlice<f32>,
17765        y_raw: u64,
17766        in_f: usize,
17767        out_f: usize,
17768    ) -> Result<(), Box<dyn std::error::Error>> {
17769        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17770            return Err("matvec_bf16_raw_out geometry".into());
17771        }
17772        let f = self.func("matvec_bf16_f32acc");
17773        let cfg = LaunchConfig {
17774            grid_dim: (out_f as u32, 1, 1),
17775            block_dim: (mmv_block(), 1, 1),
17776            shared_mem_bytes: 0,
17777        };
17778        let ini = in_f as i32;
17779        let __s_b = self.gpu.stream();
17780        let mut b = __s_b.launch_builder(&f);
17781        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17782        unsafe {
17783            b.launch(cfg)?;
17784        }
17785        Ok(())
17786    }
17787
17788    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
17789    /// UVA pointers so the caller passes persistent-static rows without holding locks).
17790    /// Exact per-element sequence of the split add + add_scaled_rows pair.
17791    pub fn add3_raw(
17792        &self,
17793        a: &CudaSlice<f32>,
17794        b: &CudaSlice<f32>,
17795        sh_raw: u64,
17796        scale_raw: u64,
17797        dst: &mut CudaSlice<f32>,
17798        n: usize,
17799    ) -> Result<(), Box<dyn std::error::Error>> {
17800        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17801            return Err("add3_raw geometry".into());
17802        }
17803        let f = self.func("add3_f32");
17804        let cfg = LaunchConfig {
17805            grid_dim: ((n as u32).div_ceil(256), 1, 1),
17806            block_dim: (256, 1, 1),
17807            shared_mem_bytes: 0,
17808        };
17809        let ni = n as i32;
17810        let __s_b = self.gpu.stream();
17811        let mut bld = __s_b.launch_builder(&f);
17812        bld.arg(a)
17813            .arg(b)
17814            .arg(&sh_raw)
17815            .arg(&scale_raw)
17816            .arg(dst)
17817            .arg(&ni);
17818        unsafe {
17819            bld.launch(cfg)?;
17820        }
17821        Ok(())
17822    }
17823
17824    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
17825    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
17826    pub fn matvec_bf16_down_addscale_into(
17827        &self,
17828        w: &CudaSlice<u8>,
17829        x: &CudaSlice<f32>,
17830        scale: &CudaSlice<f32>,
17831        dst: &mut CudaSlice<f32>,
17832        in_f: usize,
17833        out_f: usize,
17834    ) -> Result<(), Box<dyn std::error::Error>> {
17835        if w.len() != in_f * out_f * 2
17836            || x.len() < in_f
17837            || in_f % 8 != 0
17838            || dst.len() < out_f
17839            || scale.is_empty()
17840        {
17841            return Err("matvec_bf16_down_addscale geometry".into());
17842        }
17843        let f = self.func("matvec_bf16_down_addscale");
17844        let cfg = LaunchConfig {
17845            grid_dim: (out_f as u32, 1, 1),
17846            block_dim: (mmv_block(), 1, 1),
17847            shared_mem_bytes: 0,
17848        };
17849        let ini = in_f as i32;
17850        let __s_b = self.gpu.stream();
17851        let mut b = __s_b.launch_builder(&f);
17852        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17853        unsafe {
17854            b.launch(cfg)?;
17855        }
17856        Ok(())
17857    }
17858
17859    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
17860    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
17861    pub fn matvec_bf16_dual_silu_into(
17862        &self,
17863        wg: &CudaSlice<u8>,
17864        wu: &CudaSlice<u8>,
17865        x: &CudaSlice<f32>,
17866        act: &mut CudaSlice<f32>,
17867        in_f: usize,
17868        out_f: usize,
17869        limit: Option<f32>,
17870    ) -> Result<(), Box<dyn std::error::Error>> {
17871        if wg.len() != in_f * out_f * 2
17872            || wu.len() != in_f * out_f * 2
17873            || x.len() < in_f
17874            || in_f % 8 != 0
17875            || act.len() < out_f
17876        {
17877            return Err("matvec_bf16_dual_silu geometry".into());
17878        }
17879        let f = self.func("matvec_bf16_dual_silu");
17880        let cfg = LaunchConfig {
17881            grid_dim: (out_f as u32, 1, 1),
17882            block_dim: (mmv_block(), 1, 1),
17883            shared_mem_bytes: 0,
17884        };
17885        let (ini, outi) = (in_f as i32, out_f as i32);
17886        let lim = limit.unwrap_or(0.0);
17887        let __s_b = self.gpu.stream();
17888        let mut b = __s_b.launch_builder(&f);
17889        b.arg(wg)
17890            .arg(wu)
17891            .arg(x)
17892            .arg(act)
17893            .arg(&ini)
17894            .arg(&outi)
17895            .arg(&lim);
17896        unsafe {
17897            b.launch(cfg)?;
17898        }
17899        Ok(())
17900    }
17901
17902    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
17903    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
17904    #[allow(clippy::too_many_arguments)]
17905    pub fn matvec_bf16_dual_view_into(
17906        &self,
17907        wg: &cudarc::driver::CudaView<'_, u8>,
17908        wu: &cudarc::driver::CudaView<'_, u8>,
17909        x: &CudaSlice<f32>,
17910        yg: &mut CudaSlice<f32>,
17911        yu: &mut CudaSlice<f32>,
17912        in_f: usize,
17913        out_f: usize,
17914    ) -> Result<(), Box<dyn std::error::Error>> {
17915        if wg.len() != in_f * out_f * 2
17916            || wu.len() != in_f * out_f * 2
17917            || x.len() < in_f
17918            || in_f % 8 != 0
17919            || yg.len() < out_f
17920            || yu.len() < out_f
17921        {
17922            return Err(format!(
17923                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17924                wg.len(),
17925                wu.len(),
17926                x.len()
17927            )
17928            .into());
17929        }
17930        let f = self.func("matvec_bf16_dual");
17931        let cfg = LaunchConfig {
17932            grid_dim: ((2 * out_f) as u32, 1, 1),
17933            block_dim: (mmv_block(), 1, 1),
17934            shared_mem_bytes: 0,
17935        };
17936        let (ini, outi) = (in_f as i32, out_f as i32);
17937        let __s_b = self.gpu.stream();
17938        let mut b = __s_b.launch_builder(&f);
17939        b.arg(wg)
17940            .arg(wu)
17941            .arg(x)
17942            .arg(yg)
17943            .arg(yu)
17944            .arg(&ini)
17945            .arg(&outi);
17946        unsafe {
17947            b.launch(cfg)?;
17948        }
17949        Ok(())
17950    }
17951
17952    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
17953    #[allow(clippy::too_many_arguments)]
17954    pub fn matvec_bf16_dual_into(
17955        &self,
17956        wg: &CudaSlice<u8>,
17957        wu: &CudaSlice<u8>,
17958        x: &CudaSlice<f32>,
17959        yg: &mut CudaSlice<f32>,
17960        yu: &mut CudaSlice<f32>,
17961        in_f: usize,
17962        out_f: usize,
17963    ) -> Result<(), Box<dyn std::error::Error>> {
17964        if wg.len() != in_f * out_f * 2
17965            || wu.len() != in_f * out_f * 2
17966            || x.len() < in_f
17967            || in_f % 8 != 0
17968            || yg.len() < out_f
17969            || yu.len() < out_f
17970        {
17971            return Err(format!(
17972                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17973                wg.len(),
17974                wu.len(),
17975                x.len()
17976            )
17977            .into());
17978        }
17979        let f = self.func("matvec_bf16_dual");
17980        let cfg = LaunchConfig {
17981            grid_dim: ((2 * out_f) as u32, 1, 1),
17982            block_dim: (mmv_block(), 1, 1),
17983            shared_mem_bytes: 0,
17984        };
17985        let (ini, outi) = (in_f as i32, out_f as i32);
17986        let __s_b = self.gpu.stream();
17987        let mut b = __s_b.launch_builder(&f);
17988        b.arg(wg)
17989            .arg(wu)
17990            .arg(x)
17991            .arg(yg)
17992            .arg(yu)
17993            .arg(&ini)
17994            .arg(&outi);
17995        unsafe {
17996            b.launch(cfg)?;
17997        }
17998        Ok(())
17999    }
18000
18001    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
18002    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
18003    pub(crate) fn matvec_bf16_dual(
18004        &self,
18005        wg: &CudaSlice<u8>,
18006        wu: &CudaSlice<u8>,
18007        x: &CudaSlice<f32>,
18008        in_f: usize,
18009        out_f: usize,
18010    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18011        if wg.len() != in_f * out_f * 2
18012            || wu.len() != in_f * out_f * 2
18013            || x.len() < in_f
18014            || in_f % 8 != 0
18015        {
18016            return Err(format!(
18017                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18018                wg.len(),
18019                wu.len(),
18020                x.len()
18021            )
18022            .into());
18023        }
18024        let mut yg = self.alloc_uninit::<f32>(out_f)?;
18025        let mut yu = self.alloc_uninit::<f32>(out_f)?;
18026        let f = self.func("matvec_bf16_dual");
18027        let cfg = LaunchConfig {
18028            grid_dim: ((2 * out_f) as u32, 1, 1),
18029            block_dim: (mmv_block(), 1, 1),
18030            shared_mem_bytes: 0,
18031        };
18032        let (ini, outi) = (in_f as i32, out_f as i32);
18033        let __s_b = self.gpu.stream();
18034        let mut b = __s_b.launch_builder(&f);
18035        b.arg(wg)
18036            .arg(wu)
18037            .arg(x)
18038            .arg(&mut yg)
18039            .arg(&mut yu)
18040            .arg(&ini)
18041            .arg(&outi);
18042        unsafe {
18043            b.launch(cfg)?;
18044        }
18045        Ok((yg, yu))
18046    }
18047
18048    #[allow(clippy::too_many_arguments)]
18049    fn linear_bf16_chunked_inner(
18050        &self,
18051        x: &CudaSlice<f32>,
18052        data: &CudaSlice<u8>,
18053        m: usize,
18054        in_f: usize,
18055        out_f: usize,
18056        exact: bool,
18057        canonical_chunk_rows: Option<usize>,
18058    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18059        const CHUNK_BYTES: usize = 256 << 20;
18060        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
18061        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
18062        if m == 1
18063            && !exact
18064            && canonical_chunk_rows.is_none()
18065            && in_f % 8 == 0
18066            && Self::bf16_mmv_on()
18067        {
18068            return self.matvec_bf16(data, x, in_f, out_f);
18069        }
18070        let row_bytes = in_f
18071            .checked_mul(std::mem::size_of::<f32>())
18072            .ok_or("BF16 chunk row byte count overflow")?;
18073        if row_bytes == 0 || out_f == 0 {
18074            return Err("BF16 chunk dimensions must be nonzero".into());
18075        }
18076        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18077        let chunk_rows = match canonical_chunk_rows {
18078            Some(rows) if rows == 0 => {
18079                return Err("canonical BF16 chunk rows must be nonzero".into());
18080            }
18081            Some(rows) if rows > max_chunk_rows => {
18082                return Err(format!(
18083                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18084                )
18085                .into());
18086            }
18087            Some(rows) if out_f % rows != 0 => {
18088                return Err(format!(
18089                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18090                )
18091                .into());
18092            }
18093            Some(rows) => rows,
18094            None => max_chunk_rows,
18095        };
18096        if chunk_rows >= out_f {
18097            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18098            return if exact {
18099                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
18100            } else {
18101                self.linear(x, &wf32, m, in_f, out_f)
18102            };
18103        }
18104        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18105        let mut r0 = 0usize;
18106        while r0 < out_f {
18107            let rows = chunk_rows.min(out_f - r0);
18108            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
18109            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
18110            let yc = if exact {
18111                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
18112            } else {
18113                self.linear(x, &wf32, m, in_f, rows)?
18114            };
18115            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
18116            for mi in 0..m {
18117                let src = yc.slice(mi * rows..(mi + 1) * rows);
18118                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
18119                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
18120            }
18121            r0 += rows;
18122        }
18123        Ok(y)
18124    }
18125
18126    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
18127    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
18128    /// chunked BF16 numerical program instead of re-encoding the weight.
18129    pub fn linear_bf16_resident(
18130        &self,
18131        x: &CudaSlice<f32>,
18132        data: &CudaSlice<u8>,
18133        m: usize,
18134        in_f: usize,
18135        out_f: usize,
18136    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18137        if data.len() != in_f * out_f * 2 {
18138            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18139        }
18140        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
18141    }
18142
18143    /// Execute a resident BF16 projection as fixed-width output-row chunks.
18144    ///
18145    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
18146    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
18147    /// model topology rather than the active rank count.
18148    pub fn linear_bf16_resident_canonical_rows(
18149        &self,
18150        x: &CudaSlice<f32>,
18151        data: &CudaSlice<u8>,
18152        m: usize,
18153        in_f: usize,
18154        out_f: usize,
18155        canonical_chunk_rows: usize,
18156    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18157        if data.len() != in_f * out_f * 2 {
18158            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18159        }
18160        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18161    }
18162
18163    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
18164    ///
18165    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
18166    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
18167    pub fn linear_f32_resident_canonical_rows(
18168        &self,
18169        x: &CudaSlice<f32>,
18170        data: &CudaSlice<f32>,
18171        m: usize,
18172        in_f: usize,
18173        out_f: usize,
18174        canonical_chunk_rows: usize,
18175    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18176        self.linear_f32_resident_canonical_rows_inner(
18177            x,
18178            data,
18179            m,
18180            in_f,
18181            out_f,
18182            canonical_chunk_rows,
18183            false,
18184        )
18185    }
18186
18187    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
18188    ///
18189    /// The projection shapes and values are identical to
18190    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
18191    /// changes, replacing one device copy per token with one placement kernel per output chunk.
18192    pub fn linear_f32_resident_canonical_rows_strided(
18193        &self,
18194        x: &CudaSlice<f32>,
18195        data: &CudaSlice<f32>,
18196        m: usize,
18197        in_f: usize,
18198        out_f: usize,
18199        canonical_chunk_rows: usize,
18200    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18201        self.linear_f32_resident_canonical_rows_inner(
18202            x,
18203            data,
18204            m,
18205            in_f,
18206            out_f,
18207            canonical_chunk_rows,
18208            true,
18209        )
18210    }
18211
18212    fn linear_f32_resident_canonical_rows_inner(
18213        &self,
18214        x: &CudaSlice<f32>,
18215        data: &CudaSlice<f32>,
18216        m: usize,
18217        in_f: usize,
18218        out_f: usize,
18219        canonical_chunk_rows: usize,
18220        strided_output: bool,
18221    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18222        if data.len() != in_f * out_f {
18223            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18224        }
18225        if canonical_chunk_rows == 0
18226            || canonical_chunk_rows > out_f
18227            || out_f % canonical_chunk_rows != 0
18228        {
18229            return Err(format!(
18230                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18231            )
18232            .into());
18233        }
18234        if canonical_chunk_rows == out_f {
18235            return self.linear(x, data, m, in_f, out_f);
18236        }
18237
18238        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18239        let input = x.slice(0..x.len());
18240        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18241            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18242            if m == 1 {
18243                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18244                self.linear_device_into(
18245                    &input,
18246                    &weights,
18247                    &mut destination,
18248                    1,
18249                    in_f,
18250                    canonical_chunk_rows,
18251                )?;
18252                continue;
18253            }
18254            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18255            if strided_output {
18256                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18257            } else {
18258                for token in 0..m {
18259                    let source = chunk
18260                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18261                    let mut destination =
18262                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18263                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18264                }
18265            }
18266        }
18267        Ok(y)
18268    }
18269
18270    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
18271    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
18272    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
18273    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
18274    pub fn linear_f32_resident_canonical_rows_t1_into(
18275        &self,
18276        x: &CudaSlice<f32>,
18277        data: &CudaSlice<f32>,
18278        y: &mut CudaSlice<f32>,
18279        in_f: usize,
18280        out_f: usize,
18281        canonical_chunk_rows: usize,
18282    ) -> Result<(), Box<dyn std::error::Error>> {
18283        if data.len() != in_f * out_f {
18284            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18285        }
18286        if y.len() != out_f || x.len() != in_f {
18287            return Err(format!(
18288                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18289                x.len(),
18290                y.len()
18291            )
18292            .into());
18293        }
18294        if canonical_chunk_rows == 0
18295            || canonical_chunk_rows > out_f
18296            || out_f % canonical_chunk_rows != 0
18297        {
18298            return Err(format!(
18299                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18300            )
18301            .into());
18302        }
18303        let input = x.slice(0..x.len());
18304        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18305            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18306            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18307            self.linear_device_into(
18308                &input,
18309                &weights,
18310                &mut destination,
18311                1,
18312                in_f,
18313                canonical_chunk_rows,
18314            )?;
18315        }
18316        Ok(())
18317    }
18318
18319    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
18320    /// without the allocation, for workspace-resident operands.
18321    pub fn linear_t1_into(
18322        &self,
18323        x: &cudarc::driver::CudaView<'_, f32>,
18324        w: &cudarc::driver::CudaView<'_, f32>,
18325        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18326        in_f: usize,
18327        out_f: usize,
18328    ) -> Result<(), Box<dyn std::error::Error>> {
18329        self.linear_device_into(x, w, y, 1, in_f, out_f)
18330    }
18331
18332    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
18333    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
18334    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
18335    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
18336    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
18337    /// router/shexp sites and matmul_decode_exact's Float arm.
18338    pub fn linear_decode_exact(
18339        &self,
18340        x: &CudaSlice<f32>,
18341        w: &CudaSlice<f32>,
18342        m_tokens: usize,
18343        in_f: usize,
18344        out_f: usize,
18345    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18346        if m_tokens == 1 {
18347            return self.linear(x, w, 1, in_f, out_f);
18348        }
18349        let xv = self.view(x, m_tokens * in_f);
18350        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18351        for t in 0..m_tokens {
18352            let row = xv.slice(t * in_f..(t + 1) * in_f);
18353            let mut xr = self.alloc_uninit::<f32>(in_f)?;
18354            self.copy_view_into(&mut xr, 0, &row, in_f)?;
18355            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18356            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18357        }
18358        Ok(y)
18359    }
18360
18361    pub fn linear(
18362        &self,
18363        x: &CudaSlice<f32>,
18364        w: &CudaSlice<f32>,
18365        m_tokens: usize,
18366        in_f: usize,
18367        out_f: usize,
18368    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18369        self.linear_device(x, w, m_tokens, in_f, out_f)
18370    }
18371
18372    fn linear_device<I>(
18373        &self,
18374        x: &I,
18375        w: &I,
18376        m_tokens: usize,
18377        in_f: usize,
18378        out_f: usize,
18379    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18380    where
18381        I: cudarc::driver::DevicePtr<f32>,
18382    {
18383        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
18384        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
18385        Ok(c)
18386    }
18387
18388    fn linear_device_into<I, O>(
18389        &self,
18390        x: &I,
18391        w: &I,
18392        c: &mut O,
18393        m_tokens: usize,
18394        in_f: usize,
18395        out_f: usize,
18396    ) -> Result<(), Box<dyn std::error::Error>>
18397    where
18398        I: cudarc::driver::DevicePtr<f32>,
18399        O: cudarc::driver::DevicePtrMut<f32>,
18400    {
18401        use cudarc::cublaslt::{Matmul, MatmulConfig};
18402        let cfg = MatmulConfig {
18403            transa: true,
18404            transb: false,
18405            transc: false,
18406            m: out_f as u64,
18407            n: m_tokens as u64,
18408            k: in_f as u64,
18409            alpha: 1.0,
18410            lda: in_f as i64,
18411            ldb: in_f as i64,
18412            beta: 0.0,
18413            ldc: out_f as i64,
18414            stride_a: None,
18415            stride_b: None,
18416            stride_c: None,
18417            stride_bias: None,
18418            batch_size: None,
18419        };
18420        let blas = self.gpu.blas();
18421        unsafe {
18422            blas.matmul(cfg, w, x, c, None, None)?;
18423        }
18424        Ok(())
18425    }
18426
18427    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
18428    pub fn sdpa_naive(
18429        &self,
18430        q: &CudaSlice<f32>,
18431        k: &CudaSlice<f32>,
18432        v: &CudaSlice<f32>,
18433        o: &mut CudaSlice<f32>,
18434        head_dim: usize,
18435        n_head: usize,
18436        n_head_kv: usize,
18437        t: usize,
18438        t_kv: usize,
18439        scale: f32,
18440        causal: bool,
18441    ) -> Result<(), Box<dyn std::error::Error>> {
18442        let f = self.func("sdpa_naive_f32");
18443        let cfg = LaunchConfig {
18444            grid_dim: (n_head as u32, t as u32, 1),
18445            block_dim: (128, 1, 1),
18446            shared_mem_bytes: (t_kv * 4) as u32,
18447        };
18448        let (hd, nh, nhkv, ti, tkvi, cz) = (
18449            head_dim as i32,
18450            n_head as i32,
18451            n_head_kv as i32,
18452            t as i32,
18453            t_kv as i32,
18454            causal as i32,
18455        );
18456        let __s_b = self.gpu.stream();
18457        let mut b = __s_b.launch_builder(&f);
18458        b.arg(q)
18459            .arg(k)
18460            .arg(v)
18461            .arg(o)
18462            .arg(&hd)
18463            .arg(&nh)
18464            .arg(&nhkv)
18465            .arg(&ti)
18466            .arg(&tkvi)
18467            .arg(&scale)
18468            .arg(&cz);
18469        unsafe {
18470            b.launch(cfg)?;
18471        }
18472        Ok(())
18473    }
18474
18475    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
18476    /// bidirectional image islands. `span_id` labels each absolute kv position
18477    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
18478    /// reproducing the reference's non-causal image batch. window 0 = no window.
18479    #[allow(clippy::too_many_arguments)]
18480    pub fn sdpa_naive_island(
18481        &self,
18482        q: &CudaSlice<f32>,
18483        k: &CudaSlice<f32>,
18484        v: &CudaSlice<f32>,
18485        o: &mut CudaSlice<f32>,
18486        span_id: &CudaSlice<i32>,
18487        head_dim: usize,
18488        n_head: usize,
18489        n_head_kv: usize,
18490        t: usize,
18491        t_kv: usize,
18492        scale: f32,
18493        window: usize,
18494    ) -> Result<(), Box<dyn std::error::Error>> {
18495        let f = self.func("sdpa_naive_island_f32");
18496        let cfg = LaunchConfig {
18497            grid_dim: (n_head as u32, t as u32, 1),
18498            block_dim: (128, 1, 1),
18499            shared_mem_bytes: (t_kv * 4) as u32,
18500        };
18501        let (hd, nh, nhkv, ti, tkvi, wi) = (
18502            head_dim as i32,
18503            n_head as i32,
18504            n_head_kv as i32,
18505            t as i32,
18506            t_kv as i32,
18507            window as i32,
18508        );
18509        let __s_b = self.gpu.stream();
18510        let mut b = __s_b.launch_builder(&f);
18511        b.arg(q)
18512            .arg(k)
18513            .arg(v)
18514            .arg(o)
18515            .arg(span_id)
18516            .arg(&hd)
18517            .arg(&nh)
18518            .arg(&nhkv)
18519            .arg(&ti)
18520            .arg(&tkvi)
18521            .arg(&scale)
18522            .arg(&wi);
18523        unsafe {
18524            b.launch(cfg)?;
18525        }
18526        Ok(())
18527    }
18528
18529    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
18530    #[allow(clippy::too_many_arguments)]
18531    pub fn sdpa_naive_w(
18532        &self,
18533        q: &CudaSlice<f32>,
18534        k: &CudaSlice<f32>,
18535        v: &CudaSlice<f32>,
18536        o: &mut CudaSlice<f32>,
18537        head_dim: usize,
18538        n_head: usize,
18539        n_head_kv: usize,
18540        t: usize,
18541        t_kv: usize,
18542        scale: f32,
18543        causal: bool,
18544        window: usize,
18545    ) -> Result<(), Box<dyn std::error::Error>> {
18546        let f = self.func("sdpa_naive_w_f32");
18547        let cfg = LaunchConfig {
18548            grid_dim: (n_head as u32, t as u32, 1),
18549            block_dim: (128, 1, 1),
18550            shared_mem_bytes: (t_kv * 4) as u32,
18551        };
18552        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18553            head_dim as i32,
18554            n_head as i32,
18555            n_head_kv as i32,
18556            t as i32,
18557            t_kv as i32,
18558            causal as i32,
18559            window as i32,
18560        );
18561        let __s_b = self.gpu.stream();
18562        let mut b = __s_b.launch_builder(&f);
18563        b.arg(q)
18564            .arg(k)
18565            .arg(v)
18566            .arg(o)
18567            .arg(&hd)
18568            .arg(&nh)
18569            .arg(&nhkv)
18570            .arg(&ti)
18571            .arg(&tkvi)
18572            .arg(&scale)
18573            .arg(&cz)
18574            .arg(&wi);
18575        unsafe {
18576            b.launch(cfg)?;
18577        }
18578        Ok(())
18579    }
18580
18581    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
18582    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
18583    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
18584    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
18585    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
18586    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
18587    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
18588    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
18589    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
18590    #[allow(clippy::too_many_arguments)]
18591    pub fn sdpa_naive_w_lo(
18592        &self,
18593        q: &CudaSlice<f32>,
18594        k: &CudaSlice<f32>,
18595        v: &CudaSlice<f32>,
18596        o: &mut CudaSlice<f32>,
18597        head_dim: usize,
18598        n_head: usize,
18599        n_head_kv: usize,
18600        t: usize,
18601        t_kv: usize,
18602        scale: f32,
18603        causal: bool,
18604        window: usize,
18605    ) -> Result<(), Box<dyn std::error::Error>> {
18606        let kv_lo = if window > 0 {
18607            (t_kv - t + 1).saturating_sub(window)
18608        } else {
18609            0
18610        };
18611        let smem = (t_kv - kv_lo) * 4;
18612        if smem > 48 * 1024 {
18613            return Err(format!(
18614                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18615                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18616                 a window this wide needs the multi-pass long-ctx kernel"
18617            )
18618            .into());
18619        }
18620        let f = self.func("sdpa_naive_w_lo_f32");
18621        let cfg = LaunchConfig {
18622            grid_dim: (n_head as u32, t as u32, 1),
18623            block_dim: (128, 1, 1),
18624            shared_mem_bytes: smem as u32,
18625        };
18626        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18627            head_dim as i32,
18628            n_head as i32,
18629            n_head_kv as i32,
18630            t as i32,
18631            t_kv as i32,
18632            causal as i32,
18633            window as i32,
18634            kv_lo as i32,
18635        );
18636        let __s_b = self.gpu.stream();
18637        let mut b = __s_b.launch_builder(&f);
18638        b.arg(q)
18639            .arg(k)
18640            .arg(v)
18641            .arg(o)
18642            .arg(&hd)
18643            .arg(&nh)
18644            .arg(&nhkv)
18645            .arg(&ti)
18646            .arg(&tkvi)
18647            .arg(&scale)
18648            .arg(&cz)
18649            .arg(&wi)
18650            .arg(&lo);
18651        unsafe {
18652            b.launch(cfg)?;
18653        }
18654        Ok(())
18655    }
18656
18657    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
18658    pub fn sdpa_naive_view(
18659        &self,
18660        q: &CudaSlice<f32>,
18661        k: &cudarc::driver::CudaView<f32>,
18662        v: &cudarc::driver::CudaView<f32>,
18663        o: &mut CudaSlice<f32>,
18664        head_dim: usize,
18665        n_head: usize,
18666        n_head_kv: usize,
18667        t: usize,
18668        t_kv: usize,
18669        scale: f32,
18670        causal: bool,
18671    ) -> Result<(), Box<dyn std::error::Error>> {
18672        let f = self.func("sdpa_naive_f32");
18673        let cfg = LaunchConfig {
18674            grid_dim: (n_head as u32, t as u32, 1),
18675            block_dim: (128, 1, 1),
18676            shared_mem_bytes: (t_kv * 4) as u32,
18677        };
18678        let (hd, nh, nhkv, ti, tkvi, cz) = (
18679            head_dim as i32,
18680            n_head as i32,
18681            n_head_kv as i32,
18682            t as i32,
18683            t_kv as i32,
18684            causal as i32,
18685        );
18686        let __s_b = self.gpu.stream();
18687        let mut b = __s_b.launch_builder(&f);
18688        b.arg(q)
18689            .arg(k)
18690            .arg(v)
18691            .arg(o)
18692            .arg(&hd)
18693            .arg(&nh)
18694            .arg(&nhkv)
18695            .arg(&ti)
18696            .arg(&tkvi)
18697            .arg(&scale)
18698            .arg(&cz);
18699        unsafe {
18700            b.launch(cfg)?;
18701        }
18702        Ok(())
18703    }
18704
18705    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
18706    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
18707    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
18708    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
18709    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
18710    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
18711    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
18712    #[allow(clippy::too_many_arguments)]
18713    pub fn fa_dequant_kv_view_f32(
18714        &self,
18715        k: &cudarc::driver::CudaView<u8>,
18716        v: &cudarc::driver::CudaView<u8>,
18717        kf: &mut CudaSlice<f32>,
18718        vf: &mut CudaSlice<f32>,
18719        kv_dim_k: usize,
18720        kv_dim_v: usize,
18721        t_kv: usize,
18722        k_tok_bytes: usize,
18723        v_tok_bytes: usize,
18724        g: bool,
18725    ) -> Result<(), Box<dyn std::error::Error>> {
18726        let f = if g {
18727            self.func_g("fa_dequant_kv_ws_f32")
18728        } else {
18729            self.func("fa_dequant_kv_ws_f32")
18730        };
18731        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18732        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18733        let cfg = LaunchConfig {
18734            grid_dim: (nblk.max(1), 1, 1),
18735            block_dim: (256, 1, 1),
18736            shared_mem_bytes: 0,
18737        };
18738        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18739        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18740        let __s_b = self.gpu.stream();
18741        let mut b = __s_b.launch_builder(&f);
18742        b.arg(k)
18743            .arg(v)
18744            .arg(&mut *kf)
18745            .arg(&mut *vf)
18746            .arg(&kdk)
18747            .arg(&kdv)
18748            .arg(&tkvi)
18749            .arg(&ktb)
18750            .arg(&vtb);
18751        unsafe {
18752            b.launch(cfg)?;
18753        }
18754        Ok(())
18755    }
18756
18757    #[allow(clippy::too_many_arguments)]
18758    pub fn sdpa_naive_quantized_view(
18759        &self,
18760        q: &CudaSlice<f32>,
18761        k: &cudarc::driver::CudaView<u8>,
18762        v: &cudarc::driver::CudaView<u8>,
18763        o: &mut CudaSlice<f32>,
18764        head_dim: usize,
18765        n_head: usize,
18766        n_head_kv: usize,
18767        t: usize,
18768        t_kv: usize,
18769        scale: f32,
18770        causal: bool,
18771        k_tok_bytes: usize,
18772        v_tok_bytes: usize,
18773    ) -> Result<(), Box<dyn std::error::Error>> {
18774        let kv_dim = n_head_kv * head_dim;
18775        let mut kf = self.uninit(t_kv * kv_dim)?;
18776        let mut vf = self.uninit(t_kv * kv_dim)?;
18777        let f = self.func("fa_dequant_kv_ws_f32");
18778        let total = (2 * t_kv * kv_dim) as u64;
18779        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18780        let cfg = LaunchConfig {
18781            grid_dim: (nblk.max(1), 1, 1),
18782            block_dim: (256, 1, 1),
18783            shared_mem_bytes: 0,
18784        };
18785        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18786        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18787        let __s_b = self.gpu.stream();
18788        let mut b = __s_b.launch_builder(&f);
18789        b.arg(k)
18790            .arg(v)
18791            .arg(&mut kf)
18792            .arg(&mut vf)
18793            .arg(&kv_dim_i)
18794            .arg(&kv_dim_i)
18795            .arg(&t_kv_i)
18796            .arg(&k_tok_bytes_i)
18797            .arg(&v_tok_bytes_i);
18798        unsafe { b.launch(cfg)? };
18799        self.sdpa_naive(
18800            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18801        )
18802    }
18803
18804    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
18805    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
18806    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
18807    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
18808    /// unwindowed function above and produces bit-identical output at window == 0.
18809    ///
18810    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
18811    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
18812    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
18813    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
18814    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
18815    #[allow(clippy::too_many_arguments)]
18816    pub fn sdpa_naive_w_quantized_view(
18817        &self,
18818        q: &CudaSlice<f32>,
18819        k: &cudarc::driver::CudaView<u8>,
18820        v: &cudarc::driver::CudaView<u8>,
18821        o: &mut CudaSlice<f32>,
18822        head_dim: usize,
18823        n_head: usize,
18824        n_head_kv: usize,
18825        t: usize,
18826        t_kv: usize,
18827        scale: f32,
18828        causal: bool,
18829        window: usize,
18830        k_tok_bytes: usize,
18831        v_tok_bytes: usize,
18832    ) -> Result<(), Box<dyn std::error::Error>> {
18833        let kv_dim = n_head_kv * head_dim;
18834        let mut kf = self.uninit(t_kv * kv_dim)?;
18835        let mut vf = self.uninit(t_kv * kv_dim)?;
18836        let f = self.func("fa_dequant_kv_ws_f32");
18837        let total = (2 * t_kv * kv_dim) as u64;
18838        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18839        let cfg = LaunchConfig {
18840            grid_dim: (nblk.max(1), 1, 1),
18841            block_dim: (256, 1, 1),
18842            shared_mem_bytes: 0,
18843        };
18844        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18845        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18846        let __s_b = self.gpu.stream();
18847        let mut b = __s_b.launch_builder(&f);
18848        b.arg(k)
18849            .arg(v)
18850            .arg(&mut kf)
18851            .arg(&mut vf)
18852            .arg(&kv_dim_i)
18853            .arg(&kv_dim_i)
18854            .arg(&t_kv_i)
18855            .arg(&k_tok_bytes_i)
18856            .arg(&v_tok_bytes_i);
18857        unsafe { b.launch(cfg)? };
18858        self.sdpa_naive_w(
18859            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18860        )
18861    }
18862
18863    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
18864    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
18865    /// Q/K/V/O [head_dim, n_head(_kv), T].
18866    pub fn fa_prefill(
18867        &self,
18868        q: &CudaSlice<f32>,
18869        k: &CudaSlice<f32>,
18870        v: &CudaSlice<f32>,
18871        o: &mut CudaSlice<f32>,
18872        head_dim: usize,
18873        n_head: usize,
18874        n_head_kv: usize,
18875        t: usize,
18876        t_kv: usize,
18877        scale: f32,
18878        causal: bool,
18879    ) -> Result<(), Box<dyn std::error::Error>> {
18880        if portable_mma_gated() {
18881            return self.sdpa_naive(
18882                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18883            );
18884        }
18885        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
18886        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
18887        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
18888        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
18889        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
18890        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
18891        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
18892        let fa3_on = head_dim == 256
18893            && causal
18894            && t == t_kv
18895            && match std::env::var("MEMRA_FA3").as_deref() {
18896                Ok("0") => false,
18897                // The force arm consults the arch now: the bf16 stage below calls
18898                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
18899                // a portable build. Refuse at the switch, not at the lookup.
18900                Ok("1") => {
18901                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
18902                    true
18903                }
18904                _ => cfg!(memra_hopper_mma),
18905            };
18906        if fa3_on {
18907            let n = t * n_head * head_dim;
18908            let nkv = t * n_head_kv * head_dim;
18909            let mut q16 = self.alloc_u8_uninit(n * 2)?;
18910            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18911            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18912            self.f32_to_bf16_into(q, &mut q16, n)?;
18913            self.f32_to_bf16_into(k, &mut k16, nkv)?;
18914            self.f32_to_bf16_into(v, &mut v16, nkv)?;
18915            let rc = {
18916                use cudarc::driver::{DevicePtr, DevicePtrMut};
18917                let stream = self.gpu.stream();
18918                let (qp, _g1) = q16.device_ptr(&stream);
18919                let (kp, _g2) = k16.device_ptr(&stream);
18920                let (vp, _g3) = v16.device_ptr(&stream);
18921                let (op, _g4) = o.device_ptr_mut(&stream);
18922                unsafe {
18923                    memra_fa3_prefill(
18924                        qp as *const core::ffi::c_void,
18925                        kp as *const core::ffi::c_void,
18926                        vp as *const core::ffi::c_void,
18927                        op as *mut f32,
18928                        t as i32,
18929                        n_head as i32,
18930                        n_head_kv as i32,
18931                        head_dim as i32,
18932                        scale,
18933                        stream.cu_stream() as *mut core::ffi::c_void,
18934                    )
18935                }
18936            };
18937            if rc != 0 {
18938                return Err(format!("memra_fa3_prefill rc={rc}").into());
18939            }
18940            return Ok(());
18941        }
18942        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
18943        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
18944        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
18945        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
18946        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18947        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18948        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18949            const BLOCK_Q: usize = 64;
18950            const BKX: usize = 32;
18951            let f = self.func("fa_prefill_bf16_p1");
18952            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18953                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18954            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18955            f.set_attribute(
18956                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18957                shmem as i32,
18958            )?;
18959            let cfg = LaunchConfig {
18960                grid_dim: (
18961                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18962                    n_head as u32,
18963                    1,
18964                ),
18965                block_dim: (32, 4, 1),
18966                shared_mem_bytes: shmem,
18967            };
18968            let (hd, nh, nhkv, ti, tkvi, cz) = (
18969                head_dim as i32,
18970                n_head as i32,
18971                n_head_kv as i32,
18972                t as i32,
18973                t_kv as i32,
18974                causal as i32,
18975            );
18976            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18977            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18978            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18979            let __s_b = self.gpu.stream();
18980            let mut b = __s_b.launch_builder(&f);
18981            b.arg(&qb)
18982                .arg(&kb)
18983                .arg(&vb)
18984                .arg(o)
18985                .arg(&hd)
18986                .arg(&nh)
18987                .arg(&nhkv)
18988                .arg(&ti)
18989                .arg(&tkvi)
18990                .arg(&scale)
18991                .arg(&cz);
18992            unsafe {
18993                b.launch(cfg)?;
18994            }
18995            return Ok(());
18996        }
18997        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
18998        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
18999        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
19000        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
19001        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
19002        const BK: usize = 32;
19003        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
19004        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
19005        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19006        let (block_q, warps, w2_sfx): (usize, u32, &str) =
19007            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19008        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
19009        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
19010        // other head_dims to sdpa_naive before reaching here.
19011        let hd_sfx = fa_hd_suffix(head_dim)?;
19012        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19013        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
19014        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
19015        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
19016        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
19017        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
19018        let (kb16, vb16) = if bf16kv {
19019            let n = t_kv * n_head_kv * head_dim;
19020            let mut kb = self.alloc_u8_uninit(n * 2)?;
19021            let mut vb = self.alloc_u8_uninit(n * 2)?;
19022            let fcv = self.func("f32_to_bf16_bulk");
19023            let ni = n as i64;
19024            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
19025            let __s_b = self.gpu.stream();
19026            let mut b = __s_b.launch_builder(&fcv);
19027            b.arg(k).arg(&mut kb).arg(&ni);
19028            unsafe {
19029                b.launch(cfgc)?;
19030            }
19031            let __s_b = self.gpu.stream();
19032            let mut b = __s_b.launch_builder(&fcv);
19033            b.arg(v).arg(&mut vb).arg(&ni);
19034            unsafe {
19035                b.launch(cfgc)?;
19036            }
19037            (Some(kb), Some(vb))
19038        } else {
19039            (None, None)
19040        };
19041        let f = self.func(&if bf16kv {
19042            format!("fa_prefill_bf16kv_pp{hd_sfx}")
19043        } else {
19044            format!(
19045                "fa_prefill_f32{}{}{hd_sfx}",
19046                if floor { "" } else { "_pp" },
19047                if floor { "" } else { w2_sfx }
19048            )
19049        });
19050        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
19051        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
19052        let kv_stages = if bf16kv { 2 } else { 1 };
19053        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19054            + 4 * (block_q * BK + 2 * block_q)) as u32;
19055        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19056        f.set_attribute(
19057            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19058            shmem as i32,
19059        )?;
19060        let cfg = LaunchConfig {
19061            grid_dim: (
19062                (t as u32 + block_q as u32 - 1) / block_q as u32,
19063                n_head as u32,
19064                1,
19065            ),
19066            block_dim: (32, warps, 1),
19067            shared_mem_bytes: shmem,
19068        };
19069        let (hd, nh, nhkv, ti, tkvi, cz) = (
19070            head_dim as i32,
19071            n_head as i32,
19072            n_head_kv as i32,
19073            t as i32,
19074            t_kv as i32,
19075            causal as i32,
19076        );
19077        let __s_b = self.gpu.stream();
19078        let mut b = __s_b.launch_builder(&f);
19079        b.arg(q);
19080        match (&kb16, &vb16) {
19081            (Some(kb), Some(vb)) => {
19082                b.arg(kb).arg(vb);
19083            }
19084            _ => {
19085                b.arg(k).arg(v);
19086            }
19087        }
19088        b.arg(o)
19089            .arg(&hd)
19090            .arg(&nh)
19091            .arg(&nhkv)
19092            .arg(&ti)
19093            .arg(&tkvi)
19094            .arg(&scale)
19095            .arg(&cz);
19096        unsafe {
19097            b.launch(cfg)?;
19098        }
19099        Ok(())
19100    }
19101
19102    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
19103    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
19104    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
19105    #[allow(clippy::too_many_arguments)]
19106    pub fn fa_prefill_w(
19107        &self,
19108        q: &CudaSlice<f32>,
19109        k: &CudaSlice<f32>,
19110        v: &CudaSlice<f32>,
19111        o: &mut CudaSlice<f32>,
19112        head_dim: usize,
19113        n_head: usize,
19114        n_head_kv: usize,
19115        t: usize,
19116        t_kv: usize,
19117        scale: f32,
19118        causal: bool,
19119        window: usize,
19120    ) -> Result<(), Box<dyn std::error::Error>> {
19121        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
19122        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
19123        if portable_mma_gated() {
19124            return self.sdpa_naive_w(
19125                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19126            );
19127        }
19128        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
19129        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
19130        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
19131        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19132        let faw_f32 =
19133            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
19134        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19135        self.fa_prefill_w_arm(
19136            q,
19137            k,
19138            v,
19139            o,
19140            head_dim,
19141            n_head,
19142            n_head_kv,
19143            t,
19144            t_kv,
19145            scale,
19146            causal,
19147            window,
19148            floor || faw_f32,
19149            floor,
19150        )
19151    }
19152
19153    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
19154    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
19155    #[allow(clippy::too_many_arguments)]
19156    pub fn fa_prefill_w_pre(
19157        &self,
19158        qb: &CudaSlice<u8>,
19159        kb: &CudaSlice<u8>,
19160        vb: &CudaSlice<u8>,
19161        o: &mut CudaSlice<f32>,
19162        head_dim: usize,
19163        n_head: usize,
19164        n_head_kv: usize,
19165        t: usize,
19166        t_kv: usize,
19167        scale: f32,
19168        causal: bool,
19169        window: usize,
19170        v_f16: bool,
19171    ) -> Result<(), Box<dyn std::error::Error>> {
19172        const BLOCK_Q: usize = 64;
19173        const BK: usize = 32;
19174        debug_assert_eq!(head_dim, 256);
19175        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19176        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19177        if hp {
19178            const BLOCK_QH: usize = 32;
19179            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
19180            // else re-encode through the pooled scratch (stream-ordered reuse).
19181            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19182            let vh: &CudaSlice<u8> = if v_f16 {
19183                vb
19184            } else {
19185                let n = t_kv * n_head_kv * head_dim;
19186                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19187                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19188                }
19189                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19190                vguard.as_ref().unwrap()
19191            };
19192            let f = self.func("fa_prefill_w_bf16_p1h2");
19193            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19194            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19195            f.set_attribute(
19196                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19197                shmem as i32,
19198            )?;
19199            let cfg = LaunchConfig {
19200                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19201                block_dim: (32, 4, 1),
19202                shared_mem_bytes: shmem,
19203            };
19204            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19205                head_dim as i32,
19206                n_head as i32,
19207                n_head_kv as i32,
19208                t as i32,
19209                t_kv as i32,
19210                causal as i32,
19211                window as i32,
19212            );
19213            let __s_b = self.gpu.stream();
19214            let mut b = __s_b.launch_builder(&f);
19215            b.arg(qb)
19216                .arg(kb)
19217                .arg(vh)
19218                .arg(o)
19219                .arg(&hd)
19220                .arg(&nh)
19221                .arg(&nhkv)
19222                .arg(&ti)
19223                .arg(&tkvi)
19224                .arg(&scale)
19225                .arg(&cz)
19226                .arg(&wi);
19227            unsafe {
19228                b.launch(cfg)?;
19229            }
19230            return Ok(());
19231        }
19232        let f = self.func("fa_prefill_w_bf16_p1");
19233        let shmem =
19234            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19235        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19236        f.set_attribute(
19237            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19238            shmem as i32,
19239        )?;
19240        let cfg = LaunchConfig {
19241            grid_dim: (
19242                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19243                n_head as u32,
19244                1,
19245            ),
19246            block_dim: (32, 4, 1),
19247            shared_mem_bytes: shmem,
19248        };
19249        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19250            head_dim as i32,
19251            n_head as i32,
19252            n_head_kv as i32,
19253            t as i32,
19254            t_kv as i32,
19255            causal as i32,
19256            window as i32,
19257        );
19258        let __s_b = self.gpu.stream();
19259        let mut b = __s_b.launch_builder(&f);
19260        b.arg(qb)
19261            .arg(kb)
19262            .arg(vb)
19263            .arg(o)
19264            .arg(&hd)
19265            .arg(&nh)
19266            .arg(&nhkv)
19267            .arg(&ti)
19268            .arg(&tkvi)
19269            .arg(&scale)
19270            .arg(&cz)
19271            .arg(&wi);
19272        unsafe {
19273            b.launch(cfg)?;
19274        }
19275        Ok(())
19276    }
19277
19278    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
19279    #[allow(clippy::too_many_arguments)]
19280    pub fn fa_prefill_w_arm(
19281        &self,
19282        q: &CudaSlice<f32>,
19283        k: &CudaSlice<f32>,
19284        v: &CudaSlice<f32>,
19285        o: &mut CudaSlice<f32>,
19286        head_dim: usize,
19287        n_head: usize,
19288        n_head_kv: usize,
19289        t: usize,
19290        t_kv: usize,
19291        scale: f32,
19292        causal: bool,
19293        window: usize,
19294        f32_stage: bool,
19295        floor: bool,
19296    ) -> Result<(), Box<dyn std::error::Error>> {
19297        const BLOCK_Q: usize = 64;
19298        const BK: usize = 32;
19299        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19300        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
19301        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
19302        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
19303        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19304        let p1 = !floor
19305            && !f32_stage
19306            && *P1_ON.get_or_init(|| {
19307                std::env::var("MEMRA_FAW_P1")
19308                    .map(|v| v != "0")
19309                    .unwrap_or(true)
19310            });
19311        let hp =
19312            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19313        if hp {
19314            const BLOCK_QH: usize = 32;
19315            let f = self.func("fa_prefill_w_bf16_p1h2");
19316            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19317            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19318            f.set_attribute(
19319                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19320                shmem as i32,
19321            )?;
19322            let cfg = LaunchConfig {
19323                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19324                block_dim: (32, 4, 1),
19325                shared_mem_bytes: shmem,
19326            };
19327            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19328                head_dim as i32,
19329                n_head as i32,
19330                n_head_kv as i32,
19331                t as i32,
19332                t_kv as i32,
19333                causal as i32,
19334                window as i32,
19335            );
19336            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19337            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19338            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19339            let __s_b = self.gpu.stream();
19340            let mut b = __s_b.launch_builder(&f);
19341            b.arg(&qb)
19342                .arg(&kb)
19343                .arg(&vh)
19344                .arg(o)
19345                .arg(&hd)
19346                .arg(&nh)
19347                .arg(&nhkv)
19348                .arg(&ti)
19349                .arg(&tkvi)
19350                .arg(&scale)
19351                .arg(&cz)
19352                .arg(&wi);
19353            unsafe {
19354                b.launch(cfg)?;
19355            }
19356            return Ok(());
19357        }
19358        if p1 {
19359            let f = self.func("fa_prefill_w_bf16_p1");
19360            let shmem =
19361                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19362            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19363            f.set_attribute(
19364                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19365                shmem as i32,
19366            )?;
19367            let cfg = LaunchConfig {
19368                grid_dim: (
19369                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19370                    n_head as u32,
19371                    1,
19372                ),
19373                block_dim: (32, 4, 1),
19374                shared_mem_bytes: shmem,
19375            };
19376            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19377                head_dim as i32,
19378                n_head as i32,
19379                n_head_kv as i32,
19380                t as i32,
19381                t_kv as i32,
19382                causal as i32,
19383                window as i32,
19384            );
19385            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19386            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19387            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19388            let __s_b = self.gpu.stream();
19389            let mut b = __s_b.launch_builder(&f);
19390            b.arg(&qb)
19391                .arg(&kb)
19392                .arg(&vb)
19393                .arg(o)
19394                .arg(&hd)
19395                .arg(&nh)
19396                .arg(&nhkv)
19397                .arg(&ti)
19398                .arg(&tkvi)
19399                .arg(&scale)
19400                .arg(&cz)
19401                .arg(&wi);
19402            unsafe {
19403                b.launch(cfg)?;
19404            }
19405            return Ok(());
19406        }
19407        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
19408        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
19409        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19410        let g4 = !floor
19411            && !f32_stage
19412            && n_head_kv == 1
19413            && n_head % 4 == 0
19414            && *G4_ON.get_or_init(|| {
19415                std::env::var("MEMRA_FAW_G4")
19416                    .map(|v| v != "0")
19417                    .unwrap_or(true)
19418            });
19419        if g4 {
19420            const SP_M: usize = 16;
19421            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
19422            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
19423            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19424            let o2 = *O2_ON.get_or_init(|| {
19425                std::env::var("MEMRA_FAW_O2")
19426                    .map(|v| v != "0")
19427                    .unwrap_or(true)
19428            });
19429            let f = self.func(if o2 {
19430                "fa_prefill_w_bf16_g4o2"
19431            } else {
19432                "fa_prefill_w_bf16_g4"
19433            });
19434            let shmem = if o2 {
19435                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19436            } else {
19437                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19438                    as u32
19439            };
19440            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19441            f.set_attribute(
19442                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19443                shmem as i32,
19444            )?;
19445            let cfg = LaunchConfig {
19446                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19447                block_dim: (32, 4, 1),
19448                shared_mem_bytes: shmem,
19449            };
19450            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19451                head_dim as i32,
19452                n_head as i32,
19453                n_head_kv as i32,
19454                t as i32,
19455                t_kv as i32,
19456                causal as i32,
19457                window as i32,
19458            );
19459            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19460            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19461            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19462            let __s_b = self.gpu.stream();
19463            let mut b = __s_b.launch_builder(&f);
19464            b.arg(&qb)
19465                .arg(&kb)
19466                .arg(&vb)
19467                .arg(o)
19468                .arg(&hd)
19469                .arg(&nh)
19470                .arg(&nhkv)
19471                .arg(&ti)
19472                .arg(&tkvi)
19473                .arg(&scale)
19474                .arg(&cz)
19475                .arg(&wi);
19476            unsafe {
19477                b.launch(cfg)?;
19478            }
19479            return Ok(());
19480        }
19481        let f = self.func(if floor {
19482            "fa_prefill_w_f32"
19483        } else if f32_stage {
19484            "fa_prefill_w_f32_pp"
19485        } else {
19486            "fa_prefill_w_bf16_pp"
19487        });
19488        let shmem =
19489            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19490        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19491        f.set_attribute(
19492            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19493            shmem as i32,
19494        )?;
19495        let cfg = LaunchConfig {
19496            grid_dim: (
19497                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19498                n_head as u32,
19499                1,
19500            ),
19501            block_dim: (32, 4, 1),
19502            shared_mem_bytes: shmem,
19503        };
19504        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19505            head_dim as i32,
19506            n_head as i32,
19507            n_head_kv as i32,
19508            t as i32,
19509            t_kv as i32,
19510            causal as i32,
19511            window as i32,
19512        );
19513        if f32_stage {
19514            let __s_b = self.gpu.stream();
19515            let mut b = __s_b.launch_builder(&f);
19516            b.arg(q)
19517                .arg(k)
19518                .arg(v)
19519                .arg(o)
19520                .arg(&hd)
19521                .arg(&nh)
19522                .arg(&nhkv)
19523                .arg(&ti)
19524                .arg(&tkvi)
19525                .arg(&scale)
19526                .arg(&cz)
19527                .arg(&wi);
19528            unsafe {
19529                b.launch(cfg)?;
19530            }
19531        } else {
19532            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19533            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19534            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19535            let __s_b = self.gpu.stream();
19536            let mut b = __s_b.launch_builder(&f);
19537            b.arg(&qb)
19538                .arg(&kb)
19539                .arg(&vb)
19540                .arg(o)
19541                .arg(&hd)
19542                .arg(&nh)
19543                .arg(&nhkv)
19544                .arg(&ti)
19545                .arg(&tkvi)
19546                .arg(&scale)
19547                .arg(&cz)
19548                .arg(&wi);
19549            unsafe {
19550                b.launch(cfg)?;
19551            }
19552        }
19553        Ok(())
19554    }
19555
19556    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
19557    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
19558    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
19559    #[allow(clippy::too_many_arguments)]
19560    pub fn fa_prefill_hd512(
19561        &self,
19562        q: &CudaSlice<f32>,
19563        k: &CudaSlice<f32>,
19564        v: &CudaSlice<f32>,
19565        o: &mut CudaSlice<f32>,
19566        head_dim: usize,
19567        n_head: usize,
19568        n_head_kv: usize,
19569        t: usize,
19570        t_kv: usize,
19571        scale: f32,
19572        causal: bool,
19573    ) -> Result<(), Box<dyn std::error::Error>> {
19574        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
19575        if portable_mma_gated() {
19576            return self.sdpa_naive(
19577                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19578            );
19579        }
19580        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
19581        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
19582        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
19583        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
19584        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
19585        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19586        let f32_stage =
19587            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19588        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
19589        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
19590        // Own numeric config (partial-sum order) — battery-gated.
19591        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19592        let sp = !f32_stage
19593            && *SP_ON.get_or_init(|| {
19594                std::env::var("MEMRA_FA512_SP")
19595                    .map(|v| v != "0")
19596                    .unwrap_or(true)
19597            });
19598        self.fa_prefill_hd512_arm(
19599            q,
19600            k,
19601            v,
19602            o,
19603            head_dim,
19604            n_head,
19605            n_head_kv,
19606            t,
19607            t_kv,
19608            scale,
19609            causal,
19610            f32_stage,
19611            sp,
19612            sp && fa_f16pv_on(),
19613        )
19614    }
19615
19616    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
19617    #[allow(clippy::too_many_arguments)]
19618    pub fn fa_prefill_hd512_pre(
19619        &self,
19620        qb: &CudaSlice<u8>,
19621        kb: &CudaSlice<u8>,
19622        vb: &CudaSlice<u8>,
19623        o: &mut CudaSlice<f32>,
19624        head_dim: usize,
19625        n_head: usize,
19626        n_head_kv: usize,
19627        t: usize,
19628        t_kv: usize,
19629        scale: f32,
19630        causal: bool,
19631        v_f16: bool,
19632    ) -> Result<(), Box<dyn std::error::Error>> {
19633        debug_assert_eq!(head_dim, 512);
19634        const SP_M: usize = 16;
19635        const BKS: usize = 32;
19636        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
19637        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
19638        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
19639        let f16pv = fa_f16pv_on();
19640        let nw = if f16pv { fa512_wide_warps() } else { 2 };
19641        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19642        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19643        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19644        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19645            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
19646            let n = t_kv * n_head_kv * head_dim;
19647            let need = n * 2;
19648            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19649                *vguard = Some(self.alloc_uninit::<u8>(need)?);
19650            }
19651            let dst = vguard.as_mut().unwrap();
19652            self.bf16_to_f16_into(vb, n, dst)?;
19653            vguard.as_ref().unwrap()
19654        } else {
19655            vb
19656        };
19657        let f = self.func(if hp {
19658            "fa_prefill_bf16_hd512_sp16h2"
19659        } else {
19660            match (f16pv, nw) {
19661                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19662                (true, _) => "fa_prefill_bf16_hd512_sp16",
19663                _ => "fa_prefill_bf16_hd512_sp",
19664            }
19665        });
19666        let (nwarp, npart) = if hp {
19667            (4usize, 4usize)
19668        } else if nw > 2 {
19669            (nw, nw)
19670        } else {
19671            (2, 1)
19672        };
19673        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
19674        let shmem = if hp {
19675            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19676                as u32
19677        } else {
19678            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19679                + 4 * (npart * SP_M * BKS + SP_M)) as u32
19680        };
19681        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19682        f.set_attribute(
19683            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19684            shmem as i32,
19685        )?;
19686        let grid_y = if hp {
19687            (n_head / 2) as u32
19688        } else {
19689            n_head as u32
19690        };
19691        let cfg = LaunchConfig {
19692            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19693            block_dim: (32, nwarp as u32, 1),
19694            shared_mem_bytes: shmem,
19695        };
19696        let (hd, nh, nhkv, ti, tkvi, cz) = (
19697            head_dim as i32,
19698            n_head as i32,
19699            n_head_kv as i32,
19700            t as i32,
19701            t_kv as i32,
19702            causal as i32,
19703        );
19704        let __s_b = self.gpu.stream();
19705        let mut b = __s_b.launch_builder(&f);
19706        b.arg(qb)
19707            .arg(kb)
19708            .arg(vref)
19709            .arg(o)
19710            .arg(&hd)
19711            .arg(&nh)
19712            .arg(&nhkv)
19713            .arg(&ti)
19714            .arg(&tkvi)
19715            .arg(&scale)
19716            .arg(&cz);
19717        unsafe {
19718            b.launch(cfg)?;
19719        }
19720        Ok(())
19721    }
19722
19723    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
19724    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
19725    #[allow(clippy::too_many_arguments)]
19726    pub fn fa_prefill_hd512_arm(
19727        &self,
19728        q: &CudaSlice<f32>,
19729        k: &CudaSlice<f32>,
19730        v: &CudaSlice<f32>,
19731        o: &mut CudaSlice<f32>,
19732        head_dim: usize,
19733        n_head: usize,
19734        n_head_kv: usize,
19735        t: usize,
19736        t_kv: usize,
19737        scale: f32,
19738        causal: bool,
19739        f32_stage: bool,
19740        sp: bool,
19741        f16pv: bool,
19742    ) -> Result<(), Box<dyn std::error::Error>> {
19743        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19744        if sp && !f32_stage {
19745            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
19746            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
19747            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
19748            const SP_M: usize = 16;
19749            const BKS: usize = 32;
19750            let nw = if f16pv { fa512_wide_warps() } else { 2 };
19751            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19752            let f = self.func(if hp {
19753                "fa_prefill_bf16_hd512_sp16h2"
19754            } else {
19755                match (f16pv, nw) {
19756                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19757                    (true, _) => "fa_prefill_bf16_hd512_sp16",
19758                    _ => "fa_prefill_bf16_hd512_sp",
19759                }
19760            });
19761            let (nwarp, npart) = if hp {
19762                (4usize, 4usize)
19763            } else if nw > 2 {
19764                (nw, nw)
19765            } else {
19766                (2, 1)
19767            };
19768            let shmem = if hp {
19769                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19770                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19771            } else {
19772                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19773                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
19774            };
19775            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19776            f.set_attribute(
19777                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19778                shmem as i32,
19779            )?;
19780            let grid_y = if hp {
19781                (n_head / 2) as u32
19782            } else {
19783                n_head as u32
19784            };
19785            let cfg = LaunchConfig {
19786                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19787                block_dim: (32, nwarp as u32, 1),
19788                shared_mem_bytes: shmem,
19789            };
19790            let (hd, nh, nhkv, ti, tkvi, cz) = (
19791                head_dim as i32,
19792                n_head as i32,
19793                n_head_kv as i32,
19794                t as i32,
19795                t_kv as i32,
19796                causal as i32,
19797            );
19798            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19799            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19800            let vb = if f16pv {
19801                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19802            } else {
19803                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19804            };
19805            let __s_b = self.gpu.stream();
19806            let mut b = __s_b.launch_builder(&f);
19807            b.arg(&qb)
19808                .arg(&kb)
19809                .arg(&vb)
19810                .arg(o)
19811                .arg(&hd)
19812                .arg(&nh)
19813                .arg(&nhkv)
19814                .arg(&ti)
19815                .arg(&tkvi)
19816                .arg(&scale)
19817                .arg(&cz);
19818            unsafe {
19819                b.launch(cfg)?;
19820            }
19821            return Ok(());
19822        }
19823        const BLOCK_Q: usize = 32;
19824        const BK: usize = 32;
19825        const HALF: usize = 256;
19826        let f = self.func(if f32_stage {
19827            "fa_prefill_f32_hd512"
19828        } else {
19829            "fa_prefill_bf16_hd512"
19830        });
19831        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
19832        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19833            + 4 * BLOCK_Q) as u32;
19834        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19835        f.set_attribute(
19836            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19837            shmem as i32,
19838        )?;
19839        let cfg = LaunchConfig {
19840            grid_dim: (
19841                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19842                n_head as u32,
19843                2,
19844            ),
19845            block_dim: (32, 2, 1),
19846            shared_mem_bytes: shmem,
19847        };
19848        let (hd, nh, nhkv, ti, tkvi, cz) = (
19849            head_dim as i32,
19850            n_head as i32,
19851            n_head_kv as i32,
19852            t as i32,
19853            t_kv as i32,
19854            causal as i32,
19855        );
19856        if f32_stage {
19857            let __s_b = self.gpu.stream();
19858            let mut b = __s_b.launch_builder(&f);
19859            b.arg(q)
19860                .arg(k)
19861                .arg(v)
19862                .arg(o)
19863                .arg(&hd)
19864                .arg(&nh)
19865                .arg(&nhkv)
19866                .arg(&ti)
19867                .arg(&tkvi)
19868                .arg(&scale)
19869                .arg(&cz);
19870            unsafe {
19871                b.launch(cfg)?;
19872            }
19873        } else {
19874            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19875            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19876            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19877            let __s_b = self.gpu.stream();
19878            let mut b = __s_b.launch_builder(&f);
19879            b.arg(&qb)
19880                .arg(&kb)
19881                .arg(&vb)
19882                .arg(o)
19883                .arg(&hd)
19884                .arg(&nh)
19885                .arg(&nhkv)
19886                .arg(&ti)
19887                .arg(&tkvi)
19888                .arg(&scale)
19889                .arg(&cz);
19890            unsafe {
19891                b.launch(cfg)?;
19892            }
19893        }
19894        Ok(())
19895    }
19896
19897    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
19898    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
19899    /// separate f32_to_bf16 the FA entries would run).
19900    #[allow(clippy::too_many_arguments)]
19901    pub fn rope_neox2_bf16e(
19902        &self,
19903        q: &mut CudaSlice<f32>,
19904        k: &mut CudaSlice<f32>,
19905        qb: &mut CudaSlice<u8>,
19906        kb: &mut CudaSlice<u8>,
19907        pos: &CudaSlice<i32>,
19908        head_dim: usize,
19909        n_dims: usize,
19910        nh_q: usize,
19911        nh_k: usize,
19912        n_tokens: usize,
19913        base: f32,
19914        freq_scale: f32,
19915        ff: Option<&CudaSlice<f32>>,
19916    ) -> Result<(), Box<dyn std::error::Error>> {
19917        let f = self.func("rope_neox2_bf16e_f32");
19918        let rows = ((nh_q + nh_k) * n_tokens) as u32;
19919        let cfg = LaunchConfig {
19920            grid_dim: (rows, 1, 1),
19921            block_dim: ((head_dim / 2) as u32, 1, 1),
19922            shared_mem_bytes: 0,
19923        };
19924        let theta_scale = base.powf(-2.0 / n_dims as f32);
19925        let (hd, nd, nhq, nhk, nt) = (
19926            head_dim as i32,
19927            n_dims as i32,
19928            nh_q as i32,
19929            nh_k as i32,
19930            n_tokens as i32,
19931        );
19932        let __s_b = self.gpu.stream();
19933        let mut b = __s_b.launch_builder(&f);
19934        match ff {
19935            Some(t) => {
19936                b.arg(&mut *q)
19937                    .arg(&mut *k)
19938                    .arg(&mut *qb)
19939                    .arg(&mut *kb)
19940                    .arg(pos)
19941                    .arg(&hd)
19942                    .arg(&nd)
19943                    .arg(&nhq)
19944                    .arg(&nhk)
19945                    .arg(&nt)
19946                    .arg(&theta_scale)
19947                    .arg(&freq_scale)
19948                    .arg(t);
19949                unsafe {
19950                    b.launch(cfg)?;
19951                }
19952            }
19953            None => {
19954                let null: u64 = 0;
19955                b.arg(&mut *q)
19956                    .arg(&mut *k)
19957                    .arg(&mut *qb)
19958                    .arg(&mut *kb)
19959                    .arg(pos)
19960                    .arg(&hd)
19961                    .arg(&nd)
19962                    .arg(&nhq)
19963                    .arg(&nhk)
19964                    .arg(&nt)
19965                    .arg(&theta_scale)
19966                    .arg(&freq_scale)
19967                    .arg(&null);
19968                unsafe {
19969                    b.launch(cfg)?;
19970                }
19971            }
19972        }
19973        Ok(())
19974    }
19975
19976    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
19977    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
19978    pub fn f32_to_bf16(
19979        &self,
19980        x: &CudaSlice<f32>,
19981        n: usize,
19982    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19983        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19984        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19985        let f = self.func("f32_to_bf16_flat");
19986        let n_i = n as i64;
19987        let cfg = LaunchConfig {
19988            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19989            block_dim: (256, 1, 1),
19990            shared_mem_bytes: 0,
19991        };
19992        let __s_b = self.gpu.stream();
19993        let mut b = __s_b.launch_builder(&f);
19994        b.arg(x).arg(&mut y).arg(&n_i);
19995        unsafe {
19996            b.launch(cfg)?;
19997        }
19998        Ok(y)
19999    }
20000
20001    pub fn f32_to_f16(
20002        &self,
20003        x: &CudaSlice<f32>,
20004        n: usize,
20005    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20006        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20007        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20008        let f = self.func("f32_to_f16_flat");
20009        let n_i = n as i64;
20010        let cfg = LaunchConfig {
20011            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20012            block_dim: (256, 1, 1),
20013            shared_mem_bytes: 0,
20014        };
20015        let __s_b = self.gpu.stream();
20016        let mut b = __s_b.launch_builder(&f);
20017        b.arg(x).arg(&mut y).arg(&n_i);
20018        unsafe {
20019            b.launch(cfg)?;
20020        }
20021        Ok(y)
20022    }
20023
20024    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
20025    pub fn bf16_to_f16(
20026        &self,
20027        xb: &CudaSlice<u8>,
20028        n: usize,
20029    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20030        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20031        self.bf16_to_f16_into(xb, n, &mut y)?;
20032        Ok(y)
20033    }
20034
20035    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
20036    pub fn bf16_to_f16_into(
20037        &self,
20038        xb: &CudaSlice<u8>,
20039        n: usize,
20040        y: &mut CudaSlice<u8>,
20041    ) -> Result<(), Box<dyn std::error::Error>> {
20042        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
20043        assert!(y.len() >= n * 2);
20044        let f = self.func("bf16_to_f16_flat");
20045        let n2 = (n / 2) as i64;
20046        let cfg = LaunchConfig {
20047            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
20048            block_dim: (256, 1, 1),
20049            shared_mem_bytes: 0,
20050        };
20051        let __s_b = self.gpu.stream();
20052        let mut b = __s_b.launch_builder(&f);
20053        b.arg(xb).arg(y).arg(&n2);
20054        unsafe {
20055            b.launch(cfg)?;
20056        }
20057        Ok(())
20058    }
20059
20060    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
20061    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
20062    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
20063    /// head_dim in {256, 128}, bf16kv lane on.
20064    #[allow(clippy::too_many_arguments)]
20065    pub fn fa_prefill_vl8(
20066        &self,
20067        seqs: &[FaSeqVl],
20068        head_dim: usize,
20069        n_head: usize,
20070        n_head_kv: usize,
20071        scale: f32,
20072    ) -> Result<(), Box<dyn std::error::Error>> {
20073        const BK: usize = 32;
20074        let b = seqs.len();
20075        assert!(b >= 1 && b <= 8);
20076        let mut packed = [FaSeqVl::default(); 8];
20077        packed[..b].copy_from_slice(seqs);
20078        let v = FaVl8(packed);
20079        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20080        let ept = (n_head_kv * head_dim) as i32;
20081        {
20082            let f = self.func("fa_mirror_vl");
20083            let max_n = (max_t as i64) * ept as i64;
20084            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
20085            for which in 0..2i32 {
20086                let cfg = LaunchConfig {
20087                    grid_dim: (blocks, 1, b as u32),
20088                    block_dim: (256, 1, 1),
20089                    shared_mem_bytes: 0,
20090                };
20091                let __s_lb = self.gpu.stream();
20092                let mut lb = __s_lb.launch_builder(&f);
20093                lb.arg(&v).arg(&ept).arg(&which);
20094                unsafe {
20095                    lb.launch(cfg)?;
20096                }
20097            }
20098        }
20099        let hd_sfx = fa_hd_suffix(head_dim)?;
20100        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
20101        let block_q = 64usize;
20102        let kv_stages = 2usize;
20103        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20104            + 4 * (block_q * BK + 2 * block_q)) as u32;
20105        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20106        f.set_attribute(
20107            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20108            shmem as i32,
20109        )?;
20110        let cfg = LaunchConfig {
20111            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
20112            block_dim: (32, 4, 1),
20113            shared_mem_bytes: shmem,
20114        };
20115        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20116        let __s_lb = self.gpu.stream();
20117        let mut lb = __s_lb.launch_builder(&f);
20118        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
20119        unsafe {
20120            lb.launch(cfg)?;
20121        }
20122        Ok(())
20123    }
20124
20125    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
20126    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
20127    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
20128    #[allow(clippy::too_many_arguments)]
20129    pub fn attn_pre_vl8(
20130        &self,
20131        seqs: &[AttnPreVl],
20132        wq: &CudaSlice<f32>,
20133        wk: &CudaSlice<f32>,
20134        head_dim: usize,
20135        rope_dims: usize,
20136        n_head: usize,
20137        n_head_kv: usize,
20138        eps: f32,
20139        freq_base: f32,
20140        freq_scale: f32,
20141        kv_dim_k: usize,
20142        kv_dim_v: usize,
20143        k_tok_bytes: usize,
20144        v_tok_bytes: usize,
20145    ) -> Result<(), Box<dyn std::error::Error>> {
20146        let b = seqs.len();
20147        assert!(b >= 1 && b <= 8);
20148        let mut packed = [AttnPreVl::default(); 8];
20149        packed[..b].copy_from_slice(seqs);
20150        let v = AttnPreVl8(packed);
20151        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20152        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20153        {
20154            let f = self.func("q_gate_split_vl");
20155            let n = max_t * (n_head * head_dim) as u32;
20156            let cfg = LaunchConfig {
20157                grid_dim: (n.div_ceil(256), 1, b as u32),
20158                block_dim: (256, 1, 1),
20159                shared_mem_bytes: 0,
20160            };
20161            let __s_lb = self.gpu.stream();
20162            let mut lb = __s_lb.launch_builder(&f);
20163            lb.arg(&v).arg(&hd).arg(&nh);
20164            unsafe {
20165                lb.launch(cfg)?;
20166            }
20167        }
20168        {
20169            let f = self.func("attn_rms_vl");
20170            let cfg = LaunchConfig {
20171                grid_dim: (max_t * n_head as u32, 2, b as u32),
20172                block_dim: (rms_block(), 1, 1),
20173                shared_mem_bytes: 0,
20174            };
20175            let __s_lb = self.gpu.stream();
20176            let mut lb = __s_lb.launch_builder(&f);
20177            lb.arg(&v)
20178                .arg(wq)
20179                .arg(wk)
20180                .arg(&hd)
20181                .arg(&nh)
20182                .arg(&nhkv)
20183                .arg(&eps);
20184            unsafe {
20185                lb.launch(cfg)?;
20186            }
20187        }
20188        {
20189            let f = self.func("attn_rope_vl");
20190            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20191            let nd = rope_dims as i32;
20192            let cfg = LaunchConfig {
20193                grid_dim: (max_t * n_head as u32, 2, b as u32),
20194                block_dim: ((head_dim / 2) as u32, 1, 1),
20195                shared_mem_bytes: 0,
20196            };
20197            let __s_lb = self.gpu.stream();
20198            let mut lb = __s_lb.launch_builder(&f);
20199            lb.arg(&v)
20200                .arg(&hd)
20201                .arg(&nd)
20202                .arg(&nh)
20203                .arg(&nhkv)
20204                .arg(&theta_scale)
20205                .arg(&freq_scale);
20206            unsafe {
20207                lb.launch(cfg)?;
20208            }
20209        }
20210        {
20211            let f = self.func("append_kv_vl");
20212            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20213            let cfg = LaunchConfig {
20214                grid_dim: (nblk, max_t, b as u32),
20215                block_dim: (32, 1, 1),
20216                shared_mem_bytes: 0,
20217            };
20218            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20219            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20220            let __s_lb = self.gpu.stream();
20221            let mut lb = __s_lb.launch_builder(&f);
20222            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20223            unsafe {
20224                lb.launch(cfg)?;
20225            }
20226        }
20227        Ok(())
20228    }
20229
20230    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
20231    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
20232    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
20233    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
20234    pub fn fa_prefill_view(
20235        &self,
20236        q: &CudaSlice<f32>,
20237        k: &cudarc::driver::CudaView<u8>,
20238        v: &cudarc::driver::CudaView<u8>,
20239        o: &mut CudaSlice<f32>,
20240        head_dim: usize,
20241        n_head: usize,
20242        n_head_kv: usize,
20243        t: usize,
20244        t_kv: usize,
20245        scale: f32,
20246        causal: bool,
20247        k_tok_bytes: usize,
20248        v_tok_bytes: usize,
20249        g: bool,
20250    ) -> Result<(), Box<dyn std::error::Error>> {
20251        if portable_mma_gated() {
20252            return self.sdpa_naive_quantized_view(
20253                q,
20254                k,
20255                v,
20256                o,
20257                head_dim,
20258                n_head,
20259                n_head_kv,
20260                t,
20261                t_kv,
20262                scale,
20263                causal,
20264                k_tok_bytes,
20265                v_tok_bytes,
20266            );
20267        }
20268        const BLOCK_Q: usize = 64;
20269        const BK: usize = 32;
20270        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
20271        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
20272        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20273        let f = if g {
20274            self.func_g(&name)
20275        } else {
20276            self.func(&name)
20277        };
20278        let shmem =
20279            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20280        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20281        f.set_attribute(
20282            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20283            shmem as i32,
20284        )?;
20285        let cfg = LaunchConfig {
20286            grid_dim: (
20287                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20288                n_head as u32,
20289                1,
20290            ),
20291            block_dim: (32, 4, 1),
20292            shared_mem_bytes: shmem,
20293        };
20294        let (hd, nh, nhkv, ti, tkvi, cz) = (
20295            head_dim as i32,
20296            n_head as i32,
20297            n_head_kv as i32,
20298            t as i32,
20299            t_kv as i32,
20300            causal as i32,
20301        );
20302        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20303        let __s_b = self.gpu.stream();
20304        let mut b = __s_b.launch_builder(&f);
20305        b.arg(q)
20306            .arg(k)
20307            .arg(v)
20308            .arg(o)
20309            .arg(&hd)
20310            .arg(&nh)
20311            .arg(&nhkv)
20312            .arg(&ti)
20313            .arg(&tkvi)
20314            .arg(&scale)
20315            .arg(&cz)
20316            .arg(&ktb)
20317            .arg(&vtb);
20318        unsafe {
20319            b.launch(cfg)?;
20320        }
20321        Ok(())
20322    }
20323
20324    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
20325    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
20326    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
20327    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
20328    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
20329    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
20330    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
20331    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
20332    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
20333    #[allow(clippy::too_many_arguments)]
20334    pub fn fa_prefill_view_ws(
20335        &self,
20336        q: &CudaSlice<f32>,
20337        k: &cudarc::driver::CudaView<u8>,
20338        v: &cudarc::driver::CudaView<u8>,
20339        o: &mut CudaSlice<f32>,
20340        head_dim: usize,
20341        n_head: usize,
20342        n_head_kv: usize,
20343        t: usize,
20344        t_kv: usize,
20345        scale: f32,
20346        causal: bool,
20347        k_tok_bytes: usize,
20348        v_tok_bytes: usize,
20349        g: bool,
20350    ) -> Result<(), Box<dyn std::error::Error>> {
20351        if portable_mma_gated() {
20352            return self.sdpa_naive_quantized_view(
20353                q,
20354                k,
20355                v,
20356                o,
20357                head_dim,
20358                n_head,
20359                n_head_kv,
20360                t,
20361                t_kv,
20362                scale,
20363                causal,
20364                k_tok_bytes,
20365                v_tok_bytes,
20366            );
20367        }
20368        const BLOCK_Q: usize = 64;
20369        const BK: usize = 32;
20370        let kv_dim_k = n_head_kv * head_dim;
20371        let kv_dim_v = n_head_kv * head_dim;
20372        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20373        let v_ws_bytes = t_kv * kv_dim_v * 2;
20374        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
20375        let mut guard = self.prime_deqw_ws.lock().unwrap();
20376        let need_grow = match guard.as_ref() {
20377            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20378            None => true,
20379        };
20380        if need_grow {
20381            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20382            let (ck, cv) = guard
20383                .as_ref()
20384                .map(|(a, b)| (a.len(), b.len()))
20385                .unwrap_or((0, 0));
20386            *guard = Some((
20387                self.alloc_u8(grow(ck, k_ws_bytes))?,
20388                self.alloc_u8(grow(cv, v_ws_bytes))?,
20389            ));
20390        }
20391        let (kw, vw) = guard.as_mut().unwrap();
20392        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
20393        {
20394            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
20395            let f = if g {
20396                self.func_g("fa_dequant_kv_ws_bf16")
20397            } else {
20398                self.func("fa_dequant_kv_ws_bf16")
20399            };
20400            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20401            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20402            let cfg = LaunchConfig {
20403                grid_dim: (nblk.max(1), 1, 1),
20404                block_dim: (256, 1, 1),
20405                shared_mem_bytes: 0,
20406            };
20407            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20408            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20409            let __s_b = self.gpu.stream();
20410            let mut b = __s_b.launch_builder(&f);
20411            b.arg(k)
20412                .arg(v)
20413                .arg(&mut *kw)
20414                .arg(&mut *vw)
20415                .arg(&kdk)
20416                .arg(&kdv)
20417                .arg(&tkvi)
20418                .arg(&ktb)
20419                .arg(&vtb);
20420            unsafe {
20421                b.launch(cfg)?;
20422            }
20423        }
20424        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
20425        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
20426        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
20427        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
20428        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
20429        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
20430        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
20431        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20432            .map(|v| v != "0")
20433            .unwrap_or(true);
20434        {
20435            let hd_sfx = fa_hd_suffix(head_dim)?;
20436            let f = self.func(&format!(
20437                "fa_prefill_qw{}{hd_sfx}",
20438                if db { "_db" } else { "" }
20439            ));
20440            let shmem = if db {
20441                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
20442                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20443            } else {
20444                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20445            };
20446            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20447            f.set_attribute(
20448                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20449                shmem as i32,
20450            )?;
20451            let cfg = LaunchConfig {
20452                grid_dim: (
20453                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20454                    n_head as u32,
20455                    1,
20456                ),
20457                block_dim: (32, 4, 1),
20458                shared_mem_bytes: shmem,
20459            };
20460            let (hd, nh, nhkv, ti, tkvi, cz) = (
20461                head_dim as i32,
20462                n_head as i32,
20463                n_head_kv as i32,
20464                t as i32,
20465                t_kv as i32,
20466                causal as i32,
20467            );
20468            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20469            let __s_b = self.gpu.stream();
20470            let mut b = __s_b.launch_builder(&f);
20471            b.arg(q)
20472                .arg(&*kw)
20473                .arg(&*vw)
20474                .arg(o)
20475                .arg(&hd)
20476                .arg(&nh)
20477                .arg(&nhkv)
20478                .arg(&ti)
20479                .arg(&tkvi)
20480                .arg(&scale)
20481                .arg(&cz)
20482                .arg(&kdk)
20483                .arg(&kdv);
20484            unsafe {
20485                b.launch(cfg)?;
20486            }
20487        }
20488        Ok(())
20489    }
20490
20491    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
20492    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
20493    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
20494    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
20495    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
20496    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
20497    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
20498    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
20499    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
20500    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
20501    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
20502    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
20503    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
20504    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
20505    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
20506    #[allow(clippy::too_many_arguments)]
20507    pub fn fa_prefill_view_ws_w_hd128(
20508        &self,
20509        q: &CudaSlice<f32>,
20510        k: &cudarc::driver::CudaView<u8>,
20511        v: &cudarc::driver::CudaView<u8>,
20512        o: &mut CudaSlice<f32>,
20513        head_dim: usize,
20514        n_head: usize,
20515        n_head_kv: usize,
20516        t: usize,
20517        t_kv: usize,
20518        scale: f32,
20519        causal: bool,
20520        window: usize,
20521        k_tok_bytes: usize,
20522        v_tok_bytes: usize,
20523    ) -> Result<(), Box<dyn std::error::Error>> {
20524        assert_eq!(
20525            head_dim, 128,
20526            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20527        );
20528        if portable_mma_gated() {
20529            return self.sdpa_naive_w_quantized_view(
20530                q,
20531                k,
20532                v,
20533                o,
20534                head_dim,
20535                n_head,
20536                n_head_kv,
20537                t,
20538                t_kv,
20539                scale,
20540                causal,
20541                window,
20542                k_tok_bytes,
20543                v_tok_bytes,
20544            );
20545        }
20546        const BLOCK_Q: usize = 64;
20547        const BK: usize = 32;
20548        let kv_dim_k = n_head_kv * head_dim;
20549        let kv_dim_v = n_head_kv * head_dim;
20550        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20551        let v_ws_bytes = t_kv * kv_dim_v * 2;
20552        let mut guard = self.prime_deqw_ws.lock().unwrap();
20553        let need_grow = match guard.as_ref() {
20554            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20555            None => true,
20556        };
20557        if need_grow {
20558            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20559            let (ck, cv) = guard
20560                .as_ref()
20561                .map(|(a, b)| (a.len(), b.len()))
20562                .unwrap_or((0, 0));
20563            *guard = Some((
20564                self.alloc_u8(grow(ck, k_ws_bytes))?,
20565                self.alloc_u8(grow(cv, v_ws_bytes))?,
20566            ));
20567        }
20568        let (kw, vw) = guard.as_mut().unwrap();
20569        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
20570        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
20571        {
20572            let f = self.func("fa_dequant_kv_ws_bf16");
20573            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20574            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20575            let cfg = LaunchConfig {
20576                grid_dim: (nblk.max(1), 1, 1),
20577                block_dim: (256, 1, 1),
20578                shared_mem_bytes: 0,
20579            };
20580            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20581            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20582            let __s_b = self.gpu.stream();
20583            let mut b = __s_b.launch_builder(&f);
20584            b.arg(k)
20585                .arg(v)
20586                .arg(&mut *kw)
20587                .arg(&mut *vw)
20588                .arg(&kdk)
20589                .arg(&kdv)
20590                .arg(&tkvi)
20591                .arg(&ktb)
20592                .arg(&vtb);
20593            unsafe {
20594                b.launch(cfg)?;
20595            }
20596        }
20597        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
20598        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20599            .map(|v| v != "0")
20600            .unwrap_or(true);
20601        {
20602            let f = self.func(if db {
20603                "fa_prefill_qw_db_w_hd128"
20604            } else {
20605                "fa_prefill_qw_w_hd128"
20606            });
20607            let shmem = if db {
20608                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20609            } else {
20610                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20611            };
20612            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20613            f.set_attribute(
20614                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20615                shmem as i32,
20616            )?;
20617            let cfg = LaunchConfig {
20618                grid_dim: (
20619                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20620                    n_head as u32,
20621                    1,
20622                ),
20623                block_dim: (32, 4, 1),
20624                shared_mem_bytes: shmem,
20625            };
20626            let (hd, nh, nhkv, ti, tkvi, cz) = (
20627                head_dim as i32,
20628                n_head as i32,
20629                n_head_kv as i32,
20630                t as i32,
20631                t_kv as i32,
20632                causal as i32,
20633            );
20634            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20635            let __s_b = self.gpu.stream();
20636            let mut b = __s_b.launch_builder(&f);
20637            b.arg(q)
20638                .arg(&*kw)
20639                .arg(&*vw)
20640                .arg(o)
20641                .arg(&hd)
20642                .arg(&nh)
20643                .arg(&nhkv)
20644                .arg(&ti)
20645                .arg(&tkvi)
20646                .arg(&scale)
20647                .arg(&cz)
20648                .arg(&kdk)
20649                .arg(&kdv)
20650                .arg(&wnd);
20651            unsafe {
20652                b.launch(cfg)?;
20653            }
20654        }
20655        Ok(())
20656    }
20657
20658    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
20659    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
20660    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
20661    pub fn fa_decode(
20662        &self,
20663        q: &CudaSlice<f32>,
20664        k: &cudarc::driver::CudaView<u8>,
20665        v: &cudarc::driver::CudaView<u8>,
20666        o: &mut CudaSlice<f32>,
20667        head_dim: usize,
20668        n_head: usize,
20669        n_head_kv: usize,
20670        t_kv: usize,
20671        scale: f32,
20672        k_tok_bytes: usize,
20673        v_tok_bytes: usize,
20674    ) -> Result<(), Box<dyn std::error::Error>> {
20675        self.fa_decode_kvmod(
20676            q,
20677            k,
20678            v,
20679            o,
20680            head_dim,
20681            n_head,
20682            n_head_kv,
20683            t_kv,
20684            scale,
20685            k_tok_bytes,
20686            v_tok_bytes,
20687            false,
20688        )
20689    }
20690
20691    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
20692    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
20693    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
20694    #[allow(clippy::too_many_arguments)]
20695    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
20696    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
20697    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
20698    #[allow(clippy::too_many_arguments)]
20699    #[allow(clippy::too_many_arguments)]
20700    fn fa_decode_scalar_unified(
20701        &self,
20702        q: &cudarc::driver::CudaView<f32>,
20703        k: &cudarc::driver::CudaView<u8>,
20704        v: &cudarc::driver::CudaView<u8>,
20705        o: &mut cudarc::driver::CudaViewMut<f32>,
20706        head_dim: usize,
20707        n_head: usize,
20708        n_head_kv: usize,
20709        t_kv_host: usize,
20710        t_kv_dev: Option<&CudaSlice<i32>>,
20711        scale: f32,
20712        n_splits: usize,
20713        split_keys: usize,
20714        k_tok_bytes: usize,
20715        v_tok_bytes: usize,
20716        g: bool,
20717        part_o: &mut CudaSlice<f32>,
20718        part_m: &mut CudaSlice<f32>,
20719        part_l: &mut CudaSlice<f32>,
20720        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20721    ) -> Result<(), Box<dyn std::error::Error>> {
20722        let f = if g {
20723            self.func_g("fa_decode_f32")
20724        } else {
20725            self.fa_func("fa_decode_f32", head_dim)
20726        };
20727        let cfg = LaunchConfig {
20728            grid_dim: (n_head as u32, n_splits as u32, 1),
20729            block_dim: (head_dim as u32, 1, 1),
20730            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20731        };
20732        let (hd, nh, nhkv, nsp) = (
20733            head_dim as i32,
20734            n_head as i32,
20735            n_head_kv as i32,
20736            n_splits as i32,
20737        );
20738        let (ktb, vtb, tkvi, ski) = (
20739            k_tok_bytes as i64,
20740            v_tok_bytes as i64,
20741            t_kv_host as i32,
20742            split_keys as i32,
20743        );
20744        let __s_b = self.gpu.stream();
20745        let mut b = __s_b.launch_builder(&f);
20746        match t_kv_dev {
20747            Some(d) => {
20748                b.arg(q)
20749                    .arg(k)
20750                    .arg(v)
20751                    .arg(&mut *part_o)
20752                    .arg(&mut *part_m)
20753                    .arg(&mut *part_l)
20754                    .arg(&hd)
20755                    .arg(&nh)
20756                    .arg(&nhkv)
20757                    .arg(&tkvi)
20758                    .arg(d)
20759                    .arg(&scale)
20760                    .arg(&nsp)
20761                    .arg(&ski)
20762                    .arg(&ktb)
20763                    .arg(&vtb);
20764                unsafe {
20765                    b.launch(cfg)?;
20766                }
20767            }
20768            None => {
20769                let null: u64 = 0;
20770                b.arg(q)
20771                    .arg(k)
20772                    .arg(v)
20773                    .arg(&mut *part_o)
20774                    .arg(&mut *part_m)
20775                    .arg(&mut *part_l)
20776                    .arg(&hd)
20777                    .arg(&nh)
20778                    .arg(&nhkv)
20779                    .arg(&tkvi)
20780                    .arg(&null)
20781                    .arg(&scale)
20782                    .arg(&nsp)
20783                    .arg(&ski)
20784                    .arg(&ktb)
20785                    .arg(&vtb);
20786                unsafe {
20787                    b.launch(cfg)?;
20788                }
20789            }
20790        }
20791        let cfg2 = LaunchConfig {
20792            grid_dim: (n_head as u32, 1, 1),
20793            block_dim: (head_dim as u32, 1, 1),
20794            shared_mem_bytes: 0,
20795        };
20796        if let Some((oq, od)) = q8_out {
20797            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
20798            let fc = if g {
20799                self.func_g("fa_decode_combine_q8_1")
20800            } else {
20801                self.fa_func("fa_decode_combine_q8_1", head_dim)
20802            };
20803            let __s_b2 = self.gpu.stream();
20804            let mut b2 = __s_b2.launch_builder(&fc);
20805            b2.arg(&*part_o)
20806                .arg(&*part_m)
20807                .arg(&*part_l)
20808                .arg(oq)
20809                .arg(od)
20810                .arg(&hd)
20811                .arg(&nh)
20812                .arg(&nsp);
20813            unsafe {
20814                b2.launch(cfg2)?;
20815            }
20816            return Ok(());
20817        }
20818        let fc = if g {
20819            self.func_g("fa_decode_combine_f32")
20820        } else {
20821            self.fa_func("fa_decode_combine_f32", head_dim)
20822        };
20823        let __s_b2 = self.gpu.stream();
20824        let mut b2 = __s_b2.launch_builder(&fc);
20825        b2.arg(&*part_o)
20826            .arg(&*part_m)
20827            .arg(&*part_l)
20828            .arg(o)
20829            .arg(&hd)
20830            .arg(&nh)
20831            .arg(&nsp);
20832        unsafe {
20833            b2.launch(cfg2)?;
20834        }
20835        Ok(())
20836    }
20837
20838    pub fn fa_decode_kvmod(
20839        &self,
20840        q: &CudaSlice<f32>,
20841        k: &cudarc::driver::CudaView<u8>,
20842        v: &cudarc::driver::CudaView<u8>,
20843        o: &mut CudaSlice<f32>,
20844        head_dim: usize,
20845        n_head: usize,
20846        n_head_kv: usize,
20847        t_kv: usize,
20848        scale: f32,
20849        k_tok_bytes: usize,
20850        v_tok_bytes: usize,
20851        g: bool,
20852    ) -> Result<(), Box<dyn std::error::Error>> {
20853        let q_view = q.as_view();
20854        let mut o_view = o.as_view_mut();
20855        self.fa_decode_kvmod_view(
20856            &q_view,
20857            k,
20858            v,
20859            &mut o_view,
20860            head_dim,
20861            n_head,
20862            n_head_kv,
20863            t_kv,
20864            scale,
20865            k_tok_bytes,
20866            v_tok_bytes,
20867            g,
20868        )
20869    }
20870
20871    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
20872    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
20873    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
20874    /// per-session KV view and FA launch.
20875    #[allow(clippy::too_many_arguments)]
20876    pub fn fa_decode_kvmod_view(
20877        &self,
20878        q: &cudarc::driver::CudaView<f32>,
20879        k: &cudarc::driver::CudaView<u8>,
20880        v: &cudarc::driver::CudaView<u8>,
20881        o: &mut cudarc::driver::CudaViewMut<f32>,
20882        head_dim: usize,
20883        n_head: usize,
20884        n_head_kv: usize,
20885        t_kv: usize,
20886        scale: f32,
20887        k_tok_bytes: usize,
20888        v_tok_bytes: usize,
20889        g: bool,
20890    ) -> Result<(), Box<dyn std::error::Error>> {
20891        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
20892        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
20893        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
20894        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
20895        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
20896        //
20897        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
20898        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
20899        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
20900        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
20901        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
20902        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
20903        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
20904        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
20905        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
20906        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
20907        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
20908        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
20909        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
20910        // fall to the exact scalar there instead of the broken register arm.
20911        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20912        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
20913        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
20914        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
20915        if g && head_dim == 256 && !fa_v4_at(t_kv) {
20916            fa_vec = false;
20917        }
20918        let sp = fa_split_keys(t_kv, n_head_kv);
20919        let n_splits = if fa_vec {
20920            ((t_kv + sp - 1) / sp).max(1)
20921        } else {
20922            ((t_kv + 255) / 256).max(1)
20923        };
20924        let o_len = n_head * n_splits * head_dim;
20925        let ml_len = n_head * n_splits;
20926        let mut part_guard = self.fa_part_pool.lock().unwrap();
20927        if part_guard
20928            .as_ref()
20929            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20930            .unwrap_or(true)
20931        {
20932            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
20933            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
20934            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
20935            // later live allocations land at those addresses, and the next graph REPLAY writes
20936            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
20937            // output corruption began the burst after the trunk's t_kv growth first realloc'd
20938            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
20939            // the baked addresses alive (single-stream: eager writes the new buffers, replays
20940            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
20941            // (total retired < final size).
20942            let old = part_guard.take();
20943            let (co, cm) = old
20944                .as_ref()
20945                .map(|pp| (pp.0.len(), pp.1.len()))
20946                .unwrap_or((0, 0));
20947            if let Some(old) = old {
20948                self.fa_part_retired.lock().unwrap().push(old);
20949            }
20950            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20951                eprintln!(
20952                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20953                    co, o_len, cm, ml_len
20954                );
20955            }
20956            *part_guard = Some((
20957                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20958                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20959                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20960            ));
20961        }
20962        let pg = part_guard.as_mut().unwrap();
20963        self.gpu
20964            .stream()
20965            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20966        self.gpu
20967            .stream()
20968            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20969        self.gpu
20970            .stream()
20971            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20972        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20973        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20974        let (hd, nh, nhkv, tkvi, nsp) = (
20975            head_dim as i32,
20976            n_head as i32,
20977            n_head_kv as i32,
20978            t_kv as i32,
20979            n_splits as i32,
20980        );
20981        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20982        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
20983        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
20984        // silently truncating the accumulator.
20985        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20986        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
20987        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
20988        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
20989        // 178.4 -> 173.7 when 512 rode vec unconditionally).
20990        let fa512_min = fa512_min_tkv();
20991        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
20992        // g-module keeps the v4 pick (its class is not the depth-decay class).
20993        let deep = fa_vec
20994            && head_dim == 256
20995            && fa_v4_at(t_kv)
20996            && !g
20997            && fa_deep_at(t_kv)
20998            && !matches!(fa_v4_mode(), "noB3" | "stage");
20999        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21000            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
21001            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
21002            let gqa = (n_head / n_head_kv).max(1) as u32;
21003            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21004            (
21005                fv,
21006                LaunchConfig {
21007                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21008                    block_dim: (32, gqa, 1),
21009                    shared_mem_bytes: 0,
21010                },
21011            )
21012        } else if fa_vec && head_dim <= 256 {
21013            let gqa = (n_head / n_head_kv).max(1) as u32;
21014            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
21015            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
21016            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
21017            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
21018            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
21019            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
21020            // dequant each tile ONCE per block.
21021            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
21022            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
21023            // there by 12x — latency, not bandwidth, rules small KV).
21024            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21025            let smem_tkv = *SMEM_TKV.get_or_init(|| {
21026                std::env::var("MEMRA_FA_SMEM_TKV")
21027                    .ok()
21028                    .and_then(|v| v.parse().ok())
21029                    .unwrap_or_else(|| {
21030                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21031                    })
21032            });
21033            if fa_v4_at(t_kv) && head_dim == 256 {
21034                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
21035                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
21036                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
21037                let v4name = match fa_v4_mode() {
21038                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
21039                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
21040                    _ if deep => "fa_decode_vec_q_v4_deep",
21041                    _ => "fa_decode_vec_q_v4",
21042                };
21043                let fv = if g {
21044                    self.func_g(v4name)
21045                } else {
21046                    self.func(v4name)
21047                };
21048                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
21049                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
21050                let shmem = (if deep { 12160 } else { 11520 }
21051                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
21052                use cudarc::driver::sys::CUfunction_attribute_enum as A;
21053                fv.set_attribute(
21054                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21055                    shmem as i32,
21056                )?;
21057                (
21058                    fv,
21059                    LaunchConfig {
21060                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21061                        block_dim: (32, gqa, 1),
21062                        shared_mem_bytes: shmem,
21063                    },
21064                )
21065            } else if fa_v3_active(head_dim) {
21066                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
21067                // smem = sV only (half of v2's).
21068                let fv = if g {
21069                    self.func_g("fa_decode_vec_q_v3")
21070                } else {
21071                    self.func("fa_decode_vec_q_v3")
21072                };
21073                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
21074                (
21075                    fv,
21076                    LaunchConfig {
21077                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21078                        block_dim: (32, gqa, 1),
21079                        shared_mem_bytes: shmem,
21080                    },
21081                )
21082            } else if fa_v2_on() {
21083                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
21084                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
21085                // partials; same 32KB sK+sV tile as the smem twin.
21086                let fv = if g {
21087                    self.func_g("fa_decode_vec_q_v2")
21088                } else {
21089                    self.func("fa_decode_vec_q_v2")
21090                };
21091                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
21092                (
21093                    fv,
21094                    LaunchConfig {
21095                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21096                        block_dim: (32, gqa, 1),
21097                        shared_mem_bytes: shmem,
21098                    },
21099                )
21100            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
21101            {
21102                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
21103                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
21104                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
21105                let fv = if g {
21106                    self.func_g("fa_decode_vec_q_smem")
21107                } else {
21108                    self.func("fa_decode_vec_q_smem")
21109                };
21110                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
21111                use cudarc::driver::sys::CUfunction_attribute_enum as A;
21112                fv.set_attribute(
21113                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21114                    shmem as i32,
21115                )?;
21116                (
21117                    fv,
21118                    LaunchConfig {
21119                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21120                        block_dim: (32, gqa, 1),
21121                        shared_mem_bytes: shmem,
21122                    },
21123                )
21124            } else {
21125                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
21126                // dequant, zero dynamic shared memory.
21127                let fv = if g {
21128                    self.func_g("fa_decode_vec_q")
21129                } else {
21130                    self.func("fa_decode_vec_q")
21131                };
21132                (
21133                    fv,
21134                    LaunchConfig {
21135                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21136                        block_dim: (32, gqa, 1),
21137                        shared_mem_bytes: 0,
21138                    },
21139                )
21140            }
21141        } else {
21142            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
21143            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
21144            return self.fa_decode_scalar_unified(
21145                q,
21146                k,
21147                v,
21148                o,
21149                head_dim,
21150                n_head,
21151                n_head_kv,
21152                t_kv,
21153                None,
21154                scale,
21155                n_splits,
21156                if fa_vec { sp } else { 256 },
21157                k_tok_bytes,
21158                v_tok_bytes,
21159                g,
21160                part_o,
21161                part_m,
21162                part_l,
21163                None,
21164            );
21165        };
21166        let __s_b = self.gpu.stream();
21167        let mut b = __s_b.launch_builder(&f);
21168        b.arg(q)
21169            .arg(k)
21170            .arg(v)
21171            .arg(&mut *part_o)
21172            .arg(&mut *part_m)
21173            .arg(&mut *part_l)
21174            .arg(&hd)
21175            .arg(&nh)
21176            .arg(&nhkv)
21177            .arg(&tkvi)
21178            .arg(&scale)
21179            .arg(&nsp)
21180            .arg(&ktb)
21181            .arg(&vtb);
21182        unsafe {
21183            b.launch(cfg)?;
21184        }
21185        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
21186        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
21187        let (fc, cfg2) = (
21188            if g {
21189                self.func_g("fa_decode_combine_f32")
21190            } else {
21191                self.fa_func("fa_decode_combine_f32", head_dim)
21192            },
21193            LaunchConfig {
21194                grid_dim: (n_head as u32, 1, 1),
21195                block_dim: (head_dim as u32, 1, 1),
21196                shared_mem_bytes: 0,
21197            },
21198        );
21199        let __s_b2 = self.gpu.stream();
21200        let mut b2 = __s_b2.launch_builder(&fc);
21201        b2.arg(&*part_o)
21202            .arg(&*part_m)
21203            .arg(&*part_l)
21204            .arg(o)
21205            .arg(&hd)
21206            .arg(&nh)
21207            .arg(&nsp);
21208        unsafe {
21209            b2.launch(cfg2)?;
21210        }
21211        Ok(())
21212    }
21213
21214    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
21215    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
21216    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
21217    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
21218    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
21219    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
21220    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
21221    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
21222    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
21223    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
21224    #[allow(clippy::too_many_arguments)]
21225    pub fn fa_decode_batch_seqs_v4(
21226        &self,
21227        q: &CudaSlice<f32>,
21228        kv_ptrs: &cudarc::driver::CudaView<u64>,
21229        pos_seq: &CudaSlice<i32>,
21230        o: &mut CudaSlice<f32>,
21231        head_dim: usize,
21232        n_head: usize,
21233        n_head_kv: usize,
21234        b_n: usize,
21235        t_kv_max: usize,
21236        scale: f32,
21237        split_keys: usize,
21238        k_tok_bytes: usize,
21239        v_tok_bytes: usize,
21240    ) -> Result<(), Box<dyn std::error::Error>> {
21241        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21242        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21243        let o_len = b_n * n_head * n_splits_max * head_dim;
21244        let ml_len = b_n * n_head * n_splits_max;
21245        let mut part_guard = self.fa_part_pool.lock().unwrap();
21246        if part_guard
21247            .as_ref()
21248            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21249            .unwrap_or(true)
21250        {
21251            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21252            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21253            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21254            // later live allocations land at those addresses, and the next graph REPLAY writes
21255            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21256            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21257            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21258            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21259            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21260            // (total retired < final size).
21261            let old = part_guard.take();
21262            let (co, cm) = old
21263                .as_ref()
21264                .map(|pp| (pp.0.len(), pp.1.len()))
21265                .unwrap_or((0, 0));
21266            if let Some(old) = old {
21267                self.fa_part_retired.lock().unwrap().push(old);
21268            }
21269            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21270                eprintln!(
21271                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21272                    co, o_len, cm, ml_len
21273                );
21274            }
21275            *part_guard = Some((
21276                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21277                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21278                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21279            ));
21280        }
21281        let pg = part_guard.as_mut().unwrap();
21282        self.gpu
21283            .stream()
21284            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21285        self.gpu
21286            .stream()
21287            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21288        self.gpu
21289            .stream()
21290            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21291        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21292        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21293        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21294        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21295        let gqa = (n_head / n_head_kv).max(1) as u32;
21296        let f = self.func("fa_decode_vec_q_seqs_v4");
21297        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
21298        let shmem = (11520 + 32 * head_dim * 2) as u32;
21299        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21300        f.set_attribute(
21301            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21302            shmem as i32,
21303        )?;
21304        let cfg = LaunchConfig {
21305            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21306            block_dim: (32, gqa, 1),
21307            shared_mem_bytes: shmem,
21308        };
21309        {
21310            let __s_b = self.gpu.stream();
21311            let mut b = __s_b.launch_builder(&f);
21312            b.arg(q)
21313                .arg(kv_ptrs)
21314                .arg(pos_seq)
21315                .arg(&mut *part_o)
21316                .arg(&mut *part_m)
21317                .arg(&mut *part_l)
21318                .arg(&hd)
21319                .arg(&nh)
21320                .arg(&nhkv)
21321                .arg(&scale)
21322                .arg(&nspm)
21323                .arg(&spk)
21324                .arg(&ktb)
21325                .arg(&vtb);
21326            unsafe {
21327                b.launch(cfg)?;
21328            }
21329        }
21330        let fc = self.func("fa_decode_combine_seqs");
21331        let cfg2 = LaunchConfig {
21332            grid_dim: (n_head as u32, b_n as u32, 1),
21333            block_dim: (head_dim as u32, 1, 1),
21334            shared_mem_bytes: 0,
21335        };
21336        let __s_b2 = self.gpu.stream();
21337        let mut b2 = __s_b2.launch_builder(&fc);
21338        b2.arg(&*part_o)
21339            .arg(&*part_m)
21340            .arg(&*part_l)
21341            .arg(o)
21342            .arg(&hd)
21343            .arg(&nh)
21344            .arg(pos_seq)
21345            .arg(&nspm)
21346            .arg(&spk);
21347        unsafe {
21348            b2.launch(cfg2)?;
21349        }
21350        Ok(())
21351    }
21352
21353    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
21354    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
21355    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
21356    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
21357    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
21358    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
21359    #[allow(clippy::too_many_arguments)]
21360    pub fn append_kv_quantized_seqs(
21361        &self,
21362        k_rows: &CudaSlice<f32>,
21363        v_rows: &CudaSlice<f32>,
21364        kv_ptrs: &cudarc::driver::CudaView<u64>,
21365        pos_seq: &CudaSlice<i32>,
21366        b_n: usize,
21367        kv_dim_k: usize,
21368        kv_dim_v: usize,
21369        k_tok_bytes: usize,
21370        v_tok_bytes: usize,
21371    ) -> Result<(), Box<dyn std::error::Error>> {
21372        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21373        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21374        let cfg = LaunchConfig {
21375            grid_dim: (nblk, b_n as u32, 1),
21376            block_dim: (32, 1, 1),
21377            shared_mem_bytes: 0,
21378        };
21379        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21380        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21381        let __s_b = self.gpu.stream();
21382        let mut b = __s_b.launch_builder(&f);
21383        b.arg(k_rows)
21384            .arg(v_rows)
21385            .arg(kv_ptrs)
21386            .arg(pos_seq)
21387            .arg(&kdk)
21388            .arg(&kdv)
21389            .arg(&ktb)
21390            .arg(&vtb);
21391        unsafe {
21392            b.launch(cfg)?;
21393        }
21394        Ok(())
21395    }
21396
21397    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
21398    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
21399    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
21400    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
21401    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
21402    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21403        std::env::var("MEMRA_NO_FA_VEC").is_err()
21404            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21405            && base_len + 1 >= fa_vec_min_tkv()
21406            && head_dim <= 256
21407            && head_dim % 32 == 0
21408    }
21409
21410    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
21411    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
21412    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
21413    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
21414    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
21415    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
21416    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
21417    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
21418    #[allow(clippy::too_many_arguments)]
21419    pub fn fa_decode_rows(
21420        &self,
21421        q: &CudaSlice<f32>,
21422        k: &cudarc::driver::CudaView<u8>,
21423        v: &cudarc::driver::CudaView<u8>,
21424        o: &mut CudaSlice<f32>,
21425        head_dim: usize,
21426        n_head: usize,
21427        n_head_kv: usize,
21428        base_len: usize,
21429        t: usize,
21430        scale: f32,
21431        k_tok_bytes: usize,
21432        v_tok_bytes: usize,
21433        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
21434        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
21435        // keep the host arg. None is a bug for hd512 (asserted below).
21436        base_dev: Option<(&CudaSlice<i32>, i32)>,
21437        // K and V planes hold the same values (gemma globals, wv:=wk): pick
21438        // the _kv twin — V plane never read, value rides the q8_0 key dq.
21439        kv_shared: bool,
21440        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
21441        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
21442        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
21443        g: bool,
21444        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
21445        // (hd512 path) — the standalone quantize launch folds away.
21446        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21447    ) -> Result<(), Box<dyn std::error::Error>> {
21448        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21449        let t_kv_max = base_len + t; // LAST row's key bound
21450        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
21451        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
21452        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
21453        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
21454        // (parity law), so the partition is freely tunable — verify and decode move together.
21455        if head_dim == 512 {
21456            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21457            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
21458            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
21459            let v = *SP512.get_or_init(|| {
21460                std::env::var("MEMRA_FA_SP512")
21461                    .ok()
21462                    .and_then(|x| x.parse().ok())
21463                    .unwrap_or(0)
21464            });
21465            sp = if v >= 8 {
21466                v
21467            } else {
21468                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21469            };
21470        }
21471        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21472        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21473        let gqa = (n_head / n_head_kv).max(1) as u32;
21474        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
21475        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
21476        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
21477        // the different partition changes the combine's FP order (greedy tie flips at depth;
21478        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
21479        // consecutive rows by their OWN ladder value and launch once per group — each row then
21480        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
21481        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
21482        // sp override is t_kv-independent by construction).
21483        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
21484        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21485            groups.push((0, t, sp));
21486        } else {
21487            let mut r0 = 0usize;
21488            while r0 < t {
21489                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21490                let mut r1 = r0 + 1;
21491                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21492                    r1 += 1;
21493                }
21494                groups.push((r0, r1 - r0, sp_g));
21495                r0 = r1;
21496            }
21497        }
21498        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
21499        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
21500        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
21501        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21502        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21503            std::env::var("MEMRA_FA_SMEM_TKV")
21504                .ok()
21505                .and_then(|v| v.parse().ok())
21506                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21507        });
21508        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21509        let v3 = fa_v3_active(head_dim);
21510        let smem_rows =
21511            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21512        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
21513        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
21514        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
21515        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
21516        let _ = kv_shared;
21517        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
21518        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
21519        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21520        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
21521        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
21522        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
21523        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
21524        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
21525        // (kv_head, split) stages its tile once and loops the rows over it — kills the
21526        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
21527        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
21528        // shared by every hd512 caller through this wrapper (decode+verify flip together;
21529        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
21530        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
21531        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
21532        // not unpack-bound; jsonl 2026-07-14.
21533        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21534        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
21535        let tb512 = head_dim == 512
21536            && sp <= 32
21537            && n_head / n_head_kv.max(1) <= 16
21538            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21539        let fname = if tb512 {
21540            "fa_decode_vec_q_rows_v4_512_tb"
21541        } else if i2 {
21542            "fa_decode_vec_q_rows_dpl16_i2"
21543        } else if head_dim == 512 {
21544            "fa_decode_vec_q_rows_dpl16"
21545        }
21546        // gemma globals (parity law)
21547        else if v4 {
21548            "fa_decode_vec_q_rows_v4"
21549        } else if v3 {
21550            "fa_decode_vec_q_rows_v3"
21551        } else if fa_v2_on() {
21552            "fa_decode_vec_q_rows_v2"
21553        } else if smem_rows {
21554            "fa_decode_vec_q_rows_smem"
21555        } else {
21556            "fa_decode_vec_q_rows"
21557        };
21558        let f = if head_dim == 512 {
21559            self.fa_func(fname, head_dim)
21560        } else if g {
21561            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
21562            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
21563            // g-module rows against decode's g-module v4 — different programs, short-VG
21564            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
21565            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
21566            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
21567            // dq macros are format-aware.
21568            self.func_g(if smem_rows {
21569                "fa_decode_vec_q_rows"
21570            } else {
21571                fname
21572            })
21573        } else {
21574            self.func(fname)
21575        };
21576        let shmem = if tb512 {
21577            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
21578            let gk = Self::gkv_on();
21579            let sh =
21580                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21581            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21582            f.set_attribute(
21583                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21584                sh as i32,
21585            )?;
21586            sh
21587        } else if v4 || v3 || smem_rows || fa_v2_on() {
21588            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
21589            let sh = (if v4 {
21590                11520 + 32 * head_dim * if g { 1 } else { 2 }
21591            } else if v3 {
21592                32 * head_dim * 2
21593            } else {
21594                2 * 32 * head_dim * 2
21595            }) as u32;
21596            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21597            f.set_attribute(
21598                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21599                sh as i32,
21600            )?;
21601            sh
21602        } else {
21603            0
21604        };
21605        // Per-GROUP launches (single group in the common case — identical to the pre-fix
21606        // single launch there): each group gets its own partials (the rows kernel indexes
21607        // partials by its LOCAL grid.z row) and q/o row-offset views.
21608        for &(r0, t_g, sp_g) in &groups {
21609            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21610            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21611            let base_i = (base_len + r0) as i32;
21612            let o_len = t_g * n_head * n_splits_g * head_dim;
21613            let ml_len = t_g * n_head * n_splits_g;
21614            let mut part_guard = self.fa_part_pool.lock().unwrap();
21615            if part_guard
21616                .as_ref()
21617                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21618                .unwrap_or(true)
21619            {
21620                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21621                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21622                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21623                // later live allocations land at those addresses, and the next graph REPLAY writes
21624                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21625                // output corruption began the burst after the trunk's t_kv growth first realloc'd
21626                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21627                // the baked addresses alive (single-stream: eager writes the new buffers, replays
21628                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21629                // (total retired < final size).
21630                let old = part_guard.take();
21631                let (co, cm) = old
21632                    .as_ref()
21633                    .map(|pp| (pp.0.len(), pp.1.len()))
21634                    .unwrap_or((0, 0));
21635                if let Some(old) = old {
21636                    self.fa_part_retired.lock().unwrap().push(old);
21637                }
21638                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21639                    eprintln!(
21640                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21641                        co, o_len, cm, ml_len
21642                    );
21643                }
21644                *part_guard = Some((
21645                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21646                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21647                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21648                ));
21649            }
21650            let pg = part_guard.as_mut().unwrap();
21651            self.gpu
21652                .stream()
21653                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21654            self.gpu
21655                .stream()
21656                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21657            self.gpu
21658                .stream()
21659                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21660            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21661            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21662            let qv = self.view(q, t * n_head * head_dim);
21663            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21664            let cfg = LaunchConfig {
21665                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21666                block_dim: (32, gqa, 1),
21667                shared_mem_bytes: shmem,
21668            };
21669            {
21670                let __s_b = self.gpu.stream();
21671                let mut b = __s_b.launch_builder(&f);
21672                if tb512 {
21673                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
21674                    let (bd, plus) =
21675                        base_dev.expect("hd512 rows twin requires a device base counter");
21676                    let plus_g = plus + r0 as i32;
21677                    let nr = t_g as i32;
21678                    if Self::pdl_on() && Self::pdl_wb_on() {
21679                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
21680                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21681                        let s = &self.gpu.stream();
21682                        let (pq, _b0) = q_g.device_ptr(s);
21683                        let (pk, _b1) = k.device_ptr(s);
21684                        let (pv, _b2) = v.device_ptr(s);
21685                        let (po, _b3) = part_o.device_ptr_mut(s);
21686                        let (pm, _b4) = part_m.device_ptr_mut(s);
21687                        let (pl, _b5) = part_l.device_ptr_mut(s);
21688                        let (pb, _b6) = bd.device_ptr(s);
21689                        let mut ps = [
21690                            &pq as *const _ as *mut std::ffi::c_void,
21691                            &pk as *const _ as *mut _,
21692                            &pv as *const _ as *mut _,
21693                            &po as *const _ as *mut _,
21694                            &pm as *const _ as *mut _,
21695                            &pl as *const _ as *mut _,
21696                            &hd as *const _ as *mut _,
21697                            &nh as *const _ as *mut _,
21698                            &nhkv as *const _ as *mut _,
21699                            &pb as *const _ as *mut _,
21700                            &plus_g as *const _ as *mut _,
21701                            &scale as *const _ as *mut _,
21702                            &nspm as *const _ as *mut _,
21703                            &spk as *const _ as *mut _,
21704                            &ktb as *const _ as *mut _,
21705                            &vtb as *const _ as *mut _,
21706                            &nr as *const _ as *mut _,
21707                        ];
21708                        unsafe {
21709                            self.launch_pdl_flash(
21710                                Self::gkv_on(),
21711                                "fa_decode_vec_q_rows_v4_512_tb",
21712                                (n_head_kv as u32, n_splits_g as u32, 1),
21713                                (32, gqa, 1),
21714                                shmem,
21715                                &mut ps,
21716                            )?;
21717                        }
21718                    } else {
21719                        let cfg_tb = LaunchConfig {
21720                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21721                            block_dim: (32, gqa, 1),
21722                            shared_mem_bytes: shmem,
21723                        };
21724                        b.arg(&q_g)
21725                            .arg(k)
21726                            .arg(v)
21727                            .arg(&mut *part_o)
21728                            .arg(&mut *part_m)
21729                            .arg(&mut *part_l)
21730                            .arg(&hd)
21731                            .arg(&nh)
21732                            .arg(&nhkv)
21733                            .arg(bd)
21734                            .arg(&plus_g)
21735                            .arg(&scale)
21736                            .arg(&nspm)
21737                            .arg(&spk)
21738                            .arg(&ktb)
21739                            .arg(&vtb)
21740                            .arg(&nr);
21741                        unsafe {
21742                            b.launch(cfg_tb)?;
21743                        }
21744                    }
21745                } else if head_dim == 512 {
21746                    let (bd, plus) =
21747                        base_dev.expect("hd512 rows twin requires a device base counter");
21748                    let plus_g = plus + r0 as i32;
21749                    b.arg(&q_g)
21750                        .arg(k)
21751                        .arg(v)
21752                        .arg(&mut *part_o)
21753                        .arg(&mut *part_m)
21754                        .arg(&mut *part_l)
21755                        .arg(&hd)
21756                        .arg(&nh)
21757                        .arg(&nhkv)
21758                        .arg(bd)
21759                        .arg(&plus_g)
21760                        .arg(&scale)
21761                        .arg(&nspm)
21762                        .arg(&spk)
21763                        .arg(&ktb)
21764                        .arg(&vtb);
21765                    unsafe {
21766                        b.launch(cfg)?;
21767                    }
21768                } else {
21769                    b.arg(&q_g)
21770                        .arg(k)
21771                        .arg(v)
21772                        .arg(&mut *part_o)
21773                        .arg(&mut *part_m)
21774                        .arg(&mut *part_l)
21775                        .arg(&hd)
21776                        .arg(&nh)
21777                        .arg(&nhkv)
21778                        .arg(&base_i)
21779                        .arg(&scale)
21780                        .arg(&nspm)
21781                        .arg(&spk)
21782                        .arg(&ktb)
21783                        .arg(&vtb);
21784                    unsafe {
21785                        b.launch(cfg)?;
21786                    }
21787                }
21788            }
21789            let cfg2 = LaunchConfig {
21790                grid_dim: (n_head as u32, t_g as u32, 1),
21791                block_dim: (head_dim as u32, 1, 1),
21792                shared_mem_bytes: 0,
21793            };
21794            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21795            if head_dim == 512 {
21796                // device-len combine (shared by verify/eager/graph — parity by symbol): the
21797                // per-row n_splits derives from the SAME counter the rows kernel read.
21798                let (bd, plus) = base_dev.unwrap();
21799                let plus_g = plus + r0 as i32;
21800                if let Some((oq, od)) = q8_out.as_mut() {
21801                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
21802                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21803                    if Self::pdl_on() && Self::pdl_wb_on() {
21804                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
21805                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21806                        let s = &self.gpu.stream();
21807                        let (po, _g0) = part_o.device_ptr(s);
21808                        let (pm, _g1) = part_m.device_ptr(s);
21809                        let (pl, _g2) = part_l.device_ptr(s);
21810                        let (pq, _g3) = oq.device_ptr_mut(s);
21811                        let (pd, _g4) = od.device_ptr_mut(s);
21812                        let (pb, _g5) = bd.device_ptr(s);
21813                        let mut ps = [
21814                            &po as *const _ as *mut std::ffi::c_void,
21815                            &pm as *const _ as *mut _,
21816                            &pl as *const _ as *mut _,
21817                            &pq as *const _ as *mut _,
21818                            &pd as *const _ as *mut _,
21819                            &hd as *const _ as *mut _,
21820                            &nh as *const _ as *mut _,
21821                            &pb as *const _ as *mut _,
21822                            &plus_g as *const _ as *mut _,
21823                            &nspm as *const _ as *mut _,
21824                            &spk as *const _ as *mut _,
21825                        ];
21826                        unsafe {
21827                            self.launch_pdl_flash(
21828                                Self::gkv_on(),
21829                                "fa_decode_combine_rows_dc_q8_1",
21830                                cfg2.grid_dim,
21831                                cfg2.block_dim,
21832                                0,
21833                                &mut ps,
21834                            )?;
21835                        }
21836                        continue;
21837                    }
21838                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21839                    let __s_b2 = self.gpu.stream();
21840                    let mut b2 = __s_b2.launch_builder(&fc);
21841                    b2.arg(&*part_o)
21842                        .arg(&*part_m)
21843                        .arg(&*part_l)
21844                        .arg(&mut **oq)
21845                        .arg(&mut **od)
21846                        .arg(&hd)
21847                        .arg(&nh)
21848                        .arg(bd)
21849                        .arg(&plus_g)
21850                        .arg(&nspm)
21851                        .arg(&spk);
21852                    unsafe {
21853                        b2.launch(cfg2)?;
21854                    }
21855                    continue;
21856                }
21857                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21858                let __s_b2 = self.gpu.stream();
21859                let mut b2 = __s_b2.launch_builder(&fc);
21860                b2.arg(&*part_o)
21861                    .arg(&*part_m)
21862                    .arg(&*part_l)
21863                    .arg(&mut o_g)
21864                    .arg(&hd)
21865                    .arg(&nh)
21866                    .arg(bd)
21867                    .arg(&plus_g)
21868                    .arg(&nspm)
21869                    .arg(&spk);
21870                unsafe {
21871                    b2.launch(cfg2)?;
21872                }
21873            } else {
21874                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
21875                // leave the caller's pair unwritten (consumer would read garbage).
21876                assert!(
21877                    q8_out.is_none(),
21878                    "rows q8 emit requires the hd512 dc combine"
21879                );
21880                let fc = self.func("fa_decode_combine_rows");
21881                let __s_b2 = self.gpu.stream();
21882                let mut b2 = __s_b2.launch_builder(&fc);
21883                b2.arg(&*part_o)
21884                    .arg(&*part_m)
21885                    .arg(&*part_l)
21886                    .arg(&mut o_g)
21887                    .arg(&hd)
21888                    .arg(&nh)
21889                    .arg(&base_i)
21890                    .arg(&nspm)
21891                    .arg(&spk);
21892                unsafe {
21893                    b2.launch(cfg2)?;
21894                }
21895            }
21896        }
21897        Ok(())
21898    }
21899
21900    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
21901    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
21902    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
21903    #[allow(clippy::too_many_arguments)]
21904    pub fn fa_decode_rows_w(
21905        &self,
21906        q: &CudaSlice<f32>,
21907        k: &cudarc::driver::CudaView<u8>,
21908        v: &cudarc::driver::CudaView<u8>,
21909        o: &mut CudaSlice<f32>,
21910        head_dim: usize,
21911        n_head: usize,
21912        n_head_kv: usize,
21913        base_dev: &CudaSlice<i32>,
21914        base_plus: i32,
21915        t: usize,
21916        scale: f32,
21917        window: usize,
21918        k_tok_bytes: usize,
21919        v_tok_bytes: usize,
21920        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21921    ) -> Result<(), Box<dyn std::error::Error>> {
21922        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
21923        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
21924        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
21925        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
21926        debug_assert!(head_dim == 256);
21927        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
21928        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
21929        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
21930        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
21931        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
21932        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
21933        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
21934        let sp = {
21935            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21936            let v = *SPW.get_or_init(|| {
21937                std::env::var("MEMRA_FA_SPW")
21938                    .ok()
21939                    .and_then(|x| x.parse().ok())
21940                    .unwrap_or(0)
21941            });
21942            if v >= 8 {
21943                v
21944            } else {
21945                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21946            }
21947        };
21948        let n_splits_max = (window + sp - 1) / sp;
21949        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21950        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21951        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21952        let gqa = (n_head / n_head_kv).max(1) as u32;
21953        let o_len = t * n_head * n_splits_max * head_dim;
21954        let ml_len = t * n_head * n_splits_max;
21955        let mut part_guard = self.fa_part_pool.lock().unwrap();
21956        if part_guard
21957            .as_ref()
21958            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21959            .unwrap_or(true)
21960        {
21961            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21962            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21963            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21964            // later live allocations land at those addresses, and the next graph REPLAY writes
21965            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21966            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21967            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21968            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21969            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21970            // (total retired < final size).
21971            let old = part_guard.take();
21972            let (co, cm) = old
21973                .as_ref()
21974                .map(|pp| (pp.0.len(), pp.1.len()))
21975                .unwrap_or((0, 0));
21976            if let Some(old) = old {
21977                self.fa_part_retired.lock().unwrap().push(old);
21978            }
21979            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21980                eprintln!(
21981                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21982                    co, o_len, cm, ml_len
21983                );
21984            }
21985            *part_guard = Some((
21986                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21987                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21988                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21989            ));
21990        }
21991        let pg = part_guard.as_mut().unwrap();
21992        self.gpu
21993            .stream()
21994            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21995        self.gpu
21996            .stream()
21997            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21998        self.gpu
21999            .stream()
22000            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22001        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22002        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
22003        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
22004        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
22005        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
22006        // floor (deep-ctx broadcast win); register twin between.
22007        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22008        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22009            std::env::var("MEMRA_FA_SMEM_TKV")
22010                .ok()
22011                .and_then(|v| v.parse().ok())
22012                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22013        });
22014        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
22015        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
22016        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
22017        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
22018        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
22019        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22020        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
22021        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
22022        // per (lane, format-module) keeps parity structural; the old register-i2 detour
22023        // (-33%) is retired.
22024        let wg = Self::wkv_on();
22025        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
22026        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
22027        let sp2 =
22028            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
22029        if sp2 {
22030            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22031            if Self::pdl_on() && Self::pdl_wb_on() {
22032                // wave-B2b: flavor mirrors wg.
22033                use cudarc::driver::{DevicePtr, DevicePtrMut};
22034                let s = &self.gpu.stream();
22035                let (pq, _b0) = q.device_ptr(s);
22036                let (pk, _b1) = k.device_ptr(s);
22037                let (pv, _b2) = v.device_ptr(s);
22038                let (po, _b3) = part_o.device_ptr_mut(s);
22039                let (pm, _b4) = part_m.device_ptr_mut(s);
22040                let (pl, _b5) = part_l.device_ptr_mut(s);
22041                let (pb, _b6) = base_dev.device_ptr(s);
22042                let mut ps = [
22043                    &pq as *const _ as *mut std::ffi::c_void,
22044                    &pk as *const _ as *mut _,
22045                    &pv as *const _ as *mut _,
22046                    &po as *const _ as *mut _,
22047                    &pm as *const _ as *mut _,
22048                    &pl as *const _ as *mut _,
22049                    &hd as *const _ as *mut _,
22050                    &nh as *const _ as *mut _,
22051                    &nhkv as *const _ as *mut _,
22052                    &pb as *const _ as *mut _,
22053                    &base_plus as *const _ as *mut _,
22054                    &scale as *const _ as *mut _,
22055                    &nspm as *const _ as *mut _,
22056                    &spk as *const _ as *mut _,
22057                    &ktb as *const _ as *mut _,
22058                    &vtb as *const _ as *mut _,
22059                    &wini as *const _ as *mut _,
22060                ];
22061                unsafe {
22062                    self.launch_pdl_flash(
22063                        wg,
22064                        "fa_decode_vec_q_rows_v4_w_sp",
22065                        (n_head_kv as u32, n_splits_max as u32, t as u32),
22066                        (32, gqa + 1, 1),
22067                        sh,
22068                        &mut ps,
22069                    )?;
22070                }
22071            } else {
22072                let f = if wg {
22073                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
22074                } else {
22075                    self.func("fa_decode_vec_q_rows_v4_w_sp")
22076                };
22077                f.set_attribute(
22078                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22079                    sh as i32,
22080                )?;
22081                let cfg = LaunchConfig {
22082                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22083                    block_dim: (32, gqa + 1, 1),
22084                    shared_mem_bytes: sh,
22085                };
22086                let __s_b = self.gpu.stream();
22087                let mut b = __s_b.launch_builder(&f);
22088                b.arg(q)
22089                    .arg(k)
22090                    .arg(v)
22091                    .arg(&mut *part_o)
22092                    .arg(&mut *part_m)
22093                    .arg(&mut *part_l)
22094                    .arg(&hd)
22095                    .arg(&nh)
22096                    .arg(&nhkv)
22097                    .arg(base_dev)
22098                    .arg(&base_plus)
22099                    .arg(&scale)
22100                    .arg(&nspm)
22101                    .arg(&spk)
22102                    .arg(&ktb)
22103                    .arg(&vtb)
22104                    .arg(&wini);
22105                unsafe {
22106                    b.launch(cfg)?;
22107                }
22108            }
22109        } else {
22110            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
22111                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
22112                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22113                use cudarc::driver::{DevicePtr, DevicePtrMut};
22114                let s = &self.gpu.stream();
22115                let (pq, _b0) = q.device_ptr(s);
22116                let (pk, _b1) = k.device_ptr(s);
22117                let (pv, _b2) = v.device_ptr(s);
22118                let (po, _b3) = part_o.device_ptr_mut(s);
22119                let (pm, _b4) = part_m.device_ptr_mut(s);
22120                let (pl, _b5) = part_l.device_ptr_mut(s);
22121                let (pb, _b6) = base_dev.device_ptr(s);
22122                let mut ps = [
22123                    &pq as *const _ as *mut std::ffi::c_void,
22124                    &pk as *const _ as *mut _,
22125                    &pv as *const _ as *mut _,
22126                    &po as *const _ as *mut _,
22127                    &pm as *const _ as *mut _,
22128                    &pl as *const _ as *mut _,
22129                    &hd as *const _ as *mut _,
22130                    &nh as *const _ as *mut _,
22131                    &nhkv as *const _ as *mut _,
22132                    &pb as *const _ as *mut _,
22133                    &base_plus as *const _ as *mut _,
22134                    &scale as *const _ as *mut _,
22135                    &nspm as *const _ as *mut _,
22136                    &spk as *const _ as *mut _,
22137                    &ktb as *const _ as *mut _,
22138                    &vtb as *const _ as *mut _,
22139                    &wini as *const _ as *mut _,
22140                ];
22141                unsafe {
22142                    self.launch_pdl_flash(
22143                        wg,
22144                        "fa_decode_vec_q_rows_v4_w",
22145                        (n_head_kv as u32, n_splits_max as u32, t as u32),
22146                        (32, gqa, 1),
22147                        sh,
22148                        &mut ps,
22149                    )?;
22150                }
22151            } else {
22152                let pick = |name: &str| {
22153                    if wg {
22154                        self.func_g(name)
22155                    } else {
22156                        self.func(name)
22157                    }
22158                };
22159                let (f, sh) = if fa_v4_at(window) {
22160                    let f = pick("fa_decode_vec_q_rows_v4_w");
22161                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22162                } else if smem_tkv > 0 && window >= smem_tkv {
22163                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
22164                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
22165                    (
22166                        pick("fa_decode_vec_q_rows_smem_w"),
22167                        (2 * 32 * head_dim * 2) as u32,
22168                    )
22169                } else {
22170                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22171                };
22172                f.set_attribute(
22173                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22174                    sh as i32,
22175                )?;
22176                let cfg = LaunchConfig {
22177                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22178                    block_dim: (32, gqa, 1),
22179                    shared_mem_bytes: sh,
22180                };
22181                let __s_b = self.gpu.stream();
22182                let mut b = __s_b.launch_builder(&f);
22183                b.arg(q)
22184                    .arg(k)
22185                    .arg(v)
22186                    .arg(&mut *part_o)
22187                    .arg(&mut *part_m)
22188                    .arg(&mut *part_l)
22189                    .arg(&hd)
22190                    .arg(&nh)
22191                    .arg(&nhkv)
22192                    .arg(base_dev)
22193                    .arg(&base_plus)
22194                    .arg(&scale)
22195                    .arg(&nspm)
22196                    .arg(&spk)
22197                    .arg(&ktb)
22198                    .arg(&vtb)
22199                    .arg(&wini);
22200                unsafe {
22201                    b.launch(cfg)?;
22202                }
22203            }
22204        }
22205        let cfg2 = LaunchConfig {
22206            grid_dim: (n_head as u32, t as u32, 1),
22207            block_dim: (head_dim as u32, 1, 1),
22208            shared_mem_bytes: 0,
22209        };
22210        if let Some((oq, od)) = q8_out {
22211            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
22212            // consumes the pair directly; the standalone quantize launch folds away.
22213            if Self::pdl_on() && Self::pdl_wb_on() {
22214                // wave-B2: flavor mirrors the builder's wg choice.
22215                use cudarc::driver::{DevicePtr, DevicePtrMut};
22216                let s = &self.gpu.stream();
22217                let (po, _g0) = part_o.device_ptr(s);
22218                let (pm, _g1) = part_m.device_ptr(s);
22219                let (pl, _g2) = part_l.device_ptr(s);
22220                let (pq, _g3) = oq.device_ptr_mut(s);
22221                let (pd, _g4) = od.device_ptr_mut(s);
22222                let mut ps = [
22223                    &po as *const _ as *mut std::ffi::c_void,
22224                    &pm as *const _ as *mut _,
22225                    &pl as *const _ as *mut _,
22226                    &pq as *const _ as *mut _,
22227                    &pd as *const _ as *mut _,
22228                    &hd as *const _ as *mut _,
22229                    &nh as *const _ as *mut _,
22230                    &nspm as *const _ as *mut _,
22231                    &spk as *const _ as *mut _,
22232                    &wini as *const _ as *mut _,
22233                ];
22234                unsafe {
22235                    self.launch_pdl_flash(
22236                        wg,
22237                        "fa_decode_combine_rows_w_q8_1",
22238                        cfg2.grid_dim,
22239                        cfg2.block_dim,
22240                        0,
22241                        &mut ps,
22242                    )?;
22243                }
22244                return Ok(());
22245            }
22246            let fc = if wg {
22247                self.func_g("fa_decode_combine_rows_w_q8_1")
22248            } else {
22249                self.func("fa_decode_combine_rows_w_q8_1")
22250            };
22251            let __s_b2 = self.gpu.stream();
22252            let mut b2 = __s_b2.launch_builder(&fc);
22253            b2.arg(&*part_o)
22254                .arg(&*part_m)
22255                .arg(&*part_l)
22256                .arg(oq)
22257                .arg(od)
22258                .arg(&hd)
22259                .arg(&nh)
22260                .arg(&nspm)
22261                .arg(&spk)
22262                .arg(&wini);
22263            unsafe {
22264                b2.launch(cfg2)?;
22265            }
22266            return Ok(());
22267        }
22268        let fc = if wg {
22269            self.func_g("fa_decode_combine_rows_w")
22270        } else {
22271            self.func("fa_decode_combine_rows_w")
22272        };
22273        let __s_b2 = self.gpu.stream();
22274        let mut b2 = __s_b2.launch_builder(&fc);
22275        b2.arg(&*part_o)
22276            .arg(&*part_m)
22277            .arg(&*part_l)
22278            .arg(o)
22279            .arg(&hd)
22280            .arg(&nh)
22281            .arg(&nspm)
22282            .arg(&spk)
22283            .arg(&wini);
22284        unsafe {
22285            b2.launch(cfg2)?;
22286        }
22287        Ok(())
22288    }
22289
22290    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
22291    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
22292    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
22293    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
22294    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
22295    #[allow(clippy::too_many_arguments)]
22296    pub fn fa_decode_rows_dc(
22297        &self,
22298        q: &CudaSlice<f32>,
22299        k: &cudarc::driver::CudaView<u8>,
22300        v: &cudarc::driver::CudaView<u8>,
22301        o: &mut CudaSlice<f32>,
22302        head_dim: usize,
22303        n_head: usize,
22304        n_head_kv: usize,
22305        base_dev: &CudaSlice<i32>,
22306        t_kv_upper: usize,
22307        t: usize,
22308        scale: f32,
22309        k_tok_bytes: usize,
22310        v_tok_bytes: usize,
22311        base_plus: i32,
22312        g: bool,
22313    ) -> Result<(), Box<dyn std::error::Error>> {
22314        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22315        assert!(
22316            v4 || fa_v3_active(head_dim),
22317            "stream fa rows requires the v3 or v4 lane"
22318        );
22319        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22320        if v4 {
22321            let sp = fa_split_keys(t_kv_upper, n_head_kv);
22322            let n_splits_max = (t_kv_upper + sp - 1) / sp;
22323            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22324            let (nspm, spk) = (n_splits_max as i32, sp as i32);
22325            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22326            let gqa = (n_head / n_head_kv).max(1) as u32;
22327            let o_len = t * n_head * n_splits_max * head_dim;
22328            let ml_len = t * n_head * n_splits_max;
22329            let mut part_guard = self.fa_part_pool.lock().unwrap();
22330            if part_guard
22331                .as_ref()
22332                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22333                .unwrap_or(true)
22334            {
22335                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22336                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22337                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22338                // later live allocations land at those addresses, and the next graph REPLAY writes
22339                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22340                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22341                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22342                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22343                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22344                // (total retired < final size).
22345                let old = part_guard.take();
22346                let (co, cm) = old
22347                    .as_ref()
22348                    .map(|pp| (pp.0.len(), pp.1.len()))
22349                    .unwrap_or((0, 0));
22350                if let Some(old) = old {
22351                    self.fa_part_retired.lock().unwrap().push(old);
22352                }
22353                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22354                    eprintln!(
22355                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22356                        co, o_len, cm, ml_len
22357                    );
22358                }
22359                *part_guard = Some((
22360                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22361                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22362                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22363                ));
22364            }
22365            let pg = part_guard.as_mut().unwrap();
22366            self.gpu
22367                .stream()
22368                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22369            self.gpu
22370                .stream()
22371                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22372            self.gpu
22373                .stream()
22374                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22375            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22376            let f = if g {
22377                self.func_g("fa_decode_vec_q_rows_v4_dc")
22378            } else {
22379                self.func("fa_decode_vec_q_rows_v4_dc")
22380            };
22381            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22382            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22383            f.set_attribute(
22384                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22385                sh as i32,
22386            )?;
22387            let cfg = LaunchConfig {
22388                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22389                block_dim: (32, gqa, 1),
22390                shared_mem_bytes: sh,
22391            };
22392            let __s_b = self.gpu.stream();
22393            let mut b = __s_b.launch_builder(&f);
22394            b.arg(q)
22395                .arg(k)
22396                .arg(v)
22397                .arg(&mut *part_o)
22398                .arg(&mut *part_m)
22399                .arg(&mut *part_l)
22400                .arg(&hd)
22401                .arg(&nh)
22402                .arg(&nhkv)
22403                .arg(base_dev)
22404                .arg(&base_plus)
22405                .arg(&scale)
22406                .arg(&nspm)
22407                .arg(&spk)
22408                .arg(&ktb)
22409                .arg(&vtb);
22410            unsafe {
22411                b.launch(cfg)?;
22412            }
22413            let fc = self.func("fa_decode_combine_rows_dc");
22414            let cfg2 = LaunchConfig {
22415                grid_dim: (n_head as u32, t as u32, 1),
22416                block_dim: (head_dim as u32, 1, 1),
22417                shared_mem_bytes: 0,
22418            };
22419            let __s_b2 = self.gpu.stream();
22420            let mut b2 = __s_b2.launch_builder(&fc);
22421            b2.arg(&*part_o)
22422                .arg(&*part_m)
22423                .arg(&*part_l)
22424                .arg(o)
22425                .arg(&hd)
22426                .arg(&nh)
22427                .arg(base_dev)
22428                .arg(&base_plus)
22429                .arg(&nspm)
22430                .arg(&spk);
22431            unsafe {
22432                b2.launch(cfg2)?;
22433            }
22434            return Ok(());
22435        }
22436        let sp = fa_split_keys(t_kv_upper, n_head_kv);
22437        let n_splits_max = (t_kv_upper + sp - 1) / sp;
22438        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22439        let (nspm, spk) = (n_splits_max as i32, sp as i32);
22440        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22441        let gqa = (n_head / n_head_kv).max(1) as u32;
22442        let o_len = t * n_head * n_splits_max * head_dim;
22443        let ml_len = t * n_head * n_splits_max;
22444        let mut part_guard = self.fa_part_pool.lock().unwrap();
22445        if part_guard
22446            .as_ref()
22447            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22448            .unwrap_or(true)
22449        {
22450            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22451            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22452            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22453            // later live allocations land at those addresses, and the next graph REPLAY writes
22454            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22455            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22456            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22457            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22458            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22459            // (total retired < final size).
22460            let old = part_guard.take();
22461            let (co, cm) = old
22462                .as_ref()
22463                .map(|pp| (pp.0.len(), pp.1.len()))
22464                .unwrap_or((0, 0));
22465            if let Some(old) = old {
22466                self.fa_part_retired.lock().unwrap().push(old);
22467            }
22468            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22469                eprintln!(
22470                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22471                    co, o_len, cm, ml_len
22472                );
22473            }
22474            *part_guard = Some((
22475                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22476                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22477                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22478            ));
22479        }
22480        let pg = part_guard.as_mut().unwrap();
22481        self.gpu
22482            .stream()
22483            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22484        self.gpu
22485            .stream()
22486            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22487        self.gpu
22488            .stream()
22489            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22490        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22491        let f = self.func("fa_decode_vec_q_rows_v3_dc");
22492        let sh = (32 * head_dim * 2) as u32;
22493        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22494        f.set_attribute(
22495            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22496            sh as i32,
22497        )?;
22498        let cfg = LaunchConfig {
22499            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22500            block_dim: (32, gqa, 1),
22501            shared_mem_bytes: sh,
22502        };
22503        let __s_b = self.gpu.stream();
22504        let mut b = __s_b.launch_builder(&f);
22505        b.arg(q)
22506            .arg(k)
22507            .arg(v)
22508            .arg(&mut *part_o)
22509            .arg(&mut *part_m)
22510            .arg(&mut *part_l)
22511            .arg(&hd)
22512            .arg(&nh)
22513            .arg(&nhkv)
22514            .arg(base_dev)
22515            .arg(&scale)
22516            .arg(&nspm)
22517            .arg(&spk)
22518            .arg(&ktb)
22519            .arg(&vtb);
22520        unsafe {
22521            b.launch(cfg)?;
22522        }
22523        let fc = self.func("fa_decode_combine_rows_dc");
22524        let cfg2 = LaunchConfig {
22525            grid_dim: (n_head as u32, t as u32, 1),
22526            block_dim: (head_dim as u32, 1, 1),
22527            shared_mem_bytes: 0,
22528        };
22529        let plus0 = 0i32;
22530        let __s_b2 = self.gpu.stream();
22531        let mut b2 = __s_b2.launch_builder(&fc);
22532        b2.arg(&*part_o)
22533            .arg(&*part_m)
22534            .arg(&*part_l)
22535            .arg(o)
22536            .arg(&hd)
22537            .arg(&nh)
22538            .arg(base_dev)
22539            .arg(&plus0)
22540            .arg(&nspm)
22541            .arg(&spk);
22542        unsafe {
22543            b2.launch(cfg2)?;
22544        }
22545        Ok(())
22546    }
22547
22548    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
22549    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
22550    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
22551    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
22552    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
22553    ///
22554    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
22555    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
22556    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
22557    /// grouping (different but mathematically-equal log-sum-exp merge).
22558    pub fn fa_decode_dc(
22559        &self,
22560        q: &CudaSlice<f32>,
22561        k: &cudarc::driver::CudaView<u8>,
22562        v: &cudarc::driver::CudaView<u8>,
22563        o: &mut CudaSlice<f32>,
22564        head_dim: usize,
22565        n_head: usize,
22566        n_head_kv: usize,
22567        t_kv_dev: &CudaSlice<i32>,
22568        bucket_max: usize,
22569        scale: f32,
22570        k_tok_bytes: usize,
22571        v_tok_bytes: usize,
22572        g: bool,
22573    ) -> Result<(), Box<dyn std::error::Error>> {
22574        self.fa_decode_dc_q8(
22575            q,
22576            k,
22577            v,
22578            o,
22579            head_dim,
22580            n_head,
22581            n_head_kv,
22582            t_kv_dev,
22583            bucket_max,
22584            scale,
22585            k_tok_bytes,
22586            v_tok_bytes,
22587            g,
22588            None,
22589        )
22590    }
22591
22592    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
22593    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
22594    #[allow(clippy::too_many_arguments)]
22595    pub fn fa_decode_dc_q8(
22596        &self,
22597        q: &CudaSlice<f32>,
22598        k: &cudarc::driver::CudaView<u8>,
22599        v: &cudarc::driver::CudaView<u8>,
22600        o: &mut CudaSlice<f32>,
22601        head_dim: usize,
22602        n_head: usize,
22603        n_head_kv: usize,
22604        t_kv_dev: &CudaSlice<i32>,
22605        bucket_max: usize,
22606        scale: f32,
22607        k_tok_bytes: usize,
22608        v_tok_bytes: usize,
22609        g: bool,
22610        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22611    ) -> Result<(), Box<dyn std::error::Error>> {
22612        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
22613        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
22614        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
22615        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
22616        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
22617        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
22618        // 2026-07-12).
22619        let mut fa_vec =
22620            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22621        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22622            fa_vec = false;
22623        } // mirror kvmod/geom
22624        let sp = fa_split_keys(bucket_max, n_head_kv);
22625        let n_splits = if fa_vec {
22626            ((bucket_max + sp - 1) / sp).max(1)
22627        } else {
22628            ((bucket_max + 255) / 256).max(1)
22629        };
22630        let o_len = n_head * n_splits * head_dim;
22631        let ml_len = n_head * n_splits;
22632        let mut part_guard = self.fa_part_pool.lock().unwrap();
22633        if part_guard
22634            .as_ref()
22635            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22636            .unwrap_or(true)
22637        {
22638            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22639            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22640            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22641            // later live allocations land at those addresses, and the next graph REPLAY writes
22642            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22643            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22644            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22645            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22646            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22647            // (total retired < final size).
22648            let old = part_guard.take();
22649            let (co, cm) = old
22650                .as_ref()
22651                .map(|pp| (pp.0.len(), pp.1.len()))
22652                .unwrap_or((0, 0));
22653            if let Some(old) = old {
22654                self.fa_part_retired.lock().unwrap().push(old);
22655            }
22656            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22657                eprintln!(
22658                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22659                    co, o_len, cm, ml_len
22660                );
22661            }
22662            *part_guard = Some((
22663                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22664                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22665                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22666            ));
22667        }
22668        let pg = part_guard.as_mut().unwrap();
22669        self.gpu
22670            .stream()
22671            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22672        self.gpu
22673            .stream()
22674            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22675        self.gpu
22676            .stream()
22677            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22678        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22679        let (hd, nh, nhkv, nsp) = (
22680            head_dim as i32,
22681            n_head as i32,
22682            n_head_kv as i32,
22683            n_splits as i32,
22684        );
22685        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22686        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22687        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
22688        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
22689        let deep = fa_vec
22690            && head_dim == 256
22691            && fa_v4_at(bucket_max)
22692            && !g
22693            && fa_deep_at(bucket_max)
22694            && !matches!(fa_v4_mode(), "noB3" | "stage");
22695        let (f, cfg) = if fa_vec
22696            && head_dim == 512
22697            && bucket_max >= {
22698                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22699                *FA512_MIN_DC.get_or_init(|| {
22700                    std::env::var("MEMRA_FA512_MIN")
22701                        .ok()
22702                        .and_then(|v| v.parse().ok())
22703                        .unwrap_or(512)
22704                })
22705            } {
22706            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
22707            let gqa = (n_head / n_head_kv).max(1) as u32;
22708            (
22709                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22710                LaunchConfig {
22711                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22712                    block_dim: (32, gqa, 1),
22713                    shared_mem_bytes: 0,
22714                },
22715            )
22716        } else if fa_vec && head_dim == 512 {
22717            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
22718            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
22719            let q_view = q.as_view();
22720            let mut o_view = o.as_view_mut();
22721            return self.fa_decode_scalar_unified(
22722                &q_view,
22723                k,
22724                v,
22725                &mut o_view,
22726                head_dim,
22727                n_head,
22728                n_head_kv,
22729                0,
22730                Some(t_kv_dev),
22731                scale,
22732                n_splits,
22733                sp,
22734                k_tok_bytes,
22735                v_tok_bytes,
22736                g,
22737                &mut *part_o,
22738                &mut *part_m,
22739                &mut *part_l,
22740                q8_out,
22741            );
22742        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22743            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
22744            // incl the g-module route + raw-e4m3 sV sizing.
22745            let gqa = (n_head / n_head_kv).max(1) as u32;
22746            let fv = if g {
22747                self.func_g("fa_decode_vec_q_v4_dc")
22748            } else if deep {
22749                self.func("fa_decode_vec_q_v4_deep_dc")
22750            } else {
22751                self.func("fa_decode_vec_q_v4_dc")
22752            };
22753            let shmem =
22754                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22755            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22756            fv.set_attribute(
22757                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22758                shmem as i32,
22759            )?;
22760            (
22761                fv,
22762                LaunchConfig {
22763                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22764                    block_dim: (32, gqa, 1),
22765                    shared_mem_bytes: shmem,
22766                },
22767            )
22768        } else if fa_vec && fa_v3_active(head_dim) {
22769            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
22770            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
22771            let gqa = (n_head / n_head_kv).max(1) as u32;
22772            let fv = if g {
22773                self.func_g("fa_decode_vec_q_v3_dc")
22774            } else {
22775                self.func("fa_decode_vec_q_v3_dc")
22776            };
22777            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
22778            (
22779                fv,
22780                LaunchConfig {
22781                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22782                    block_dim: (32, gqa, 1),
22783                    shared_mem_bytes: shmem,
22784                },
22785            )
22786        } else if fa_vec && fa_v2_on() {
22787            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
22788            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
22789            // a numeric config; eager, rows-verify and graph all switch together).
22790            let gqa = (n_head / n_head_kv).max(1) as u32;
22791            let fv = if g {
22792                self.func_g("fa_decode_vec_q_v2_dc")
22793            } else {
22794                self.func("fa_decode_vec_q_v2_dc")
22795            };
22796            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22797            (
22798                fv,
22799                LaunchConfig {
22800                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22801                    block_dim: (32, gqa, 1),
22802                    shared_mem_bytes: shmem,
22803                },
22804            )
22805        } else if fa_vec {
22806            let gqa = (n_head / n_head_kv).max(1) as u32;
22807            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
22808            let fv = if g {
22809                self.func_g("fa_decode_vec_q_dc")
22810            } else {
22811                self.func("fa_decode_vec_q_dc")
22812            };
22813            (
22814                fv,
22815                LaunchConfig {
22816                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22817                    block_dim: (32, gqa, 1),
22818                    shared_mem_bytes: 0,
22819                },
22820            )
22821        } else {
22822            let q_view = q.as_view();
22823            let mut o_view = o.as_view_mut();
22824            return self.fa_decode_scalar_unified(
22825                &q_view,
22826                k,
22827                v,
22828                &mut o_view,
22829                head_dim,
22830                n_head,
22831                n_head_kv,
22832                0,
22833                Some(t_kv_dev),
22834                scale,
22835                n_splits,
22836                if fa_vec { sp } else { 256 },
22837                k_tok_bytes,
22838                v_tok_bytes,
22839                g,
22840                &mut *part_o,
22841                &mut *part_m,
22842                &mut *part_l,
22843                q8_out,
22844            );
22845        };
22846        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
22847        let __s_b = self.gpu.stream();
22848        let mut b = __s_b.launch_builder(&f);
22849        b.arg(q)
22850            .arg(k)
22851            .arg(v)
22852            .arg(&mut *part_o)
22853            .arg(&mut *part_m)
22854            .arg(&mut *part_l)
22855            .arg(&hd)
22856            .arg(&nh)
22857            .arg(&nhkv)
22858            .arg(t_kv_dev)
22859            .arg(&scale)
22860            .arg(&nsp)
22861            .arg(&ski)
22862            .arg(&ktb)
22863            .arg(&vtb);
22864        unsafe {
22865            b.launch(cfg)?;
22866        }
22867        let cfg2 = LaunchConfig {
22868            grid_dim: (n_head as u32, 1, 1),
22869            block_dim: (head_dim as u32, 1, 1),
22870            shared_mem_bytes: 0,
22871        };
22872        if let Some((oq, od)) = q8_out {
22873            let fc = if g {
22874                self.func_g("fa_decode_combine_q8_1")
22875            } else {
22876                self.fa_func("fa_decode_combine_q8_1", head_dim)
22877            };
22878            let __s_b2 = self.gpu.stream();
22879            let mut b2 = __s_b2.launch_builder(&fc);
22880            b2.arg(&*part_o)
22881                .arg(&*part_m)
22882                .arg(&*part_l)
22883                .arg(oq)
22884                .arg(od)
22885                .arg(&hd)
22886                .arg(&nh)
22887                .arg(&nsp);
22888            unsafe {
22889                b2.launch(cfg2)?;
22890            }
22891            return Ok(());
22892        }
22893        let fc = if g {
22894            self.func_g("fa_decode_combine_f32")
22895        } else {
22896            self.fa_func("fa_decode_combine_f32", head_dim)
22897        };
22898        let __s_b2 = self.gpu.stream();
22899        let mut b2 = __s_b2.launch_builder(&fc);
22900        b2.arg(&*part_o)
22901            .arg(&*part_m)
22902            .arg(&*part_l)
22903            .arg(o)
22904            .arg(&hd)
22905            .arg(&nh)
22906            .arg(&nsp);
22907        unsafe {
22908            b2.launch(cfg2)?;
22909        }
22910        Ok(())
22911    }
22912
22913    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
22914    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
22915    /// at equal rows.
22916    #[allow(clippy::too_many_arguments)]
22917    pub fn append_kv_quantized_dcw(
22918        &self,
22919        k_row: &CudaSlice<f32>,
22920        v_row: &CudaSlice<f32>,
22921        kc: &mut CudaSlice<u8>,
22922        vc: &mut CudaSlice<u8>,
22923        len_dev: &CudaSlice<i32>,
22924        base_dev: Option<&CudaSlice<i32>>,
22925        kv_dim_k: usize,
22926        kv_dim_v: usize,
22927        k_tok_bytes: usize,
22928        v_tok_bytes: usize,
22929    ) -> Result<(), Box<dyn std::error::Error>> {
22930        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22931        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22932        let cfg = LaunchConfig {
22933            grid_dim: (nblk, 1, 1),
22934            block_dim: (32, 1, 1),
22935            shared_mem_bytes: 0,
22936        };
22937        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22938        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22939        let null: u64 = 0;
22940        let __s_b = self.gpu.stream();
22941        let mut b = __s_b.launch_builder(&f);
22942        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22943        match base_dev {
22944            Some(base) => {
22945                b.arg(base);
22946            }
22947            None => {
22948                b.arg(&null);
22949            }
22950        }
22951        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22952        unsafe {
22953            b.launch(cfg)?;
22954        }
22955        Ok(())
22956    }
22957
22958    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
22959    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22960        let f = self.func("inc_i32");
22961        let cfg = LaunchConfig {
22962            grid_dim: (1, 1, 1),
22963            block_dim: (1, 1, 1),
22964            shared_mem_bytes: 0,
22965        };
22966        let __s_b = self.gpu.stream();
22967        let mut b = __s_b.launch_builder(&f);
22968        b.arg(counter);
22969        unsafe {
22970            b.launch(cfg)?;
22971        }
22972        Ok(())
22973    }
22974
22975    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
22976    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
22977    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
22978    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
22979    /// kernel class on this lane); callers keep eager below the vec floor and for any other
22980    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
22981    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
22982    /// alive across bucket growth.
22983    #[allow(clippy::too_many_arguments)]
22984    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
22985    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
22986    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
22987    fn fa_part_pool_grow(
22988        &self,
22989        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22990        o_len: usize,
22991        ml_len: usize,
22992    ) -> Result<(), Box<dyn std::error::Error>> {
22993        if part_guard
22994            .as_ref()
22995            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22996            .unwrap_or(true)
22997        {
22998            let old = part_guard.take();
22999            let (co, cm) = old
23000                .as_ref()
23001                .map(|pp| (pp.0.len(), pp.1.len()))
23002                .unwrap_or((0, 0));
23003            if let Some(old) = old {
23004                self.fa_part_retired.lock().unwrap().push(old);
23005            }
23006            *part_guard = Some((
23007                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23008                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23009                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23010            ));
23011        }
23012        Ok(())
23013    }
23014
23015    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
23016    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
23017    pub fn fa_dcw_pool_ensure(
23018        &self,
23019        head_dim: usize,
23020        n_head: usize,
23021        n_head_kv: usize,
23022        bucket_max: usize,
23023    ) -> Result<(), Box<dyn std::error::Error>> {
23024        let sp = fa_split_keys(bucket_max, n_head_kv);
23025        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23026        let o_len = n_head * n_splits * head_dim;
23027        let ml_len = n_head * n_splits;
23028        let mut part_guard = self.fa_part_pool.lock().unwrap();
23029        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
23030    }
23031
23032    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
23033    /// appended; one launch walks the KV stream once with two query rows (per-row causal
23034    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
23035    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
23036    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
23037    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
23038    /// outputs (the head gate fuses into the combine as in the t=1 path).
23039    #[allow(clippy::too_many_arguments)]
23040    pub fn fa_decode_dcw2(
23041        &self,
23042        q2: &CudaSlice<f32>,
23043        k_ring: &cudarc::driver::CudaView<u8>,
23044        v_ring: &cudarc::driver::CudaView<u8>,
23045        o2: &mut CudaSlice<f32>,
23046        head_dim: usize,
23047        n_head: usize,
23048        n_head_kv: usize,
23049        len_dev: &CudaSlice<i32>,
23050        base_dev: Option<&CudaSlice<i32>>,
23051        window: usize,
23052        bucket_max: usize,
23053        scale: f32,
23054        k_tok_bytes: usize,
23055        v_tok_bytes: usize,
23056        gate2: &CudaSlice<f32>,
23057    ) -> Result<(), Box<dyn std::error::Error>> {
23058        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23059        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23060            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
23061        }
23062        let sp = fa_split_keys(bucket_max, n_head_kv);
23063        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23064        // Partials for BOTH rows: row-major halves.
23065        let o_len = 2 * n_head * n_splits * head_dim;
23066        let ml_len = 2 * n_head * n_splits;
23067        let mut part_guard = self.fa_part_pool.lock().unwrap();
23068        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23069        let pg = part_guard.as_mut().unwrap();
23070        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23071        let (hd, nh, nhkv, nsp) = (
23072            head_dim as i32,
23073            n_head as i32,
23074            n_head_kv as i32,
23075            n_splits as i32,
23076        );
23077        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23078        let (ski, win) = (sp as i32, window as i32);
23079        let gqa = (n_head / n_head_kv).max(1) as u32;
23080        let smem = (32 * head_dim * 2) as u32;
23081        let f = self.func("fa_decode_vec_q_v3_dcw2");
23082        let cfg = LaunchConfig {
23083            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23084            block_dim: (32, gqa, 1),
23085            shared_mem_bytes: smem,
23086        };
23087        let null: u64 = 0;
23088        {
23089            let __s_b = self.gpu.stream();
23090            let mut b = __s_b.launch_builder(&f);
23091            b.arg(q2)
23092                .arg(k_ring)
23093                .arg(v_ring)
23094                .arg(&mut *part_o)
23095                .arg(&mut *part_m)
23096                .arg(&mut *part_l)
23097                .arg(&hd)
23098                .arg(&nh)
23099                .arg(&nhkv)
23100                .arg(len_dev);
23101            match base_dev {
23102                Some(base) => {
23103                    b.arg(base);
23104                }
23105                None => {
23106                    b.arg(&null);
23107                }
23108            }
23109            b.arg(&win)
23110                .arg(&scale)
23111                .arg(&nsp)
23112                .arg(&ski)
23113                .arg(&ktb)
23114                .arg(&vtb);
23115            unsafe {
23116                b.launch(cfg)?;
23117            }
23118        }
23119        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
23120        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
23121        // one launch covers both rows with the exact t=1 program per (row, head).
23122        let fc = {
23123            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23124            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23125                self.func("fa_decode_combine_gate_f32_s")
23126            } else {
23127                self.func("fa_decode_combine_gate_f32")
23128            }
23129        };
23130        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23131        let nh2 = (2 * n_head) as i32;
23132        let cfg2 = LaunchConfig {
23133            grid_dim: ((2 * n_head) as u32, 1, 1),
23134            block_dim: (head_dim as u32, 1, 1),
23135            shared_mem_bytes: if combine_shared {
23136                (2 * n_splits * 4) as u32
23137            } else {
23138                0
23139            },
23140        };
23141        let __s_b2 = self.gpu.stream();
23142        let mut b2 = __s_b2.launch_builder(&fc);
23143        b2.arg(&*part_o)
23144            .arg(&*part_m)
23145            .arg(&*part_l)
23146            .arg(gate2)
23147            .arg(o2)
23148            .arg(&hd)
23149            .arg(&nh2)
23150            .arg(&nsp);
23151        unsafe {
23152            b2.launch(cfg2)?;
23153        }
23154        Ok(())
23155    }
23156
23157    pub fn fa_decode_dcw(
23158        &self,
23159        q: &CudaSlice<f32>,
23160        k_ring: &cudarc::driver::CudaView<u8>,
23161        v_ring: &cudarc::driver::CudaView<u8>,
23162        o: &mut CudaSlice<f32>,
23163        head_dim: usize,
23164        n_head: usize,
23165        n_head_kv: usize,
23166        len_dev: &CudaSlice<i32>,
23167        base_dev: Option<&CudaSlice<i32>>,
23168        window: usize,
23169        bucket_max: usize,
23170        scale: f32,
23171        k_tok_bytes: usize,
23172        v_tok_bytes: usize,
23173        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
23174        // one launch saved); `o` then receives the GATED output and the caller skips its
23175        // attn_head_gate call.
23176        fused_gate: Option<&CudaSlice<f32>>,
23177    ) -> Result<(), Box<dyn std::error::Error>> {
23178        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23179        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23180            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"
23181                .into());
23182        }
23183        let sp = fa_split_keys(bucket_max, n_head_kv);
23184        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23185        let o_len = n_head * n_splits * head_dim;
23186        let ml_len = n_head * n_splits;
23187        let mut part_guard = self.fa_part_pool.lock().unwrap();
23188        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23189        let pg = part_guard.as_mut().unwrap();
23190        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
23191        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
23192        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
23193        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
23194        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23195        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
23196        // finds the attention children BY their three-memset signature and updates the
23197        // memset widths per bucket — capturing without them silently kills retargeting
23198        // (battery-v8 token drift, 2026-08-21).
23199        let memset_on = *MEMSET_ON
23200            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23201            || crate::tp::token_graph_building();
23202        if memset_on {
23203            self.gpu
23204                .stream()
23205                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23206            self.gpu
23207                .stream()
23208                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23209            self.gpu
23210                .stream()
23211                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23212        }
23213        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23214        let (hd, nh, nhkv, nsp) = (
23215            head_dim as i32,
23216            n_head as i32,
23217            n_head_kv as i32,
23218            n_splits as i32,
23219        );
23220        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23221        let (ski, win) = (sp as i32, window as i32);
23222        let gqa = (n_head / n_head_kv).max(1) as u32;
23223        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
23224        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
23225        // see fa_dec_v3_walk_u). Same launch geometry.
23226        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23227        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23228        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23229            Ok("2") => 2,
23230            Ok("1") => 1,
23231            _ => 0,
23232        });
23233        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
23234        // permission-blocked in this container and the module params are not exposed, so this
23235        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
23236        // prints cumulative cycle shares every 430 launches.
23237        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23238        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23239        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23240            std::sync::Mutex::new(None);
23241        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
23242        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
23243        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
23244        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23245        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23246            && (n_head / n_head_kv) % 2 == 0
23247            && (n_head / n_head_kv) >= 2;
23248        let f = if fprof {
23249            self.func("fa_decode_vec_q_v3_dcw_prof")
23250        } else if hs2 {
23251            self.func("fa_decode_vec_q_v3_dcw_hs2")
23252        } else if hoist == 2 {
23253            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
23254            self.func("fa_decode_vec_q_v3_dcw_hc")
23255        } else if hoist == 1 {
23256            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
23257            self.func("fa_decode_vec_q_v3_dcw_h")
23258        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23259            self.func("fa_decode_vec_q_v3_dcw_u8")
23260        } else {
23261            self.func("fa_decode_vec_q_v3_dcw")
23262        };
23263        let cfg = LaunchConfig {
23264            grid_dim: if hs2 {
23265                ((2 * n_head_kv) as u32, n_splits as u32, 1)
23266            } else {
23267                (n_head_kv as u32, n_splits as u32, 1)
23268            },
23269            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23270            shared_mem_bytes: smem,
23271        };
23272        let null: u64 = 0;
23273        let __s_b = self.gpu.stream();
23274        let mut b = __s_b.launch_builder(&f);
23275        b.arg(q)
23276            .arg(k_ring)
23277            .arg(v_ring)
23278            .arg(&mut *part_o)
23279            .arg(&mut *part_m)
23280            .arg(&mut *part_l)
23281            .arg(&hd)
23282            .arg(&nh)
23283            .arg(&nhkv)
23284            .arg(len_dev);
23285        match base_dev {
23286            Some(base) => {
23287                b.arg(base);
23288            }
23289            None => {
23290                b.arg(&null);
23291            }
23292        }
23293        b.arg(&win)
23294            .arg(&scale)
23295            .arg(&nsp)
23296            .arg(&ski)
23297            .arg(&ktb)
23298            .arg(&vtb);
23299        if fprof {
23300            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
23301            if guard
23302                .as_ref()
23303                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
23304            {
23305                *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
23306            }
23307            let (_, buf) = guard.as_mut().expect("armed above");
23308            b.arg(&*buf);
23309            unsafe {
23310                b.launch(cfg)?;
23311            }
23312            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23313            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
23314            if n % 430 == 0 {
23315                self.stream().synchronize()?;
23316                let h = self.dtoh_u64(buf)?;
23317                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
23318                let tot: u64 = h[..6].iter().sum();
23319                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
23320                for (i, name) in phases.iter().enumerate() {
23321                    let pct = if tot > 0 {
23322                        h[i] as f64 / tot as f64 * 100.0
23323                    } else {
23324                        0.0
23325                    };
23326                    line.push_str(&format!(" {name}={pct:.1}%"));
23327                }
23328                if h[6] > 0 {
23329                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
23330                }
23331                eprintln!("{line}");
23332            }
23333        } else {
23334            unsafe {
23335                b.launch(cfg)?;
23336            }
23337        }
23338        let mut combine_shared = false;
23339        let fc = if fused_gate.is_some() {
23340            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
23341            // n_splits-deep dependent global load chain every thread used to walk twice).
23342            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23343            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23344                combine_shared = true;
23345                self.func("fa_decode_combine_gate_f32_s")
23346            } else {
23347                self.func("fa_decode_combine_gate_f32")
23348            }
23349        } else {
23350            self.fa_func("fa_decode_combine_f32", head_dim)
23351        };
23352        let cfg2 = LaunchConfig {
23353            grid_dim: (n_head as u32, 1, 1),
23354            block_dim: (head_dim as u32, 1, 1),
23355            shared_mem_bytes: if combine_shared {
23356                (2 * n_splits * 4) as u32
23357            } else {
23358                0
23359            },
23360        };
23361        let __s_b2 = self.gpu.stream();
23362        let mut b2 = __s_b2.launch_builder(&fc);
23363        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
23364        if let Some(gate_row) = fused_gate {
23365            b2.arg(gate_row);
23366        }
23367        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
23368        unsafe {
23369            b2.launch(cfg2)?;
23370        }
23371        Ok(())
23372    }
23373
23374    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
23375    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
23376    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
23377    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
23378    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
23379    pub fn fa_geom_eager(
23380        &self,
23381        t_kv: usize,
23382        head_dim: usize,
23383        n_head_kv: usize,
23384        g: bool,
23385    ) -> (bool, usize) {
23386        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
23387        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
23388        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
23389        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
23390        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
23391        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
23392        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
23393        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
23394        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
23395        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
23396        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
23397        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
23398        // family; everything else falls to the g-module scalar.
23399        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
23400        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
23401        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
23402        if g && head_dim == 256 && !fa_v4_at(t_kv) {
23403            fa_vec = false;
23404        }
23405        let sp = fa_split_keys(t_kv, n_head_kv);
23406        let n_splits = if fa_vec {
23407            ((t_kv + sp - 1) / sp).max(1)
23408        } else {
23409            ((t_kv + 255) / 256).max(1)
23410        };
23411        (fa_vec, n_splits)
23412    }
23413
23414    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
23415    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
23416    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
23417    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
23418    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
23419    pub fn fa_bucket_key(
23420        &self,
23421        t_kv: usize,
23422        head_dim: usize,
23423        n_head_kv: usize,
23424        g: bool,
23425    ) -> (bool, usize) {
23426        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
23427    }
23428
23429    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
23430    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
23431    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
23432    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
23433    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
23434    /// device data) — every per-step varying scalar must come from a device counter. Returns the
23435    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
23436    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
23437    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
23438    /// replays (transients returning to the pool get reused by unrelated work and corrupt
23439    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
23440    pub fn capture_graph_retained<F>(
23441        &self,
23442        step: F,
23443    ) -> Result<
23444        (
23445            cudarc::driver::CudaGraph,
23446            Vec<Box<dyn std::any::Any + Send>>,
23447        ),
23448        Box<dyn std::error::Error>,
23449    >
23450    where
23451        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23452    {
23453        use cudarc::driver::sys::CUgraphInstantiate_flags;
23454        self.capture_graph_retained_flags(
23455            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23456            step,
23457        )
23458    }
23459
23460    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
23461    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
23462    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
23463    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
23464    pub fn capture_graph_retained_flags<F>(
23465        &self,
23466        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
23467        mut step: F,
23468    ) -> Result<
23469        (
23470            cudarc::driver::CudaGraph,
23471            Vec<Box<dyn std::any::Any + Send>>,
23472        ),
23473        Box<dyn std::error::Error>,
23474    >
23475    where
23476        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23477    {
23478        use cudarc::driver::sys::CUstreamCaptureMode;
23479        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
23480        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
23481        // while the capture region is open become dead copy NODES replayed every launch
23482        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
23483        // warmup runs allocate the same transient sequence at the same pool addresses, so
23484        // retaining the warmup clones preserves the draft-graph fix without polluting the
23485        // captured graph.
23486        self.capture_keep.lock().unwrap().clear();
23487        let was_tracking = self.gpu.ctx.is_event_tracking();
23488        if was_tracking {
23489            unsafe {
23490                self.gpu.ctx.disable_event_tracking();
23491            }
23492        }
23493        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23494            self.capture_keep_on
23495                .store(true, std::sync::atomic::Ordering::Relaxed);
23496            let w = (|| {
23497                step(self)?;
23498                step(self)
23499            })();
23500            self.capture_keep_on
23501                .store(false, std::sync::atomic::Ordering::Relaxed);
23502            w?;
23503            self.gpu.stream().synchronize()?;
23504            self.gpu
23505                .stream()
23506                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23507            let r = step(self);
23508            let g = self.gpu.stream().end_capture(flags);
23509            r?;
23510            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23511            graph.upload()?;
23512            Ok(graph)
23513        };
23514        let result = run();
23515        self.capture_keep_on
23516            .store(false, std::sync::atomic::Ordering::Relaxed);
23517        if was_tracking {
23518            unsafe {
23519                self.gpu.ctx.enable_event_tracking();
23520            }
23521        }
23522        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
23523        Ok((result?, keeper))
23524    }
23525
23526    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
23527    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
23528    /// alloc-free with persistent operands, and their bodies carry device side effects
23529    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
23530    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
23531    pub fn capture_graph_retained_nowarm<F>(
23532        &self,
23533        mut step: F,
23534    ) -> Result<
23535        (
23536            cudarc::driver::CudaGraph,
23537            Vec<Box<dyn std::any::Any + Send>>,
23538        ),
23539        Box<dyn std::error::Error>,
23540    >
23541    where
23542        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23543    {
23544        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23545        let was_tracking = self.gpu.ctx.is_event_tracking();
23546        if was_tracking {
23547            unsafe {
23548                self.gpu.ctx.disable_event_tracking();
23549            }
23550        }
23551        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23552            self.gpu.stream().synchronize()?;
23553            self.gpu
23554                .stream()
23555                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23556            let r = step(self);
23557            let g = self.gpu.stream().end_capture(
23558                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23559            );
23560            r?;
23561            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23562            graph.upload()?;
23563            Ok(graph)
23564        };
23565        let result = run();
23566        if was_tracking {
23567            unsafe {
23568                self.gpu.ctx.enable_event_tracking();
23569            }
23570        }
23571        Ok((result?, Vec::new()))
23572    }
23573
23574    pub fn capture_graph<F>(
23575        &self,
23576        mut step: F,
23577    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
23578    where
23579        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23580    {
23581        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23582        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
23583        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
23584        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
23585        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
23586        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
23587        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
23588        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
23589        let was_tracking = self.gpu.ctx.is_event_tracking();
23590        if was_tracking {
23591            unsafe {
23592                self.gpu.ctx.disable_event_tracking();
23593            }
23594        }
23595        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
23596        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
23597        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
23598        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
23599        // measure that scan's real cost on the generic path. Diagnostic door only; the
23600        // default stays AUTO_FREE until a measured A/B justifies moving it.
23601        let iflag = {
23602            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23603            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23604                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
23605                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
23606                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23607                Ok("priority") => {
23608                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23609                }
23610                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23611            })
23612        };
23613        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
23614        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
23615        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
23616        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
23617        // eager step executions and are node-count-invariant. Printing the split bounds the
23618        // refactor's ceiling instead of assuming it.
23619        let ct = {
23620            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23621            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23622        };
23623        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
23624        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
23625        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
23626        // chased, and node-count-invariant, so no capture-body refactor could touch it.
23627        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
23628        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
23629        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
23630        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
23631        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
23632        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
23633        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
23634        // grow and never frees, resident counters/scratch, cache set in place), and the
23635        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
23636        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
23637        // settling and pool mapping. Arbitrated adversarially, not by taste:
23638        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
23639        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
23640        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
23641        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
23642        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
23643        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
23644        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
23645        let warmups = {
23646            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23647            *W.get_or_init(|| {
23648                std::env::var("MEMRA_GRAPH_WARMUPS")
23649                    .ok()
23650                    .and_then(|v| v.parse().ok())
23651                    .filter(|n| *n >= 1)
23652                    .unwrap_or(1)
23653            })
23654        };
23655        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23656            let t_w = std::time::Instant::now();
23657            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
23658            for _ in 0..warmups {
23659                step(self)?;
23660            }
23661            self.gpu.stream().synchronize()?;
23662            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23663            // capture the third run.
23664            let t_c = std::time::Instant::now();
23665            self.gpu
23666                .stream()
23667                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23668            // If the body errors mid-capture, end the capture before propagating so the stream isn't
23669            // left in a capturing state.
23670            let r = step(self);
23671            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23672            let t_i = std::time::Instant::now();
23673            let g = self.gpu.stream().end_capture(iflag);
23674            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23675            r?;
23676            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23677            let t_u = std::time::Instant::now();
23678            graph.upload()?;
23679            if ct {
23680                println!(
23681                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
23682                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
23683                    t_u.elapsed().as_secs_f64() * 1e3
23684                );
23685            }
23686            Ok(graph)
23687        };
23688        let result = run();
23689        if was_tracking {
23690            unsafe {
23691                self.gpu.ctx.enable_event_tracking();
23692            }
23693        }
23694        result
23695    }
23696
23697    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
23698    pub fn gdn_scan_s128_view(
23699        &self,
23700        q: &CudaSlice<f32>,
23701        k: &CudaSlice<f32>,
23702        v: &CudaSlice<f32>,
23703        g: &CudaSlice<f32>,
23704        beta: &CudaSlice<f32>,
23705        state_in: &cudarc::driver::CudaView<f32>,
23706        state_out: &mut cudarc::driver::CudaViewMut<f32>,
23707        o: &mut CudaSlice<f32>,
23708        n_head: usize,
23709        t: usize,
23710        scale: f32,
23711    ) -> Result<(), Box<dyn std::error::Error>> {
23712        let f = self.func("gdn_scan_s128");
23713        const S_V: u32 = 128;
23714        const WARP: u32 = 32;
23715        const COLS: u32 = 4;
23716        let cfg = LaunchConfig {
23717            grid_dim: (n_head as u32, 1, S_V / COLS),
23718            block_dim: (WARP, COLS, 1),
23719            shared_mem_bytes: 0,
23720        };
23721        let (h, ti) = (n_head as i32, t as i32);
23722        let __s_b = self.gpu.stream();
23723        let mut b = __s_b.launch_builder(&f);
23724        b.arg(q)
23725            .arg(k)
23726            .arg(v)
23727            .arg(g)
23728            .arg(beta)
23729            .arg(state_in)
23730            .arg(state_out)
23731            .arg(o)
23732            .arg(&h)
23733            .arg(&ti)
23734            .arg(&scale);
23735        unsafe {
23736            b.launch(cfg)?;
23737        }
23738        Ok(())
23739    }
23740
23741    /// conv1d where the input is a CudaView (resident conv state assembled in place).
23742    pub fn ssm_conv1d_view(
23743        &self,
23744        x: &cudarc::driver::CudaView<f32>,
23745        w: &CudaSlice<f32>,
23746        y: &mut CudaSlice<f32>,
23747        conv_dim: usize,
23748        t: usize,
23749        d_conv: usize,
23750        silu: bool,
23751    ) -> Result<(), Box<dyn std::error::Error>> {
23752        let f = self.func("ssm_conv1d_silu_f32");
23753        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
23754        let cfg = LaunchConfig {
23755            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23756            block_dim: (256, 1, 1),
23757            shared_mem_bytes: 0,
23758        };
23759        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23760        let __s_b = self.gpu.stream();
23761        let mut b = __s_b.launch_builder(&f);
23762        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23763        unsafe {
23764            b.launch(cfg)?;
23765        }
23766        Ok(())
23767    }
23768
23769    /// Depthwise causal conv1d + optional SiLU.
23770    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
23771    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
23772    /// FUSED prefill conv (token-major input, zero left-state): replaces
23773    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
23774    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
23775    pub fn ssm_conv1d_tm(
23776        &self,
23777        qkv_tm: &CudaSlice<f32>,
23778        w: &CudaSlice<f32>,
23779        y: &mut CudaSlice<f32>,
23780        conv_dim: usize,
23781        t: usize,
23782        d_conv: usize,
23783    ) -> Result<(), Box<dyn std::error::Error>> {
23784        let f = self.func("ssm_conv1d_tm_f32");
23785        let cfg = LaunchConfig {
23786            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23787            block_dim: (256, 1, 1),
23788            shared_mem_bytes: 0,
23789        };
23790        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23791        let __s_b = self.gpu.stream();
23792        let mut b = __s_b.launch_builder(&f);
23793        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23794        unsafe {
23795            b.launch(cfg)?;
23796        }
23797        Ok(())
23798    }
23799
23800    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
23801    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
23802    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
23803    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
23804    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
23805    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
23806    /// columns; the final ring == what T sequential decode ring rolls leave).
23807    pub fn ssm_conv1d_tm_state(
23808        &self,
23809        qkv_tm: &CudaSlice<f32>,
23810        conv_state: &mut CudaSlice<f32>,
23811        w: &CudaSlice<f32>,
23812        y: &mut CudaSlice<f32>,
23813        conv_dim: usize,
23814        t: usize,
23815        d_conv: usize,
23816    ) -> Result<(), Box<dyn std::error::Error>> {
23817        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23818    }
23819
23820    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
23821    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
23822    #[allow(clippy::too_many_arguments)]
23823    pub fn ssm_conv1d_tm_state_pad(
23824        &self,
23825        qkv_tm: &CudaSlice<f32>,
23826        conv_state: &mut CudaSlice<f32>,
23827        w: &CudaSlice<f32>,
23828        y: &mut CudaSlice<f32>,
23829        conv_dim: usize,
23830        t: usize,
23831        d_conv: usize,
23832        pad_len: Option<&CudaSlice<i32>>,
23833    ) -> Result<(), Box<dyn std::error::Error>> {
23834        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23835        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23836        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23837        // cloning first keeps the ordering trivially correct under any future stream split.
23838        let ring_old = if t < d_conv - 1 {
23839            Some(self.clone_dtod(conv_state)?)
23840        } else {
23841            None
23842        };
23843        {
23844            let f = self.func("ssm_conv1d_tm_state_f32");
23845            let cfg = LaunchConfig {
23846                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23847                block_dim: (256, 1, 1),
23848                shared_mem_bytes: 0,
23849            };
23850            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23851            let __s_b = self.gpu.stream();
23852            let mut b = __s_b.launch_builder(&f);
23853            b.arg(qkv_tm)
23854                .arg(&*conv_state)
23855                .arg(w)
23856                .arg(y)
23857                .arg(&cd)
23858                .arg(&ti)
23859                .arg(&dc);
23860            unsafe {
23861                b.launch(cfg)?;
23862            }
23863        }
23864        match (ring_old, pad_len) {
23865            (None, Some(len_d)) => {
23866                let f = self.func("ssm_conv_ring_update_dev_f32");
23867                let n = conv_dim * (d_conv - 1);
23868                let cfg = LaunchConfig::for_num_elems(n as u32);
23869                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23870                let __s_b = self.gpu.stream();
23871                let mut b = __s_b.launch_builder(&f);
23872                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23873                unsafe {
23874                    b.launch(cfg)?;
23875                }
23876            }
23877            (None, None) => {
23878                let f = self.func("ssm_conv_ring_update_f32");
23879                let n = conv_dim * (d_conv - 1);
23880                let cfg = LaunchConfig::for_num_elems(n as u32);
23881                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23882                let __s_b = self.gpu.stream();
23883                let mut b = __s_b.launch_builder(&f);
23884                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23885                unsafe {
23886                    b.launch(cfg)?;
23887                }
23888            }
23889            (Some(old), _) => {
23890                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23891            }
23892        }
23893        Ok(())
23894    }
23895
23896    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
23897    pub fn ssm_conv1d_tm_state_pad_v(
23898        &self,
23899        qkv_tm: &cudarc::driver::CudaView<f32>,
23900        conv_state: &mut CudaSlice<f32>,
23901        w: &CudaSlice<f32>,
23902        y: &mut CudaSlice<f32>,
23903        conv_dim: usize,
23904        t: usize,
23905        d_conv: usize,
23906        pad_len: Option<&CudaSlice<i32>>,
23907    ) -> Result<(), Box<dyn std::error::Error>> {
23908        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23909        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23910        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23911        // cloning first keeps the ordering trivially correct under any future stream split.
23912        let ring_old = if t < d_conv - 1 {
23913            Some(self.clone_dtod(conv_state)?)
23914        } else {
23915            None
23916        };
23917        {
23918            let f = self.func("ssm_conv1d_tm_state_f32");
23919            let cfg = LaunchConfig {
23920                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23921                block_dim: (256, 1, 1),
23922                shared_mem_bytes: 0,
23923            };
23924            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23925            let __s_b = self.gpu.stream();
23926            let mut b = __s_b.launch_builder(&f);
23927            b.arg(qkv_tm)
23928                .arg(&*conv_state)
23929                .arg(w)
23930                .arg(y)
23931                .arg(&cd)
23932                .arg(&ti)
23933                .arg(&dc);
23934            unsafe {
23935                b.launch(cfg)?;
23936            }
23937        }
23938        match (ring_old, pad_len) {
23939            (None, Some(len_d)) => {
23940                let f = self.func("ssm_conv_ring_update_dev_f32");
23941                let n = conv_dim * (d_conv - 1);
23942                let cfg = LaunchConfig::for_num_elems(n as u32);
23943                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23944                let __s_b = self.gpu.stream();
23945                let mut b = __s_b.launch_builder(&f);
23946                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23947                unsafe {
23948                    b.launch(cfg)?;
23949                }
23950            }
23951            (None, None) => {
23952                let f = self.func("ssm_conv_ring_update_f32");
23953                let n = conv_dim * (d_conv - 1);
23954                let cfg = LaunchConfig::for_num_elems(n as u32);
23955                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23956                let __s_b = self.gpu.stream();
23957                let mut b = __s_b.launch_builder(&f);
23958                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23959                unsafe {
23960                    b.launch(cfg)?;
23961                }
23962            }
23963            (Some(_), _) => unreachable!(
23964                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23965            ),
23966        }
23967        Ok(())
23968    }
23969
23970    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
23971    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
23972    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
23973    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
23974    pub fn ssm_conv_ring_rebuild(
23975        &self,
23976        qkv_tm: &CudaSlice<f32>,
23977        ring_old: &CudaSlice<f32>,
23978        conv_state: &mut CudaSlice<f32>,
23979        conv_dim: usize,
23980        tc: usize,
23981        d_conv: usize,
23982    ) -> Result<(), Box<dyn std::error::Error>> {
23983        let f = self.func("ssm_conv_ring_rebuild_f32");
23984        let n = conv_dim * (d_conv - 1);
23985        let cfg = LaunchConfig::for_num_elems(n as u32);
23986        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23987        let __s_b = self.gpu.stream();
23988        let mut b = __s_b.launch_builder(&f);
23989        b.arg(qkv_tm)
23990            .arg(ring_old)
23991            .arg(conv_state)
23992            .arg(&cd)
23993            .arg(&ti)
23994            .arg(&dc);
23995        unsafe {
23996            b.launch(cfg)?;
23997        }
23998        Ok(())
23999    }
24000
24001    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
24002    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
24003    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
24004    /// the argmax + run-spec gates are the authority.
24005    #[allow(clippy::too_many_arguments)]
24006    pub fn gdn_prep_decode(
24007        &self,
24008        conv_out: &CudaSlice<f32>,
24009        beta_raw: &CudaSlice<f32>,
24010        alpha: &CudaSlice<f32>,
24011        dt_bias: &CudaSlice<f32>,
24012        a: &CudaSlice<f32>,
24013        q_l2: &mut CudaSlice<f32>,
24014        k_l2: &mut CudaSlice<f32>,
24015        v_g: &mut CudaSlice<f32>,
24016        beta: &mut CudaSlice<f32>,
24017        g_log: &mut CudaSlice<f32>,
24018        d_state: usize,
24019        num_v: usize,
24020        num_k: usize,
24021        key_dim: usize,
24022        eps: f32,
24023    ) -> Result<(), Box<dyn std::error::Error>> {
24024        let f = self.func("gdn_prep_decode_f32");
24025        let cfg = LaunchConfig {
24026            grid_dim: (num_v as u32, 1, 1),
24027            block_dim: (32, 4, 1),
24028            shared_mem_bytes: 0,
24029        };
24030        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24031        let __s_b = self.gpu.stream();
24032        let mut b = __s_b.launch_builder(&f);
24033        b.arg(conv_out)
24034            .arg(beta_raw)
24035            .arg(alpha)
24036            .arg(dt_bias)
24037            .arg(a)
24038            .arg(q_l2)
24039            .arg(k_l2)
24040            .arg(v_g)
24041            .arg(beta)
24042            .arg(g_log)
24043            .arg(&ds)
24044            .arg(&nv)
24045            .arg(&nk)
24046            .arg(&kd)
24047            .arg(&eps);
24048        unsafe {
24049            b.launch(cfg)?;
24050        }
24051        Ok(())
24052    }
24053
24054    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
24055    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
24056    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
24057    #[allow(clippy::too_many_arguments)]
24058    pub fn ssm_conv1d_gdn(
24059        &self,
24060        qkv_tm: &CudaSlice<f32>,
24061        w: &CudaSlice<f32>,
24062        q_g: &mut CudaSlice<f32>,
24063        k_g: &mut CudaSlice<f32>,
24064        v_g: &mut CudaSlice<f32>,
24065        conv_dim: usize,
24066        t: usize,
24067        d_conv: usize,
24068        d_state: usize,
24069        num_v: usize,
24070        num_k: usize,
24071        key_dim: usize,
24072    ) -> Result<(), Box<dyn std::error::Error>> {
24073        let f = self.func("ssm_conv1d_gdn_f32");
24074        let cfg = LaunchConfig {
24075            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24076            block_dim: (256, 1, 1),
24077            shared_mem_bytes: 0,
24078        };
24079        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24080        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24081        let __s_b = self.gpu.stream();
24082        let mut b = __s_b.launch_builder(&f);
24083        b.arg(qkv_tm)
24084            .arg(w)
24085            .arg(q_g)
24086            .arg(k_g)
24087            .arg(v_g)
24088            .arg(&cd)
24089            .arg(&ti)
24090            .arg(&dc)
24091            .arg(&ds)
24092            .arg(&nv)
24093            .arg(&nk)
24094            .arg(&kd);
24095        unsafe {
24096            b.launch(cfg)?;
24097        }
24098        Ok(())
24099    }
24100
24101    pub fn ssm_conv1d(
24102        &self,
24103        x: &CudaSlice<f32>,
24104        w: &CudaSlice<f32>,
24105        y: &mut CudaSlice<f32>,
24106        conv_dim: usize,
24107        t: usize,
24108        d_conv: usize,
24109        silu: bool,
24110    ) -> Result<(), Box<dyn std::error::Error>> {
24111        let f = self.func("ssm_conv1d_silu_f32");
24112        let cfg = LaunchConfig {
24113            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24114            block_dim: (256, 1, 1),
24115            shared_mem_bytes: 0,
24116        };
24117        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24118        let __s_b = self.gpu.stream();
24119        let mut b = __s_b.launch_builder(&f);
24120        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24121        unsafe {
24122            b.launch(cfg)?;
24123        }
24124        Ok(())
24125    }
24126
24127    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
24128    /// o:[128,H,T]. Single sequence.
24129    pub fn gdn_scan_s128(
24130        &self,
24131        q: &CudaSlice<f32>,
24132        k: &CudaSlice<f32>,
24133        v: &CudaSlice<f32>,
24134        g: &CudaSlice<f32>,
24135        beta: &CudaSlice<f32>,
24136        state_in: &CudaSlice<f32>,
24137        state_out: &mut CudaSlice<f32>,
24138        o: &mut CudaSlice<f32>,
24139        n_head: usize,
24140        t: usize,
24141        scale: f32,
24142    ) -> Result<(), Box<dyn std::error::Error>> {
24143        let f = self.func("gdn_scan_s128");
24144        const S_V: u32 = 128;
24145        const WARP: u32 = 32;
24146        const COLS_PER_BLOCK: u32 = 4;
24147        let cfg = LaunchConfig {
24148            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
24149            block_dim: (WARP, COLS_PER_BLOCK, 1),
24150            shared_mem_bytes: 0,
24151        };
24152        let (h, ti) = (n_head as i32, t as i32);
24153        let __s_b = self.gpu.stream();
24154        let mut b = __s_b.launch_builder(&f);
24155        b.arg(q)
24156            .arg(k)
24157            .arg(v)
24158            .arg(g)
24159            .arg(beta)
24160            .arg(state_in)
24161            .arg(state_out)
24162            .arg(o)
24163            .arg(&h)
24164            .arg(&ti)
24165            .arg(&scale);
24166        unsafe {
24167            b.launch(cfg)?;
24168        }
24169        Ok(())
24170    }
24171
24172    // ==== B2' batched decode state ops (decode_batch.rs) ====
24173    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
24174    // Bodies are the single-seq kernels per sequence — bit-identical per row.
24175
24176    #[allow(clippy::too_many_arguments)]
24177    pub fn ssm_conv1d_fused_decode_b(
24178        &self,
24179        qkv_cols: &CudaSlice<f32>,
24180        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24181        w: &CudaSlice<f32>,
24182        conv_outs: &mut CudaSlice<f32>,
24183        conv_dim: usize,
24184        d_conv: usize,
24185        b_n: usize,
24186    ) -> Result<(), Box<dyn std::error::Error>> {
24187        let f = self.func("ssm_conv1d_fused_decode_b_f32");
24188        let cfg = LaunchConfig {
24189            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24190            block_dim: (256, 1, 1),
24191            shared_mem_bytes: 0,
24192        };
24193        let (cd, dc) = (conv_dim as i32, d_conv as i32);
24194        let __s_b = self.gpu.stream();
24195        let mut b = __s_b.launch_builder(&f);
24196        b.arg(qkv_cols)
24197            .arg(conv_state_ptrs)
24198            .arg(w)
24199            .arg(conv_outs)
24200            .arg(&cd)
24201            .arg(&dc);
24202        unsafe {
24203            b.launch(cfg)?;
24204        }
24205        Ok(())
24206    }
24207
24208    #[allow(clippy::too_many_arguments)]
24209    pub fn gdn_prep_decode_b(
24210        &self,
24211        conv_outs: &CudaSlice<f32>,
24212        beta_raws: &CudaSlice<f32>,
24213        alphas: &CudaSlice<f32>,
24214        dt_bias: &CudaSlice<f32>,
24215        a: &CudaSlice<f32>,
24216        q_l2: &mut CudaSlice<f32>,
24217        k_l2: &mut CudaSlice<f32>,
24218        v_g: &mut CudaSlice<f32>,
24219        beta: &mut CudaSlice<f32>,
24220        g_log: &mut CudaSlice<f32>,
24221        d_state: usize,
24222        num_v: usize,
24223        num_k: usize,
24224        key_dim: usize,
24225        eps: f32,
24226        conv_dim: usize,
24227        b_n: usize,
24228    ) -> Result<(), Box<dyn std::error::Error>> {
24229        let f = self.func("gdn_prep_decode_b_f32");
24230        let cfg = LaunchConfig {
24231            grid_dim: (num_v as u32, 1, b_n as u32),
24232            block_dim: (32, 4, 1),
24233            shared_mem_bytes: 0,
24234        };
24235        let (ds, nv, nk, kd, cd) = (
24236            d_state as i32,
24237            num_v as i32,
24238            num_k as i32,
24239            key_dim as i32,
24240            conv_dim as i32,
24241        );
24242        let __s_b = self.gpu.stream();
24243        let mut b = __s_b.launch_builder(&f);
24244        b.arg(conv_outs)
24245            .arg(beta_raws)
24246            .arg(alphas)
24247            .arg(dt_bias)
24248            .arg(a)
24249            .arg(q_l2)
24250            .arg(k_l2)
24251            .arg(v_g)
24252            .arg(beta)
24253            .arg(g_log)
24254            .arg(&ds)
24255            .arg(&nv)
24256            .arg(&nk)
24257            .arg(&kd)
24258            .arg(&eps)
24259            .arg(&cd);
24260        unsafe {
24261            b.launch(cfg)?;
24262        }
24263        Ok(())
24264    }
24265
24266    #[allow(clippy::too_many_arguments)]
24267    pub fn gdn_scan_s128_batched(
24268        &self,
24269        q: &CudaSlice<f32>,
24270        k: &CudaSlice<f32>,
24271        v: &CudaSlice<f32>,
24272        g: &CudaSlice<f32>,
24273        beta: &CudaSlice<f32>,
24274        state_in_ptrs: &cudarc::driver::CudaView<u64>,
24275        state_out_ptrs: &cudarc::driver::CudaView<u64>,
24276        o: &mut CudaSlice<f32>,
24277        n_head: usize,
24278        b_n: usize,
24279        scale: f32,
24280    ) -> Result<(), Box<dyn std::error::Error>> {
24281        let f = self.func("gdn_scan_s128_b");
24282        const S_V: u32 = 128;
24283        const WARP: u32 = 32;
24284        const COLS_PER_BLOCK: u32 = 4;
24285        let cfg = LaunchConfig {
24286            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24287            block_dim: (WARP, COLS_PER_BLOCK, 1),
24288            shared_mem_bytes: 0,
24289        };
24290        let h = n_head as i32;
24291        let __s_b = self.gpu.stream();
24292        let mut b = __s_b.launch_builder(&f);
24293        b.arg(q)
24294            .arg(k)
24295            .arg(v)
24296            .arg(g)
24297            .arg(beta)
24298            .arg(state_in_ptrs)
24299            .arg(state_out_ptrs)
24300            .arg(o)
24301            .arg(&h)
24302            .arg(&scale);
24303        unsafe {
24304            b.launch(cfg)?;
24305        }
24306        Ok(())
24307    }
24308
24309    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
24310    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
24311    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
24312    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
24313    /// numeric class; only the pointer arithmetic moved host-side.
24314    #[allow(clippy::too_many_arguments)]
24315    pub fn ssm_conv1d_fused_decode_b_view(
24316        &self,
24317        qkv_cols: &cudarc::driver::CudaView<f32>,
24318        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24319        w: &CudaSlice<f32>,
24320        conv_outs: &mut CudaSlice<f32>,
24321        conv_dim: usize,
24322        d_conv: usize,
24323        b_n: usize,
24324    ) -> Result<(), Box<dyn std::error::Error>> {
24325        let f = self.func("ssm_conv1d_fused_decode_b_f32");
24326        let cfg = LaunchConfig {
24327            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24328            block_dim: (256, 1, 1),
24329            shared_mem_bytes: 0,
24330        };
24331        let (cd, dc) = (conv_dim as i32, d_conv as i32);
24332        let __s_b = self.gpu.stream();
24333        let mut b = __s_b.launch_builder(&f);
24334        b.arg(qkv_cols)
24335            .arg(conv_state_ptrs)
24336            .arg(w)
24337            .arg(conv_outs)
24338            .arg(&cd)
24339            .arg(&dc);
24340        unsafe {
24341            b.launch(cfg)?;
24342        }
24343        Ok(())
24344    }
24345
24346    #[allow(clippy::too_many_arguments)]
24347    pub fn gdn_prep_decode_b_view(
24348        &self,
24349        conv_outs: &CudaSlice<f32>,
24350        beta_raws: &cudarc::driver::CudaView<f32>,
24351        alphas: &cudarc::driver::CudaView<f32>,
24352        dt_bias: &CudaSlice<f32>,
24353        a: &CudaSlice<f32>,
24354        q_l2: &mut CudaSlice<f32>,
24355        k_l2: &mut CudaSlice<f32>,
24356        v_g: &mut CudaSlice<f32>,
24357        beta: &mut CudaSlice<f32>,
24358        g_log: &mut CudaSlice<f32>,
24359        d_state: usize,
24360        num_v: usize,
24361        num_k: usize,
24362        key_dim: usize,
24363        eps: f32,
24364        conv_dim: usize,
24365        b_n: usize,
24366    ) -> Result<(), Box<dyn std::error::Error>> {
24367        let f = self.func("gdn_prep_decode_b_f32");
24368        let cfg = LaunchConfig {
24369            grid_dim: (num_v as u32, 1, b_n as u32),
24370            block_dim: (32, 4, 1),
24371            shared_mem_bytes: 0,
24372        };
24373        let (ds, nv, nk, kd, cd) = (
24374            d_state as i32,
24375            num_v as i32,
24376            num_k as i32,
24377            key_dim as i32,
24378            conv_dim as i32,
24379        );
24380        let __s_b = self.gpu.stream();
24381        let mut b = __s_b.launch_builder(&f);
24382        b.arg(conv_outs)
24383            .arg(beta_raws)
24384            .arg(alphas)
24385            .arg(dt_bias)
24386            .arg(a)
24387            .arg(q_l2)
24388            .arg(k_l2)
24389            .arg(v_g)
24390            .arg(beta)
24391            .arg(g_log)
24392            .arg(&ds)
24393            .arg(&nv)
24394            .arg(&nk)
24395            .arg(&kd)
24396            .arg(&eps)
24397            .arg(&cd);
24398        unsafe {
24399            b.launch(cfg)?;
24400        }
24401        Ok(())
24402    }
24403
24404    #[allow(clippy::too_many_arguments)]
24405    pub fn gdn_scan_s128_batched_view(
24406        &self,
24407        q: &CudaSlice<f32>,
24408        k: &CudaSlice<f32>,
24409        v: &CudaSlice<f32>,
24410        g: &CudaSlice<f32>,
24411        beta: &CudaSlice<f32>,
24412        state_in_ptrs: &cudarc::driver::CudaView<u64>,
24413        state_out_ptrs: &cudarc::driver::CudaView<u64>,
24414        o: &mut cudarc::driver::CudaViewMut<f32>,
24415        n_head: usize,
24416        b_n: usize,
24417        scale: f32,
24418    ) -> Result<(), Box<dyn std::error::Error>> {
24419        let f = self.func("gdn_scan_s128_b");
24420        const S_V: u32 = 128;
24421        const WARP: u32 = 32;
24422        const COLS_PER_BLOCK: u32 = 4;
24423        let cfg = LaunchConfig {
24424            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24425            block_dim: (WARP, COLS_PER_BLOCK, 1),
24426            shared_mem_bytes: 0,
24427        };
24428        let h = n_head as i32;
24429        let __s_b = self.gpu.stream();
24430        let mut b = __s_b.launch_builder(&f);
24431        b.arg(q)
24432            .arg(k)
24433            .arg(v)
24434            .arg(g)
24435            .arg(beta)
24436            .arg(state_in_ptrs)
24437            .arg(state_out_ptrs)
24438            .arg(o)
24439            .arg(&h)
24440            .arg(&scale);
24441        unsafe {
24442            b.launch(cfg)?;
24443        }
24444        Ok(())
24445    }
24446
24447    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
24448    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
24449    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
24450    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
24451    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
24452    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
24453    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
24454    /// identity law); prime_cache/forward/forward_last are the only callers.
24455    pub fn gdn_chunked_enabled() -> bool {
24456        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24457        *E.get_or_init(|| {
24458            std::env::var("MEMRA_GDN_CHUNKED")
24459                .map(|v| v != "0")
24460                .unwrap_or(true)
24461        })
24462    }
24463
24464    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
24465    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
24466    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
24467    /// of 32 in [32, 128] (kernel row mappings require it).
24468    pub fn gdn_chunk_size() -> usize {
24469        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24470        *C.get_or_init(|| {
24471            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
24472                .ok()
24473                .and_then(|v| v.parse().ok())
24474                .unwrap_or(32);
24475            c.clamp(32, 128) / 32 * 32
24476        })
24477    }
24478
24479    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
24480    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
24481    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
24482    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
24483    #[allow(clippy::too_many_arguments)]
24484    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
24485    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
24486    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
24487    #[allow(clippy::too_many_arguments)]
24488    pub fn gdn_chunk_k123(
24489        &self,
24490        q: &CudaSlice<f32>,
24491        k: &CudaSlice<f32>,
24492        v: &CudaSlice<f32>,
24493        g: &CudaSlice<f32>,
24494        beta: &CudaSlice<f32>,
24495        wb16: Option<&mut CudaSlice<u8>>,
24496        n_head: usize,
24497        t: usize,
24498        c: usize,
24499        hk: usize,
24500        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
24501    ) -> Result<
24502        (
24503            CudaSlice<f32>,
24504            CudaSlice<f32>,
24505            CudaSlice<f32>,
24506            CudaSlice<f32>,
24507        ),
24508        Box<dyn std::error::Error>,
24509    > {
24510        const D: usize = 128;
24511        let h = n_head;
24512        let nc = (t + c - 1) / c;
24513        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24514        let mut gcum = self.uninit(t * h)?;
24515        let mut a = self.uninit(nc * h * c * c)?;
24516        let mut p = self.uninit(nc * h * c * c)?;
24517        let mut u = self.uninit(nc * h * c * D)?;
24518        let mut w = self.uninit(nc * h * c * D)?;
24519        {
24520            // K1
24521            let f = self.func("gdn_chunk_cumgate_f32");
24522            let cfg = LaunchConfig {
24523                grid_dim: (nc as u32, h as u32, 1),
24524                block_dim: (32, 1, 1),
24525                shared_mem_bytes: 0,
24526            };
24527            let __s_b = self.gpu.stream();
24528            let mut b = __s_b.launch_builder(&f);
24529            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
24530            unsafe {
24531                b.launch(cfg)?;
24532            }
24533        }
24534        if let Some((qb, kb, pb)) = k2w {
24535            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
24536            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
24537            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
24538            let f = self.func("gdn_k2_wgmma");
24539            let cfg = LaunchConfig {
24540                grid_dim: (nc as u32, h as u32, 1),
24541                block_dim: (128, 1, 1),
24542                shared_mem_bytes: 0,
24543            };
24544            let hki = hk as i32;
24545            let __s_b = self.gpu.stream();
24546            let mut b = __s_b.launch_builder(&f);
24547            b.arg(qb)
24548                .arg(kb)
24549                .arg(&gcum)
24550                .arg(beta)
24551                .arg(&mut a)
24552                .arg(&mut *pb)
24553                .arg(&hi)
24554                .arg(&ti)
24555                .arg(&ci)
24556                .arg(&hki);
24557            unsafe {
24558                b.launch(cfg)?;
24559            }
24560        } else if c <= 64 && !portable_mma_gated() {
24561            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
24562            let f = self.func("gdn_chunk_attn_f32");
24563            f.set_attribute(
24564                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24565                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24566            )?;
24567            let jt = ((c + 31) / 32) as u32;
24568            let cfg = LaunchConfig {
24569                grid_dim: (nc as u32, h as u32, jt),
24570                block_dim: (256, 1, 1),
24571                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24572            };
24573            let hki = hk as i32;
24574            let __s_b = self.gpu.stream();
24575            let mut b = __s_b.launch_builder(&f);
24576            b.arg(q)
24577                .arg(k)
24578                .arg(&gcum)
24579                .arg(beta)
24580                .arg(&mut a)
24581                .arg(&mut p)
24582                .arg(&hi)
24583                .arg(&ti)
24584                .arg(&ci)
24585                .arg(&hki);
24586            unsafe {
24587                b.launch(cfg)?;
24588            }
24589        } else {
24590            // K2 generic (C = 128, or the portable target's low-smem fallback)
24591            assert!(
24592                hk == h,
24593                "generic K2 is broadcast-only (de-broadcast rides C==32)"
24594            );
24595            let f = self.func("gdn_chunk_attn_g_f32");
24596            let cfg = LaunchConfig {
24597                grid_dim: (nc as u32, h as u32, 1),
24598                block_dim: (32, 8, 1),
24599                shared_mem_bytes: 0,
24600            };
24601            let __s_b = self.gpu.stream();
24602            let mut b = __s_b.launch_builder(&f);
24603            b.arg(q)
24604                .arg(k)
24605                .arg(&gcum)
24606                .arg(beta)
24607                .arg(&mut a)
24608                .arg(&mut p)
24609                .arg(&hi)
24610                .arg(&ti)
24611                .arg(&ci);
24612            unsafe {
24613                b.launch(cfg)?;
24614            }
24615        }
24616        {
24617            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
24618            let cfg = LaunchConfig {
24619                grid_dim: (nc as u32, h as u32, 1),
24620                block_dim: (256, 1, 1),
24621                shared_mem_bytes: 0,
24622            };
24623            match c {
24624                32 | 64 => {
24625                    let f = self.func(if c == 32 {
24626                        "gdn_chunk_solve32_f32"
24627                    } else {
24628                        "gdn_chunk_solve64_f32"
24629                    });
24630                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
24631                    let wb: u64 = match wb16 {
24632                        Some(d) => self.addr_u8(d),
24633                        None => 0,
24634                    };
24635                    let hki = hk as i32;
24636                    let __s_b = self.gpu.stream();
24637                    let mut b = __s_b.launch_builder(&f);
24638                    b.arg(v)
24639                        .arg(k)
24640                        .arg(&a)
24641                        .arg(&gcum)
24642                        .arg(&mut u)
24643                        .arg(&mut w)
24644                        .arg(&wb)
24645                        .arg(&hi)
24646                        .arg(&ti)
24647                        .arg(&hki);
24648                    unsafe {
24649                        b.launch(cfg)?;
24650                    }
24651                }
24652                _ => {
24653                    assert!(hk == h, "generic K3 is broadcast-only");
24654                    let f = self.func("gdn_chunk_solve_f32");
24655                    let __s_b = self.gpu.stream();
24656                    let mut b = __s_b.launch_builder(&f);
24657                    b.arg(v)
24658                        .arg(k)
24659                        .arg(&a)
24660                        .arg(&gcum)
24661                        .arg(&mut u)
24662                        .arg(&mut w)
24663                        .arg(&hi)
24664                        .arg(&ti)
24665                        .arg(&ci);
24666                    unsafe {
24667                        b.launch(cfg)?;
24668                    }
24669                }
24670            }
24671        }
24672        Ok((gcum, p, u, w))
24673    }
24674
24675    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
24676    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
24677    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
24678    pub fn gdn_db_on() -> bool {
24679        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24680    }
24681
24682    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
24683    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
24684    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
24685    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
24686    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
24687    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
24688    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
24689    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
24690    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24691        !portable_mma_gated()
24692            && c == 32
24693            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24694                Ok("1") => true,
24695                Ok("0") => false,
24696                _ => gdn_mma_default_on(),
24697            }
24698    }
24699
24700    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
24701    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
24702    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
24703    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
24704    /// force would silently produce garbage. Required since the sm_120a mma default
24705    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
24706    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24707        cfg!(memra_hopper_mma)
24708            && self.gdn_mma_enabled(c)
24709            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24710    }
24711
24712    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
24713    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
24714    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
24715    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
24716    #[allow(clippy::too_many_arguments)]
24717    pub fn ssm_conv1d_gdn_state_pad(
24718        &self,
24719        qkv_tm: &cudarc::driver::CudaView<f32>,
24720        conv_state: &mut CudaSlice<f32>,
24721        w: &CudaSlice<f32>,
24722        q_g: &mut CudaSlice<f32>,
24723        k_g: &mut CudaSlice<f32>,
24724        v_g: &mut CudaSlice<f32>,
24725        conv_dim: usize,
24726        t: usize,
24727        d_conv: usize,
24728        d_state: usize,
24729        num_v: usize,
24730        num_k: usize,
24731        key_dim: usize,
24732        hk: usize,
24733        pad_len: Option<&CudaSlice<i32>>,
24734    ) -> Result<(), Box<dyn std::error::Error>> {
24735        assert!(
24736            t >= d_conv - 1,
24737            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24738        );
24739        {
24740            let f = self.func("ssm_conv1d_gdn_state_f32");
24741            let cfg = LaunchConfig {
24742                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24743                block_dim: (256, 1, 1),
24744                shared_mem_bytes: 0,
24745            };
24746            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24747            let (ds, nv, nk, kd, hki) = (
24748                d_state as i32,
24749                num_v as i32,
24750                num_k as i32,
24751                key_dim as i32,
24752                hk as i32,
24753            );
24754            let __s_b = self.gpu.stream();
24755            let mut b = __s_b.launch_builder(&f);
24756            b.arg(qkv_tm)
24757                .arg(&*conv_state)
24758                .arg(w)
24759                .arg(q_g)
24760                .arg(k_g)
24761                .arg(v_g)
24762                .arg(&cd)
24763                .arg(&ti)
24764                .arg(&dc)
24765                .arg(&ds)
24766                .arg(&nv)
24767                .arg(&nk)
24768                .arg(&kd)
24769                .arg(&hki);
24770            unsafe {
24771                b.launch(cfg)?;
24772            }
24773        }
24774        match pad_len {
24775            Some(len_d) => {
24776                let f = self.func("ssm_conv_ring_update_dev_f32");
24777                let n = conv_dim * (d_conv - 1);
24778                let cfg = LaunchConfig::for_num_elems(n as u32);
24779                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24780                let __s_b = self.gpu.stream();
24781                let mut b = __s_b.launch_builder(&f);
24782                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24783                unsafe {
24784                    b.launch(cfg)?;
24785                }
24786            }
24787            None => {
24788                let f = self.func("ssm_conv_ring_update_f32");
24789                let n = conv_dim * (d_conv - 1);
24790                let cfg = LaunchConfig::for_num_elems(n as u32);
24791                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24792                let __s_b = self.gpu.stream();
24793                let mut b = __s_b.launch_builder(&f);
24794                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24795                unsafe {
24796                    b.launch(cfg)?;
24797                }
24798            }
24799        }
24800        Ok(())
24801    }
24802
24803    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
24804    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
24805    /// K2/K3 can write them.
24806    pub fn gdn_chunk_alloc(
24807        &self,
24808        n_head: usize,
24809        t: usize,
24810        c: usize,
24811        hk: usize,
24812    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24813        const D: usize = 128;
24814        assert!(
24815            c == 32,
24816            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24817        );
24818        let h = n_head;
24819        let nc = (t + c - 1) / c;
24820        Ok(GdnChunkBufs {
24821            gcum: self.uninit(t * h)?,
24822            a: self.uninit(nc * h * c * c)?,
24823            p: self.uninit(nc * h * c * c)?,
24824            u: self.uninit(nc * h * c * D)?,
24825            w: self.uninit(nc * h * c * D)?,
24826            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24827            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24828            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24829            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24830            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24831            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24832            o: self.uninit(D * h * t)?,
24833            t,
24834            nc,
24835        })
24836    }
24837
24838    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
24839    pub fn f32_to_bf16_v(
24840        &self,
24841        x: &cudarc::driver::CudaView<f32>,
24842        dst: &mut CudaSlice<u8>,
24843        n: usize,
24844    ) -> Result<(), Box<dyn std::error::Error>> {
24845        let f = self.func("f32_to_bf16_bulk");
24846        let ni = n as i64;
24847        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24848        let __s_b = self.gpu.stream();
24849        let mut b = __s_b.launch_builder(&f);
24850        b.arg(x).arg(dst).arg(&ni);
24851        unsafe {
24852            b.launch(cfg)?;
24853        }
24854        Ok(())
24855    }
24856
24857    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
24858    pub fn f32_to_bf16_into(
24859        &self,
24860        x: &CudaSlice<f32>,
24861        dst: &mut CudaSlice<u8>,
24862        n: usize,
24863    ) -> Result<(), Box<dyn std::error::Error>> {
24864        let f = self.func("f32_to_bf16_bulk");
24865        let ni = n as i64;
24866        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24867        let __s_b = self.gpu.stream();
24868        let mut b = __s_b.launch_builder(&f);
24869        b.arg(x).arg(dst).arg(&ni);
24870        unsafe {
24871            b.launch(cfg)?;
24872        }
24873        Ok(())
24874    }
24875
24876    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
24877    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
24878    pub fn gdn_chunk_k123_vl8(
24879        &self,
24880        seqs: &[GdnSeqVl],
24881        n_head: usize,
24882        hk: usize,
24883        wq: Option<&GdnWVl8>,
24884    ) -> Result<(), Box<dyn std::error::Error>> {
24885        let b = seqs.len();
24886        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24887        let mut packed = [GdnSeqVl::default(); 8];
24888        packed[..b].copy_from_slice(seqs);
24889        let v = GdnVl8(packed);
24890        let (hi, ci) = (n_head as i32, 32i32);
24891        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24892        {
24893            let f = self.func("gdn_chunk_cumgate_vl");
24894            let cfg = LaunchConfig {
24895                grid_dim: (max_nc, n_head as u32, b as u32),
24896                block_dim: (32, 1, 1),
24897                shared_mem_bytes: 0,
24898            };
24899            let __s_lb = self.gpu.stream();
24900            let mut lb = __s_lb.launch_builder(&f);
24901            lb.arg(&v).arg(&hi).arg(&ci);
24902            unsafe {
24903                lb.launch(cfg)?;
24904            }
24905        }
24906        let hki = hk as i32;
24907        if let Some(w) = wq {
24908            // K2-wgmma vl twin (writes A + pre-masked Pb16)
24909            let f = self.func("gdn_k2_wgmma_vl");
24910            let cfg = LaunchConfig {
24911                grid_dim: (max_nc, n_head as u32, b as u32),
24912                block_dim: (128, 1, 1),
24913                shared_mem_bytes: 0,
24914            };
24915            let __s_lb = self.gpu.stream();
24916            let mut lb = __s_lb.launch_builder(&f);
24917            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24918            unsafe {
24919                lb.launch(cfg)?;
24920            }
24921        } else {
24922            let f = self.func("gdn_chunk_attn_vl");
24923            f.set_attribute(
24924                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24925                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24926            )?;
24927            let cfg = LaunchConfig {
24928                grid_dim: (max_nc, n_head as u32, b as u32),
24929                block_dim: (256, 1, 1),
24930                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24931            };
24932            let __s_lb = self.gpu.stream();
24933            let mut lb = __s_lb.launch_builder(&f);
24934            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24935            unsafe {
24936                lb.launch(cfg)?;
24937            }
24938        }
24939        {
24940            let f = self.func("gdn_chunk_solve32_vl");
24941            let cfg = LaunchConfig {
24942                grid_dim: (max_nc, n_head as u32, b as u32),
24943                block_dim: (256, 1, 1),
24944                shared_mem_bytes: 0,
24945            };
24946            let __s_lb = self.gpu.stream();
24947            let mut lb = __s_lb.launch_builder(&f);
24948            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24949            unsafe {
24950                lb.launch(cfg)?;
24951            }
24952        }
24953        Ok(())
24954    }
24955
24956    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
24957    /// fused gate-prep, 5 launches for every sequence (per-element math identical
24958    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
24959    #[allow(clippy::too_many_arguments)]
24960    pub fn gdn_prep_vl8(
24961        &self,
24962        seqs: &[GdnPrepVl],
24963        conv_w: &CudaSlice<f32>,
24964        dt_bias: &CudaSlice<f32>,
24965        a: &CudaSlice<f32>,
24966        conv_dim: usize,
24967        d_conv: usize,
24968        d_state: usize,
24969        num_v: usize,
24970        num_k: usize,
24971        key_dim: usize,
24972        hk: usize,
24973        eps: f32,
24974    ) -> Result<(), Box<dyn std::error::Error>> {
24975        let b = seqs.len();
24976        assert!(b >= 1 && b <= 8);
24977        let mut packed = [GdnPrepVl::default(); 8];
24978        packed[..b].copy_from_slice(seqs);
24979        let v = GdnPrepVl8(packed);
24980        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24981        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24982        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24983        assert!(
24984            conv_fuse || hk == num_v,
24985            "de-broadcast requires the fused conv"
24986        );
24987        if conv_fuse {
24988            let f = self.func("ssm_conv1d_gdn_state_vl");
24989            let cfg = LaunchConfig {
24990                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24991                block_dim: (256, 1, 1),
24992                shared_mem_bytes: 0,
24993            };
24994            let (dsi, nvi, nki, kdi, hki) = (
24995                d_state as i32,
24996                num_v as i32,
24997                num_k as i32,
24998                key_dim as i32,
24999                hk as i32,
25000            );
25001            let __s_lb = self.gpu.stream();
25002            let mut lb = __s_lb.launch_builder(&f);
25003            lb.arg(&v)
25004                .arg(conv_w)
25005                .arg(&cdi)
25006                .arg(&dci)
25007                .arg(&dsi)
25008                .arg(&nvi)
25009                .arg(&nki)
25010                .arg(&kdi)
25011                .arg(&hki);
25012            unsafe {
25013                lb.launch(cfg)?;
25014            }
25015        } else {
25016            let f = self.func("ssm_conv1d_tm_state_vl");
25017            let cfg = LaunchConfig {
25018                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25019                block_dim: (256, 1, 1),
25020                shared_mem_bytes: 0,
25021            };
25022            let __s_lb = self.gpu.stream();
25023            let mut lb = __s_lb.launch_builder(&f);
25024            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
25025            unsafe {
25026                lb.launch(cfg)?;
25027            }
25028        }
25029        {
25030            let f = self.func("ssm_conv_ring_update_vl");
25031            let n = (conv_dim * (d_conv - 1)) as u32;
25032            let cfg = LaunchConfig {
25033                grid_dim: (n.div_ceil(256), 1, b as u32),
25034                block_dim: (256, 1, 1),
25035                shared_mem_bytes: 0,
25036            };
25037            let __s_lb = self.gpu.stream();
25038            let mut lb = __s_lb.launch_builder(&f);
25039            lb.arg(&v).arg(&cdi).arg(&dci);
25040            unsafe {
25041                lb.launch(cfg)?;
25042            }
25043        }
25044        if !conv_fuse {
25045            let f = self.func("qkv_to_gdn_repack_vl");
25046            let n = max_t * (num_v * d_state) as u32;
25047            let cfg = LaunchConfig {
25048                grid_dim: (n.div_ceil(256), 1, b as u32),
25049                block_dim: (256, 1, 1),
25050                shared_mem_bytes: 0,
25051            };
25052            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25053            let __s_lb = self.gpu.stream();
25054            let mut lb = __s_lb.launch_builder(&f);
25055            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
25056            unsafe {
25057                lb.launch(cfg)?;
25058            }
25059        }
25060        if Self::l2_v2_on(d_state) {
25061            let f = self.func("gdn_l2_v2_vl");
25062            let cfg = LaunchConfig {
25063                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
25064                block_dim: (256, 1, 1),
25065                shared_mem_bytes: 0,
25066            };
25067            let (dsi, nvi) = (d_state as i32, hk as i32);
25068            let __s_lb = self.gpu.stream();
25069            let mut lb = __s_lb.launch_builder(&f);
25070            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25071            unsafe {
25072                lb.launch(cfg)?;
25073            }
25074        } else {
25075            let f = self.func("gdn_l2_vl");
25076            let cfg = LaunchConfig {
25077                grid_dim: (max_t * hk as u32, 2, b as u32),
25078                block_dim: (256, 1, 1),
25079                shared_mem_bytes: 0,
25080            };
25081            let (dsi, nvi) = (d_state as i32, hk as i32);
25082            let __s_lb = self.gpu.stream();
25083            let mut lb = __s_lb.launch_builder(&f);
25084            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25085            unsafe {
25086                lb.launch(cfg)?;
25087            }
25088        }
25089        {
25090            let f = self.func("gdn_gate_prep_vl");
25091            let n = max_t * num_v as u32;
25092            let cfg = LaunchConfig {
25093                grid_dim: (n.div_ceil(256), 1, b as u32),
25094                block_dim: (256, 1, 1),
25095                shared_mem_bytes: 0,
25096            };
25097            let nvi = num_v as i32;
25098            let __s_lb = self.gpu.stream();
25099            let mut lb = __s_lb.launch_builder(&f);
25100            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
25101            unsafe {
25102                lb.launch(cfg)?;
25103            }
25104        }
25105        Ok(())
25106    }
25107
25108    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
25109    pub fn gdn_mirror_vl8(
25110        &self,
25111        seqs: &[GdnSeqVl],
25112        n_head: usize,
25113        which: i32,
25114        hk: usize,
25115    ) -> Result<(), Box<dyn std::error::Error>> {
25116        let b = seqs.len();
25117        assert!(b >= 1 && b <= 8);
25118        let mut packed = [GdnSeqVl::default(); 8];
25119        packed[..b].copy_from_slice(seqs);
25120        let v = GdnVl8(packed);
25121        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
25122        let max_n = seqs
25123            .iter()
25124            .map(|s| {
25125                if which == 0 {
25126                    s.t as i64 * ept as i64
25127                } else {
25128                    s.nc as i64 * ept as i64 * 32
25129                }
25130            })
25131            .max()
25132            .unwrap();
25133        let f = self.func("gdn_mirror_vl");
25134        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25135        let cfg = LaunchConfig {
25136            grid_dim: (blocks, 1, b as u32),
25137            block_dim: (256, 1, 1),
25138            shared_mem_bytes: 0,
25139        };
25140        let __s_lb = self.gpu.stream();
25141        let mut lb = __s_lb.launch_builder(&f);
25142        lb.arg(&v).arg(&ept).arg(&which);
25143        unsafe {
25144            lb.launch(cfg)?;
25145        }
25146        Ok(())
25147    }
25148
25149    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
25150    pub fn gdn_tail_vl8(
25151        &self,
25152        seqs: &[GdnPrepVl],
25153        norm_w: &CudaSlice<f32>,
25154        d_state: usize,
25155        num_v: usize,
25156        eps: f32,
25157    ) -> Result<(), Box<dyn std::error::Error>> {
25158        let b = seqs.len();
25159        assert!(b >= 1 && b <= 8);
25160        let mut packed = [GdnPrepVl::default(); 8];
25161        packed[..b].copy_from_slice(seqs);
25162        let v = GdnPrepVl8(packed);
25163        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25164        let f = self.func("gated_rmsnorm_f16out_vl");
25165        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25166        let cfg = LaunchConfig {
25167            grid_dim: (max_t * num_v as u32, 1, b as u32),
25168            block_dim: (128, 1, 1),
25169            shared_mem_bytes: 0,
25170        };
25171        let (dsi, nvi) = (d_state as i32, num_v as i32);
25172        let __s_lb = self.gpu.stream();
25173        let mut lb = __s_lb.launch_builder(&f);
25174        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25175        unsafe {
25176            lb.launch(cfg)?;
25177        }
25178        Ok(())
25179    }
25180
25181    /// Raw device address helpers for the varlen by-value arg struct (single-stream
25182    /// launches; every buffer outlives the call — the f16 FFI discipline).
25183    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25184        use cudarc::driver::DevicePtr;
25185        let s = self.gpu.stream();
25186        let (p, _g) = x.device_ptr(&s);
25187        p as u64
25188    }
25189    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25190        use cudarc::driver::DevicePtrMut;
25191        let s = self.gpu.stream();
25192        let (p, _g) = x.device_ptr_mut(&s);
25193        p as u64
25194    }
25195    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25196        use cudarc::driver::DevicePtr;
25197        let s = self.gpu.stream();
25198        let (p, _g) = x.device_ptr(&s);
25199        p as u64
25200    }
25201    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25202        use cudarc::driver::DevicePtr;
25203        let s = self.gpu.stream();
25204        let (p, _g) = x.device_ptr(&s);
25205        p as u64
25206    }
25207
25208    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
25209    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
25210    /// launches, so this is strictly bit-gateable against them).
25211    pub fn gdn_chunk_vl8(
25212        &self,
25213        seqs: &[GdnSeqVl],
25214        n_head: usize,
25215        scale: f32,
25216        hk: usize,
25217        wq: Option<&GdnWVl8>,
25218    ) -> Result<(), Box<dyn std::error::Error>> {
25219        const NSPLIT: u32 = 4;
25220        let b = seqs.len();
25221        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25222        let mut packed = [GdnSeqVl::default(); 8];
25223        packed[..b].copy_from_slice(seqs);
25224        let v = GdnVl8(packed);
25225        let (hi, ci) = (n_head as i32, 32i32);
25226        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25227        let hki = hk as i32;
25228        if let Some(w) = wq {
25229            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
25230            let f = self.func("gdn_k45_wgmma_vl");
25231            let cfg = LaunchConfig {
25232                grid_dim: (n_head as u32, NSPLIT, b as u32),
25233                block_dim: (256, 1, 1),
25234                shared_mem_bytes: 0,
25235            };
25236            let __s_lb = self.gpu.stream();
25237            let mut lb = __s_lb.launch_builder(&f);
25238            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25239            unsafe {
25240                lb.launch(cfg)?;
25241            }
25242            let _ = max_nc;
25243            return Ok(());
25244        }
25245        {
25246            let f = self.func("gdn_chunk_state_mma_vl");
25247            let cfg = LaunchConfig {
25248                grid_dim: (n_head as u32, NSPLIT, b as u32),
25249                block_dim: (256, 1, 1),
25250                shared_mem_bytes: 0,
25251            };
25252            let __s_lb = self.gpu.stream();
25253            let mut lb = __s_lb.launch_builder(&f);
25254            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25255            unsafe {
25256                lb.launch(cfg)?;
25257            }
25258        }
25259        {
25260            let f = self.func("gdn_chunk_output_mma_vl");
25261            let cfg = LaunchConfig {
25262                grid_dim: (max_nc, n_head as u32, b as u32),
25263                block_dim: (256, 1, 1),
25264                shared_mem_bytes: 0,
25265            };
25266            let __s_lb = self.gpu.stream();
25267            let mut lb = __s_lb.launch_builder(&f);
25268            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25269            unsafe {
25270                lb.launch(cfg)?;
25271            }
25272        }
25273        Ok(())
25274    }
25275    pub fn gdn_scan_chunked(
25276        &self,
25277        q: &CudaSlice<f32>,
25278        k: &CudaSlice<f32>,
25279        v: &CudaSlice<f32>,
25280        g: &CudaSlice<f32>,
25281        beta: &CudaSlice<f32>,
25282        kb16_pre: Option<&CudaSlice<u8>>,
25283        qb16_pre: Option<&CudaSlice<u8>>,
25284        state_in: &CudaSlice<f32>,
25285        state_out: &mut CudaSlice<f32>,
25286        o: &mut CudaSlice<f32>,
25287        n_head: usize,
25288        t: usize,
25289        scale: f32,
25290        c: usize,
25291        hk: usize,
25292    ) -> Result<(), Box<dyn std::error::Error>> {
25293        const D: usize = 128;
25294        const NSPLIT: u32 = 4;
25295        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25296        let h = n_head;
25297        let nc = (t + c - 1) / c;
25298        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25299        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
25300        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
25301        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
25302        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
25303        let gdn_mma_pre = !portable_mma_gated()
25304            && c == 32
25305            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25306                Ok("1") => true,
25307                Ok("0") => false,
25308                _ => gdn_mma_default_on(),
25309            };
25310        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
25311            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
25312        } else {
25313            None
25314        };
25315        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
25316        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
25317        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
25318        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
25319        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
25320            && gdn_mma_pre
25321            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
25322        let nk = t * hk * D;
25323        let mut kb16_local: Option<CudaSlice<u8>> = None;
25324        if gdn_mma_pre && kb16_pre.is_none() {
25325            let mut kb = self.alloc_u8_uninit(nk * 2)?;
25326            let f = self.func("f32_to_bf16_bulk");
25327            let n2 = nk as i64;
25328            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25329            let __s_b = self.gpu.stream();
25330            let mut b = __s_b.launch_builder(&f);
25331            b.arg(k).arg(&mut kb).arg(&n2);
25332            unsafe {
25333                b.launch(cfg2)?;
25334            }
25335            kb16_local = Some(kb);
25336        }
25337        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
25338        if let Some(kb) = kb16_pre {
25339            assert!(kb.len() >= nk * 2, "kb16_pre too small");
25340        }
25341        let mut qb16: Option<CudaSlice<u8>> = None;
25342        let mut pb16: Option<CudaSlice<u8>> = None;
25343        if gdn_wgmma_pre {
25344            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
25345            // the standalone bulk cvt only serves callers without the prep mirror.
25346            if qb16_pre.is_none() {
25347                let mut qb = self.alloc_u8_uninit(nk * 2)?;
25348                let f = self.func("f32_to_bf16_bulk");
25349                let n2 = nk as i64;
25350                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25351                let __s_b = self.gpu.stream();
25352                let mut b = __s_b.launch_builder(&f);
25353                b.arg(q).arg(&mut qb).arg(&n2);
25354                unsafe {
25355                    b.launch(cfg2)?;
25356                }
25357                qb16 = Some(qb);
25358            } else if let Some(qb) = qb16_pre {
25359                assert!(qb.len() >= nk * 2, "qb16_pre too small");
25360            }
25361            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
25362        }
25363        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
25364        let k2w = if gdn_wgmma_pre {
25365            Some((
25366                *qb16_ref0.as_ref().unwrap(),
25367                *kb16_ref0.as_ref().unwrap(),
25368                pb16.as_mut().unwrap(),
25369            ))
25370        } else {
25371            None
25372        };
25373        let (gcum, p, u, w) =
25374            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
25375        let _ = &w;
25376        let mut y = self.uninit(nc * h * c * D)?;
25377        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
25378        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
25379        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
25380        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
25381        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
25382        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
25383        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
25384        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
25385        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
25386        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
25387        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
25388        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
25389        // sites must agree or the pre-work arms while the scan takes the scalar route.
25390        let gdn_mma = !portable_mma_gated()
25391            && c == 32
25392            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25393                Ok("1") => true,
25394                Ok("0") => false,
25395                _ => gdn_mma_default_on(),
25396            };
25397        if gdn_mma {
25398            let wb16 = wb16_pre
25399                .take()
25400                .expect("mma path pre-allocates wb16 (K3 store fold)");
25401            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
25402            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
25403            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
25404            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
25405            // pass runs inside the persistent-M kernel; Y and Ssnap are never
25406            // materialized. New numeric class (gk folds into k^T instead of ys) —
25407            // explicit opt-in until the state-carry battery promotes it. Env read per
25408            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
25409            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
25410            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
25411            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
25412            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
25413            if gdn_wgmma_pre {
25414                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
25415                let qb16 = qb16_ref0.unwrap();
25416                let pb16 = pb16.as_ref().unwrap();
25417                {
25418                    let f = self.func("gdn_k45_wgmma");
25419                    let cfg = LaunchConfig {
25420                        grid_dim: (h as u32, 4, 1),
25421                        block_dim: (256, 1, 1),
25422                        shared_mem_bytes: 0,
25423                    };
25424                    let hki = hk as i32;
25425                    let __s_b = self.gpu.stream();
25426                    let mut b = __s_b.launch_builder(&f);
25427                    b.arg(kb16_ref)
25428                        .arg(&gcum)
25429                        .arg(beta)
25430                        .arg(&u)
25431                        .arg(&wb16)
25432                        .arg(qb16)
25433                        .arg(pb16)
25434                        .arg(o)
25435                        .arg(&scale)
25436                        .arg(state_in)
25437                        .arg(&mut *state_out)
25438                        .arg(&hi)
25439                        .arg(&ti)
25440                        .arg(&ci)
25441                        .arg(&hki);
25442                    unsafe {
25443                        b.launch(cfg)?;
25444                    }
25445                }
25446                return Ok(());
25447            }
25448            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
25449            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
25450            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
25451            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
25452            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
25453            {
25454                let f = self.func("gdn_chunk_state_mma");
25455                let cfg = LaunchConfig {
25456                    grid_dim: (h as u32, NSPLIT, 1),
25457                    block_dim: (256, 1, 1),
25458                    shared_mem_bytes: 0,
25459                };
25460                let hki = hk as i32;
25461                let __s_b = self.gpu.stream();
25462                let mut b = __s_b.launch_builder(&f);
25463                b.arg(kb16_ref)
25464                    .arg(&gcum)
25465                    .arg(beta)
25466                    .arg(&u)
25467                    .arg(&wb16)
25468                    .arg(&mut y16)
25469                    .arg(&mut ssnap16)
25470                    .arg(state_in)
25471                    .arg(&mut *state_out)
25472                    .arg(&hi)
25473                    .arg(&ti)
25474                    .arg(&ci)
25475                    .arg(&hki);
25476                unsafe {
25477                    b.launch(cfg)?;
25478                }
25479            }
25480            {
25481                // K5-mma (bf16 St/Y consumers)
25482                let f = self.func("gdn_chunk_output_mma");
25483                let jt = ((c + 31) / 32) as u32;
25484                let cfg = LaunchConfig {
25485                    grid_dim: (nc as u32, h as u32, jt),
25486                    block_dim: (256, 1, 1),
25487                    shared_mem_bytes: 0,
25488                };
25489                let hki = hk as i32;
25490                let __s_b = self.gpu.stream();
25491                let mut b = __s_b.launch_builder(&f);
25492                b.arg(q)
25493                    .arg(&gcum)
25494                    .arg(&p)
25495                    .arg(&y16)
25496                    .arg(&ssnap16)
25497                    .arg(o)
25498                    .arg(&hi)
25499                    .arg(&ti)
25500                    .arg(&ci)
25501                    .arg(&scale)
25502                    .arg(&hki);
25503                unsafe {
25504                    b.launch(cfg)?;
25505                }
25506            }
25507            return Ok(());
25508        }
25509        {
25510            // K4 (sequential over chunks inside; blocks col-partition the state)
25511            let f = self.func("gdn_chunk_state_f32");
25512            let cfg = LaunchConfig {
25513                grid_dim: (h as u32, NSPLIT, 1),
25514                block_dim: (256, 1, 1),
25515                shared_mem_bytes: 0,
25516            };
25517            let __s_b = self.gpu.stream();
25518            let mut b = __s_b.launch_builder(&f);
25519            b.arg(k)
25520                .arg(&gcum)
25521                .arg(beta)
25522                .arg(&u)
25523                .arg(&w)
25524                .arg(&mut y)
25525                .arg(&mut ssnap)
25526                .arg(state_in)
25527                .arg(&mut *state_out)
25528                .arg(&hi)
25529                .arg(&ti)
25530                .arg(&ci);
25531            unsafe {
25532                b.launch(cfg)?;
25533            }
25534        }
25535        {
25536            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
25537            let f = self.func("gdn_chunk_output_f32");
25538            let jt = ((c + 31) / 32) as u32;
25539            let cfg = LaunchConfig {
25540                grid_dim: (nc as u32, h as u32, jt),
25541                block_dim: (256, 1, 1),
25542                shared_mem_bytes: 0,
25543            };
25544            let __s_b = self.gpu.stream();
25545            let mut b = __s_b.launch_builder(&f);
25546            b.arg(q)
25547                .arg(&gcum)
25548                .arg(&p)
25549                .arg(&y)
25550                .arg(&ssnap)
25551                .arg(o)
25552                .arg(&hi)
25553                .arg(&ti)
25554                .arg(&ci)
25555                .arg(&scale);
25556            unsafe {
25557                b.launch(cfg)?;
25558            }
25559        }
25560        Ok(())
25561    }
25562
25563    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
25564    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
25565    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
25566    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
25567    ///
25568    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
25569    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
25570    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
25571    #[allow(clippy::too_many_arguments)]
25572    #[allow(clippy::too_many_arguments)]
25573    pub fn gdn_scan_prefill(
25574        &self,
25575        q: &CudaSlice<f32>,
25576        k: &CudaSlice<f32>,
25577        v: &CudaSlice<f32>,
25578        g: &CudaSlice<f32>,
25579        beta: &CudaSlice<f32>,
25580        kb16_pre: Option<&CudaSlice<u8>>,
25581        qb16_pre: Option<&CudaSlice<u8>>,
25582        state_in: &CudaSlice<f32>,
25583        state_out: &mut CudaSlice<f32>,
25584        o: &mut CudaSlice<f32>,
25585        n_head: usize,
25586        t: usize,
25587        scale: f32,
25588        hk: usize,
25589    ) -> Result<(), Box<dyn std::error::Error>> {
25590        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
25591            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
25592            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25593        }
25594        if Self::gdn_chunked_enabled() && t >= 16 {
25595            self.gdn_scan_chunked(
25596                q,
25597                k,
25598                v,
25599                g,
25600                beta,
25601                kb16_pre,
25602                qb16_pre,
25603                state_in,
25604                state_out,
25605                o,
25606                n_head,
25607                t,
25608                scale,
25609                Self::gdn_chunk_size(),
25610                hk,
25611            )
25612        } else {
25613            assert!(
25614                hk == n_head,
25615                "s128 scan is broadcast-only (prep guarantees by predicate)"
25616            );
25617            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25618        }
25619    }
25620
25621    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
25622    #[allow(clippy::too_many_arguments)]
25623    fn gdn_scan_diff(
25624        &self,
25625        q: &CudaSlice<f32>,
25626        k: &CudaSlice<f32>,
25627        v: &CudaSlice<f32>,
25628        g: &CudaSlice<f32>,
25629        beta: &CudaSlice<f32>,
25630        state_in: &CudaSlice<f32>,
25631        state_out: &mut CudaSlice<f32>,
25632        o: &mut CudaSlice<f32>,
25633        n_head: usize,
25634        t: usize,
25635        scale: f32,
25636    ) -> Result<(), Box<dyn std::error::Error>> {
25637        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25638        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25639        let mut o_c = self.uninit(o.len())?;
25640        let mut st_c = self.uninit(state_out.len())?;
25641        self.gdn_scan_chunked(
25642            q,
25643            k,
25644            v,
25645            g,
25646            beta,
25647            None,
25648            None,
25649            state_in,
25650            &mut st_c,
25651            &mut o_c,
25652            n_head,
25653            t,
25654            scale,
25655            Self::gdn_chunk_size(),
25656            n_head,
25657        )?;
25658        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25659        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25660        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25661        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25662            let mut max_abs = 0f32;
25663            let mut max_rel = 0f32;
25664            let mut sum_rel = 0f64;
25665            for (x, y) in a.iter().zip(b) {
25666                let ad = (x - y).abs();
25667                let rel = ad / x.abs().max(y.abs()).max(1e-3);
25668                if ad > max_abs {
25669                    max_abs = ad;
25670                }
25671                if rel > max_rel {
25672                    max_rel = rel;
25673                }
25674                sum_rel += rel as f64;
25675            }
25676            (max_abs, max_rel, sum_rel / a.len() as f64)
25677        };
25678        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25679        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25680        println!(
25681            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25682                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25683            Self::gdn_chunk_size()
25684        );
25685        Ok(())
25686    }
25687
25688    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
25689    pub fn gdn_glog(
25690        &self,
25691        alpha: &CudaSlice<f32>,
25692        dt_bias: &CudaSlice<f32>,
25693        a: &CudaSlice<f32>,
25694        g_log: &mut CudaSlice<f32>,
25695        n_head: usize,
25696        t: usize,
25697    ) -> Result<(), Box<dyn std::error::Error>> {
25698        let f = self.func("gdn_glog_f32");
25699        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25700        let (h, ti) = (n_head as i32, t as i32);
25701        let __s_b = self.gpu.stream();
25702        let mut b = __s_b.launch_builder(&f);
25703        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25704        unsafe {
25705            b.launch(cfg)?;
25706        }
25707        Ok(())
25708    }
25709
25710    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
25711    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
25712    pub fn sigmoid_v(
25713        &self,
25714        x: &cudarc::driver::CudaView<f32>,
25715        y: &mut CudaSlice<f32>,
25716        n: usize,
25717    ) -> Result<(), Box<dyn std::error::Error>> {
25718        let f = self.func("sigmoid_f32");
25719        let cfg = LaunchConfig::for_num_elems(n as u32);
25720        let ni = n as i32;
25721        let __s_b = self.gpu.stream();
25722        let mut b = __s_b.launch_builder(&f);
25723        b.arg(x).arg(y).arg(&ni);
25724        unsafe {
25725            b.launch(cfg)?;
25726        }
25727        Ok(())
25728    }
25729
25730    pub fn gdn_glog_v(
25731        &self,
25732        alpha: &cudarc::driver::CudaView<f32>,
25733        dt_bias: &CudaSlice<f32>,
25734        a: &CudaSlice<f32>,
25735        g_log: &mut CudaSlice<f32>,
25736        n_head: usize,
25737        t: usize,
25738    ) -> Result<(), Box<dyn std::error::Error>> {
25739        let f = self.func("gdn_glog_f32");
25740        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25741        let (h, ti) = (n_head as i32, t as i32);
25742        let __s_b = self.gpu.stream();
25743        let mut b = __s_b.launch_builder(&f);
25744        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25745        unsafe {
25746            b.launch(cfg)?;
25747        }
25748        Ok(())
25749    }
25750
25751    pub fn sigmoid(
25752        &self,
25753        x: &CudaSlice<f32>,
25754        y: &mut CudaSlice<f32>,
25755        n: usize,
25756    ) -> Result<(), Box<dyn std::error::Error>> {
25757        let f = self.func("sigmoid_f32");
25758        let cfg = LaunchConfig::for_num_elems(n as u32);
25759        let ni = n as i32;
25760        let __s_b = self.gpu.stream();
25761        let mut b = __s_b.launch_builder(&f);
25762        b.arg(x).arg(y).arg(&ni);
25763        unsafe {
25764            b.launch(cfg)?;
25765        }
25766        Ok(())
25767    }
25768
25769    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
25770    /// (replaces sigmoid + mul + convert). Bit-identical class.
25771    pub fn sig_mul_f16out(
25772        &self,
25773        a: &CudaSlice<f32>,
25774        g: &CudaSlice<f32>,
25775        dst: &mut CudaSlice<f32>,
25776        dst16: &mut CudaSlice<u8>,
25777        n: usize,
25778    ) -> Result<(), Box<dyn std::error::Error>> {
25779        let f = self.func("sig_mul_f16out_f32");
25780        let cfg = LaunchConfig::for_num_elems(n as u32);
25781        let ni = n as i32;
25782        let __s_b = self.gpu.stream();
25783        let mut b = __s_b.launch_builder(&f);
25784        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25785        unsafe {
25786            b.launch(cfg)?;
25787        }
25788        Ok(())
25789    }
25790
25791    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
25792    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
25793    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
25794    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
25795    ///
25796    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
25797    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
25798    /// applies the wrong number of distinct gate values.
25799    #[allow(clippy::too_many_arguments)]
25800    pub fn attn_head_gate(
25801        &self,
25802        a: &CudaSlice<f32>,
25803        g: &CudaSlice<f32>,
25804        dst: &mut CudaSlice<f32>,
25805        dst16: Option<&mut CudaSlice<u8>>,
25806        head_dim: usize,
25807        n_head: usize,
25808        t: usize,
25809    ) -> Result<(), Box<dyn std::error::Error>> {
25810        let f = self.func("attn_head_gate_f32");
25811        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25812        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25813        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
25814        let d16: u64 = match dst16 {
25815            Some(d) => self.addr_u8(d),
25816            None => 0,
25817        };
25818        let __s_b = self.gpu.stream();
25819        let mut b = __s_b.launch_builder(&f);
25820        b.arg(a)
25821            .arg(g)
25822            .arg(dst)
25823            .arg(&d16)
25824            .arg(&hd)
25825            .arg(&nh)
25826            .arg(&ti);
25827        unsafe {
25828            b.launch(cfg)?;
25829        }
25830        Ok(())
25831    }
25832
25833    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
25834    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
25835    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
25836    ///
25837    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
25838    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
25839    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
25840    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
25841    #[allow(clippy::too_many_arguments)]
25842    pub fn swiglu_clamped_mul_scaled(
25843        &self,
25844        gate: &CudaSlice<f32>,
25845        up: &CudaSlice<f32>,
25846        gs: f32,
25847        us: f32,
25848        limit: f32,
25849        dst: &mut CudaSlice<f32>,
25850        n: usize,
25851    ) -> Result<(), Box<dyn std::error::Error>> {
25852        debug_assert!(
25853            limit > 1e-6,
25854            "swiglu_clamped needs a live limit; use silu_mul_scaled"
25855        );
25856        let f = self.func("swiglu_clamped_mul_scaled_f32");
25857        let cfg = LaunchConfig::for_num_elems(n as u32);
25858        let ni = n as i32;
25859        let __s_b = self.gpu.stream();
25860        let mut b = __s_b.launch_builder(&f);
25861        b.arg(gate)
25862            .arg(up)
25863            .arg(&gs)
25864            .arg(&us)
25865            .arg(&limit)
25866            .arg(dst)
25867            .arg(&ni);
25868        unsafe {
25869            b.launch(cfg)?;
25870        }
25871        Ok(())
25872    }
25873
25874    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
25875    pub fn gated_rmsnorm(
25876        &self,
25877        o: &CudaSlice<f32>,
25878        w: &CudaSlice<f32>,
25879        z: &CudaSlice<f32>,
25880        dst: &mut CudaSlice<f32>,
25881        ncols: usize,
25882        nrows: usize,
25883        eps: f32,
25884    ) -> Result<(), Box<dyn std::error::Error>> {
25885        let f = self.func("gated_rmsnorm_f32");
25886        let cfg = LaunchConfig {
25887            grid_dim: (nrows as u32, 1, 1),
25888            block_dim: (128, 1, 1),
25889            shared_mem_bytes: 0,
25890        };
25891        let (nc, e) = (ncols as i32, eps);
25892        let __s_b = self.gpu.stream();
25893        let mut b = __s_b.launch_builder(&f);
25894        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25895        unsafe {
25896            b.launch(cfg)?;
25897        }
25898        Ok(())
25899    }
25900
25901    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
25902    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
25903    pub fn gated_rmsnorm_f16out(
25904        &self,
25905        o: &CudaSlice<f32>,
25906        w: &CudaSlice<f32>,
25907        z: &CudaSlice<f32>,
25908        dst: &mut CudaSlice<f32>,
25909        dst16: &mut CudaSlice<u8>,
25910        ncols: usize,
25911        nrows: usize,
25912        eps: f32,
25913    ) -> Result<(), Box<dyn std::error::Error>> {
25914        let f = self.func("gated_rmsnorm_f16out_f32");
25915        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25916        let cfg = LaunchConfig {
25917            grid_dim: (nrows as u32, 1, 1),
25918            block_dim: (128, 1, 1),
25919            shared_mem_bytes: 0,
25920        };
25921        let (nc, e) = (ncols as i32, eps);
25922        let __s_b = self.gpu.stream();
25923        let mut b = __s_b.launch_builder(&f);
25924        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25925        unsafe {
25926            b.launch(cfg)?;
25927        }
25928        Ok(())
25929    }
25930
25931    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
25932    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
25933    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
25934    #[allow(clippy::too_many_arguments)]
25935    pub fn add_rms_norm_zq8(
25936        &self,
25937        a: &CudaSlice<f32>,
25938        b_in: &CudaSlice<f32>,
25939        w: &CudaSlice<f32>,
25940        res: &mut CudaSlice<f32>,
25941        z: &mut CudaSlice<f32>,
25942        ncols: usize,
25943        nrows: usize,
25944        eps: f32,
25945    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25946        assert!(ncols % 32 == 0);
25947        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25948        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25949        let f = self.func("add_rms_norm_zq8");
25950        let cfg = LaunchConfig {
25951            grid_dim: (nrows as u32, 1, 1),
25952            block_dim: (1024, 1, 1),
25953            shared_mem_bytes: 0,
25954        };
25955        let (nc, ep) = (ncols as i32, eps);
25956        let __s_b = self.gpu.stream();
25957        let mut b = __s_b.launch_builder(&f);
25958        b.arg(a)
25959            .arg(b_in)
25960            .arg(w)
25961            .arg(res)
25962            .arg(z)
25963            .arg(&mut q)
25964            .arg(&mut d)
25965            .arg(&nc)
25966            .arg(&ep);
25967        unsafe {
25968            b.launch(cfg)?;
25969        }
25970        Ok((q, d))
25971    }
25972
25973    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
25974    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
25975    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
25976    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
25977    pub fn gated_rmsnorm_zv(
25978        &self,
25979        o: &CudaSlice<f32>,
25980        w: &CudaSlice<f32>,
25981        z: &cudarc::driver::CudaView<f32>,
25982        dst: &mut CudaSlice<f32>,
25983        ncols: usize,
25984        nrows: usize,
25985        eps: f32,
25986    ) -> Result<(), Box<dyn std::error::Error>> {
25987        let f = self.func("gated_rmsnorm_f32");
25988        let cfg = LaunchConfig {
25989            grid_dim: (nrows as u32, 1, 1),
25990            block_dim: (128, 1, 1),
25991            shared_mem_bytes: 0,
25992        };
25993        let (nc, e) = (ncols as i32, eps);
25994        let __s_b = self.gpu.stream();
25995        let mut b = __s_b.launch_builder(&f);
25996        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25997        unsafe {
25998            b.launch(cfg)?;
25999        }
26000        Ok(())
26001    }
26002
26003    pub fn gated_rmsnorm_f16out_zv(
26004        &self,
26005        o: &CudaSlice<f32>,
26006        w: &CudaSlice<f32>,
26007        z: &cudarc::driver::CudaView<f32>,
26008        dst: &mut CudaSlice<f32>,
26009        dst16: &mut CudaSlice<u8>,
26010        ncols: usize,
26011        nrows: usize,
26012        eps: f32,
26013    ) -> Result<(), Box<dyn std::error::Error>> {
26014        let f = self.func("gated_rmsnorm_f16out_f32");
26015        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
26016        let cfg = LaunchConfig {
26017            grid_dim: (nrows as u32, 1, 1),
26018            block_dim: (128, 1, 1),
26019            shared_mem_bytes: 0,
26020        };
26021        let (nc, e) = (ncols as i32, eps);
26022        let __s_b = self.gpu.stream();
26023        let mut b = __s_b.launch_builder(&f);
26024        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26025        unsafe {
26026            b.launch(cfg)?;
26027        }
26028        Ok(())
26029    }
26030
26031    pub fn gated_rmsnorm_q8_1(
26032        &self,
26033        o: &CudaSlice<f32>,
26034        w: &CudaSlice<f32>,
26035        z: &CudaSlice<f32>,
26036        ncols: usize,
26037        nrows: usize,
26038        eps: f32,
26039    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26040        assert!(ncols % 32 == 0);
26041        let f = self.func("gated_rmsnorm_q8_1");
26042        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
26043        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26044        let cfg = LaunchConfig {
26045            grid_dim: (nrows as u32, 1, 1),
26046            block_dim: (128, 1, 1),
26047            shared_mem_bytes: 0,
26048        };
26049        let (nc, ep) = (ncols as i32, eps);
26050        let __s_b = self.gpu.stream();
26051        let mut b = __s_b.launch_builder(&f);
26052        b.arg(o)
26053            .arg(w)
26054            .arg(z)
26055            .arg(&mut out_q)
26056            .arg(&mut out_d)
26057            .arg(&nc)
26058            .arg(&ep);
26059        unsafe {
26060            b.launch(cfg)?;
26061        }
26062        Ok((out_q, out_d))
26063    }
26064
26065    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
26066    pub fn transpose(
26067        &self,
26068        inp: &CudaSlice<f32>,
26069        rows: usize,
26070        cols: usize,
26071    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26072        let f = self.func("transpose_f32");
26073        let mut out = self.zeros(rows * cols)?;
26074        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
26075        let (r, c) = (rows as i32, cols as i32);
26076        let __s_b = self.gpu.stream();
26077        let mut b = __s_b.launch_builder(&f);
26078        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
26079        unsafe {
26080            b.launch(cfg)?;
26081        }
26082        Ok(out)
26083    }
26084
26085    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
26086    pub fn repeat_heads(
26087        &self,
26088        inp: &CudaSlice<f32>,
26089        out: &mut CudaSlice<f32>,
26090        head_dim: usize,
26091        n_in: usize,
26092        n_out: usize,
26093        t: usize,
26094    ) -> Result<(), Box<dyn std::error::Error>> {
26095        let f = self.func("repeat_heads_f32");
26096        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
26097        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
26098        let __s_b = self.gpu.stream();
26099        let mut b = __s_b.launch_builder(&f);
26100        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
26101        unsafe {
26102            b.launch(cfg)?;
26103        }
26104        Ok(())
26105    }
26106
26107    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
26108    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
26109    ///
26110    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
26111    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
26112    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
26113    pub fn q_gate_split(
26114        &self,
26115        qf: &CudaSlice<f32>,
26116        q_out: &mut CudaSlice<f32>,
26117        gate_out: &mut CudaSlice<f32>,
26118        head_dim: usize,
26119        n_head: usize,
26120        t: usize,
26121    ) -> Result<(), Box<dyn std::error::Error>> {
26122        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
26123        let out_need = head_dim * n_head * t;
26124        if q_out.len() < out_need || gate_out.len() < out_need {
26125            return Err(format!(
26126                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
26127                q_out.len(),
26128                gate_out.len()
26129            )
26130            .into());
26131        }
26132        let f = self.func("q_gate_split_f32");
26133        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26134        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26135        let __s_b = self.gpu.stream();
26136        let mut b = __s_b.launch_builder(&f);
26137        b.arg(qf)
26138            .arg(q_out)
26139            .arg(gate_out)
26140            .arg(&hd)
26141            .arg(&nh)
26142            .arg(&ti);
26143        unsafe {
26144            b.launch(cfg)?;
26145        }
26146        Ok(())
26147    }
26148
26149    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
26150    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
26151    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
26152    pub fn qkv_to_gdn_repack(
26153        &self,
26154        conv_out: &CudaSlice<f32>,
26155        q_g: &mut CudaSlice<f32>,
26156        k_g: &mut CudaSlice<f32>,
26157        v_g: &mut CudaSlice<f32>,
26158        d_state: usize,
26159        num_v: usize,
26160        num_k: usize,
26161        key_dim: usize,
26162        t: usize,
26163    ) -> Result<(), Box<dyn std::error::Error>> {
26164        let f = self.func("qkv_to_gdn_repack_f32");
26165        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26166        let (ds, nv, nk, kd, ti) = (
26167            d_state as i32,
26168            num_v as i32,
26169            num_k as i32,
26170            key_dim as i32,
26171            t as i32,
26172        );
26173        let __s_b = self.gpu.stream();
26174        let mut b = __s_b.launch_builder(&f);
26175        b.arg(conv_out)
26176            .arg(q_g)
26177            .arg(k_g)
26178            .arg(v_g)
26179            .arg(&ds)
26180            .arg(&nv)
26181            .arg(&nk)
26182            .arg(&kd)
26183            .arg(&ti);
26184        unsafe {
26185            b.launch(cfg)?;
26186        }
26187        Ok(())
26188    }
26189
26190    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
26191    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
26192    pub fn conv_left_pad(
26193        &self,
26194        src: &CudaSlice<f32>,
26195        dst: &mut CudaSlice<f32>,
26196        conv_dim: usize,
26197        t: usize,
26198        pad: usize,
26199    ) -> Result<(), Box<dyn std::error::Error>> {
26200        let f = self.func("conv_left_pad_f32");
26201        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26202        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26203        let __s_b = self.gpu.stream();
26204        let mut b = __s_b.launch_builder(&f);
26205        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26206        unsafe {
26207            b.launch(cfg)?;
26208        }
26209        Ok(())
26210    }
26211
26212    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
26213    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
26214    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
26215    pub fn conv_assemble_and_roll(
26216        &self,
26217        qkv_col: &CudaSlice<f32>,
26218        conv_state: &mut CudaSlice<f32>,
26219        conv_in: &mut CudaSlice<f32>,
26220        conv_dim: usize,
26221        pad: usize,
26222    ) -> Result<(), Box<dyn std::error::Error>> {
26223        let f = self.func("conv_assemble_and_roll_f32");
26224        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26225        let (cd, p) = (conv_dim as i32, pad as i32);
26226        let __s_b = self.gpu.stream();
26227        let mut b = __s_b.launch_builder(&f);
26228        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26229        unsafe {
26230            b.launch(cfg)?;
26231        }
26232        Ok(())
26233    }
26234
26235    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
26236    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
26237    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
26238    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
26239    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
26240    pub fn ssm_conv1d_fused_decode(
26241        &self,
26242        qkv_col: &CudaSlice<f32>,
26243        conv_state: &mut CudaSlice<f32>,
26244        w: &CudaSlice<f32>,
26245        conv_out: &mut CudaSlice<f32>,
26246        conv_dim: usize,
26247        d_conv: usize,
26248    ) -> Result<(), Box<dyn std::error::Error>> {
26249        let f = self.func("ssm_conv1d_fused_decode_f32");
26250        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26251        let (cd, dc) = (conv_dim as i32, d_conv as i32);
26252        let __s_b = self.gpu.stream();
26253        let mut b = __s_b.launch_builder(&f);
26254        b.arg(qkv_col)
26255            .arg(conv_state)
26256            .arg(w)
26257            .arg(conv_out)
26258            .arg(&cd)
26259            .arg(&dc);
26260        unsafe {
26261            b.launch(cfg)?;
26262        }
26263        Ok(())
26264    }
26265
26266    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
26267    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
26268    pub fn slice_range(
26269        &self,
26270        src: &CudaSlice<f32>,
26271        start: usize,
26272        len: usize,
26273    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26274        let host = self.gpu.stream().clone_dtoh(src)?;
26275        self.gpu.stream().synchronize()?;
26276        Ok(self.htod(&host[start..start + len])?)
26277    }
26278}
26279
26280#[cfg(test)]
26281mod target_dispatch_tests {
26282    use super::legacy_quant_gemm_allowed;
26283
26284    #[test]
26285    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26286        // sm_120a native lane
26287        assert!(legacy_quant_gemm_allowed(false, false, false));
26288        assert!(!legacy_quant_gemm_allowed(false, false, true));
26289        // pure portable lane (sm_89): gated
26290        assert!(!legacy_quant_gemm_allowed(true, false, false));
26291        assert!(!legacy_quant_gemm_allowed(true, false, true));
26292        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
26293        assert!(legacy_quant_gemm_allowed(true, true, false));
26294        assert!(!legacy_quant_gemm_allowed(true, true, true));
26295    }
26296
26297    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26298    #[test]
26299    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26300        assert!(!legacy_quant_gemm_allowed(
26301            cfg!(memra_portable_cuda),
26302            cfg!(memra_hopper_mma),
26303            false
26304        ));
26305    }
26306
26307    #[cfg(memra_hopper_mma)]
26308    #[test]
26309    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
26310        assert!(legacy_quant_gemm_allowed(
26311            cfg!(memra_portable_cuda),
26312            cfg!(memra_hopper_mma),
26313            false
26314        ));
26315        assert!(super::portable_mma_gated() == false);
26316    }
26317}
26318
26319/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
26320/// inherent methods (inherent methods win name resolution, so no recursion).
26321impl memra_kv::KvDev for Engine {
26322    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26323        Engine::zeros(self, n)
26324    }
26325    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26326        Engine::uninit(self, n)
26327    }
26328    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
26329        Engine::alloc_u8(self, n)
26330    }
26331    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
26332        Engine::htod_i32(self, v)
26333    }
26334    fn clone_dtod(
26335        &self,
26336        src: &CudaSlice<f32>,
26337    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26338        Engine::clone_dtod(self, src)
26339    }
26340    fn copy_into(
26341        &self,
26342        dst: &mut CudaSlice<f32>,
26343        off: usize,
26344        src: &CudaSlice<f32>,
26345        len: usize,
26346    ) -> Result<(), Box<dyn std::error::Error>> {
26347        Engine::copy_into(self, dst, off, src, len)
26348    }
26349    fn set_i32_one(
26350        &self,
26351        d: &mut CudaSlice<i32>,
26352        v: i32,
26353    ) -> Result<(), Box<dyn std::error::Error>> {
26354        Engine::set_i32_one(self, d, v)
26355    }
26356}
26357
26358#[cfg(test)]
26359mod fused_gate_bounds_tests {
26360    use super::*;
26361
26362    /// The fused `[q|gate]` split's read-site guard, on the device.
26363    ///
26364    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
26365    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
26366    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
26367    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
26368    /// `FusedQGateExtent` before the launch.
26369    ///
26370    /// Catch demonstration for this test (guard temporarily removed, then restored):
26371    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
26372    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
26373    /// the call returns `Err`. Receipt in the lane report.
26374    #[test]
26375    #[ignore = "requires a CUDA GPU"]
26376    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
26377        let e = Engine::new(0).unwrap();
26378        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
26379        let fused = 2 * head_dim * n_head * t;
26380        let out_n = head_dim * n_head * t;
26381
26382        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
26383        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
26384        let mut q = e.uninit(out_n).unwrap();
26385        let mut gate = e.uninit(out_n).unwrap();
26386        let err = e
26387            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
26388            .expect_err("half-width wq must be refused, not read past")
26389            .to_string();
26390        assert!(err.contains("NO fused gate"), "{err}");
26391        assert!(err.contains(&format!("{fused}")), "{err}");
26392
26393        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
26394        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
26395        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
26396        let wide = e.htod(&host).unwrap();
26397        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
26398            .expect("full-width wq splits");
26399        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
26400        for tok in 0..t {
26401            for hh in 0..n_head {
26402                for d in 0..head_dim {
26403                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
26404                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
26405                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
26406                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
26407                }
26408            }
26409        }
26410
26411        // undersized destinations are refused too (the other half of the extent contract)
26412        let mut small = e.uninit(out_n - 1).unwrap();
26413        assert!(
26414            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
26415                .is_err()
26416        );
26417    }
26418}
26419
26420/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
26421/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
26422/// any launch, so the refusal is testable without a device.
26423#[cfg(test)]
26424mod fused_rope_width_tests {
26425    use super::Engine;
26426
26427    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
26428    /// safetensors route derives the same), which is why the fusion is legal there today.
26429    #[test]
26430    fn full_width_is_accepted() {
26431        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
26432        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
26433        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
26434    }
26435
26436    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
26437    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
26438    ///
26439    /// ```text
26440    /// attention.key_length     512   rope.dimension_count     512   (global class)
26441    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
26442    /// ```
26443    ///
26444    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
26445    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
26446    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
26447    /// instead of a silently over-rotated head.
26448    #[test]
26449    fn gemma4_official_artifact_widths_pass() {
26450        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
26451        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
26452    }
26453
26454    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
26455    /// with no `n_dims`, silently rotating the pass-through band.
26456    #[test]
26457    fn partial_rotary_is_refused_with_the_geometry_named() {
26458        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
26459        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
26460            .expect_err("partial rotary must refuse");
26461        let msg = err.to_string();
26462        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
26463        assert!(msg.contains("n_rot 64"), "{msg}");
26464        assert!(msg.contains("head_dim 256"), "{msg}");
26465        assert!(
26466            msg.contains("64..256"),
26467            "names the band it would corrupt: {msg}"
26468        );
26469        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
26470        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
26471        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
26472        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
26473    }
26474}