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/// The default dynamic-shared-memory launch bound the naive SDPA family lives under: past
14/// `T_kv * 4 > 48KB` (T_kv > 12288) the smem kernel cannot launch — the measured
15/// dspark/full-attn long-ctx crash class. `sdpa_naive` dispatches to the byte-identical
16/// gmem-scores twin above this line.
17const SDPA_NAIVE_SMEM_MAX: usize = 48 * 1024;
18
19/// Guard on the gmem twin's `n_head * T * T_kv * 4`-byte scores workspace. The shapes that
20/// legitimately hit the smem bound are tall-KV blocks (T <= draft block size), which land in
21/// the tens of MB; 1 GiB refuses a square T==T_kv misuse before it silently eats the card.
22const SDPA_NAIVE_GMEM_WS_MAX: usize = 1 << 30;
23
24#[cfg(debug_assertions)]
25pub(crate) fn debug_assert_tensor_stream_device<T>(
26    tensor: &CudaSlice<T>,
27    stream: &CudaStream,
28    site: &str,
29) {
30    let tensor_dev = tensor.ordinal();
31    let stream_dev = stream.context().ordinal();
32    assert_eq!(
33        tensor_dev, stream_dev,
34        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
35    );
36}
37
38fn ensure_tensor_stream_device<T>(
39    tensor: &impl DeviceSlice<T>,
40    stream: &CudaStream,
41    site: &str,
42) -> Result<(), Box<dyn std::error::Error>> {
43    let tensor_dev = tensor.stream().context().ordinal();
44    let stream_dev = stream.context().ordinal();
45    if tensor_dev != stream_dev {
46        return Err(format!(
47            "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
48             stream on dev{stream_dev}"
49        )
50        .into());
51    }
52    Ok(())
53}
54
55pub use memra_gguf;
56pub use memra_runtime;
57
58pub mod forward;
59pub mod hybrid;
60pub mod hybrid_forward;
61pub mod model;
62pub mod sigrouter_contract;
63pub mod vision;
64pub mod vision_gemma;
65pub mod vision_pre;
66/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
67/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
68pub mod cache {
69    pub use memra_kv::*;
70}
71pub mod decode;
72pub mod decode_batch;
73pub mod dflash;
74pub mod eagle;
75pub mod gemma_spec;
76pub mod graph_update;
77/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
78/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
79/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
80pub mod mla;
81pub mod moesd;
82pub mod parallel;
83pub mod plan_backend;
84pub mod pp;
85pub mod round_stream;
86pub mod spec;
87pub mod tp;
88pub use memra_sampling as sampler;
89
90/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
91/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
92/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
93/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
94/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
95///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
96///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
97///                     stream sync per projection (round-47 ledgered defect).
98///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
99///                     construction, zero syncs, f32 C with the act row-scale folded in.
100/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
101/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
102/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
103/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
104/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
105/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
106///
107/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
108/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
109/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
110/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
111/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
112/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
113/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
114/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
115///
116/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
117/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
118/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
119/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
120/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
121/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
122/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
123///
124/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
125/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
126/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
127/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
128/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
129/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
130/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
131/// the k-quant-only admission survives as the rollback seam, not the default.
132/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
133pub fn moe_f16g_mode() -> u8 {
134    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
135    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
136        Ok("0") => 0,
137        Ok("2") => 2,
138        Ok("3") => 3,
139        Ok(_) => 1,
140        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
141        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
142        Err(_) => 2,
143    })
144}
145/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
146/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
147/// (shape_sel, cross) for the FFI:
148///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
149///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
150///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
151///                         back to 32x64 in-launcher when the device/in_f can't take it).
152///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
153///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
154///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
155///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
156///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
157///                         verdict was stale).
158pub fn moe_f16g_sk_params() -> (i32, i32) {
159    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
160    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
161        Ok("0") => (-1, 0),
162        Ok("32") => (0, i32::MAX),
163        Ok("128") => (0, 1),
164        _ => {
165            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
166                .ok()
167                .and_then(|v| v.parse().ok())
168                .unwrap_or(64);
169            (0, cross)
170        }
171    })
172}
173/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
174/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
175/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
176/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
177/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
178/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
179/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
180/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
181/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
182/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
183pub fn moe_f16g_direct_on(qtype: i32) -> bool {
184    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
185    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
186        Ok("0") => 0,
187        Ok("kq") => 1,
188        _ => 2,
189    });
190    match m {
191        0 => false,
192        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
193        _ => true,
194    }
195}
196/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
197/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
198/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
199/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
200/// stage under q35's routing skew. Bit-identical to every other sk form by construction
201/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
202/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
203/// tail. in_f % 64 != 0 falls back in-launcher.
204pub fn moe_f16g_tail_on() -> bool {
205    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
207}
208
209/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
210/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
211/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
212/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
213/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
214/// still opens this door for A/B.
215pub fn moe_f16g_gemma_on() -> bool {
216    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
217    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
218}
219
220/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
221/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
222/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
223pub fn moe_fuse_actq_on() -> bool {
224    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
225    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
226}
227
228/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
229/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
230/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
231/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
232/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
233/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
234/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
235/// verify already use (dispatch parity, one router kernel for every t).
236/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
237pub fn router_prefill_exact_on() -> bool {
238    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
239    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
240}
241
242pub fn router_kernel_on() -> bool {
243    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244    *ON.get_or_init(|| {
245        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
246        if !on {
247            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
248        }
249        on
250    })
251}
252
253/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
254/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
255/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
256/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
257/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
258/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
259/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
260/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
261/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
262/// seam, perf-only: bits are equal by the kernel-check gate).
263/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
264/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
265/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
266/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
267pub const ROUTER_BATCH_MIN_T: usize = 8;
268pub fn router_batch_on() -> bool {
269    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
271}
272mod cpu_experts;
273#[cfg(memra_cutlass)]
274pub mod cutlass_ffi;
275pub mod dsv4_ffi;
276pub mod dsv4_gpu;
277pub mod f16_ffi;
278pub mod fp8_ffi;
279pub mod mmq_ffi;
280pub mod moe_cache;
281pub mod prime_graph;
282pub mod spill;
283mod spill_pread;
284
285// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
286// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
287// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
288// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
289// broke every machine that wasn't the build machine. Same bytes, same module image;
290// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
291const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
292const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
293const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
294const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
295const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
296const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
297/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
298const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
299
300/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
301/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
302/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
303/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
304/// compile-time default (zero behavior change).
305fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
306    assert!(
307        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
308        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
309    );
310    match std::env::var("MEMRA_GEMM_FATBIN") {
311        Ok(path) => std::borrow::Cow::Owned(
312            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
313        ),
314        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
315    }
316}
317
318/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
319/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
320/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
321/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
322/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
323/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
324pub(crate) const fn portable_mma_gated() -> bool {
325    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
326}
327
328/// Refuse an env force that would reach a kernel THIS BUILD DOES NOT CONTAIN.
329///
330/// Doors of the shape `MEMRA_X=1 => true` are arch-blind: they were written so an operator could
331/// force a promoted path on, and the default arm (`cfg!(memra_hopper_mma)` or similar) is the only
332/// thing that consulted the arch. On a portable build the forced path then reaches
333/// `Engine::func`, which resolves lazily and ends in `panic!("kernel {name} not in any fatbin")` —
334/// a confusing crash naming a kernel the operator never heard of, several frames from the switch
335/// they actually flipped.
336///
337/// Found 2026-08-23 by tools/fatbin-lookup-census.py, which listed 20 looked-up kernels absent
338/// from the sm_89 fatbins. 18 of those turned out to be correctly unreachable (the GDN varlen
339/// chain is gated through `gdn_mma_enabled`, which starts with `!portable_mma_gated()`); these
340/// env doors were the two that were genuinely reachable, and only by explicit operator action.
341///
342/// Same shape and same message style as `gemm_fatbin_bytes`'s refusal above — one idiom for
343/// "this switch cannot work on this build", so it fails at the switch instead of at the lookup.
344#[track_caller]
345pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
346    assert!(
347        !portable_mma_gated(),
348        "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
349         portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
350    );
351}
352
353/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
354/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
355/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
356/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
357/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
358/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
359/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
360/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
361pub(crate) const fn gdn_mma_default_on() -> bool {
362    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
363}
364
365/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
366const fn konst_eq(a: &str, b: &str) -> bool {
367    let (a, b) = (a.as_bytes(), b.as_bytes());
368    if a.len() != b.len() {
369        return false;
370    }
371    let mut i = 0;
372    while i < a.len() {
373        if a[i] != b[i] {
374            return false;
375        }
376        i += 1;
377    }
378    true
379}
380
381/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
382/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
383/// in a pure helper so the dispatch guard can be regression-tested without constructing an
384/// Engine or allocating a GPU tensor.
385const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
386    (!portable_cuda || hopper_mma) && !no_gemm
387}
388
389// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
390// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
391// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
392// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
393// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
394// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
395// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
396const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
397const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
398const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
399const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
400const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
401
402/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
403/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
404pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
405
406/// The flash_attn fatbin matching the selected KV formats.
407fn flash_fatbin_bytes() -> &'static [u8] {
408    match kv_cache_formats() {
409        ("q8_0", "q5_1") => FLASH_FATBIN,
410        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
411        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
412        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
413        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
414        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
415        other => unreachable!("kv_cache_formats returned {other:?}"),
416    }
417}
418
419/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
420/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
421/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
422/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
423/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
424/// defaults (zero behavior change).
425fn k1_launch_override() -> Option<(u32, u32, u32)> {
426    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
427    *K1.get_or_init(|| {
428        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
429        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
430        match p.as_slice() {
431            [bm, bn, w] => Some((*bm, *bn, *w)),
432            _ => None,
433        }
434    })
435}
436
437/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
438/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
439/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
440/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
441/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
442/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
443pub(crate) fn wgmma_gemm_enabled() -> bool {
444    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
445    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
446}
447
448/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
449/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
450/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
451/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
452/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
453/// the split count changes the combine's FP summation order, and the spec verify's batched forward
454/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
455/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
456/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
457/// adaptive retries (any retry MUST pass run-spec self-consistency first).
458/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
459/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
460/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
461/// between eager decode and the verify (the spec-exactness law).
462pub const FA_VEC_MIN_TKV: usize = 96;
463/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
464/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
465/// which moves the crossover — sweep per model, adopt per the battery.
466pub fn fa_vec_min_tkv() -> usize {
467    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
468    *V.get_or_init(|| {
469        std::env::var("MEMRA_FA_VEC_MIN")
470            .ok()
471            .and_then(|v| v.parse().ok())
472            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
473    })
474}
475
476/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
477/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
478/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
479///
480/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
481/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
482/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
483/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
484/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
485/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
486pub fn fa_f16pv_on() -> bool {
487    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
488    *ON.get_or_init(|| {
489        std::env::var("MEMRA_FA_F16PV")
490            .map(|v| v != "0")
491            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
492    })
493}
494
495/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
496/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
497/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
498pub fn fa512_hp_on() -> bool {
499    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
500    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
501}
502
503/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
504/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
505/// accumulation. Even n_head and even GQA group required (guarded per call).
506pub fn faw_hp_on() -> bool {
507    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
508    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
509}
510
511/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
512/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
513/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
514pub fn fa512_wide_warps() -> usize {
515    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
516    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
517        Ok("1") => 4,
518        _ => 2,
519    })
520}
521
522/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
523/// and the gemma global-layer rows/parity call sites.
524pub fn fa512_min_tkv() -> usize {
525    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
526    *FA512_MIN.get_or_init(|| {
527        std::env::var("MEMRA_FA512_MIN")
528            .ok()
529            .and_then(|v| v.parse().ok())
530            .unwrap_or(512)
531    })
532}
533/// Per-model crossover default, set at model load BEFORE the first decode (per-model
534/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
535/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
536pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
537    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
538/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
539/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
540/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
541pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
542/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
543/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
544/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
545/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
546/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
547pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
548    std::sync::atomic::AtomicBool::new(false);
549/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
550/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
551/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
552/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
553/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
554/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
555pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
556    std::sync::atomic::AtomicBool::new(true);
557pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
558    std::sync::atomic::AtomicUsize::new(16);
559/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
560/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
561/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
562/// latency-bound at 256 threads — 7us/launch measured).
563pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
564/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
565pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
566/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
567/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
568/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
569/// explicit numerical-form seam. mmq_ffi reads this before the env.
570pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
571/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
572/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
573pub use memra_kv::KV_FP8_FORCE;
574/// bf16 matvec family block size (MEMRA_MMV_BLOCK, default 128, clamped to [64, 256] and a
575/// multiple of 32 — the f32acc twin's shared reduce caps at 256). NUMERIC-CLASS knob: the
576/// per-thread stride and reduction order change with the block, same acceptance class as
577/// MEMRA_RMS_BLOCK (fresh-tape identity + battery at the pinned value).
578pub(crate) fn mmv_block() -> u32 {
579    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
580    *V.get_or_init(|| {
581        std::env::var("MEMRA_MMV_BLOCK")
582            .ok()
583            .and_then(|v| v.parse().ok())
584            .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
585            .unwrap_or(128)
586    })
587}
588
589/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
590/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
591/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
592/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
593pub(crate) fn sig_expf_dev_on() -> bool {
594    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
595    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
596}
597
598pub(crate) fn topk_fast_on() -> bool {
599    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
600    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
601}
602
603pub(crate) fn rms_block() -> u32 {
604    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
605    *V.get_or_init(|| {
606        std::env::var("MEMRA_RMS_BLOCK")
607            .ok()
608            .and_then(|v| v.parse().ok())
609            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
610    })
611}
612
613pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
614    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
615    if let Some(forced) = *S.get_or_init(|| {
616        std::env::var("MEMRA_FA_SPLIT")
617            .ok()
618            .and_then(|v| v.parse().ok())
619            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
620    }) {
621        return forced;
622    }
623    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
624    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
625    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
626    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
627    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
628    //
629    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
630    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
631    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
632    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
633    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
634    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
635    // rig-divergence law: this branch is measured on 188 SMs only).
636    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
637    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
638    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
639    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
640    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
641        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
642    {
643        return if t_kv <= 8192 {
644            16
645        } else if t_kv <= 16384 {
646            64
647        } else {
648            128
649        };
650    }
651    let big_rig = fa_sm_count() >= 128;
652    if big_rig {
653        let _ = n_head_kv;
654        if t_kv <= 2048 {
655            // MEMRA_FA_SP_SHORT=N: the SHORT rung only (the SWA layers' capped t_kv lands
656            // here on step37: 33 of 45 layers at t_kv=512). At 16 the tile loop runs
657            // HALF-EMPTY (FA_DEC_TILE=32 -> nt=16 per split), so the V staging pass moves a
658            // half tile per iteration and the combine carries 2x the partials; 32 makes each
659            // split exactly one full tile. A global MEMRA_FA_SPLIT cannot isolate this — it
660            // moves the deep-ctx rung too, where more splits measured worse.
661            // NUMERIC-CLASS door (key partition -> different per-split partials/combine):
662            // new tape + battery, exactly like every other split-ladder change.
663            static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
664            if let Some(sp) = *SHORT.get_or_init(|| {
665                std::env::var("MEMRA_FA_SP_SHORT")
666                    .ok()
667                    .and_then(|v| v.parse().ok())
668                    .filter(|&s: &usize| s >= 8 && s % 8 == 0)
669            }) {
670                return sp;
671            }
672            16
673        } else if t_kv <= 16384 {
674            64
675        } else {
676            128
677        }
678    } else if n_head_kv <= 4 {
679        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
680        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
681        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
682        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
683        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
684        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
685        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
686        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
687        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
688        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
689        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
690        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
691        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
692        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
693        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
694        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
695        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
696        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
697        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
698        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
699        if t_kv <= 512 {
700            8
701        } else if t_kv <= 16384 {
702            64
703        } else {
704            128
705        }
706    } else {
707        if t_kv <= 8192 {
708            32
709        } else if t_kv <= 16384 {
710            64
711        } else {
712            128
713        }
714    }
715}
716
717/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
718/// same attribute Engine::batched_variant reads).
719pub(crate) fn fa_sm_count() -> i32 {
720    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
721    *N.get_or_init(|| {
722        cudarc::driver::result::init().ok();
723        cudarc::driver::result::device::get(0)
724            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
725                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
726            .unwrap_or(82)
727    })
728}
729
730/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
731/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
732/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
733fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
734    match head_dim {
735        256 => Ok(""),
736        128 => Ok("_hd128"),
737        d => Err(format!(
738            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
739                          callers must gate to sdpa_naive"
740        )
741        .into()),
742    }
743}
744
745/// Quant type codes matching qmatvec.cu QType enum.
746pub const QT_Q8_0: i32 = 0;
747pub const QT_Q4_K: i32 = 1;
748pub const QT_Q6_K: i32 = 2;
749pub const QT_Q5_K: i32 = 3;
750pub const QT_Q3_K: i32 = 4;
751pub const QT_IQ4_XS: i32 = 5;
752pub const QT_IQ3_S: i32 = 6;
753pub const QT_NVFP4: i32 = 7;
754/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
755/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
756/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
757/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
758/// — ONE weight copy total, no Q8_0 re-encode duplicate.
759pub const QT_F8_E4M3: i32 = 10;
760/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
761/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
762pub const QT_NVFP4_RP: i32 = 9;
763/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
764pub const QT_F32: i32 = 8;
765pub const QT_BF16: i32 = 11;
766pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
767/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
768/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
769/// dp4a/MMQ implementation exists.
770pub const QT_Q2_K: i32 = 13;
771/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
772/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
773/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
774/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
775/// scalar `scale` field is 1.0 by the layout contract.
776///
777/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
778/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
779/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
780/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
781/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
782/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
783/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
784/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
785/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
786pub const QT_F8_E4M3_BLK: i32 = 14;
787
788/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
789pub struct Engine {
790    pub gpu: memra_runtime::Gpu,
791    module: Arc<CudaModule>,
792    hybrid: Arc<CudaModule>,
793    qmatvec: Arc<CudaModule>,
794    flash: Arc<CudaModule>,
795    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
796    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
797    /// Lazy: loaded on first global-format use; None until then.
798    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
799    gemm: Arc<CudaModule>,
800    router: Arc<CudaModule>,
801    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
802    sample: Arc<CudaModule>,
803    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
804    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
805    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
806    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
807    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
808    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
809    /// the single largest block. The cache still owns every address for its full lifetime.
810    moe_cache_layout: Mutex<Option<Vec<usize>>>,
811    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
812    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
813    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
814    /// verify between replays) reuse their addresses and the replay reads/writes live memory
815    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
816    capture_keep_on: std::sync::atomic::AtomicBool,
817    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
818    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
819    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
820    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
821    verify_exact: std::sync::atomic::AtomicBool,
822    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
823    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
824    pub copy_stream: Arc<CudaStream>,
825    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
826    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
827    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
828    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
829    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
830    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
831    #[cfg(memra_cutlass)]
832    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
833    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
834    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
835    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
836    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
837    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
838    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
839    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
840    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
841    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
842    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
843    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
844    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
845    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
846    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
847    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
848    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
849    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
850    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
851    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
852    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
853    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
854    /// before capture under the generate_graph tracking-off window so it carries no events).
855    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
856    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
857    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
858    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
859    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
860    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
861    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
862    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
863    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
864    router_stage: Mutex<Option<PinnedStage>>,
865}
866
867/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
868/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
869/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
870/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
871/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
872/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
873/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
874/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
875/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
876fn fa_v2_on() -> bool {
877    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
878    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
879    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
880    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
881    // + graph bit-identity green on all three models.
882    std::env::var("MEMRA_FA_V2")
883        .map(|v| v != "0")
884        .unwrap_or(true)
885}
886
887/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
888/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
889/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
890/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
891/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
892/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
893/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
894pub(crate) fn fa_v3_on() -> bool {
895    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
896    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
897    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
898    std::env::var("MEMRA_FA_V3")
899        .map(|v| v != "0")
900        .unwrap_or(true)
901}
902
903/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
904/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
905/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
906/// predicate so the twins can never diverge.
907fn fa_v4_mode() -> &'static str {
908    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
909    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
910}
911fn fa_v4_on() -> bool {
912    fa_v4_mode() != "0"
913} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
914/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
915/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
916/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
917/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
918/// stays kernel-family-identical to decode at the same t_kv.
919/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
920/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
921pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
922    std::sync::atomic::AtomicUsize::new(1024);
923pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
924    std::sync::atomic::AtomicUsize::new(usize::MAX);
925pub fn fa_v4_at_pub(t_kv: usize) -> bool {
926    fa_v4_at(t_kv)
927}
928fn fa_v4_at(t_kv: usize) -> bool {
929    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
930    let mx = *M.get_or_init(|| {
931        std::env::var("MEMRA_FA_V4_MAX")
932            .ok()
933            .and_then(|v| v.parse().ok())
934            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
935    });
936    fa_v4_on() && t_kv < mx
937}
938/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
939/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
940/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
941/// (same split partition, same softmax/accumulation order, same partials/combine) and only
942/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
943/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
944/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
945/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
946/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
947/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
948/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
949/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
950/// within one process (the v2/v3 pattern).
951pub const FA_DEEP_MIN_DEFAULT: usize = 0;
952fn fa_deep_at(t_kv: usize) -> bool {
953    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
954        return false;
955    }
956    let min = std::env::var("MEMRA_FA_DEEP_MIN")
957        .ok()
958        .and_then(|v| v.parse().ok())
959        .unwrap_or(FA_DEEP_MIN_DEFAULT);
960    t_kv >= min
961}
962/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
963pub fn fa_deep_at_pub(t_kv: usize) -> bool {
964    fa_deep_at(t_kv)
965}
966
967fn fa_v3_active(head_dim: usize) -> bool {
968    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
969    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
970    fa_v3_on()
971        && head_dim % 128 == 0
972        && kv_cache_formats() == ("q8_0", "q5_1")
973        && !Engine::kv_fp8_on()
974}
975
976/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
977/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
978/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
979/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
980/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
981/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
982/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
983pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
984    std::env::var("MEMRA_NO_FA_VEC").is_err()
985        && t_kv >= fa_vec_min_tkv()
986        && head_dim == 256
987        && fa_v4_at(t_kv)
988        && !matches!(fa_v4_mode(), "noB3" | "stage")
989        && !Engine::kv_fp8_on()
990}
991/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
992pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
993    fa_split_keys(t_kv, n_head_kv)
994}
995
996/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
997/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
998/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
999/// so we allocate through `result::malloc_host` with flags=0 directly.
1000struct PinnedStage {
1001    ptr: *mut u8,
1002    cap: usize,
1003}
1004unsafe impl Send for PinnedStage {}
1005impl PinnedStage {
1006    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1007        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1008        Ok(PinnedStage { ptr, cap })
1009    }
1010}
1011impl Drop for PinnedStage {
1012    fn drop(&mut self) {
1013        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1014    }
1015}
1016
1017/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
1018/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
1019pub const ARGMAX_NB: usize = 256;
1020
1021/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
1022pub(crate) use memra_fa3_vl as fa3_vl_raw;
1023
1024unsafe extern "C" {
1025    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
1026    fn memra_fa3_prefill(
1027        q16: *const core::ffi::c_void,
1028        k16: *const core::ffi::c_void,
1029        v16: *const core::ffi::c_void,
1030        o: *mut f32,
1031        t: i32,
1032        h: i32,
1033        hkv: i32,
1034        d: i32,
1035        scale: f32,
1036        stream: *mut core::ffi::c_void,
1037    ) -> i32;
1038    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
1039    pub(crate) fn memra_fa3_vl(
1040        q16s: *const *const core::ffi::c_void,
1041        k16s: *const *const core::ffi::c_void,
1042        v16s: *const *const core::ffi::c_void,
1043        os: *const *mut f32,
1044        ts: *const i32,
1045        b: i32,
1046        h: i32,
1047        hkv: i32,
1048        d: i32,
1049        scale: f32,
1050        stream: *mut core::ffi::c_void,
1051    ) -> i32;
1052}
1053
1054/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1055/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1056/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1057/// (slots are never re-allocated), so passing raw values is stable across the launch.
1058#[repr(C)]
1059#[derive(Clone, Copy)]
1060pub struct WPtr8(pub [u64; 8]);
1061unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1062
1063/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1064/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1065/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1066/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1067#[repr(C)]
1068#[derive(Clone, Copy, Default)]
1069pub struct GdnSeqVl {
1070    pub kb16: u64,
1071    pub gcum: u64,
1072    pub beta: u64,
1073    pub u: u64,
1074    pub wb16: u64,
1075    pub y: u64,
1076    pub ssnap: u64,
1077    pub state_in: u64,
1078    pub state_out: u64,
1079    pub q: u64,
1080    pub p: u64,
1081    pub o: u64,
1082    pub k: u64,
1083    pub v: u64,
1084    pub g: u64,
1085    pub a: u64,
1086    pub w: u64,
1087    pub t: i32,
1088    pub nc: i32,
1089}
1090unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1091#[repr(C)]
1092#[derive(Clone, Copy)]
1093pub struct GdnVl8(pub [GdnSeqVl; 8]);
1094unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1095
1096/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1097/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1098#[repr(C)]
1099#[derive(Clone, Copy, Default)]
1100pub struct GdnWVl {
1101    pub qb16: u64,
1102    pub pb16: u64,
1103}
1104unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1105#[repr(C)]
1106#[derive(Clone, Copy)]
1107pub struct GdnWVl8(pub [GdnWVl; 8]);
1108unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1109
1110/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1111#[repr(C)]
1112#[derive(Clone, Copy, Default)]
1113pub struct GdnPrepVl {
1114    pub qkv: u64,
1115    pub conv_state: u64,
1116    pub conv_out: u64,
1117    pub q_g: u64,
1118    pub k_g: u64,
1119    pub v_g: u64,
1120    pub q_l2: u64,
1121    pub k_l2: u64,
1122    pub beta_raw: u64,
1123    pub alpha: u64,
1124    pub beta: u64,
1125    pub g_log: u64,
1126    pub o: u64,
1127    pub z: u64,
1128    pub gn: u64,
1129    pub gn16: u64,
1130    pub kb16: u64,
1131    pub qb16: u64,
1132    pub t: i32,
1133    pub pad: i32,
1134}
1135unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1136#[repr(C)]
1137#[derive(Clone, Copy)]
1138pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1139unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1140
1141/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1142#[repr(C)]
1143#[derive(Clone, Copy, Default)]
1144pub struct FaSeqVl {
1145    pub q: u64,
1146    pub k16: u64,
1147    pub v16: u64,
1148    pub o: u64,
1149    pub kf: u64,
1150    pub vf: u64,
1151    pub t: i32,
1152    pub pad: i32,
1153}
1154unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1155#[repr(C)]
1156#[derive(Clone, Copy)]
1157pub struct FaVl8(pub [FaSeqVl; 8]);
1158unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1159
1160/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1161#[repr(C)]
1162#[derive(Clone, Copy, Default)]
1163pub struct AttnPreVl {
1164    pub qf: u64,
1165    pub kf: u64,
1166    pub vf: u64,
1167    pub q: u64,
1168    pub gate: u64,
1169    pub qn: u64,
1170    pub kn: u64,
1171    pub kc: u64,
1172    pub vc: u64,
1173    pub t: i32,
1174    pub pad: i32,
1175}
1176unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1177#[repr(C)]
1178#[derive(Clone, Copy)]
1179pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1180unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1181
1182/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1183/// varlen K1-K5 chain fills them).
1184pub struct GdnChunkBufs {
1185    pub gcum: CudaSlice<f32>,
1186    pub a: CudaSlice<f32>,
1187    pub p: CudaSlice<f32>,
1188    pub u: CudaSlice<f32>,
1189    pub w: CudaSlice<f32>,
1190    pub kb16: CudaSlice<u8>,
1191    pub wb16: CudaSlice<u8>,
1192    pub y16: CudaSlice<u8>,
1193    pub ssnap16: CudaSlice<u8>,
1194    pub qb16: CudaSlice<u8>,
1195    pub pb16: CudaSlice<u8>,
1196    pub o: CudaSlice<f32>,
1197    pub t: usize,
1198    pub nc: usize,
1199}
1200
1201/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1202#[repr(C)]
1203#[derive(Clone, Copy)]
1204pub struct F32x8(pub [f32; 8]);
1205unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1206
1207/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1208/// process. Bench binaries read it right after the call to print gen-only throughput without the
1209/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1210pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1211
1212impl Engine {
1213    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1214        let gpu = memra_runtime::Gpu::new(ordinal)?;
1215        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1216        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1217        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1218        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1219            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1220            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1221                .and_then(|d| unsafe {
1222                    Ok((
1223                        cudarc::driver::result::device::get_attribute(
1224                            d,
1225                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1226                        )?,
1227                        cudarc::driver::result::device::get_attribute(
1228                            d,
1229                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1230                        )?,
1231                    ))
1232                })
1233                .unwrap_or((0, 0));
1234            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1235            let ok = matches!(
1236                (built, maj, min),
1237                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1238            );
1239            if !ok {
1240                return Err(format!(
1241                    "memra was built for sm_{built} but device {ordinal} reports compute \
1242                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1243                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1244                )
1245                .into());
1246            }
1247        }
1248        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1249        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1250        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1251        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1252        unsafe {
1253            use cudarc::driver::sys;
1254            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1255            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1256            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1257                let mut thresh: u64 = u64::MAX;
1258                let _ = sys::cuMemPoolSetAttribute(
1259                    pool,
1260                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1261                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1262                );
1263            }
1264        }
1265        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1266        let hybrid = gpu
1267            .ctx
1268            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1269        let qmatvec = gpu
1270            .ctx
1271            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1272        let flash = gpu
1273            .ctx
1274            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1275        let gemm = gpu
1276            .ctx
1277            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1278        let router = gpu
1279            .ctx
1280            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1281        let sample = gpu
1282            .ctx
1283            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1284        let copy_stream = gpu.ctx.new_stream()?;
1285        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1286        // cudarc is in multi-stream mode (main stream +
1287        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1288        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1289        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1290        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
1291        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1292        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1293        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1294        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1295        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1296        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1297        // implicit event tracking.
1298        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1299        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1300        if std::env::var("MEMRA_EVT")
1301            .map(|v| v == "1")
1302            .unwrap_or(false)
1303        {
1304            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1305        } else {
1306            unsafe {
1307                gpu.ctx.disable_event_tracking();
1308            }
1309        }
1310        Ok(Self {
1311            gpu,
1312            module,
1313            hybrid,
1314            qmatvec,
1315            flash,
1316            flash_g: std::sync::OnceLock::new(),
1317            gemm,
1318            router,
1319            sample,
1320            moe_cache: Mutex::new(None),
1321            moe_cache_layout: Mutex::new(None),
1322            copy_stream,
1323            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1324            verify_exact: std::sync::atomic::AtomicBool::new(false),
1325            capture_keep: Mutex::new(Vec::new()),
1326            argmax_partials: Mutex::new(None),
1327            prime_deqw_ws: Mutex::new(None),
1328            router_stage: Mutex::new(None),
1329            fp8_scratch: Mutex::new(None),
1330            fa_vf16_scratch: Mutex::new(None),
1331            fa_part_pool: Mutex::new(None),
1332            fa_part_retired: Mutex::new(Vec::new()),
1333            fn_cache: Mutex::new(Default::default()),
1334            f16_scratch: Mutex::new(None),
1335            #[cfg(memra_cutlass)]
1336            cutlass_scratch: Mutex::new(None),
1337        })
1338    }
1339
1340    pub fn ctx(&self) -> &Arc<CudaContext> {
1341        &self.gpu.ctx
1342    }
1343
1344    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1345    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1346    ///
1347    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1348    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1349    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1350    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1351    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1352    ///
1353    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1354    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1355    /// under-count headroom does not belong in a gate that queues real work, but the honest
1356    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1357    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1358    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1359    ///
1360    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1361    pub fn pool_cached_bytes(&self) -> usize {
1362        let (reserved, used) = self.pool_reserved_used();
1363        reserved.saturating_sub(used)
1364    }
1365
1366    /// Bytes the driver's per-device CUDA GRAPH memory pool currently holds RESERVED
1367    /// (`cuDeviceGetGraphMemAttribute` RESERVED_MEM_CURRENT) — the backing store of every
1368    /// captured alloc node, which on this engine means the dspark verify-graph pool
1369    /// (decode/step graphs bake pre-allocated buffers and own no alloc nodes). This memory
1370    /// is DISTINCT from the async pool above: `mem_get_info`'s `free` already excludes it,
1371    /// it is never released back (the vgraph pool has no eviction by design), and it GROWS
1372    /// as new (segment, vt)/(vt, rung, hi) keys capture — the growth is what
1373    /// `dspark_vg_admission_debt` charges at admission. Returns 0 if the attribute cannot
1374    /// be queried (never a false headroom claim, matching `pool_cached_bytes`).
1375    pub fn device_graph_mem_reserved(&self) -> usize {
1376        use cudarc::driver::sys as cus;
1377        let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1378            return 0;
1379        };
1380        let mut bytes: u64 = 0;
1381        let rc = unsafe {
1382            cus::cuDeviceGetGraphMemAttribute(
1383                dev,
1384                cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1385                &mut bytes as *mut u64 as *mut std::ffi::c_void,
1386            )
1387        };
1388        if rc == cus::cudaError_enum::CUDA_SUCCESS {
1389            bytes as usize
1390        } else {
1391            0
1392        }
1393    }
1394
1395    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1396    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1397    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1398    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1399    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1400    /// (0, 0) if the pool cannot be queried.
1401    pub fn pool_reserved_used(&self) -> (usize, usize) {
1402        use cudarc::driver::sys;
1403        unsafe {
1404            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1405            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1406                != sys::CUresult::CUDA_SUCCESS
1407            {
1408                return (0, 0);
1409            }
1410            let (mut reserved, mut used) = (0u64, 0u64);
1411            if sys::cuMemPoolGetAttribute(
1412                pool,
1413                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1414                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1415            ) != sys::CUresult::CUDA_SUCCESS
1416            {
1417                return (0, 0);
1418            }
1419            if sys::cuMemPoolGetAttribute(
1420                pool,
1421                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1422                &mut used as *mut u64 as *mut core::ffi::c_void,
1423            ) != sys::CUresult::CUDA_SUCCESS
1424            {
1425                return (0, 0);
1426            }
1427            (reserved as usize, used as usize)
1428        }
1429    }
1430
1431    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1432    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1433    pub fn stream(&self) -> Arc<CudaStream> {
1434        self.gpu.stream()
1435    }
1436    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1437    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1438    pub fn gkv_on() -> bool {
1439        memra_kv::gkv_on()
1440    }
1441
1442    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1443    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1444    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1445    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1446    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1447    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1448    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1449    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1450    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1451    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1452    /// ON for both — no acceptance cost measured.
1453    pub fn wkv_on() -> bool {
1454        memra_kv::wkv_on()
1455    }
1456
1457    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1458    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1459    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1460    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1461    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1462    pub fn kv_fp8_on() -> bool {
1463        memra_kv::kv_fp8_on()
1464    }
1465
1466    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1467    /// when the fp8-globals arm is on; everything else from the default flash module.
1468    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1469        if head_dim == 512 && Self::gkv_on() {
1470            self.func_g(name)
1471        } else {
1472            self.func(name)
1473        }
1474    }
1475
1476    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1477    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1478    /// per-format fatbins; fall back to the base modules for those.
1479    fn func_g(&self, name: &str) -> CudaFunction {
1480        let m = self.flash_g.get_or_init(|| {
1481            self.gpu
1482                .ctx
1483                .load_module(cudarc::nvrtc::Ptx::from_binary(
1484                    FLASH_FATBIN_KF8VF8.to_vec(),
1485                ))
1486                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1487        });
1488        let key = format!("g:{name}");
1489        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1490            return f.clone();
1491        }
1492        let f = match m.load_function(name) {
1493            Ok(f) => f,
1494            Err(_) => self.func(name),
1495        };
1496        self.fn_cache.lock().unwrap().insert(key, f.clone());
1497        f
1498    }
1499
1500    fn func(&self, name: &str) -> CudaFunction {
1501        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1502        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1503        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1504            return f.clone();
1505        }
1506        let f = self
1507            .module
1508            .load_function(name)
1509            .or_else(|_| self.hybrid.load_function(name))
1510            .or_else(|_| self.qmatvec.load_function(name))
1511            .or_else(|_| self.flash.load_function(name))
1512            .or_else(|_| self.gemm.load_function(name))
1513            .or_else(|_| self.router.load_function(name))
1514            .or_else(|_| self.sample.load_function(name))
1515            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1516        self.fn_cache
1517            .lock()
1518            .unwrap()
1519            .insert(name.to_string(), f.clone());
1520        f
1521    }
1522
1523    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1524    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1525    pub fn scatter_trim_logits(
1526        &self,
1527        src: &CudaSlice<f32>,
1528        d2t: &CudaSlice<u32>,
1529        dst: &mut CudaSlice<f32>,
1530        d_vocab: usize,
1531        n_vocab: usize,
1532    ) -> Result<(), Box<dyn std::error::Error>> {
1533        let f1 = self.func("scatter_trim_logits_f32");
1534        let f2 = self.func("scatter_trim_logits_pass2_f32");
1535        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1536        let cfg1 = LaunchConfig {
1537            grid_dim: (256, 1, 1),
1538            block_dim: (256, 1, 1),
1539            shared_mem_bytes: 0,
1540        };
1541        let __s_b1 = self.gpu.stream();
1542        let mut b1 = __s_b1.launch_builder(&f1);
1543        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1544        unsafe {
1545            b1.launch(cfg1)?;
1546        }
1547        let cfg2 = LaunchConfig {
1548            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1549            block_dim: (256, 1, 1),
1550            shared_mem_bytes: 0,
1551        };
1552        let __s_b2 = self.gpu.stream();
1553        let mut b2 = __s_b2.launch_builder(&f2);
1554        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1555        unsafe {
1556            b2.launch(cfg2)?;
1557        }
1558        Ok(())
1559    }
1560
1561    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1562    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1563
1564    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1565    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1566    #[allow(clippy::too_many_arguments)]
1567    pub fn filter_stats(
1568        &self,
1569        x: &CudaSlice<f32>,
1570        row_stride: usize,
1571        rows: &CudaSlice<i32>,
1572        out_th: &mut CudaSlice<f32>,
1573        out_z: &mut CudaSlice<f32>,
1574        out_max: &mut CudaSlice<f32>,
1575        n: usize,
1576        nrow: usize,
1577        temp: f32,
1578        top_k: i32,
1579        top_p: f32,
1580        min_p: f32,
1581    ) -> Result<(), Box<dyn std::error::Error>> {
1582        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1583        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1584        // L2-resident, so the extra passes are near-free while the per-thread selection list
1585        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1586        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1587        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1588        //
1589        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1590        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1591        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1592        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1593        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1594        // Admission: cooperative grid must co-reside (16*nrow blocks vs SM count).
1595        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1596        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1597        let coop_on =
1598            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1599        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1600        if coop_on && 16 * nrow <= self.sm_count() as usize {
1601            let f = self.func("filter_stats_coop_f32");
1602            let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1603            let cfg = LaunchConfig {
1604                grid_dim: (16, nrow as u32, 1),
1605                block_dim: (512, 1, 1),
1606                shared_mem_bytes: 0,
1607            };
1608            let __s_b = self.gpu.stream();
1609            let mut b = __s_b.launch_builder(&f);
1610            b.arg(x)
1611                .arg(&rs)
1612                .arg(rows)
1613                .arg(&mut *out_th)
1614                .arg(&mut *out_z)
1615                .arg(&mut *out_max)
1616                .arg(&mut ws)
1617                .arg(&ni)
1618                .arg(&nr)
1619                .arg(&temp)
1620                .arg(&top_k)
1621                .arg(&top_p)
1622                .arg(&min_p);
1623            unsafe {
1624                b.launch_cooperative(cfg)?;
1625            }
1626            return Ok(());
1627        }
1628        let f = self.func("filter_stats_f32");
1629        let cfg = LaunchConfig {
1630            grid_dim: (nrow as u32, 1, 1),
1631            block_dim: (1024, 1, 1),
1632            shared_mem_bytes: 0,
1633        };
1634        let __s_b = self.gpu.stream();
1635        let mut b = __s_b.launch_builder(&f);
1636        b.arg(x)
1637            .arg(&rs)
1638            .arg(rows)
1639            .arg(&mut *out_th)
1640            .arg(&mut *out_z)
1641            .arg(&mut *out_max)
1642            .arg(&ni)
1643            .arg(&nr)
1644            .arg(&temp)
1645            .arg(&top_k)
1646            .arg(&top_p)
1647            .arg(&min_p);
1648        unsafe {
1649            b.launch(cfg)?;
1650        }
1651        Ok(())
1652    }
1653
1654    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1655    #[allow(clippy::too_many_arguments)]
1656    pub fn softmax_gather_filtered(
1657        &self,
1658        x: &CudaSlice<f32>,
1659        row_stride: usize,
1660        ids: &CudaSlice<u32>,
1661        rows: &CudaSlice<i32>,
1662        th: &CudaSlice<f32>,
1663        z: &CudaSlice<f32>,
1664        out: &mut CudaSlice<f32>,
1665        n: usize,
1666        npair: usize,
1667        temp: f32,
1668    ) -> Result<(), Box<dyn std::error::Error>> {
1669        let f = self.func("softmax_gather_filtered_f32");
1670        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1671        let cfg = LaunchConfig {
1672            grid_dim: (npair as u32, 1, 1),
1673            block_dim: (256, 1, 1),
1674            shared_mem_bytes: 0,
1675        };
1676        let __s_b = self.gpu.stream();
1677        let mut b = __s_b.launch_builder(&f);
1678        b.arg(x)
1679            .arg(&rs)
1680            .arg(ids)
1681            .arg(rows)
1682            .arg(th)
1683            .arg(z)
1684            .arg(&mut *out)
1685            .arg(&ni)
1686            .arg(&np)
1687            .arg(&temp);
1688        unsafe {
1689            b.launch(cfg)?;
1690        }
1691        Ok(())
1692    }
1693
1694    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1695    #[allow(clippy::too_many_arguments)]
1696    pub fn residual_sample_filtered(
1697        &self,
1698        p: &CudaSlice<f32>,
1699        q: Option<&CudaSlice<f32>>,
1700        n: usize,
1701        temp: f32,
1702        seed: u64,
1703        stream_pos: u32,
1704        p_stats: (f32, f32, f32),
1705        q_stats: (f32, f32, f32),
1706        out_tok: &mut CudaSlice<u32>,
1707    ) -> Result<(), Box<dyn std::error::Error>> {
1708        let f = self.func("residual_sample_filtered_f32");
1709        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1710        let has_q: i32 = q.is_some() as i32;
1711        let qbuf = q.unwrap_or(p);
1712        let (pm, pth, pz) = p_stats;
1713        let (qm, qth, qz) = q_stats;
1714        let cfg = LaunchConfig {
1715            grid_dim: (1, 1, 1),
1716            block_dim: (1024, 1, 1),
1717            shared_mem_bytes: 0,
1718        };
1719        let __s_b = self.gpu.stream();
1720        let mut b = __s_b.launch_builder(&f);
1721        b.arg(p)
1722            .arg(qbuf)
1723            .arg(&has_q)
1724            .arg(&ni)
1725            .arg(&temp)
1726            .arg(&slo)
1727            .arg(&shi)
1728            .arg(&stream_pos)
1729            .arg(&pm)
1730            .arg(&pth)
1731            .arg(&pz)
1732            .arg(&qm)
1733            .arg(&qth)
1734            .arg(&qz)
1735            .arg(&mut *out_tok);
1736        unsafe {
1737            b.launch(cfg)?;
1738        }
1739        Ok(())
1740    }
1741
1742    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1743    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1744    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1745    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1746    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1747    #[allow(clippy::too_many_arguments)]
1748    pub fn residual_sample_sparse_q(
1749        &self,
1750        p: &CudaSlice<f32>,
1751        cand_ids: &CudaSlice<u32>,
1752        q_probs: &CudaSlice<f32>,
1753        n_cand: usize,
1754        n: usize,
1755        temp: f32,
1756        seed: u64,
1757        stream_pos: u32,
1758        p_stats: (f32, f32, f32),
1759        out_tok: &mut CudaSlice<u32>,
1760    ) -> Result<(), Box<dyn std::error::Error>> {
1761        assert!(
1762            n_cand >= 1 && n_cand <= 32,
1763            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1764        );
1765        let f = self.func("residual_sample_sparse_q_f32");
1766        let (ni, nc) = (n as i32, n_cand as i32);
1767        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1768        let (pm, pth, pz) = p_stats;
1769        let cfg = LaunchConfig {
1770            grid_dim: (1, 1, 1),
1771            block_dim: (1024, 1, 1),
1772            shared_mem_bytes: 0,
1773        };
1774        let __s_b = self.gpu.stream();
1775        let mut b = __s_b.launch_builder(&f);
1776        b.arg(p)
1777            .arg(cand_ids)
1778            .arg(q_probs)
1779            .arg(&nc)
1780            .arg(&ni)
1781            .arg(&temp)
1782            .arg(&slo)
1783            .arg(&shi)
1784            .arg(&stream_pos)
1785            .arg(&pm)
1786            .arg(&pth)
1787            .arg(&pz)
1788            .arg(&mut *out_tok);
1789        unsafe {
1790            b.launch(cfg)?;
1791        }
1792        Ok(())
1793    }
1794
1795    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1796    #[allow(clippy::too_many_arguments)]
1797    pub fn gumbel_perturb_filtered(
1798        &self,
1799        x: &CudaSlice<f32>,
1800        y: &mut CudaSlice<f32>,
1801        n: usize,
1802        seed: u64,
1803        stream_pos: u32,
1804        temp: f32,
1805        row_max: f32,
1806        th: f32,
1807    ) -> Result<(), Box<dyn std::error::Error>> {
1808        let f = self.func("gumbel_perturb_filtered_f32");
1809        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1810        let cfg = LaunchConfig {
1811            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1812            block_dim: (256, 1, 1),
1813            shared_mem_bytes: 0,
1814        };
1815        let __s_b = self.gpu.stream();
1816        let mut b = __s_b.launch_builder(&f);
1817        b.arg(x)
1818            .arg(&mut *y)
1819            .arg(&ni)
1820            .arg(&slo)
1821            .arg(&shi)
1822            .arg(&stream_pos)
1823            .arg(&temp)
1824            .arg(&row_max)
1825            .arg(&th);
1826        unsafe {
1827            b.launch(cfg)?;
1828        }
1829        Ok(())
1830    }
1831
1832    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1833    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1834    /// filtered rejection sampling exact for the penalized target.
1835    #[allow(clippy::too_many_arguments)]
1836    pub fn penalize_logits(
1837        &self,
1838        x: &mut CudaSlice<f32>,
1839        hist: &CudaSlice<u32>,
1840        n_hist: usize,
1841        rep: f32,
1842        freq: f32,
1843        present: f32,
1844        n: usize,
1845    ) -> Result<(), Box<dyn std::error::Error>> {
1846        if n_hist == 0 {
1847            return Ok(());
1848        }
1849        let f = self.func("penalize_logits_f32");
1850        let (nh, ni) = (n_hist as i32, n as i32);
1851        let cfg = LaunchConfig {
1852            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1853            block_dim: (128, 1, 1),
1854            shared_mem_bytes: 0,
1855        };
1856        let __s_b = self.gpu.stream();
1857        let mut b = __s_b.launch_builder(&f);
1858        b.arg(&mut *x)
1859            .arg(hist)
1860            .arg(&nh)
1861            .arg(&rep)
1862            .arg(&freq)
1863            .arg(&present)
1864            .arg(&ni);
1865        unsafe {
1866            b.launch(cfg)?;
1867        }
1868        Ok(())
1869    }
1870
1871    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1872    #[allow(clippy::too_many_arguments)]
1873    pub fn penalize_logits_rows(
1874        &self,
1875        x: &mut CudaSlice<f32>,
1876        hist: &CudaSlice<u32>,
1877        n_hist: usize,
1878        rep: f32,
1879        freq: f32,
1880        present: f32,
1881        n: usize,
1882        nrow: usize,
1883    ) -> Result<(), Box<dyn std::error::Error>> {
1884        if n_hist == 0 || nrow == 0 {
1885            return Ok(());
1886        }
1887        let f = self.func("penalize_logits_rows_f32");
1888        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1889        let cfg = LaunchConfig {
1890            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1891            block_dim: (128, 1, 1),
1892            shared_mem_bytes: 0,
1893        };
1894        let __s_b = self.gpu.stream();
1895        let mut b = __s_b.launch_builder(&f);
1896        b.arg(&mut *x)
1897            .arg(hist)
1898            .arg(&nh)
1899            .arg(&rep)
1900            .arg(&freq)
1901            .arg(&present)
1902            .arg(&ni)
1903            .arg(&nr);
1904        unsafe {
1905            b.launch(cfg)?;
1906        }
1907        Ok(())
1908    }
1909
1910    /// Heterogeneous serving-batch penalties over host-maintained sparse window counts.
1911    /// `offsets[r]..offsets[r+1]` indexes the unique positive-count `(id,count)` entries for logits row
1912    /// `rows[r]`; each row may carry independent repetition/frequency/presence coefficients.
1913    /// One thread owns one distinct logit, so the kernel needs neither atomics nor the
1914    /// history-squared dedup scan used by the speculative raw-history oracle.
1915    #[allow(clippy::too_many_arguments)]
1916    pub fn penalize_logits_sparse_rows(
1917        &self,
1918        x: &mut CudaSlice<f32>,
1919        ids: &[u32],
1920        counts: &[u32],
1921        offsets: &[i32],
1922        rows: &[i32],
1923        reps: &[f32],
1924        freqs: &[f32],
1925        presents: &[f32],
1926        n: usize,
1927    ) -> Result<(), Box<dyn std::error::Error>> {
1928        let nrow = rows.len();
1929        if nrow == 0 {
1930            return Ok(());
1931        }
1932        let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
1933        let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
1934        let entry_count =
1935            i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
1936        if ids.len() != counts.len()
1937            || offsets.len() != nrow + 1
1938            || reps.len() != nrow
1939            || freqs.len() != nrow
1940            || presents.len() != nrow
1941            || offsets.first().copied() != Some(0)
1942            || offsets.last().copied() != Some(entry_count)
1943        {
1944            return Err("sparse penalty row metadata shape mismatch".into());
1945        }
1946        if counts.contains(&0) {
1947            return Err("sparse penalty counts must be positive".into());
1948        }
1949        let mut max_len = 0usize;
1950        for pair in offsets.windows(2) {
1951            if pair[0] < 0 || pair[1] < pair[0] {
1952                return Err("sparse penalty offsets must be monotonic".into());
1953            }
1954            max_len = max_len.max((pair[1] - pair[0]) as usize);
1955        }
1956        if max_len == 0 {
1957            return Ok(());
1958        }
1959
1960        let mut seen = std::collections::HashSet::with_capacity(ids.len());
1961        for (r, &row) in rows.iter().enumerate() {
1962            if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
1963                return Err("sparse penalty row index exceeds logits shape".into());
1964            }
1965            let begin = offsets[r] as usize;
1966            let end = offsets[r + 1] as usize;
1967            for &id in &ids[begin..end] {
1968                if id as usize >= n {
1969                    return Err("sparse penalty token id exceeds logits row".into());
1970                }
1971                if !seen.insert((row, id)) {
1972                    return Err("sparse penalty entries must be unique per logits row".into());
1973                }
1974            }
1975        }
1976
1977        // SAFETY: the checks above establish every invariant of the launch-only helper.
1978        unsafe {
1979            self.penalize_logits_sparse_rows_unchecked(
1980                x, ids, counts, offsets, rows, reps, freqs, presents, n,
1981            )
1982        }
1983    }
1984
1985    /// Launch-only form for the serving hot path, whose `HashMap`-backed producer already
1986    /// guarantees unique ids and whose rows are enumerated from the live batch.
1987    ///
1988    /// # Safety
1989    ///
1990    /// Shapes must match the safe wrapper, offsets must be monotonic and in bounds, every row
1991    /// must index `x`, and each `(row,id)` pair must occur at most once. Token ids outside the
1992    /// logits row are safe no-ops because the kernel bounds-checks them before computing `x`.
1993    #[allow(clippy::too_many_arguments)]
1994    pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
1995        &self,
1996        x: &mut CudaSlice<f32>,
1997        ids: &[u32],
1998        counts: &[u32],
1999        offsets: &[i32],
2000        rows: &[i32],
2001        reps: &[f32],
2002        freqs: &[f32],
2003        presents: &[f32],
2004        n: usize,
2005    ) -> Result<(), Box<dyn std::error::Error>> {
2006        let nrow = rows.len();
2007        if nrow == 0 {
2008            return Ok(());
2009        }
2010        let max_len = offsets
2011            .windows(2)
2012            .map(|pair| (pair[1] - pair[0]) as usize)
2013            .max()
2014            .unwrap_or(0);
2015        if max_len == 0 {
2016            return Ok(());
2017        }
2018        let ids_d = self.htod_u32_v(ids)?;
2019        let counts_d = self.htod_u32_v(counts)?;
2020        let offsets_d = self.htod_i32(offsets)?;
2021        let rows_d = self.htod_i32(rows)?;
2022        let reps_d = self.htod(reps)?;
2023        let freqs_d = self.htod(freqs)?;
2024        let presents_d = self.htod(presents)?;
2025        let f = self.func("penalize_logits_sparse_rows_f32");
2026        let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2027        let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2028        let cfg = LaunchConfig {
2029            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2030            block_dim: (128, 1, 1),
2031            shared_mem_bytes: 0,
2032        };
2033        let __s_b = self.gpu.stream();
2034        let mut b = __s_b.launch_builder(&f);
2035        b.arg(&mut *x)
2036            .arg(&ids_d)
2037            .arg(&counts_d)
2038            .arg(&offsets_d)
2039            .arg(&rows_d)
2040            .arg(&reps_d)
2041            .arg(&freqs_d)
2042            .arg(&presents_d)
2043            .arg(&ni)
2044            .arg(&nr);
2045        unsafe {
2046            b.launch(cfg)?;
2047        }
2048        Ok(())
2049    }
2050
2051    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
2052    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
2053    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
2054    /// is the within-round evolving penalty state block drafting needs: verify row r's
2055    /// target is penalized by every token committed before it INCLUDING same-round
2056    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
2057    /// approximation this exists to replace on the dspark route.
2058    #[allow(clippy::too_many_arguments)]
2059    pub fn penalize_logits_rows_inc(
2060        &self,
2061        x: &mut CudaSlice<f32>,
2062        hist: &CudaSlice<u32>,
2063        n_hist0: usize,
2064        rep: f32,
2065        freq: f32,
2066        present: f32,
2067        n: usize,
2068        nrow: usize,
2069        win: usize,
2070    ) -> Result<(), Box<dyn std::error::Error>> {
2071        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2072            return Ok(());
2073        }
2074        debug_assert!(
2075            hist.len() >= n_hist0 + nrow - 1,
2076            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2077        );
2078        let f = self.func("penalize_logits_rows_inc_f32");
2079        let max_len = win.min(n_hist0 + nrow - 1).max(1);
2080        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2081        let cfg = LaunchConfig {
2082            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2083            block_dim: (128, 1, 1),
2084            shared_mem_bytes: 0,
2085        };
2086        let __s_b = self.gpu.stream();
2087        let mut b = __s_b.launch_builder(&f);
2088        b.arg(&mut *x)
2089            .arg(hist)
2090            .arg(&nh)
2091            .arg(&rep)
2092            .arg(&freq)
2093            .arg(&present)
2094            .arg(&ni)
2095            .arg(&nr)
2096            .arg(&wi);
2097        unsafe {
2098            b.launch(cfg)?;
2099        }
2100        Ok(())
2101    }
2102
2103    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
2104    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
2105    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
2106    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
2107    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
2108    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
2109    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
2110    pub fn wpf_level() -> u32 {
2111        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2112        *ON.get_or_init(|| {
2113            std::env::var("MEMRA_WPF")
2114                .ok()
2115                .and_then(|v| v.parse().ok())
2116                .unwrap_or(1)
2117        })
2118    }
2119
2120    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
2121    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
2122    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
2123    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
2124    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
2125    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
2126    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
2127    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
2128    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
2129    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
2130    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
2131    pub fn set_verify_exact(&self, on: bool) {
2132        self.verify_exact
2133            .store(on, std::sync::atomic::Ordering::Relaxed);
2134    }
2135    pub(crate) fn verify_exact_on(&self) -> bool {
2136        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2137    }
2138
2139    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
2140    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
2141    pub fn qkv_append_on() -> bool {
2142        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2143        *ON.get_or_init(|| {
2144            std::env::var("MEMRA_QKV_APPEND")
2145                .map(|v| v != "0")
2146                .unwrap_or(true)
2147        })
2148    }
2149
2150    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
2151    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
2152    pub fn pdl_wb_on() -> bool {
2153        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2154        *ON.get_or_init(|| {
2155            std::env::var("MEMRA_PDL_WB")
2156                .map(|v| v != "0")
2157                .unwrap_or(true)
2158        })
2159    }
2160
2161    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
2162    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
2163    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
2164    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
2165    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
2166    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
2167    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
2168    pub fn norm_ilp_on() -> bool {
2169        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2170        *ON.get_or_init(|| {
2171            std::env::var("MEMRA_NORM_ILP")
2172                .map(|v| v != "0")
2173                .unwrap_or(true)
2174        })
2175    }
2176
2177    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
2178    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
2179    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
2180    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
2181    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
2182    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
2183    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
2184    pub fn tk_ffn_dual_on() -> bool {
2185        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2186        *ON.get_or_init(|| {
2187            std::env::var("MEMRA_TK_FFN_DUAL")
2188                .map(|v| v != "0")
2189                .unwrap_or(true)
2190        })
2191    }
2192
2193    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
2194    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
2195    /// per-model no-harm bisect knob.
2196    pub fn pdl_mmvq_on() -> bool {
2197        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2198        *ON.get_or_init(|| {
2199            std::env::var("MEMRA_PDL_MMVQ")
2200                .map(|v| v != "0")
2201                .unwrap_or(true)
2202        })
2203    }
2204
2205    pub fn pdl_on() -> bool {
2206        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2207        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2208    }
2209
2210    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
2211    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
2212    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
2213    /// on the producer before any read), bit-identical by construction.
2214    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
2215    pub fn pdl_nvfp4q8_on() -> bool {
2216        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2217        *ON.get_or_init(|| {
2218            std::env::var("MEMRA_PDL_NVFP4")
2219                .map(|v| v != "0")
2220                .unwrap_or(true)
2221        })
2222    }
2223
2224    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
2225    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
2226    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
2227    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
2228    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
2229    fn q40_mr1_on() -> bool {
2230        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2231        match *Q40MR.get_or_init(|| {
2232            std::env::var("MEMRA_Q40_MR")
2233                .ok()
2234                .and_then(|v| v.parse().ok())
2235        }) {
2236            Some(v) => v == 1,
2237            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2238        }
2239    }
2240
2241    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
2242    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
2243    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
2244    /// writes wrong bytes silently.
2245    fn pdl_func_flash(
2246        &self,
2247        g: bool,
2248        name: &'static str,
2249    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2250        use cudarc::driver::sys as cu;
2251        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
2252        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
2253        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
2254        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
2255        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
2256        // this engine's CUcontext; single-context runs behave exactly as before.
2257        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2258            std::sync::Mutex::new(None);
2259        static FNS: std::sync::Mutex<
2260            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2261        > = std::sync::Mutex::new(None);
2262        let ctx_key = self.ctx().cu_ctx() as usize;
2263        if let Some(&f) = FNS
2264            .lock()
2265            .unwrap()
2266            .get_or_insert_with(Default::default)
2267            .get(&(ctx_key, g, name))
2268        {
2269            return Ok(f as cu::CUfunction);
2270        }
2271        let module = {
2272            let mut mods = MODS.lock().unwrap();
2273            let map = mods.get_or_insert_with(Default::default);
2274            match map.get(&(ctx_key, g)) {
2275                Some(&m) => m,
2276                None => {
2277                    let m = self.pdl_load_module_in_ctx(if g {
2278                        FLASH_FATBIN_KF8VF8
2279                    } else {
2280                        FLASH_FATBIN
2281                    })?;
2282                    map.insert((ctx_key, g), m);
2283                    m
2284                }
2285            }
2286        };
2287        let cname = std::ffi::CString::new(name)?;
2288        let mut f: cu::CUfunction = std::ptr::null_mut();
2289        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2290        if r != cu::CUresult::CUDA_SUCCESS {
2291            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2292        }
2293        FNS.lock()
2294            .unwrap()
2295            .get_or_insert_with(Default::default)
2296            .insert((ctx_key, g, name), f as usize);
2297        Ok(f)
2298    }
2299
2300    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2301    /// the module to the thread's CURRENT context — a remote-stage engine must not
2302    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2303    /// current context before returning.
2304    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2305        use cudarc::driver::sys as cu;
2306        let mut prev: cu::CUcontext = std::ptr::null_mut();
2307        unsafe {
2308            cu::cuCtxGetCurrent(&mut prev).result()?;
2309        }
2310        self.ctx().bind_to_thread()?;
2311        let mut m: cu::CUmodule = std::ptr::null_mut();
2312        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2313        let restore = if prev.is_null() {
2314            cu::CUresult::CUDA_SUCCESS
2315        } else {
2316            unsafe { cu::cuCtxSetCurrent(prev) }
2317        };
2318        if r != cu::CUresult::CUDA_SUCCESS {
2319            return Err(format!("pdl module load: {r:?}").into());
2320        }
2321        if restore != cu::CUresult::CUDA_SUCCESS {
2322            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2323        }
2324        Ok(m as usize)
2325    }
2326
2327    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
2328    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
2329    pub fn raw_kernel_function(
2330        &self,
2331        name: &'static str,
2332    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2333        self.pdl_func(name)
2334    }
2335
2336    fn pdl_func(
2337        &self,
2338        name: &'static str,
2339    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2340        use cudarc::driver::sys as cu;
2341        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2342        // are context-scoped; key everything by this engine's CUcontext).
2343        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2344            std::sync::Mutex::new(None);
2345        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2346        // duplicate module, loaded lazily on the first kernels-module miss.
2347        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2348            std::sync::Mutex::new(None);
2349        static FNS: std::sync::Mutex<
2350            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2351        > = std::sync::Mutex::new(None);
2352        let ctx_key = self.ctx().cu_ctx() as usize;
2353        if let Some(&f) = FNS
2354            .lock()
2355            .unwrap()
2356            .get_or_insert_with(Default::default)
2357            .get(&(ctx_key, name))
2358        {
2359            return Ok(f as cu::CUfunction);
2360        }
2361        let module = {
2362            let mut mods = MODULES.lock().unwrap();
2363            let map = mods.get_or_insert_with(Default::default);
2364            match map.get(&ctx_key) {
2365                Some(&m) => m,
2366                None => {
2367                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2368                    map.insert(ctx_key, m);
2369                    m
2370                }
2371            }
2372        };
2373        let cname = std::ffi::CString::new(name)?;
2374        let mut f: cu::CUfunction = std::ptr::null_mut();
2375        let mut r =
2376            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2377        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2378            let qmodule = {
2379                let mut mods = QMODULES.lock().unwrap();
2380                let map = mods.get_or_insert_with(Default::default);
2381                match map.get(&ctx_key) {
2382                    Some(&m) => m,
2383                    None => {
2384                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2385                        map.insert(ctx_key, m);
2386                        m
2387                    }
2388                }
2389            };
2390            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2391        }
2392        if r != cu::CUresult::CUDA_SUCCESS {
2393            return Err(format!("pdl_func {name}: {r:?}").into());
2394        }
2395        FNS.lock()
2396            .unwrap()
2397            .get_or_insert_with(Default::default)
2398            .insert((ctx_key, name), f as usize);
2399        Ok(f)
2400    }
2401
2402    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2403    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2404    ///
2405    /// # Safety
2406    /// `params` must match the kernel's exact parameter list (order, types, count) —
2407    /// a mismatch corrupts the launch silently.
2408    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2409    /// builder path's fa_func/func_g choice exactly).
2410    ///
2411    /// # Safety
2412    /// Same contract as `launch_pdl`.
2413    unsafe fn launch_pdl_flash(
2414        &self,
2415        g: bool,
2416        name: &'static str,
2417        grid: (u32, u32, u32),
2418        block: (u32, u32, u32),
2419        smem: u32,
2420        params: &mut [*mut std::ffi::c_void],
2421    ) -> Result<(), Box<dyn std::error::Error>> {
2422        use cudarc::driver::sys as cu;
2423        let f = self.pdl_func_flash(g, name)?;
2424        if smem > 0 {
2425            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2426            let r =
2427                unsafe {
2428                    cu::cuFuncSetAttribute(f,
2429                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2430                smem as i32)
2431                };
2432            if r != cu::CUresult::CUDA_SUCCESS {
2433                return Err(format!("pdl smem attr {name}: {r:?}").into());
2434            }
2435        }
2436        let mut attr = cu::CUlaunchAttribute {
2437            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2438            pad: [0; 4],
2439            value: cu::CUlaunchAttributeValue {
2440                programmaticStreamSerializationAllowed: 1,
2441            },
2442        };
2443        let cfg = cu::CUlaunchConfig {
2444            gridDimX: grid.0,
2445            gridDimY: grid.1,
2446            gridDimZ: grid.2,
2447            blockDimX: block.0,
2448            blockDimY: block.1,
2449            blockDimZ: block.2,
2450            sharedMemBytes: smem,
2451            hStream: self.gpu.stream().cu_stream(),
2452            attrs: &mut attr,
2453            numAttrs: 1,
2454        };
2455        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2456        if r != cu::CUresult::CUDA_SUCCESS {
2457            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2458        }
2459        Ok(())
2460    }
2461
2462    unsafe fn launch_pdl(
2463        &self,
2464        name: &'static str,
2465        grid: (u32, u32, u32),
2466        block: (u32, u32, u32),
2467        params: &mut [*mut std::ffi::c_void],
2468    ) -> Result<(), Box<dyn std::error::Error>> {
2469        use cudarc::driver::sys as cu;
2470        let f = self.pdl_func(name)?;
2471        let mut attr = cu::CUlaunchAttribute {
2472            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2473            pad: [0; 4],
2474            value: cu::CUlaunchAttributeValue {
2475                programmaticStreamSerializationAllowed: 1,
2476            },
2477        };
2478        let cfg = cu::CUlaunchConfig {
2479            gridDimX: grid.0,
2480            gridDimY: grid.1,
2481            gridDimZ: grid.2,
2482            blockDimX: block.0,
2483            blockDimY: block.1,
2484            blockDimZ: block.2,
2485            sharedMemBytes: 0,
2486            hStream: self.gpu.stream().cu_stream(),
2487            attrs: &mut attr,
2488            numAttrs: 1,
2489        };
2490        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2491        if r != cu::CUresult::CUDA_SUCCESS {
2492            return Err(format!("launch_pdl {name}: {r:?}").into());
2493        }
2494        Ok(())
2495    }
2496
2497    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2498    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2499    pub fn prefetch_weight_l2(
2500        &self,
2501        w: &crate::model::GpuTensor,
2502    ) -> Result<(), Box<dyn std::error::Error>> {
2503        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2504            let p = rp4.as_ref().unwrap_or(bytes);
2505            self.prefetch_l2(p, p.len())?;
2506        }
2507        Ok(())
2508    }
2509
2510    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2511    /// by the DEVICE token id at tok[idx] into f32.
2512    pub fn gather_row_bf16(
2513        &self,
2514        table: &CudaSlice<u8>,
2515        tok: &CudaSlice<u32>,
2516        idx: usize,
2517        dst: &mut CudaSlice<f32>,
2518        ncols: usize,
2519    ) -> Result<(), Box<dyn std::error::Error>> {
2520        let f = self.func("gather_row_bf16_f32");
2521        let cfg = LaunchConfig {
2522            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2523            block_dim: (256, 1, 1),
2524            shared_mem_bytes: 0,
2525        };
2526        let (nc, ix) = (ncols as i32, idx as i32);
2527        let __s_b = self.gpu.stream();
2528        let mut b = __s_b.launch_builder(&f);
2529        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2530        unsafe {
2531            b.launch(cfg)?;
2532        }
2533        Ok(())
2534    }
2535
2536    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2537    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2538    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2539    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2540    /// finish(1).
2541    #[allow(clippy::too_many_arguments)]
2542    pub fn dflash2_dynconv(
2543        &self,
2544        x: &CudaSlice<f32>,
2545        dyn_: &CudaSlice<f32>,
2546        base: &CudaSlice<f32>,
2547        out: &mut CudaSlice<f32>,
2548        rows: usize,
2549        hidden: usize,
2550        group_size: usize,
2551        ksize: usize,
2552        half: usize,
2553    ) -> Result<(), Box<dyn std::error::Error>> {
2554        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2555        let f = self.func("dflash2_dynconv_f32");
2556        let n = rows * hidden;
2557        let cfg = LaunchConfig {
2558            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2559            block_dim: (256, 1, 1),
2560            shared_mem_bytes: 0,
2561        };
2562        let (ri, hi, gi, ki, hf) = (
2563            rows as i32,
2564            hidden as i32,
2565            group_size as i32,
2566            ksize as i32,
2567            half as i32,
2568        );
2569        let __s_b = self.gpu.stream();
2570        let mut b = __s_b.launch_builder(&f);
2571        b.arg(x)
2572            .arg(dyn_)
2573            .arg(base)
2574            .arg(out)
2575            .arg(&ri)
2576            .arg(&hi)
2577            .arg(&gi)
2578            .arg(&ki)
2579            .arg(&hf);
2580        unsafe {
2581            b.launch(cfg)?;
2582        }
2583        Ok(())
2584    }
2585
2586    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2587    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2588    /// value-descending, ties to the lower index.
2589    pub fn topk_rows(
2590        &self,
2591        logits: &CudaSlice<f32>,
2592        n_rows: usize,
2593        n_cols: usize,
2594        k: usize,
2595    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2596        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2597        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2598        let f = self.func("topk_rows_f32");
2599        let nth = 256usize;
2600        let mut vals = self.uninit(n_rows * k)?;
2601        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2602        let cfg = LaunchConfig {
2603            grid_dim: (n_rows as u32, 1, 1),
2604            block_dim: (nth as u32, 1, 1),
2605            shared_mem_bytes: (nth * k * 8) as u32,
2606        };
2607        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2608        let __s_b = self.gpu.stream();
2609        let mut b = __s_b.launch_builder(&f);
2610        b.arg(logits)
2611            .arg(&nr)
2612            .arg(&nc)
2613            .arg(&ki)
2614            .arg(&mut vals)
2615            .arg(&mut idxs);
2616        unsafe {
2617            b.launch(cfg)?;
2618        }
2619        Ok((vals, idxs))
2620    }
2621
2622    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2623    pub fn add_row_inplace(
2624        &self,
2625        logits: &mut CudaSlice<f32>,
2626        bias: &CudaSlice<f32>,
2627        n: usize,
2628        row_off: usize,
2629    ) -> Result<(), Box<dyn std::error::Error>> {
2630        let f = self.func("add_row_inplace_f32");
2631        let cfg = LaunchConfig {
2632            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2633            block_dim: (256, 1, 1),
2634            shared_mem_bytes: 0,
2635        };
2636        let (ni, off) = (n as i32, row_off as i64);
2637        let __s_b = self.gpu.stream();
2638        let mut b = __s_b.launch_builder(&f);
2639        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2640        unsafe {
2641            b.launch(cfg)?;
2642        }
2643        Ok(())
2644    }
2645
2646    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2647    pub fn prefetch_l2(
2648        &self,
2649        p: &CudaSlice<u8>,
2650        n: usize,
2651    ) -> Result<(), Box<dyn std::error::Error>> {
2652        let f = self.func("prefetch_l2_bytes");
2653        let lines = n.div_ceil(128);
2654        let ni = n as i64;
2655        let cfg = LaunchConfig {
2656            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2657            block_dim: (256, 1, 1),
2658            shared_mem_bytes: 0,
2659        };
2660        let __s_b = self.gpu.stream();
2661        let mut b = __s_b.launch_builder(&f);
2662        b.arg(p).arg(&ni);
2663        unsafe {
2664            b.launch(cfg)?;
2665        }
2666        Ok(())
2667    }
2668
2669    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2670    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2671    pub fn router_gemv(
2672        &self,
2673        w: &CudaSlice<f32>,
2674        x: &CudaSlice<f32>,
2675        n_embd: usize,
2676        n_experts: usize,
2677        t: usize,
2678    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2679        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2680        // stream differs) — too small to justify a numeric config change; deleted.
2681        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2682        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2683        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2684        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2685            Ok("0") => false,
2686            Ok(_) => true,
2687            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2688        };
2689        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2690        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2691        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2692        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2693        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2694        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2695        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2696        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2697        // (perf-only, bits equal).
2698        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2699        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2700    }
2701
2702    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2703    /// force both forms; `batch` requires `w8`).
2704    pub fn router_gemv_form(
2705        &self,
2706        w: &CudaSlice<f32>,
2707        x: &CudaSlice<f32>,
2708        n_embd: usize,
2709        n_experts: usize,
2710        t: usize,
2711        w8: bool,
2712        batch: bool,
2713    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2714        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2715        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2716        let f = if batch {
2717            self.func("router_gemv_f32_w8_batch")
2718        } else if w8 {
2719            self.func("router_gemv_f32_w8")
2720        } else {
2721            self.func("router_gemv_f32")
2722        };
2723        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2724        let cfg = if batch {
2725            LaunchConfig {
2726                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2727                block_dim: (32, 8, 1),
2728                shared_mem_bytes: 0,
2729            }
2730        } else {
2731            LaunchConfig {
2732                grid_dim: (n_experts as u32, t as u32, 1),
2733                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2734                shared_mem_bytes: 0,
2735            }
2736        };
2737        let __s_b = self.gpu.stream();
2738        let mut b = __s_b.launch_builder(&f);
2739        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2740        unsafe {
2741            b.launch(cfg)?;
2742        }
2743        Ok(y)
2744    }
2745
2746    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
2747    /// buffer — token-graph alloc-free.
2748    pub fn router_gemv_into(
2749        &self,
2750        w: &CudaSlice<f32>,
2751        x: &CudaSlice<f32>,
2752        y: &mut CudaSlice<f32>,
2753        n_embd: usize,
2754        n_experts: usize,
2755        t: usize,
2756    ) -> Result<(), Box<dyn std::error::Error>> {
2757        if y.len() < t * n_experts {
2758            return Err("router_gemv_into output too small".into());
2759        }
2760        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2761            Ok("0") => false,
2762            Ok(_) => true,
2763            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2764        };
2765        let f = if w8 {
2766            self.func("router_gemv_f32_w8")
2767        } else {
2768            self.func("router_gemv_f32")
2769        };
2770        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2771        let cfg = LaunchConfig {
2772            grid_dim: (n_experts as u32, t as u32, 1),
2773            block_dim: (32, if w8 { 8 } else { 1 }, 1),
2774            shared_mem_bytes: 0,
2775        };
2776        let __s_b = self.gpu.stream();
2777        let mut b = __s_b.launch_builder(&f);
2778        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2779        unsafe {
2780            b.launch(cfg)?;
2781        }
2782        Ok(())
2783    }
2784
2785    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2786    pub fn rows_permute(
2787        &self,
2788        src: &CudaSlice<f32>,
2789        idx: &CudaSlice<i32>,
2790        nrows: usize,
2791        ncols: usize,
2792    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2793        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2794        let f = self.func("rows_permute_f32");
2795        let (nc, nr) = (ncols as i32, nrows as i32);
2796        let cfg = LaunchConfig {
2797            grid_dim: (nrows as u32, 1, 1),
2798            block_dim: (256, 1, 1),
2799            shared_mem_bytes: 0,
2800        };
2801        let __s_b = self.gpu.stream();
2802        let mut b = __s_b.launch_builder(&f);
2803        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2804        unsafe {
2805            b.launch(cfg)?;
2806        }
2807        Ok(dst)
2808    }
2809
2810    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2811    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2812    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2813    /// decode chain and the small-t spec-verify chain match per row by construction.
2814    pub fn sigmoid_dot_rows(
2815        &self,
2816        x: &CudaSlice<f32>,
2817        w: &CudaSlice<f32>,
2818        n_embd: usize,
2819        t: usize,
2820    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2821        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2822        // config; same class as MEMRA_ROUTER_V2).
2823        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2824        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2825            let gs = self.linear(x, w, t, n_embd, 1)?;
2826            let mut g = self.uninit(t)?;
2827            self.sigmoid(&gs, &mut g, t)?;
2828            return Ok(g);
2829        }
2830        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2831        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2832        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2833        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2834        // flags doctrine; this per-token form serves every t.
2835        let mut g = self.alloc_uninit::<f32>(t)?;
2836        let f = self.func("sigmoid_dot_rows_f32");
2837        let (ne, ti) = (n_embd as i32, t as i32);
2838        let cfg = LaunchConfig {
2839            grid_dim: (t as u32, 1, 1),
2840            block_dim: (32, 8, 1),
2841            shared_mem_bytes: 0,
2842        };
2843        let __s_b = self.gpu.stream();
2844        let mut b = __s_b.launch_builder(&f);
2845        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2846        unsafe {
2847            b.launch(cfg)?;
2848        }
2849        Ok(g)
2850    }
2851
2852    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
2853    pub fn sigmoid_dot_rows_into(
2854        &self,
2855        x: &CudaSlice<f32>,
2856        w: &CudaSlice<f32>,
2857        g: &mut CudaSlice<f32>,
2858        n_embd: usize,
2859        t: usize,
2860    ) -> Result<(), Box<dyn std::error::Error>> {
2861        if g.len() < t {
2862            return Err("sigmoid_dot_rows_into output too small".into());
2863        }
2864        let f = self.func("sigmoid_dot_rows_f32");
2865        let (ne, ti) = (n_embd as i32, t as i32);
2866        let cfg = LaunchConfig {
2867            grid_dim: (t as u32, 1, 1),
2868            block_dim: (32, 8, 1),
2869            shared_mem_bytes: 0,
2870        };
2871        let __s_b = self.gpu.stream();
2872        let mut b = __s_b.launch_builder(&f);
2873        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2874        unsafe {
2875            b.launch(cfg)?;
2876        }
2877        Ok(())
2878    }
2879
2880    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2881    pub fn spec_rollback_stream(
2882        &self,
2883        len_ptrs: &CudaSlice<u64>,
2884        pos_start: &CudaSlice<i32>,
2885        acc: &CudaSlice<u32>,
2886        base: usize,
2887        n_rows: usize,
2888    ) -> Result<(), Box<dyn std::error::Error>> {
2889        let f = self.func("spec_rollback_stream");
2890        let (b, nr) = (base as i32, n_rows as i32);
2891        let cfg = LaunchConfig {
2892            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2893            block_dim: (64, 1, 1),
2894            shared_mem_bytes: 0,
2895        };
2896        let __s_bl = self.gpu.stream();
2897        let mut bl = __s_bl.launch_builder(&f);
2898        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2899        unsafe {
2900            bl.launch(cfg)?;
2901        }
2902        Ok(())
2903    }
2904
2905    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2906    pub fn plain_tok_ring(
2907        &self,
2908        vam: &CudaSlice<u32>,
2909        pos_start: &CudaSlice<i32>,
2910        base: usize,
2911        ring: &mut CudaSlice<u32>,
2912    ) -> Result<(), Box<dyn std::error::Error>> {
2913        let f = self.func("plain_tok_ring");
2914        let (b, cap) = (base as i32, ring.len() as i32);
2915        let cfg = LaunchConfig {
2916            grid_dim: (1, 1, 1),
2917            block_dim: (32, 1, 1),
2918            shared_mem_bytes: 0,
2919        };
2920        let __s_bl = self.gpu.stream();
2921        let mut bl = __s_bl.launch_builder(&f);
2922        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2923        unsafe {
2924            bl.launch(cfg)?;
2925        }
2926        Ok(())
2927    }
2928
2929    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2930    pub fn spec_ring_commit(
2931        &self,
2932        vtok: &CudaSlice<u32>,
2933        acc: &CudaSlice<u32>,
2934        brk: &CudaSlice<u32>,
2935        ring: &mut CudaSlice<u32>,
2936        pend: &mut CudaSlice<u32>,
2937    ) -> Result<(), Box<dyn std::error::Error>> {
2938        let f = self.func("spec_ring_commit");
2939        let cfg = LaunchConfig {
2940            grid_dim: (1, 1, 1),
2941            block_dim: (32, 1, 1),
2942            shared_mem_bytes: 0,
2943        };
2944        let __s_b = self.gpu.stream();
2945        let mut b = __s_b.launch_builder(&f);
2946        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2947        unsafe {
2948            b.launch(cfg)?;
2949        }
2950        Ok(())
2951    }
2952    pub fn i32_copy_add(
2953        &self,
2954        src: &CudaSlice<i32>,
2955        dst: &mut CudaSlice<i32>,
2956        delta: i32,
2957    ) -> Result<(), Box<dyn std::error::Error>> {
2958        let f = self.func("i32_copy_add");
2959        let cfg = LaunchConfig {
2960            grid_dim: (1, 1, 1),
2961            block_dim: (32, 1, 1),
2962            shared_mem_bytes: 0,
2963        };
2964        let __s_b = self.gpu.stream();
2965        let mut b = __s_b.launch_builder(&f);
2966        b.arg(src).arg(dst).arg(&delta);
2967        unsafe {
2968            b.launch(cfg)?;
2969        }
2970        Ok(())
2971    }
2972    pub fn u32_copy(
2973        &self,
2974        src: &CudaSlice<u32>,
2975        dst: &mut CudaSlice<u32>,
2976    ) -> Result<(), Box<dyn std::error::Error>> {
2977        let f = self.func("u32_copy");
2978        let cfg = LaunchConfig {
2979            grid_dim: (1, 1, 1),
2980            block_dim: (32, 1, 1),
2981            shared_mem_bytes: 0,
2982        };
2983        let __s_b = self.gpu.stream();
2984        let mut b = __s_b.launch_builder(&f);
2985        b.arg(src).arg(dst);
2986        unsafe {
2987            b.launch(cfg)?;
2988        }
2989        Ok(())
2990    }
2991
2992    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2993    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2994    /// caps acceptance exactly like drafting fewer tokens).
2995    pub fn spec_adapt_k(
2996        &self,
2997        acc: &CudaSlice<u32>,
2998        brk: &mut CudaSlice<u32>,
2999        floor: usize,
3000        cap: usize,
3001    ) -> Result<(), Box<dyn std::error::Error>> {
3002        let f = self.func("spec_adapt_k");
3003        let (fl, cp) = (floor as i32, cap as i32);
3004        let cfg = LaunchConfig {
3005            grid_dim: (1, 1, 1),
3006            block_dim: (32, 1, 1),
3007            shared_mem_bytes: 0,
3008        };
3009        let __s_b = self.gpu.stream();
3010        let mut b = __s_b.launch_builder(&f);
3011        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3012        unsafe {
3013            b.launch(cfg)?;
3014        }
3015        Ok(())
3016    }
3017
3018    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
3019    pub fn spec_accept_greedy_dc(
3020        &self,
3021        preds: &CudaSlice<u32>,
3022        vtok: &CudaSlice<u32>,
3023        last_pred: &CudaSlice<u32>,
3024        brk: &CudaSlice<u32>,
3025        out: &mut CudaSlice<u32>,
3026    ) -> Result<(), Box<dyn std::error::Error>> {
3027        let f = self.func("spec_accept_greedy_dc");
3028        let cfg = LaunchConfig {
3029            grid_dim: (1, 1, 1),
3030            block_dim: (32, 1, 1),
3031            shared_mem_bytes: 0,
3032        };
3033        let __s_b = self.gpu.stream();
3034        let mut b = __s_b.launch_builder(&f);
3035        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3036        unsafe {
3037            b.launch(cfg)?;
3038        }
3039        Ok(())
3040    }
3041
3042    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
3043    pub fn pos_iota(
3044        &self,
3045        pos0: &CudaSlice<i32>,
3046        out: &mut CudaSlice<i32>,
3047        t: usize,
3048    ) -> Result<(), Box<dyn std::error::Error>> {
3049        let f = self.func("pos_iota_i32");
3050        let ti = t as i32;
3051        let cfg = LaunchConfig {
3052            grid_dim: (1, 1, 1),
3053            block_dim: (t.max(1) as u32, 1, 1),
3054            shared_mem_bytes: 0,
3055        };
3056        let __s_b = self.gpu.stream();
3057        let mut b = __s_b.launch_builder(&f);
3058        b.arg(pos0).arg(out).arg(&ti);
3059        unsafe {
3060            b.launch(cfg)?;
3061        }
3062        Ok(())
3063    }
3064    #[allow(clippy::too_many_arguments)]
3065    pub fn append_kv_quantized_rows_dc(
3066        &self,
3067        k_rows: &CudaSlice<f32>,
3068        v_rows: &CudaSlice<f32>,
3069        kc: &mut CudaSlice<u8>,
3070        vc: &mut CudaSlice<u8>,
3071        t0_dev: &CudaSlice<i32>,
3072        t: usize,
3073        kv_dim_k: usize,
3074        kv_dim_v: usize,
3075        k_tok_bytes: usize,
3076        v_tok_bytes: usize,
3077        g: bool,
3078    ) -> Result<(), Box<dyn std::error::Error>> {
3079        let f = if g {
3080            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3081        } else {
3082            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3083        };
3084        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3085        let cfg = LaunchConfig {
3086            grid_dim: (nblk, t as u32, 1),
3087            block_dim: (32, 1, 1),
3088            shared_mem_bytes: 0,
3089        };
3090        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3091        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3092        let __s_b = self.gpu.stream();
3093        let mut b = __s_b.launch_builder(&f);
3094        b.arg(k_rows)
3095            .arg(v_rows)
3096            .arg(kc)
3097            .arg(vc)
3098            .arg(t0_dev)
3099            .arg(&kdk)
3100            .arg(&kdv)
3101            .arg(&ktb)
3102            .arg(&vtb);
3103        unsafe {
3104            b.launch(cfg)?;
3105        }
3106        Ok(())
3107    }
3108
3109    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
3110    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
3111    #[allow(clippy::too_many_arguments)]
3112    pub fn append_kv_quantized_row_dc_inc(
3113        &self,
3114        k_row: &CudaSlice<f32>,
3115        v_row: &CudaSlice<f32>,
3116        kc: &mut CudaSlice<u8>,
3117        vc: &mut CudaSlice<u8>,
3118        t0_dev: &mut CudaSlice<i32>,
3119        kv_dim_k: usize,
3120        kv_dim_v: usize,
3121        k_tok_bytes: usize,
3122        v_tok_bytes: usize,
3123        g: bool,
3124    ) -> Result<(), Box<dyn std::error::Error>> {
3125        let f = if g {
3126            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3127        } else {
3128            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3129        };
3130        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3131        let cfg = LaunchConfig {
3132            grid_dim: (1, 1, 1),
3133            block_dim: (nthreads, 1, 1),
3134            shared_mem_bytes: 0,
3135        };
3136        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3137        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3138        let __s_b = self.gpu.stream();
3139        let mut b = __s_b.launch_builder(&f);
3140        b.arg(k_row)
3141            .arg(v_row)
3142            .arg(kc)
3143            .arg(vc)
3144            .arg(t0_dev)
3145            .arg(&kdk)
3146            .arg(&kdv)
3147            .arg(&ktb)
3148            .arg(&vtb);
3149        unsafe {
3150            b.launch(cfg)?;
3151        }
3152        Ok(())
3153    }
3154
3155    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
3156    pub fn pack_tok_p(
3157        &self,
3158        tok: &CudaSlice<u32>,
3159        p: &CudaSlice<f32>,
3160        out: &mut CudaSlice<u32>,
3161        slot: usize,
3162    ) -> Result<(), Box<dyn std::error::Error>> {
3163        let f = self.func("pack_tok_p");
3164        let sl = slot as i32;
3165        let cfg = LaunchConfig {
3166            grid_dim: (1, 1, 1),
3167            block_dim: (32, 1, 1),
3168            shared_mem_bytes: 0,
3169        };
3170        let __s_b = self.gpu.stream();
3171        let mut b = __s_b.launch_builder(&f);
3172        b.arg(tok).arg(p).arg(out).arg(&sl);
3173        unsafe {
3174            b.launch(cfg)?;
3175        }
3176        Ok(())
3177    }
3178    pub fn tok_map_u32(
3179        &self,
3180        tok: &mut CudaSlice<u32>,
3181        map: &CudaSlice<u32>,
3182    ) -> Result<(), Box<dyn std::error::Error>> {
3183        let f = self.func("tok_map_u32");
3184        let cfg = LaunchConfig {
3185            grid_dim: (1, 1, 1),
3186            block_dim: (32, 1, 1),
3187            shared_mem_bytes: 0,
3188        };
3189        let __s_b = self.gpu.stream();
3190        let mut b = __s_b.launch_builder(&f);
3191        b.arg(tok).arg(map);
3192        unsafe {
3193            b.launch(cfg)?;
3194        }
3195        Ok(())
3196    }
3197
3198    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
3199    #[allow(clippy::too_many_arguments)]
3200    pub fn spec_assemble_verify(
3201        &self,
3202        tokp: &CudaSlice<u32>,
3203        pend: &CudaSlice<u32>,
3204        d2t: Option<&CudaSlice<u32>>,
3205        vtok: &mut CudaSlice<u32>,
3206        brk: &mut CudaSlice<u32>,
3207        p_min: f32,
3208        k: usize,
3209        pmin0: bool,
3210    ) -> Result<(), Box<dyn std::error::Error>> {
3211        let f = self.func("spec_assemble_verify");
3212        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3213        let cfg = LaunchConfig {
3214            grid_dim: (1, 1, 1),
3215            block_dim: (32, 1, 1),
3216            shared_mem_bytes: 0,
3217        };
3218        let __s_b = self.gpu.stream();
3219        let mut b = __s_b.launch_builder(&f);
3220        match d2t {
3221            Some(m) => {
3222                b.arg(tokp)
3223                    .arg(pend)
3224                    .arg(m)
3225                    .arg(vtok)
3226                    .arg(brk)
3227                    .arg(&p_min)
3228                    .arg(&ki)
3229                    .arg(&pm);
3230                unsafe {
3231                    b.launch(cfg)?;
3232                }
3233            }
3234            None => {
3235                let null: u64 = 0;
3236                b.arg(tokp)
3237                    .arg(pend)
3238                    .arg(&null)
3239                    .arg(vtok)
3240                    .arg(brk)
3241                    .arg(&p_min)
3242                    .arg(&ki)
3243                    .arg(&pm);
3244                unsafe {
3245                    b.launch(cfg)?;
3246                }
3247            }
3248        }
3249        Ok(())
3250    }
3251
3252    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
3253    #[allow(clippy::too_many_arguments)]
3254    pub fn ssm_conv_ring_rebuild_dc(
3255        &self,
3256        qkv_tm: &CudaSlice<f32>,
3257        ring_old: &CudaSlice<f32>,
3258        conv_state: &mut CudaSlice<f32>,
3259        conv_dim: usize,
3260        acc: &CudaSlice<u32>,
3261        base: usize,
3262        t_v: usize,
3263        d_conv: usize,
3264    ) -> Result<(), Box<dyn std::error::Error>> {
3265        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3266        let n = conv_dim * (d_conv - 1);
3267        let cfg = LaunchConfig::for_num_elems(n as u32);
3268        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3269        let __s_b = self.gpu.stream();
3270        let mut b = __s_b.launch_builder(&f);
3271        b.arg(qkv_tm)
3272            .arg(ring_old)
3273            .arg(conv_state)
3274            .arg(&cd)
3275            .arg(acc)
3276            .arg(&b0)
3277            .arg(&tv)
3278            .arg(&dc);
3279        unsafe {
3280            b.launch(cfg)?;
3281        }
3282        Ok(())
3283    }
3284    #[allow(clippy::too_many_arguments)]
3285    pub fn gdn_scan_s128_dc(
3286        &self,
3287        q: &CudaSlice<f32>,
3288        k: &CudaSlice<f32>,
3289        v: &CudaSlice<f32>,
3290        g: &CudaSlice<f32>,
3291        beta: &CudaSlice<f32>,
3292        state_in: &CudaSlice<f32>,
3293        state_out: &mut CudaSlice<f32>,
3294        o: &mut CudaSlice<f32>,
3295        n_head: usize,
3296        acc: &CudaSlice<u32>,
3297        base: usize,
3298        t_v: usize,
3299        scale: f32,
3300    ) -> Result<(), Box<dyn std::error::Error>> {
3301        let f = self.func("gdn_scan_s128_dc");
3302        const S_V: u32 = 128;
3303        const WARP: u32 = 32;
3304        const COLS_PER_BLOCK: u32 = 4;
3305        let cfg = LaunchConfig {
3306            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3307            block_dim: (WARP, COLS_PER_BLOCK, 1),
3308            shared_mem_bytes: 0,
3309        };
3310        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3311        let __s_b = self.gpu.stream();
3312        let mut b = __s_b.launch_builder(&f);
3313        b.arg(q)
3314            .arg(k)
3315            .arg(v)
3316            .arg(g)
3317            .arg(beta)
3318            .arg(state_in)
3319            .arg(state_out)
3320            .arg(o)
3321            .arg(&h)
3322            .arg(acc)
3323            .arg(&b0)
3324            .arg(&tv)
3325            .arg(&scale);
3326        unsafe {
3327            b.launch(cfg)?;
3328        }
3329        Ok(())
3330    }
3331
3332    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
3333    pub fn spec_rollback_kv(
3334        &self,
3335        len_ptrs: &CudaSlice<u64>,
3336        saved: &CudaSlice<i32>,
3337        acc: &CudaSlice<u32>,
3338        base: usize,
3339        n_layer: usize,
3340    ) -> Result<(), Box<dyn std::error::Error>> {
3341        let f = self.func("spec_rollback_kv");
3342        let (b, nl) = (base as i32, n_layer as i32);
3343        let cfg = LaunchConfig {
3344            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3345            block_dim: (64, 1, 1),
3346            shared_mem_bytes: 0,
3347        };
3348        let __s_bl = self.gpu.stream();
3349        let mut bl = __s_bl.launch_builder(&f);
3350        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3351        unsafe {
3352            bl.launch(cfg)?;
3353        }
3354        Ok(())
3355    }
3356
3357    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3358    pub fn spec_fork_valid(
3359        &self,
3360        acc: &CudaSlice<u32>,
3361        optimistic_pending: u32,
3362        valid: &mut CudaSlice<u32>,
3363    ) -> Result<(), Box<dyn std::error::Error>> {
3364        let f = self.func("spec_fork_valid");
3365        let cfg = LaunchConfig {
3366            grid_dim: (1, 1, 1),
3367            block_dim: (1, 1, 1),
3368            shared_mem_bytes: 0,
3369        };
3370        let __s_bl = self.gpu.stream();
3371        let mut bl = __s_bl.launch_builder(&f);
3372        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3373        unsafe {
3374            bl.launch(cfg)?;
3375        }
3376        Ok(())
3377    }
3378
3379    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3380    pub fn spec_fork_reconcile_kv(
3381        &self,
3382        len_ptrs: &CudaSlice<u64>,
3383        saved: &CudaSlice<i32>,
3384        acc: &CudaSlice<u32>,
3385        valid: &CudaSlice<u32>,
3386        base: usize,
3387        n_layer: usize,
3388    ) -> Result<(), Box<dyn std::error::Error>> {
3389        let f = self.func("spec_fork_reconcile_kv");
3390        let (b, nl) = (base as i32, n_layer as i32);
3391        let cfg = LaunchConfig {
3392            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3393            block_dim: (64, 1, 1),
3394            shared_mem_bytes: 0,
3395        };
3396        let __s_bl = self.gpu.stream();
3397        let mut bl = __s_bl.launch_builder(&f);
3398        bl.arg(len_ptrs)
3399            .arg(saved)
3400            .arg(acc)
3401            .arg(valid)
3402            .arg(&b)
3403            .arg(&nl);
3404        unsafe {
3405            bl.launch(cfg)?;
3406        }
3407        Ok(())
3408    }
3409
3410    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3411    pub fn spec_fork_restore_f32(
3412        &self,
3413        snapshot: &CudaSlice<f32>,
3414        state: &mut CudaSlice<f32>,
3415        valid: &CudaSlice<u32>,
3416    ) -> Result<(), Box<dyn std::error::Error>> {
3417        assert_eq!(
3418            snapshot.len(),
3419            state.len(),
3420            "fork recurrent snapshot shape mismatch"
3421        );
3422        let f = self.func("spec_fork_restore_f32");
3423        let n = state.len() as i32;
3424        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3425        let cfg = LaunchConfig {
3426            grid_dim: (blocks, 1, 1),
3427            block_dim: (256, 1, 1),
3428            shared_mem_bytes: 0,
3429        };
3430        let __s_bl = self.gpu.stream();
3431        let mut bl = __s_bl.launch_builder(&f);
3432        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3433        unsafe {
3434            bl.launch(cfg)?;
3435        }
3436        Ok(())
3437    }
3438
3439    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3440    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3441    pub fn spec_seed_gather(
3442        &self,
3443        vx: &CudaSlice<f32>,
3444        fill_prev: &CudaSlice<f32>,
3445        acc: &CudaSlice<u32>,
3446        h_seed: &mut CudaSlice<f32>,
3447        base: usize,
3448        n_embd: usize,
3449    ) -> Result<(), Box<dyn std::error::Error>> {
3450        let f = self.func("spec_seed_gather");
3451        let (b, ne) = (base as i32, n_embd as i32);
3452        let cfg = LaunchConfig {
3453            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3454            block_dim: (256, 1, 1),
3455            shared_mem_bytes: 0,
3456        };
3457        let __s_bl = self.gpu.stream();
3458        let mut bl = __s_bl.launch_builder(&f);
3459        bl.arg(vx)
3460            .arg(fill_prev)
3461            .arg(acc)
3462            .arg(h_seed)
3463            .arg(&b)
3464            .arg(&ne);
3465        unsafe {
3466            bl.launch(cfg)?;
3467        }
3468        Ok(())
3469    }
3470
3471    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3472    pub fn spec_accept_greedy(
3473        &self,
3474        preds: &CudaSlice<u32>,
3475        draft: &CudaSlice<u32>,
3476        last_pred: u32,
3477        base: usize,
3478        k_round: usize,
3479        out: &mut CudaSlice<u32>,
3480    ) -> Result<(), Box<dyn std::error::Error>> {
3481        let f = self.func("spec_accept_greedy");
3482        let (b, k) = (base as i32, k_round as i32);
3483        let cfg = LaunchConfig {
3484            grid_dim: (1, 1, 1),
3485            block_dim: (32, 1, 1),
3486            shared_mem_bytes: 0,
3487        };
3488        let __s_bl = self.gpu.stream();
3489        let mut bl = __s_bl.launch_builder(&f);
3490        bl.arg(preds)
3491            .arg(draft)
3492            .arg(&last_pred)
3493            .arg(&b)
3494            .arg(&k)
3495            .arg(out);
3496        unsafe {
3497            bl.launch(cfg)?;
3498        }
3499        Ok(())
3500    }
3501
3502    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3503    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3504    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3505
3506    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3507    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3508    pub fn gumbel_perturb(
3509        &self,
3510        x: &CudaSlice<f32>,
3511        y: &mut CudaSlice<f32>,
3512        n: usize,
3513        seed: u64,
3514        stream_pos: u32,
3515        temp: f32,
3516    ) -> Result<(), Box<dyn std::error::Error>> {
3517        let f = self.func("gumbel_perturb_f32");
3518        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3519        let cfg = LaunchConfig {
3520            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3521            block_dim: (256, 1, 1),
3522            shared_mem_bytes: 0,
3523        };
3524        let __s_b = self.gpu.stream();
3525        let mut b = __s_b.launch_builder(&f);
3526        b.arg(x)
3527            .arg(&mut *y)
3528            .arg(&ni)
3529            .arg(&slo)
3530            .arg(&shi)
3531            .arg(&stream_pos)
3532            .arg(&temp);
3533        unsafe {
3534            b.launch(cfg)?;
3535        }
3536        Ok(())
3537    }
3538
3539    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3540    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3541    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3542    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3543    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3544    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3545    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3546    pub fn mask_logits_col(
3547        &self,
3548        logits: &mut CudaSlice<f32>,
3549        mask: &CudaSlice<u32>,
3550        col: usize,
3551        n: usize,
3552        mask_words: usize,
3553    ) -> Result<(), Box<dyn std::error::Error>> {
3554        let f = self.func("mask_logits_f32");
3555        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3556        let cfg = LaunchConfig {
3557            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3558            block_dim: (256, 1, 1),
3559            shared_mem_bytes: 0,
3560        };
3561        let __s_b = self.gpu.stream();
3562        let mut b = __s_b.launch_builder(&f);
3563        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3564        unsafe {
3565            b.launch(cfg)?;
3566        }
3567        Ok(())
3568    }
3569
3570    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3571    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3572    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3573    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3574    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3575    /// pointer-invariance IS the serving isolation contract for sampled rows.
3576    pub fn gumbel_perturb_col(
3577        &self,
3578        x: &CudaSlice<f32>,
3579        col: usize,
3580        y: &mut CudaSlice<f32>,
3581        n: usize,
3582        seed: u64,
3583        stream_pos: u32,
3584        temp: f32,
3585    ) -> Result<(), Box<dyn std::error::Error>> {
3586        let f = self.func("gumbel_perturb_f32");
3587        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3588        let col_view = x.slice(col * n..(col + 1) * n);
3589        let cfg = LaunchConfig {
3590            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3591            block_dim: (256, 1, 1),
3592            shared_mem_bytes: 0,
3593        };
3594        let __s_b = self.gpu.stream();
3595        let mut b = __s_b.launch_builder(&f);
3596        b.arg(&col_view)
3597            .arg(&mut *y)
3598            .arg(&ni)
3599            .arg(&slo)
3600            .arg(&shi)
3601            .arg(&stream_pos)
3602            .arg(&temp);
3603        unsafe {
3604            b.launch(cfg)?;
3605        }
3606        Ok(())
3607    }
3608
3609    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3610    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3611    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3612    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3613    /// the serving isolation contract for sampled rows).
3614    #[allow(clippy::too_many_arguments)]
3615    pub fn gumbel_perturb_filtered_col(
3616        &self,
3617        x: &CudaSlice<f32>,
3618        col: usize,
3619        y: &mut CudaSlice<f32>,
3620        n: usize,
3621        seed: u64,
3622        stream_pos: u32,
3623        temp: f32,
3624        stat_max: &CudaSlice<f32>,
3625        stat_th: &CudaSlice<f32>,
3626        stat_idx: usize,
3627    ) -> Result<(), Box<dyn std::error::Error>> {
3628        let f = self.func("gumbel_perturb_filtered_col_f32");
3629        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3630        let (ci, si) = (col as i32, stat_idx as i32);
3631        let cfg = LaunchConfig {
3632            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3633            block_dim: (256, 1, 1),
3634            shared_mem_bytes: 0,
3635        };
3636        let __s_b = self.gpu.stream();
3637        let mut b = __s_b.launch_builder(&f);
3638        b.arg(x)
3639            .arg(&ci)
3640            .arg(&mut *y)
3641            .arg(&ni)
3642            .arg(&slo)
3643            .arg(&shi)
3644            .arg(&stream_pos)
3645            .arg(&temp)
3646            .arg(stat_max)
3647            .arg(stat_th)
3648            .arg(&si);
3649        unsafe {
3650            b.launch(cfg)?;
3651        }
3652        Ok(())
3653    }
3654
3655    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3656    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3657    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3658    /// reads it (counter is data, not state — graph-replay-safe).
3659    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3660        let f = self.func("memra_sctr_inc");
3661        let cfg = LaunchConfig {
3662            grid_dim: (1, 1, 1),
3663            block_dim: (1, 1, 1),
3664            shared_mem_bytes: 0,
3665        };
3666        let __s_b = self.gpu.stream();
3667        let mut b = __s_b.launch_builder(&f);
3668        b.arg(&mut *ctr);
3669        unsafe {
3670            b.launch(cfg)?;
3671        }
3672        Ok(())
3673    }
3674
3675    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3676    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3677    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3678    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3679    pub fn gumbel_perturb_ctr(
3680        &self,
3681        x: &CudaSlice<f32>,
3682        y: &mut CudaSlice<f32>,
3683        n: usize,
3684        seed: u64,
3685        ctr: &CudaSlice<u32>,
3686        temp: f32,
3687    ) -> Result<(), Box<dyn std::error::Error>> {
3688        let f = self.func("gumbel_perturb_ctr_f32");
3689        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3690        let cfg = LaunchConfig {
3691            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3692            block_dim: (256, 1, 1),
3693            shared_mem_bytes: 0,
3694        };
3695        let __s_b = self.gpu.stream();
3696        let mut b = __s_b.launch_builder(&f);
3697        b.arg(x)
3698            .arg(&mut *y)
3699            .arg(&ni)
3700            .arg(&slo)
3701            .arg(&shi)
3702            .arg(ctr)
3703            .arg(&temp);
3704        unsafe {
3705            b.launch(cfg)?;
3706        }
3707        Ok(())
3708    }
3709
3710    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3711    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3712    /// (smallest-index tie-break — matches the argmax-gate contract).
3713    pub fn softmax_gather(
3714        &self,
3715        x: &CudaSlice<f32>,
3716        row_stride: usize,
3717        ids: &CudaSlice<u32>,
3718        rows: &CudaSlice<i32>,
3719        out: &mut CudaSlice<f32>,
3720        n: usize,
3721        npair: usize,
3722        temp: f32,
3723    ) -> Result<(), Box<dyn std::error::Error>> {
3724        let f = self.func("softmax_gather_f32");
3725        let (ni, rs) = (n as i32, row_stride as i64);
3726        let np = npair as i32;
3727        let cfg = LaunchConfig {
3728            grid_dim: (npair as u32, 1, 1),
3729            block_dim: (256, 1, 1),
3730            shared_mem_bytes: 0,
3731        };
3732        let __s_b = self.gpu.stream();
3733        let mut b = __s_b.launch_builder(&f);
3734        b.arg(x)
3735            .arg(&rs)
3736            .arg(ids)
3737            .arg(rows)
3738            .arg(&mut *out)
3739            .arg(&ni)
3740            .arg(&np)
3741            .arg(&temp);
3742        unsafe {
3743            b.launch(cfg)?;
3744        }
3745        Ok(())
3746    }
3747
3748    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3749    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3750    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3751    pub fn residual_sample(
3752        &self,
3753        p: &CudaSlice<f32>,
3754        q: Option<&CudaSlice<f32>>,
3755        n: usize,
3756        temp: f32,
3757        seed: u64,
3758        stream_pos: u32,
3759        out_tok: &mut CudaSlice<u32>,
3760    ) -> Result<(), Box<dyn std::error::Error>> {
3761        let f = self.func("residual_sample_f32");
3762        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3763        let nth = 1024u32;
3764        let cfg = LaunchConfig {
3765            grid_dim: (1, 1, 1),
3766            block_dim: (nth, 1, 1),
3767            shared_mem_bytes: 0,
3768        };
3769        let has_q: i32 = q.is_some() as i32;
3770        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3771        let __s_b = self.gpu.stream();
3772        let mut b = __s_b.launch_builder(&f);
3773        b.arg(p)
3774            .arg(qbuf)
3775            .arg(&has_q)
3776            .arg(&ni)
3777            .arg(&temp)
3778            .arg(&slo)
3779            .arg(&shi)
3780            .arg(&stream_pos)
3781            .arg(&mut *out_tok);
3782        unsafe {
3783            b.launch(cfg)?;
3784        }
3785        Ok(())
3786    }
3787
3788    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3789    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3790    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3791    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3792    pub fn with_moe_cache<R>(
3793        &self,
3794        max_block_bytes: usize,
3795        f: impl FnOnce(
3796            &mut crate::moe_cache::MoeSlotCache,
3797            &Engine,
3798        ) -> Result<R, Box<dyn std::error::Error>>,
3799    ) -> Result<R, Box<dyn std::error::Error>> {
3800        let mut guard = self.moe_cache.lock().unwrap();
3801        if guard.is_none() {
3802            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3803        }
3804        let cache = guard.as_mut().unwrap();
3805        f(cache, self)
3806    }
3807
3808    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3809    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3810    pub fn freeze_moe_cache(&self) {
3811        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3812            cache.freeze();
3813        }
3814    }
3815
3816    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3817    /// Never constructs a cache.
3818    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3819        self.moe_cache
3820            .lock()
3821            .unwrap()
3822            .as_ref()
3823            .map(crate::moe_cache::MoeSlotCache::export_residency)
3824    }
3825
3826    pub(crate) fn moe_cache_frozen(&self) -> bool {
3827        self.moe_cache
3828            .lock()
3829            .unwrap()
3830            .as_ref()
3831            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3832    }
3833
3834    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3835    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3836    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3837    /// while leaving the profiling warmup's established batched behavior untouched.
3838    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3839    /// tokenwise arm anyway.)
3840    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3841        crate::cpu_experts::configured()
3842            && self.moe_cache_frozen()
3843            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3844    }
3845
3846    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3847    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3848        assert!(
3849            self.moe_cache.lock().unwrap().is_none(),
3850            "MoE cache layout configured after cache construction"
3851        );
3852        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3853    }
3854
3855    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3856        self.moe_cache_layout.lock().unwrap().clone()
3857    }
3858
3859    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3860    pub fn moe_cache_enabled() -> bool {
3861        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3862    }
3863
3864    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3865    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3866    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3867        let guard = self.moe_cache.lock().unwrap();
3868        guard
3869            .as_ref()
3870            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3871    }
3872
3873    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3874    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3875    /// callers compare a before/after snapshot around a decode window.
3876    pub fn cpu_expert_stats(
3877        &self,
3878    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3879        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3880    }
3881
3882    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3883    /// the backend tail that resident-GPU expert work did not hide.
3884    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3885        crate::cpu_experts::predictor_stats()
3886    }
3887
3888    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3889        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3890    }
3891
3892    /// CPU-routed expert selections grouped by how many of their three projections were already
3893    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3894    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3895        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3896    }
3897
3898    /// Positioned-read proof-backend counters:
3899    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3900    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3901        let guard = self.moe_cache.lock().unwrap();
3902        guard
3903            .as_ref()
3904            .and_then(|cache| cache.pread_stats())
3905            .map(|stats| {
3906                (
3907                    stats.reads,
3908                    stats.bytes,
3909                    stats.read_errors,
3910                    stats.short_reads,
3911                    stats.fallbacks,
3912                    stats.buffer_waits,
3913                    stats.ring_full,
3914                )
3915            })
3916    }
3917
3918    /// Spill configuration values that warned and substituted their documented defaults.
3919    pub fn spill_config_fallbacks(&self) -> u64 {
3920        crate::spill_pread::config_fallbacks()
3921    }
3922
3923    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3924    pub fn moe_cache_reset_counters(&self) {
3925        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3926            c.reset_counters();
3927        }
3928    }
3929
3930    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3931        Ok(self.gpu.stream().clone_htod(v)?)
3932    }
3933
3934    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3935    /// past the final q4_0 block through their aligned window — the bytes never reach a
3936    /// result (funnelshift discards them) but must be mapped memory.
3937    pub fn htod_bytes_padded(
3938        &self,
3939        v: &[u8],
3940        pad: usize,
3941    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3942        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3943        {
3944            let mut view = d.slice_mut(0..v.len());
3945            self.gpu.stream().memcpy_htod(v, &mut view)?;
3946        }
3947        Ok(d)
3948    }
3949
3950    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3951    pub fn copy_into(
3952        &self,
3953        dst: &mut CudaSlice<f32>,
3954        off: usize,
3955        src: &CudaSlice<f32>,
3956        len: usize,
3957    ) -> Result<(), Box<dyn std::error::Error>> {
3958        let mut view = dst.slice_mut(off..off + len);
3959        self.gpu
3960            .stream()
3961            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3962        Ok(())
3963    }
3964
3965    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3966    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3967    pub fn copy_u8_into(
3968        &self,
3969        dst: &mut CudaSlice<u8>,
3970        off: usize,
3971        src: &CudaSlice<u8>,
3972        len: usize,
3973    ) -> Result<(), Box<dyn std::error::Error>> {
3974        let mut view = dst.slice_mut(off..off + len);
3975        self.gpu
3976            .stream()
3977            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3978        Ok(())
3979    }
3980
3981    /// D2D byte-range copy with explicit source and destination offsets.
3982    pub fn copy_u8_range_into(
3983        &self,
3984        dst: &mut CudaSlice<u8>,
3985        dst_off: usize,
3986        src: &CudaSlice<u8>,
3987        src_off: usize,
3988        len: usize,
3989    ) -> Result<(), Box<dyn std::error::Error>> {
3990        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3991        self.gpu
3992            .stream()
3993            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3994        Ok(())
3995    }
3996
3997    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3998    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3999    /// keeping the audited attention range contiguous without changing its absolute start.
4000    pub fn prepare_kv_append(
4001        &self,
4002        kv: &mut crate::cache::KvLayer,
4003        retain_from: usize,
4004        append_rows: usize,
4005    ) -> Result<usize, Box<dyn std::error::Error>> {
4006        let Some(plan) = kv
4007            .ring
4008            .as_ref()
4009            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4010            .transpose()?
4011        else {
4012            return Ok(kv.len);
4013        };
4014        match plan {
4015            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4016            crate::cache::KvRingAppend::Rebase {
4017                src_row,
4018                keep_rows,
4019                new_base,
4020                write_row,
4021            } => {
4022                if keep_rows > 0 {
4023                    let k_len = keep_rows * kv.k_tok_bytes;
4024                    let v_len = keep_rows * kv.v_tok_bytes;
4025                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4026                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4027                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4028                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4029                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4030                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4031                }
4032                kv.ring.as_mut().unwrap().apply_rebase(new_base);
4033                Ok(write_row)
4034            }
4035        }
4036    }
4037
4038    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
4039    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
4040    pub fn htod_u8_into(
4041        &self,
4042        dst: &mut CudaSlice<u8>,
4043        off: usize,
4044        src: &[u8],
4045    ) -> Result<(), Box<dyn std::error::Error>> {
4046        let mut view = dst.slice_mut(off..off + src.len());
4047        self.gpu.stream().memcpy_htod(src, &mut view)?;
4048        Ok(())
4049    }
4050
4051    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4052        b.slice(0..len)
4053    }
4054
4055    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
4056    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
4057    pub fn view_u8_range<'a>(
4058        &self,
4059        b: &'a CudaSlice<u8>,
4060        start: usize,
4061        end: usize,
4062    ) -> cudarc::driver::CudaView<'a, u8> {
4063        b.slice(start..end)
4064    }
4065    pub fn view_u8<'a>(
4066        &self,
4067        b: &'a CudaSlice<u8>,
4068        len: usize,
4069    ) -> cudarc::driver::CudaView<'a, u8> {
4070        b.slice(0..len)
4071    }
4072
4073    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
4074    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
4075    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
4076    pub fn append_kv_quantized(
4077        &self,
4078        k_row: &CudaSlice<f32>,
4079        v_row: &CudaSlice<f32>,
4080        kc: &mut CudaSlice<u8>,
4081        vc: &mut CudaSlice<u8>,
4082        t: usize,
4083        kv_dim_k: usize,
4084        kv_dim_v: usize,
4085        k_tok_bytes: usize,
4086        v_tok_bytes: usize,
4087        g: bool,
4088    ) -> Result<(), Box<dyn std::error::Error>> {
4089        let f = if g {
4090            self.func_g("append_quantize_kv_q8_0_q5_1")
4091        } else {
4092            self.func("append_quantize_kv_q8_0_q5_1")
4093        };
4094        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4095        let cfg = LaunchConfig {
4096            grid_dim: (nblk, 1, 1),
4097            block_dim: (32, 1, 1),
4098            shared_mem_bytes: 0,
4099        };
4100        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4101        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4102        let __s_b = self.gpu.stream();
4103        let mut b = __s_b.launch_builder(&f);
4104        b.arg(k_row)
4105            .arg(v_row)
4106            .arg(kc)
4107            .arg(vc)
4108            .arg(&ti)
4109            .arg(&kdk)
4110            .arg(&kdv)
4111            .arg(&ktb)
4112            .arg(&vtb);
4113        unsafe {
4114            b.launch(cfg)?;
4115        }
4116        Ok(())
4117    }
4118
4119    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
4120    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
4121    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
4122    pub fn append_kv_quantized_dc(
4123        &self,
4124        k_row: &CudaSlice<f32>,
4125        v_row: &CudaSlice<f32>,
4126        kc: &mut CudaSlice<u8>,
4127        vc: &mut CudaSlice<u8>,
4128        t_dev: &CudaSlice<i32>,
4129        kv_dim_k: usize,
4130        kv_dim_v: usize,
4131        k_tok_bytes: usize,
4132        v_tok_bytes: usize,
4133        g: bool,
4134    ) -> Result<(), Box<dyn std::error::Error>> {
4135        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4136        let (kdk, kdv) = (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        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
4139        if Self::pdl_on() && Self::pdl_wb_on() {
4140            use cudarc::driver::{DevicePtr, DevicePtrMut};
4141            let s = &self.gpu.stream();
4142            let (pk, _g0) = k_row.device_ptr(s);
4143            let (pv, _g1) = v_row.device_ptr(s);
4144            let (pkc, _g2) = kc.device_ptr_mut(s);
4145            let (pvc, _g3) = vc.device_ptr_mut(s);
4146            let (pt, _g4) = t_dev.device_ptr(s);
4147            let mut ps = [
4148                &pk as *const _ as *mut std::ffi::c_void,
4149                &pv as *const _ as *mut _,
4150                &pkc as *const _ as *mut _,
4151                &pvc as *const _ as *mut _,
4152                &pt as *const _ as *mut _,
4153                &kdk as *const _ as *mut _,
4154                &kdv as *const _ as *mut _,
4155                &ktb as *const _ as *mut _,
4156                &vtb as *const _ as *mut _,
4157            ];
4158            unsafe {
4159                self.launch_pdl_flash(
4160                    g,
4161                    "append_quantize_kv_q8_0_q5_1_dc",
4162                    (nblk, 1, 1),
4163                    (32, 1, 1),
4164                    0,
4165                    &mut ps,
4166                )?;
4167            }
4168            return Ok(());
4169        }
4170        let f = if g {
4171            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4172        } else {
4173            self.func("append_quantize_kv_q8_0_q5_1_dc")
4174        };
4175        let cfg = LaunchConfig {
4176            grid_dim: (nblk, 1, 1),
4177            block_dim: (32, 1, 1),
4178            shared_mem_bytes: 0,
4179        };
4180        let __s_b = self.gpu.stream();
4181        let mut b = __s_b.launch_builder(&f);
4182        b.arg(k_row)
4183            .arg(v_row)
4184            .arg(kc)
4185            .arg(vc)
4186            .arg(t_dev)
4187            .arg(&kdk)
4188            .arg(&kdv)
4189            .arg(&ktb)
4190            .arg(&vtb);
4191        unsafe {
4192            b.launch(cfg)?;
4193        }
4194        Ok(())
4195    }
4196
4197    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
4198    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
4199    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
4200    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
4201    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
4202    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
4203    #[allow(clippy::too_many_arguments)]
4204    pub fn append_kv_quantized_rows(
4205        &self,
4206        k_rows: &CudaSlice<f32>,
4207        v_rows: &CudaSlice<f32>,
4208        kc: &mut CudaSlice<u8>,
4209        vc: &mut CudaSlice<u8>,
4210        t0: usize,
4211        t: usize,
4212        kv_dim_k: usize,
4213        kv_dim_v: usize,
4214        k_tok_bytes: usize,
4215        v_tok_bytes: usize,
4216        g: bool,
4217    ) -> Result<(), Box<dyn std::error::Error>> {
4218        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4219            for i in 0..t {
4220                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4221                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4222                self.append_kv_quantized_view(
4223                    &k_row,
4224                    &v_row,
4225                    kc,
4226                    vc,
4227                    t0 + i,
4228                    kv_dim_k,
4229                    kv_dim_v,
4230                    k_tok_bytes,
4231                    v_tok_bytes,
4232                    g,
4233                )?;
4234            }
4235            return Ok(());
4236        }
4237        let f = if g {
4238            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4239        } else {
4240            self.func("append_quantize_kv_q8_0_q5_1_rows")
4241        };
4242        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4243        let cfg = LaunchConfig {
4244            grid_dim: (nblk, t as u32, 1),
4245            block_dim: (32, 1, 1),
4246            shared_mem_bytes: 0,
4247        };
4248        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4249        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4250        let __s_b = self.gpu.stream();
4251        let mut b = __s_b.launch_builder(&f);
4252        b.arg(k_rows)
4253            .arg(v_rows)
4254            .arg(kc)
4255            .arg(vc)
4256            .arg(&t0i)
4257            .arg(&kdk)
4258            .arg(&kdv)
4259            .arg(&ktb)
4260            .arg(&vtb);
4261        unsafe {
4262            b.launch(cfg)?;
4263        }
4264        Ok(())
4265    }
4266
4267    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
4268    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
4269    /// later, inside a captured graph) without a host round-trip.
4270    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4271        let f = self.func("inc_i32");
4272        let cfg = LaunchConfig {
4273            grid_dim: (1, 1, 1),
4274            block_dim: (1, 1, 1),
4275            shared_mem_bytes: 0,
4276        };
4277        let __s_b = self.gpu.stream();
4278        let mut b = __s_b.launch_builder(&f);
4279        b.arg(p);
4280        unsafe {
4281            b.launch(cfg)?;
4282        }
4283        Ok(())
4284    }
4285
4286    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
4287    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
4288    pub fn append_kv_quantized_view(
4289        &self,
4290        k_row: &cudarc::driver::CudaView<f32>,
4291        v_row: &cudarc::driver::CudaView<f32>,
4292        kc: &mut CudaSlice<u8>,
4293        vc: &mut CudaSlice<u8>,
4294        t: usize,
4295        kv_dim_k: usize,
4296        kv_dim_v: usize,
4297        k_tok_bytes: usize,
4298        v_tok_bytes: usize,
4299        g: bool,
4300    ) -> Result<(), Box<dyn std::error::Error>> {
4301        let stream = self.gpu.stream();
4302        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4303        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4304        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4305        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4306        let f = if g {
4307            self.func_g("append_quantize_kv_q8_0_q5_1")
4308        } else {
4309            self.func("append_quantize_kv_q8_0_q5_1")
4310        };
4311        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4312        let cfg = LaunchConfig {
4313            grid_dim: (nblk, 1, 1),
4314            block_dim: (32, 1, 1),
4315            shared_mem_bytes: 0,
4316        };
4317        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4318        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4319        let mut b = stream.launch_builder(&f);
4320        b.arg(k_row)
4321            .arg(v_row)
4322            .arg(kc)
4323            .arg(vc)
4324            .arg(&ti)
4325            .arg(&kdk)
4326            .arg(&kdv)
4327            .arg(&ktb)
4328            .arg(&vtb);
4329        unsafe {
4330            b.launch(cfg)?;
4331        }
4332        Ok(())
4333    }
4334
4335    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
4336    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
4337    pub fn copy_view_into(
4338        &self,
4339        dst: &mut CudaSlice<f32>,
4340        off: usize,
4341        src: &cudarc::driver::CudaView<f32>,
4342        len: usize,
4343    ) -> Result<(), Box<dyn std::error::Error>> {
4344        let mut view = dst.slice_mut(off..off + len);
4345        self.gpu
4346            .stream()
4347            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4348        Ok(())
4349    }
4350
4351    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
4352    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
4353    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
4354    pub fn clone_dtod(
4355        &self,
4356        src: &CudaSlice<f32>,
4357    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4358        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4359        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4360        Ok(dst)
4361    }
4362
4363    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4364    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4365    pub fn dtod_copy_view(
4366        &self,
4367        src: &cudarc::driver::CudaView<f32>,
4368        dst: &mut CudaSlice<f32>,
4369    ) -> Result<(), Box<dyn std::error::Error>> {
4370        self.gpu.stream().memcpy_dtod(src, dst)?;
4371        Ok(())
4372    }
4373
4374    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4375    pub fn dtod_copy_view_i8(
4376        &self,
4377        src: &cudarc::driver::CudaView<i8>,
4378        dst: &mut CudaSlice<i8>,
4379    ) -> Result<(), Box<dyn std::error::Error>> {
4380        self.gpu.stream().memcpy_dtod(src, dst)?;
4381        Ok(())
4382    }
4383
4384    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4385    pub fn dtod_copy_into(
4386        &self,
4387        src: &CudaSlice<f32>,
4388        dst: &mut CudaSlice<f32>,
4389        offset: usize,
4390    ) -> Result<(), Box<dyn std::error::Error>> {
4391        let n = src.len();
4392        let mut dv = dst.slice_mut(offset..offset + n);
4393        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4394        Ok(())
4395    }
4396
4397    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4398    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4399    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4400    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4401    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4402    pub fn copy_batch_uniform_f32(
4403        &self,
4404        table: &CudaSlice<u64>,
4405        n: usize,
4406        words: usize,
4407    ) -> Result<(), Box<dyn std::error::Error>> {
4408        if n == 0 || words == 0 {
4409            return Ok(());
4410        }
4411        debug_assert!(
4412            table.len() >= 2 * n,
4413            "pointer table must hold n srcs + n dsts"
4414        );
4415        let f = self.func("copy_batch_uniform_f32");
4416        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4417        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4418        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4419        let (ni, wi) = (n as i32, words as i32);
4420        let cfg = LaunchConfig {
4421            grid_dim: (chunks, n as u32, 1),
4422            block_dim: (256, 1, 1),
4423            shared_mem_bytes: 0,
4424        };
4425        let __s = self.gpu.stream();
4426        let mut b = __s.launch_builder(&f);
4427        b.arg(table).arg(&ni).arg(&wi);
4428        unsafe {
4429            b.launch(cfg)?;
4430        }
4431        Ok(())
4432    }
4433
4434    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4435    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4436    pub fn htod_u64_into(
4437        &self,
4438        v: &[u64],
4439        dst: &mut CudaSlice<u64>,
4440    ) -> Result<(), Box<dyn std::error::Error>> {
4441        let mut view = dst.slice_mut(0..v.len());
4442        self.gpu.stream().memcpy_htod(v, &mut view)?;
4443        Ok(())
4444    }
4445
4446    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4447    /// device pointer-table entry at run time, so a captured graph follows the gdn
4448    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4449    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4450    pub fn copy_indirect_src_f32(
4451        &self,
4452        src_entry: &cudarc::driver::CudaView<u64>,
4453        dst: &mut CudaSlice<f32>,
4454        dst_off: usize,
4455        words: usize,
4456    ) -> Result<(), Box<dyn std::error::Error>> {
4457        let f = self.func("copy_indirect_src_f32");
4458        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4459        let wi = words as i32;
4460        let cfg = LaunchConfig {
4461            grid_dim: (chunks, 1, 1),
4462            block_dim: (256, 1, 1),
4463            shared_mem_bytes: 0,
4464        };
4465        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4466        let __s = self.gpu.stream();
4467        let mut b = __s.launch_builder(&f);
4468        b.arg(src_entry).arg(&mut dv).arg(&wi);
4469        unsafe {
4470            b.launch(cfg)?;
4471        }
4472        Ok(())
4473    }
4474
4475    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4476    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4477        self.alloc_uninit::<i8>(n)
4478    }
4479
4480    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4481    pub fn qmatvec(
4482        &self,
4483        w: &CudaSlice<u8>,
4484        x: &CudaSlice<f32>,
4485        m: usize,
4486        in_f: usize,
4487        out_f: usize,
4488        qtype: i32,
4489        row_bytes: usize,
4490    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4491        let f = self.func("qmatvec_f32");
4492        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4493        let cfg = LaunchConfig {
4494            grid_dim: (out_f as u32, m as u32, 1),
4495            block_dim: (256, 1, 1),
4496            shared_mem_bytes: 0,
4497        };
4498        let (inf, outf, mi, qt, rb) =
4499            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4500        let __s_b = self.gpu.stream();
4501        let mut b = __s_b.launch_builder(&f);
4502        b.arg(w)
4503            .arg(x)
4504            .arg(&mut y)
4505            .arg(&inf)
4506            .arg(&outf)
4507            .arg(&mi)
4508            .arg(&qt)
4509            .arg(&rb);
4510        unsafe {
4511            b.launch(cfg)?;
4512        }
4513        Ok(y)
4514    }
4515
4516    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4517    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4518        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4519        self.keep_if_capturing(&s);
4520        Ok(s)
4521    }
4522
4523    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4524    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4525    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4526    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4527        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4528        self.keep_if_capturing(&s);
4529        Ok(s)
4530    }
4531
4532    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4533    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4534    pub fn memset_zeros_view(
4535        &self,
4536        dst: &mut cudarc::driver::CudaViewMut<f32>,
4537    ) -> Result<(), Box<dyn std::error::Error>> {
4538        self.gpu.stream().memset_zeros(dst)?;
4539        Ok(())
4540    }
4541
4542    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4543    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4544    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4545    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4546    /// stream would require an event).
4547    pub fn stage_expert(
4548        &self,
4549        host_bytes: &[u8],
4550        scratch: &mut CudaSlice<u8>,
4551        off: usize,
4552    ) -> Result<(), Box<dyn std::error::Error>> {
4553        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4554        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4555        Ok(())
4556    }
4557
4558    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4559    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4560    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4561    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4562    /// One CTA per token row, 256 threads (one per expert).
4563    pub fn moe_router_topk(
4564        &self,
4565        logits: &CudaSlice<f32>,
4566        t: usize,
4567        n_expert: usize,
4568        n_used: usize,
4569    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4570        let f = self.func("moe_router_topk_f32");
4571        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4572        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4573        let cfg = LaunchConfig {
4574            grid_dim: (t as u32, 1, 1),
4575            block_dim: (n_expert as u32, 1, 1),
4576            shared_mem_bytes: 0,
4577        };
4578        let (ne, nu) = (n_expert as i32, n_used as i32);
4579        let __s_b = self.gpu.stream();
4580        let mut b = __s_b.launch_builder(&f);
4581        b.arg(logits)
4582            .arg(&mut sel_idx)
4583            .arg(&mut sel_w)
4584            .arg(&ne)
4585            .arg(&nu);
4586        unsafe {
4587            b.launch(cfg)?;
4588        }
4589        Ok((sel_idx, sel_w))
4590    }
4591
4592    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4593    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4594    pub fn moe_router_topk_scaled(
4595        &self,
4596        logits: &CudaSlice<f32>,
4597        t: usize,
4598        n_expert: usize,
4599        n_used: usize,
4600        ex_scale: &CudaSlice<f32>,
4601    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4602        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4603        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4604        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4605        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4606        let f = self.func("moe_router_topk_scaled_f32");
4607        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4608        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4609        let cfg = LaunchConfig {
4610            grid_dim: (t as u32, 1, 1),
4611            block_dim: (n_expert as u32, 1, 1),
4612            shared_mem_bytes: 0,
4613        };
4614        let (ne, nu) = (n_expert as i32, n_used as i32);
4615        let __s_b = self.gpu.stream();
4616        let mut b = __s_b.launch_builder(&f);
4617        b.arg(logits)
4618            .arg(&mut sel_idx)
4619            .arg(&mut sel_w)
4620            .arg(&ne)
4621            .arg(&nu)
4622            .arg(ex_scale);
4623        unsafe {
4624            b.launch(cfg)?;
4625        }
4626        Ok((sel_idx, sel_w))
4627    }
4628
4629    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4630    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4631    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4632    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4633    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4634    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4635    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4636    pub fn moe_router_topk_host(
4637        &self,
4638        logits: &CudaSlice<f32>,
4639        t: usize,
4640        n_expert: usize,
4641        n_used: usize,
4642    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4643        let f = self.func("moe_router_topk_f32");
4644        let n = t * n_used;
4645        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4646        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4647        let cfg = LaunchConfig {
4648            grid_dim: (t as u32, 1, 1),
4649            block_dim: (n_expert as u32, 1, 1),
4650            shared_mem_bytes: 0,
4651        };
4652        let (ne, nu) = (n_expert as i32, n_used as i32);
4653        let __s_b = self.gpu.stream();
4654        let mut b = __s_b.launch_builder(&f);
4655        b.arg(logits)
4656            .arg(&mut sel_idx)
4657            .arg(&mut sel_w)
4658            .arg(&ne)
4659            .arg(&nu);
4660        unsafe {
4661            b.launch(cfg)?;
4662        }
4663        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4664        let bytes = n * 8;
4665        let mut guard = self.router_stage.lock().unwrap();
4666        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4667            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4668        }
4669        let stage = guard.as_mut().unwrap();
4670        let (si, sw) = unsafe {
4671            (
4672                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4673                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4674            )
4675        };
4676        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4677        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4678        self.gpu.stream().synchronize()?; // ONE sync for both
4679        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4680    }
4681
4682    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4683    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4684    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4685    #[allow(clippy::too_many_arguments)]
4686    pub fn moe_router_sigmoid_topk(
4687        &self,
4688        logits: &CudaSlice<f32>,
4689        t: usize,
4690        n_expert: usize,
4691        n_used: usize,
4692        active_count: usize,
4693        correction_bias: &CudaSlice<f32>,
4694        active: &CudaSlice<u8>,
4695        scaling_factor: f32,
4696        route_norm: bool,
4697    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4698        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4699        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4700            return Err(format!(
4701                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4702            )
4703            .into());
4704        }
4705        if logits.len() < t * n_expert
4706            || correction_bias.len() != n_expert
4707            || active.len() != n_expert
4708        {
4709            return Err(format!(
4710                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4711                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4712            ).into());
4713        }
4714        let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4715            // Latency twin of the dexp arm (identical outputs): barrier-lean top-k
4716            // over the dexp scoring class. Composes the two doors it rides.
4717            self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4718        } else if crate::sig_expf_dev_on() {
4719            self.func("moe_router_sigmoid_topk_f32_dexp")
4720        } else if crate::topk_fast_on() {
4721            self.func("moe_router_sigmoid_topk_f32_fast")
4722        } else {
4723            self.func("moe_router_sigmoid_topk_f32")
4724        };
4725        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4726        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4727        let threads = n_expert.div_ceil(32) * 32;
4728        let cfg = LaunchConfig {
4729            grid_dim: (t as u32, 1, 1),
4730            block_dim: (threads as u32, 1, 1),
4731            shared_mem_bytes: 0,
4732        };
4733        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4734        let __s_b = self.gpu.stream();
4735        let mut b = __s_b.launch_builder(&f);
4736        b.arg(logits)
4737            .arg(correction_bias)
4738            .arg(active)
4739            .arg(&mut sel_idx)
4740            .arg(&mut sel_w)
4741            .arg(&ne)
4742            .arg(&nu)
4743            .arg(&scaling_factor)
4744            .arg(&rn);
4745        unsafe {
4746            b.launch(cfg)?;
4747        }
4748        Ok((sel_idx, sel_w))
4749    }
4750
4751    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
4752    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
4753    #[allow(clippy::too_many_arguments)]
4754    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
4755    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
4756    /// the model engine can wait on it with a same-device stream memop.
4757    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4758        if ptr == 0 {
4759            return Err("ring_flag_raw: unarmed flag".into());
4760        }
4761        let f = self.func("memra_ring_flag");
4762        let cfg = LaunchConfig {
4763            grid_dim: (1, 1, 1),
4764            block_dim: (32, 1, 1),
4765            shared_mem_bytes: 0,
4766        };
4767        let __s_b = self.gpu.stream();
4768        let mut b = __s_b.launch_builder(&f);
4769        b.arg(&ptr).arg(&value);
4770        unsafe {
4771            b.launch(cfg)?;
4772        }
4773        Ok(())
4774    }
4775
4776    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
4777    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
4778    pub fn moe_sel_w_mirror(
4779        &self,
4780        sel_src: &CudaSlice<i32>,
4781        w_src: &CudaSlice<f32>,
4782        sel_dst: &mut CudaSlice<i32>,
4783        w_dst: &mut CudaSlice<f32>,
4784        n: usize,
4785    ) -> Result<(), Box<dyn std::error::Error>> {
4786        if n == 0
4787            || n > 32
4788            || sel_src.len() < n
4789            || w_src.len() < n
4790            || sel_dst.len() < n
4791            || w_dst.len() < n
4792        {
4793            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4794        }
4795        let f = self.func("moe_sel_w_mirror");
4796        let cfg = LaunchConfig {
4797            grid_dim: (1, 1, 1),
4798            block_dim: (32, 1, 1),
4799            shared_mem_bytes: 0,
4800        };
4801        let ni = n as i32;
4802        let __s_b = self.gpu.stream();
4803        let mut b = __s_b.launch_builder(&f);
4804        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
4805        unsafe {
4806            b.launch(cfg)?;
4807        }
4808        Ok(())
4809    }
4810
4811    pub fn moe_router_sigmoid_topk_into(
4812        &self,
4813        logits: &CudaSlice<f32>,
4814        t: usize,
4815        n_expert: usize,
4816        n_used: usize,
4817        active_count: usize,
4818        correction_bias: &CudaSlice<f32>,
4819        active: &CudaSlice<u8>,
4820        scaling_factor: f32,
4821        route_norm: bool,
4822        sel_idx: &mut CudaSlice<i32>,
4823        sel_w: &mut CudaSlice<f32>,
4824    ) -> Result<(), Box<dyn std::error::Error>> {
4825        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4826        if n_expert == 0
4827            || n_expert > 1024
4828            || n_used == 0
4829            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
4830            || n_used > n_expert
4831            || logits.len() < t * n_expert
4832            || correction_bias.len() != n_expert
4833            || active.len() != n_expert
4834            || sel_idx.len() < t * n_used
4835            || sel_w.len() < t * n_used
4836        {
4837            return Err("sigmoid router _into geometry mismatch".into());
4838        }
4839        let f = if crate::sig_expf_dev_on() {
4840            self.func("moe_router_sigmoid_topk_f32_dexp")
4841        } else if crate::topk_fast_on() {
4842            self.func("moe_router_sigmoid_topk_f32_fast")
4843        } else {
4844            self.func("moe_router_sigmoid_topk_f32")
4845        };
4846        let threads = n_expert.div_ceil(32) * 32;
4847        let cfg = LaunchConfig {
4848            grid_dim: (t as u32, 1, 1),
4849            block_dim: (threads as u32, 1, 1),
4850            shared_mem_bytes: 0,
4851        };
4852        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4853        let __s_b = self.gpu.stream();
4854        let mut b = __s_b.launch_builder(&f);
4855        b.arg(logits)
4856            .arg(correction_bias)
4857            .arg(active)
4858            .arg(&mut *sel_idx)
4859            .arg(&mut *sel_w)
4860            .arg(&ne)
4861            .arg(&nu)
4862            .arg(&scaling_factor)
4863            .arg(&rn);
4864        unsafe {
4865            b.launch(cfg)?;
4866        }
4867        Ok(())
4868    }
4869
4870    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4871    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4872    #[allow(clippy::too_many_arguments)]
4873    pub fn moe_router_sigmoid_topk_host(
4874        &self,
4875        logits: &CudaSlice<f32>,
4876        t: usize,
4877        n_expert: usize,
4878        n_used: usize,
4879        active_count: usize,
4880        correction_bias: &CudaSlice<f32>,
4881        active: &CudaSlice<u8>,
4882        scaling_factor: f32,
4883        route_norm: bool,
4884    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4885        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4886            logits,
4887            t,
4888            n_expert,
4889            n_used,
4890            active_count,
4891            correction_bias,
4892            active,
4893            scaling_factor,
4894            route_norm,
4895        )?;
4896        let n = t * n_used;
4897        let bytes = n * 8;
4898        let mut guard = self.router_stage.lock().unwrap();
4899        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4900            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4901        }
4902        let stage = guard.as_mut().unwrap();
4903        let (si, sw) = unsafe {
4904            (
4905                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4906                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4907            )
4908        };
4909        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4910        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4911        self.gpu.stream().synchronize()?;
4912        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4913    }
4914
4915    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4916    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4917    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4918    pub fn stage_expert_async(
4919        &self,
4920        host_bytes: &[u8],
4921        scratch: &mut CudaSlice<u8>,
4922        off: usize,
4923    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4924        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4925        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4926        Ok(self.copy_stream.record_event(None)?)
4927    }
4928
4929    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4930    pub fn compute_wait(
4931        &self,
4932        ev: &cudarc::driver::CudaEvent,
4933    ) -> Result<(), Box<dyn std::error::Error>> {
4934        self.gpu.stream().wait(ev)?;
4935        Ok(())
4936    }
4937
4938    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4939    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4940    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4941    /// CudaView base+offset pointer is honored by the launch arg.
4942    pub fn qmatvec_view(
4943        &self,
4944        w: &CudaSlice<u8>,
4945        range: std::ops::Range<usize>,
4946        x: &cudarc::driver::CudaView<f32>,
4947        m: usize,
4948        in_f: usize,
4949        out_f: usize,
4950        qtype: i32,
4951        row_bytes: usize,
4952    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4953        let f = self.func("qmatvec_f32");
4954        let wv = w.slice(range); // CudaView<u8>, offset honored
4955        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4956        let cfg = LaunchConfig {
4957            grid_dim: (out_f as u32, m as u32, 1),
4958            block_dim: (256, 1, 1),
4959            shared_mem_bytes: 0,
4960        };
4961        let (inf, outf, mi, qt, rb) =
4962            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4963        let __s_b = self.gpu.stream();
4964        let mut b = __s_b.launch_builder(&f);
4965        b.arg(&wv)
4966            .arg(x)
4967            .arg(&mut y)
4968            .arg(&inf)
4969            .arg(&outf)
4970            .arg(&mi)
4971            .arg(&qt)
4972            .arg(&rb);
4973        unsafe {
4974            b.launch(cfg)?;
4975        }
4976        Ok(y)
4977    }
4978
4979    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4980    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4981    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4982    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4983    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4984    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4985    #[allow(clippy::too_many_arguments)]
4986    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4987    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4988    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4989    pub fn moe_gate_up_silu8_q8(
4990        &self,
4991        gp: WPtr8,
4992        up: WPtr8,
4993        aq: &CudaSlice<i8>,
4994        ad: &CudaSlice<f32>,
4995        in_f: usize,
4996        n_ff: usize,
4997        n_used: usize,
4998        qt_g: i32,
4999        qt_u: i32,
5000        rb_g: usize,
5001        rb_u: usize,
5002    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5003        let f = self.func("moe_gate_up_silu8_q8");
5004        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5005        let cfg = LaunchConfig {
5006            grid_dim: (n_ff as u32, n_used as u32, 1),
5007            block_dim: (32, 1, 1),
5008            shared_mem_bytes: 0,
5009        };
5010        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5011        let __s_b = self.gpu.stream();
5012        let mut b = __s_b.launch_builder(&f);
5013        b.arg(&gp)
5014            .arg(&up)
5015            .arg(aq)
5016            .arg(ad)
5017            .arg(&mut act)
5018            .arg(&inf)
5019            .arg(&nff)
5020            .arg(&qt_g)
5021            .arg(&qt_u)
5022            .arg(&rbg)
5023            .arg(&rbu);
5024        unsafe {
5025            b.launch(cfg)?;
5026        }
5027        Ok(act)
5028    }
5029
5030    #[allow(clippy::too_many_arguments)]
5031    pub fn moe_down8_fma_q8(
5032        &self,
5033        dp: WPtr8,
5034        w: F32x8,
5035        aq2: &CudaSlice<i8>,
5036        ad2: &CudaSlice<f32>,
5037        dst: &mut cudarc::driver::CudaViewMut<f32>,
5038        in_f: usize,
5039        out_f: usize,
5040        n_used: usize,
5041        qt: i32,
5042        rb: usize,
5043    ) -> Result<(), Box<dyn std::error::Error>> {
5044        let f = self.func("moe_down8_fma_q8");
5045        let cfg = LaunchConfig {
5046            grid_dim: (out_f as u32, 1, 1),
5047            block_dim: (32, 1, 1),
5048            shared_mem_bytes: 0,
5049        };
5050        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5051        let __s_b = self.gpu.stream();
5052        let mut b = __s_b.launch_builder(&f);
5053        b.arg(&dp)
5054            .arg(&w)
5055            .arg(aq2)
5056            .arg(ad2)
5057            .arg(dst)
5058            .arg(&inf)
5059            .arg(&outf)
5060            .arg(&nu)
5061            .arg(&qt)
5062            .arg(&rbi);
5063        unsafe {
5064            b.launch(cfg)?;
5065        }
5066        Ok(())
5067    }
5068
5069    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
5070    pub fn qmatvec_expert_q8(
5071        &self,
5072        w: &CudaSlice<u8>,
5073        range: std::ops::Range<usize>,
5074        aq: &CudaSlice<i8>,
5075        ad: &CudaSlice<f32>,
5076        m: usize,
5077        in_f: usize,
5078        out_f: usize,
5079        qtype: i32,
5080        row_bytes: usize,
5081    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5082        let f = self.func("qmatvec_expert_q8");
5083        let wv = w.slice(range);
5084        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5085        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
5086        let cfg = LaunchConfig {
5087            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5088            block_dim: (32, ROWS, 1),
5089            shared_mem_bytes: 0,
5090        };
5091        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5092        let __s_b = self.gpu.stream();
5093        let mut b = __s_b.launch_builder(&f);
5094        b.arg(&wv)
5095            .arg(aq)
5096            .arg(ad)
5097            .arg(&mut y)
5098            .arg(&inf)
5099            .arg(&outf)
5100            .arg(&mi)
5101            .arg(&qtype)
5102            .arg(&rbi);
5103        unsafe {
5104            b.launch(cfg)?;
5105        }
5106        Ok(y)
5107    }
5108
5109    pub fn moe_gate_up_silu8(
5110        &self,
5111        gp: WPtr8,
5112        up: WPtr8,
5113        x: &cudarc::driver::CudaView<f32>,
5114        in_f: usize,
5115        n_ff: usize,
5116        n_used: usize,
5117        qt_g: i32,
5118        qt_u: i32,
5119        rb_g: usize,
5120        rb_u: usize,
5121    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5122        let f = self.func("moe_gate_up_silu8_f32");
5123        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
5124        let cfg = LaunchConfig {
5125            grid_dim: (n_ff as u32, n_used as u32, 1),
5126            block_dim: (256, 1, 1),
5127            shared_mem_bytes: 0,
5128        };
5129        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5130        let __s_b = self.gpu.stream();
5131        let mut b = __s_b.launch_builder(&f);
5132        b.arg(&gp)
5133            .arg(&up)
5134            .arg(x)
5135            .arg(&mut act)
5136            .arg(&inf)
5137            .arg(&nff)
5138            .arg(&qt_g)
5139            .arg(&qt_u)
5140            .arg(&rbg)
5141            .arg(&rbu);
5142        unsafe {
5143            b.launch(cfg)?;
5144        }
5145        Ok(act)
5146    }
5147
5148    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
5149    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
5150    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
5151    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
5152    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
5153    #[allow(clippy::too_many_arguments)]
5154    pub fn moe_down8_fma_into(
5155        &self,
5156        dp: WPtr8,
5157        w: F32x8,
5158        act: &CudaSlice<f32>,
5159        dst: &mut cudarc::driver::CudaViewMut<f32>,
5160        in_f: usize,
5161        out_f: usize,
5162        n_used: usize,
5163        qt: i32,
5164        rb: usize,
5165    ) -> Result<(), Box<dyn std::error::Error>> {
5166        let f = self.func("moe_down8_fma_f32");
5167        let cfg = LaunchConfig {
5168            grid_dim: (out_f as u32, 1, 1),
5169            block_dim: (256, 1, 1),
5170            shared_mem_bytes: 0,
5171        };
5172        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5173        let __s_b = self.gpu.stream();
5174        let mut b = __s_b.launch_builder(&f);
5175        b.arg(&dp)
5176            .arg(&w)
5177            .arg(act)
5178            .arg(dst)
5179            .arg(&inf)
5180            .arg(&outf)
5181            .arg(&nu)
5182            .arg(&qt)
5183            .arg(&rbv);
5184        unsafe {
5185            b.launch(cfg)?;
5186        }
5187        Ok(())
5188    }
5189
5190    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
5191    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
5192    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
5193    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
5194    #[allow(clippy::too_many_arguments)]
5195    /// dp4a q8 twin of the _dev pair (resident-experts arc).
5196    ///
5197    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
5198    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
5199    /// down's FMA chain stays slot-ordered serial). Seams:
5200    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
5201    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
5202    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
5203    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
5204    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
5205    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
5206    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
5207    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
5208    ///                       only) | w8h2 (h2 x slot-parallel)
5209    #[allow(clippy::too_many_arguments)]
5210    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
5211    #[allow(clippy::too_many_arguments)]
5212    pub fn moe_pairs_matvec_q8(
5213        &self,
5214        table: &CudaSlice<u64>,
5215        proj: i32,
5216        pair_tok: &CudaSlice<i32>,
5217        pair_ex: &CudaSlice<i32>,
5218        aq: &CudaSlice<i8>,
5219        ad: &CudaSlice<f32>,
5220        in_f: usize,
5221        out_f: usize,
5222        n_expert: usize,
5223        n_pairs: usize,
5224        qtype: i32,
5225        row_bytes: usize,
5226    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5227        let f = self.func("moe_pairs_matvec_q8");
5228        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5229        const ROWS: u32 = 4;
5230        let cfg = LaunchConfig {
5231            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5232            block_dim: (32, ROWS, 1),
5233            shared_mem_bytes: 0,
5234        };
5235        let (inf, outf, ne, np, rbi) = (
5236            in_f as i32,
5237            out_f as i32,
5238            n_expert as i32,
5239            n_pairs as i32,
5240            row_bytes as i64,
5241        );
5242        let __s_b = self.gpu.stream();
5243        let mut b = __s_b.launch_builder(&f);
5244        b.arg(table)
5245            .arg(&proj)
5246            .arg(pair_tok)
5247            .arg(pair_ex)
5248            .arg(aq)
5249            .arg(ad)
5250            .arg(&mut y)
5251            .arg(&inf)
5252            .arg(&outf)
5253            .arg(&ne)
5254            .arg(&np)
5255            .arg(&qtype)
5256            .arg(&rbi);
5257        unsafe {
5258            b.launch(cfg)?;
5259        }
5260        Ok(y)
5261    }
5262
5263    /// Expert-major pair matvec (weight-reuse across each expert's token group).
5264    #[allow(clippy::too_many_arguments)]
5265    pub fn moe_pairs_matvec_q8_em(
5266        &self,
5267        table: &CudaSlice<u64>,
5268        proj: i32,
5269        ex_ids: &CudaSlice<i32>,
5270        ex_off: &CudaSlice<i32>,
5271        ex_pairs: &CudaSlice<i32>,
5272        pair_tok: &CudaSlice<i32>,
5273        aq: &CudaSlice<i8>,
5274        ad: &CudaSlice<f32>,
5275        in_f: usize,
5276        out_f: usize,
5277        n_expert: usize,
5278        n_active: usize,
5279        n_pairs: usize,
5280        qtype: i32,
5281        row_bytes: usize,
5282    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5283        let f = self.func("moe_pairs_matvec_q8_em");
5284        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5285        const ROWS: u32 = 4;
5286        let cfg = LaunchConfig {
5287            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5288            block_dim: (32, ROWS, 1),
5289            shared_mem_bytes: 0,
5290        };
5291        let (inf, outf, ne, na, rbi) = (
5292            in_f as i32,
5293            out_f as i32,
5294            n_expert as i32,
5295            n_active as i32,
5296            row_bytes as i64,
5297        );
5298        let __s_b = self.gpu.stream();
5299        let mut b = __s_b.launch_builder(&f);
5300        b.arg(table)
5301            .arg(&proj)
5302            .arg(ex_ids)
5303            .arg(ex_off)
5304            .arg(ex_pairs)
5305            .arg(pair_tok)
5306            .arg(aq)
5307            .arg(ad)
5308            .arg(&mut y)
5309            .arg(&inf)
5310            .arg(&outf)
5311            .arg(&ne)
5312            .arg(&na)
5313            .arg(&qtype)
5314            .arg(&rbi);
5315        unsafe {
5316            b.launch(cfg)?;
5317        }
5318        Ok(y)
5319    }
5320
5321    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
5322    // weight group once per (row,group) then dp4a's across the expert's token group.
5323    #[allow(clippy::too_many_arguments)]
5324    pub fn moe_pairs_matvec_q8_dec(
5325        &self,
5326        table: &CudaSlice<u64>,
5327        proj: i32,
5328        ex_ids: &CudaSlice<i32>,
5329        ex_off: &CudaSlice<i32>,
5330        ex_pairs: &CudaSlice<i32>,
5331        pair_tok: &CudaSlice<i32>,
5332        aq: &CudaSlice<i8>,
5333        ad: &CudaSlice<f32>,
5334        in_f: usize,
5335        out_f: usize,
5336        n_expert: usize,
5337        n_active: usize,
5338        n_pairs: usize,
5339        qtype: i32,
5340        row_bytes: usize,
5341    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5342        let f = self.func("moe_pairs_matvec_q8_dec");
5343        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5344        const ROWS: u32 = 4;
5345        let cfg = LaunchConfig {
5346            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5347            block_dim: (32, ROWS, 1),
5348            shared_mem_bytes: 0,
5349        };
5350        let (inf, outf, ne, na, rbi) = (
5351            in_f as i32,
5352            out_f as i32,
5353            n_expert as i32,
5354            n_active as i32,
5355            row_bytes as i64,
5356        );
5357        let __s_b = self.gpu.stream();
5358        let mut b = __s_b.launch_builder(&f);
5359        b.arg(table)
5360            .arg(&proj)
5361            .arg(ex_ids)
5362            .arg(ex_off)
5363            .arg(ex_pairs)
5364            .arg(pair_tok)
5365            .arg(aq)
5366            .arg(ad)
5367            .arg(&mut y)
5368            .arg(&inf)
5369            .arg(&outf)
5370            .arg(&ne)
5371            .arg(&na)
5372            .arg(&qtype)
5373            .arg(&rbi);
5374        unsafe {
5375            b.launch(cfg)?;
5376        }
5377        Ok(y)
5378    }
5379
5380    pub fn moe_pairs_gelu_mul(
5381        &self,
5382        gate: &CudaSlice<f32>,
5383        up: &CudaSlice<f32>,
5384        n: usize,
5385    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5386        let f = self.func("moe_pairs_gelu_mul");
5387        let mut act = self.alloc_uninit::<f32>(n)?;
5388        let cfg = LaunchConfig::for_num_elems(n as u32);
5389        let nl = n as i64;
5390        let __s_b = self.gpu.stream();
5391        let mut b = __s_b.launch_builder(&f);
5392        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5393        unsafe {
5394            b.launch(cfg)?;
5395        }
5396        Ok(act)
5397    }
5398
5399    pub fn moe_pairs_silu_mul(
5400        &self,
5401        gate: &CudaSlice<f32>,
5402        up: &CudaSlice<f32>,
5403        n: usize,
5404    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5405        let f = self.func("moe_pairs_silu_mul");
5406        let mut act = self.alloc_uninit::<f32>(n)?;
5407        let cfg = LaunchConfig::for_num_elems(n as u32);
5408        let nl = n as i64;
5409        let __s_b = self.gpu.stream();
5410        let mut b = __s_b.launch_builder(&f);
5411        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5412        unsafe {
5413            b.launch(cfg)?;
5414        }
5415        Ok(act)
5416    }
5417
5418    #[allow(clippy::too_many_arguments)]
5419    pub fn moe_pairs_scatter(
5420        &self,
5421        y_down: &CudaSlice<f32>,
5422        pair_w: &CudaSlice<f32>,
5423        tok_pair_off: &CudaSlice<i32>,
5424        tok_pair_ids: &CudaSlice<i32>,
5425        moe_out: &mut CudaSlice<f32>,
5426        t: usize,
5427        n_embd: usize,
5428    ) -> Result<(), Box<dyn std::error::Error>> {
5429        let f = self.func("moe_pairs_scatter");
5430        let cfg = LaunchConfig {
5431            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5432            block_dim: (256, 1, 1),
5433            shared_mem_bytes: 0,
5434        };
5435        let ne = n_embd as i32;
5436        let __s_b = self.gpu.stream();
5437        let mut b = __s_b.launch_builder(&f);
5438        b.arg(y_down)
5439            .arg(pair_w)
5440            .arg(tok_pair_off)
5441            .arg(tok_pair_ids)
5442            .arg(moe_out)
5443            .arg(&ne);
5444        unsafe {
5445            b.launch(cfg)?;
5446        }
5447        Ok(())
5448    }
5449
5450    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
5451    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
5452    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
5453    #[allow(clippy::too_many_arguments)]
5454    pub fn moe_gate_up_gelu8_dev_q8(
5455        &self,
5456        table: &CudaSlice<u64>,
5457        sel: &cudarc::driver::CudaView<i32>,
5458        aq: &CudaSlice<i8>,
5459        ad: &CudaSlice<f32>,
5460        in_f: usize,
5461        n_ff: usize,
5462        n_used: usize,
5463        n_expert: usize,
5464        qt_g: i32,
5465        qt_u: i32,
5466        rb_g: usize,
5467        rb_u: usize,
5468    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5469        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5470        let (inf, nff, ne, rbg, rbu) = (
5471            in_f as i32,
5472            n_ff as i32,
5473            n_expert as i32,
5474            rb_g as i64,
5475            rb_u as i64,
5476        );
5477        let f = self.func("moe_gate_up_gelu8_dev_q8");
5478        let cfg = LaunchConfig {
5479            grid_dim: (n_ff as u32, n_used as u32, 1),
5480            block_dim: (32, 1, 1),
5481            shared_mem_bytes: 0,
5482        };
5483        let __s_b = self.gpu.stream();
5484        let mut b = __s_b.launch_builder(&f);
5485        b.arg(table)
5486            .arg(sel)
5487            .arg(aq)
5488            .arg(ad)
5489            .arg(&mut act)
5490            .arg(&inf)
5491            .arg(&nff)
5492            .arg(&ne)
5493            .arg(&qt_g)
5494            .arg(&qt_u)
5495            .arg(&rbg)
5496            .arg(&rbu);
5497        unsafe {
5498            b.launch(cfg)?;
5499        }
5500        Ok(act)
5501    }
5502
5503    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5504    #[allow(clippy::too_many_arguments)]
5505    pub fn moe_gate_up_gelu8_dev_q8_rows(
5506        &self,
5507        table: &CudaSlice<u64>,
5508        sel: &CudaSlice<i32>,
5509        aq: &CudaSlice<i8>,
5510        ad: &CudaSlice<f32>,
5511        t: usize,
5512        in_f: usize,
5513        n_ff: usize,
5514        n_used: usize,
5515        n_expert: usize,
5516        qt_g: i32,
5517        qt_u: i32,
5518        rb_g: usize,
5519        rb_u: usize,
5520    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5521        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5522        let (inf, nff, ne, rbg, rbu, nu) = (
5523            in_f as i32,
5524            n_ff as i32,
5525            n_expert as i32,
5526            rb_g as i64,
5527            rb_u as i64,
5528            n_used as i32,
5529        );
5530        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5531        let cfg = LaunchConfig {
5532            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5533            block_dim: (32, 1, 1),
5534            shared_mem_bytes: 0,
5535        };
5536        let __s_b = self.gpu.stream();
5537        let mut b = __s_b.launch_builder(&f);
5538        b.arg(table)
5539            .arg(sel)
5540            .arg(aq)
5541            .arg(ad)
5542            .arg(&mut act)
5543            .arg(&inf)
5544            .arg(&nff)
5545            .arg(&ne)
5546            .arg(&qt_g)
5547            .arg(&qt_u)
5548            .arg(&rbg)
5549            .arg(&rbu)
5550            .arg(&nu);
5551        unsafe {
5552            b.launch(cfg)?;
5553        }
5554        Ok(act)
5555    }
5556
5557    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5558    #[allow(clippy::too_many_arguments)]
5559    pub fn moe_gate_up_gelu8_dev_q8_csr(
5560        &self,
5561        table: &CudaSlice<u64>,
5562        sel: &CudaSlice<i32>,
5563        aq: &CudaSlice<i8>,
5564        ad: &CudaSlice<f32>,
5565        n_pairs: usize,
5566        in_f: usize,
5567        n_ff: usize,
5568        n_used: usize,
5569        n_expert: usize,
5570        qt_g: i32,
5571        qt_u: i32,
5572        rb_g: usize,
5573        rb_u: usize,
5574    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5575        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5576        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5577            in_f as i32,
5578            n_ff as i32,
5579            n_expert as i32,
5580            rb_g as i64,
5581            rb_u as i64,
5582            n_used as i32,
5583            n_pairs as i32,
5584        );
5585        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5586        let cfg = LaunchConfig {
5587            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5588            block_dim: (32, 1, 1),
5589            shared_mem_bytes: 0,
5590        };
5591        let __s_b = self.gpu.stream();
5592        let mut b = __s_b.launch_builder(&f);
5593        b.arg(table)
5594            .arg(sel)
5595            .arg(aq)
5596            .arg(ad)
5597            .arg(&mut act)
5598            .arg(&inf)
5599            .arg(&nff)
5600            .arg(&ne)
5601            .arg(&qt_g)
5602            .arg(&qt_u)
5603            .arg(&rbg)
5604            .arg(&rbu)
5605            .arg(&nu)
5606            .arg(&npi);
5607        unsafe {
5608            b.launch(cfg)?;
5609        }
5610        Ok(act)
5611    }
5612
5613    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5614    #[allow(clippy::too_many_arguments)]
5615    pub fn moe_down8_fma_dev_q8_rows_g(
5616        &self,
5617        table: &CudaSlice<u64>,
5618        sel: &CudaSlice<i32>,
5619        w: &CudaSlice<f32>,
5620        aq2: &CudaSlice<i8>,
5621        ad2: &CudaSlice<f32>,
5622        dst: &mut CudaSlice<f32>,
5623        t: usize,
5624        in_f: usize,
5625        out_f: usize,
5626        n_used: usize,
5627        n_expert: usize,
5628        qt: i32,
5629        rb: usize,
5630    ) -> Result<(), Box<dyn std::error::Error>> {
5631        let (inf, outf, nu, ne, rbi) = (
5632            in_f as i32,
5633            out_f as i32,
5634            n_used as i32,
5635            n_expert as i32,
5636            rb as i64,
5637        );
5638        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5639        // eight warps, then replay the original slot-ordered FMA chain. Every
5640        // other shape retains the generic one-warp rows kernel.
5641        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5642        let f = self.func(if step_b1_w8 {
5643            "moe_down8_fma_dev_q8_rows_w8"
5644        } else {
5645            "moe_down8_fma_dev_q8_rows_g"
5646        });
5647        let cfg = LaunchConfig {
5648            grid_dim: (out_f as u32, 1, t as u32),
5649            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5650            shared_mem_bytes: 0,
5651        };
5652        let __s_b = self.gpu.stream();
5653        let mut b = __s_b.launch_builder(&f);
5654        b.arg(table)
5655            .arg(sel)
5656            .arg(w)
5657            .arg(aq2)
5658            .arg(ad2)
5659            .arg(dst)
5660            .arg(&inf)
5661            .arg(&outf)
5662            .arg(&nu)
5663            .arg(&ne)
5664            .arg(&qt)
5665            .arg(&rbi);
5666        unsafe {
5667            b.launch(cfg)?;
5668        }
5669        Ok(())
5670    }
5671
5672    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5673    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5674    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5675    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5676        let (out_f, in_f) = (2048usize, 2816usize);
5677        let nblk = in_f / 32;
5678        let mut seed = 0x9E3779B97F4A7C15u64;
5679        let mut rng = move || {
5680            seed = seed
5681                .wrapping_mul(6364136223846793005)
5682                .wrapping_add(1442695040888963407);
5683            (seed >> 33) as u8
5684        };
5685        let mut w = vec![0u8; out_f * nblk * 18];
5686        for b in w.iter_mut() {
5687            *b = rng();
5688        }
5689        for r in 0..out_f {
5690            for g in 0..nblk {
5691                let off = (r * nblk + g) * 18;
5692                w[off] = 0x00;
5693                w[off + 1] = 0x2C; // sane half d
5694            }
5695        }
5696        let qplane = out_f * nblk * 16;
5697        let mut wrp = vec![0u8; w.len()];
5698        for r in 0..out_f {
5699            for g in 0..nblk {
5700                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5701                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5702                    .copy_from_slice(&src[0..2]);
5703                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5704            }
5705        }
5706        let w_d = self.htod_bytes(&w)?;
5707        let wrp_d = self.htod_bytes(&wrp)?;
5708        let mut aq = vec![0i8; m * in_f];
5709        for v in aq.iter_mut() {
5710            *v = rng() as i8;
5711        }
5712        let aq_d = self.htod_i8(&aq)?;
5713        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5714        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5715        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5716        const RPB: u32 = 4;
5717        let cfg = LaunchConfig {
5718            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5719            block_dim: (32, RPB, 1),
5720            shared_mem_bytes: 0,
5721        };
5722        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5723        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5724        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5725        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5726        {
5727            let __s_b = self.gpu.stream();
5728            let mut b = __s_b.launch_builder(&fb);
5729            b.arg(&w_d)
5730                .arg(&aq_d)
5731                .arg(&ad_d)
5732                .arg(&mut y0)
5733                .arg(&inf)
5734                .arg(&outf)
5735                .arg(&mi)
5736                .arg(&rb);
5737            unsafe {
5738                b.launch(cfg)?;
5739            }
5740            let __s_b = self.gpu.stream();
5741            let mut b = __s_b.launch_builder(&fr);
5742            b.arg(&wrp_d)
5743                .arg(&aq_d)
5744                .arg(&ad_d)
5745                .arg(&mut y1)
5746                .arg(&inf)
5747                .arg(&outf)
5748                .arg(&mi)
5749                .arg(&qp);
5750            unsafe {
5751                b.launch(cfg)?;
5752            }
5753        }
5754        self.gpu.stream().synchronize()?;
5755        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5756        let nd = h0
5757            .iter()
5758            .zip(&h1)
5759            .filter(|(a, b)| a.to_bits() != b.to_bits())
5760            .count();
5761        if nd != 0 {
5762            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5763        }
5764        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5765            self.gpu.stream().synchronize()?;
5766            let t0 = std::time::Instant::now();
5767            for _ in 0..500 {
5768                if rp {
5769                    let __s_b = self.gpu.stream();
5770                    let mut b = __s_b.launch_builder(&fr);
5771                    b.arg(&wrp_d)
5772                        .arg(&aq_d)
5773                        .arg(&ad_d)
5774                        .arg(&mut y1)
5775                        .arg(&inf)
5776                        .arg(&outf)
5777                        .arg(&mi)
5778                        .arg(&qp);
5779                    unsafe {
5780                        b.launch(cfg)?;
5781                    }
5782                } else {
5783                    let __s_b = self.gpu.stream();
5784                    let mut b = __s_b.launch_builder(&fb);
5785                    b.arg(&w_d)
5786                        .arg(&aq_d)
5787                        .arg(&ad_d)
5788                        .arg(&mut y0)
5789                        .arg(&inf)
5790                        .arg(&outf)
5791                        .arg(&mi)
5792                        .arg(&rb);
5793                    unsafe {
5794                        b.launch(cfg)?;
5795                    }
5796                }
5797            }
5798            self.gpu.stream().synchronize()?;
5799            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5800        };
5801        let _ = time(false)?;
5802        let _ = time(true)?; // warm
5803        Ok((time(false)?, time(true)?))
5804    }
5805
5806    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5807    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5808    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5809    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5810    pub fn build_q4_rp4(
5811        &self,
5812        t: &mut crate::model::GpuTensor,
5813    ) -> Result<(), Box<dyn std::error::Error>> {
5814        use crate::model::GpuTensor;
5815        let GpuTensor::Quant {
5816            bytes,
5817            qtype,
5818            row_bytes,
5819            ne,
5820            rp4,
5821            ..
5822        } = t
5823        else {
5824            return Ok(());
5825        };
5826        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5827            return Ok(());
5828        }
5829        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5830        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5831            return Ok(());
5832        }
5833        let nblk = in_f / 32;
5834        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5835        let f = self.func("q4_0_split_rp_build");
5836        let n = (out_f * nblk) as i32;
5837        let cfg = LaunchConfig {
5838            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5839            block_dim: (256, 1, 1),
5840            shared_mem_bytes: 0,
5841        };
5842        let (of, nb) = (out_f as i32, nblk as i32);
5843        let _ = n;
5844        let __s_b = self.gpu.stream();
5845        let mut b = __s_b.launch_builder(&f);
5846        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5847        unsafe {
5848            b.launch(cfg)?;
5849        }
5850        *rp4 = Some(dst);
5851        Ok(())
5852    }
5853
5854    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5855    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5856    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5857    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5858    pub fn build_q8_rp4(
5859        &self,
5860        t: &mut crate::model::GpuTensor,
5861    ) -> Result<(), Box<dyn std::error::Error>> {
5862        use crate::model::GpuTensor;
5863        let GpuTensor::Quant {
5864            bytes,
5865            qtype,
5866            row_bytes,
5867            ne,
5868            rp4,
5869            ..
5870        } = t
5871        else {
5872            return Ok(());
5873        };
5874        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5875            return Ok(());
5876        }
5877        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5878        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5879            return Ok(());
5880        }
5881        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5882        Ok(())
5883    }
5884
5885    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5886    /// mirror without a GpuTensor (same kernel the loader path above uses).
5887    pub fn build_q8_rp4_raw(
5888        &self,
5889        bytes: &CudaSlice<u8>,
5890        in_f: usize,
5891        out_f: usize,
5892    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5893        assert!(in_f % 32 == 0);
5894        let nblk = in_f / 32;
5895        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5896        let f = self.func("q8_0_split_rp_build");
5897        let cfg = LaunchConfig {
5898            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5899            block_dim: (256, 1, 1),
5900            shared_mem_bytes: 0,
5901        };
5902        let (of, nb) = (out_f as i32, nblk as i32);
5903        let __s_b = self.gpu.stream();
5904        let mut b = __s_b.launch_builder(&f);
5905        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5906        unsafe {
5907            b.launch(cfg)?;
5908        }
5909        Ok(dst)
5910    }
5911
5912    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5913    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5914    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5915    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5916    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5917    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5918    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5919    pub fn build_q4k_rp4(
5920        &self,
5921        t: &mut crate::model::GpuTensor,
5922    ) -> Result<(), Box<dyn std::error::Error>> {
5923        use crate::model::GpuTensor;
5924        let GpuTensor::Quant {
5925            bytes,
5926            qtype,
5927            row_bytes,
5928            ne,
5929            rp4,
5930            ..
5931        } = t
5932        else {
5933            return Ok(());
5934        };
5935        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5936            return Ok(());
5937        }
5938        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5939        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5940            return Ok(());
5941        }
5942        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5943        Ok(())
5944    }
5945
5946    pub fn build_q6k_rp4(
5947        &self,
5948        t: &mut crate::model::GpuTensor,
5949    ) -> Result<(), Box<dyn std::error::Error>> {
5950        use crate::model::GpuTensor;
5951        let GpuTensor::Quant {
5952            bytes,
5953            qtype,
5954            row_bytes,
5955            ne,
5956            rp4,
5957            ..
5958        } = t
5959        else {
5960            return Ok(());
5961        };
5962        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5963            return Ok(());
5964        }
5965        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5966        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5967            return Ok(());
5968        }
5969        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5970        Ok(())
5971    }
5972
5973    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5974    pub fn build_kq_rp4_raw(
5975        &self,
5976        bytes: &CudaSlice<u8>,
5977        in_f: usize,
5978        out_f: usize,
5979        qtype: i32,
5980    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5981        assert!(in_f % 256 == 0);
5982        let nsbk = in_f / 256;
5983        let (sb_bytes, kname) = match qtype {
5984            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5985            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5986            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5987        };
5988        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5989        let f = self.func(kname);
5990        let cfg = LaunchConfig {
5991            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5992            block_dim: (256, 1, 1),
5993            shared_mem_bytes: 0,
5994        };
5995        let (of, nb) = (out_f as i32, nsbk as i32);
5996        let __s_b = self.gpu.stream();
5997        let mut b = __s_b.launch_builder(&f);
5998        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5999        unsafe {
6000            b.launch(cfg)?;
6001        }
6002        Ok(dst)
6003    }
6004
6005    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
6006    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
6007    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
6008    pub fn kqrp_enabled() -> bool {
6009        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6010        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6011            Ok("0") => false,
6012            Ok(_) => true,
6013            Err(_) => cfg!(memra_hopper_mma),
6014        })
6015    }
6016
6017    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
6018    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
6019    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
6020    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
6021    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
6022    pub fn build_q4_rp_swap(
6023        &self,
6024        t: &mut crate::model::GpuTensor,
6025    ) -> Result<bool, Box<dyn std::error::Error>> {
6026        use crate::model::GpuTensor;
6027        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
6028        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
6029        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
6030        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
6031        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
6032        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
6033        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
6034        // this fn's OWN builder serves may ever be swapped; everything else refuses
6035        // here, regardless of walk ordering.
6036        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6037            return Ok(false);
6038        }
6039        self.build_q4_rp4(t)?;
6040        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
6041        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6042            return Ok(false);
6043        };
6044        match rp4.take() {
6045            Some(split) => {
6046                *bytes = split; // the GGUF-layout buffer drops here
6047                *rp = true;
6048                Ok(true)
6049            }
6050            None => Ok(false),
6051        }
6052    }
6053
6054    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
6055    pub fn q4rp_enabled() -> bool {
6056        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6057        *ON.get_or_init(|| {
6058            std::env::var("MEMRA_Q4RP")
6059                .map(|v| v != "0")
6060                .unwrap_or(true)
6061        })
6062    }
6063
6064    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
6065    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
6066    pub fn copy_rows_strided(
6067        &self,
6068        src: &CudaSlice<f32>,
6069        dst: &mut CudaSlice<f32>,
6070        row_elems: usize,
6071        n_rows: usize,
6072        src_stride: usize,
6073        src_off: usize,
6074    ) -> Result<(), Box<dyn std::error::Error>> {
6075        let f = self.func("copy_rows_strided_f32");
6076        let cfg = LaunchConfig {
6077            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6078            block_dim: (256, 1, 1),
6079            shared_mem_bytes: 0,
6080        };
6081        let (re, nr) = (row_elems as i32, n_rows as i32);
6082        let (st, off) = (src_stride as i64, src_off as i64);
6083        let __s_b = self.gpu.stream();
6084        let mut b = __s_b.launch_builder(&f);
6085        b.arg(src)
6086            .arg(&mut *dst)
6087            .arg(&re)
6088            .arg(&nr)
6089            .arg(&st)
6090            .arg(&off);
6091        unsafe {
6092            b.launch(cfg)?;
6093        }
6094        Ok(())
6095    }
6096
6097    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
6098    ///
6099    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
6100    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
6101    /// one peer copy per token.
6102    pub fn place_rows_strided(
6103        &self,
6104        src: &CudaSlice<f32>,
6105        dst: &mut CudaSlice<f32>,
6106        row_elems: usize,
6107        n_rows: usize,
6108        dst_stride: usize,
6109        dst_off: usize,
6110    ) -> Result<(), Box<dyn std::error::Error>> {
6111        if row_elems == 0 || n_rows == 0 {
6112            return Err("strided row placement requires nonzero rows and row width".into());
6113        }
6114        let src_len = n_rows
6115            .checked_mul(row_elems)
6116            .ok_or("strided row placement source size overflow")?;
6117        let dst_len = n_rows
6118            .checked_sub(1)
6119            .and_then(|rows| rows.checked_mul(dst_stride))
6120            .and_then(|base| base.checked_add(dst_off))
6121            .and_then(|base| base.checked_add(row_elems))
6122            .ok_or("strided row placement destination size overflow")?;
6123        let row_end = dst_off
6124            .checked_add(row_elems)
6125            .ok_or("strided row placement row size overflow")?;
6126        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6127            return Err(format!(
6128                "strided row placement geometry mismatch: src={} need_src={src_len} \
6129                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6130                 dst_stride={dst_stride} dst_off={dst_off}",
6131                src.len(),
6132                dst.len(),
6133            )
6134            .into());
6135        }
6136        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6137            return Err("strided row placement exceeds CUDA kernel geometry".into());
6138        }
6139        let f = self.func("place_rows_strided_f32");
6140        let cfg = LaunchConfig {
6141            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6142            block_dim: (256, 1, 1),
6143            shared_mem_bytes: 0,
6144        };
6145        let (re, nr) = (row_elems as i32, n_rows as i32);
6146        let (st, off) = (dst_stride as i64, dst_off as i64);
6147        let __s_b = self.gpu.stream();
6148        let mut b = __s_b.launch_builder(&f);
6149        b.arg(src)
6150            .arg(&mut *dst)
6151            .arg(&re)
6152            .arg(&nr)
6153            .arg(&st)
6154            .arg(&off);
6155        unsafe {
6156            b.launch(cfg)?;
6157        }
6158        Ok(())
6159    }
6160
6161    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
6162    pub fn u32_set_k(
6163        &self,
6164        dst: &mut CudaSlice<u32>,
6165        v: u32,
6166        idx: usize,
6167    ) -> Result<(), Box<dyn std::error::Error>> {
6168        let f = self.func("u32_set_k");
6169        let cfg = LaunchConfig {
6170            grid_dim: (1, 1, 1),
6171            block_dim: (1, 1, 1),
6172            shared_mem_bytes: 0,
6173        };
6174        let ii = idx as i32;
6175        let __s_b = self.gpu.stream();
6176        let mut b = __s_b.launch_builder(&f);
6177        b.arg(dst).arg(&v).arg(&ii);
6178        unsafe {
6179            b.launch(cfg)?;
6180        }
6181        Ok(())
6182    }
6183
6184    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
6185    pub fn i32_add_k(
6186        &self,
6187        d: &mut CudaSlice<i32>,
6188        v: i32,
6189    ) -> Result<(), Box<dyn std::error::Error>> {
6190        let f = self.func("i32_add_k");
6191        let cfg = LaunchConfig {
6192            grid_dim: (1, 1, 1),
6193            block_dim: (32, 1, 1),
6194            shared_mem_bytes: 0,
6195        };
6196        let __s_b = self.gpu.stream();
6197        let mut b = __s_b.launch_builder(&f);
6198        b.arg(d).arg(&v);
6199        unsafe {
6200            b.launch(cfg)?;
6201        }
6202        Ok(())
6203    }
6204
6205    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
6206    pub fn i32_iota_from(
6207        &self,
6208        ctr: &CudaSlice<i32>,
6209        dst: &mut CudaSlice<i32>,
6210        n: usize,
6211    ) -> Result<(), Box<dyn std::error::Error>> {
6212        let f = self.func("i32_iota_from");
6213        let cfg = LaunchConfig::for_num_elems(n as u32);
6214        let ni = n as i32;
6215        let __s_b = self.gpu.stream();
6216        let mut b = __s_b.launch_builder(&f);
6217        b.arg(ctr).arg(dst).arg(&ni);
6218        unsafe {
6219            b.launch(cfg)?;
6220        }
6221        Ok(())
6222    }
6223
6224    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
6225    pub fn u32_map_k(
6226        &self,
6227        buf: &mut CudaSlice<u32>,
6228        map: &CudaSlice<u32>,
6229        idx: usize,
6230    ) -> Result<(), Box<dyn std::error::Error>> {
6231        let f = self.func("u32_map_k");
6232        let cfg = LaunchConfig {
6233            grid_dim: (1, 1, 1),
6234            block_dim: (1, 1, 1),
6235            shared_mem_bytes: 0,
6236        };
6237        let ii = idx as i32;
6238        let __s_b = self.gpu.stream();
6239        let mut b = __s_b.launch_builder(&f);
6240        b.arg(buf).arg(map).arg(&ii);
6241        unsafe {
6242            b.launch(cfg)?;
6243        }
6244        Ok(())
6245    }
6246
6247    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
6248    #[allow(clippy::too_many_arguments)]
6249    pub fn u32_pack2(
6250        &self,
6251        a: &CudaSlice<u32>,
6252        off_a: usize,
6253        n1: usize,
6254        b_in: &CudaSlice<u32>,
6255        n2: usize,
6256        out: &mut CudaSlice<u32>,
6257    ) -> Result<(), Box<dyn std::error::Error>> {
6258        let f = self.func("u32_pack2");
6259        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6260        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6261        let __s_b = self.gpu.stream();
6262        let mut b = __s_b.launch_builder(&f);
6263        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6264        unsafe {
6265            b.launch(cfg)?;
6266        }
6267        Ok(())
6268    }
6269
6270    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
6271    pub fn moe_w_exscale(
6272        &self,
6273        w: &mut CudaSlice<f32>,
6274        sel: &CudaSlice<i32>,
6275        s: &CudaSlice<f32>,
6276        n: usize,
6277    ) -> Result<(), Box<dyn std::error::Error>> {
6278        let f = self.func("moe_w_exscale");
6279        let cfg = LaunchConfig::for_num_elems(n as u32);
6280        let ni = n as i32;
6281        let __s_b = self.gpu.stream();
6282        let mut b = __s_b.launch_builder(&f);
6283        b.arg(w).arg(sel).arg(s).arg(&ni);
6284        unsafe {
6285            b.launch(cfg)?;
6286        }
6287        Ok(())
6288    }
6289
6290    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
6291    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
6292    pub fn moe_w_scale_by_expert(
6293        &self,
6294        w: &mut CudaSlice<f32>,
6295        sel: &CudaSlice<i32>,
6296        macros: &CudaSlice<f32>,
6297        n_expert: usize,
6298        n: usize,
6299    ) -> Result<(), Box<dyn std::error::Error>> {
6300        let f = self.func("moe_w_scale_by_expert");
6301        let cfg = LaunchConfig {
6302            grid_dim: (n.div_ceil(64) as u32, 1, 1),
6303            block_dim: (64, 1, 1),
6304            shared_mem_bytes: 0,
6305        };
6306        let (ne, nn) = (n_expert as i32, n as i32);
6307        let __s_b = self.gpu.stream();
6308        let mut b = __s_b.launch_builder(&f);
6309        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6310        unsafe {
6311            b.launch(cfg)?;
6312        }
6313        Ok(())
6314    }
6315
6316    pub fn moe_gate_up_silu8_dev_q8(
6317        &self,
6318        table: &CudaSlice<u64>,
6319        sel: &cudarc::driver::CudaView<i32>,
6320        aq: &CudaSlice<i8>,
6321        ad: &CudaSlice<f32>,
6322        in_f: usize,
6323        n_ff: usize,
6324        n_used: usize,
6325        n_expert: usize,
6326        qt_g: i32,
6327        qt_u: i32,
6328        rb_g: usize,
6329        rb_u: usize,
6330        macros: &CudaSlice<f32>,
6331    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6332        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6333        let (mode, wpb) = GU.get_or_init(|| {
6334            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6335            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6336                .ok()
6337                .and_then(|v| v.parse().ok())
6338                .unwrap_or(4u32)
6339                .clamp(1, 16);
6340            (mode, wpb)
6341        });
6342        let (mode, wpb) = (mode.as_str(), *wpb);
6343        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6344        let (inf, nff, ne, rbg, rbu) = (
6345            in_f as i32,
6346            n_ff as i32,
6347            n_expert as i32,
6348            rb_g as i64,
6349            rb_u as i64,
6350        );
6351        let (f, cfg) = match mode {
6352            "1" | "2" | "4" => {
6353                let rpw: u32 = mode.parse().unwrap();
6354                let f = self.func(match rpw {
6355                    1 => "moe_gate_up_silu8_dev_q8_r1",
6356                    2 => "moe_gate_up_silu8_dev_q8_r2",
6357                    _ => "moe_gate_up_silu8_dev_q8_r4",
6358                });
6359                let rows_per_block = (rpw * wpb) as usize;
6360                let gx = n_ff.div_ceil(rows_per_block) as u32;
6361                (
6362                    f,
6363                    LaunchConfig {
6364                        grid_dim: (gx, n_used as u32, 1),
6365                        block_dim: (32, wpb, 1),
6366                        shared_mem_bytes: 0,
6367                    },
6368                )
6369            }
6370            "j8" if n_used <= 32 => (
6371                self.func("moe_gate_up_silu8_dev_q8_j8"),
6372                LaunchConfig {
6373                    grid_dim: (n_ff as u32, 1, 1),
6374                    block_dim: (32, n_used as u32, 1),
6375                    shared_mem_bytes: 0,
6376                },
6377            ),
6378            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
6379            "vsm2" => {
6380                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6381                let sh = (rb_g + rb_u) as u32;
6382                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6383                f.set_attribute(
6384                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6385                    sh as i32,
6386                )?;
6387                (
6388                    f,
6389                    LaunchConfig {
6390                        grid_dim: (n_ff as u32, n_used as u32, 1),
6391                        block_dim: (32, 1, 1),
6392                        shared_mem_bytes: sh,
6393                    },
6394                )
6395            }
6396            "vsm" => {
6397                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6398                let sh = (rb_g + rb_u) as u32;
6399                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6400                f.set_attribute(
6401                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6402                    sh as i32,
6403                )?;
6404                (
6405                    f,
6406                    LaunchConfig {
6407                        grid_dim: (n_ff as u32, n_used as u32, 1),
6408                        block_dim: (32, 1, 1),
6409                        shared_mem_bytes: sh,
6410                    },
6411                )
6412            }
6413            "sg" => (
6414                self.func("moe_gate_up_silu8_dev_q8_sg"),
6415                LaunchConfig {
6416                    grid_dim: (n_ff as u32, n_used as u32, 1),
6417                    block_dim: (32, 1, 1),
6418                    shared_mem_bytes: 0,
6419                },
6420            ),
6421            "j8sg" if n_used <= 32 => (
6422                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6423                LaunchConfig {
6424                    grid_dim: (n_ff as u32, 1, 1),
6425                    block_dim: (32, n_used as u32, 1),
6426                    shared_mem_bytes: 0,
6427                },
6428            ),
6429            "u64" if in_f == 2048 => (
6430                self.func("moe_gate_up_silu8_dev_q8_u64"),
6431                LaunchConfig {
6432                    grid_dim: (n_ff as u32, n_used as u32, 1),
6433                    block_dim: (32, 1, 1),
6434                    shared_mem_bytes: 0,
6435                },
6436            ),
6437            "gs4" if in_f == 2048 => (
6438                self.func("moe_gate_up_silu8_dev_q8_gs4"),
6439                LaunchConfig {
6440                    grid_dim: (n_ff as u32, n_used as u32, 1),
6441                    block_dim: (32, 4, 1),
6442                    shared_mem_bytes: 0,
6443                },
6444            ),
6445            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
6446            "v" | "" => (
6447                self.func("moe_gate_up_silu8_dev_q8_v"),
6448                LaunchConfig {
6449                    grid_dim: (n_ff as u32, n_used as u32, 1),
6450                    block_dim: (32, 1, 1),
6451                    shared_mem_bytes: 0,
6452                },
6453            ),
6454            "s2" => (
6455                self.func("moe_gate_up_silu8_dev_q8_s2"),
6456                LaunchConfig {
6457                    grid_dim: (n_ff as u32, n_used as u32, 1),
6458                    block_dim: (32, 2, 1),
6459                    shared_mem_bytes: 0,
6460                },
6461            ),
6462            "s2z" => {
6463                let rz = wpb.min(16); // s2z smem tile is [16][2]
6464                (
6465                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
6466                    LaunchConfig {
6467                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6468                        block_dim: (32, 2, rz),
6469                        shared_mem_bytes: 0,
6470                    },
6471                )
6472            }
6473            _ => (
6474                self.func("moe_gate_up_silu8_dev_q8"),
6475                LaunchConfig {
6476                    grid_dim: (n_ff as u32, n_used as u32, 1),
6477                    block_dim: (32, 1, 1),
6478                    shared_mem_bytes: 0,
6479                },
6480            ),
6481        };
6482        let __s_b = self.gpu.stream();
6483        let mut b = __s_b.launch_builder(&f);
6484        b.arg(table)
6485            .arg(sel)
6486            .arg(aq)
6487            .arg(ad)
6488            .arg(&mut act)
6489            .arg(&inf)
6490            .arg(&nff)
6491            .arg(&ne)
6492            .arg(&qt_g)
6493            .arg(&qt_u)
6494            .arg(&rbg)
6495            .arg(&rbu)
6496            .arg(macros);
6497        unsafe {
6498            b.launch(cfg)?;
6499        }
6500        Ok(act)
6501    }
6502
6503    #[allow(clippy::too_many_arguments)]
6504    pub fn moe_down8_fma_dev_q8(
6505        &self,
6506        table: &CudaSlice<u64>,
6507        sel: &cudarc::driver::CudaView<i32>,
6508        w: &cudarc::driver::CudaView<f32>,
6509        aq2: &CudaSlice<i8>,
6510        ad2: &CudaSlice<f32>,
6511        dst: &mut cudarc::driver::CudaViewMut<f32>,
6512        in_f: usize,
6513        out_f: usize,
6514        n_used: usize,
6515        n_expert: usize,
6516        qt: i32,
6517        rb: usize,
6518    ) -> Result<(), Box<dyn std::error::Error>> {
6519        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6520        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6521        let (inf, outf, nu, ne, rbi) = (
6522            in_f as i32,
6523            out_f as i32,
6524            n_used as i32,
6525            n_expert as i32,
6526            rb as i64,
6527        );
6528        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
6529        // the h2 twins are nsb==16 (in_f==512) shape-gated.
6530        let (f, cfg) = match mode.as_str() {
6531            m @ ("1" | "2" | "4") if n_used <= 8 => {
6532                let rpw: usize = m.parse().unwrap();
6533                let f = self.func(match rpw {
6534                    1 => "moe_down8_fma_dev_q8_w8r1",
6535                    2 => "moe_down8_fma_dev_q8_w8r2",
6536                    _ => "moe_down8_fma_dev_q8_w8r4",
6537                });
6538                (
6539                    f,
6540                    LaunchConfig {
6541                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6542                        block_dim: (32, n_used as u32, 1),
6543                        shared_mem_bytes: 0,
6544                    },
6545                )
6546            }
6547            "h2" if in_f == 512 => (
6548                self.func("moe_down8_fma_dev_q8_h2"),
6549                LaunchConfig {
6550                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6551                    block_dim: (32, 1, 1),
6552                    shared_mem_bytes: 0,
6553                },
6554            ),
6555            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6556            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6557            "" if in_f == 704 && n_used <= 8 => (
6558                self.func("moe_down8_fma_dev_q8_w8r2"),
6559                LaunchConfig {
6560                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6561                    block_dim: (32, n_used as u32, 1),
6562                    shared_mem_bytes: 0,
6563                },
6564            ),
6565            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6566            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6567            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6568            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6569                self.func("moe_down8_fma_dev_q8_w8h2v"),
6570                LaunchConfig {
6571                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6572                    block_dim: (32, n_used as u32, 1),
6573                    shared_mem_bytes: 0,
6574                },
6575            ),
6576            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6577                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6578                LaunchConfig {
6579                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6580                    block_dim: (32, n_used as u32, 1),
6581                    shared_mem_bytes: 0,
6582                },
6583            ),
6584            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6585                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6586                LaunchConfig {
6587                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6588                    block_dim: (32, n_used as u32, 1),
6589                    shared_mem_bytes: 0,
6590                },
6591            ),
6592            "w8h2" if in_f == 512 && n_used <= 8 => (
6593                self.func("moe_down8_fma_dev_q8_w8h2"),
6594                LaunchConfig {
6595                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6596                    block_dim: (32, n_used as u32, 1),
6597                    shared_mem_bytes: 0,
6598                },
6599            ),
6600            _ => (
6601                self.func("moe_down8_fma_dev_q8"),
6602                LaunchConfig {
6603                    grid_dim: (out_f as u32, 1, 1),
6604                    block_dim: (32, 1, 1),
6605                    shared_mem_bytes: 0,
6606                },
6607            ),
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(w)
6614            .arg(aq2)
6615            .arg(ad2)
6616            .arg(dst)
6617            .arg(&inf)
6618            .arg(&outf)
6619            .arg(&nu)
6620            .arg(&ne)
6621            .arg(&qt)
6622            .arg(&rbi);
6623        unsafe {
6624            b.launch(cfg)?;
6625        }
6626        Ok(())
6627    }
6628
6629    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6630    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6631    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6632    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6633    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6634    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6635    #[allow(clippy::too_many_arguments)]
6636    pub fn moe_gate_up_silu8_dev_q8_rows(
6637        &self,
6638        table: &CudaSlice<u64>,
6639        sel: &CudaSlice<i32>,
6640        aq: &CudaSlice<i8>,
6641        ad: &CudaSlice<f32>,
6642        t: usize,
6643        in_f: usize,
6644        n_ff: usize,
6645        n_used: usize,
6646        n_expert: usize,
6647        qt_g: i32,
6648        qt_u: i32,
6649        rb_g: usize,
6650        rb_u: usize,
6651        macros: &CudaSlice<f32>,
6652    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6653        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6654        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6655        let cfg = LaunchConfig {
6656            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6657            block_dim: (32, 1, 1),
6658            shared_mem_bytes: 0,
6659        };
6660        let (inf, nff, ne, nu, rbg, rbu) = (
6661            in_f as i32,
6662            n_ff as i32,
6663            n_expert as i32,
6664            n_used as i32,
6665            rb_g as i64,
6666            rb_u as i64,
6667        );
6668        let __s_b = self.gpu.stream();
6669        let mut b = __s_b.launch_builder(&f);
6670        b.arg(table)
6671            .arg(sel)
6672            .arg(aq)
6673            .arg(ad)
6674            .arg(&mut act)
6675            .arg(&inf)
6676            .arg(&nff)
6677            .arg(&ne)
6678            .arg(&qt_g)
6679            .arg(&qt_u)
6680            .arg(&rbg)
6681            .arg(&rbu)
6682            .arg(&nu)
6683            .arg(macros);
6684        unsafe {
6685            b.launch(cfg)?;
6686        }
6687        Ok(act)
6688    }
6689
6690    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6691    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6692    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6693    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6694    #[allow(clippy::too_many_arguments)]
6695    pub fn moe_down8_fma_dev_q8_rows(
6696        &self,
6697        table: &CudaSlice<u64>,
6698        sel: &CudaSlice<i32>,
6699        w: &CudaSlice<f32>,
6700        aq2: &CudaSlice<i8>,
6701        ad2: &CudaSlice<f32>,
6702        dst: &mut CudaSlice<f32>,
6703        t: usize,
6704        in_f: usize,
6705        out_f: usize,
6706        n_used: usize,
6707        n_expert: usize,
6708        qt: i32,
6709        rb: usize,
6710    ) -> Result<(), Box<dyn std::error::Error>> {
6711        assert!(
6712            in_f == 512 && n_used <= 8,
6713            "down rows twin is w8h2v shape-gated"
6714        );
6715        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6716        let cfg = LaunchConfig {
6717            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6718            block_dim: (32, n_used as u32, 1),
6719            shared_mem_bytes: 0,
6720        };
6721        let (inf, outf, nu, ne, rbi) = (
6722            in_f as i32,
6723            out_f as i32,
6724            n_used as i32,
6725            n_expert as i32,
6726            rb as i64,
6727        );
6728        let __s_b = self.gpu.stream();
6729        let mut b = __s_b.launch_builder(&f);
6730        b.arg(table)
6731            .arg(sel)
6732            .arg(w)
6733            .arg(aq2)
6734            .arg(ad2)
6735            .arg(dst)
6736            .arg(&inf)
6737            .arg(&outf)
6738            .arg(&nu)
6739            .arg(&ne)
6740            .arg(&qt)
6741            .arg(&rbi);
6742        unsafe {
6743            b.launch(cfg)?;
6744        }
6745        Ok(())
6746    }
6747
6748    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6749    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6750    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6751    #[allow(clippy::too_many_arguments)]
6752    pub fn moe_gate_up_silu8_dev_q8_csr(
6753        &self,
6754        table: &CudaSlice<u64>,
6755        sel: &CudaSlice<i32>,
6756        aq: &CudaSlice<i8>,
6757        ad: &CudaSlice<f32>,
6758        n_pairs: usize,
6759        in_f: usize,
6760        n_ff: usize,
6761        n_used: usize,
6762        n_expert: usize,
6763        qt_g: i32,
6764        qt_u: i32,
6765        rb_g: usize,
6766        rb_u: usize,
6767    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6768        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6769        // host gate guarantees qt_g == qt_u within a supported class.
6770        let f = if qt_g == crate::QT_NVFP4 {
6771            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6772        } else {
6773            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6774        };
6775        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6776        let cfg = LaunchConfig {
6777            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6778            block_dim: (32, 1, 1),
6779            shared_mem_bytes: 0,
6780        };
6781        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6782            in_f as i32,
6783            n_ff as i32,
6784            n_expert as i32,
6785            n_used as i32,
6786            n_pairs as i32,
6787            rb_g as i64,
6788            rb_u as i64,
6789        );
6790        let __s_b = self.gpu.stream();
6791        let mut b = __s_b.launch_builder(&f);
6792        b.arg(table)
6793            .arg(sel)
6794            .arg(aq)
6795            .arg(ad)
6796            .arg(&mut act)
6797            .arg(&inf)
6798            .arg(&nff)
6799            .arg(&ne)
6800            .arg(&qt_g)
6801            .arg(&qt_u)
6802            .arg(&rbg)
6803            .arg(&rbu)
6804            .arg(&nu)
6805            .arg(&npi);
6806        unsafe {
6807            b.launch(cfg)?;
6808        }
6809        Ok(act)
6810    }
6811
6812    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6813    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6814    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6815    #[allow(clippy::too_many_arguments)]
6816    pub fn moe_down8_fma_dev_q8_variant(
6817        &self,
6818        variant: &str,
6819        table: &CudaSlice<u64>,
6820        sel: &cudarc::driver::CudaView<i32>,
6821        w: &cudarc::driver::CudaView<f32>,
6822        aq2: &CudaSlice<i8>,
6823        ad2: &CudaSlice<f32>,
6824        dst: &mut cudarc::driver::CudaViewMut<f32>,
6825        in_f: usize,
6826        out_f: usize,
6827        n_used: usize,
6828        n_expert: usize,
6829        qt: i32,
6830        rb: usize,
6831    ) -> Result<(), Box<dyn std::error::Error>> {
6832        let (inf, outf, nu, ne, rbi) = (
6833            in_f as i32,
6834            out_f as i32,
6835            n_used as i32,
6836            n_expert as i32,
6837            rb as i64,
6838        );
6839        let (f, cfg) = match variant {
6840            "w8h2" | "w8h2v" => (
6841                self.func(if variant == "w8h2" {
6842                    "moe_down8_fma_dev_q8_w8h2"
6843                } else {
6844                    "moe_down8_fma_dev_q8_w8h2v"
6845                }),
6846                LaunchConfig {
6847                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6848                    block_dim: (32, n_used as u32, 1),
6849                    shared_mem_bytes: 0,
6850                },
6851            ),
6852            "w8h2r2" | "w8h2r2v" => (
6853                self.func(if variant == "w8h2r2" {
6854                    "moe_down8_fma_dev_q8_w8h2r2"
6855                } else {
6856                    "moe_down8_fma_dev_q8_w8h2r2v"
6857                }),
6858                LaunchConfig {
6859                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6860                    block_dim: (32, n_used as u32, 1),
6861                    shared_mem_bytes: 0,
6862                },
6863            ),
6864            _ => (
6865                self.func("moe_down8_fma_dev_q8"),
6866                LaunchConfig {
6867                    grid_dim: (out_f as u32, 1, 1),
6868                    block_dim: (32, 1, 1),
6869                    shared_mem_bytes: 0,
6870                },
6871            ),
6872        };
6873        let __s_b = self.gpu.stream();
6874        let mut b = __s_b.launch_builder(&f);
6875        b.arg(table)
6876            .arg(sel)
6877            .arg(w)
6878            .arg(aq2)
6879            .arg(ad2)
6880            .arg(dst)
6881            .arg(&inf)
6882            .arg(&outf)
6883            .arg(&nu)
6884            .arg(&ne)
6885            .arg(&qt)
6886            .arg(&rbi);
6887        unsafe {
6888            b.launch(cfg)?;
6889        }
6890        Ok(())
6891    }
6892
6893    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6894    #[allow(clippy::too_many_arguments)]
6895    pub fn moe_gate_up_silu8_dev_q8_variant(
6896        &self,
6897        variant: &str,
6898        table: &CudaSlice<u64>,
6899        sel: &cudarc::driver::CudaView<i32>,
6900        aq: &CudaSlice<i8>,
6901        ad: &CudaSlice<f32>,
6902        in_f: usize,
6903        n_ff: usize,
6904        n_used: usize,
6905        n_expert: usize,
6906        qt_g: i32,
6907        qt_u: i32,
6908        rb_g: usize,
6909        rb_u: usize,
6910    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6911        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6912        let (inf, nff, ne, rbg, rbu) = (
6913            in_f as i32,
6914            n_ff as i32,
6915            n_expert as i32,
6916            rb_g as i64,
6917            rb_u as i64,
6918        );
6919        let f = self.func(if variant == "v" {
6920            "moe_gate_up_silu8_dev_q8_v"
6921        } else {
6922            "moe_gate_up_silu8_dev_q8"
6923        });
6924        let cfg = LaunchConfig {
6925            grid_dim: (n_ff as u32, n_used as u32, 1),
6926            block_dim: (32, 1, 1),
6927            shared_mem_bytes: 0,
6928        };
6929        let __s_b = self.gpu.stream();
6930        let mut b = __s_b.launch_builder(&f);
6931        b.arg(table)
6932            .arg(sel)
6933            .arg(aq)
6934            .arg(ad)
6935            .arg(&mut act)
6936            .arg(&inf)
6937            .arg(&nff)
6938            .arg(&ne)
6939            .arg(&qt_g)
6940            .arg(&qt_u)
6941            .arg(&rbg)
6942            .arg(&rbu);
6943        unsafe {
6944            b.launch(cfg)?;
6945        }
6946        Ok(act)
6947    }
6948
6949    pub fn moe_gate_up_silu8_dev(
6950        &self,
6951        table: &CudaSlice<u64>,
6952        sel: &cudarc::driver::CudaView<i32>,
6953        x: &cudarc::driver::CudaView<f32>,
6954        in_f: usize,
6955        n_ff: usize,
6956        n_used: usize,
6957        n_expert: usize,
6958        qt_g: i32,
6959        qt_u: i32,
6960        rb_g: usize,
6961        rb_u: usize,
6962        macros: &CudaSlice<f32>,
6963    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6964        let f = self.func("moe_gate_up_silu8_dev");
6965        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6966        let cfg = LaunchConfig {
6967            grid_dim: (n_ff as u32, n_used as u32, 1),
6968            block_dim: (256, 1, 1),
6969            shared_mem_bytes: 0,
6970        };
6971        let (inf, nff, ne, rbg, rbu) = (
6972            in_f as i32,
6973            n_ff as i32,
6974            n_expert as i32,
6975            rb_g as i64,
6976            rb_u as i64,
6977        );
6978        let __s_b = self.gpu.stream();
6979        let mut b = __s_b.launch_builder(&f);
6980        b.arg(table)
6981            .arg(sel)
6982            .arg(x)
6983            .arg(&mut act)
6984            .arg(&inf)
6985            .arg(&nff)
6986            .arg(&ne)
6987            .arg(&qt_g)
6988            .arg(&qt_u)
6989            .arg(&rbg)
6990            .arg(&rbu)
6991            .arg(macros);
6992        unsafe {
6993            b.launch(cfg)?;
6994        }
6995        Ok(act)
6996    }
6997
6998    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6999    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
7000    #[allow(clippy::too_many_arguments)]
7001    pub fn moe_down8_fma_dev(
7002        &self,
7003        table: &CudaSlice<u64>,
7004        sel: &cudarc::driver::CudaView<i32>,
7005        w: &cudarc::driver::CudaView<f32>,
7006        act: &CudaSlice<f32>,
7007        dst: &mut cudarc::driver::CudaViewMut<f32>,
7008        in_f: usize,
7009        out_f: usize,
7010        n_used: usize,
7011        n_expert: usize,
7012        qt: i32,
7013        rb: usize,
7014    ) -> Result<(), Box<dyn std::error::Error>> {
7015        let f = self.func("moe_down8_fma_dev");
7016        let cfg = LaunchConfig {
7017            grid_dim: (out_f as u32, 1, 1),
7018            block_dim: (256, 1, 1),
7019            shared_mem_bytes: 0,
7020        };
7021        let (inf, outf, nu, ne, rbv) = (
7022            in_f as i32,
7023            out_f as i32,
7024            n_used as i32,
7025            n_expert as i32,
7026            rb as i64,
7027        );
7028        let __s_b = self.gpu.stream();
7029        let mut b = __s_b.launch_builder(&f);
7030        b.arg(table)
7031            .arg(sel)
7032            .arg(w)
7033            .arg(act)
7034            .arg(dst)
7035            .arg(&inf)
7036            .arg(&outf)
7037            .arg(&nu)
7038            .arg(&ne)
7039            .arg(&qt)
7040            .arg(&rbv);
7041        unsafe {
7042            b.launch(cfg)?;
7043        }
7044        Ok(())
7045    }
7046
7047    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
7048    pub fn axpy_into(
7049        &self,
7050        src: &CudaSlice<f32>,
7051        alpha: f32,
7052        dst: &mut cudarc::driver::CudaViewMut<f32>,
7053        n: usize,
7054    ) -> Result<(), Box<dyn std::error::Error>> {
7055        let f = self.func("axpy_f32");
7056        let cfg = LaunchConfig::for_num_elems(n as u32);
7057        let (a, ni) = (alpha, n as i32);
7058        let __s_b = self.gpu.stream();
7059        let mut b = __s_b.launch_builder(&f);
7060        b.arg(src).arg(dst).arg(&a).arg(&ni);
7061        unsafe {
7062            b.launch(cfg)?;
7063        }
7064        Ok(())
7065    }
7066
7067    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
7068    pub fn axpy_host_into(
7069        &self,
7070        src: &cudarc::driver::CudaView<'_, f32>,
7071        alpha: f32,
7072        dst: &mut cudarc::driver::CudaViewMut<f32>,
7073        n: usize,
7074    ) -> Result<(), Box<dyn std::error::Error>> {
7075        let f = self.func("axpy_host_f32");
7076        let cfg = LaunchConfig::for_num_elems(n as u32);
7077        let (a, ni) = (alpha, n as i32);
7078        let __s_b = self.gpu.stream();
7079        let mut b = __s_b.launch_builder(&f);
7080        b.arg(src).arg(dst).arg(&a).arg(&ni);
7081        unsafe {
7082            b.launch(cfg)?;
7083        }
7084        Ok(())
7085    }
7086
7087    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
7088    pub fn add_scaled_rows(
7089        &self,
7090        src: &CudaSlice<f32>,
7091        scale: &CudaSlice<f32>,
7092        dst: &mut CudaSlice<f32>,
7093        ncols: usize,
7094        nrows: usize,
7095    ) -> Result<(), Box<dyn std::error::Error>> {
7096        let f = self.func("add_scaled_rows_f32");
7097        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7098        let (nc, nr) = (ncols as i32, nrows as i32);
7099        let __s_b = self.gpu.stream();
7100        let mut b = __s_b.launch_builder(&f);
7101        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7102        unsafe {
7103            b.launch(cfg)?;
7104        }
7105        Ok(())
7106    }
7107
7108    // ======== A2 GROUPED MoE PREFILL KERNELS ========
7109
7110    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
7111    pub fn gather_rows(
7112        &self,
7113        src: &CudaSlice<f32>,
7114        idx: &CudaSlice<i32>,
7115        dst: &mut CudaSlice<f32>,
7116        ncols: usize,
7117        m_e: usize,
7118    ) -> Result<(), Box<dyn std::error::Error>> {
7119        let f = self.func("gather_rows_f32");
7120        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7121        let (nc, me) = (ncols as i32, m_e as i32);
7122        let __s_b = self.gpu.stream();
7123        let mut b = __s_b.launch_builder(&f);
7124        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7125        unsafe {
7126            b.launch(cfg)?;
7127        }
7128        Ok(())
7129    }
7130
7131    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
7132    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
7133    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
7134    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
7135    pub fn scatter_slot(
7136        &self,
7137        src: &CudaSlice<f32>,
7138        tok_idx: &CudaSlice<i32>,
7139        slot_idx: &CudaSlice<i32>,
7140        weight: &CudaSlice<f32>,
7141        dst: &mut CudaSlice<f32>,
7142        wbuf: &mut CudaSlice<f32>,
7143        ncols: usize,
7144        n_used: usize,
7145        m_e: usize,
7146    ) -> Result<(), Box<dyn std::error::Error>> {
7147        let f = self.func("scatter_add_slot_f32");
7148        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7149        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7150        let __s_b = self.gpu.stream();
7151        let mut b = __s_b.launch_builder(&f);
7152        b.arg(src)
7153            .arg(tok_idx)
7154            .arg(slot_idx)
7155            .arg(weight)
7156            .arg(dst)
7157            .arg(wbuf)
7158            .arg(&nc)
7159            .arg(&nu)
7160            .arg(&me);
7161        unsafe {
7162            b.launch(cfg)?;
7163        }
7164        Ok(())
7165    }
7166
7167    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
7168    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
7169    /// Uses FMA for bit-identity with the sequential axpy path.
7170    pub fn reduce_slots(
7171        &self,
7172        slots: &CudaSlice<f32>,
7173        wbuf: &CudaSlice<f32>,
7174        dst: &mut CudaSlice<f32>,
7175        ncols: usize,
7176        n_used: usize,
7177        t: usize,
7178    ) -> Result<(), Box<dyn std::error::Error>> {
7179        let f = self.func("reduce_slots_f32");
7180        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7181        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7182        let __s_b = self.gpu.stream();
7183        let mut b = __s_b.launch_builder(&f);
7184        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7185        unsafe {
7186            b.launch(cfg)?;
7187        }
7188        Ok(())
7189    }
7190
7191    /// Canonical slot-order reduction with separately rounded multiply and add.
7192    ///
7193    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
7194    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
7195    pub fn reduce_slots_host(
7196        &self,
7197        slots: &CudaSlice<f32>,
7198        wbuf: &CudaSlice<f32>,
7199        dst: &mut CudaSlice<f32>,
7200        ncols: usize,
7201        n_used: usize,
7202        t: usize,
7203    ) -> Result<(), Box<dyn std::error::Error>> {
7204        let f = self.func("reduce_slots_host_f32");
7205        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7206        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7207        let __s_b = self.gpu.stream();
7208        let mut b = __s_b.launch_builder(&f);
7209        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7210        unsafe {
7211            b.launch(cfg)?;
7212        }
7213        Ok(())
7214    }
7215
7216    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
7217    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
7218    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
7219    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
7220    /// GPU time, ~half of it redundant re-quantization of the same row.
7221    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
7222    pub fn quantize_q8_1_view(
7223        &self,
7224        x: &cudarc::driver::CudaView<f32>,
7225        m: usize,
7226        in_f: usize,
7227    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7228        let f = self.func("quantize_q8_1");
7229        let nblk = in_f / 32;
7230        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7231        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7232        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7233        let (inf, mi) = (in_f as i32, m as i32);
7234        let __s_b = self.gpu.stream();
7235        let mut b = __s_b.launch_builder(&f);
7236        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7237        unsafe {
7238            b.launch(cfg)?;
7239        }
7240        Ok((q, d))
7241    }
7242
7243    pub fn quantize_q8_1(
7244        &self,
7245        x: &CudaSlice<f32>,
7246        m: usize,
7247        in_f: usize,
7248    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7249        let nblk = in_f / 32;
7250        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
7251        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
7252        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
7253        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7254        let (inf, mi) = (in_f as i32, m as i32);
7255        if Self::pdl_on() && Self::pdl_wb_on() {
7256            {
7257                use cudarc::driver::{DevicePtr, DevicePtrMut};
7258                let s = &self.gpu.stream();
7259                let (px, _g0) = x.device_ptr(s);
7260                let (pq, _g1) = q.device_ptr_mut(s);
7261                let (pd, _g2) = d.device_ptr_mut(s);
7262                let mut ps = [
7263                    &px as *const _ as *mut std::ffi::c_void,
7264                    &pq as *const _ as *mut _,
7265                    &pd as *const _ as *mut _,
7266                    &inf as *const _ as *mut _,
7267                    &mi as *const _ as *mut _,
7268                ];
7269                unsafe {
7270                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7271                }
7272            }
7273            return Ok((q, d));
7274        }
7275        let f = self.func("quantize_q8_1");
7276        let __s_b = self.gpu.stream();
7277        let mut b = __s_b.launch_builder(&f);
7278        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7279        unsafe {
7280            b.launch(cfg)?;
7281        }
7282        Ok((q, d))
7283    }
7284
7285    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
7286    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
7287    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
7288    pub fn quantize_fp4_act(
7289        &self,
7290        x: &CudaSlice<f32>,
7291        m: usize,
7292        in_f: usize,
7293    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7294        let f = self.func("quantize_fp4_act");
7295        let nb16 = in_f / 16;
7296        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
7297        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
7298        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
7299        let (inf, mi) = (in_f as i32, m as i32);
7300        let __s_b = self.gpu.stream();
7301        let mut b = __s_b.launch_builder(&f);
7302        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7303        unsafe {
7304            b.launch(cfg)?;
7305        }
7306        Ok((aq4, ad4))
7307    }
7308
7309    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
7310    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
7311    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
7312    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
7313    pub fn qmatvec_gemm_nvfp4_fp4(
7314        &self,
7315        bytes: &CudaSlice<u8>,
7316        x: &CudaSlice<f32>,
7317        m: usize,
7318        in_f: usize,
7319        out_f: usize,
7320        row_bytes: usize,
7321        scale: f32,
7322    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7323        assert!(
7324            in_f % 64 == 0,
7325            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7326        );
7327        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7328        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7329        if scale != 1.0 {
7330            self.scale_inplace(&mut y, scale, m * out_f)?;
7331        }
7332        Ok(y)
7333    }
7334
7335    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
7336    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
7337    fn fp4_gemm_launch(
7338        &self,
7339        bytes: &CudaSlice<u8>,
7340        aq4: &CudaSlice<u32>,
7341        ad4: &CudaSlice<u8>,
7342        m: usize,
7343        in_f: usize,
7344        out_f: usize,
7345        row_bytes: usize,
7346    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7347        let f = self.func("qmatvec_gemm_nvfp4_fp4");
7348        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7349        const BM: u32 = 64;
7350        const BN: u32 = 256;
7351        let cfg = LaunchConfig {
7352            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7353            block_dim: (32, 4, 1),
7354            shared_mem_bytes: 0,
7355        };
7356        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7357        let __s_b = self.gpu.stream();
7358        let mut b = __s_b.launch_builder(&f);
7359        b.arg(bytes)
7360            .arg(aq4)
7361            .arg(ad4)
7362            .arg(&mut y)
7363            .arg(&inf)
7364            .arg(&outf)
7365            .arg(&mi)
7366            .arg(&rb);
7367        unsafe {
7368            b.launch(cfg)?;
7369        }
7370        Ok(y)
7371    }
7372
7373    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
7374    pub fn qmatvec_gemm_nvfp4_fp4_raw(
7375        &self,
7376        bytes: &CudaSlice<u8>,
7377        x: &CudaSlice<f32>,
7378        m: usize,
7379        in_f: usize,
7380        out_f: usize,
7381        row_bytes: usize,
7382    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7383        assert!(
7384            in_f % 64 == 0,
7385            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7386        );
7387        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7388        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7389    }
7390
7391    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
7392    pub fn qmatvec_q8_0_fast(
7393        &self,
7394        w: &CudaSlice<u8>,
7395        x: &CudaSlice<f32>,
7396        m: usize,
7397        in_f: usize,
7398        out_f: usize,
7399        row_bytes: usize,
7400    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7401        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7402        let f = self.func("qmatvec_q8_0_dp4a");
7403        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7404        let cfg = LaunchConfig {
7405            grid_dim: (out_f as u32, m as u32, 1),
7406            block_dim: (128, 1, 1),
7407            shared_mem_bytes: 0,
7408        };
7409        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7410        let __s_b = self.gpu.stream();
7411        let mut b = __s_b.launch_builder(&f);
7412        b.arg(w)
7413            .arg(&aq)
7414            .arg(&ad)
7415            .arg(&mut y)
7416            .arg(&inf)
7417            .arg(&outf)
7418            .arg(&mi)
7419            .arg(&rb);
7420        unsafe {
7421            b.launch(cfg)?;
7422        }
7423        Ok(y)
7424    }
7425
7426    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7427    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7428    pub fn qmatvec_q4_K_fast(
7429        &self,
7430        w: &CudaSlice<u8>,
7431        x: &CudaSlice<f32>,
7432        m: usize,
7433        in_f: usize,
7434        out_f: usize,
7435        row_bytes: usize,
7436    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7437        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7438        let f = self.func("qmatvec_q4_K_dp4a");
7439        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7440        let cfg = LaunchConfig {
7441            grid_dim: (out_f as u32, m as u32, 1),
7442            block_dim: (128, 1, 1),
7443            shared_mem_bytes: 0,
7444        };
7445        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7446        let __s_b = self.gpu.stream();
7447        let mut b = __s_b.launch_builder(&f);
7448        b.arg(w)
7449            .arg(&aq)
7450            .arg(&ad)
7451            .arg(&mut y)
7452            .arg(&inf)
7453            .arg(&outf)
7454            .arg(&mi)
7455            .arg(&rb);
7456        unsafe {
7457            b.launch(cfg)?;
7458        }
7459        Ok(y)
7460    }
7461
7462    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7463    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7464    pub fn qmatvec_q6_K_fast(
7465        &self,
7466        w: &CudaSlice<u8>,
7467        x: &CudaSlice<f32>,
7468        m: usize,
7469        in_f: usize,
7470        out_f: usize,
7471        row_bytes: usize,
7472    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7473        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7474        let f = self.func("qmatvec_q6_K_dp4a");
7475        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7476        let cfg = LaunchConfig {
7477            grid_dim: (out_f as u32, m as u32, 1),
7478            block_dim: (128, 1, 1),
7479            shared_mem_bytes: 0,
7480        };
7481        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7482        let __s_b = self.gpu.stream();
7483        let mut b = __s_b.launch_builder(&f);
7484        b.arg(w)
7485            .arg(&aq)
7486            .arg(&ad)
7487            .arg(&mut y)
7488            .arg(&inf)
7489            .arg(&outf)
7490            .arg(&mi)
7491            .arg(&rb);
7492        unsafe {
7493            b.launch(cfg)?;
7494        }
7495        Ok(y)
7496    }
7497
7498    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7499    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7500    pub fn qmatvec_q5_K_fast(
7501        &self,
7502        w: &CudaSlice<u8>,
7503        x: &CudaSlice<f32>,
7504        m: usize,
7505        in_f: usize,
7506        out_f: usize,
7507        row_bytes: usize,
7508    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7509        self.qmatvec_dp4a_named(
7510            "qmatvec_q5_K_dp4a",
7511            &w.slice(0..w.len()),
7512            x,
7513            m,
7514            in_f,
7515            out_f,
7516            row_bytes,
7517        )
7518    }
7519    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7520    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7521    pub fn qmatvec_q3_K_fast(
7522        &self,
7523        w: &CudaSlice<u8>,
7524        x: &CudaSlice<f32>,
7525        m: usize,
7526        in_f: usize,
7527        out_f: usize,
7528        row_bytes: usize,
7529    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7530        self.qmatvec_dp4a_named(
7531            "qmatvec_q3_K_dp4a",
7532            &w.slice(0..w.len()),
7533            x,
7534            m,
7535            in_f,
7536            out_f,
7537            row_bytes,
7538        )
7539    }
7540    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
7541    pub fn qmatvec_nvfp4_fast_rp(
7542        &self,
7543        w: &CudaSlice<u8>,
7544        x: &CudaSlice<f32>,
7545        m: usize,
7546        in_f: usize,
7547        out_f: usize,
7548        row_bytes: usize,
7549    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7550        assert!(
7551            in_f % 64 == 0,
7552            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7553        );
7554        self.qmatvec_dp4a_named(
7555            "qmatvec_nvfp4_dp4a_rp",
7556            &w.slice(0..w.len()),
7557            x,
7558            m,
7559            in_f,
7560            out_f,
7561            row_bytes,
7562        )
7563    }
7564    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
7565    pub fn qmatvec_nvfp4_fast(
7566        &self,
7567        w: &cudarc::driver::CudaView<'_, u8>,
7568        x: &CudaSlice<f32>,
7569        m: usize,
7570        in_f: usize,
7571        out_f: usize,
7572        row_bytes: usize,
7573    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7574        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
7575        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
7576        assert!(
7577            in_f % 64 == 0,
7578            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7579        );
7580        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7581    }
7582    /// v2-layout twin of `qmatvec_nvfp4_fast` for the slot-major expert banks
7583    /// (MEMRA_NVFP4_BANK_V2) — bit-identical per row, coalesced reads.
7584    pub fn qmatvec_nvfp4_fast_v2(
7585        &self,
7586        w: &cudarc::driver::CudaView<'_, u8>,
7587        x: &CudaSlice<f32>,
7588        m: usize,
7589        in_f: usize,
7590        out_f: usize,
7591        row_bytes: usize,
7592    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7593        assert!(
7594            in_f % 64 == 0,
7595            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7596        );
7597        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7598    }
7599    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
7600    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7601    pub fn qmatvec_iq4_XS_fast(
7602        &self,
7603        w: &CudaSlice<u8>,
7604        x: &CudaSlice<f32>,
7605        m: usize,
7606        in_f: usize,
7607        out_f: usize,
7608        row_bytes: usize,
7609    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7610        self.qmatvec_dp4a_named(
7611            "qmatvec_iq4_XS_dp4a",
7612            &w.slice(0..w.len()),
7613            x,
7614            m,
7615            in_f,
7616            out_f,
7617            row_bytes,
7618        )
7619    }
7620
7621    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
7622    fn qmatvec_dp4a_named(
7623        &self,
7624        name: &str,
7625        w: &cudarc::driver::CudaView<'_, u8>,
7626        x: &CudaSlice<f32>,
7627        m: usize,
7628        in_f: usize,
7629        out_f: usize,
7630        row_bytes: usize,
7631    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7632        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7633        let f = self.func(name);
7634        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7635        let cfg = LaunchConfig {
7636            grid_dim: (out_f as u32, m as u32, 1),
7637            block_dim: (128, 1, 1),
7638            shared_mem_bytes: 0,
7639        };
7640        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7641        let __s_b = self.gpu.stream();
7642        let mut b = __s_b.launch_builder(&f);
7643        b.arg(w)
7644            .arg(&aq)
7645            .arg(&ad)
7646            .arg(&mut y)
7647            .arg(&inf)
7648            .arg(&outf)
7649            .arg(&mi)
7650            .arg(&rb);
7651        unsafe {
7652            b.launch(cfg)?;
7653        }
7654        Ok(y)
7655    }
7656
7657    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
7658    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
7659    /// its output); this entry exists so a routed-expert program can quantize one activation
7660    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
7661    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
7662    #[allow(clippy::too_many_arguments)]
7663    pub fn qmatvec_nvfp4_fast_prequant_into(
7664        &self,
7665        w: &CudaSlice<u8>,
7666        aq: &CudaSlice<i8>,
7667        ad: &CudaSlice<f32>,
7668        y: &mut CudaSlice<f32>,
7669        m: usize,
7670        in_f: usize,
7671        out_f: usize,
7672        row_bytes: usize,
7673    ) -> Result<(), Box<dyn std::error::Error>> {
7674        assert!(
7675            in_f % 64 == 0,
7676            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7677        );
7678        if y.len() < m * out_f {
7679            return Err(format!(
7680                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7681                y.len()
7682            )
7683            .into());
7684        }
7685        let f = self.func("qmatvec_nvfp4_dp4a");
7686        let cfg = LaunchConfig {
7687            grid_dim: (out_f as u32, m as u32, 1),
7688            block_dim: (128, 1, 1),
7689            shared_mem_bytes: 0,
7690        };
7691        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7692        let __s_b = self.gpu.stream();
7693        let mut b = __s_b.launch_builder(&f);
7694        b.arg(w)
7695            .arg(aq)
7696            .arg(ad)
7697            .arg(y)
7698            .arg(&inf)
7699            .arg(&outf)
7700            .arg(&mi)
7701            .arg(&rb);
7702        unsafe {
7703            b.launch(cfg)?;
7704        }
7705        Ok(())
7706    }
7707
7708    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
7709    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
7710    #[allow(clippy::too_many_arguments)]
7711    pub fn matvec_f32_qkv_into(
7712        &self,
7713        wq: &CudaSlice<f32>,
7714        wk: &CudaSlice<f32>,
7715        wv: &CudaSlice<f32>,
7716        wg: &CudaSlice<f32>,
7717        x: &CudaSlice<f32>,
7718        yq: &mut CudaSlice<f32>,
7719        yk: &mut CudaSlice<f32>,
7720        yv: &mut CudaSlice<f32>,
7721        yg: &mut CudaSlice<f32>,
7722        in_f: usize,
7723        out_q: usize,
7724        out_kv: usize,
7725        out_g: usize,
7726    ) -> Result<(), Box<dyn std::error::Error>> {
7727        if in_f % 4 != 0
7728            || wq.len() != out_q * in_f
7729            || wk.len() != out_kv * in_f
7730            || wv.len() != out_kv * in_f
7731            || wg.len() < out_g * in_f
7732            || x.len() < in_f
7733            || yq.len() < out_q
7734            || yk.len() < out_kv
7735            || yv.len() < out_kv
7736            || (out_g > 0 && yg.len() < out_g)
7737        {
7738            return Err(format!(
7739                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7740                 wq={} wk={} wv={} wg={}",
7741                wq.len(),
7742                wk.len(),
7743                wv.len(),
7744                wg.len()
7745            )
7746            .into());
7747        }
7748        let f = self.func("matvec_f32_qkv");
7749        let cfg = LaunchConfig {
7750            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7751            block_dim: (128, 1, 1),
7752            shared_mem_bytes: 0,
7753        };
7754        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7755        let __s_b = self.gpu.stream();
7756        let mut b = __s_b.launch_builder(&f);
7757        b.arg(wq)
7758            .arg(wk)
7759            .arg(wv)
7760            .arg(wg)
7761            .arg(x)
7762            .arg(yq)
7763            .arg(yk)
7764            .arg(yv)
7765            .arg(yg)
7766            .arg(&inf)
7767            .arg(&oq)
7768            .arg(&okv)
7769            .arg(&og);
7770        unsafe {
7771            b.launch(cfg)?;
7772        }
7773        Ok(())
7774    }
7775
7776    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
7777    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
7778    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
7779    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
7780    /// kernel — the batching only removes host launch latency.
7781    #[allow(clippy::too_many_arguments)]
7782    /// FUSION #2a: gate+up sweeps in one launch (v2 banks only; identical geometry both
7783    /// banks, caller-guarded). Per-row bit-identical to two qmatvec_nvfp4_sel_into calls.
7784    #[allow(clippy::too_many_arguments)]
7785    pub fn qmatvec_nvfp4_sel_gu_into(
7786        &self,
7787        gate_bank: &CudaSlice<u8>,
7788        up_bank: &CudaSlice<u8>,
7789        sel: &CudaSlice<i32>,
7790        aq: &CudaSlice<i8>,
7791        ad: &CudaSlice<f32>,
7792        yg: &mut CudaSlice<f32>,
7793        yu: &mut CudaSlice<f32>,
7794        n_sel: usize,
7795        in_f: usize,
7796        out_f: usize,
7797        row_bytes: usize,
7798        expert_stride: usize,
7799    ) -> Result<(), Box<dyn std::error::Error>> {
7800        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7801        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7802            return Err("NVFP4 gu sel geometry".into());
7803        }
7804        // MEMRA_SEL_GU_RPW=2|4: multirow twin (activation group read once, reused across
7805        // RPW rows' gate+up dots) — bit-identical per row, one block per RPW rows.
7806        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7807        let rpw = *RPW.get_or_init(|| {
7808            std::env::var("MEMRA_SEL_GU_RPW")
7809                .ok()
7810                .and_then(|v| v.parse().ok())
7811                .filter(|r| *r == 2 || *r == 4)
7812                .unwrap_or(1)
7813        });
7814        let rpw = if out_f % rpw == 0 { rpw } else { 1 };
7815        // MEMRA_SEL_GU_WPR=1: warp-per-row (NUMERIC-CLASS — per-row reduction order changes;
7816        // acceptance is the argmax gate + battery, the QKV_FUSED/BF16_MMV class).
7817        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7818        let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
7819        let f = self.func(match (wpr, rpw) {
7820            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
7821            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
7822            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
7823            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
7824        });
7825        let cfg = LaunchConfig {
7826            grid_dim: if wpr {
7827                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
7828            } else if rpw == 1 {
7829                ((2 * out_f) as u32, n_sel as u32, 1)
7830            } else {
7831                ((out_f / rpw) as u32, n_sel as u32, 1)
7832            },
7833            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
7834            shared_mem_bytes: 0,
7835        };
7836        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7837        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7838        let (ars, adrs) = (0i64, 0i64);
7839        let __s_b = self.gpu.stream();
7840        let mut b = __s_b.launch_builder(&f);
7841        b.arg(gate_bank)
7842            .arg(up_bank)
7843            .arg(sel)
7844            .arg(aq)
7845            .arg(ad)
7846            .arg(yg)
7847            .arg(yu)
7848            .arg(&inf)
7849            .arg(&outf)
7850            .arg(&ns)
7851            .arg(&rb)
7852            .arg(&es)
7853            .arg(&ars)
7854            .arg(&adrs);
7855        unsafe {
7856            b.launch(cfg)?;
7857        }
7858        Ok(())
7859    }
7860
7861    /// MEMRA_SEL_DOWN8=1: the DOWN sweep and the route-weight combine in ONE launch
7862    /// (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm ported to the
7863    /// NVFP4 banks). Block = (32, n_sel): one warp per slot instead of one warp per
7864    /// (row, slot), and the n_sel x out_f partial buffer disappears. Bit-identical to
7865    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce
7866    /// tree, same slot-ordered chain. Requires the v2 banks and nsb <= 32 (the fit-block
7867    /// class the reduce identity is argued at).
7868    #[allow(clippy::too_many_arguments)]
7869    pub fn qmatvec_nvfp4_sel_down8_into(
7870        &self,
7871        bank: &CudaSlice<u8>,
7872        sel: &CudaSlice<i32>,
7873        aq: &CudaSlice<i8>,
7874        ad: &CudaSlice<f32>,
7875        route_w: &CudaSlice<f32>,
7876        md: &CudaSlice<f32>,
7877        dst: &mut CudaSlice<f32>,
7878        n_sel: usize,
7879        in_f: usize,
7880        out_f: usize,
7881        row_bytes: usize,
7882        expert_stride: usize,
7883        act_row_stride: usize,
7884        ad_row_stride: usize,
7885    ) -> Result<(), Box<dyn std::error::Error>> {
7886        if in_f % 64 != 0
7887            || n_sel == 0
7888            || n_sel > 8
7889            || (in_f >> 5) > 32
7890            || dst.len() < out_f
7891            || sel.len() < n_sel
7892            || route_w.len() < n_sel
7893        {
7894            return Err(format!(
7895                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
7896                dst.len()
7897            )
7898            .into());
7899        }
7900        if !crate::tp::nvfp4_bank_v2_on() {
7901            return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
7902        }
7903        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
7904        let cfg = LaunchConfig {
7905            grid_dim: (out_f as u32, 1, 1),
7906            block_dim: (32, n_sel as u32, 1),
7907            shared_mem_bytes: 0,
7908        };
7909        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7910        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7911        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7912        let __s_b = self.gpu.stream();
7913        let mut b = __s_b.launch_builder(&f);
7914        b.arg(bank)
7915            .arg(sel)
7916            .arg(aq)
7917            .arg(ad)
7918            .arg(route_w)
7919            .arg(md)
7920            .arg(dst)
7921            .arg(&inf)
7922            .arg(&outf)
7923            .arg(&ns)
7924            .arg(&rb)
7925            .arg(&es)
7926            .arg(&ars)
7927            .arg(&adrs);
7928        unsafe {
7929            b.launch(cfg)?;
7930        }
7931        Ok(())
7932    }
7933
7934    /// T-ROW twin of the down8 fusion (spec verify + batched serving MoE): one block per
7935    /// (output row, token row) = the exact t=1 down8 program per token — bit-identical
7936    /// per row to its own down8/axpy pair at any t.
7937    #[allow(clippy::too_many_arguments)]
7938    pub fn qmatvec_nvfp4_sel_down8_rows_into(
7939        &self,
7940        bank: &CudaSlice<u8>,
7941        sel: &CudaSlice<i32>,
7942        aq: &CudaSlice<i8>,
7943        ad: &CudaSlice<f32>,
7944        route_w: &CudaSlice<f32>,
7945        md: &CudaSlice<f32>,
7946        dst: &mut CudaSlice<f32>,
7947        t: usize,
7948        n_sel_col: usize,
7949        in_f: usize,
7950        out_f: usize,
7951        row_bytes: usize,
7952        expert_stride: usize,
7953        act_row_stride: usize,
7954        ad_row_stride: usize,
7955    ) -> Result<(), Box<dyn std::error::Error>> {
7956        let n_sel = t * n_sel_col;
7957        if in_f % 64 != 0
7958            || n_sel_col == 0
7959            || n_sel_col > 8
7960            || t == 0
7961            || t > 64
7962            || (in_f >> 5) > 32
7963            || dst.len() < t * out_f
7964            || sel.len() < n_sel
7965            || route_w.len() < n_sel
7966        {
7967            return Err(format!(
7968                "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
7969                dst.len()
7970            )
7971            .into());
7972        }
7973        if !crate::tp::nvfp4_bank_v2_on() {
7974            return Err(
7975                "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
7976            );
7977        }
7978        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
7979        let cfg = LaunchConfig {
7980            grid_dim: (out_f as u32, t as u32, 1),
7981            block_dim: (32, n_sel_col as u32, 1),
7982            shared_mem_bytes: 0,
7983        };
7984        let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
7985        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7986        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7987        let __s_b = self.gpu.stream();
7988        let mut b = __s_b.launch_builder(&f);
7989        b.arg(bank)
7990            .arg(sel)
7991            .arg(aq)
7992            .arg(ad)
7993            .arg(route_w)
7994            .arg(md)
7995            .arg(dst)
7996            .arg(&inf)
7997            .arg(&outf)
7998            .arg(&nsc)
7999            .arg(&rb)
8000            .arg(&es)
8001            .arg(&ars)
8002            .arg(&adrs);
8003        unsafe {
8004            b.launch(cfg)?;
8005        }
8006        Ok(())
8007    }
8008
8009    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
8010    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
8011    #[allow(clippy::too_many_arguments)]
8012    pub fn qmatvec_nvfp4_sel_gu_ep_into(
8013        &self,
8014        gate_bank: &CudaSlice<u8>,
8015        up_bank: &CudaSlice<u8>,
8016        sel: &CudaSlice<i32>,
8017        aq: &CudaSlice<i8>,
8018        ad: &CudaSlice<f32>,
8019        yg: &mut CudaSlice<f32>,
8020        yu: &mut CudaSlice<f32>,
8021        n_sel: usize,
8022        in_f: usize,
8023        out_f: usize,
8024        row_bytes: usize,
8025        expert_stride: usize,
8026        owner: usize,
8027    ) -> Result<(), Box<dyn std::error::Error>> {
8028        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8029        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8030            return Err("NVFP4 gu ep geometry".into());
8031        }
8032        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8033        let cfg = LaunchConfig {
8034            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8035            block_dim: (128, 1, 1),
8036            shared_mem_bytes: 0,
8037        };
8038        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8039        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8040        let (ars, adrs) = (0i64, 0i64);
8041        let __s_b = self.gpu.stream();
8042        let mut b = __s_b.launch_builder(&f);
8043        b.arg(gate_bank)
8044            .arg(up_bank)
8045            .arg(sel)
8046            .arg(aq)
8047            .arg(ad)
8048            .arg(yg)
8049            .arg(yu)
8050            .arg(&inf)
8051            .arg(&outf)
8052            .arg(&ns)
8053            .arg(&rb)
8054            .arg(&es)
8055            .arg(&ars)
8056            .arg(&adrs)
8057            .arg(&own);
8058        unsafe {
8059            b.launch(cfg)?;
8060        }
8061        Ok(())
8062    }
8063
8064    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
8065    #[allow(clippy::too_many_arguments)]
8066    pub fn silu_mul_scaled_q8_1_sel_ep_into(
8067        &self,
8068        gate: &CudaSlice<f32>,
8069        up: &CudaSlice<f32>,
8070        gmac: &CudaSlice<f32>,
8071        umac: &CudaSlice<f32>,
8072        sel: &CudaSlice<i32>,
8073        limit: Option<f32>,
8074        out_q: &mut CudaSlice<i8>,
8075        out_d: &mut CudaSlice<f32>,
8076        n_per: usize,
8077        n_sel: usize,
8078        owner: usize,
8079    ) -> Result<(), Box<dyn std::error::Error>> {
8080        if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8081            return Err("NVFP4 silu ep geometry".into());
8082        }
8083        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8084        let warps = n_sel * n_per / 32;
8085        let cfg = LaunchConfig {
8086            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8087            block_dim: (128, 1, 1),
8088            shared_mem_bytes: 0,
8089        };
8090        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8091        let (lim, has) = match limit {
8092            Some(l) => (l, 1i32),
8093            None => (0.0f32, 0i32),
8094        };
8095        let __s_b = self.gpu.stream();
8096        let mut b = __s_b.launch_builder(&f);
8097        b.arg(gate)
8098            .arg(up)
8099            .arg(gmac)
8100            .arg(umac)
8101            .arg(sel)
8102            .arg(&lim)
8103            .arg(&has)
8104            .arg(out_q)
8105            .arg(out_d)
8106            .arg(&np)
8107            .arg(&ns)
8108            .arg(&own);
8109        unsafe {
8110            b.launch(cfg)?;
8111        }
8112        Ok(())
8113    }
8114
8115    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
8116    #[allow(clippy::too_many_arguments)]
8117    pub fn qmatvec_nvfp4_sel_down8_ep_into(
8118        &self,
8119        bank: &CudaSlice<u8>,
8120        sel: &CudaSlice<i32>,
8121        aq: &CudaSlice<i8>,
8122        ad: &CudaSlice<f32>,
8123        route_w: &CudaSlice<f32>,
8124        md: &CudaSlice<f32>,
8125        dst: &mut CudaSlice<f32>,
8126        n_sel: usize,
8127        in_f: usize,
8128        out_f: usize,
8129        row_bytes: usize,
8130        expert_stride: usize,
8131        act_row_stride: usize,
8132        ad_row_stride: usize,
8133        owner: usize,
8134    ) -> Result<(), Box<dyn std::error::Error>> {
8135        if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8136            return Err("NVFP4 down8 ep geometry".into());
8137        }
8138        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8139        let cfg = LaunchConfig {
8140            grid_dim: (out_f as u32, 1, 1),
8141            block_dim: (32, n_sel as u32, 1),
8142            shared_mem_bytes: 0,
8143        };
8144        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8145        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8146        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8147        let __s_b = self.gpu.stream();
8148        let mut b = __s_b.launch_builder(&f);
8149        b.arg(bank)
8150            .arg(sel)
8151            .arg(aq)
8152            .arg(ad)
8153            .arg(route_w)
8154            .arg(md)
8155            .arg(dst)
8156            .arg(&inf)
8157            .arg(&outf)
8158            .arg(&ns)
8159            .arg(&rb)
8160            .arg(&es)
8161            .arg(&ars)
8162            .arg(&adrs)
8163            .arg(&own);
8164        unsafe {
8165            b.launch(cfg)?;
8166        }
8167        Ok(())
8168    }
8169
8170    pub fn qmatvec_nvfp4_sel_into(
8171        &self,
8172        bank: &CudaSlice<u8>,
8173        sel: &CudaSlice<i32>,
8174        aq: &CudaSlice<i8>,
8175        ad: &CudaSlice<f32>,
8176        y: &mut CudaSlice<f32>,
8177        n_sel: usize,
8178        in_f: usize,
8179        out_f: usize,
8180        row_bytes: usize,
8181        expert_stride: usize,
8182        act_row_stride: usize,
8183        ad_row_stride: usize,
8184    ) -> Result<(), Box<dyn std::error::Error>> {
8185        assert!(
8186            in_f % 64 == 0,
8187            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8188        );
8189        if y.len() < n_sel * out_f || sel.len() < n_sel {
8190            return Err(format!(
8191                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8192                y.len(),
8193                sel.len()
8194            )
8195            .into());
8196        }
8197        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
8198        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
8199        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
8200        // sequential-rows variant was flat). Default stays the single-row form.
8201        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
8202        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
8203        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8204        let mode = *MR.get_or_init(|| {
8205            if crate::tp::nvfp4_bank_v2_on() {
8206                3
8207            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8208                2
8209            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8210                1
8211            } else {
8212                0
8213            }
8214        });
8215        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8216        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
8217        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
8218        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
8219        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8220        let v2s = mode == 3
8221            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8222            && row_bytes % 16 == 0
8223            && in_f <= 4096;
8224        let f = match (mode, v2s) {
8225            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8226            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8227            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8228            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8229            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8230        };
8231        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
8232        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
8233        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
8234        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
8235        let nsb = in_f >> 5;
8236        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8237            32
8238        } else if mode == 1 {
8239            512
8240        } else {
8241            128
8242        };
8243        let cfg = LaunchConfig {
8244            grid_dim: (
8245                if v2s {
8246                    (out_f as u32).div_ceil(8)
8247                } else {
8248                    match mode {
8249                        2 => (out_f as u32).div_ceil(16),
8250                        1 => (out_f as u32).div_ceil(4),
8251                        _ => out_f as u32,
8252                    }
8253                },
8254                n_sel as u32,
8255                1,
8256            ),
8257            block_dim: (fit_block, 1, 1),
8258            shared_mem_bytes: 0,
8259        };
8260        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8261        let (rb, es, ars, adrs) = (
8262            row_bytes as i64,
8263            expert_stride as i64,
8264            act_row_stride as i64,
8265            ad_row_stride as i64,
8266        );
8267        let __s_b = self.gpu.stream();
8268        let mut b = __s_b.launch_builder(&f);
8269        b.arg(bank)
8270            .arg(sel)
8271            .arg(aq)
8272            .arg(ad)
8273            .arg(y)
8274            .arg(&inf)
8275            .arg(&outf)
8276            .arg(&ns)
8277            .arg(&rb)
8278            .arg(&es)
8279            .arg(&ars)
8280            .arg(&adrs);
8281        unsafe {
8282            b.launch(cfg)?;
8283        }
8284        Ok(())
8285    }
8286
8287    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
8288    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
8289    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
8290    /// takes the plain SiLU kernel.
8291    #[allow(clippy::too_many_arguments)]
8292    pub fn silu_mul_scaled_q8_1_sel_into(
8293        &self,
8294        gate: &CudaSlice<f32>,
8295        up: &CudaSlice<f32>,
8296        gmac: &CudaSlice<f32>,
8297        umac: &CudaSlice<f32>,
8298        sel: &CudaSlice<i32>,
8299        limit: Option<f32>,
8300        out_q: &mut CudaSlice<i8>,
8301        out_d: &mut CudaSlice<f32>,
8302        n_per: usize,
8303        n_sel: usize,
8304    ) -> Result<(), Box<dyn std::error::Error>> {
8305        let n = n_per * n_sel;
8306        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8307            return Err(format!(
8308                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8309                out_q.len(),
8310                out_d.len()
8311            )
8312            .into());
8313        }
8314        if let Some(limit) = limit {
8315            if limit <= 1e-6 {
8316                return Err(format!(
8317                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8318                )
8319                .into());
8320            }
8321            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8322            let cfg = LaunchConfig::for_num_elems(n as u32);
8323            let (np, ns) = (n_per as i32, n_sel as i32);
8324            let __s_b = self.gpu.stream();
8325            let mut b = __s_b.launch_builder(&f);
8326            b.arg(gate)
8327                .arg(up)
8328                .arg(gmac)
8329                .arg(umac)
8330                .arg(sel)
8331                .arg(&limit)
8332                .arg(out_q)
8333                .arg(out_d)
8334                .arg(&np)
8335                .arg(&ns);
8336            unsafe {
8337                b.launch(cfg)?;
8338            }
8339            return Ok(());
8340        }
8341        let f = self.func("silu_mul_scaled_q8_1_sel");
8342        let cfg = LaunchConfig::for_num_elems(n as u32);
8343        let (np, ns) = (n_per as i32, n_sel as i32);
8344        let __s_b = self.gpu.stream();
8345        let mut b = __s_b.launch_builder(&f);
8346        b.arg(gate)
8347            .arg(up)
8348            .arg(gmac)
8349            .arg(umac)
8350            .arg(sel)
8351            .arg(out_q)
8352            .arg(out_d)
8353            .arg(&np)
8354            .arg(&ns);
8355        unsafe {
8356            b.launch(cfg)?;
8357        }
8358        Ok(())
8359    }
8360
8361    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8362        Ok(self.gpu.stream().clone_htod(v)?)
8363    }
8364    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8365        Ok(self.gpu.stream().clone_htod(v)?)
8366    }
8367    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
8368    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8369        Ok(self.gpu.stream().clone_htod(v)?)
8370    }
8371    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8372        Ok(self.gpu.stream().clone_htod(v)?)
8373    }
8374    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
8375    pub fn dtoh_view(
8376        &self,
8377        d: &cudarc::driver::CudaView<f32>,
8378    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8379        let v = self.gpu.stream().clone_dtoh(d)?;
8380        self.gpu.stream().synchronize()?;
8381        Ok(v)
8382    }
8383    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8384        let v = self.gpu.stream().clone_dtoh(d)?;
8385        self.gpu.stream().synchronize()?;
8386        Ok(v)
8387    }
8388    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
8389    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
8390    /// issuing them together avoids a second stream synchronization in every trunk layer.
8391    pub fn dtoh_pair(
8392        &self,
8393        a: &CudaSlice<f32>,
8394        b: &CudaSlice<f32>,
8395    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8396        let av = self.gpu.stream().clone_dtoh(a)?;
8397        let bv = self.gpu.stream().clone_dtoh(b)?;
8398        self.gpu.stream().synchronize()?;
8399        Ok((av, bv))
8400    }
8401    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
8402    /// cross a shape-sensitive host boundary.
8403    pub fn dtoh_pair_views(
8404        &self,
8405        a: &cudarc::driver::CudaView<f32>,
8406        b: &cudarc::driver::CudaView<f32>,
8407    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8408        let av = self.gpu.stream().clone_dtoh(a)?;
8409        let bv = self.gpu.stream().clone_dtoh(b)?;
8410        self.gpu.stream().synchronize()?;
8411        Ok((av, bv))
8412    }
8413    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
8414    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8415        let v = self.gpu.stream().clone_dtoh(d)?;
8416        self.gpu.stream().synchronize()?;
8417        Ok(v)
8418    }
8419    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
8420    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8421        let v = self.gpu.stream().clone_dtoh(d)?;
8422        self.gpu.stream().synchronize()?;
8423        Ok(v)
8424    }
8425    pub fn dtoh_u8_view(
8426        &self,
8427        d: &cudarc::driver::CudaView<u8>,
8428    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8429        let v = self.gpu.stream().clone_dtoh(d)?;
8430        self.gpu.stream().synchronize()?;
8431        Ok(v)
8432    }
8433    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8434        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8435        self.keep_if_capturing(&s);
8436        Ok(s)
8437    }
8438
8439    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
8440    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
8441    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
8442    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
8443    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
8444    /// back (or kept resident for graph replay). Returns the device token buffer.
8445    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
8446    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
8447    pub fn prob_of_token_device(
8448        &self,
8449        logits: &CudaSlice<f32>,
8450        tok: &CudaSlice<u32>,
8451        n_vocab: usize,
8452    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8453        let nb = ARGMAX_NB;
8454        let mut part = self.alloc_uninit::<f32>(nb)?;
8455        let mut p = self.alloc_uninit::<f32>(1)?;
8456        let f1 = self.func("prob_of_token_partial_f32");
8457        let cfg1 = LaunchConfig {
8458            grid_dim: (nb as u32, 1, 1),
8459            block_dim: (256, 1, 1),
8460            shared_mem_bytes: 0,
8461        };
8462        let nv = n_vocab as i32;
8463        let __s_b1 = self.gpu.stream();
8464        let mut b1 = __s_b1.launch_builder(&f1);
8465        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8466        unsafe {
8467            b1.launch(cfg1)?;
8468        }
8469        let f2 = self.func("prob_of_token_final_f32");
8470        let cfg2 = LaunchConfig {
8471            grid_dim: (1, 1, 1),
8472            block_dim: (256, 1, 1),
8473            shared_mem_bytes: 0,
8474        };
8475        let nbi = nb as i32;
8476        let __s_b2 = self.gpu.stream();
8477        let mut b2 = __s_b2.launch_builder(&f2);
8478        b2.arg(&part).arg(&mut p).arg(&nbi);
8479        unsafe {
8480            b2.launch(cfg2)?;
8481        }
8482        Ok(p)
8483    }
8484
8485    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
8486    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
8487    /// where the host reads the p-min confidence between replays. Same kernels, same math.
8488    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
8489    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
8490    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
8491    pub fn prob_of_token_device_col(
8492        &self,
8493        logits: &CudaSlice<f32>,
8494        tok_all: &CudaSlice<u32>,
8495        tok_idx: usize,
8496        p_out: &mut CudaSlice<f32>,
8497        p_idx: usize,
8498        n_vocab: usize,
8499    ) -> Result<(), Box<dyn std::error::Error>> {
8500        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8501        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8502        let nb = ARGMAX_NB;
8503        let mut part = self.alloc_uninit::<f32>(nb)?;
8504        let f1 = self.func("prob_of_token_partial_f32");
8505        let cfg1 = LaunchConfig {
8506            grid_dim: (nb as u32, 1, 1),
8507            block_dim: (256, 1, 1),
8508            shared_mem_bytes: 0,
8509        };
8510        let nv = n_vocab as i32;
8511        let __s_b1 = self.gpu.stream();
8512        let mut b1 = __s_b1.launch_builder(&f1);
8513        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8514        unsafe {
8515            b1.launch(cfg1)?;
8516        }
8517        let f2 = self.func("prob_of_token_final_f32");
8518        let cfg2 = LaunchConfig {
8519            grid_dim: (1, 1, 1),
8520            block_dim: (256, 1, 1),
8521            shared_mem_bytes: 0,
8522        };
8523        let nbi = nb as i32;
8524        let __s_b2 = self.gpu.stream();
8525        let mut b2 = __s_b2.launch_builder(&f2);
8526        b2.arg(&part).arg(&mut p_v).arg(&nbi);
8527        unsafe {
8528            b2.launch(cfg2)?;
8529        }
8530        Ok(())
8531    }
8532
8533    pub fn prob_of_token_device_into(
8534        &self,
8535        logits: &CudaSlice<f32>,
8536        tok: &CudaSlice<u32>,
8537        p_out: &mut CudaSlice<f32>,
8538        n_vocab: usize,
8539    ) -> Result<(), Box<dyn std::error::Error>> {
8540        let nb = ARGMAX_NB;
8541        let mut part = self.alloc_uninit::<f32>(nb)?;
8542        let f1 = self.func("prob_of_token_partial_f32");
8543        let cfg1 = LaunchConfig {
8544            grid_dim: (nb as u32, 1, 1),
8545            block_dim: (256, 1, 1),
8546            shared_mem_bytes: 0,
8547        };
8548        let nv = n_vocab as i32;
8549        let __s_b1 = self.gpu.stream();
8550        let mut b1 = __s_b1.launch_builder(&f1);
8551        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8552        unsafe {
8553            b1.launch(cfg1)?;
8554        }
8555        let f2 = self.func("prob_of_token_final_f32");
8556        let cfg2 = LaunchConfig {
8557            grid_dim: (1, 1, 1),
8558            block_dim: (256, 1, 1),
8559            shared_mem_bytes: 0,
8560        };
8561        let nbi = nb as i32;
8562        let __s_b2 = self.gpu.stream();
8563        let mut b2 = __s_b2.launch_builder(&f2);
8564        b2.arg(&part).arg(p_out).arg(&nbi);
8565        unsafe {
8566            b2.launch(cfg2)?;
8567        }
8568        Ok(())
8569    }
8570
8571    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
8572    /// (graph-constant params, device-varying index). Capture-safe.
8573    pub fn u32_hist_append(
8574        &self,
8575        tok: &CudaSlice<u32>,
8576        hist: &mut CudaSlice<u32>,
8577        idx: &mut CudaSlice<i32>,
8578    ) -> Result<(), Box<dyn std::error::Error>> {
8579        let f = self.func("u32_hist_append");
8580        let cfg = LaunchConfig {
8581            grid_dim: (1, 1, 1),
8582            block_dim: (32, 1, 1),
8583            shared_mem_bytes: 0,
8584        };
8585        let __s_b = self.gpu.stream();
8586        let mut b = __s_b.launch_builder(&f);
8587        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8588        unsafe {
8589            b.launch(cfg)?;
8590        }
8591        Ok(())
8592    }
8593
8594    pub fn argmax_token_device(
8595        &self,
8596        logits: &CudaSlice<f32>,
8597        n_vocab: usize,
8598    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8599        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8600        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8601        Ok(tok)
8602    }
8603    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
8604    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
8605    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
8606    /// pointer is baked once and the token id never round-trips to host inside steady state. The
8607    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
8608    /// captured passes bake fixed addresses.
8609    pub fn argmax_token_device_into(
8610        &self,
8611        logits: &CudaSlice<f32>,
8612        tok: &mut CudaSlice<u32>,
8613        n_vocab: usize,
8614    ) -> Result<(), Box<dyn std::error::Error>> {
8615        let nb = ARGMAX_NB;
8616        let f1 = self.func("argmax_partial_f32");
8617        let f2 = self.func("argmax_final_f32");
8618        let mut guard = self.argmax_partials.lock().unwrap();
8619        if guard.is_none() {
8620            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8621            // buffers carry no cudarc events (illegal inside capture).
8622            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8623            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8624            *guard = Some((pv, pi));
8625        }
8626        let (part_v, part_i) = guard.as_mut().unwrap();
8627        let nv = n_vocab as i32;
8628        let nbi = nb as i32;
8629        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8630        let cfg1 = LaunchConfig {
8631            grid_dim: (nb as u32, 1, 1),
8632            block_dim: (256, 1, 1),
8633            shared_mem_bytes: 0,
8634        };
8635        let __s_b1 = self.gpu.stream();
8636        let mut b1 = __s_b1.launch_builder(&f1);
8637        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8638        unsafe {
8639            b1.launch(cfg1)?;
8640        }
8641        // pass 2: one block reduces NB partials -> token_out[0].
8642        let cfg2 = LaunchConfig {
8643            grid_dim: (1, 1, 1),
8644            block_dim: (256, 1, 1),
8645            shared_mem_bytes: 0,
8646        };
8647        let __s_b2 = self.gpu.stream();
8648        let mut b2 = __s_b2.launch_builder(&f2);
8649        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8650        unsafe {
8651            b2.launch(cfg2)?;
8652        }
8653        Ok(())
8654    }
8655    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8656    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8657    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8658    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8659    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8660    pub fn argmax_token_device_col(
8661        &self,
8662        logits: &CudaSlice<f32>,
8663        col: usize,
8664        n_vocab: usize,
8665        toks: &mut CudaSlice<u32>,
8666        out_idx: usize,
8667    ) -> Result<(), Box<dyn std::error::Error>> {
8668        let nb = ARGMAX_NB;
8669        let f1 = self.func("argmax_partial_f32");
8670        let f2 = self.func("argmax_final_f32");
8671        let mut guard = self.argmax_partials.lock().unwrap();
8672        if guard.is_none() {
8673            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8674            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8675            *guard = Some((pv, pi));
8676        }
8677        let (part_v, part_i) = guard.as_mut().unwrap();
8678        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8679        let nv = n_vocab as i32;
8680        let nbi = nb as i32;
8681        let cfg1 = LaunchConfig {
8682            grid_dim: (nb as u32, 1, 1),
8683            block_dim: (256, 1, 1),
8684            shared_mem_bytes: 0,
8685        };
8686        let __s_b1 = self.gpu.stream();
8687        let mut b1 = __s_b1.launch_builder(&f1);
8688        b1.arg(&col_view)
8689            .arg(&mut *part_v)
8690            .arg(&mut *part_i)
8691            .arg(&nv);
8692        unsafe {
8693            b1.launch(cfg1)?;
8694        }
8695        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8696        let cfg2 = LaunchConfig {
8697            grid_dim: (1, 1, 1),
8698            block_dim: (256, 1, 1),
8699            shared_mem_bytes: 0,
8700        };
8701        let __s_b2 = self.gpu.stream();
8702        let mut b2 = __s_b2.launch_builder(&f2);
8703        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8704        unsafe {
8705            b2.launch(cfg2)?;
8706        }
8707        Ok(())
8708    }
8709    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8710    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8711        Ok(self.gpu.stream().clone_htod(v)?)
8712    }
8713    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8714        let v = self.gpu.stream().clone_dtoh(d)?;
8715        self.gpu.stream().synchronize()?;
8716        Ok(v)
8717    }
8718
8719    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8720        let v = self.gpu.stream().clone_dtoh(d)?;
8721        self.gpu.stream().synchronize()?;
8722        Ok(v)
8723    }
8724    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8725    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8726    /// contents change every step, the address must not, so a captured graph can read it).
8727    pub fn htod_u32_into(
8728        &self,
8729        dst: &mut CudaSlice<u32>,
8730        src: &[u32],
8731    ) -> Result<(), Box<dyn std::error::Error>> {
8732        let mut view = dst.slice_mut(0..src.len());
8733        self.gpu.stream().memcpy_htod(src, &mut view)?;
8734        Ok(())
8735    }
8736
8737    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8738    /// table without changing the device address its reconcile kernel consumes.
8739    pub fn htod_i32_into(
8740        &self,
8741        dst: &mut CudaSlice<i32>,
8742        src: &[i32],
8743    ) -> Result<(), Box<dyn std::error::Error>> {
8744        let mut view = dst.slice_mut(0..src.len());
8745        self.gpu.stream().memcpy_htod(src, &mut view)?;
8746        Ok(())
8747    }
8748
8749    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8750        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8751        self.keep_if_capturing(&s);
8752        Ok(s)
8753    }
8754    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8755    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8756    pub fn embed_gather_device_into(
8757        &self,
8758        embd: &CudaSlice<u8>,
8759        token_d: &CudaSlice<u32>,
8760        x_out: &mut CudaSlice<f32>,
8761        n_embd: usize,
8762        qtype: i32,
8763        row_bytes: usize,
8764    ) -> Result<(), Box<dyn std::error::Error>> {
8765        let f = self.func("embed_gather_u32");
8766        let cfg = LaunchConfig {
8767            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8768            block_dim: (256, 1, 1),
8769            shared_mem_bytes: 0,
8770        };
8771        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8772        let __s_b = self.gpu.stream();
8773        let mut b = __s_b.launch_builder(&f);
8774        b.arg(embd)
8775            .arg(token_d)
8776            .arg(x_out)
8777            .arg(&ne)
8778            .arg(&qt)
8779            .arg(&rb);
8780        unsafe {
8781            b.launch(cfg)?;
8782        }
8783        Ok(())
8784    }
8785    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
8786    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8787        let v = self.gpu.stream().clone_dtoh(d)?;
8788        self.gpu.stream().synchronize()?;
8789        Ok(v[0])
8790    }
8791    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
8792    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
8793    /// the counter value after the throwaway capture warmups corrupt it.
8794    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
8795    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
8796    /// copy (fine at stream-idle boundaries, poison mid-round).
8797    pub fn i32_set_k(
8798        &self,
8799        dst: &mut CudaSlice<i32>,
8800        v: i32,
8801    ) -> Result<(), Box<dyn std::error::Error>> {
8802        let f = self.func("i32_set_k");
8803        let cfg = LaunchConfig {
8804            grid_dim: (1, 1, 1),
8805            block_dim: (1, 1, 1),
8806            shared_mem_bytes: 0,
8807        };
8808        let idx = 0i32;
8809        let __s_b = self.gpu.stream();
8810        let mut b = __s_b.launch_builder(&f);
8811        b.arg(dst).arg(&v).arg(&idx);
8812        unsafe {
8813            b.launch(cfg)?;
8814        }
8815        Ok(())
8816    }
8817
8818    pub fn set_i32_one(
8819        &self,
8820        d: &mut CudaSlice<i32>,
8821        v: i32,
8822    ) -> Result<(), Box<dyn std::error::Error>> {
8823        self.gpu.stream().memcpy_htod(&[v], d)?;
8824        Ok(())
8825    }
8826    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
8827    /// during priming / capture-state restore.
8828    pub fn set_u32_one(
8829        &self,
8830        d: &mut CudaSlice<u32>,
8831        v: u32,
8832    ) -> Result<(), Box<dyn std::error::Error>> {
8833        self.gpu.stream().memcpy_htod(&[v], d)?;
8834        Ok(())
8835    }
8836    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
8837    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8838        let v = self.gpu.stream().clone_dtoh(d)?;
8839        self.gpu.stream().synchronize()?;
8840        Ok(v[0])
8841    }
8842    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
8843    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8844        Ok(self.gpu.stream().clone_htod(bytes)?)
8845    }
8846    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
8847    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
8848    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
8849    pub fn embed_gather_device(
8850        &self,
8851        embd: &CudaSlice<u8>,
8852        token_d: &CudaSlice<u32>,
8853        n_embd: usize,
8854        qtype: i32,
8855        row_bytes: usize,
8856    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8857        let f = self.func("embed_gather_u32");
8858        let mut x = self.alloc_uninit::<f32>(n_embd)?;
8859        let cfg = LaunchConfig {
8860            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8861            block_dim: (256, 1, 1),
8862            shared_mem_bytes: 0,
8863        };
8864        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8865        let __s_b = self.gpu.stream();
8866        let mut b = __s_b.launch_builder(&f);
8867        b.arg(embd)
8868            .arg(token_d)
8869            .arg(&mut x)
8870            .arg(&ne)
8871            .arg(&qt)
8872            .arg(&rb);
8873        unsafe {
8874            b.launch(cfg)?;
8875        }
8876        Ok(x)
8877    }
8878
8879    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
8880    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
8881    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
8882    pub fn embed_gather_device_t(
8883        &self,
8884        embd: &CudaSlice<u8>,
8885        tokens: &[u32],
8886        n_embd: usize,
8887        qtype: i32,
8888        row_bytes: usize,
8889    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8890        let t = tokens.len();
8891        let tok_d = self.gpu.stream().clone_htod(tokens)?;
8892        let f = self.func("embed_gather_u32_t");
8893        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8894        let cfg = LaunchConfig {
8895            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8896            block_dim: (256, 1, 1),
8897            shared_mem_bytes: 0,
8898        };
8899        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8900        let __s_b = self.gpu.stream();
8901        let mut b = __s_b.launch_builder(&f);
8902        b.arg(embd)
8903            .arg(&tok_d)
8904            .arg(&mut x)
8905            .arg(&ne)
8906            .arg(&qt)
8907            .arg(&rb)
8908            .arg(&ti);
8909        unsafe {
8910            b.launch(cfg)?;
8911        }
8912        Ok(x)
8913    }
8914
8915    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
8916    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
8917    /// as embed_gather_device_t — bit-identical rows.
8918    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
8919    pub fn embed_gather_device_tv(
8920        &self,
8921        embd: &CudaSlice<u8>,
8922        tok_v: &cudarc::driver::CudaView<u32>,
8923        t: usize,
8924        n_embd: usize,
8925        qtype: i32,
8926        row_bytes: usize,
8927    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8928        let f = self.func("embed_gather_u32_t");
8929        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8930        let cfg = LaunchConfig {
8931            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8932            block_dim: (256, 1, 1),
8933            shared_mem_bytes: 0,
8934        };
8935        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8936        let __s_b = self.gpu.stream();
8937        let mut b = __s_b.launch_builder(&f);
8938        b.arg(embd)
8939            .arg(tok_v)
8940            .arg(&mut x)
8941            .arg(&ne)
8942            .arg(&qt)
8943            .arg(&rb)
8944            .arg(&ti);
8945        unsafe {
8946            b.launch(cfg)?;
8947        }
8948        Ok(x)
8949    }
8950
8951    pub fn embed_gather_device_td(
8952        &self,
8953        embd: &CudaSlice<u8>,
8954        tok_d: &CudaSlice<u32>,
8955        t: usize,
8956        n_embd: usize,
8957        qtype: i32,
8958        row_bytes: usize,
8959    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8960        let f = self.func("embed_gather_u32_t");
8961        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8962        let cfg = LaunchConfig {
8963            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8964            block_dim: (256, 1, 1),
8965            shared_mem_bytes: 0,
8966        };
8967        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8968        let __s_b = self.gpu.stream();
8969        let mut b = __s_b.launch_builder(&f);
8970        b.arg(embd)
8971            .arg(tok_d)
8972            .arg(&mut x)
8973            .arg(&ne)
8974            .arg(&qt)
8975            .arg(&rb)
8976            .arg(&ti);
8977        unsafe {
8978            b.launch(cfg)?;
8979        }
8980        Ok(x)
8981    }
8982
8983    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
8984    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
8985    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
8986    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
8987    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
8988    #[inline]
8989    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
8990    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8991        if self
8992            .capture_keep_on
8993            .load(std::sync::atomic::Ordering::Relaxed)
8994        {
8995            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8996        }
8997    }
8998
8999    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9000        &self,
9001        n: usize,
9002    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9003        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9004        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
9005        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
9006        // not cover engine-internal buffers). Debug-only: massive launch overhead.
9007        {
9008            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9009            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9010                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
9011                use cudarc::driver::DevicePtrMut;
9012                let n_bytes = s.len() * std::mem::size_of::<T>();
9013                let stream = self.gpu.stream();
9014                let (p_, _g) = s.device_ptr_mut(&stream);
9015                unsafe {
9016                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9017                        .result()?;
9018                }
9019            }
9020        }
9021        self.keep_if_capturing(&s);
9022        Ok(s)
9023    }
9024
9025    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
9026    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
9027    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
9028    /// consumers alloc through this (m=1 decode arms).
9029    pub fn uninit_q8_pair(
9030        &self,
9031        n: usize,
9032    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9033        Ok((
9034            self.alloc_uninit::<i8>(n)?,
9035            self.alloc_uninit::<f32>(n / 32)?,
9036        ))
9037    }
9038
9039    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9040        self.alloc_uninit::<f32>(n)
9041    }
9042
9043    /// i8 uninitialized scratch (same contract as `uninit`).
9044    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9045        self.alloc_uninit::<i8>(n)
9046    }
9047
9048    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
9049    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
9050    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
9051    #[allow(clippy::too_many_arguments)]
9052    pub fn rms_norm3(
9053        &self,
9054        x: &CudaSlice<f32>,
9055        w0: &CudaSlice<f32>,
9056        w1: &CudaSlice<f32>,
9057        w2: &CudaSlice<f32>,
9058        d0: &mut CudaSlice<f32>,
9059        d1: &mut CudaSlice<f32>,
9060        d2: &mut CudaSlice<f32>,
9061        ncols: usize,
9062        nrows: usize,
9063        eps: f32,
9064    ) -> Result<(), Box<dyn std::error::Error>> {
9065        let f = self.func("rms_norm3_f32");
9066        let cfg = LaunchConfig {
9067            grid_dim: (nrows as u32, 1, 1),
9068            block_dim: (rms_block(), 1, 1),
9069            shared_mem_bytes: 0,
9070        };
9071        let (nc, e) = (ncols as i32, eps);
9072        let __s_b = self.gpu.stream();
9073        let mut b = __s_b.launch_builder(&f);
9074        b.arg(x)
9075            .arg(w0)
9076            .arg(w1)
9077            .arg(w2)
9078            .arg(d0)
9079            .arg(d1)
9080            .arg(d2)
9081            .arg(&nc)
9082            .arg(&e);
9083        unsafe {
9084            b.launch(cfg)?;
9085        }
9086        Ok(())
9087    }
9088
9089    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
9090    #[allow(clippy::too_many_arguments)]
9091    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
9092    /// piggybacks on the same conditions.
9093    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9094        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9095        *WARP_ON.get_or_init(|| {
9096            std::env::var("MEMRA_QKVNORM_W")
9097                .map(|v| v != "0")
9098                .unwrap_or(true)
9099        }) && ncols % 4 == 0
9100            && rows >= 64
9101    }
9102
9103    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
9104    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
9105    #[allow(clippy::too_many_arguments)]
9106    pub fn rms_norm_qkv_w4b(
9107        &self,
9108        q: &CudaSlice<f32>,
9109        k: &CudaSlice<f32>,
9110        v: &CudaSlice<f32>,
9111        wq: &CudaSlice<f32>,
9112        wk: &CudaSlice<f32>,
9113        wv: &CudaSlice<f32>,
9114        dq: &mut CudaSlice<f32>,
9115        dk: &mut CudaSlice<f32>,
9116        dv: &mut CudaSlice<f32>,
9117        dvb: &mut CudaSlice<u8>,
9118        ncols: usize,
9119        rq: usize,
9120        rk: usize,
9121        eps: f32,
9122        vf16: bool,
9123    ) -> Result<(), Box<dyn std::error::Error>> {
9124        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9125        let f = self.func("rms_norm_qkv_w4b_f32");
9126        let rows = (rq + 2 * rk) as u32;
9127        let cfg = LaunchConfig {
9128            grid_dim: (rows.div_ceil(8), 1, 1),
9129            block_dim: (256, 1, 1),
9130            shared_mem_bytes: 0,
9131        };
9132        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9133        let vf = vf16 as i32;
9134        let __s_b = self.gpu.stream();
9135        let mut b = __s_b.launch_builder(&f);
9136        b.arg(q)
9137            .arg(k)
9138            .arg(v)
9139            .arg(wq)
9140            .arg(wk)
9141            .arg(wv)
9142            .arg(dq)
9143            .arg(dk)
9144            .arg(dv)
9145            .arg(&mut *dvb)
9146            .arg(&nc)
9147            .arg(&rqi)
9148            .arg(&rki)
9149            .arg(&rvi)
9150            .arg(&e)
9151            .arg(&vf);
9152        unsafe {
9153            b.launch(cfg)?;
9154        }
9155        Ok(())
9156    }
9157
9158    pub fn rms_norm_qkv(
9159        &self,
9160        q: &CudaSlice<f32>,
9161        k: &CudaSlice<f32>,
9162        v: &CudaSlice<f32>,
9163        wq: &CudaSlice<f32>,
9164        wk: &CudaSlice<f32>,
9165        wv: &CudaSlice<f32>,
9166        dq: &mut CudaSlice<f32>,
9167        dk: &mut CudaSlice<f32>,
9168        dv: &mut CudaSlice<f32>,
9169        ncols: usize,
9170        rq: usize,
9171        rk: usize,
9172        eps: f32,
9173    ) -> Result<(), Box<dyn std::error::Error>> {
9174        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
9175        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
9176        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
9177        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9178        let warp_on = *WARP_ON.get_or_init(|| {
9179            std::env::var("MEMRA_QKVNORM_W")
9180                .map(|v| v != "0")
9181                .unwrap_or(true)
9182        });
9183        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
9184        // replay numerics are untouched on every model; only prefill depth takes the new config.
9185        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9186            let f = self.func("rms_norm_qkv_w4_f32");
9187            let rows = (rq + 2 * rk) as u32;
9188            let cfg = LaunchConfig {
9189                grid_dim: (rows.div_ceil(8), 1, 1),
9190                block_dim: (256, 1, 1),
9191                shared_mem_bytes: 0,
9192            };
9193            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9194            let __s_b = self.gpu.stream();
9195            let mut b = __s_b.launch_builder(&f);
9196            b.arg(q)
9197                .arg(k)
9198                .arg(v)
9199                .arg(wq)
9200                .arg(wk)
9201                .arg(wv)
9202                .arg(dq)
9203                .arg(dk)
9204                .arg(dv)
9205                .arg(&nc)
9206                .arg(&rqi)
9207                .arg(&rki)
9208                .arg(&rvi)
9209                .arg(&e);
9210            unsafe {
9211                b.launch(cfg)?;
9212            }
9213            return Ok(());
9214        }
9215        let f = self.func("rms_norm_qkv_f32");
9216        let grid = (rq + 2 * rk) as u32;
9217        let cfg = LaunchConfig {
9218            grid_dim: (grid, 1, 1),
9219            block_dim: (rms_block(), 1, 1),
9220            shared_mem_bytes: 0,
9221        };
9222        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9223        let __s_b = self.gpu.stream();
9224        let mut b = __s_b.launch_builder(&f);
9225        b.arg(q)
9226            .arg(k)
9227            .arg(v)
9228            .arg(wq)
9229            .arg(wk)
9230            .arg(wv)
9231            .arg(dq)
9232            .arg(dk)
9233            .arg(dv)
9234            .arg(&nc)
9235            .arg(&rqi)
9236            .arg(&rki)
9237            .arg(&e);
9238        unsafe {
9239            b.launch(cfg)?;
9240        }
9241        Ok(())
9242    }
9243
9244    /// gemma4 fused pair of rms_norms over two different inputs (same width).
9245    #[allow(clippy::too_many_arguments)]
9246    pub fn rms_norm2x(
9247        &self,
9248        a: &CudaSlice<f32>,
9249        bb: &CudaSlice<f32>,
9250        wa: &CudaSlice<f32>,
9251        wb: &CudaSlice<f32>,
9252        da: &mut CudaSlice<f32>,
9253        db: &mut CudaSlice<f32>,
9254        ncols: usize,
9255        nrows: usize,
9256        eps: f32,
9257    ) -> Result<(), Box<dyn std::error::Error>> {
9258        let f = self.func("rms_norm2x_f32");
9259        let cfg = LaunchConfig {
9260            grid_dim: (2 * nrows as u32, 1, 1),
9261            block_dim: (rms_block(), 1, 1),
9262            shared_mem_bytes: 0,
9263        };
9264        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9265        let __s_b = self.gpu.stream();
9266        let mut b = __s_b.launch_builder(&f);
9267        b.arg(a)
9268            .arg(bb)
9269            .arg(wa)
9270            .arg(wb)
9271            .arg(da)
9272            .arg(db)
9273            .arg(&nc)
9274            .arg(&nr)
9275            .arg(&e);
9276        unsafe {
9277            b.launch(cfg)?;
9278        }
9279        Ok(())
9280    }
9281
9282    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
9283    pub fn softcap(
9284        &self,
9285        y: &mut CudaSlice<f32>,
9286        cap: f32,
9287        n: usize,
9288    ) -> Result<(), Box<dyn std::error::Error>> {
9289        let f = self.func("softcap_f32");
9290        let cfg = LaunchConfig::for_num_elems(n as u32);
9291        let ni = n as i32;
9292        let __s_b = self.gpu.stream();
9293        let mut b = __s_b.launch_builder(&f);
9294        b.arg(y).arg(&cap).arg(&ni);
9295        unsafe {
9296            b.launch(cfg)?;
9297        }
9298        Ok(())
9299    }
9300
9301    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
9302    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
9303    pub fn mask_ids_rows(
9304        &self,
9305        y: &mut CudaSlice<f32>,
9306        ids: &CudaSlice<i32>,
9307        n_ids: usize,
9308        n_vocab: usize,
9309        t: usize,
9310    ) -> Result<(), Box<dyn std::error::Error>> {
9311        let f = self.func("mask_ids_rows_f32");
9312        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9313        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9314        let __s_b = self.gpu.stream();
9315        let mut b = __s_b.launch_builder(&f);
9316        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9317        unsafe {
9318            b.launch(cfg)?;
9319        }
9320        Ok(())
9321    }
9322
9323    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
9324    #[allow(clippy::too_many_arguments)]
9325    pub fn add_scale_rms_norm(
9326        &self,
9327        a: &CudaSlice<f32>,
9328        b_in: &CudaSlice<f32>,
9329        c: f32,
9330        w: &CudaSlice<f32>,
9331        res: &mut CudaSlice<f32>,
9332        dst: &mut CudaSlice<f32>,
9333        ncols: usize,
9334        nrows: usize,
9335        eps: f32,
9336    ) -> Result<(), Box<dyn std::error::Error>> {
9337        let f = self.func("add_scale_rms_norm_f32");
9338        let cfg = LaunchConfig {
9339            grid_dim: (nrows as u32, 1, 1),
9340            block_dim: (rms_block(), 1, 1),
9341            shared_mem_bytes: 0,
9342        };
9343        let (nc, e2) = (ncols as i32, eps);
9344        let __s_b = self.gpu.stream();
9345        let mut b = __s_b.launch_builder(&f);
9346        b.arg(a)
9347            .arg(b_in)
9348            .arg(&c)
9349            .arg(w)
9350            .arg(res)
9351            .arg(dst)
9352            .arg(&nc)
9353            .arg(&e2);
9354        unsafe {
9355            b.launch(cfg)?;
9356        }
9357        Ok(())
9358    }
9359
9360    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
9361    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
9362    #[allow(clippy::too_many_arguments)]
9363    pub fn add_scale_rms_norm_q8_1(
9364        &self,
9365        a: &CudaSlice<f32>,
9366        b_in: &CudaSlice<f32>,
9367        c: f32,
9368        w: &CudaSlice<f32>,
9369        res: &mut CudaSlice<f32>,
9370        ncols: usize,
9371        nrows: usize,
9372        eps: f32,
9373    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9374        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9375        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9376        let (nc, e2) = (ncols as i32, eps);
9377        if Self::pdl_on() && Self::pdl_wb_on() {
9378            {
9379                use cudarc::driver::{DevicePtr, DevicePtrMut};
9380                let s = &self.gpu.stream();
9381                let (pa, _g0) = a.device_ptr(s);
9382                let (pb, _g1) = b_in.device_ptr(s);
9383                let (pw, _g2) = w.device_ptr(s);
9384                let (pr, _g3) = res.device_ptr_mut(s);
9385                let (pq, _g4) = out_q.device_ptr_mut(s);
9386                let (pd, _g5) = out_d.device_ptr_mut(s);
9387                let mut ps = [
9388                    &pa as *const _ as *mut std::ffi::c_void,
9389                    &pb as *const _ as *mut _,
9390                    &c as *const _ as *mut _,
9391                    &pw as *const _ as *mut _,
9392                    &pr as *const _ as *mut _,
9393                    &pq as *const _ as *mut _,
9394                    &pd as *const _ as *mut _,
9395                    &nc as *const _ as *mut _,
9396                    &e2 as *const _ as *mut _,
9397                ];
9398                unsafe {
9399                    self.launch_pdl(
9400                        "add_scale_rms_norm_q8_1",
9401                        (nrows as u32, 1, 1),
9402                        (rms_block(), 1, 1),
9403                        &mut ps,
9404                    )?;
9405                }
9406            }
9407            return Ok((out_q, out_d));
9408        }
9409        let f = self.func("add_scale_rms_norm_q8_1");
9410        let cfg = LaunchConfig {
9411            grid_dim: (nrows as u32, 1, 1),
9412            block_dim: (rms_block(), 1, 1),
9413            shared_mem_bytes: 0,
9414        };
9415        let __s_b = self.gpu.stream();
9416        let mut b = __s_b.launch_builder(&f);
9417        b.arg(a)
9418            .arg(b_in)
9419            .arg(&c)
9420            .arg(w)
9421            .arg(res)
9422            .arg(&mut out_q)
9423            .arg(&mut out_d)
9424            .arg(&nc)
9425            .arg(&e2);
9426        unsafe {
9427            b.launch(cfg)?;
9428        }
9429        Ok((out_q, out_d))
9430    }
9431
9432    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
9433    #[allow(clippy::too_many_arguments)]
9434    pub fn add_scale_rms_norm_q8_1_into(
9435        &self,
9436        a: &CudaSlice<f32>,
9437        b_in: &CudaSlice<f32>,
9438        c: f32,
9439        w: &CudaSlice<f32>,
9440        res: &mut CudaSlice<f32>,
9441        ncols: usize,
9442        nrows: usize,
9443        eps: f32,
9444        out_q: &mut CudaSlice<i8>,
9445        out_d: &mut CudaSlice<f32>,
9446    ) -> Result<(), Box<dyn std::error::Error>> {
9447        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9448        let (nc, e2) = (ncols as i32, eps);
9449        if Self::pdl_on() && Self::pdl_wb_on() {
9450            use cudarc::driver::{DevicePtr, DevicePtrMut};
9451            let s = &self.gpu.stream();
9452            let (pa, _g0) = a.device_ptr(s);
9453            let (pb, _g1) = b_in.device_ptr(s);
9454            let (pw, _g2) = w.device_ptr(s);
9455            let (pr, _g3) = res.device_ptr_mut(s);
9456            let (pq, _g4) = out_q.device_ptr_mut(s);
9457            let (pd, _g5) = out_d.device_ptr_mut(s);
9458            let mut ps = [
9459                &pa as *const _ as *mut std::ffi::c_void,
9460                &pb as *const _ as *mut _,
9461                &c as *const _ as *mut _,
9462                &pw as *const _ as *mut _,
9463                &pr as *const _ as *mut _,
9464                &pq as *const _ as *mut _,
9465                &pd as *const _ as *mut _,
9466                &nc as *const _ as *mut _,
9467                &e2 as *const _ as *mut _,
9468            ];
9469            unsafe {
9470                self.launch_pdl(
9471                    "add_scale_rms_norm_q8_1",
9472                    (nrows as u32, 1, 1),
9473                    (rms_block(), 1, 1),
9474                    &mut ps,
9475                )?;
9476            }
9477            return Ok(());
9478        }
9479        let f = self.func("add_scale_rms_norm_q8_1");
9480        let cfg = LaunchConfig {
9481            grid_dim: (nrows as u32, 1, 1),
9482            block_dim: (rms_block(), 1, 1),
9483            shared_mem_bytes: 0,
9484        };
9485        let __s_b = self.gpu.stream();
9486        let mut b = __s_b.launch_builder(&f);
9487        b.arg(a)
9488            .arg(b_in)
9489            .arg(&c)
9490            .arg(w)
9491            .arg(res)
9492            .arg(&mut *out_q)
9493            .arg(&mut *out_d)
9494            .arg(&nc)
9495            .arg(&e2);
9496        unsafe {
9497            b.launch(cfg)?;
9498        }
9499        Ok(())
9500    }
9501
9502    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
9503    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
9504    #[allow(clippy::too_many_arguments)]
9505    pub fn rms_pre_add_scale_rms_norm_q8_1(
9506        &self,
9507        a: &CudaSlice<f32>,
9508        wa: &CudaSlice<f32>,
9509        b_in: &CudaSlice<f32>,
9510        c: f32,
9511        w: &CudaSlice<f32>,
9512        res: &mut CudaSlice<f32>,
9513        ncols: usize,
9514        nrows: usize,
9515        eps: f32,
9516    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9517        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9518        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9519        let (nc, e2) = (ncols as i32, eps);
9520        if Self::pdl_on() {
9521            {
9522                use cudarc::driver::{DevicePtr, DevicePtrMut};
9523                let s = &self.gpu.stream();
9524                let (pa, _g0) = a.device_ptr(s);
9525                let (pwa, _g1) = wa.device_ptr(s);
9526                let (pb, _g2) = b_in.device_ptr(s);
9527                let (pw, _g3) = w.device_ptr(s);
9528                let (pr, _g4) = res.device_ptr_mut(s);
9529                let (pq, _g5) = out_q.device_ptr_mut(s);
9530                let (pd, _g6) = out_d.device_ptr_mut(s);
9531                let mut ps = [
9532                    &pa as *const _ as *mut std::ffi::c_void,
9533                    &pwa as *const _ as *mut _,
9534                    &pb as *const _ as *mut _,
9535                    &c as *const _ as *mut _,
9536                    &pw as *const _ as *mut _,
9537                    &pr as *const _ as *mut _,
9538                    &pq as *const _ as *mut _,
9539                    &pd as *const _ as *mut _,
9540                    &nc as *const _ as *mut _,
9541                    &e2 as *const _ as *mut _,
9542                ];
9543                unsafe {
9544                    self.launch_pdl(
9545                        "rms_pre_add_scale_rms_norm_q8_1",
9546                        (nrows as u32, 1, 1),
9547                        (rms_block(), 1, 1),
9548                        &mut ps,
9549                    )?;
9550                }
9551            }
9552            return Ok((out_q, out_d));
9553        }
9554        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9555        let cfg = LaunchConfig {
9556            grid_dim: (nrows as u32, 1, 1),
9557            block_dim: (rms_block(), 1, 1),
9558            shared_mem_bytes: 0,
9559        };
9560        let __s_b = self.gpu.stream();
9561        let mut b = __s_b.launch_builder(&f);
9562        b.arg(a)
9563            .arg(wa)
9564            .arg(b_in)
9565            .arg(&c)
9566            .arg(w)
9567            .arg(res)
9568            .arg(&mut out_q)
9569            .arg(&mut out_d)
9570            .arg(&nc)
9571            .arg(&e2);
9572        unsafe {
9573            b.launch(cfg)?;
9574        }
9575        Ok((out_q, out_d))
9576    }
9577
9578    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
9579    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
9580    pub fn gelu_tanh_mul_q8_1(
9581        &self,
9582        gate: &CudaSlice<f32>,
9583        up: &cudarc::driver::CudaView<f32>,
9584        act: &mut CudaSlice<f32>,
9585        ncols: usize,
9586        nrows: usize,
9587    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9588        debug_assert!(ncols % 128 == 0);
9589        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9590        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9591        let nc = ncols as i32;
9592        if Self::pdl_on() {
9593            {
9594                use cudarc::driver::{DevicePtr, DevicePtrMut};
9595                let s = &self.gpu.stream();
9596                let (pg, _g0) = gate.device_ptr(s);
9597                let (pu, _g1) = up.device_ptr(s);
9598                let (pact, _g2) = act.device_ptr_mut(s);
9599                let (pq, _g3) = out_q.device_ptr_mut(s);
9600                let (pd, _g4) = out_d.device_ptr_mut(s);
9601                let mut ps = [
9602                    &pg as *const _ as *mut std::ffi::c_void,
9603                    &pu as *const _ as *mut _,
9604                    &pact as *const _ as *mut _,
9605                    &pq as *const _ as *mut _,
9606                    &pd as *const _ as *mut _,
9607                    &nc as *const _ as *mut _,
9608                ];
9609                unsafe {
9610                    self.launch_pdl(
9611                        "gelu_tanh_mul_q8_1",
9612                        (nrows as u32, 1, 1),
9613                        (rms_block(), 1, 1),
9614                        &mut ps,
9615                    )?;
9616                }
9617            }
9618            return Ok((out_q, out_d));
9619        }
9620        let f = self.func("gelu_tanh_mul_q8_1");
9621        let cfg = LaunchConfig {
9622            grid_dim: (nrows as u32, 1, 1),
9623            block_dim: (rms_block(), 1, 1),
9624            shared_mem_bytes: 0,
9625        };
9626        let __s_b = self.gpu.stream();
9627        let mut b = __s_b.launch_builder(&f);
9628        b.arg(gate)
9629            .arg(up)
9630            .arg(act)
9631            .arg(&mut out_q)
9632            .arg(&mut out_d)
9633            .arg(&nc);
9634        unsafe {
9635            b.launch(cfg)?;
9636        }
9637        Ok((out_q, out_d))
9638    }
9639
9640    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9641    #[allow(clippy::too_many_arguments)]
9642    pub fn gelu_tanh_mul_q8_1_into(
9643        &self,
9644        gate: &CudaSlice<f32>,
9645        up: &cudarc::driver::CudaView<f32>,
9646        act: &mut CudaSlice<f32>,
9647        ncols: usize,
9648        nrows: usize,
9649        out_q: &mut CudaSlice<i8>,
9650        out_d: &mut CudaSlice<f32>,
9651    ) -> Result<(), Box<dyn std::error::Error>> {
9652        debug_assert!(ncols % 128 == 0);
9653        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9654        let nc = ncols as i32;
9655        if Self::pdl_on() {
9656            use cudarc::driver::{DevicePtr, DevicePtrMut};
9657            let s = &self.gpu.stream();
9658            let (pg, _g0) = gate.device_ptr(s);
9659            let (pu, _g1) = up.device_ptr(s);
9660            let (pact, _g2) = act.device_ptr_mut(s);
9661            let (pq, _g3) = out_q.device_ptr_mut(s);
9662            let (pd, _g4) = out_d.device_ptr_mut(s);
9663            let mut ps = [
9664                &pg as *const _ as *mut std::ffi::c_void,
9665                &pu as *const _ as *mut _,
9666                &pact as *const _ as *mut _,
9667                &pq as *const _ as *mut _,
9668                &pd as *const _ as *mut _,
9669                &nc as *const _ as *mut _,
9670            ];
9671            unsafe {
9672                self.launch_pdl(
9673                    "gelu_tanh_mul_q8_1",
9674                    (nrows as u32, 1, 1),
9675                    (rms_block(), 1, 1),
9676                    &mut ps,
9677                )?;
9678            }
9679            return Ok(());
9680        }
9681        let f = self.func("gelu_tanh_mul_q8_1");
9682        let cfg = LaunchConfig {
9683            grid_dim: (nrows as u32, 1, 1),
9684            block_dim: (rms_block(), 1, 1),
9685            shared_mem_bytes: 0,
9686        };
9687        let __s_b = self.gpu.stream();
9688        let mut b = __s_b.launch_builder(&f);
9689        b.arg(gate)
9690            .arg(up)
9691            .arg(&mut *act)
9692            .arg(&mut *out_q)
9693            .arg(&mut *out_d)
9694            .arg(&nc);
9695        unsafe {
9696            b.launch(cfg)?;
9697        }
9698        Ok(())
9699    }
9700
9701    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9702    #[allow(clippy::too_many_arguments)]
9703    pub fn add_rms_norm3_q8z(
9704        &self,
9705        a: &CudaSlice<f32>,
9706        b_in: &CudaSlice<f32>,
9707        w0: &CudaSlice<f32>,
9708        w1: &CudaSlice<f32>,
9709        w2: &CudaSlice<f32>,
9710        res: &mut CudaSlice<f32>,
9711        out1: &mut CudaSlice<f32>,
9712        ncols: usize,
9713        nrows: usize,
9714        eps: f32,
9715    ) -> Result<
9716        (
9717            (CudaSlice<i8>, CudaSlice<f32>),
9718            (CudaSlice<i8>, CudaSlice<f32>),
9719        ),
9720        Box<dyn std::error::Error>,
9721    > {
9722        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9723        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9724        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9725        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9726        let f = self.func("add_rms_norm3_q8z_f32");
9727        let cfg = LaunchConfig {
9728            grid_dim: (nrows as u32, 1, 1),
9729            block_dim: (rms_block(), 1, 1),
9730            shared_mem_bytes: 0,
9731        };
9732        let (nc, e2) = (ncols as i32, eps);
9733        let __s_b = self.gpu.stream();
9734        let mut b = __s_b.launch_builder(&f);
9735        b.arg(a)
9736            .arg(b_in)
9737            .arg(w0)
9738            .arg(w1)
9739            .arg(w2)
9740            .arg(res)
9741            .arg(&mut q0)
9742            .arg(&mut d0)
9743            .arg(out1)
9744            .arg(&mut q2)
9745            .arg(&mut d2)
9746            .arg(&nc)
9747            .arg(&e2);
9748        unsafe {
9749            b.launch(cfg)?;
9750        }
9751        Ok(((q0, d0), (q2, d2)))
9752    }
9753
9754    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9755    #[allow(clippy::too_many_arguments)]
9756    pub fn add_rms_norm3(
9757        &self,
9758        a: &CudaSlice<f32>,
9759        b_in: &CudaSlice<f32>,
9760        w0: &CudaSlice<f32>,
9761        w1: &CudaSlice<f32>,
9762        w2: &CudaSlice<f32>,
9763        res: &mut CudaSlice<f32>,
9764        d0: &mut CudaSlice<f32>,
9765        d1: &mut CudaSlice<f32>,
9766        d2: &mut CudaSlice<f32>,
9767        ncols: usize,
9768        nrows: usize,
9769        eps: f32,
9770    ) -> Result<(), Box<dyn std::error::Error>> {
9771        let f = self.func("add_rms_norm3_f32");
9772        let cfg = LaunchConfig {
9773            grid_dim: (nrows as u32, 1, 1),
9774            block_dim: (rms_block(), 1, 1),
9775            shared_mem_bytes: 0,
9776        };
9777        let (nc, e2) = (ncols as i32, eps);
9778        let __s_b = self.gpu.stream();
9779        let mut b = __s_b.launch_builder(&f);
9780        b.arg(a)
9781            .arg(b_in)
9782            .arg(w0)
9783            .arg(w1)
9784            .arg(w2)
9785            .arg(res)
9786            .arg(d0)
9787            .arg(d1)
9788            .arg(d2)
9789            .arg(&nc)
9790            .arg(&e2);
9791        unsafe {
9792            b.launch(cfg)?;
9793        }
9794        Ok(())
9795    }
9796
9797    /// dst = (a + b) * c (residual add + layer scale, one launch).
9798    pub fn add_scale(
9799        &self,
9800        a: &CudaSlice<f32>,
9801        b_in: &CudaSlice<f32>,
9802        c: f32,
9803        dst: &mut CudaSlice<f32>,
9804        n: usize,
9805    ) -> Result<(), Box<dyn std::error::Error>> {
9806        let f = self.func("add_scale_f32");
9807        let cfg = LaunchConfig::for_num_elems(n as u32);
9808        let ni = n as i32;
9809        let __s_b = self.gpu.stream();
9810        let mut b = __s_b.launch_builder(&f);
9811        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9812        unsafe {
9813            b.launch(cfg)?;
9814        }
9815        Ok(())
9816    }
9817
9818    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
9819    pub fn layer_norm_bias(
9820        &self,
9821        x: &CudaSlice<f32>,
9822        w: &CudaSlice<f32>,
9823        b: &CudaSlice<f32>,
9824        dst: &mut CudaSlice<f32>,
9825        ncols: usize,
9826        nrows: usize,
9827        eps: f32,
9828    ) -> Result<(), Box<dyn std::error::Error>> {
9829        let f = self.func("layer_norm_bias_f32");
9830        let (nc, e) = (ncols as i32, eps);
9831        let cfg = LaunchConfig {
9832            grid_dim: (nrows as u32, 1, 1),
9833            block_dim: (256, 1, 1),
9834            shared_mem_bytes: 0,
9835        };
9836        let __s_b = self.gpu.stream();
9837        let mut lb = __s_b.launch_builder(&f);
9838        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9839        unsafe {
9840            lb.launch(cfg)?;
9841        }
9842        Ok(())
9843    }
9844
9845    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
9846    pub fn gelu_tanh(
9847        &self,
9848        x: &CudaSlice<f32>,
9849        dst: &mut CudaSlice<f32>,
9850        n: usize,
9851    ) -> Result<(), Box<dyn std::error::Error>> {
9852        let f = self.func("gelu_tanh_f32");
9853        let ni = n as i64;
9854        let cfg = LaunchConfig {
9855            grid_dim: (n.div_ceil(256) as u32, 1, 1),
9856            block_dim: (256, 1, 1),
9857            shared_mem_bytes: 0,
9858        };
9859        let __s_b = self.gpu.stream();
9860        let mut lb = __s_b.launch_builder(&f);
9861        lb.arg(x).arg(&mut *dst).arg(&ni);
9862        unsafe {
9863            lb.launch(cfg)?;
9864        }
9865        Ok(())
9866    }
9867
9868    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
9869    pub fn row_softmax(
9870        &self,
9871        x: &mut CudaSlice<f32>,
9872        ncols: usize,
9873        nrows: usize,
9874    ) -> Result<(), Box<dyn std::error::Error>> {
9875        let f = self.func("row_softmax_f32");
9876        let nc = ncols as i32;
9877        let cfg = LaunchConfig {
9878            grid_dim: (nrows as u32, 1, 1),
9879            block_dim: (256, 1, 1),
9880            shared_mem_bytes: 0,
9881        };
9882        let __s_b = self.gpu.stream();
9883        let mut lb = __s_b.launch_builder(&f);
9884        lb.arg(&mut *x).arg(&nc);
9885        unsafe {
9886            lb.launch(cfg)?;
9887        }
9888        Ok(())
9889    }
9890
9891    pub fn rms_norm(
9892        &self,
9893        x: &CudaSlice<f32>,
9894        w: &CudaSlice<f32>,
9895        dst: &mut CudaSlice<f32>,
9896        ncols: usize,
9897        nrows: usize,
9898        eps: f32,
9899    ) -> Result<(), Box<dyn std::error::Error>> {
9900        let (nc, e) = (ncols as i32, eps);
9901        let kname = if Self::norm_ilp_on() {
9902            "rms_norm_f32_v2"
9903        } else {
9904            "rms_norm_f32"
9905        };
9906        if Self::pdl_on() && Self::pdl_wb_on() {
9907            use cudarc::driver::{DevicePtr, DevicePtrMut};
9908            let s = &self.gpu.stream();
9909            let (px, _g0) = x.device_ptr(s);
9910            let (pw, _g1) = w.device_ptr(s);
9911            let (pd, _g2) = dst.device_ptr_mut(s);
9912            let mut ps = [
9913                &px as *const _ as *mut std::ffi::c_void,
9914                &pw as *const _ as *mut _,
9915                &pd as *const _ as *mut _,
9916                &nc as *const _ as *mut _,
9917                &e as *const _ as *mut _,
9918            ];
9919            unsafe {
9920                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9921            }
9922            return Ok(());
9923        }
9924        let f = self.func(kname);
9925        let cfg = LaunchConfig {
9926            grid_dim: (nrows as u32, 1, 1),
9927            block_dim: (rms_block(), 1, 1),
9928            shared_mem_bytes: 0,
9929        };
9930        let __s_b = self.gpu.stream();
9931        let mut b = __s_b.launch_builder(&f);
9932        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9933        unsafe {
9934            b.launch(cfg)?;
9935        }
9936        Ok(())
9937    }
9938
9939    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
9940    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
9941    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
9942    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
9943    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
9944    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
9945    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
9946    pub fn rms_norm_decode(
9947        &self,
9948        x: &CudaSlice<f32>,
9949        w: &CudaSlice<f32>,
9950        dst: &mut CudaSlice<f32>,
9951        ncols: usize,
9952        nrows: usize,
9953        eps: f32,
9954    ) -> Result<(), Box<dyn std::error::Error>> {
9955        let f = self.func(if Self::norm_ilp_on() {
9956            "rms_norm_f32_v2"
9957        } else {
9958            "rms_norm_f32"
9959        });
9960        let cfg = LaunchConfig {
9961            grid_dim: (nrows as u32, 1, 1),
9962            block_dim: (1024, 1, 1),
9963            shared_mem_bytes: 0,
9964        };
9965        let (nc, e) = (ncols as i32, eps);
9966        let __s_b = self.gpu.stream();
9967        let mut b = __s_b.launch_builder(&f);
9968        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9969        unsafe {
9970            b.launch(cfg)?;
9971        }
9972        Ok(())
9973    }
9974
9975    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
9976    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
9977    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
9978    pub fn rms_norm_q8_1(
9979        &self,
9980        x: &CudaSlice<f32>,
9981        w: &CudaSlice<f32>,
9982        ncols: usize,
9983        nrows: usize,
9984        eps: f32,
9985    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9986        let nblk = ncols / 32;
9987        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9988        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9989        let (nc, e) = (ncols as i32, eps);
9990        if Self::pdl_on() {
9991            {
9992                use cudarc::driver::{DevicePtr, DevicePtrMut};
9993                let s = &self.gpu.stream();
9994                let (px, _g0) = x.device_ptr(s);
9995                let (pw, _g1) = w.device_ptr(s);
9996                let (pq, _g2) = q.device_ptr_mut(s);
9997                let (pd, _g3) = d.device_ptr_mut(s);
9998                let mut ps = [
9999                    &px as *const _ as *mut std::ffi::c_void,
10000                    &pw as *const _ as *mut _,
10001                    &pq as *const _ as *mut _,
10002                    &pd as *const _ as *mut _,
10003                    &nc as *const _ as *mut _,
10004                    &e as *const _ as *mut _,
10005                ];
10006                unsafe {
10007                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10008                }
10009            }
10010            return Ok((q, d));
10011        }
10012        let f = self.func("rms_norm_q8_1");
10013        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
10014        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
10015        let cfg = LaunchConfig {
10016            grid_dim: (nrows as u32, 1, 1),
10017            block_dim: (1024, 1, 1),
10018            shared_mem_bytes: 0,
10019        };
10020        let __s_b = self.gpu.stream();
10021        let mut b = __s_b.launch_builder(&f);
10022        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10023        unsafe {
10024            b.launch(cfg)?;
10025        }
10026        Ok((q, d))
10027    }
10028
10029    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
10030    /// PDL arm), caller-owned outputs.
10031    pub fn rms_norm_q8_1_into(
10032        &self,
10033        x: &CudaSlice<f32>,
10034        w: &CudaSlice<f32>,
10035        ncols: usize,
10036        nrows: usize,
10037        eps: f32,
10038        q: &mut CudaSlice<i8>,
10039        d: &mut CudaSlice<f32>,
10040    ) -> Result<(), Box<dyn std::error::Error>> {
10041        let nblk = ncols / 32;
10042        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10043        let (nc, e) = (ncols as i32, eps);
10044        if Self::pdl_on() {
10045            use cudarc::driver::{DevicePtr, DevicePtrMut};
10046            let s = &self.gpu.stream();
10047            let (px, _g0) = x.device_ptr(s);
10048            let (pw, _g1) = w.device_ptr(s);
10049            let (pq, _g2) = q.device_ptr_mut(s);
10050            let (pd, _g3) = d.device_ptr_mut(s);
10051            let mut ps = [
10052                &px as *const _ as *mut std::ffi::c_void,
10053                &pw as *const _ as *mut _,
10054                &pq as *const _ as *mut _,
10055                &pd as *const _ as *mut _,
10056                &nc as *const _ as *mut _,
10057                &e as *const _ as *mut _,
10058            ];
10059            unsafe {
10060                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10061            }
10062            return Ok(());
10063        }
10064        let f = self.func("rms_norm_q8_1");
10065        let cfg = LaunchConfig {
10066            grid_dim: (nrows as u32, 1, 1),
10067            block_dim: (1024, 1, 1),
10068            shared_mem_bytes: 0,
10069        };
10070        let __s_b = self.gpu.stream();
10071        let mut b = __s_b.launch_builder(&f);
10072        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10073        unsafe {
10074            b.launch(cfg)?;
10075        }
10076        Ok(())
10077    }
10078
10079    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
10080    pub fn quantize_q8_1_into(
10081        &self,
10082        x: &CudaSlice<f32>,
10083        m: usize,
10084        in_f: usize,
10085        q: &mut CudaSlice<i8>,
10086        d: &mut CudaSlice<f32>,
10087    ) -> Result<(), Box<dyn std::error::Error>> {
10088        let nblk = in_f / 32;
10089        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10090        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10091        let (inf, mi) = (in_f as i32, m as i32);
10092        if Self::pdl_on() && Self::pdl_wb_on() {
10093            use cudarc::driver::{DevicePtr, DevicePtrMut};
10094            let s = &self.gpu.stream();
10095            let (px, _g0) = x.device_ptr(s);
10096            let (pq, _g1) = q.device_ptr_mut(s);
10097            let (pd, _g2) = d.device_ptr_mut(s);
10098            let mut ps = [
10099                &px as *const _ as *mut std::ffi::c_void,
10100                &pq as *const _ as *mut _,
10101                &pd as *const _ as *mut _,
10102                &inf as *const _ as *mut _,
10103                &mi as *const _ as *mut _,
10104            ];
10105            unsafe {
10106                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10107            }
10108            return Ok(());
10109        }
10110        let f = self.func("quantize_q8_1");
10111        let __s_b = self.gpu.stream();
10112        let mut b = __s_b.launch_builder(&f);
10113        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10114        unsafe {
10115            b.launch(cfg)?;
10116        }
10117        Ok(())
10118    }
10119
10120    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
10121    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
10122    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
10123    pub fn add_rms_norm_q8_1(
10124        &self,
10125        a: &CudaSlice<f32>,
10126        b_in: &CudaSlice<f32>,
10127        w: &CudaSlice<f32>,
10128        res: &mut CudaSlice<f32>,
10129        ncols: usize,
10130        nrows: usize,
10131        eps: f32,
10132    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10133        let nblk = ncols / 32;
10134        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10135        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10136        let f = self.func("add_rms_norm_q8_1");
10137        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
10138        let cfg = LaunchConfig {
10139            grid_dim: (nrows as u32, 1, 1),
10140            block_dim: (1024, 1, 1),
10141            shared_mem_bytes: 0,
10142        };
10143        let (nc, e) = (ncols as i32, eps);
10144        let __s_bld = self.gpu.stream();
10145        let mut bld = __s_bld.launch_builder(&f);
10146        bld.arg(a)
10147            .arg(b_in)
10148            .arg(w)
10149            .arg(res)
10150            .arg(&mut q)
10151            .arg(&mut d)
10152            .arg(&nc)
10153            .arg(&e);
10154        unsafe {
10155            bld.launch(cfg)?;
10156        }
10157        Ok((q, d))
10158    }
10159
10160    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
10161    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
10162    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
10163    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
10164    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
10165    #[allow(clippy::too_many_arguments)]
10166    pub fn join_add_rms_norm_raw(
10167        &self,
10168        a0_raw: u64,
10169        a1_raw: u64,
10170        x: &CudaSlice<f32>,
10171        w: &CudaSlice<f32>,
10172        res: &mut CudaSlice<f32>,
10173        dst: &mut CudaSlice<f32>,
10174        ncols: usize,
10175        eps: f32,
10176    ) -> Result<(), Box<dyn std::error::Error>> {
10177        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10178            return Err("join_add_rms_norm geometry".into());
10179        }
10180        let f = self.func("join_add_rms_norm_f32");
10181        let cfg = LaunchConfig {
10182            grid_dim: (1, 1, 1),
10183            block_dim: (rms_block(), 1, 1),
10184            shared_mem_bytes: 0,
10185        };
10186        let (nc, e) = (ncols as i32, eps);
10187        let __s_b = self.gpu.stream();
10188        let mut b = __s_b.launch_builder(&f);
10189        b.arg(&a0_raw)
10190            .arg(&a1_raw)
10191            .arg(x)
10192            .arg(w)
10193            .arg(&mut *res)
10194            .arg(&mut *dst)
10195            .arg(&nc)
10196            .arg(&e);
10197        unsafe {
10198            b.launch(cfg)?;
10199        }
10200        Ok(())
10201    }
10202
10203    pub fn add_rms_norm(
10204        &self,
10205        a: &CudaSlice<f32>,
10206        b: &CudaSlice<f32>,
10207        w: &CudaSlice<f32>,
10208        res: &mut CudaSlice<f32>,
10209        dst: &mut CudaSlice<f32>,
10210        ncols: usize,
10211        nrows: usize,
10212        eps: f32,
10213    ) -> Result<(), Box<dyn std::error::Error>> {
10214        let (nc, e) = (ncols as i32, eps);
10215        let kname = if Self::norm_ilp_on() {
10216            "add_rms_norm_f32_v2"
10217        } else {
10218            "add_rms_norm_f32"
10219        };
10220        if Self::pdl_on() && Self::pdl_wb_on() {
10221            use cudarc::driver::{DevicePtr, DevicePtrMut};
10222            let s = &self.gpu.stream();
10223            let (pa, _g0) = a.device_ptr(s);
10224            let (pb, _g1) = b.device_ptr(s);
10225            let (pw, _g2) = w.device_ptr(s);
10226            let (pr, _g3) = res.device_ptr_mut(s);
10227            let (pd, _g4) = dst.device_ptr_mut(s);
10228            let mut ps = [
10229                &pa as *const _ as *mut std::ffi::c_void,
10230                &pb as *const _ as *mut _,
10231                &pw as *const _ as *mut _,
10232                &pr as *const _ as *mut _,
10233                &pd as *const _ as *mut _,
10234                &nc as *const _ as *mut _,
10235                &e as *const _ as *mut _,
10236            ];
10237            unsafe {
10238                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10239            }
10240            return Ok(());
10241        }
10242        let f = self.func(kname);
10243        let cfg = LaunchConfig {
10244            grid_dim: (nrows as u32, 1, 1),
10245            block_dim: (rms_block(), 1, 1),
10246            shared_mem_bytes: 0,
10247        };
10248        let __s_b2 = self.gpu.stream();
10249        let mut b2 = __s_b2.launch_builder(&f);
10250        b2.arg(a)
10251            .arg(b)
10252            .arg(w)
10253            .arg(&mut *res)
10254            .arg(&mut *dst)
10255            .arg(&nc)
10256            .arg(&e);
10257        unsafe {
10258            b2.launch(cfg)?;
10259        }
10260        Ok(())
10261    }
10262
10263    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
10264    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
10265    #[allow(clippy::too_many_arguments)]
10266    pub fn rms_pre_add_rms_norm(
10267        &self,
10268        a: &CudaSlice<f32>,
10269        wa: &CudaSlice<f32>,
10270        b: &CudaSlice<f32>,
10271        w: &CudaSlice<f32>,
10272        res: &mut CudaSlice<f32>,
10273        dst: &mut CudaSlice<f32>,
10274        ncols: usize,
10275        nrows: usize,
10276        eps: f32,
10277    ) -> Result<(), Box<dyn std::error::Error>> {
10278        let f = self.func("rms_pre_add_rms_norm_f32");
10279        let cfg = LaunchConfig {
10280            grid_dim: (nrows as u32, 1, 1),
10281            block_dim: (rms_block(), 1, 1),
10282            shared_mem_bytes: 0,
10283        };
10284        let (nc, e) = (ncols as i32, eps);
10285        let __s_b2 = self.gpu.stream();
10286        let mut b2 = __s_b2.launch_builder(&f);
10287        b2.arg(a)
10288            .arg(wa)
10289            .arg(b)
10290            .arg(w)
10291            .arg(&mut *res)
10292            .arg(&mut *dst)
10293            .arg(&nc)
10294            .arg(&e);
10295        unsafe {
10296            b2.launch(cfg)?;
10297        }
10298        Ok(())
10299    }
10300
10301    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
10302    #[allow(clippy::too_many_arguments)]
10303    pub fn rms_pre_add_rms_norm_q8z(
10304        &self,
10305        a: &CudaSlice<f32>,
10306        wa: &CudaSlice<f32>,
10307        b: &CudaSlice<f32>,
10308        w: &CudaSlice<f32>,
10309        res: &mut CudaSlice<f32>,
10310        dst: &mut CudaSlice<f32>,
10311        ncols: usize,
10312        nrows: usize,
10313        eps: f32,
10314    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10315        debug_assert!(ncols % 128 == 0);
10316        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10317        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10318        let (nc, e) = (ncols as i32, eps);
10319        if Self::pdl_on() {
10320            {
10321                use cudarc::driver::{DevicePtr, DevicePtrMut};
10322                let s = &self.gpu.stream();
10323                let (pa, _g0) = a.device_ptr(s);
10324                let (pwa, _g1) = wa.device_ptr(s);
10325                let (pb, _g2) = b.device_ptr(s);
10326                let (pw, _g3) = w.device_ptr(s);
10327                let (pr, _g4) = res.device_ptr_mut(s);
10328                let (pdst, _g5) = dst.device_ptr_mut(s);
10329                let (pq, _g6) = out_q.device_ptr_mut(s);
10330                let (pd, _g7) = out_d.device_ptr_mut(s);
10331                let mut ps = [
10332                    &pa as *const _ as *mut std::ffi::c_void,
10333                    &pwa as *const _ as *mut _,
10334                    &pb as *const _ as *mut _,
10335                    &pw as *const _ as *mut _,
10336                    &pr as *const _ as *mut _,
10337                    &pdst as *const _ as *mut _,
10338                    &pq as *const _ as *mut _,
10339                    &pd as *const _ as *mut _,
10340                    &nc as *const _ as *mut _,
10341                    &e as *const _ as *mut _,
10342                ];
10343                unsafe {
10344                    self.launch_pdl(
10345                        "rms_pre_add_rms_norm_q8z_f32",
10346                        (nrows as u32, 1, 1),
10347                        (rms_block(), 1, 1),
10348                        &mut ps,
10349                    )?;
10350                }
10351            }
10352            return Ok((out_q, out_d));
10353        }
10354        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10355        let cfg = LaunchConfig {
10356            grid_dim: (nrows as u32, 1, 1),
10357            block_dim: (rms_block(), 1, 1),
10358            shared_mem_bytes: 0,
10359        };
10360        let __s_b2 = self.gpu.stream();
10361        let mut b2 = __s_b2.launch_builder(&f);
10362        b2.arg(a)
10363            .arg(wa)
10364            .arg(b)
10365            .arg(w)
10366            .arg(&mut *res)
10367            .arg(&mut *dst)
10368            .arg(&mut out_q)
10369            .arg(&mut out_d)
10370            .arg(&nc)
10371            .arg(&e);
10372        unsafe {
10373            b2.launch(cfg)?;
10374        }
10375        Ok((out_q, out_d))
10376    }
10377
10378    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
10379    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
10380    /// body must stay attribute-free (the fused2_into precedent).
10381    #[allow(clippy::too_many_arguments)]
10382    pub fn rms_pre_add_rms_norm_q8z_into(
10383        &self,
10384        a: &CudaSlice<f32>,
10385        wa: &CudaSlice<f32>,
10386        b: &CudaSlice<f32>,
10387        w: &CudaSlice<f32>,
10388        res: &mut CudaSlice<f32>,
10389        dst: &mut CudaSlice<f32>,
10390        ncols: usize,
10391        nrows: usize,
10392        eps: f32,
10393        out_q: &mut CudaSlice<i8>,
10394        out_d: &mut CudaSlice<f32>,
10395    ) -> Result<(), Box<dyn std::error::Error>> {
10396        debug_assert!(ncols % 128 == 0);
10397        let (nc, e) = (ncols as i32, eps);
10398        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10399        let cfg = LaunchConfig {
10400            grid_dim: (nrows as u32, 1, 1),
10401            block_dim: (rms_block(), 1, 1),
10402            shared_mem_bytes: 0,
10403        };
10404        let __s_b = self.gpu.stream();
10405        let mut b2 = __s_b.launch_builder(&f);
10406        b2.arg(a)
10407            .arg(wa)
10408            .arg(b)
10409            .arg(w)
10410            .arg(&mut *res)
10411            .arg(&mut *dst)
10412            .arg(&mut *out_q)
10413            .arg(&mut *out_d)
10414            .arg(&nc)
10415            .arg(&e);
10416        unsafe {
10417            b2.launch(cfg)?;
10418        }
10419        Ok(())
10420    }
10421
10422    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
10423    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
10424    #[allow(clippy::too_many_arguments)]
10425    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10426        &self,
10427        a: &CudaSlice<f32>,
10428        wa: &CudaSlice<f32>,
10429        b_in: &CudaSlice<f32>,
10430        c: f32,
10431        w: &CudaSlice<f32>,
10432        res: &mut CudaSlice<f32>,
10433        ncols: usize,
10434        nrows: usize,
10435        eps: f32,
10436        out_q: &mut CudaSlice<i8>,
10437        out_d: &mut CudaSlice<f32>,
10438    ) -> Result<(), Box<dyn std::error::Error>> {
10439        debug_assert!(ncols % 128 == 0);
10440        let (nc, e2) = (ncols as i32, eps);
10441        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10442        let cfg = LaunchConfig {
10443            grid_dim: (nrows as u32, 1, 1),
10444            block_dim: (rms_block(), 1, 1),
10445            shared_mem_bytes: 0,
10446        };
10447        let __s_b = self.gpu.stream();
10448        let mut b2 = __s_b.launch_builder(&f);
10449        b2.arg(a)
10450            .arg(wa)
10451            .arg(b_in)
10452            .arg(&c)
10453            .arg(w)
10454            .arg(&mut *res)
10455            .arg(&mut *out_q)
10456            .arg(&mut *out_d)
10457            .arg(&nc)
10458            .arg(&e2);
10459        unsafe {
10460            b2.launch(cfg)?;
10461        }
10462        Ok(())
10463    }
10464
10465    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
10466    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
10467    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
10468    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
10469    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
10470    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
10471    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
10472    pub fn g4_pnfold_on() -> bool {
10473        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10474        *ON.get_or_init(|| {
10475            std::env::var("MEMRA_G4_PNFOLD")
10476                .map(|v| v != "0")
10477                .unwrap_or(true)
10478        })
10479    }
10480
10481    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
10482    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
10483    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
10484    pub fn build_q4_out_concat3(
10485        &self,
10486        w0: &crate::model::GpuTensor,
10487        w1: &crate::model::GpuTensor,
10488        w2: &crate::model::GpuTensor,
10489    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10490        use crate::model::GpuTensor;
10491        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10492            match w {
10493                GpuTensor::Quant {
10494                    qtype,
10495                    row_bytes,
10496                    rp,
10497                    ..
10498                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10499                _ => None,
10500            }
10501        };
10502        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10503        else {
10504            return Ok(None);
10505        };
10506        if rb0 != rb1
10507            || rb0 != rb2
10508            || w0.in_features() != w1.in_features()
10509            || w0.in_features() != w2.in_features()
10510        {
10511            return Ok(None);
10512        }
10513        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10514            match w {
10515                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10516                _ => unreachable!(),
10517            }
10518        }
10519        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10520        let total = rb0 * (o0 + o1 + o2);
10521        let mut cat = self.alloc_u8(total)?;
10522        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10523        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10524        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10525        Ok(Some(GpuTensor::Quant {
10526            bytes: cat,
10527            qtype: QT_Q4_0,
10528            row_bytes: rb0,
10529            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10530            scale: 1.0,
10531            rp: false,
10532            #[cfg(memra_cutlass)]
10533            cutlass: None,
10534            fp8: None,
10535            blk: None,
10536            rp4: None,
10537            f16: None,
10538        }))
10539    }
10540
10541    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
10542    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
10543    ///
10544    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
10545    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
10546    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
10547    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
10548    ///
10549    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
10550    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
10551    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
10552    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
10553    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
10554    ///
10555    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
10556    /// width. A future partial-rotary caller fails at its first launch with the geometry named
10557    /// instead of serving quietly wrong logits.
10558    fn full_width_rope_only(
10559        kernel: &str,
10560        n_rot: usize,
10561        head_dim: usize,
10562    ) -> Result<(), Box<dyn std::error::Error>> {
10563        if n_rot == head_dim {
10564            return Ok(());
10565        }
10566        Err(format!(
10567            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10568             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10569             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10570             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10571             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10572        )
10573        .into())
10574    }
10575
10576    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
10577    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10578    /// ([`Engine::full_width_rope_only`]).
10579    #[allow(clippy::too_many_arguments)]
10580    pub fn rms_norm_qkv_rope_cat(
10581        &self,
10582        qkv: &CudaSlice<f32>,
10583        wq: &CudaSlice<f32>,
10584        wk: &CudaSlice<f32>,
10585        wv: &CudaSlice<f32>,
10586        q: &mut CudaSlice<f32>,
10587        k: &mut CudaSlice<f32>,
10588        v: &mut CudaSlice<f32>,
10589        head_dim: usize,
10590        n_rot: usize,
10591        rq: usize,
10592        rk: usize,
10593        pos: &CudaSlice<i32>,
10594        nh_q: usize,
10595        nh_k: usize,
10596        base: f32,
10597        freq_scale: f32,
10598        ff: Option<&CudaSlice<f32>>,
10599        eps: f32,
10600    ) -> Result<(), Box<dyn std::error::Error>> {
10601        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10602        let rows = rq + rk + rk;
10603        let theta_scale = base.powf(-2.0 / head_dim as f32);
10604        let (nc, rqi, rki, nhq, nhk) = (
10605            head_dim as i32,
10606            rq as i32,
10607            rk as i32,
10608            nh_q as i32,
10609            nh_k as i32,
10610        );
10611        if Self::pdl_on() {
10612            use cudarc::driver::{DevicePtr, DevicePtrMut};
10613            let s = &self.gpu.stream();
10614            let (pqkv, _g0) = qkv.device_ptr(s);
10615            let (pwq, _g1) = wq.device_ptr(s);
10616            let (pwk, _g2) = wk.device_ptr(s);
10617            let (pwv, _g3) = wv.device_ptr(s);
10618            let (pq, _g4) = q.device_ptr_mut(s);
10619            let (pk, _g5) = k.device_ptr_mut(s);
10620            let (pv, _g6) = v.device_ptr_mut(s);
10621            let (ppos, _g7) = pos.device_ptr(s);
10622            let (pff, _g8) = match ff {
10623                Some(t) => {
10624                    let (p, g) = t.device_ptr(s);
10625                    (p, Some(g))
10626                }
10627                None => (0, None),
10628            };
10629            let mut ps = [
10630                &pqkv as *const _ as *mut std::ffi::c_void,
10631                &pwq as *const _ as *mut _,
10632                &pwk as *const _ as *mut _,
10633                &pwv as *const _ as *mut _,
10634                &pq as *const _ as *mut _,
10635                &pk as *const _ as *mut _,
10636                &pv as *const _ as *mut _,
10637                &nc as *const _ as *mut _,
10638                &rqi as *const _ as *mut _,
10639                &rki as *const _ as *mut _,
10640                &ppos as *const _ as *mut _,
10641                &nhq as *const _ as *mut _,
10642                &nhk as *const _ as *mut _,
10643                &theta_scale as *const _ as *mut _,
10644                &freq_scale as *const _ as *mut _,
10645                &pff as *const _ as *mut _,
10646                &eps as *const _ as *mut _,
10647            ];
10648            unsafe {
10649                self.launch_pdl(
10650                    "rms_norm_qkv_rope_cat_f32",
10651                    (rows as u32, 1, 1),
10652                    (rms_block(), 1, 1),
10653                    &mut ps,
10654                )?;
10655            }
10656            return Ok(());
10657        }
10658        let f = self.func("rms_norm_qkv_rope_cat_f32");
10659        let cfg = LaunchConfig {
10660            grid_dim: (rows as u32, 1, 1),
10661            block_dim: (rms_block(), 1, 1),
10662            shared_mem_bytes: 0,
10663        };
10664        let __s_b = self.gpu.stream();
10665        let mut b = __s_b.launch_builder(&f);
10666        match ff {
10667            Some(t) => {
10668                b.arg(qkv)
10669                    .arg(wq)
10670                    .arg(wk)
10671                    .arg(wv)
10672                    .arg(&mut *q)
10673                    .arg(&mut *k)
10674                    .arg(&mut *v)
10675                    .arg(&nc)
10676                    .arg(&rqi)
10677                    .arg(&rki)
10678                    .arg(pos)
10679                    .arg(&nhq)
10680                    .arg(&nhk)
10681                    .arg(&theta_scale)
10682                    .arg(&freq_scale)
10683                    .arg(t)
10684                    .arg(&eps);
10685                unsafe {
10686                    b.launch(cfg)?;
10687                }
10688            }
10689            None => {
10690                let null: u64 = 0;
10691                b.arg(qkv)
10692                    .arg(wq)
10693                    .arg(wk)
10694                    .arg(wv)
10695                    .arg(&mut *q)
10696                    .arg(&mut *k)
10697                    .arg(&mut *v)
10698                    .arg(&nc)
10699                    .arg(&rqi)
10700                    .arg(&rki)
10701                    .arg(pos)
10702                    .arg(&nhq)
10703                    .arg(&nhk)
10704                    .arg(&theta_scale)
10705                    .arg(&freq_scale)
10706                    .arg(&null)
10707                    .arg(&eps);
10708                unsafe {
10709                    b.launch(cfg)?;
10710                }
10711            }
10712        }
10713        Ok(())
10714    }
10715
10716    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10717    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10718    /// ([`Engine::full_width_rope_only`]).
10719    #[allow(clippy::too_many_arguments)]
10720    pub fn rms_norm_qkv_rope(
10721        &self,
10722        q0: &CudaSlice<f32>,
10723        k0: &CudaSlice<f32>,
10724        v0: &CudaSlice<f32>,
10725        wq: &CudaSlice<f32>,
10726        wk: &CudaSlice<f32>,
10727        wv: &CudaSlice<f32>,
10728        q: &mut CudaSlice<f32>,
10729        k: &mut CudaSlice<f32>,
10730        v: &mut CudaSlice<f32>,
10731        head_dim: usize,
10732        n_rot: usize,
10733        rq: usize,
10734        rk: usize,
10735        pos: &CudaSlice<i32>,
10736        nh_q: usize,
10737        nh_k: usize,
10738        base: f32,
10739        freq_scale: f32,
10740        ff: Option<&CudaSlice<f32>>,
10741        eps: f32,
10742    ) -> Result<(), Box<dyn std::error::Error>> {
10743        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10744        let f = self.func("rms_norm_qkv_rope_f32");
10745        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10746        let cfg = LaunchConfig {
10747            grid_dim: (rows as u32, 1, 1),
10748            block_dim: (rms_block(), 1, 1),
10749            shared_mem_bytes: 0,
10750        };
10751        let theta_scale = base.powf(-2.0 / head_dim as f32);
10752        let (nc, rqi, rki, nhq, nhk) = (
10753            head_dim as i32,
10754            rq as i32,
10755            rk as i32,
10756            nh_q as i32,
10757            nh_k as i32,
10758        );
10759        let __s_b = self.gpu.stream();
10760        let mut b = __s_b.launch_builder(&f);
10761        match ff {
10762            Some(t) => {
10763                b.arg(q0)
10764                    .arg(k0)
10765                    .arg(v0)
10766                    .arg(wq)
10767                    .arg(wk)
10768                    .arg(wv)
10769                    .arg(&mut *q)
10770                    .arg(&mut *k)
10771                    .arg(&mut *v)
10772                    .arg(&nc)
10773                    .arg(&rqi)
10774                    .arg(&rki)
10775                    .arg(pos)
10776                    .arg(&nhq)
10777                    .arg(&nhk)
10778                    .arg(&theta_scale)
10779                    .arg(&freq_scale)
10780                    .arg(t)
10781                    .arg(&eps);
10782                unsafe {
10783                    b.launch(cfg)?;
10784                }
10785            }
10786            None => {
10787                let null: u64 = 0;
10788                b.arg(q0)
10789                    .arg(k0)
10790                    .arg(v0)
10791                    .arg(wq)
10792                    .arg(wk)
10793                    .arg(wv)
10794                    .arg(&mut *q)
10795                    .arg(&mut *k)
10796                    .arg(&mut *v)
10797                    .arg(&nc)
10798                    .arg(&rqi)
10799                    .arg(&rki)
10800                    .arg(pos)
10801                    .arg(&nhq)
10802                    .arg(&nhk)
10803                    .arg(&theta_scale)
10804                    .arg(&freq_scale)
10805                    .arg(&null)
10806                    .arg(&eps);
10807                unsafe {
10808                    b.launch(cfg)?;
10809                }
10810            }
10811        }
10812        Ok(())
10813    }
10814
10815    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
10816    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
10817    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
10818    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10819    /// ([`Engine::full_width_rope_only`]).
10820    #[allow(clippy::too_many_arguments)]
10821    pub fn rms_norm_qkv_rope_append_dc(
10822        &self,
10823        q0: &CudaSlice<f32>,
10824        k0: &CudaSlice<f32>,
10825        v0: &CudaSlice<f32>,
10826        wq: &CudaSlice<f32>,
10827        wk: &CudaSlice<f32>,
10828        wv: &CudaSlice<f32>,
10829        q: &mut CudaSlice<f32>,
10830        k: &mut CudaSlice<f32>,
10831        v: &mut CudaSlice<f32>,
10832        head_dim: usize,
10833        n_rot: usize,
10834        rq: usize,
10835        rk: usize,
10836        pos: &CudaSlice<i32>,
10837        nh_q: usize,
10838        nh_k: usize,
10839        base: f32,
10840        freq_scale: f32,
10841        ff: Option<&CudaSlice<f32>>,
10842        eps: f32,
10843        kc: &mut CudaSlice<u8>,
10844        vc: &mut CudaSlice<u8>,
10845        t_dev: &CudaSlice<i32>,
10846        k_tok_bytes: usize,
10847        v_tok_bytes: usize,
10848        g: bool,
10849    ) -> Result<(), Box<dyn std::error::Error>> {
10850        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10851        let rows = rq + rk + rk;
10852        let theta_scale = base.powf(-2.0 / head_dim as f32);
10853        let (nc, rqi, rki, nhq, nhk) = (
10854            head_dim as i32,
10855            rq as i32,
10856            rk as i32,
10857            nh_q as i32,
10858            nh_k as i32,
10859        );
10860        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10861        if Self::pdl_on() && Self::pdl_wb_on() {
10862            use cudarc::driver::{DevicePtr, DevicePtrMut};
10863            let s = &self.gpu.stream();
10864            let (p0, _a0) = q0.device_ptr(s);
10865            let (p1, _a1) = k0.device_ptr(s);
10866            let (p2, _a2) = v0.device_ptr(s);
10867            let (pwq, _a3) = wq.device_ptr(s);
10868            let (pwk, _a4) = wk.device_ptr(s);
10869            let (pwv, _a5) = wv.device_ptr(s);
10870            let (pq, _a6) = q.device_ptr_mut(s);
10871            let (pk, _a7) = k.device_ptr_mut(s);
10872            let (pv, _a8) = v.device_ptr_mut(s);
10873            let (pp, _a9) = pos.device_ptr(s);
10874            let pff: u64 = match ff {
10875                Some(t) => {
10876                    let (p, _gg) = t.device_ptr(s);
10877                    p as u64
10878                }
10879                None => 0,
10880            };
10881            let (pkc, _a10) = kc.device_ptr_mut(s);
10882            let (pvc, _a11) = vc.device_ptr_mut(s);
10883            let (pt, _a12) = t_dev.device_ptr(s);
10884            let mut ps = [
10885                &p0 as *const _ as *mut std::ffi::c_void,
10886                &p1 as *const _ as *mut _,
10887                &p2 as *const _ as *mut _,
10888                &pwq as *const _ as *mut _,
10889                &pwk as *const _ as *mut _,
10890                &pwv as *const _ as *mut _,
10891                &pq as *const _ as *mut _,
10892                &pk as *const _ as *mut _,
10893                &pv as *const _ as *mut _,
10894                &nc as *const _ as *mut _,
10895                &rqi as *const _ as *mut _,
10896                &rki as *const _ as *mut _,
10897                &pp as *const _ as *mut _,
10898                &nhq as *const _ as *mut _,
10899                &nhk as *const _ as *mut _,
10900                &theta_scale as *const _ as *mut _,
10901                &freq_scale as *const _ as *mut _,
10902                &pff as *const _ as *mut _,
10903                &eps as *const _ as *mut _,
10904                &pkc as *const _ as *mut _,
10905                &pvc as *const _ as *mut _,
10906                &pt as *const _ as *mut _,
10907                &ktb as *const _ as *mut _,
10908                &vtb as *const _ as *mut _,
10909            ];
10910            unsafe {
10911                self.launch_pdl_flash(
10912                    g,
10913                    "rms_norm_qkv_rope_append_dc_f32",
10914                    (rows as u32, 1, 1),
10915                    (rms_block(), 1, 1),
10916                    0,
10917                    &mut ps,
10918                )?;
10919            }
10920            return Ok(());
10921        }
10922        let f = if g {
10923            self.func_g("rms_norm_qkv_rope_append_dc_f32")
10924        } else {
10925            self.func("rms_norm_qkv_rope_append_dc_f32")
10926        };
10927        let cfg = LaunchConfig {
10928            grid_dim: (rows as u32, 1, 1),
10929            block_dim: (rms_block(), 1, 1),
10930            shared_mem_bytes: 0,
10931        };
10932        let __s_b = self.gpu.stream();
10933        let mut b = __s_b.launch_builder(&f);
10934        match ff {
10935            Some(t) => {
10936                b.arg(q0)
10937                    .arg(k0)
10938                    .arg(v0)
10939                    .arg(wq)
10940                    .arg(wk)
10941                    .arg(wv)
10942                    .arg(&mut *q)
10943                    .arg(&mut *k)
10944                    .arg(&mut *v)
10945                    .arg(&nc)
10946                    .arg(&rqi)
10947                    .arg(&rki)
10948                    .arg(pos)
10949                    .arg(&nhq)
10950                    .arg(&nhk)
10951                    .arg(&theta_scale)
10952                    .arg(&freq_scale)
10953                    .arg(t)
10954                    .arg(&eps)
10955                    .arg(&mut *kc)
10956                    .arg(&mut *vc)
10957                    .arg(t_dev)
10958                    .arg(&ktb)
10959                    .arg(&vtb);
10960                unsafe {
10961                    b.launch(cfg)?;
10962                }
10963            }
10964            None => {
10965                let null: u64 = 0;
10966                b.arg(q0)
10967                    .arg(k0)
10968                    .arg(v0)
10969                    .arg(wq)
10970                    .arg(wk)
10971                    .arg(wv)
10972                    .arg(&mut *q)
10973                    .arg(&mut *k)
10974                    .arg(&mut *v)
10975                    .arg(&nc)
10976                    .arg(&rqi)
10977                    .arg(&rki)
10978                    .arg(pos)
10979                    .arg(&nhq)
10980                    .arg(&nhk)
10981                    .arg(&theta_scale)
10982                    .arg(&freq_scale)
10983                    .arg(&null)
10984                    .arg(&eps)
10985                    .arg(&mut *kc)
10986                    .arg(&mut *vc)
10987                    .arg(t_dev)
10988                    .arg(&ktb)
10989                    .arg(&vtb);
10990                unsafe {
10991                    b.launch(cfg)?;
10992                }
10993            }
10994        }
10995        Ok(())
10996    }
10997
10998    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
10999    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
11000    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
11001    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
11002    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
11003    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
11004    /// `head_dim` ([`Engine::full_width_rope_only`]).
11005    #[allow(clippy::too_many_arguments)]
11006    pub fn rms_norm_qkv_rope_append(
11007        &self,
11008        q0: &CudaSlice<f32>,
11009        k0: &CudaSlice<f32>,
11010        v0: &CudaSlice<f32>,
11011        wq: &CudaSlice<f32>,
11012        wk: &CudaSlice<f32>,
11013        wv: &CudaSlice<f32>,
11014        q: &mut CudaSlice<f32>,
11015        k: &mut CudaSlice<f32>,
11016        v: &mut CudaSlice<f32>,
11017        head_dim: usize,
11018        n_rot: usize,
11019        rq: usize,
11020        rk: usize,
11021        pos: &CudaSlice<i32>,
11022        nh_q: usize,
11023        nh_k: usize,
11024        base: f32,
11025        freq_scale: f32,
11026        ff: Option<&CudaSlice<f32>>,
11027        eps: f32,
11028        kc: &mut CudaSlice<u8>,
11029        vc: &mut CudaSlice<u8>,
11030        t: usize,
11031        k_tok_bytes: usize,
11032        v_tok_bytes: usize,
11033        g: bool,
11034    ) -> Result<(), Box<dyn std::error::Error>> {
11035        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11036        let rows = rq + rk + rk;
11037        let theta_scale = base.powf(-2.0 / head_dim as f32);
11038        let (nc, rqi, rki, nhq, nhk) = (
11039            head_dim as i32,
11040            rq as i32,
11041            rk as i32,
11042            nh_q as i32,
11043            nh_k as i32,
11044        );
11045        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11046        let ti = t as i32;
11047        if Self::pdl_on() && Self::pdl_wb_on() {
11048            use cudarc::driver::{DevicePtr, DevicePtrMut};
11049            let s = &self.gpu.stream();
11050            let (p0, _a0) = q0.device_ptr(s);
11051            let (p1, _a1) = k0.device_ptr(s);
11052            let (p2, _a2) = v0.device_ptr(s);
11053            let (pwq, _a3) = wq.device_ptr(s);
11054            let (pwk, _a4) = wk.device_ptr(s);
11055            let (pwv, _a5) = wv.device_ptr(s);
11056            let (pq, _a6) = q.device_ptr_mut(s);
11057            let (pk, _a7) = k.device_ptr_mut(s);
11058            let (pv, _a8) = v.device_ptr_mut(s);
11059            let (pp, _a9) = pos.device_ptr(s);
11060            let pff: u64 = match ff {
11061                Some(t) => {
11062                    let (p, _gg) = t.device_ptr(s);
11063                    p as u64
11064                }
11065                None => 0,
11066            };
11067            let (pkc, _a10) = kc.device_ptr_mut(s);
11068            let (pvc, _a11) = vc.device_ptr_mut(s);
11069            let mut ps = [
11070                &p0 as *const _ as *mut std::ffi::c_void,
11071                &p1 as *const _ as *mut _,
11072                &p2 as *const _ as *mut _,
11073                &pwq as *const _ as *mut _,
11074                &pwk as *const _ as *mut _,
11075                &pwv as *const _ as *mut _,
11076                &pq as *const _ as *mut _,
11077                &pk as *const _ as *mut _,
11078                &pv as *const _ as *mut _,
11079                &nc as *const _ as *mut _,
11080                &rqi as *const _ as *mut _,
11081                &rki as *const _ as *mut _,
11082                &pp as *const _ as *mut _,
11083                &nhq as *const _ as *mut _,
11084                &nhk as *const _ as *mut _,
11085                &theta_scale as *const _ as *mut _,
11086                &freq_scale as *const _ as *mut _,
11087                &pff as *const _ as *mut _,
11088                &eps as *const _ as *mut _,
11089                &pkc as *const _ as *mut _,
11090                &pvc as *const _ as *mut _,
11091                &ti as *const _ as *mut _,
11092                &ktb as *const _ as *mut _,
11093                &vtb as *const _ as *mut _,
11094            ];
11095            unsafe {
11096                self.launch_pdl_flash(
11097                    g,
11098                    "rms_norm_qkv_rope_append_f32",
11099                    (rows as u32, 1, 1),
11100                    (rms_block(), 1, 1),
11101                    0,
11102                    &mut ps,
11103                )?;
11104            }
11105            return Ok(());
11106        }
11107        let f = if g {
11108            self.func_g("rms_norm_qkv_rope_append_f32")
11109        } else {
11110            self.func("rms_norm_qkv_rope_append_f32")
11111        };
11112        let cfg = LaunchConfig {
11113            grid_dim: (rows as u32, 1, 1),
11114            block_dim: (rms_block(), 1, 1),
11115            shared_mem_bytes: 0,
11116        };
11117        let __s_b = self.gpu.stream();
11118        let mut b = __s_b.launch_builder(&f);
11119        let null: u64 = 0;
11120        b.arg(q0)
11121            .arg(k0)
11122            .arg(v0)
11123            .arg(wq)
11124            .arg(wk)
11125            .arg(wv)
11126            .arg(&mut *q)
11127            .arg(&mut *k)
11128            .arg(&mut *v)
11129            .arg(&nc)
11130            .arg(&rqi)
11131            .arg(&rki)
11132            .arg(pos)
11133            .arg(&nhq)
11134            .arg(&nhk)
11135            .arg(&theta_scale)
11136            .arg(&freq_scale);
11137        match ff {
11138            Some(t) => {
11139                b.arg(t);
11140            }
11141            None => {
11142                b.arg(&null);
11143            }
11144        }
11145        b.arg(&eps)
11146            .arg(&mut *kc)
11147            .arg(&mut *vc)
11148            .arg(&ti)
11149            .arg(&ktb)
11150            .arg(&vtb);
11151        unsafe {
11152            b.launch(cfg)?;
11153        }
11154        Ok(())
11155    }
11156
11157    pub fn add_q8_1(
11158        &self,
11159        a: &CudaSlice<f32>,
11160        b: &CudaSlice<f32>,
11161        res: &mut CudaSlice<f32>,
11162        ncols: usize,
11163        nrows: usize,
11164    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11165        debug_assert!(ncols % 128 == 0);
11166        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11167        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11168        let f = self.func("add_q8_1_f32");
11169        let cfg = LaunchConfig {
11170            grid_dim: (nrows as u32, 1, 1),
11171            block_dim: (rms_block(), 1, 1),
11172            shared_mem_bytes: 0,
11173        };
11174        let nc = ncols as i32;
11175        let __s_b2 = self.gpu.stream();
11176        let mut b2 = __s_b2.launch_builder(&f);
11177        b2.arg(a)
11178            .arg(b)
11179            .arg(&mut *res)
11180            .arg(&mut out_q)
11181            .arg(&mut out_d)
11182            .arg(&nc);
11183        unsafe {
11184            b2.launch(cfg)?;
11185        }
11186        Ok((out_q, out_d))
11187    }
11188
11189    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
11190    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
11191    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
11192    pub fn rms_pre_add_q8_1(
11193        &self,
11194        a: &CudaSlice<f32>,
11195        wa: &CudaSlice<f32>,
11196        b: &CudaSlice<f32>,
11197        res: &mut CudaSlice<f32>,
11198        ncols: usize,
11199        nrows: usize,
11200        eps: f32,
11201    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11202        debug_assert!(ncols % 128 == 0);
11203        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11204        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11205        let f = self.func("rms_pre_add_q8_1_f32");
11206        let cfg = LaunchConfig {
11207            grid_dim: (nrows as u32, 1, 1),
11208            block_dim: (rms_block(), 1, 1),
11209            shared_mem_bytes: 0,
11210        };
11211        let (nc, ep) = (ncols as i32, eps);
11212        let __s_b2 = self.gpu.stream();
11213        let mut b2 = __s_b2.launch_builder(&f);
11214        b2.arg(a)
11215            .arg(wa)
11216            .arg(b)
11217            .arg(&mut *res)
11218            .arg(&mut out_q)
11219            .arg(&mut out_d)
11220            .arg(&nc)
11221            .arg(&ep);
11222        unsafe {
11223            b2.launch(cfg)?;
11224        }
11225        Ok((out_q, out_d))
11226    }
11227
11228    /// L2 norm per row (head_dim), no weight.
11229    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
11230    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
11231    pub fn l2_v2_on(ncols: usize) -> bool {
11232        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11233    }
11234
11235    pub fn l2_norm_pp(
11236        &self,
11237        x: &CudaSlice<f32>,
11238        dst: &mut CudaSlice<f32>,
11239        dst16: Option<&mut CudaSlice<u8>>,
11240        ncols: usize,
11241        nrows: usize,
11242        eps: f32,
11243    ) -> Result<(), Box<dyn std::error::Error>> {
11244        if Self::l2_v2_on(ncols) {
11245            let f = self.func("l2_norm_pp_v2_f32");
11246            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
11247            let cfg = LaunchConfig {
11248                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11249                block_dim: (256, 1, 1),
11250                shared_mem_bytes: 0,
11251            };
11252            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11253            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
11254            let d16: u64 = match dst16 {
11255                Some(d) => self.addr_u8(d),
11256                None => 0,
11257            };
11258            let __s_b = self.gpu.stream();
11259            let mut b = __s_b.launch_builder(&f);
11260            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11261            unsafe {
11262                b.launch(cfg)?;
11263            }
11264            return Ok(());
11265        }
11266        self.l2_norm(x, dst, ncols, nrows, eps)
11267    }
11268
11269    pub fn l2_norm(
11270        &self,
11271        x: &CudaSlice<f32>,
11272        dst: &mut CudaSlice<f32>,
11273        ncols: usize,
11274        nrows: usize,
11275        eps: f32,
11276    ) -> Result<(), Box<dyn std::error::Error>> {
11277        let f = self.func("l2_norm_f32");
11278        let cfg = LaunchConfig {
11279            grid_dim: (nrows as u32, 1, 1),
11280            block_dim: (256, 1, 1),
11281            shared_mem_bytes: 0,
11282        };
11283        let (nc, e) = (ncols as i32, eps);
11284        let __s_b = self.gpu.stream();
11285        let mut b = __s_b.launch_builder(&f);
11286        b.arg(x).arg(dst).arg(&nc).arg(&e);
11287        unsafe {
11288            b.launch(cfg)?;
11289        }
11290        Ok(())
11291    }
11292
11293    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
11294    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
11295    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
11296    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
11297    /// propagate through gdn_scan and flip argmax on marginal logits.
11298    pub fn l2_norm_decode(
11299        &self,
11300        x: &CudaSlice<f32>,
11301        dst: &mut CudaSlice<f32>,
11302        ncols: usize,
11303        nrows: usize,
11304        eps: f32,
11305    ) -> Result<(), Box<dyn std::error::Error>> {
11306        let f = self.func("l2_norm_f32");
11307        let cfg = LaunchConfig {
11308            grid_dim: (nrows as u32, 1, 1),
11309            block_dim: (32, 1, 1),
11310            shared_mem_bytes: 0,
11311        };
11312        let (nc, e) = (ncols as i32, eps);
11313        let __s_b = self.gpu.stream();
11314        let mut b = __s_b.launch_builder(&f);
11315        b.arg(x).arg(dst).arg(&nc).arg(&e);
11316        unsafe {
11317            b.launch(cfg)?;
11318        }
11319        Ok(())
11320    }
11321
11322    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
11323    pub fn rope_neox(
11324        &self,
11325        x: &mut CudaSlice<f32>,
11326        pos: &CudaSlice<i32>,
11327        head_dim: usize,
11328        n_dims: usize,
11329        n_heads: usize,
11330        n_tokens: usize,
11331        freq_base: f32,
11332        freq_scale: f32,
11333    ) -> Result<(), Box<dyn std::error::Error>> {
11334        let f = self.func("rope_neox_f32");
11335        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11336        let grid = (n_heads * n_tokens) as u32;
11337        let cfg = LaunchConfig {
11338            grid_dim: (grid, 1, 1),
11339            block_dim: ((head_dim / 2) as u32, 1, 1),
11340            shared_mem_bytes: 0,
11341        };
11342        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11343        let __s_b = self.gpu.stream();
11344        let mut b = __s_b.launch_builder(&f);
11345        b.arg(x)
11346            .arg(pos)
11347            .arg(&hd)
11348            .arg(&nd)
11349            .arg(&nh)
11350            .arg(&theta_scale)
11351            .arg(&freq_scale);
11352        unsafe {
11353            b.launch(cfg)?;
11354        }
11355        Ok(())
11356    }
11357
11358    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
11359    pub fn rope_neox_ff(
11360        &self,
11361        x: &mut CudaSlice<f32>,
11362        pos: &CudaSlice<i32>,
11363        head_dim: usize,
11364        n_dims: usize,
11365        n_heads: usize,
11366        n_tokens: usize,
11367        freq_base: f32,
11368        freq_scale: f32,
11369        ff: &CudaSlice<f32>,
11370    ) -> Result<(), Box<dyn std::error::Error>> {
11371        let f = self.func("rope_neox_ff_f32");
11372        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11373        let grid = (n_heads * n_tokens) as u32;
11374        let cfg = LaunchConfig {
11375            grid_dim: (grid, 1, 1),
11376            block_dim: ((head_dim / 2) as u32, 1, 1),
11377            shared_mem_bytes: 0,
11378        };
11379        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11380        let __s_b = self.gpu.stream();
11381        let mut b = __s_b.launch_builder(&f);
11382        b.arg(x)
11383            .arg(pos)
11384            .arg(&hd)
11385            .arg(&nd)
11386            .arg(&nh)
11387            .arg(&theta_scale)
11388            .arg(&freq_scale)
11389            .arg(ff);
11390        unsafe {
11391            b.launch(cfg)?;
11392        }
11393        Ok(())
11394    }
11395
11396    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
11397    #[allow(clippy::too_many_arguments)]
11398    pub fn rope_neox2(
11399        &self,
11400        q: &mut CudaSlice<f32>,
11401        k: &mut CudaSlice<f32>,
11402        pos: &CudaSlice<i32>,
11403        head_dim: usize,
11404        n_dims: usize,
11405        nh_q: usize,
11406        nh_k: usize,
11407        n_tokens: usize,
11408        freq_base: f32,
11409        freq_scale: f32,
11410        ff: Option<&CudaSlice<f32>>,
11411    ) -> Result<(), Box<dyn std::error::Error>> {
11412        let f = self.func("rope_neox2_f32");
11413        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11414        let grid = ((nh_q + nh_k) * n_tokens) as u32;
11415        let cfg = LaunchConfig {
11416            grid_dim: (grid, 1, 1),
11417            block_dim: ((head_dim / 2) as u32, 1, 1),
11418            shared_mem_bytes: 0,
11419        };
11420        let (hd, nd, nq, nk, nt) = (
11421            head_dim as i32,
11422            n_dims as i32,
11423            nh_q as i32,
11424            nh_k as i32,
11425            n_tokens as i32,
11426        );
11427        let __s_b = self.gpu.stream();
11428        let mut b = __s_b.launch_builder(&f);
11429        b.arg(q)
11430            .arg(k)
11431            .arg(pos)
11432            .arg(&hd)
11433            .arg(&nd)
11434            .arg(&nq)
11435            .arg(&nk)
11436            .arg(&nt)
11437            .arg(&theta_scale)
11438            .arg(&freq_scale);
11439        match ff {
11440            Some(ffv) => {
11441                b.arg(ffv);
11442                unsafe {
11443                    b.launch(cfg)?;
11444                }
11445            }
11446            None => {
11447                let null: u64 = 0;
11448                b.arg(&null);
11449                unsafe {
11450                    b.launch(cfg)?;
11451                }
11452            }
11453        }
11454        Ok(())
11455    }
11456
11457    /// gemma4 R1: dst = GELU_tanh(gate) * up.
11458    pub fn gelu_tanh_mul(
11459        &self,
11460        gate: &CudaSlice<f32>,
11461        up: &CudaSlice<f32>,
11462        dst: &mut CudaSlice<f32>,
11463        n: usize,
11464    ) -> Result<(), Box<dyn std::error::Error>> {
11465        let f = self.func("gelu_tanh_mul_f32");
11466        let cfg = LaunchConfig::for_num_elems(n as u32);
11467        let ni = n as i32;
11468        let __s_b = self.gpu.stream();
11469        let mut b = __s_b.launch_builder(&f);
11470        b.arg(gate).arg(up).arg(dst).arg(&ni);
11471        unsafe {
11472            b.launch(cfg)?;
11473        }
11474        Ok(())
11475    }
11476
11477    pub fn silu_mul(
11478        &self,
11479        gate: &CudaSlice<f32>,
11480        up: &CudaSlice<f32>,
11481        dst: &mut CudaSlice<f32>,
11482        n: usize,
11483    ) -> Result<(), Box<dyn std::error::Error>> {
11484        let f = self.func("silu_mul_f32");
11485        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11486        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11487        let ni = n as i32;
11488        let __s_b = self.gpu.stream();
11489        let mut b = __s_b.launch_builder(&f);
11490        b.arg(gate).arg(up).arg(dst).arg(&ni);
11491        unsafe {
11492            b.launch(cfg)?;
11493        }
11494        Ok(())
11495    }
11496
11497    /// SwiGLU twin using Memra's host-matching expf transcription.
11498    pub fn silu_mul_host_expf(
11499        &self,
11500        gate: &CudaSlice<f32>,
11501        up: &CudaSlice<f32>,
11502        dst: &mut CudaSlice<f32>,
11503        n: usize,
11504    ) -> Result<(), Box<dyn std::error::Error>> {
11505        let f = self.func("silu_mul_host_expf_f32");
11506        let cfg = LaunchConfig::for_num_elems(n as u32);
11507        let ni = n as i32;
11508        let __s_b = self.gpu.stream();
11509        let mut b = __s_b.launch_builder(&f);
11510        b.arg(gate).arg(up).arg(dst).arg(&ni);
11511        unsafe {
11512            b.launch(cfg)?;
11513        }
11514        Ok(())
11515    }
11516
11517    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
11518    pub fn silu_clamped_mul_host_expf(
11519        &self,
11520        gate: &CudaSlice<f32>,
11521        up: &CudaSlice<f32>,
11522        limit: f32,
11523        dst: &mut CudaSlice<f32>,
11524        n: usize,
11525    ) -> Result<(), Box<dyn std::error::Error>> {
11526        if !limit.is_finite() || limit <= 0.0 {
11527            return Err(
11528                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11529            );
11530        }
11531        let f = self.func("silu_clamped_mul_host_expf_f32");
11532        let cfg = LaunchConfig::for_num_elems(n as u32);
11533        let ni = n as i32;
11534        let __s_b = self.gpu.stream();
11535        let mut b = __s_b.launch_builder(&f);
11536        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11537        unsafe {
11538            b.launch(cfg)?;
11539        }
11540        Ok(())
11541    }
11542
11543    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
11544    /// for the down projection — kills the standalone convert pass. Bit-identical class.
11545    pub fn silu_mul_f16out(
11546        &self,
11547        gate: &CudaSlice<f32>,
11548        up: &CudaSlice<f32>,
11549        dst: &mut CudaSlice<f32>,
11550        dst16: &mut CudaSlice<u8>,
11551        n: usize,
11552    ) -> Result<(), Box<dyn std::error::Error>> {
11553        let f = self.func("silu_mul_f16out_f32");
11554        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11555        let ni = n as i32;
11556        let __s_b = self.gpu.stream();
11557        let mut b = __s_b.launch_builder(&f);
11558        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11559        unsafe {
11560            b.launch(cfg)?;
11561        }
11562        Ok(())
11563    }
11564
11565    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
11566    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
11567    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
11568    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
11569    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
11570    /// launches per dense FFN layer (the gate+up post-matmul scales).
11571    pub fn silu_mul_scaled(
11572        &self,
11573        gate: &CudaSlice<f32>,
11574        up: &CudaSlice<f32>,
11575        gs: f32,
11576        us: f32,
11577        dst: &mut CudaSlice<f32>,
11578        n: usize,
11579    ) -> Result<(), Box<dyn std::error::Error>> {
11580        let f = self.func("silu_mul_scaled_f32");
11581        let cfg = LaunchConfig::for_num_elems(n as u32);
11582        let ni = n as i32;
11583        let (gsf, usf) = (gs, us);
11584        let __s_b = self.gpu.stream();
11585        let mut b = __s_b.launch_builder(&f);
11586        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11587        unsafe {
11588            b.launch(cfg)?;
11589        }
11590        Ok(())
11591    }
11592
11593    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
11594    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
11595    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
11596    #[allow(clippy::too_many_arguments)]
11597    pub fn swigluoai_mul_scaled(
11598        &self,
11599        gate: &CudaSlice<f32>,
11600        up: &CudaSlice<f32>,
11601        gs: f32,
11602        us: f32,
11603        alpha: f32,
11604        limit: f32,
11605        dst: &mut CudaSlice<f32>,
11606        n: usize,
11607    ) -> Result<(), Box<dyn std::error::Error>> {
11608        let f = self.func("swigluoai_mul_scaled_f32");
11609        let cfg = LaunchConfig::for_num_elems(n as u32);
11610        let ni = n as i32;
11611        let __s_b = self.gpu.stream();
11612        let mut b = __s_b.launch_builder(&f);
11613        b.arg(gate)
11614            .arg(up)
11615            .arg(&gs)
11616            .arg(&us)
11617            .arg(&alpha)
11618            .arg(&limit)
11619            .arg(dst)
11620            .arg(&ni);
11621        unsafe {
11622            b.launch(cfg)?;
11623        }
11624        Ok(())
11625    }
11626
11627    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11628    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11629    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11630    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11631    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11632    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11633    /// n must be a multiple of 32 (n_ff always is).
11634    pub fn silu_mul_scaled_q8_1(
11635        &self,
11636        gate: &CudaSlice<f32>,
11637        up: &CudaSlice<f32>,
11638        gs: f32,
11639        us: f32,
11640        n: usize,
11641    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11642        let f = self.func("silu_mul_scaled_q8_1");
11643        let nblk = n / 32;
11644        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11645        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11646        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11647        let cfg = LaunchConfig::for_num_elems(n as u32);
11648        let (gsf, usf, ni) = (gs, us, n as i32);
11649        let __s_b = self.gpu.stream();
11650        let mut b = __s_b.launch_builder(&f);
11651        b.arg(gate)
11652            .arg(up)
11653            .arg(&gsf)
11654            .arg(&usf)
11655            .arg(&mut aq)
11656            .arg(&mut ad)
11657            .arg(&ni);
11658        unsafe {
11659            b.launch(cfg)?;
11660        }
11661        Ok((aq, ad))
11662    }
11663
11664    pub fn add(
11665        &self,
11666        a: &CudaSlice<f32>,
11667        b_in: &CudaSlice<f32>,
11668        dst: &mut CudaSlice<f32>,
11669        n: usize,
11670    ) -> Result<(), Box<dyn std::error::Error>> {
11671        let f = self.func("add_f32");
11672        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11673        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11674        let ni = n as i32;
11675        let __s_bld = self.gpu.stream();
11676        let mut bld = __s_bld.launch_builder(&f);
11677        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11678        unsafe {
11679            bld.launch(cfg)?;
11680        }
11681        Ok(())
11682    }
11683
11684    pub fn mul(
11685        &self,
11686        a: &CudaSlice<f32>,
11687        b_in: &CudaSlice<f32>,
11688        dst: &mut CudaSlice<f32>,
11689        n: usize,
11690    ) -> Result<(), Box<dyn std::error::Error>> {
11691        let f = self.func("mul_f32");
11692        let cfg = LaunchConfig::for_num_elems(n as u32);
11693        let ni = n as i32;
11694        let __s_bld = self.gpu.stream();
11695        let mut bld = __s_bld.launch_builder(&f);
11696        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11697        unsafe {
11698            bld.launch(cfg)?;
11699        }
11700        Ok(())
11701    }
11702
11703    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11704    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11705    pub fn matmul(
11706        &self,
11707        w: &crate::model::GpuTensor,
11708        x: &CudaSlice<f32>,
11709        m: usize,
11710    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11711        use crate::model::GpuTensor;
11712        let in_f = w.in_features();
11713        let out_f = w.out_features();
11714        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11715        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11716        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11717        // gives nothing). Quantize the activation once here then call the GEMM.
11718        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11719        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11720        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11721        #[allow(non_snake_case)]
11722        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11723        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11724        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11725            usize::MAX
11726        } else {
11727            16usize
11728        };
11729
11730        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11731        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11732        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11733        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11734        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11735        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11736        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11737        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11738        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11739        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11740        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11741        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11742        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11743        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11744        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11745        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11746        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11747        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11748        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11749        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11750        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11751        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11752        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11753        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11754        if m >= GEMM_M_THRESHOLD {
11755            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11756                return Ok(y);
11757            }
11758            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
11759            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
11760            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
11761            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
11762            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
11763            // tile defaults differently by operand source.
11764            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11765                return Ok(y);
11766            }
11767            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
11768            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
11769            if let Some(y) = self.try_f16_gemm(w, x, m)? {
11770                return Ok(y);
11771            }
11772        }
11773        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
11774        // m threshold the rest of this method uses:
11775        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
11776        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
11777        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
11778        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
11779        //     across every tier by construction with no batched twin needed.
11780        //
11781        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
11782        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
11783        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
11784        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
11785        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
11786        // arms is what makes sure it never gets there.
11787        if let GpuTensor::Quant { qtype, .. } = w {
11788            if *qtype == QT_F8_E4M3_BLK {
11789                if m >= GEMM_M_THRESHOLD {
11790                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11791                        return Ok(y);
11792                    }
11793                }
11794                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11795                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11796                    return Ok(y);
11797                }
11798            }
11799        }
11800        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11801            return self.qmatvec_mmq(w, x, m);
11802        }
11803        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11804            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11805            return self.qmatvec_gemm(w, &aq, &ad, m);
11806        }
11807        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
11808        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
11809        if m >= GEMM_M_THRESHOLD {
11810            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11811                return Ok(y);
11812            }
11813        }
11814        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
11815        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
11816        // to Stage-A f32-dequant (the correctness oracle path).
11817        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11818        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
11819        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
11820        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
11821        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
11822        if m == 1 && fast {
11823            if let GpuTensor::Quant {
11824                bytes,
11825                qtype,
11826                row_bytes,
11827                rp,
11828                rp4,
11829                scale,
11830                ..
11831            } = w
11832            {
11833                if self.mmvq_supports(*qtype) {
11834                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
11835                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
11836                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
11837                    let (bytes, rp) = match rp4 {
11838                        Some(m4) => (m4, true),
11839                        None => (bytes, *rp),
11840                    };
11841                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11842                    return self.qmatvec_mmvq(
11843                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11844                    );
11845                }
11846            }
11847        }
11848        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
11849        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
11850        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
11851        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
11852        // block below. MEMRA_NO_BATCHED -> per-m path.
11853        //
11854        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
11855        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
11856        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
11857        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
11858        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
11859        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
11860        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
11861        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
11862        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
11863        if (2..=16).contains(&m)
11864            && fast
11865            && std::env::var("MEMRA_NO_BATCHED").is_err()
11866            && (m <= 4 || Self::b8_enabled())
11867        {
11868            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
11869            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
11870            // is present (rp4) — the mirror pick below then routes to the _rp family.
11871            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
11872            // because the native e4m3 row layout is already aligned and needs no mirror.
11873            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
11874            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
11875            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
11876            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
11877            let m_ok = m <= 8
11878                || matches!(w, GpuTensor::Quant { qtype, .. }
11879                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11880                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11881            if m_ok {
11882                if let GpuTensor::Quant {
11883                    bytes,
11884                    qtype,
11885                    row_bytes,
11886                    rp,
11887                    rp4,
11888                    ..
11889                } = w
11890                {
11891                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11892                        let (bytes, rp) = match rp4 {
11893                            Some(m4) => (m4, true),
11894                            None => (bytes, *rp),
11895                        };
11896                        let mcols = Self::batched_mcols(m);
11897                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11898                        let mut y = self.qmatvec_mmvq_batched(
11899                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11900                        )?;
11901                        if let GpuTensor::Quant { scale, .. } = w {
11902                            if *scale != 1.0 {
11903                                self.scale_inplace(&mut y, *scale, m * out_f)?;
11904                            }
11905                        }
11906                        return Ok(y);
11907                    }
11908                }
11909            }
11910        }
11911        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
11912        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
11913        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
11914        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
11915        // for this dtype, so the generic match below must never see it under `fast`.
11916        if fast {
11917            if let GpuTensor::Quant {
11918                bytes,
11919                qtype,
11920                row_bytes,
11921                scale,
11922                ..
11923            } = w
11924            {
11925                if *qtype == QT_F8_E4M3 {
11926                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11927                    return self.qmatvec_mmvq(
11928                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11929                    );
11930                }
11931            }
11932        }
11933        let mut y = match w {
11934            GpuTensor::Quant {
11935                bytes,
11936                qtype,
11937                row_bytes,
11938                ..
11939            } if fast && *qtype == QT_Q8_0 => {
11940                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11941            }
11942            GpuTensor::Quant {
11943                bytes,
11944                qtype,
11945                row_bytes,
11946                ..
11947            } if fast && *qtype == QT_Q4_K => {
11948                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11949            }
11950            GpuTensor::Quant {
11951                bytes,
11952                qtype,
11953                row_bytes,
11954                ..
11955            } if fast && *qtype == QT_Q6_K => {
11956                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11957            }
11958            GpuTensor::Quant {
11959                bytes,
11960                qtype,
11961                row_bytes,
11962                ..
11963            } if fast && *qtype == QT_Q5_K => {
11964                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11965            }
11966            GpuTensor::Quant {
11967                bytes,
11968                qtype,
11969                row_bytes,
11970                ..
11971            } if fast && *qtype == QT_Q3_K => {
11972                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11973            }
11974            GpuTensor::Quant {
11975                bytes,
11976                qtype,
11977                row_bytes,
11978                rp,
11979                ..
11980            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11981                if *rp {
11982                    "qmatvec_nvfp4_dp4a_rp"
11983                } else {
11984                    "qmatvec_nvfp4_dp4a"
11985                },
11986                &bytes.slice(0..bytes.len()),
11987                x,
11988                m,
11989                in_f,
11990                out_f,
11991                *row_bytes,
11992            )?,
11993            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
11994            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
11995            // anomaly (research/kat-anomaly-20260802/).
11996            GpuTensor::Quant {
11997                bytes,
11998                qtype,
11999                row_bytes,
12000                ..
12001            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12002                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12003            }
12004            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
12005            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
12006            // without first writing the matching kernel, or func() will panic
12007            // "kernel ... not in any fatbin".
12008            GpuTensor::Quant {
12009                bytes,
12010                qtype,
12011                row_bytes,
12012                rp,
12013                ..
12014            } =>
12015            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
12016            // deq(row,j) form cannot address the planes; same value/product order).
12017            {
12018                self.qmatvec(
12019                    bytes,
12020                    x,
12021                    m,
12022                    in_f,
12023                    out_f,
12024                    if *rp && *qtype == QT_NVFP4 {
12025                        QT_NVFP4_RP
12026                    } else {
12027                        *qtype
12028                    },
12029                    *row_bytes,
12030                )?
12031            }
12032            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12033            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
12034            // cuBLASLt f32 GEMV as the Float arm.
12035            GpuTensor::FloatBf16 { data, .. } => {
12036                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
12037                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
12038                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
12039                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
12040                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
12041                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
12042                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12043                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12044                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12045                    y
12046                } else {
12047                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12048                }
12049            }
12050        };
12051        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
12052        if let GpuTensor::Quant { scale, .. } = w {
12053            if *scale != 1.0 {
12054                self.scale_inplace(&mut y, *scale, m * out_f)?;
12055            }
12056        }
12057        Ok(y)
12058    }
12059
12060    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
12061    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
12062    ///
12063    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
12064    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
12065    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
12066    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
12067    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
12068    /// path must not pay an env lookup for a flag that is off.
12069    pub fn stage_a_raw_needed() -> bool {
12070        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12071        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12072    }
12073
12074    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
12075    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
12076    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12077        use crate::model::GpuTensor;
12078        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12079            return false;
12080        }
12081        match w {
12082            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
12083            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
12084            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
12085            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
12086            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
12087            // block class has no fused twin yet, so each of its projections takes its own launch.
12088            GpuTensor::Quant { qtype, .. } => {
12089                matches!(
12090                    *qtype,
12091                    QT_Q8_0
12092                        | QT_Q4_K
12093                        | QT_Q6_K
12094                        | QT_Q5_K
12095                        | QT_Q3_K
12096                        | QT_NVFP4
12097                        | QT_F8_E4M3
12098                        | QT_F8_E4M3_BLK
12099                        | QT_Q4_0
12100                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12101            }
12102            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12103        }
12104    }
12105
12106    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
12107    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
12108    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
12109    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
12110    pub fn matmul_pre(
12111        &self,
12112        w: &crate::model::GpuTensor,
12113        aq: &CudaSlice<i8>,
12114        ad: &CudaSlice<f32>,
12115        x_fallback: &CudaSlice<f32>,
12116        m: usize,
12117    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12118        use crate::model::GpuTensor;
12119        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
12120        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
12121        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
12122        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
12123        // rc=30013 dig, 2026-07-31).
12124        let x_raw_ok = x_fallback.len() >= m * w.in_features();
12125        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
12126        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
12127        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12128            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12129                return Ok(y);
12130            }
12131            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
12132            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
12133            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12134                return Ok(y);
12135            }
12136            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
12137            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12138                return Ok(y);
12139            }
12140        }
12141        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
12142        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
12143        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
12144        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
12145        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
12146        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12147            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12148                return Ok(y);
12149            }
12150        }
12151        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12152            return Ok(y);
12153        }
12154        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
12155        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
12156        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
12157        // aq/ad.
12158        if m >= 16
12159            && w.out_features() >= 128
12160            && self.mmq_supports(w)
12161            && !self.verify_exact_on()
12162            && x_raw_ok
12163        {
12164            return self.qmatvec_mmq(w, x_fallback, m);
12165        }
12166        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
12167        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
12168        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12169            if let Some(y) =
12170                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12171            {
12172                return Ok(y);
12173            }
12174        }
12175        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
12176        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
12177        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12178            return self.qmatvec_gemm(w, aq, ad, m);
12179        }
12180        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
12181        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
12182        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
12183        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
12184        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
12185        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
12186        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
12187        // which reads `m * in_f` floats out of a 0-byte allocation ->
12188        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
12189        // it poisons the context, so every LATER request in that process fails with an unrelated
12190        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
12191        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
12192        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
12193        // dense artifact and left the arm with no working truth instrument.
12194        //
12195        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
12196        // strictly better than an illegal address surfacing later at an unrelated sync point, and
12197        // an oracle that cannot run must say so rather than corrupt the context it runs in.
12198        if !self.uses_q8_1_fast(w) {
12199            if !x_raw_ok {
12200                return Err(format!(
12201                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12202                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12203                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
12204                     activation (see Engine::rms_norm_decode, which is bit-identical to \
12205                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12206                    x_fallback.len(),
12207                    m,
12208                    w.in_features(),
12209                    m * w.in_features()
12210                )
12211                .into());
12212            }
12213            return self.matmul(w, x_fallback, m);
12214        }
12215        let in_f = w.in_features();
12216        let out_f = w.out_features();
12217        let (bytes, qtype, row_bytes, scale, rp) = match w {
12218            GpuTensor::Quant {
12219                bytes,
12220                qtype,
12221                row_bytes,
12222                scale,
12223                rp,
12224                ..
12225            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12226            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12227        };
12228        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
12229        // the dp4a/oracle tails below keep the raw GGUF bytes.
12230        let (mbytes, mrp) = match w {
12231            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12232            _ => (bytes, rp),
12233        };
12234        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
12235        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
12236        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
12237        if m == 1 && self.mmvq_supports(qtype) {
12238            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12239        }
12240        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
12241        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
12242        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
12243        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
12244        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
12245        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
12246        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
12247        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
12248        // m=5..8 on the old per-m path (b8-tier-only seam).
12249        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
12250        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
12251        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
12252        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12253            && std::env::var("MEMRA_NO_BATCHED").is_err()
12254            && (m <= 4 || Self::b8_enabled())
12255            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
12256            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
12257            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
12258            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12259                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12260        {
12261            let mcols = Self::batched_mcols(m);
12262            return self.qmatvec_mmvq_batched(
12263                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12264            );
12265        }
12266        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
12267        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
12268        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
12269        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
12270        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
12271        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12272            let (b2, r2) = if qtype == QT_Q4_0 {
12273                (mbytes, mrp)
12274            } else {
12275                (bytes, rp)
12276            };
12277            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12278        }
12279        let name = match qtype {
12280            QT_Q8_0 => "qmatvec_q8_0_dp4a",
12281            QT_Q4_K => "qmatvec_q4_K_dp4a",
12282            QT_Q6_K => "qmatvec_q6_K_dp4a",
12283            QT_Q5_K => "qmatvec_q5_K_dp4a",
12284            QT_Q3_K => "qmatvec_q3_K_dp4a",
12285            QT_NVFP4 => {
12286                if rp {
12287                    "qmatvec_nvfp4_dp4a_rp"
12288                } else {
12289                    "qmatvec_nvfp4_dp4a"
12290                }
12291            }
12292            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12293            _ => unreachable!(),
12294        };
12295        let f = self.func(name);
12296        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
12297        let cfg = LaunchConfig {
12298            grid_dim: (out_f as u32, m as u32, 1),
12299            block_dim: (128, 1, 1),
12300            shared_mem_bytes: 0,
12301        };
12302        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12303        let __s_b = self.gpu.stream();
12304        let mut b = __s_b.launch_builder(&f);
12305        b.arg(bytes)
12306            .arg(aq)
12307            .arg(ad)
12308            .arg(&mut y)
12309            .arg(&inf)
12310            .arg(&outf)
12311            .arg(&mi)
12312            .arg(&rb);
12313        unsafe {
12314            b.launch(cfg)?;
12315        }
12316        if scale != 1.0 {
12317            self.scale_inplace(&mut y, scale, m * out_f)?;
12318        }
12319        Ok(y)
12320    }
12321
12322    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
12323    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
12324    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
12325    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
12326    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
12327    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
12328    /// reduce as m=1); this method just forces that path unconditionally.
12329    pub fn matmul_decode_exact(
12330        &self,
12331        w: &crate::model::GpuTensor,
12332        x: &CudaSlice<f32>,
12333        m: usize,
12334    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12335        use crate::model::GpuTensor;
12336        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
12337        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
12338        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
12339        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
12340        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
12341        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
12342        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
12343        if let GpuTensor::Float { data, .. } = w {
12344            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12345        }
12346        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
12347        // float linear (same n-independent reduction contract as the Float arm above).
12348        if let GpuTensor::FloatBf16 { data, .. } = w {
12349            let (in_f, out_f) = (w.in_features(), w.out_features());
12350            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
12351            // contract — the whole-weight f32 dequant disappears too).
12352            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12353                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12354                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12355                return Ok(y);
12356            }
12357            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12358        }
12359        if !self.uses_q8_1_fast(w) {
12360            return self.matmul(w, x, m);
12361        }
12362        let in_f = w.in_features();
12363        let out_f = w.out_features();
12364        let (bytes, qtype, row_bytes, scale, rp) = match w {
12365            GpuTensor::Quant {
12366                bytes,
12367                qtype,
12368                row_bytes,
12369                scale,
12370                rp,
12371                ..
12372            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12373            _ => return self.matmul(w, x, m),
12374        };
12375        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
12376        // which does its own mirror pick).
12377        let (bytes, rp) = match w {
12378            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12379            _ => (bytes, rp),
12380        };
12381        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12382        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
12383        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
12384        // (token,row) by construction, which is exactly what this method exists to guarantee.
12385        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12386            return Ok(y);
12387        }
12388        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
12389        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
12390        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
12391        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
12392        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
12393        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
12394        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
12395        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
12396        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12397            && std::env::var("MEMRA_NO_BATCHED").is_err()
12398            && (m <= 4 || Self::b8_enabled())
12399            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
12400            // no mirror precondition, `rp` selects the layout only.
12401            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12402                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12403        {
12404            let mcols = Self::batched_mcols(m);
12405            return self.qmatvec_mmvq_batched(
12406                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12407            );
12408        }
12409        if self.mmvq_supports(qtype) {
12410            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
12411            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
12412            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12413        }
12414        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
12415        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
12416        self.matmul_pre(w, &aq, &ad, x, m)
12417    }
12418
12419    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
12420    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
12421    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
12422    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
12423    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
12424    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
12425    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
12426    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
12427    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
12428    pub fn matmul_decode_exact_pre(
12429        &self,
12430        w: &crate::model::GpuTensor,
12431        aq: &CudaSlice<i8>,
12432        ad: &CudaSlice<f32>,
12433        m: usize,
12434    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12435        use crate::model::GpuTensor;
12436        debug_assert!(
12437            self.uses_q8_1_fast(w),
12438            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12439        );
12440        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
12441        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12442            return Ok(y);
12443        }
12444        let in_f = w.in_features();
12445        let out_f = w.out_features();
12446        let (bytes, qtype, row_bytes, scale, rp) = match w {
12447            GpuTensor::Quant {
12448                bytes,
12449                qtype,
12450                row_bytes,
12451                scale,
12452                rp,
12453                ..
12454            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12455            _ => {
12456                return Err(
12457                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12458                );
12459            }
12460        };
12461        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
12462        let (bytes, rp) = match w {
12463            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12464            _ => (bytes, rp),
12465        };
12466        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
12467        if (2..=16).contains(&m)
12468            && self.batched_supports(qtype)
12469            && self.mmvq_supports(qtype)
12470            && std::env::var("MEMRA_NO_BATCHED").is_err()
12471            && (m <= 4 || Self::b8_enabled())
12472            && (m <= 8
12473                || qtype == QT_Q4_0
12474                || qtype == QT_Q6_K
12475                || qtype == QT_F8_E4M3
12476                || qtype == QT_NVFP4
12477                || qtype == QT_Q4_K
12478                || qtype == QT_Q5_K
12479                || qtype == QT_Q8_0)
12480        {
12481            let mcols = Self::batched_mcols(m);
12482            return self.qmatvec_mmvq_batched(
12483                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12484            );
12485        }
12486        if self.mmvq_supports(qtype) {
12487            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12488        }
12489        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
12490        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
12491        let x0 = self.zeros(0)?;
12492        self.matmul_pre(w, aq, ad, &x0, m)
12493    }
12494
12495    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
12496    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
12497    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
12498    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
12499    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
12500    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
12501    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
12502    /// per-tensor path.
12503    pub fn matmul_decode_exact_dual_pre(
12504        &self,
12505        w0: &crate::model::GpuTensor,
12506        w1: &crate::model::GpuTensor,
12507        aq: &CudaSlice<i8>,
12508        ad: &CudaSlice<f32>,
12509        m: usize,
12510    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12511    {
12512        use crate::model::GpuTensor;
12513        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12514        let on = *ON.get_or_init(|| {
12515            std::env::var("MEMRA_SPEC_DUAL_T")
12516                .map(|v| v != "0")
12517                .unwrap_or(true)
12518        });
12519        if !on
12520            || !(2..=7).contains(&m)
12521            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12522            || !self.uses_q8_1_fast(w0)
12523            || !self.uses_q8_1_fast(w1)
12524        {
12525            return Ok(None);
12526        }
12527        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
12528        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
12529        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
12530        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
12531        if !self.mmvq_supports(QT_NVFP4) {
12532            return Ok(None);
12533        }
12534        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12535        if w1.in_features() != in_f || w1.out_features() != out_f {
12536            return Ok(None);
12537        }
12538        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12539            (
12540                GpuTensor::Quant {
12541                    bytes: b0,
12542                    qtype: q0,
12543                    row_bytes: rb0,
12544                    scale: s0,
12545                    rp: rp0,
12546                    rp4: None,
12547                    ..
12548                },
12549                GpuTensor::Quant {
12550                    bytes: b1,
12551                    qtype: q1,
12552                    row_bytes: rb1,
12553                    scale: s1,
12554                    rp: rp1,
12555                    rp4: None,
12556                    ..
12557                },
12558            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12559                (b0, b1, *rb0, *s0, *s1, *rp0)
12560            }
12561            _ => return Ok(None),
12562        };
12563        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
12564        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
12565        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12566        {
12567            return Ok(None);
12568        }
12569        let (y0, y1) =
12570            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12571        Ok(Some(((y0, s0), (y1, s1))))
12572    }
12573
12574    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
12575    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
12576    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
12577    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
12578    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
12579    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
12580    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
12581    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
12582    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
12583    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
12584    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
12585    pub fn matmul_decode_exact_group4_pre(
12586        &self,
12587        ws: [&crate::model::GpuTensor; 4],
12588        aq: &CudaSlice<i8>,
12589        ad: &CudaSlice<f32>,
12590        m: usize,
12591    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12592        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12593        let on = *ON.get_or_init(|| {
12594            std::env::var("MEMRA_TK_GDN_GROUP")
12595                .map(|v| v != "0")
12596                .unwrap_or(true)
12597        });
12598        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12599    }
12600
12601    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
12602    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
12603    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
12604    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
12605    pub fn matmul_decode_exact_group3_pre(
12606        &self,
12607        ws: [&crate::model::GpuTensor; 3],
12608        aq: &CudaSlice<i8>,
12609        ad: &CudaSlice<f32>,
12610        m: usize,
12611    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12612        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12613        let on = *ON.get_or_init(|| {
12614            std::env::var("MEMRA_TK_FA_GROUP")
12615                .map(|v| v != "0")
12616                .unwrap_or(true)
12617        });
12618        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12619    }
12620
12621    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
12622    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
12623    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
12624    fn matmul_decode_exact_group_pre(
12625        &self,
12626        ws: &[&crate::model::GpuTensor],
12627        aq: &CudaSlice<i8>,
12628        ad: &CudaSlice<f32>,
12629        m: usize,
12630        on: bool,
12631        tag: &'static str,
12632    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12633        use crate::model::GpuTensor;
12634        if !on
12635            || !(2..=16).contains(&m)
12636            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12637            || (m > 4 && !Self::b8_enabled())
12638            || !self.mmvq_supports(QT_NVFP4)
12639            || !self.batched_supports(QT_NVFP4)
12640        {
12641            return Ok(None);
12642        }
12643        let in_f = ws[0].in_features();
12644        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12645        for w in ws {
12646            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12647                return Ok(None);
12648            }
12649            match w {
12650                GpuTensor::Quant {
12651                    bytes,
12652                    qtype,
12653                    scale,
12654                    rp: true,
12655                    rp4: None,
12656                    ..
12657                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12658                    parts.push((bytes, w.out_features(), *scale));
12659                }
12660                _ => return Ok(None),
12661            }
12662        }
12663        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12664        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12665        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12666        let mcols = if (5..=7).contains(&m) && b567 {
12667            m
12668        } else {
12669            Self::batched_mcols(m)
12670        };
12671        let kname: &'static str = match mcols {
12672            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12673            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12674            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12675            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12676            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12677            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12678            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12679            _ => return Ok(None),
12680        };
12681        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12682        // the second door's print on the slice-D battery — key the once-set by tag.
12683        if std::env::var("MEMRA_DEBUG").is_ok() {
12684            use std::sync::Mutex;
12685            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12686            let mut seen = SEEN.lock().unwrap();
12687            if !seen.contains(&tag) {
12688                seen.push(tag);
12689                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12690            }
12691        }
12692        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12693        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12694        let total: usize = parts.iter().map(|p| p.1).sum();
12695        let three = parts.len() == 3;
12696        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12697        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12698        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12699        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12700        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12701        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12702        let cfg = LaunchConfig {
12703            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12704            block_dim: (32, ROWS_PER_BLOCK, 1),
12705            shared_mem_bytes: 0,
12706        };
12707        let (inf, mi) = (in_f as i32, m as i32);
12708        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12709        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12710        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12711        let s3 = if three { 1.0f32 } else { parts[3].2 };
12712        let w3 = if three { parts[0].0 } else { parts[3].0 };
12713        let f = self.func(kname);
12714        let __s_b = self.gpu.stream();
12715        let mut b = __s_b.launch_builder(&f);
12716        b.arg(parts[0].0)
12717            .arg(parts[1].0)
12718            .arg(parts[2].0)
12719            .arg(w3)
12720            .arg(aq)
12721            .arg(ad)
12722            .arg(&mut y0)
12723            .arg(&mut y1)
12724            .arg(&mut y2)
12725            .arg(&mut y3)
12726            .arg(&inf)
12727            .arg(&n0)
12728            .arg(&n1)
12729            .arg(&n2)
12730            .arg(&n3)
12731            .arg(&mi)
12732            .arg(&s0)
12733            .arg(&s1)
12734            .arg(&s2)
12735            .arg(&s3);
12736        unsafe {
12737            b.launch(cfg)?;
12738        }
12739        Ok(Some(if three {
12740            vec![y0, y1, y2]
12741        } else {
12742            vec![y0, y1, y2, y3]
12743        }))
12744    }
12745
12746    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12747    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12748    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12749    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12750    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12751    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12752    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12753    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12754    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12755    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12756    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
12757    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
12758    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
12759    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
12760    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
12761    pub fn matmul_decode_exact_dual(
12762        &self,
12763        w0: &crate::model::GpuTensor,
12764        w1: &crate::model::GpuTensor,
12765        x: &CudaSlice<f32>,
12766        m: usize,
12767    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12768        use crate::model::GpuTensor;
12769        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12770        let on = *ON.get_or_init(|| {
12771            std::env::var("MEMRA_SPEC_DUAL_T")
12772                .map(|v| v != "0")
12773                .unwrap_or(true)
12774        });
12775        if !on
12776            || !(2..=4).contains(&m)
12777            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12778            || !self.uses_q8_1_fast(w0)
12779            || !self.uses_q8_1_fast(w1)
12780        {
12781            return Ok(None);
12782        }
12783        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
12784        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
12785        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
12786        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
12787        if !self.mmvq_supports(QT_NVFP4) {
12788            return Ok(None);
12789        }
12790        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12791        if w1.in_features() != in_f || w1.out_features() != out_f {
12792            return Ok(None);
12793        }
12794        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12795            (
12796                GpuTensor::Quant {
12797                    bytes: b0,
12798                    qtype: q0,
12799                    row_bytes: rb0,
12800                    scale: s0,
12801                    rp: rp0,
12802                    rp4: None,
12803                    ..
12804                },
12805                GpuTensor::Quant {
12806                    bytes: b1,
12807                    qtype: q1,
12808                    row_bytes: rb1,
12809                    scale: s1,
12810                    rp: rp1,
12811                    rp4: None,
12812                    ..
12813                },
12814            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12815                (b0, b1, *rb0, *s0, *s1, *rp0)
12816            }
12817            _ => return Ok(None),
12818        };
12819        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
12820        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
12821        if std::env::var("MEMRA_DEBUG").is_ok() {
12822            static ONCE: std::sync::Once = std::sync::Once::new();
12823            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12824        }
12825        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12826        let (y0, y1) =
12827            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12828        let mut y0 = y0;
12829        let mut y1 = y1;
12830        if s0 != 1.0 {
12831            self.scale_inplace(&mut y0, s0, m * out_f)?;
12832        }
12833        if s1 != 1.0 {
12834            self.scale_inplace(&mut y1, s1, m * out_f)?;
12835        }
12836        Ok(Some((y0, y1)))
12837    }
12838
12839    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
12840    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
12841    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
12842    /// twins (both buffers must be the repacked layout).
12843    #[allow(clippy::too_many_arguments)]
12844    pub fn qmatvec_batched_dual_raw(
12845        &self,
12846        b0: &CudaSlice<u8>,
12847        b1: &CudaSlice<u8>,
12848        aq: &CudaSlice<i8>,
12849        ad: &CudaSlice<f32>,
12850        m: usize,
12851        in_f: usize,
12852        out_f: usize,
12853        row_bytes: usize,
12854        rp: bool,
12855    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12856        const ROWS_PER_BLOCK: u32 = 4;
12857        let mcols = Self::batched_mcols(m);
12858        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
12859        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
12860        let tiny_rp1 = rp
12861            && mcols == 4
12862            && out_f <= 128
12863            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12864        let (name, rows_per_block) = if tiny_rp1 {
12865            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12866        } else {
12867            match (mcols, rp, m) {
12868                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12869                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12870                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12871                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12872                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12873                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12874                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12875                _ => {
12876                    return Err(
12877                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12878                    );
12879                }
12880            }
12881        };
12882        let f = self.func(name);
12883        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12884        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12885        let cfg = LaunchConfig {
12886            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12887            block_dim: (32, ROWS_PER_BLOCK, 1),
12888            shared_mem_bytes: 0,
12889        };
12890        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12891        let __s_b = self.gpu.stream();
12892        let mut b = __s_b.launch_builder(&f);
12893        b.arg(b0)
12894            .arg(b1)
12895            .arg(aq)
12896            .arg(ad)
12897            .arg(&mut y0)
12898            .arg(&mut y1)
12899            .arg(&inf)
12900            .arg(&outf)
12901            .arg(&mi)
12902            .arg(&rb);
12903        unsafe {
12904            b.launch(cfg)?;
12905        }
12906        Ok((y0, y1))
12907    }
12908
12909    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
12910    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
12911    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
12912    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
12913    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
12914    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
12915    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
12916    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
12917    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
12918    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
12919    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
12920    pub fn matmul_pre_dual_noscale(
12921        &self,
12922        w0: &crate::model::GpuTensor,
12923        w1: &crate::model::GpuTensor,
12924        aq: &CudaSlice<i8>,
12925        ad: &CudaSlice<f32>,
12926        m: usize,
12927    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12928    {
12929        use crate::model::GpuTensor;
12930        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12931            return Ok(None);
12932        }
12933        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
12934        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
12935        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
12936        // would mix dispatch families across the pair — the exact class `q8_fused_params`
12937        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
12938        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
12939        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
12940        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
12941        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
12942        if !self.mmvq_supports(QT_NVFP4) {
12943            return Ok(None);
12944        }
12945        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12946        if w1.in_features() != in_f || w1.out_features() != out_f {
12947            return Ok(None);
12948        }
12949        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
12950        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
12951        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
12952        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
12953        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
12954        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
12955        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
12956        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
12957        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
12958        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
12959        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
12960        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
12961        let no_mirror =
12962            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12963        if self.q8_ffn_fuse2_on()
12964            && no_mirror(w0)
12965            && no_mirror(w1)
12966            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12967        {
12968            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12969            return Ok(Some(((y0, 1.0), (y1, 1.0))));
12970        }
12971        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
12972        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
12973        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
12974        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
12975        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
12976        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
12977        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
12978        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
12979        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
12980        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12981            let (y0, y1) =
12982                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12983            return Ok(Some(((y0, p0.3), (y1, p1.3))));
12984        }
12985        let (b0, q0, rb0, s0, rp0) = match w0 {
12986            GpuTensor::Quant {
12987                bytes,
12988                qtype,
12989                row_bytes,
12990                scale,
12991                rp,
12992                ..
12993            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12994            _ => return Ok(None),
12995        };
12996        let (b1, q1, rb1, s1, rp1) = match w1 {
12997            GpuTensor::Quant {
12998                bytes,
12999                qtype,
13000                row_bytes,
13001                scale,
13002                rp,
13003                ..
13004            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13005            _ => return Ok(None),
13006        };
13007        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13008            return Ok(None);
13009        }
13010        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13011        const RPW: u32 = 2;
13012        let rows_per_block = ROWS_PER_BLOCK * RPW;
13013        let f = self.func(if rp0 {
13014            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13015        } else {
13016            "qmatvec_nvfp4_mmvq_dual_mr2"
13017        });
13018        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13019        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13020        let cfg = LaunchConfig {
13021            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13022            block_dim: (32, ROWS_PER_BLOCK, 1),
13023            shared_mem_bytes: 0,
13024        };
13025        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13026        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
13027        // yscale args stay 1.0 here (they exist for the single-tensor callers).
13028        let one = 1.0f32;
13029        let __s_b = self.gpu.stream();
13030        let mut b = __s_b.launch_builder(&f);
13031        b.arg(b0)
13032            .arg(b1)
13033            .arg(aq)
13034            .arg(ad)
13035            .arg(&mut y0)
13036            .arg(&mut y1)
13037            .arg(&inf)
13038            .arg(&outf)
13039            .arg(&mi)
13040            .arg(&rb)
13041            .arg(&one)
13042            .arg(&one);
13043        unsafe {
13044            b.launch(cfg)?;
13045        }
13046        Ok(Some(((y0, s0), (y1, s1))))
13047    }
13048
13049    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
13050    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
13051    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
13052    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
13053    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
13054    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
13055    /// back to the three singles.
13056    #[allow(clippy::too_many_arguments)]
13057    pub fn matmul_nvfp4_fused3(
13058        &self,
13059        w0: &crate::model::GpuTensor,
13060        w1: &crate::model::GpuTensor,
13061        w2: &crate::model::GpuTensor,
13062        aq: &CudaSlice<i8>,
13063        ad: &CudaSlice<f32>,
13064        m: usize,
13065    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13066    {
13067        use crate::model::GpuTensor;
13068        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13069        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
13070        // verbatim, weight rows read once for all m columns, bit-identical per
13071        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
13072        // segments would re-read the weight per row" note described the grid.y=m lift,
13073        // which this twin deliberately is NOT.
13074        if !self.mmvq_supports(QT_NVFP4)
13075            || !self.uses_q8_1_fast(w0)
13076            || !self.uses_q8_1_fast(w1)
13077            || !self.uses_q8_1_fast(w2)
13078        {
13079            return Ok(None);
13080        }
13081        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
13082        // door — same family and bit-identity law as the fused4 delegate above.
13083        if (9..=16).contains(&m) {
13084            return Ok(
13085                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13086                    Some(mut ys) => {
13087                        let y2 = ys.pop().unwrap();
13088                        let y1 = ys.pop().unwrap();
13089                        let y0 = ys.pop().unwrap();
13090                        Some((y0, y1, y2))
13091                    }
13092                    None => None,
13093                },
13094            );
13095        }
13096        if !(1..=8).contains(&m) {
13097            return Ok(None);
13098        }
13099        if m > 1 {
13100            let in_f = w0.in_features();
13101            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13102                || !self.batched_supports(QT_NVFP4)
13103                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13104                || (m > 4 && !Self::b8_enabled())
13105                || in_f % 512 != 0
13106                || in_f / 64 > 272
13107            {
13108                return Ok(None);
13109            }
13110        }
13111        let unpack = |w: &crate::model::GpuTensor| match w {
13112            GpuTensor::Quant {
13113                bytes,
13114                qtype,
13115                scale,
13116                rp,
13117                ..
13118            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13119            _ => None,
13120        };
13121        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13122            return Ok(None);
13123        };
13124        let in_f = w0.in_features();
13125        if w1.in_features() != in_f || w2.in_features() != in_f {
13126            return Ok(None);
13127        }
13128        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13129        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13130        const RPW: u32 = 2;
13131        let rows_pb = ROWS_PER_BLOCK * RPW;
13132        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13133        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13134        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13135        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13136        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13137        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13138        // only dereferenced for the launch-arg build inside this call.
13139        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13140        if m > 1 {
13141            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
13142            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13143                return Ok(None);
13144            }
13145            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13146            let cfg = LaunchConfig {
13147                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13148                block_dim: (32, ROWS_PER_BLOCK, 1),
13149                shared_mem_bytes: 0,
13150            };
13151            let __s_b = self.gpu.stream();
13152            let mut b = __s_b.launch_builder(&f);
13153            b.arg(b0)
13154                .arg(b1)
13155                .arg(b2)
13156                .arg(aq)
13157                .arg(ad)
13158                .arg(&mut y0)
13159                .arg(&mut y1)
13160                .arg(&mut y2)
13161                .arg(&inf)
13162                .arg(&oi0)
13163                .arg(&oi1)
13164                .arg(&oi2)
13165                .arg(&mi);
13166            unsafe {
13167                b.launch(cfg)?;
13168            }
13169            return Ok(Some((y0, y1, y2)));
13170        }
13171        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13172        let cfg = LaunchConfig {
13173            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13174            block_dim: (32, ROWS_PER_BLOCK, 1),
13175            shared_mem_bytes: 0,
13176        };
13177        let __s_b = self.gpu.stream();
13178        let mut b = __s_b.launch_builder(&f);
13179        b.arg(b0)
13180            .arg(b1)
13181            .arg(b2)
13182            .arg(aq)
13183            .arg(ad)
13184            .arg(&mut y0)
13185            .arg(&mut y1)
13186            .arg(&mut y2)
13187            .arg(&inf)
13188            .arg(&oi0)
13189            .arg(&oi1)
13190            .arg(&oi2)
13191            .arg(&mi)
13192            .arg(&p0.1)
13193            .arg(&p1.1)
13194            .arg(&p2.1);
13195        unsafe {
13196            b.launch(cfg)?;
13197        }
13198        Ok(Some((y0, y1, y2)))
13199    }
13200
13201    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
13202    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
13203    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
13204    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
13205    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
13206    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
13207    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
13208    /// same-binary interleaved A/B arm.
13209    pub fn matmul_nvfp4_fused2(
13210        &self,
13211        w0: &crate::model::GpuTensor,
13212        w1: &crate::model::GpuTensor,
13213        aq: &CudaSlice<i8>,
13214        ad: &CudaSlice<f32>,
13215        m: usize,
13216    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13217        use crate::model::GpuTensor;
13218        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13219        let off =
13220            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13221        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
13222        // read serves all m rows); the fused segments would re-read the weight per row.
13223        if off
13224            || m != 1
13225            || !self.mmvq_supports(QT_NVFP4)
13226            || !self.uses_q8_1_fast(w0)
13227            || !self.uses_q8_1_fast(w1)
13228        {
13229            return Ok(None);
13230        }
13231        let unpack = |w: &crate::model::GpuTensor| match w {
13232            GpuTensor::Quant {
13233                bytes,
13234                qtype,
13235                scale,
13236                rp,
13237                ..
13238            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13239            _ => None,
13240        };
13241        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13242            return Ok(None);
13243        };
13244        let in_f = w0.in_features();
13245        if w1.in_features() != in_f {
13246            return Ok(None);
13247        }
13248        let (o0, o1) = (w0.out_features(), w1.out_features());
13249        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13250        const RPW: u32 = 2;
13251        let rows_pb = ROWS_PER_BLOCK * RPW;
13252        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13253        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13254        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13255        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13256        let cfg = LaunchConfig {
13257            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13258            block_dim: (32, ROWS_PER_BLOCK, 1),
13259            shared_mem_bytes: 0,
13260        };
13261        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13262        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13263        // only dereferenced for the launch-arg build inside this call.
13264        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13265        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
13266        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
13267        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13268            {
13269                use cudarc::driver::{DevicePtr, DevicePtrMut};
13270                let s = &self.gpu.stream();
13271                let (pw0, _g0) = b0.device_ptr(s);
13272                let (pw1, _g1) = b1.device_ptr(s);
13273                let (paq, _g2) = aq.device_ptr(s);
13274                let (pad, _g3) = ad.device_ptr(s);
13275                let (py0, _g4) = y0.device_ptr_mut(s);
13276                let (py1, _g5) = y1.device_ptr_mut(s);
13277                let (s0, s1) = (p0.1, p1.1);
13278                let mut ps = [
13279                    &pw0 as *const _ as *mut std::ffi::c_void,
13280                    &pw1 as *const _ as *mut _,
13281                    &paq as *const _ as *mut _,
13282                    &pad as *const _ as *mut _,
13283                    &py0 as *const _ as *mut _,
13284                    &py1 as *const _ as *mut _,
13285                    &inf as *const _ as *mut _,
13286                    &oi0 as *const _ as *mut _,
13287                    &oi1 as *const _ as *mut _,
13288                    &mi as *const _ as *mut _,
13289                    &s0 as *const _ as *mut _,
13290                    &s1 as *const _ as *mut _,
13291                ];
13292                unsafe {
13293                    self.launch_pdl(
13294                        "qmatvec_nvfp4_mmvq_fused2_rp",
13295                        cfg.grid_dim,
13296                        cfg.block_dim,
13297                        &mut ps,
13298                    )?;
13299                }
13300            }
13301            return Ok(Some((y0, y1)));
13302        }
13303        let __s_b = self.gpu.stream();
13304        let mut b = __s_b.launch_builder(&f);
13305        b.arg(b0)
13306            .arg(b1)
13307            .arg(aq)
13308            .arg(ad)
13309            .arg(&mut y0)
13310            .arg(&mut y1)
13311            .arg(&inf)
13312            .arg(&oi0)
13313            .arg(&oi1)
13314            .arg(&mi)
13315            .arg(&p0.1)
13316            .arg(&p1.1);
13317        unsafe {
13318            b.launch(cfg)?;
13319        }
13320        Ok(Some((y0, y1)))
13321    }
13322
13323    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
13324    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
13325    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
13326    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
13327    pub fn matmul_nvfp4_fused2_into(
13328        &self,
13329        w0: &crate::model::GpuTensor,
13330        w1: &crate::model::GpuTensor,
13331        aq: &CudaSlice<i8>,
13332        ad: &CudaSlice<f32>,
13333        y0: &mut CudaSlice<f32>,
13334        y1: &mut CudaSlice<f32>,
13335    ) -> Result<bool, Box<dyn std::error::Error>> {
13336        use crate::model::GpuTensor;
13337        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13338        let off =
13339            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13340        if off
13341            || !self.mmvq_supports(QT_NVFP4)
13342            || !self.uses_q8_1_fast(w0)
13343            || !self.uses_q8_1_fast(w1)
13344        {
13345            return Ok(false);
13346        }
13347        let unpack = |w: &crate::model::GpuTensor| match w {
13348            GpuTensor::Quant {
13349                bytes,
13350                qtype,
13351                scale,
13352                rp,
13353                ..
13354            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13355            _ => None,
13356        };
13357        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13358            return Ok(false);
13359        };
13360        let in_f = w0.in_features();
13361        if w1.in_features() != in_f {
13362            return Ok(false);
13363        }
13364        let (o0, o1) = (w0.out_features(), w1.out_features());
13365        if y0.len() < o0 || y1.len() < o1 {
13366            return Ok(false);
13367        }
13368        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13369        const RPW: u32 = 2;
13370        let rows_pb = ROWS_PER_BLOCK * RPW;
13371        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13372        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13373        let cfg = LaunchConfig {
13374            grid_dim: (nb(o0) + nb(o1), 1, 1),
13375            block_dim: (32, ROWS_PER_BLOCK, 1),
13376            shared_mem_bytes: 0,
13377        };
13378        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13379        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13380        // only dereferenced for the launch-arg build inside this call.
13381        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13382        let __s_b = self.gpu.stream();
13383        let mut b = __s_b.launch_builder(&f);
13384        b.arg(b0)
13385            .arg(b1)
13386            .arg(aq)
13387            .arg(ad)
13388            .arg(&mut *y0)
13389            .arg(&mut *y1)
13390            .arg(&inf)
13391            .arg(&oi0)
13392            .arg(&oi1)
13393            .arg(&mi)
13394            .arg(&p0.1)
13395            .arg(&p1.1);
13396        unsafe {
13397            b.launch(cfg)?;
13398        }
13399        Ok(true)
13400    }
13401
13402    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
13403    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
13404    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
13405    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
13406    #[allow(clippy::type_complexity)]
13407    pub fn matmul_nvfp4_fused4(
13408        &self,
13409        w0: &crate::model::GpuTensor,
13410        w1: &crate::model::GpuTensor,
13411        w2: &crate::model::GpuTensor,
13412        w3: &crate::model::GpuTensor,
13413        aq: &CudaSlice<i8>,
13414        ad: &CudaSlice<f32>,
13415        m: usize,
13416    ) -> Result<
13417        Option<(
13418            CudaSlice<f32>,
13419            CudaSlice<f32>,
13420            CudaSlice<f32>,
13421            CudaSlice<f32>,
13422        )>,
13423        Box<dyn std::error::Error>,
13424    > {
13425        use crate::model::GpuTensor;
13426        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
13427        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13428        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
13429        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
13430        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
13431        // Admission mirrors the singles' batched gates below.
13432        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13433            || !self.mmvq_supports(QT_NVFP4)
13434            || !self.uses_q8_1_fast(w0)
13435            || !self.uses_q8_1_fast(w1)
13436            || !self.uses_q8_1_fast(w2)
13437            || !self.uses_q8_1_fast(w3)
13438        {
13439            return Ok(None);
13440        }
13441        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
13442        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
13443        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
13444        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
13445        if (9..=16).contains(&m) {
13446            return Ok(
13447                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13448                    Some(mut ys) => {
13449                        let y3 = ys.pop().unwrap();
13450                        let y2 = ys.pop().unwrap();
13451                        let y1 = ys.pop().unwrap();
13452                        let y0 = ys.pop().unwrap();
13453                        Some((y0, y1, y2, y3))
13454                    }
13455                    None => None,
13456                },
13457            );
13458        }
13459        if !(1..=8).contains(&m) {
13460            return Ok(None);
13461        }
13462        if m > 1 {
13463            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
13464            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
13465            let in_f = w0.in_features();
13466            if !self.batched_supports(QT_NVFP4)
13467                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13468                || (m > 4 && !Self::b8_enabled())
13469                || in_f % 512 != 0
13470                || in_f / 64 > 272
13471            {
13472                return Ok(None);
13473            }
13474        }
13475        let unpack = |w: &crate::model::GpuTensor| match w {
13476            GpuTensor::Quant {
13477                bytes,
13478                qtype,
13479                scale,
13480                rp,
13481                ..
13482            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13483            _ => None,
13484        };
13485        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13486            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13487        else {
13488            return Ok(None);
13489        };
13490        let in_f = w0.in_features();
13491        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13492            return Ok(None);
13493        }
13494        let (o0, o1, o2, o3) = (
13495            w0.out_features(),
13496            w1.out_features(),
13497            w2.out_features(),
13498            w3.out_features(),
13499        );
13500        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13501        const RPW: u32 = 2;
13502        let rows_pb = ROWS_PER_BLOCK * RPW;
13503        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13504        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13505        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13506        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13507        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13508        let (inf, oi0, oi1, oi2, oi3, mi) = (
13509            in_f as i32,
13510            o0 as i32,
13511            o1 as i32,
13512            o2 as i32,
13513            o3 as i32,
13514            m as i32,
13515        );
13516        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13517        // only dereferenced for the launch-arg build inside this call.
13518        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13519        if m > 1 {
13520            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
13521            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
13522            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13523                return Ok(None);
13524            }
13525            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13526            let cfg = LaunchConfig {
13527                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13528                block_dim: (32, ROWS_PER_BLOCK, 1),
13529                shared_mem_bytes: 0,
13530            };
13531            let __s_b = self.gpu.stream();
13532            let mut b = __s_b.launch_builder(&f);
13533            b.arg(b0)
13534                .arg(b1)
13535                .arg(b2)
13536                .arg(b3)
13537                .arg(aq)
13538                .arg(ad)
13539                .arg(&mut y0)
13540                .arg(&mut y1)
13541                .arg(&mut y2)
13542                .arg(&mut y3)
13543                .arg(&inf)
13544                .arg(&oi0)
13545                .arg(&oi1)
13546                .arg(&oi2)
13547                .arg(&oi3)
13548                .arg(&mi);
13549            unsafe {
13550                b.launch(cfg)?;
13551            }
13552            return Ok(Some((y0, y1, y2, y3)));
13553        }
13554        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13555        let cfg = LaunchConfig {
13556            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13557            block_dim: (32, ROWS_PER_BLOCK, 1),
13558            shared_mem_bytes: 0,
13559        };
13560        let __s_b = self.gpu.stream();
13561        let mut b = __s_b.launch_builder(&f);
13562        b.arg(b0)
13563            .arg(b1)
13564            .arg(b2)
13565            .arg(b3)
13566            .arg(aq)
13567            .arg(ad)
13568            .arg(&mut y0)
13569            .arg(&mut y1)
13570            .arg(&mut y2)
13571            .arg(&mut y3)
13572            .arg(&inf)
13573            .arg(&oi0)
13574            .arg(&oi1)
13575            .arg(&oi2)
13576            .arg(&oi3)
13577            .arg(&mi)
13578            .arg(&p0.1)
13579            .arg(&p1.1)
13580            .arg(&p2.1)
13581            .arg(&p3.1);
13582        unsafe {
13583            b.launch(cfg)?;
13584        }
13585        Ok(Some((y0, y1, y2, y3)))
13586    }
13587
13588    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
13589    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
13590    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
13591    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
13592    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
13593    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
13594    /// back to the per-tensor path.
13595    pub fn matmul_q8_fused2(
13596        &self,
13597        w0: &crate::model::GpuTensor,
13598        w1: &crate::model::GpuTensor,
13599        aq: &CudaSlice<i8>,
13600        ad: &CudaSlice<f32>,
13601    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13602        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
13603        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
13604        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
13605        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
13606        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
13607        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13608            return Ok(Some(self.e4m3_fused2_core(
13609                p0.0,
13610                p1.0,
13611                aq,
13612                ad,
13613                w0.in_features(),
13614                p0.1,
13615                p1.1,
13616                p0.2,
13617                p0.3,
13618                p1.3,
13619            )?));
13620        }
13621        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13622            return Ok(None);
13623        };
13624        Ok(Some(self.q8_fused2_core(
13625            p0.0,
13626            p1.0,
13627            aq,
13628            ad,
13629            w0.in_features(),
13630            p0.1,
13631            p1.1,
13632            p0.2,
13633        )?))
13634    }
13635
13636    #[allow(clippy::too_many_arguments)]
13637    fn q8_fused2_core(
13638        &self,
13639        b0: &CudaSlice<u8>,
13640        b1: &CudaSlice<u8>,
13641        aq: &CudaSlice<i8>,
13642        ad: &CudaSlice<f32>,
13643        in_f: usize,
13644        out0: usize,
13645        out1: usize,
13646        row_bytes: usize,
13647    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13648        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13649        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13650        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13651        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13652        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13653        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13654        let cfg = LaunchConfig {
13655            grid_dim: (nb0 + nb1, 1, 1),
13656            block_dim: (32, ROWS_PER_BLOCK, 1),
13657            shared_mem_bytes: 0,
13658        };
13659        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13660        let __s_b = self.gpu.stream();
13661        let mut b = __s_b.launch_builder(&f);
13662        b.arg(b0)
13663            .arg(b1)
13664            .arg(aq)
13665            .arg(ad)
13666            .arg(&mut y0)
13667            .arg(&mut y1)
13668            .arg(&inf)
13669            .arg(&o0)
13670            .arg(&o1)
13671            .arg(&rbl);
13672        unsafe {
13673            b.launch(cfg)?;
13674        }
13675        Ok((y0, y1))
13676    }
13677
13678    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13679    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13680    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13681    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13682    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13683    pub fn matmul_q8_fused2_x(
13684        &self,
13685        w0: &crate::model::GpuTensor,
13686        w1: &crate::model::GpuTensor,
13687        x: &CudaSlice<f32>,
13688    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13689        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13690            return Ok(None);
13691        }
13692        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13693            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13694            return Ok(Some(self.e4m3_fused2_core(
13695                p0.0,
13696                p1.0,
13697                &aq,
13698                &ad,
13699                w0.in_features(),
13700                p0.1,
13701                p1.1,
13702                p0.2,
13703                p0.3,
13704                p1.3,
13705            )?));
13706        }
13707        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13708            return Ok(None);
13709        };
13710        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13711        Ok(Some(self.q8_fused2_core(
13712            p0.0,
13713            p1.0,
13714            &aq,
13715            &ad,
13716            w0.in_features(),
13717            p0.1,
13718            p1.1,
13719            p0.2,
13720        )?))
13721    }
13722
13723    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13724    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13725    #[allow(clippy::too_many_arguments)]
13726    pub fn qmatvec_q8_fused2_raw(
13727        &self,
13728        b0: &CudaSlice<u8>,
13729        b1: &CudaSlice<u8>,
13730        x: &CudaSlice<f32>,
13731        in_f: usize,
13732        out0: usize,
13733        out1: usize,
13734        row_bytes: usize,
13735    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13736        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13737        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13738    }
13739
13740    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13741    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13742    /// (tensor,row) to three separate m=1 MMVQ launches.
13743    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13744    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13745    pub fn matmul_q4_fused3(
13746        &self,
13747        w0: &crate::model::GpuTensor,
13748        w1: &crate::model::GpuTensor,
13749        w2: &crate::model::GpuTensor,
13750        aq: &CudaSlice<i8>,
13751        ad: &CudaSlice<f32>,
13752    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13753    {
13754        use crate::model::GpuTensor;
13755        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13756            match w {
13757                GpuTensor::Quant {
13758                    qtype, row_bytes, ..
13759                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13760                _ => None,
13761            }
13762        };
13763        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13764            return Ok(None);
13765        };
13766        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13767            return Ok(None);
13768        }
13769        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
13770        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
13771        // the separate matvecs (each routes its own rp).
13772        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13773            match w {
13774                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13775                    Some(m) => (m, true),
13776                    None => (bytes, *rp),
13777                },
13778                _ => unreachable!(),
13779            }
13780        }
13781        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13782        if rp0 != rp1 || rp1 != rp2 {
13783            return Ok(None);
13784        }
13785        let rp = rp0;
13786        let rpb: u32 = 4;
13787        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
13788        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
13789        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
13790        let mr1 = rp && Self::q40_mr1_on();
13791        let nb = |o: usize| {
13792            if mr1 {
13793                (o as u32).div_ceil(rpb)
13794            } else {
13795                (o as u32).div_ceil(2).div_ceil(rpb)
13796            }
13797        };
13798        let grid = nb(o0) + nb(o1) + nb(o2);
13799        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13800        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13801        let mut y2 = self.alloc_uninit::<f32>(o2)?;
13802        let f = self.func(if mr1 {
13803            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13804        } else if rp {
13805            "qmatvec_q4_0_mmvq_fused3_rp"
13806        } else {
13807            "qmatvec_q4_0_mmvq_fused3"
13808        });
13809        let cfg = LaunchConfig {
13810            grid_dim: (grid, 1, 1),
13811            block_dim: (32, rpb, 1),
13812            shared_mem_bytes: 0,
13813        };
13814        let inf = w0.in_features() as i32;
13815        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13816        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13817        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
13818        // variant may take the programmatic-serialization launch.
13819        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13820            {
13821                use cudarc::driver::{DevicePtr, DevicePtrMut};
13822                let s = &self.gpu.stream();
13823                let (p0, _g0) = b0.device_ptr(s);
13824                let (p1, _g1) = b1.device_ptr(s);
13825                let (p2, _g2) = b2.device_ptr(s);
13826                let (paq, _g3) = aq.device_ptr(s);
13827                let (pad, _g4) = ad.device_ptr(s);
13828                let (py0, _g5) = y0.device_ptr_mut(s);
13829                let (py1, _g6) = y1.device_ptr_mut(s);
13830                let (py2, _g7) = y2.device_ptr_mut(s);
13831                let mut ps = [
13832                    &p0 as *const _ as *mut std::ffi::c_void,
13833                    &p1 as *const _ as *mut _,
13834                    &p2 as *const _ as *mut _,
13835                    &paq as *const _ as *mut _,
13836                    &pad as *const _ as *mut _,
13837                    &py0 as *const _ as *mut _,
13838                    &py1 as *const _ as *mut _,
13839                    &py2 as *const _ as *mut _,
13840                    &inf as *const _ as *mut _,
13841                    &oo0 as *const _ as *mut _,
13842                    &oo1 as *const _ as *mut _,
13843                    &oo2 as *const _ as *mut _,
13844                    &r0 as *const _ as *mut _,
13845                    &r1 as *const _ as *mut _,
13846                    &r2 as *const _ as *mut _,
13847                ];
13848                unsafe {
13849                    self.launch_pdl(
13850                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13851                        (grid, 1, 1),
13852                        (32, rpb, 1),
13853                        &mut ps,
13854                    )?;
13855                }
13856            }
13857            return Ok(Some((y0, y1, y2)));
13858        }
13859        let __s_b = self.gpu.stream();
13860        let mut b = __s_b.launch_builder(&f);
13861        b.arg(b0)
13862            .arg(b1)
13863            .arg(b2)
13864            .arg(aq)
13865            .arg(ad)
13866            .arg(&mut y0)
13867            .arg(&mut y1)
13868            .arg(&mut y2)
13869            .arg(&inf)
13870            .arg(&oo0)
13871            .arg(&oo1)
13872            .arg(&oo2)
13873            .arg(&r0)
13874            .arg(&r1)
13875            .arg(&r2);
13876        unsafe {
13877            b.launch(cfg)?;
13878        }
13879        Ok(Some((y0, y1, y2)))
13880    }
13881
13882    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13883    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
13884    #[allow(clippy::too_many_arguments)]
13885    pub fn matmul_q4_fused3_into(
13886        &self,
13887        w0: &crate::model::GpuTensor,
13888        w1: &crate::model::GpuTensor,
13889        w2: &crate::model::GpuTensor,
13890        aq: &CudaSlice<i8>,
13891        ad: &CudaSlice<f32>,
13892        y0: &mut CudaSlice<f32>,
13893        y1: &mut CudaSlice<f32>,
13894        y2: &mut CudaSlice<f32>,
13895    ) -> Result<bool, Box<dyn std::error::Error>> {
13896        use crate::model::GpuTensor;
13897        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13898            match w {
13899                GpuTensor::Quant {
13900                    qtype, row_bytes, ..
13901                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13902                _ => None,
13903            }
13904        };
13905        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13906            return Ok(false);
13907        };
13908        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13909            return Ok(false);
13910        }
13911        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13912            match w {
13913                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13914                    Some(m) => (m, true),
13915                    None => (bytes, *rp),
13916                },
13917                _ => unreachable!(),
13918            }
13919        }
13920        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13921        if rp0 != rp1 || rp1 != rp2 {
13922            return Ok(false);
13923        }
13924        let rp = rp0;
13925        let rpb: u32 = 4;
13926        let mr1 = rp && Self::q40_mr1_on();
13927        let nb = |o: usize| {
13928            if mr1 {
13929                (o as u32).div_ceil(rpb)
13930            } else {
13931                (o as u32).div_ceil(2).div_ceil(rpb)
13932            }
13933        };
13934        let grid = nb(o0) + nb(o1) + nb(o2);
13935        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13936        let f = self.func(if mr1 {
13937            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13938        } else if rp {
13939            "qmatvec_q4_0_mmvq_fused3_rp"
13940        } else {
13941            "qmatvec_q4_0_mmvq_fused3"
13942        });
13943        let cfg = LaunchConfig {
13944            grid_dim: (grid, 1, 1),
13945            block_dim: (32, rpb, 1),
13946            shared_mem_bytes: 0,
13947        };
13948        let inf = w0.in_features() as i32;
13949        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13950        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13951        // PDL wave-A: identical to the owned twin (capture-lane parity).
13952        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13953            use cudarc::driver::{DevicePtr, DevicePtrMut};
13954            let s = &self.gpu.stream();
13955            let (p0, _g0) = b0.device_ptr(s);
13956            let (p1, _g1) = b1.device_ptr(s);
13957            let (p2, _g2) = b2.device_ptr(s);
13958            let (paq, _g3) = aq.device_ptr(s);
13959            let (pad, _g4) = ad.device_ptr(s);
13960            let (py0, _g5) = y0.device_ptr_mut(s);
13961            let (py1, _g6) = y1.device_ptr_mut(s);
13962            let (py2, _g7) = y2.device_ptr_mut(s);
13963            let mut ps = [
13964                &p0 as *const _ as *mut std::ffi::c_void,
13965                &p1 as *const _ as *mut _,
13966                &p2 as *const _ as *mut _,
13967                &paq as *const _ as *mut _,
13968                &pad as *const _ as *mut _,
13969                &py0 as *const _ as *mut _,
13970                &py1 as *const _ as *mut _,
13971                &py2 as *const _ as *mut _,
13972                &inf as *const _ as *mut _,
13973                &oo0 as *const _ as *mut _,
13974                &oo1 as *const _ as *mut _,
13975                &oo2 as *const _ as *mut _,
13976                &r0 as *const _ as *mut _,
13977                &r1 as *const _ as *mut _,
13978                &r2 as *const _ as *mut _,
13979            ];
13980            unsafe {
13981                self.launch_pdl(
13982                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13983                    (grid, 1, 1),
13984                    (32, rpb, 1),
13985                    &mut ps,
13986                )?;
13987            }
13988            return Ok(true);
13989        }
13990        let __s_b = self.gpu.stream();
13991        let mut b = __s_b.launch_builder(&f);
13992        b.arg(b0)
13993            .arg(b1)
13994            .arg(b2)
13995            .arg(aq)
13996            .arg(ad)
13997            .arg(&mut *y0)
13998            .arg(&mut *y1)
13999            .arg(&mut *y2)
14000            .arg(&inf)
14001            .arg(&oo0)
14002            .arg(&oo1)
14003            .arg(&oo2)
14004            .arg(&r0)
14005            .arg(&r1)
14006            .arg(&r2);
14007        unsafe {
14008            b.launch(cfg)?;
14009        }
14010        Ok(true)
14011    }
14012
14013    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
14014    pub fn matmul_q4_fused2(
14015        &self,
14016        w0: &crate::model::GpuTensor,
14017        w1: &crate::model::GpuTensor,
14018        aq: &CudaSlice<i8>,
14019        ad: &CudaSlice<f32>,
14020    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14021        use crate::model::GpuTensor;
14022        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14023            match w {
14024                GpuTensor::Quant {
14025                    qtype, row_bytes, ..
14026                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14027                _ => None,
14028            }
14029        };
14030        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14031            return Ok(None);
14032        };
14033        if w0.in_features() != w1.in_features() {
14034            return Ok(None);
14035        }
14036        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
14037        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14038            match w {
14039                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14040                    Some(m) => (m, true),
14041                    None => (bytes, *rp),
14042                },
14043                _ => unreachable!(),
14044            }
14045        }
14046        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14047        if rp0 != rp1 {
14048            return Ok(None);
14049        }
14050        let rp = rp0;
14051        let rpb: u32 = 4;
14052        // mr1 twin — see matmul_q4_fused3.
14053        let mr1 = rp && Self::q40_mr1_on();
14054        let nb = |o: usize| {
14055            if mr1 {
14056                (o as u32).div_ceil(rpb)
14057            } else {
14058                (o as u32).div_ceil(2).div_ceil(rpb)
14059            }
14060        };
14061        let grid = nb(o0) + nb(o1);
14062        let mut y0 = self.alloc_uninit::<f32>(o0)?;
14063        let mut y1 = self.alloc_uninit::<f32>(o1)?;
14064        let f = self.func(if mr1 {
14065            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14066        } else if rp {
14067            "qmatvec_q4_0_mmvq_fused2_rp"
14068        } else {
14069            "qmatvec_q4_0_mmvq_fused2"
14070        });
14071        let cfg = LaunchConfig {
14072            grid_dim: (grid, 1, 1),
14073            block_dim: (32, rpb, 1),
14074            shared_mem_bytes: 0,
14075        };
14076        let inf = w0.in_features() as i32;
14077        let (oo0, oo1) = (o0 as i32, o1 as i32);
14078        let (r0, r1) = (rb0 as i64, rb1 as i64);
14079        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
14080        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14081            {
14082                use cudarc::driver::{DevicePtr, DevicePtrMut};
14083                let s = &self.gpu.stream();
14084                let (p0, _g0) = b0.device_ptr(s);
14085                let (p1, _g1) = b1.device_ptr(s);
14086                let (paq, _g2) = aq.device_ptr(s);
14087                let (pad, _g3) = ad.device_ptr(s);
14088                let (py0, _g4) = y0.device_ptr_mut(s);
14089                let (py1, _g5) = y1.device_ptr_mut(s);
14090                let mut ps = [
14091                    &p0 as *const _ as *mut std::ffi::c_void,
14092                    &p1 as *const _ as *mut _,
14093                    &paq as *const _ as *mut _,
14094                    &pad as *const _ as *mut _,
14095                    &py0 as *const _ as *mut _,
14096                    &py1 as *const _ as *mut _,
14097                    &inf as *const _ as *mut _,
14098                    &oo0 as *const _ as *mut _,
14099                    &oo1 as *const _ as *mut _,
14100                    &r0 as *const _ as *mut _,
14101                    &r1 as *const _ as *mut _,
14102                ];
14103                unsafe {
14104                    self.launch_pdl(
14105                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14106                        (grid, 1, 1),
14107                        (32, rpb, 1),
14108                        &mut ps,
14109                    )?;
14110                }
14111            }
14112            return Ok(Some((y0, y1)));
14113        }
14114        let __s_b = self.gpu.stream();
14115        let mut b = __s_b.launch_builder(&f);
14116        b.arg(b0)
14117            .arg(b1)
14118            .arg(aq)
14119            .arg(ad)
14120            .arg(&mut y0)
14121            .arg(&mut y1)
14122            .arg(&inf)
14123            .arg(&oo0)
14124            .arg(&oo1)
14125            .arg(&r0)
14126            .arg(&r1);
14127        unsafe {
14128            b.launch(cfg)?;
14129        }
14130        Ok(Some((y0, y1)))
14131    }
14132
14133    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
14134    pub fn matmul_q4_fused2_into(
14135        &self,
14136        w0: &crate::model::GpuTensor,
14137        w1: &crate::model::GpuTensor,
14138        aq: &CudaSlice<i8>,
14139        ad: &CudaSlice<f32>,
14140        y0: &mut CudaSlice<f32>,
14141        y1: &mut CudaSlice<f32>,
14142    ) -> Result<bool, Box<dyn std::error::Error>> {
14143        use crate::model::GpuTensor;
14144        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14145            match w {
14146                GpuTensor::Quant {
14147                    qtype, row_bytes, ..
14148                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14149                _ => None,
14150            }
14151        };
14152        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14153            return Ok(false);
14154        };
14155        if w0.in_features() != w1.in_features() {
14156            return Ok(false);
14157        }
14158        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14159            match w {
14160                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14161                    Some(m) => (m, true),
14162                    None => (bytes, *rp),
14163                },
14164                _ => unreachable!(),
14165            }
14166        }
14167        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14168        if rp0 != rp1 {
14169            return Ok(false);
14170        }
14171        let rp = rp0;
14172        let rpb: u32 = 4;
14173        let mr1 = rp && Self::q40_mr1_on();
14174        let nb = |o: usize| {
14175            if mr1 {
14176                (o as u32).div_ceil(rpb)
14177            } else {
14178                (o as u32).div_ceil(2).div_ceil(rpb)
14179            }
14180        };
14181        let grid = nb(o0) + nb(o1);
14182        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14183        let f = self.func(if mr1 {
14184            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14185        } else if rp {
14186            "qmatvec_q4_0_mmvq_fused2_rp"
14187        } else {
14188            "qmatvec_q4_0_mmvq_fused2"
14189        });
14190        let cfg = LaunchConfig {
14191            grid_dim: (grid, 1, 1),
14192            block_dim: (32, rpb, 1),
14193            shared_mem_bytes: 0,
14194        };
14195        let inf = w0.in_features() as i32;
14196        let (oo0, oo1) = (o0 as i32, o1 as i32);
14197        let (r0, r1) = (rb0 as i64, rb1 as i64);
14198        // PDL wave-A: identical to the owned twin (capture-lane parity).
14199        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14200            use cudarc::driver::{DevicePtr, DevicePtrMut};
14201            let s = &self.gpu.stream();
14202            let (p0, _g0) = b0.device_ptr(s);
14203            let (p1, _g1) = b1.device_ptr(s);
14204            let (paq, _g2) = aq.device_ptr(s);
14205            let (pad, _g3) = ad.device_ptr(s);
14206            let (py0, _g4) = y0.device_ptr_mut(s);
14207            let (py1, _g5) = y1.device_ptr_mut(s);
14208            let mut ps = [
14209                &p0 as *const _ as *mut std::ffi::c_void,
14210                &p1 as *const _ as *mut _,
14211                &paq as *const _ as *mut _,
14212                &pad as *const _ as *mut _,
14213                &py0 as *const _ as *mut _,
14214                &py1 as *const _ as *mut _,
14215                &inf as *const _ as *mut _,
14216                &oo0 as *const _ as *mut _,
14217                &oo1 as *const _ as *mut _,
14218                &r0 as *const _ as *mut _,
14219                &r1 as *const _ as *mut _,
14220            ];
14221            unsafe {
14222                self.launch_pdl(
14223                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14224                    (grid, 1, 1),
14225                    (32, rpb, 1),
14226                    &mut ps,
14227                )?;
14228            }
14229            return Ok(true);
14230        }
14231        let __s_b = self.gpu.stream();
14232        let mut b = __s_b.launch_builder(&f);
14233        b.arg(b0)
14234            .arg(b1)
14235            .arg(aq)
14236            .arg(ad)
14237            .arg(&mut *y0)
14238            .arg(&mut *y1)
14239            .arg(&inf)
14240            .arg(&oo0)
14241            .arg(&oo1)
14242            .arg(&r0)
14243            .arg(&r1);
14244        unsafe {
14245            b.launch(cfg)?;
14246        }
14247        Ok(true)
14248    }
14249
14250    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
14251    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
14252    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
14253    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
14254    pub fn matmul_q4_fused2_batched(
14255        &self,
14256        w0: &crate::model::GpuTensor,
14257        w1: &crate::model::GpuTensor,
14258        aq: &CudaSlice<i8>,
14259        ad: &CudaSlice<f32>,
14260        m: usize,
14261    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14262        use crate::model::GpuTensor;
14263        if m < 2 || m > 8 {
14264            return Ok(None);
14265        }
14266        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14267            match w {
14268                GpuTensor::Quant {
14269                    qtype, row_bytes, ..
14270                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14271                _ => None,
14272            }
14273        };
14274        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14275            return Ok(None);
14276        };
14277        if w0.in_features() != w1.in_features() {
14278            return Ok(None);
14279        }
14280        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14281            match w {
14282                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14283                    Some(mr) => (mr, true),
14284                    None => (bytes, *rp),
14285                },
14286                _ => unreachable!(),
14287            }
14288        }
14289        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14290        if !rp0 || !rp1 {
14291            return Ok(None);
14292        }
14293        let mcols = Self::batched_mcols(m);
14294        let rpb: u32 = 4;
14295        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14296        let grid = nb(o0) + nb(o1);
14297        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14298        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14299        let f = self.func(match mcols {
14300            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14301            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14302            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14303        });
14304        let cfg = LaunchConfig {
14305            grid_dim: (grid, 1, 1),
14306            block_dim: (32, rpb, 1),
14307            shared_mem_bytes: 0,
14308        };
14309        let inf = w0.in_features() as i32;
14310        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14311        let rb = rb0 as i64;
14312        let __s_b = self.gpu.stream();
14313        let mut b = __s_b.launch_builder(&f);
14314        b.arg(b0)
14315            .arg(b1)
14316            .arg(aq)
14317            .arg(ad)
14318            .arg(&mut y0)
14319            .arg(&mut y1)
14320            .arg(&inf)
14321            .arg(&oo0)
14322            .arg(&oo1)
14323            .arg(&mi)
14324            .arg(&rb);
14325        unsafe {
14326            b.launch(cfg)?;
14327        }
14328        Ok(Some((y0, y1)))
14329    }
14330
14331    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
14332    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
14333    #[allow(clippy::too_many_arguments)]
14334    pub fn matmul_q4_fused3_batched(
14335        &self,
14336        w0: &crate::model::GpuTensor,
14337        w1: &crate::model::GpuTensor,
14338        w2: &crate::model::GpuTensor,
14339        aq: &CudaSlice<i8>,
14340        ad: &CudaSlice<f32>,
14341        m: usize,
14342    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14343    {
14344        use crate::model::GpuTensor;
14345        if m < 2 || m > 8 {
14346            return Ok(None);
14347        }
14348        let q4 = |w: &GpuTensor| -> Option<usize> {
14349            match w {
14350                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14351                _ => None,
14352            }
14353        };
14354        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14355            return Ok(None);
14356        };
14357        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14358            return Ok(None);
14359        }
14360        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14361            match w {
14362                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14363                    Some(mr) => (mr, true),
14364                    None => (bytes, *rp),
14365                },
14366                _ => unreachable!(),
14367            }
14368        }
14369        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14370        if !rp0 || !rp1 || !rp2 {
14371            return Ok(None);
14372        }
14373        let mcols = Self::batched_mcols(m);
14374        let rpb: u32 = 4;
14375        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14376        let grid = nb(o0) + nb(o1) + nb(o2);
14377        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14378        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14379        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14380        let f = self.func(match mcols {
14381            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14382            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14383            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14384        });
14385        let cfg = LaunchConfig {
14386            grid_dim: (grid, 1, 1),
14387            block_dim: (32, rpb, 1),
14388            shared_mem_bytes: 0,
14389        };
14390        let inf = w0.in_features() as i32;
14391        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14392        let rb = 0i64;
14393        let __s_b = self.gpu.stream();
14394        let mut b = __s_b.launch_builder(&f);
14395        b.arg(b0)
14396            .arg(b1)
14397            .arg(b2)
14398            .arg(aq)
14399            .arg(ad)
14400            .arg(&mut y0)
14401            .arg(&mut y1)
14402            .arg(&mut y2)
14403            .arg(&inf)
14404            .arg(&oo0)
14405            .arg(&oo1)
14406            .arg(&oo2)
14407            .arg(&mi)
14408            .arg(&rb);
14409        unsafe {
14410            b.launch(cfg)?;
14411        }
14412        Ok(Some((y0, y1, y2)))
14413    }
14414
14415    pub fn matmul_q8_fused3(
14416        &self,
14417        w0: &crate::model::GpuTensor,
14418        w1: &crate::model::GpuTensor,
14419        w2: &crate::model::GpuTensor,
14420        aq: &CudaSlice<i8>,
14421        ad: &CudaSlice<f32>,
14422    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14423    {
14424        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
14425        // are per-tensor FP8, so native residency without this arm meant three separate launches.
14426        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14427            return Ok(Some(self.e4m3_fused3_core(
14428                p0.0,
14429                p1.0,
14430                p2.0,
14431                aq,
14432                ad,
14433                w0.in_features(),
14434                p0.1,
14435                p1.1,
14436                p2.1,
14437                p0.2,
14438                p0.3,
14439                p1.3,
14440                p2.3,
14441            )?));
14442        }
14443        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14444            return Ok(None);
14445        };
14446        Ok(Some(self.q8_fused3_core(
14447            p0.0,
14448            p1.0,
14449            p2.0,
14450            aq,
14451            ad,
14452            w0.in_features(),
14453            p0.1,
14454            p1.1,
14455            p2.1,
14456            p0.2,
14457        )?))
14458    }
14459
14460    #[allow(clippy::too_many_arguments)]
14461    fn q8_fused3_core(
14462        &self,
14463        b0: &CudaSlice<u8>,
14464        b1: &CudaSlice<u8>,
14465        b2: &CudaSlice<u8>,
14466        aq: &CudaSlice<i8>,
14467        ad: &CudaSlice<f32>,
14468        in_f: usize,
14469        out0: usize,
14470        out1: usize,
14471        out2: usize,
14472        row_bytes: usize,
14473    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14474        const ROWS_PER_BLOCK: u32 = 4;
14475        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14476        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14477        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14478        let f = self.func("qmatvec_q8_0_mmvq_fused3");
14479        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14480        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14481        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14482        let cfg = LaunchConfig {
14483            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14484            block_dim: (32, ROWS_PER_BLOCK, 1),
14485            shared_mem_bytes: 0,
14486        };
14487        let (inf, o0, o1, o2, rbl) = (
14488            in_f as i32,
14489            out0 as i32,
14490            out1 as i32,
14491            out2 as i32,
14492            row_bytes as i64,
14493        );
14494        let __s_b = self.gpu.stream();
14495        let mut b = __s_b.launch_builder(&f);
14496        b.arg(b0)
14497            .arg(b1)
14498            .arg(b2)
14499            .arg(aq)
14500            .arg(ad)
14501            .arg(&mut y0)
14502            .arg(&mut y1)
14503            .arg(&mut y2)
14504            .arg(&inf)
14505            .arg(&o0)
14506            .arg(&o1)
14507            .arg(&o2)
14508            .arg(&rbl);
14509        unsafe {
14510            b.launch(cfg)?;
14511        }
14512        Ok((y0, y1, y2))
14513    }
14514
14515    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
14516    #[allow(clippy::too_many_arguments)]
14517    pub fn qmatvec_q8_fused3_raw(
14518        &self,
14519        b0: &CudaSlice<u8>,
14520        b1: &CudaSlice<u8>,
14521        b2: &CudaSlice<u8>,
14522        x: &CudaSlice<f32>,
14523        in_f: usize,
14524        out0: usize,
14525        out1: usize,
14526        out2: usize,
14527        row_bytes: usize,
14528    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14529        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14530        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14531    }
14532
14533    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
14534    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
14535    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
14536    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
14537    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
14538    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
14539    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
14540    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
14541    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
14542    /// twin must not introduce a batched program the reference path would not run).
14543    pub fn matmul_q8_fused2_t(
14544        &self,
14545        w0: &crate::model::GpuTensor,
14546        w1: &crate::model::GpuTensor,
14547        aq: &CudaSlice<i8>,
14548        ad: &CudaSlice<f32>,
14549        m: usize,
14550    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14551        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
14552        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
14553        // fuses too — same template body, still bit-identical to the two _b8 launches.
14554        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14555            return Ok(None);
14556        }
14557        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
14558        // so the fused b8 launch would introduce a batched program the reference path would not run.
14559        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14560            if m > 4 && !Self::b8_enabled() {
14561                return Ok(None);
14562            }
14563            return Ok(Some(self.e4m3_fused2_t_core(
14564                p0.0,
14565                p1.0,
14566                aq,
14567                ad,
14568                m,
14569                w0.in_features(),
14570                p0.1,
14571                p1.1,
14572                p0.2,
14573                p0.3,
14574                p1.3,
14575            )?));
14576        }
14577        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14578            return Ok(None);
14579        };
14580        Ok(Some(self.q8_fused2_t_core(
14581            p0.0,
14582            p1.0,
14583            aq,
14584            ad,
14585            m,
14586            w0.in_features(),
14587            p0.1,
14588            p1.1,
14589            p0.2,
14590        )?))
14591    }
14592
14593    #[allow(clippy::too_many_arguments)]
14594    fn q8_fused2_t_core(
14595        &self,
14596        b0: &CudaSlice<u8>,
14597        b1: &CudaSlice<u8>,
14598        aq: &CudaSlice<i8>,
14599        ad: &CudaSlice<f32>,
14600        m: usize,
14601        in_f: usize,
14602        out0: usize,
14603        out1: usize,
14604        row_bytes: usize,
14605    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14606        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14607        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14608        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14609        let f = self.func(match Self::batched_mcols(m) {
14610            2 => "qmatvec_q8_0_mmvq_fused2_b2",
14611            4 => "qmatvec_q8_0_mmvq_fused2_b4",
14612            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
14613            _ => "qmatvec_q8_0_mmvq_fused2_b8",
14614        });
14615        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14616        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14617        let cfg = LaunchConfig {
14618            grid_dim: (nb0 + nb1, 1, 1),
14619            block_dim: (32, ROWS_PER_BLOCK, 1),
14620            shared_mem_bytes: 0,
14621        };
14622        let (inf, o0, o1, mi, rbl) = (
14623            in_f as i32,
14624            out0 as i32,
14625            out1 as i32,
14626            m as i32,
14627            row_bytes as i64,
14628        );
14629        let __s_b = self.gpu.stream();
14630        let mut b = __s_b.launch_builder(&f);
14631        b.arg(b0)
14632            .arg(b1)
14633            .arg(aq)
14634            .arg(ad)
14635            .arg(&mut y0)
14636            .arg(&mut y1)
14637            .arg(&inf)
14638            .arg(&o0)
14639            .arg(&o1)
14640            .arg(&mi)
14641            .arg(&rbl);
14642        unsafe {
14643            b.launch(cfg)?;
14644        }
14645        Ok((y0, y1))
14646    }
14647
14648    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14649    /// q8_1 quant of the [m, in_f] activation), no env gating.
14650    #[allow(clippy::too_many_arguments)]
14651    pub fn qmatvec_q8_fused2_t_raw(
14652        &self,
14653        b0: &CudaSlice<u8>,
14654        b1: &CudaSlice<u8>,
14655        x: &CudaSlice<f32>,
14656        m: usize,
14657        in_f: usize,
14658        out0: usize,
14659        out1: usize,
14660        row_bytes: usize,
14661    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14662        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14663        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14664    }
14665
14666    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14667    /// `matmul_q8_fused2_t` with three ranges.
14668    #[allow(clippy::too_many_arguments)]
14669    pub fn matmul_q8_fused3_t(
14670        &self,
14671        w0: &crate::model::GpuTensor,
14672        w1: &crate::model::GpuTensor,
14673        w2: &crate::model::GpuTensor,
14674        aq: &CudaSlice<i8>,
14675        ad: &CudaSlice<f32>,
14676        m: usize,
14677    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14678    {
14679        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14680            return Ok(None);
14681        }
14682        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14683            return Ok(Some(self.e4m3_fused3_t_core(
14684                p0.0,
14685                p1.0,
14686                p2.0,
14687                aq,
14688                ad,
14689                m,
14690                w0.in_features(),
14691                p0.1,
14692                p1.1,
14693                p2.1,
14694                p0.2,
14695                p0.3,
14696                p1.3,
14697                p2.3,
14698            )?));
14699        }
14700        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14701            return Ok(None);
14702        };
14703        Ok(Some(self.q8_fused3_t_core(
14704            p0.0,
14705            p1.0,
14706            p2.0,
14707            aq,
14708            ad,
14709            m,
14710            w0.in_features(),
14711            p0.1,
14712            p1.1,
14713            p2.1,
14714            p0.2,
14715        )?))
14716    }
14717
14718    #[allow(clippy::too_many_arguments)]
14719    fn q8_fused3_t_core(
14720        &self,
14721        b0: &CudaSlice<u8>,
14722        b1: &CudaSlice<u8>,
14723        b2: &CudaSlice<u8>,
14724        aq: &CudaSlice<i8>,
14725        ad: &CudaSlice<f32>,
14726        m: usize,
14727        in_f: usize,
14728        out0: usize,
14729        out1: usize,
14730        out2: usize,
14731        row_bytes: usize,
14732    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14733        const ROWS_PER_BLOCK: u32 = 4;
14734        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14735        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14736        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14737        let f = self.func(if Self::batched_mcols(m) == 2 {
14738            "qmatvec_q8_0_mmvq_fused3_b2"
14739        } else {
14740            "qmatvec_q8_0_mmvq_fused3_b4"
14741        });
14742        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14743        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14744        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14745        let cfg = LaunchConfig {
14746            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14747            block_dim: (32, ROWS_PER_BLOCK, 1),
14748            shared_mem_bytes: 0,
14749        };
14750        let (inf, o0, o1, o2, mi, rbl) = (
14751            in_f as i32,
14752            out0 as i32,
14753            out1 as i32,
14754            out2 as i32,
14755            m as i32,
14756            row_bytes as i64,
14757        );
14758        let __s_b = self.gpu.stream();
14759        let mut b = __s_b.launch_builder(&f);
14760        b.arg(b0)
14761            .arg(b1)
14762            .arg(b2)
14763            .arg(aq)
14764            .arg(ad)
14765            .arg(&mut y0)
14766            .arg(&mut y1)
14767            .arg(&mut y2)
14768            .arg(&inf)
14769            .arg(&o0)
14770            .arg(&o1)
14771            .arg(&o2)
14772            .arg(&mi)
14773            .arg(&rbl);
14774        unsafe {
14775            b.launch(cfg)?;
14776        }
14777        Ok((y0, y1, y2))
14778    }
14779
14780    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
14781    #[allow(clippy::too_many_arguments)]
14782    pub fn qmatvec_q8_fused3_t_raw(
14783        &self,
14784        b0: &CudaSlice<u8>,
14785        b1: &CudaSlice<u8>,
14786        b2: &CudaSlice<u8>,
14787        x: &CudaSlice<f32>,
14788        m: usize,
14789        in_f: usize,
14790        out0: usize,
14791        out1: usize,
14792        out2: usize,
14793        row_bytes: usize,
14794    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14795        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14796        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14797    }
14798
14799    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
14800    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
14801    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
14802    pub fn q8_ffn_fuse2_on(&self) -> bool {
14803        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14804        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14805    }
14806
14807    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
14808    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
14809    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
14810    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
14811    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
14812    #[allow(clippy::type_complexity)]
14813    fn q8_fused_params<'w, const N: usize>(
14814        &self,
14815        ws: &[&'w crate::model::GpuTensor; N],
14816    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14817        use crate::model::GpuTensor;
14818        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14819            return None;
14820        }
14821        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14822            return None;
14823        }
14824        let in_f = ws[0].in_features();
14825        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14826        for (i, w) in ws.iter().enumerate() {
14827            match w {
14828                GpuTensor::Quant {
14829                    bytes,
14830                    qtype,
14831                    row_bytes,
14832                    scale,
14833                    ..
14834                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14835                    out[i] = Some((bytes, w.out_features(), *row_bytes))
14836                }
14837                _ => return None,
14838            }
14839        }
14840        Some(out.map(|o| o.unwrap()))
14841    }
14842
14843    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
14844    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
14845    pub fn e4m3_dual_on(&self) -> bool {
14846        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14847        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14848    }
14849
14850    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
14851    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
14852    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
14853    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
14854    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
14855    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
14856    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
14857    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
14858    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
14859    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
14860    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
14861    #[allow(clippy::type_complexity)]
14862    fn e4m3_fused_params<'w, const N: usize>(
14863        &self,
14864        ws: &[&'w crate::model::GpuTensor; N],
14865    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14866        use crate::model::GpuTensor;
14867        if !self.e4m3_dual_on() {
14868            return None;
14869        }
14870        let in_f = ws[0].in_features();
14871        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14872        for (i, w) in ws.iter().enumerate() {
14873            match w {
14874                GpuTensor::Quant {
14875                    bytes,
14876                    qtype,
14877                    row_bytes,
14878                    scale,
14879                    rp,
14880                    rp4,
14881                    ..
14882                } if *qtype == QT_F8_E4M3
14883                    && w.in_features() == in_f
14884                    && *row_bytes == in_f
14885                    && !*rp
14886                    && rp4.is_none() =>
14887                {
14888                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14889                }
14890                _ => return None,
14891            }
14892        }
14893        Some(out.map(|o| o.unwrap()))
14894    }
14895
14896    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
14897    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
14898    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
14899    #[allow(clippy::too_many_arguments)]
14900    fn e4m3_fused2_core(
14901        &self,
14902        b0: &CudaSlice<u8>,
14903        b1: &CudaSlice<u8>,
14904        aq: &CudaSlice<i8>,
14905        ad: &CudaSlice<f32>,
14906        in_f: usize,
14907        out0: usize,
14908        out1: usize,
14909        row_bytes: usize,
14910        ws0: f32,
14911        ws1: f32,
14912    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14913        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14914        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14915        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14916        let f = self.func("qmatvec_e4m3_mmvq_fused2");
14917        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14918        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14919        let cfg = LaunchConfig {
14920            grid_dim: (nb0 + nb1, 1, 1),
14921            block_dim: (32, ROWS_PER_BLOCK, 1),
14922            shared_mem_bytes: 0,
14923        };
14924        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14925        let __s_b = self.gpu.stream();
14926        let mut b = __s_b.launch_builder(&f);
14927        b.arg(b0)
14928            .arg(b1)
14929            .arg(aq)
14930            .arg(ad)
14931            .arg(&mut y0)
14932            .arg(&mut y1)
14933            .arg(&inf)
14934            .arg(&o0)
14935            .arg(&o1)
14936            .arg(&rbl)
14937            .arg(&ws0)
14938            .arg(&ws1);
14939        unsafe {
14940            b.launch(cfg)?;
14941        }
14942        Ok((y0, y1))
14943    }
14944
14945    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
14946    #[allow(clippy::too_many_arguments)]
14947    fn e4m3_fused3_core(
14948        &self,
14949        b0: &CudaSlice<u8>,
14950        b1: &CudaSlice<u8>,
14951        b2: &CudaSlice<u8>,
14952        aq: &CudaSlice<i8>,
14953        ad: &CudaSlice<f32>,
14954        in_f: usize,
14955        out0: usize,
14956        out1: usize,
14957        out2: usize,
14958        row_bytes: usize,
14959        ws0: f32,
14960        ws1: f32,
14961        ws2: f32,
14962    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14963        const ROWS_PER_BLOCK: u32 = 4;
14964        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14965        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14966        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14967        let f = self.func("qmatvec_e4m3_mmvq_fused3");
14968        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14969        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14970        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14971        let cfg = LaunchConfig {
14972            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14973            block_dim: (32, ROWS_PER_BLOCK, 1),
14974            shared_mem_bytes: 0,
14975        };
14976        let (inf, o0, o1, o2, rbl) = (
14977            in_f as i32,
14978            out0 as i32,
14979            out1 as i32,
14980            out2 as i32,
14981            row_bytes as i64,
14982        );
14983        let __s_b = self.gpu.stream();
14984        let mut b = __s_b.launch_builder(&f);
14985        b.arg(b0)
14986            .arg(b1)
14987            .arg(b2)
14988            .arg(aq)
14989            .arg(ad)
14990            .arg(&mut y0)
14991            .arg(&mut y1)
14992            .arg(&mut y2)
14993            .arg(&inf)
14994            .arg(&o0)
14995            .arg(&o1)
14996            .arg(&o2)
14997            .arg(&rbl)
14998            .arg(&ws0)
14999            .arg(&ws1)
15000            .arg(&ws2);
15001        unsafe {
15002            b.launch(cfg)?;
15003        }
15004        Ok((y0, y1, y2))
15005    }
15006
15007    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
15008    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
15009    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
15010    #[allow(clippy::too_many_arguments)]
15011    fn e4m3_fused2_t_core(
15012        &self,
15013        b0: &CudaSlice<u8>,
15014        b1: &CudaSlice<u8>,
15015        aq: &CudaSlice<i8>,
15016        ad: &CudaSlice<f32>,
15017        m: usize,
15018        in_f: usize,
15019        out0: usize,
15020        out1: usize,
15021        row_bytes: usize,
15022        ws0: f32,
15023        ws1: f32,
15024    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15025        const ROWS_PER_BLOCK: u32 = 4;
15026        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15027        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15028        let f = self.func(match Self::batched_mcols(m) {
15029            2 => "qmatvec_e4m3_mmvq_fused2_b2",
15030            4 => "qmatvec_e4m3_mmvq_fused2_b4",
15031            _ => "qmatvec_e4m3_mmvq_fused2_b8",
15032        });
15033        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15034        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15035        let cfg = LaunchConfig {
15036            grid_dim: (nb0 + nb1, 1, 1),
15037            block_dim: (32, ROWS_PER_BLOCK, 1),
15038            shared_mem_bytes: 0,
15039        };
15040        let (inf, o0, o1, mi, rbl) = (
15041            in_f as i32,
15042            out0 as i32,
15043            out1 as i32,
15044            m as i32,
15045            row_bytes as i64,
15046        );
15047        let __s_b = self.gpu.stream();
15048        let mut b = __s_b.launch_builder(&f);
15049        b.arg(b0)
15050            .arg(b1)
15051            .arg(aq)
15052            .arg(ad)
15053            .arg(&mut y0)
15054            .arg(&mut y1)
15055            .arg(&inf)
15056            .arg(&o0)
15057            .arg(&o1)
15058            .arg(&mi)
15059            .arg(&rbl);
15060        unsafe {
15061            b.launch(cfg)?;
15062        }
15063        if ws0 != 1.0 {
15064            self.scale_inplace(&mut y0, ws0, m * out0)?;
15065        }
15066        if ws1 != 1.0 {
15067            self.scale_inplace(&mut y1, ws1, m * out1)?;
15068        }
15069        Ok((y0, y1))
15070    }
15071
15072    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
15073    #[allow(clippy::too_many_arguments)]
15074    fn e4m3_fused3_t_core(
15075        &self,
15076        b0: &CudaSlice<u8>,
15077        b1: &CudaSlice<u8>,
15078        b2: &CudaSlice<u8>,
15079        aq: &CudaSlice<i8>,
15080        ad: &CudaSlice<f32>,
15081        m: usize,
15082        in_f: usize,
15083        out0: usize,
15084        out1: usize,
15085        out2: usize,
15086        row_bytes: usize,
15087        ws0: f32,
15088        ws1: f32,
15089        ws2: f32,
15090    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15091        const ROWS_PER_BLOCK: u32 = 4;
15092        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15093        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15094        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15095        let f = self.func(if Self::batched_mcols(m) == 2 {
15096            "qmatvec_e4m3_mmvq_fused3_b2"
15097        } else {
15098            "qmatvec_e4m3_mmvq_fused3_b4"
15099        });
15100        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15101        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15102        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15103        let cfg = LaunchConfig {
15104            grid_dim: (nb0 + nb1 + nb2, 1, 1),
15105            block_dim: (32, ROWS_PER_BLOCK, 1),
15106            shared_mem_bytes: 0,
15107        };
15108        let (inf, o0, o1, o2, mi, rbl) = (
15109            in_f as i32,
15110            out0 as i32,
15111            out1 as i32,
15112            out2 as i32,
15113            m as i32,
15114            row_bytes as i64,
15115        );
15116        let __s_b = self.gpu.stream();
15117        let mut b = __s_b.launch_builder(&f);
15118        b.arg(b0)
15119            .arg(b1)
15120            .arg(b2)
15121            .arg(aq)
15122            .arg(ad)
15123            .arg(&mut y0)
15124            .arg(&mut y1)
15125            .arg(&mut y2)
15126            .arg(&inf)
15127            .arg(&o0)
15128            .arg(&o1)
15129            .arg(&o2)
15130            .arg(&mi)
15131            .arg(&rbl);
15132        unsafe {
15133            b.launch(cfg)?;
15134        }
15135        if ws0 != 1.0 {
15136            self.scale_inplace(&mut y0, ws0, m * out0)?;
15137        }
15138        if ws1 != 1.0 {
15139            self.scale_inplace(&mut y1, ws1, m * out1)?;
15140        }
15141        if ws2 != 1.0 {
15142            self.scale_inplace(&mut y2, ws2, m * out2)?;
15143        }
15144        Ok((y0, y1, y2))
15145    }
15146
15147    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
15148    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
15149    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
15150    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
15151    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
15152    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
15153    ///
15154    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
15155    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
15156    pub fn qmatvec_e4m3_blk_mmvq(
15157        &self,
15158        bytes: &CudaSlice<u8>,
15159        aq: &CudaSlice<i8>,
15160        ad: &CudaSlice<f32>,
15161        scales: &CudaSlice<f32>,
15162        m: usize,
15163        in_f: usize,
15164        out_f: usize,
15165        row_bytes: usize,
15166        scale_cols: usize,
15167    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15168        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
15169        self.qmatvec_e4m3_blk_mmvq_into(
15170            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15171        )?;
15172        Ok(y)
15173    }
15174
15175    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
15176    #[allow(clippy::too_many_arguments)]
15177    pub fn qmatvec_e4m3_blk_mmvq_into(
15178        &self,
15179        bytes: &CudaSlice<u8>,
15180        aq: &CudaSlice<i8>,
15181        ad: &CudaSlice<f32>,
15182        scales: &CudaSlice<f32>,
15183        m: usize,
15184        in_f: usize,
15185        out_f: usize,
15186        row_bytes: usize,
15187        scale_cols: usize,
15188        y: &mut CudaSlice<f32>,
15189    ) -> Result<(), Box<dyn std::error::Error>> {
15190        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15191        let f = self.func("qmatvec_e4m3_blk_mmvq");
15192        let cfg = LaunchConfig {
15193            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15194            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
15195            shared_mem_bytes: 0,                // warp-only reduce
15196        };
15197        let (inf, outf, mi, rb, sc) = (
15198            in_f as i32,
15199            out_f as i32,
15200            m as i32,
15201            row_bytes as i64,
15202            scale_cols as i32,
15203        );
15204        let __s_b = self.gpu.stream();
15205        let mut b = __s_b.launch_builder(&f);
15206        b.arg(bytes)
15207            .arg(aq)
15208            .arg(ad)
15209            .arg(scales)
15210            .arg(&mut *y)
15211            .arg(&inf)
15212            .arg(&outf)
15213            .arg(&mi)
15214            .arg(&rb)
15215            .arg(&sc);
15216        unsafe {
15217            b.launch(cfg)?;
15218        }
15219        Ok(())
15220    }
15221
15222    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
15223    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
15224    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
15225    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
15226    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
15227    #[allow(clippy::too_many_arguments)]
15228    pub fn qmatvec_e4m3_blk_mmvq_batched(
15229        &self,
15230        bytes: &CudaSlice<u8>,
15231        aq: &CudaSlice<i8>,
15232        ad: &CudaSlice<f32>,
15233        scales: &CudaSlice<f32>,
15234        m: usize,
15235        in_f: usize,
15236        out_f: usize,
15237        row_bytes: usize,
15238        scale_cols: usize,
15239        mcols: usize,
15240    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15241        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15242        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15243        let name = match mcols {
15244            2 => "qmatvec_e4m3_blk_mmvq_b2",
15245            4 => "qmatvec_e4m3_blk_mmvq_b4",
15246            8 => "qmatvec_e4m3_blk_mmvq_b8",
15247            16 => "qmatvec_e4m3_blk_mmvq_b16",
15248            _ => {
15249                return Err(
15250                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15251                );
15252            }
15253        };
15254        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15255        let f = self.func(name);
15256        let cfg = LaunchConfig {
15257            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15258            block_dim: (32, ROWS_PER_BLOCK, 1),
15259            shared_mem_bytes: 0,
15260        };
15261        let (inf, outf, mi, rb, sc) = (
15262            in_f as i32,
15263            out_f as i32,
15264            m as i32,
15265            row_bytes as i64,
15266            scale_cols as i32,
15267        );
15268        let __s_b = self.gpu.stream();
15269        let mut b = __s_b.launch_builder(&f);
15270        b.arg(bytes)
15271            .arg(aq)
15272            .arg(ad)
15273            .arg(scales)
15274            .arg(&mut y)
15275            .arg(&inf)
15276            .arg(&outf)
15277            .arg(&mi)
15278            .arg(&rb)
15279            .arg(&sc);
15280        unsafe {
15281            b.launch(cfg)?;
15282        }
15283        Ok(y)
15284    }
15285
15286    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
15287    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
15288    #[allow(clippy::too_many_arguments)]
15289    pub fn qmatvec_e4m3_blk_batched_raw(
15290        &self,
15291        bytes: &CudaSlice<u8>,
15292        x: &CudaSlice<f32>,
15293        scales: &CudaSlice<f32>,
15294        m: usize,
15295        in_f: usize,
15296        out_f: usize,
15297        row_bytes: usize,
15298        scale_cols: usize,
15299        mcols: usize,
15300    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15301        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15302        self.qmatvec_e4m3_blk_mmvq_batched(
15303            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15304        )
15305    }
15306
15307    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
15308    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
15309    #[allow(clippy::too_many_arguments)]
15310    pub fn qmatvec_e4m3_blk_mmvq_raw(
15311        &self,
15312        bytes: &CudaSlice<u8>,
15313        x: &CudaSlice<f32>,
15314        scales: &CudaSlice<f32>,
15315        m: usize,
15316        in_f: usize,
15317        out_f: usize,
15318        row_bytes: usize,
15319        scale_cols: usize,
15320    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15321        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15322        self.qmatvec_e4m3_blk_mmvq(
15323            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15324        )
15325    }
15326
15327    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
15328    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
15329    #[allow(clippy::too_many_arguments)]
15330    pub fn qmatvec_e4m3_fused2_raw(
15331        &self,
15332        b0: &CudaSlice<u8>,
15333        b1: &CudaSlice<u8>,
15334        x: &CudaSlice<f32>,
15335        in_f: usize,
15336        out0: usize,
15337        out1: usize,
15338        row_bytes: usize,
15339        ws0: f32,
15340        ws1: f32,
15341    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15342        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15343        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15344    }
15345
15346    #[allow(clippy::too_many_arguments)]
15347    pub fn qmatvec_e4m3_fused3_raw(
15348        &self,
15349        b0: &CudaSlice<u8>,
15350        b1: &CudaSlice<u8>,
15351        b2: &CudaSlice<u8>,
15352        x: &CudaSlice<f32>,
15353        in_f: usize,
15354        out0: usize,
15355        out1: usize,
15356        out2: usize,
15357        row_bytes: usize,
15358        ws0: f32,
15359        ws1: f32,
15360        ws2: f32,
15361    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15362        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15363        self.e4m3_fused3_core(
15364            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15365        )
15366    }
15367
15368    #[allow(clippy::too_many_arguments)]
15369    pub fn qmatvec_e4m3_fused2_t_raw(
15370        &self,
15371        b0: &CudaSlice<u8>,
15372        b1: &CudaSlice<u8>,
15373        x: &CudaSlice<f32>,
15374        m: usize,
15375        in_f: usize,
15376        out0: usize,
15377        out1: usize,
15378        row_bytes: usize,
15379        ws0: f32,
15380        ws1: f32,
15381    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15382        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15383        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15384    }
15385
15386    #[allow(clippy::too_many_arguments)]
15387    pub fn qmatvec_e4m3_fused3_t_raw(
15388        &self,
15389        b0: &CudaSlice<u8>,
15390        b1: &CudaSlice<u8>,
15391        b2: &CudaSlice<u8>,
15392        x: &CudaSlice<f32>,
15393        m: usize,
15394        in_f: usize,
15395        out0: usize,
15396        out1: usize,
15397        out2: usize,
15398        row_bytes: usize,
15399        ws0: f32,
15400        ws1: f32,
15401        ws2: f32,
15402    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15403        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15404        self.e4m3_fused3_t_core(
15405            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15406        )
15407    }
15408
15409    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
15410    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
15411    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
15412    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
15413    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
15414    ///
15415    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
15416    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
15417    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
15418    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
15419    fn try_e4m3_blk_pre(
15420        &self,
15421        w: &crate::model::GpuTensor,
15422        aq: &CudaSlice<i8>,
15423        ad: &CudaSlice<f32>,
15424        m: usize,
15425    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15426        use crate::model::GpuTensor;
15427        if let GpuTensor::Quant {
15428            bytes,
15429            qtype,
15430            row_bytes,
15431            blk: Some(g),
15432            ..
15433        } = w
15434        {
15435            if *qtype == QT_F8_E4M3_BLK {
15436                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
15437                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
15438                // below, so the decode-exactness contract is preserved at every width. Gated by
15439                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
15440                // one rollback door covers every dtype's batched tier.
15441                if (2..=16).contains(&m)
15442                    && std::env::var("MEMRA_NO_BATCHED").is_err()
15443                    && (m <= 4 || Self::b8_enabled())
15444                {
15445                    let mcols = Self::batched_mcols(m);
15446                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15447                        bytes,
15448                        aq,
15449                        ad,
15450                        &g.scales,
15451                        m,
15452                        w.in_features(),
15453                        w.out_features(),
15454                        *row_bytes,
15455                        g.cols,
15456                        mcols,
15457                    )?));
15458                }
15459                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15460                    bytes,
15461                    aq,
15462                    ad,
15463                    &g.scales,
15464                    m,
15465                    w.in_features(),
15466                    w.out_features(),
15467                    *row_bytes,
15468                    g.cols,
15469                )?));
15470            }
15471        }
15472        Ok(None)
15473    }
15474
15475    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
15476    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
15477    ///
15478    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
15479    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
15480    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
15481    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
15482    /// prefill keeps the floor's arithmetic and the floor's kernels.
15483    ///
15484    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
15485    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
15486    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
15487    /// (projection, prefill call) and frees immediately.
15488    ///
15489    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
15490    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
15491    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
15492    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
15493    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
15494    /// single-variable comparison instead of a two-variable one.
15495    ///
15496    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
15497    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
15498    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
15499    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
15500    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
15501    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
15502    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
15503    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
15504    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
15505    ///
15506    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
15507    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
15508    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
15509    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
15510    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
15511    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
15512    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
15513    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
15514    /// because v2's denominator had its slab already resident while this class's floor must build it
15515    /// every call; same tile, opposite sign, because the question changed.
15516    ///
15517    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
15518    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
15519    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
15520    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
15521    fn try_e4m3_blk_prefill(
15522        &self,
15523        w: &crate::model::GpuTensor,
15524        x: &CudaSlice<f32>,
15525        m: usize,
15526    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15527        use crate::model::GpuTensor;
15528        let GpuTensor::Quant {
15529            bytes,
15530            qtype,
15531            blk: Some(g),
15532            ..
15533        } = w
15534        else {
15535            return Ok(None);
15536        };
15537        if *qtype != QT_F8_E4M3_BLK {
15538            return Ok(None);
15539        }
15540        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
15541        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
15542        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
15543        // through to the dequant below when they do, never silently produce nothing.
15544        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15545            return Ok(Some(y));
15546        }
15547        let (in_f, out_f) = (w.in_features(), w.out_features());
15548        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15549        let tmp = GpuTensor::Quant {
15550            bytes: slab,
15551            qtype: QT_Q8_0,
15552            row_bytes: in_f / 32 * 34,
15553            ne: vec![in_f as u64, out_f as u64],
15554            scale: 1.0,
15555            rp: false,
15556            #[cfg(memra_cutlass)]
15557            cutlass: None,
15558            fp8: None,
15559            blk: None,
15560            f16: None,
15561            rp4: None,
15562        };
15563        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
15564        Ok(Some(self.matmul(&tmp, x, m)?))
15565    }
15566
15567    pub fn matmul_pre_noscale(
15568        &self,
15569        w: &crate::model::GpuTensor,
15570        aq: &CudaSlice<i8>,
15571        ad: &CudaSlice<f32>,
15572        m: usize,
15573    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15574        use crate::model::GpuTensor;
15575        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
15576        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
15577        // rather than let the tail below refuse and cost the caller a re-dispatch.
15578        if m == 1 {
15579            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15580                return Ok(Some((y, 1.0)));
15581            }
15582        }
15583        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
15584        if m != 1 || !self.uses_q8_1_fast(w) {
15585            return Ok(None);
15586        }
15587        let in_f = w.in_features();
15588        let out_f = w.out_features();
15589        let (bytes, qtype, row_bytes, scale, rp) = match w {
15590            GpuTensor::Quant {
15591                bytes,
15592                qtype,
15593                row_bytes,
15594                scale,
15595                rp,
15596                ..
15597            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15598            _ => return Ok(None),
15599        };
15600        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
15601        if self.mmvq_supports(qtype) {
15602            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
15603            let (mbytes, mrp) = match w {
15604                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15605                _ => (bytes, rp),
15606            };
15607            let y = self.qmatvec_mmvq(
15608                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
15609            )?;
15610            return Ok(Some((y, scale)));
15611        }
15612        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
15613        let name = match qtype {
15614            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15615            QT_Q4_K => "qmatvec_q4_K_dp4a",
15616            QT_Q6_K => "qmatvec_q6_K_dp4a",
15617            QT_Q5_K => "qmatvec_q5_K_dp4a",
15618            QT_Q3_K => "qmatvec_q3_K_dp4a",
15619            QT_NVFP4 => {
15620                if rp {
15621                    "qmatvec_nvfp4_dp4a_rp"
15622                } else {
15623                    "qmatvec_nvfp4_dp4a"
15624                }
15625            }
15626            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15627            _ => return Ok(None),
15628        };
15629        let f = self.func(name);
15630        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15631        let cfg = LaunchConfig {
15632            grid_dim: (out_f as u32, m as u32, 1),
15633            block_dim: (128, 1, 1),
15634            shared_mem_bytes: 0,
15635        };
15636        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15637        let __s_b = self.gpu.stream();
15638        let mut b = __s_b.launch_builder(&f);
15639        b.arg(bytes)
15640            .arg(aq)
15641            .arg(ad)
15642            .arg(&mut y)
15643            .arg(&inf)
15644            .arg(&outf)
15645            .arg(&mi)
15646            .arg(&rb);
15647        unsafe {
15648            b.launch(cfg)?;
15649        }
15650        Ok(Some((y, scale)))
15651    }
15652
15653    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15654    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15655    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15656        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15657        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15658        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15659        // is a pure function of the dtype — the decode-parity law holds under every env.
15660        if qtype == QT_F8_E4M3 {
15661            return true;
15662        }
15663        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15664            return false;
15665        }
15666        matches!(
15667            qtype,
15668            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15669        )
15670    }
15671
15672    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15673    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15674    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15675    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15676    pub fn qmatvec_mmvq(
15677        &self,
15678        bytes: &CudaSlice<u8>,
15679        aq: &CudaSlice<i8>,
15680        ad: &CudaSlice<f32>,
15681        m: usize,
15682        in_f: usize,
15683        out_f: usize,
15684        qtype: i32,
15685        row_bytes: usize,
15686        scale: f32,
15687        rp: bool,
15688    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15689        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15690        self.qmatvec_mmvq_into(
15691            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15692        )?;
15693        Ok(y)
15694    }
15695
15696    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15697    #[allow(clippy::too_many_arguments)]
15698    pub fn qmatvec_mmvq_into(
15699        &self,
15700        bytes: &CudaSlice<u8>,
15701        aq: &CudaSlice<i8>,
15702        ad: &CudaSlice<f32>,
15703        m: usize,
15704        in_f: usize,
15705        out_f: usize,
15706        qtype: i32,
15707        row_bytes: usize,
15708        scale: f32,
15709        rp: bool,
15710        y: &mut CudaSlice<f32>,
15711    ) -> Result<(), Box<dyn std::error::Error>> {
15712        debug_assert!(y.len() >= m * out_f);
15713        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15714        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15715        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15716        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15717        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15718        if qtype == QT_Q8_0
15719            && rp
15720            && m == 1
15721            && out_f >= 64
15722            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15723            && {
15724                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15725                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15726            }
15727        {
15728            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15729            let cfg = LaunchConfig {
15730                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15731                block_dim: (32, 2, 1),
15732                shared_mem_bytes: 0,
15733            };
15734            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15735            let __s_b = self.gpu.stream();
15736            let mut b = __s_b.launch_builder(&f);
15737            b.arg(bytes)
15738                .arg(aq)
15739                .arg(ad)
15740                .arg(&mut *y)
15741                .arg(&inf)
15742                .arg(&outf)
15743                .arg(&mi)
15744                .arg(&rb);
15745            unsafe {
15746                b.launch(cfg)?;
15747            }
15748            if scale != 1.0 {
15749                self.scale_inplace(y, scale, out_f)?;
15750            }
15751            return Ok(());
15752        }
15753        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15754        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15755        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15756        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
15757        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
15758        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
15759        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
15760        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
15761        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15762            2
15763        } else {
15764            1
15765        };
15766        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
15767        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
15768        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
15769        // valid-window interleaved, bit-identical per row — same dot program).
15770        if m == 1 && qtype == QT_Q4_0 {
15771            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15772            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
15773            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
15774            mr = *Q40MR.get_or_init(|| {
15775                std::env::var("MEMRA_Q40_MR")
15776                    .ok()
15777                    .and_then(|v| v.parse().ok())
15778                    .unwrap_or(1)
15779            });
15780        }
15781        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
15782        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
15783        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
15784        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
15785        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
15786        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
15787        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
15788        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
15789        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
15790        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
15791        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15792        let q5_force = q5_mode.as_deref() == Some("2");
15793        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
15794        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
15795        let q5_il = qtype == QT_Q5_K
15796            && m == 1
15797            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15798        if q5_il && !q5_force && out_f > 65536 {
15799            mr = 1;
15800        }
15801        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
15802        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
15803        if qtype == QT_Q4_0 && rp && mr != 1 {
15804            mr = 2;
15805        }
15806        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
15807        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
15808        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
15809        if qtype == QT_Q8_0 && rp {
15810            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15811            mr = *Q80MR.get_or_init(|| {
15812                std::env::var("MEMRA_Q80_MR")
15813                    .ok()
15814                    .and_then(|v| v.parse().ok())
15815                    .unwrap_or(1)
15816            });
15817        }
15818        let name = match (qtype, mr, rp) {
15819            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15820            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15821            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15822            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15823            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15824            (QT_Q5_K, 2, _) => {
15825                if q5_il {
15826                    "qmatvec_q5_K_mmvq_mr2_il"
15827                } else {
15828                    "qmatvec_q5_K_mmvq_mr2"
15829                }
15830            }
15831            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15832            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
15833            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
15834            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
15835            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
15836            (QT_Q8_0, _, true)
15837                if in_f % 1024 == 0 && {
15838                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15839                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15840                } =>
15841            {
15842                "qmatvec_q8_0_mmvq_rpca"
15843            }
15844            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15845            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15846            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
15847            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
15848            // reach a GGUF-layout kernel or vice versa.
15849            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15850            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15851            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15852            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15853            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15854            (QT_Q5_K, _, _) => {
15855                if q5_il {
15856                    "qmatvec_q5_K_mmvq_il"
15857                } else {
15858                    "qmatvec_q5_K_mmvq"
15859                }
15860            }
15861            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15862            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15863            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15864            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15865        };
15866        let f = self.func(name);
15867        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
15868        let rows_per_block = ROWS_PER_BLOCK * mr;
15869        let cfg = LaunchConfig {
15870            grid_dim: (
15871                (out_f as u32 + rows_per_block - 1) / rows_per_block,
15872                m as u32,
15873                1,
15874            ),
15875            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
15876            shared_mem_bytes: 0,                // warp-only reduce at m=1
15877        };
15878        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15879        let __s_b = self.gpu.stream();
15880        let mut b = __s_b.launch_builder(&f);
15881        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
15882        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
15883        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
15884        // weight_scale). Other mmvq kernels keep the 8-arg signature.
15885        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15886            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
15887            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
15888            if Self::pdl_on()
15889                && Self::pdl_mmvq_on()
15890                && Self::pdl_nvfp4q8_on()
15891                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15892            {
15893                use cudarc::driver::{DevicePtr, DevicePtrMut};
15894                let s = &self.gpu.stream();
15895                let (pw, _g0) = bytes.device_ptr(s);
15896                let (paq, _g1) = aq.device_ptr(s);
15897                let (pad, _g2) = ad.device_ptr(s);
15898                let (py, _g3) = y.device_ptr_mut(s);
15899                let mut ps = [
15900                    &pw as *const _ as *mut std::ffi::c_void,
15901                    &paq as *const _ as *mut _,
15902                    &pad as *const _ as *mut _,
15903                    &py as *const _ as *mut _,
15904                    &inf as *const _ as *mut _,
15905                    &outf as *const _ as *mut _,
15906                    &mi as *const _ as *mut _,
15907                    &rb as *const _ as *mut _,
15908                    &scale as *const _ as *mut _,
15909                ];
15910                unsafe {
15911                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15912                }
15913                return Ok(());
15914            }
15915            b.arg(bytes)
15916                .arg(aq)
15917                .arg(ad)
15918                .arg(&mut *y)
15919                .arg(&inf)
15920                .arg(&outf)
15921                .arg(&mi)
15922                .arg(&rb)
15923                .arg(&scale);
15924            unsafe {
15925                b.launch(cfg)?;
15926            }
15927        } else if Self::pdl_on()
15928            && Self::pdl_mmvq_on()
15929            && (matches!(
15930                name,
15931                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15932            ) || (Self::pdl_nvfp4q8_on()
15933                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15934        {
15935            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
15936            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
15937            // names may take this launch (unmarked kernels would read unordered).
15938            {
15939                use cudarc::driver::{DevicePtr, DevicePtrMut};
15940                let s = &self.gpu.stream();
15941                let (pw, _g0) = bytes.device_ptr(s);
15942                let (paq, _g1) = aq.device_ptr(s);
15943                let (pad, _g2) = ad.device_ptr(s);
15944                let (py, _g3) = y.device_ptr_mut(s);
15945                let mut ps = [
15946                    &pw as *const _ as *mut std::ffi::c_void,
15947                    &paq as *const _ as *mut _,
15948                    &pad as *const _ as *mut _,
15949                    &py as *const _ as *mut _,
15950                    &inf as *const _ as *mut _,
15951                    &outf as *const _ as *mut _,
15952                    &mi as *const _ as *mut _,
15953                    &rb as *const _ as *mut _,
15954                ];
15955                unsafe {
15956                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15957                }
15958            }
15959            if scale != 1.0 {
15960                self.scale_inplace(y, scale, m * out_f)?;
15961            }
15962        } else {
15963            b.arg(bytes)
15964                .arg(aq)
15965                .arg(ad)
15966                .arg(&mut *y)
15967                .arg(&inf)
15968                .arg(&outf)
15969                .arg(&mi)
15970                .arg(&rb);
15971            unsafe {
15972                b.launch(cfg)?;
15973            }
15974            if scale != 1.0 {
15975                self.scale_inplace(y, scale, m * out_f)?;
15976            }
15977        }
15978        Ok(())
15979    }
15980
15981    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
15982    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
15983    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
15984    pub fn qmatvec_mmvq_raw(
15985        &self,
15986        bytes: &CudaSlice<u8>,
15987        x: &CudaSlice<f32>,
15988        m: usize,
15989        in_f: usize,
15990        out_f: usize,
15991        qtype: i32,
15992        row_bytes: usize,
15993        rp: bool,
15994    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15995        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15996        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15997    }
15998
15999    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
16000    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
16001    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
16002    pub fn batched_supports(&self, qtype: i32) -> bool {
16003        matches!(
16004            qtype,
16005            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16006        )
16007    }
16008
16009    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
16010    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
16011    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
16012    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
16013    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
16014    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
16015    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
16016    pub fn iq_fast_enabled() -> bool {
16017        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16018        *ON.get_or_init(|| {
16019            std::env::var("MEMRA_IQ_FAST")
16020                .map(|v| v != "0")
16021                .unwrap_or(true)
16022        })
16023    }
16024
16025    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
16026    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
16027    pub fn b8_enabled() -> bool {
16028        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16029        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16030    }
16031
16032    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
16033    pub fn batched_mcols(m: usize) -> usize {
16034        if m == 2 {
16035            2
16036        } else if m <= 4 {
16037            4
16038        } else if m <= 8 {
16039            8
16040        } else {
16041            16
16042        }
16043    }
16044
16045    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
16046    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
16047    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
16048    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
16049    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16050        Some(match (qtype, mcols) {
16051            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16052            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16053            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16054            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
16055            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
16056            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
16057            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
16058            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
16059            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16060            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16061            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16062            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16063            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
16064            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
16065            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16066            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16067            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16068            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16069            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
16070            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
16071            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16072            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16073            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16074            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16075            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16076            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16077            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16078            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16079            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
16080            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
16081            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
16082            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16083            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16084            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16085            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16086            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
16087            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
16088            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
16089            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16090            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16091            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16092            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16093            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16094            _ => return None,
16095        })
16096    }
16097
16098    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
16099    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
16100    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
16101    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
16102    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
16103    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
16104    ///
16105    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
16106    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
16107    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
16108    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
16109    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
16110    /// msweep on all six 27B shapes (2026-07-03):
16111    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
16112    ///          it applies for b4 (-3..-14%), never loses;
16113    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
16114    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
16115    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
16116    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
16117    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
16118    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
16119    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
16120    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
16121    /// b2: in_f>=6144 -> r2, else base.
16122    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
16123    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
16124    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
16125    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
16126    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
16127    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
16128    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
16129    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
16130    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
16131    /// Device SM count (cached) — grid-fill policy input.
16132    pub fn sm_count(&self) -> i32 {
16133        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16134        *SMS.get_or_init(|| {
16135            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16136            self.gpu
16137                .ctx
16138                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16139                .unwrap_or(82)
16140        })
16141    }
16142
16143    pub fn batched_variant(
16144        &self,
16145        _m: usize,
16146        in_f: usize,
16147        out_f: usize,
16148        qtype: i32,
16149        row_bytes: usize,
16150        mcols: usize,
16151        rp: bool,
16152    ) -> &'static str {
16153        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
16154        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
16155        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
16156        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
16157        if qtype == QT_Q8_0 {
16158            return if rp { "rp" } else { "base" };
16159        }
16160        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16161        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16162            Ok("base") => "base",
16163            Ok("pf") => "pf",
16164            Ok("r2") => "r2",
16165            Ok("r2w8") => "r2w8",
16166            Ok("pfr2") => "pfr2",
16167            Ok("ca") => "ca",
16168            Ok("car2") => "car2",
16169            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
16170            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
16171            Ok("rp") => "rp",
16172            Ok("rpr2") => "rpr2",
16173            Ok("rpr2w8") => "rpr2w8",
16174            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
16175            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
16176            Ok("rpca") => "rpca",
16177            Ok("rpcar2") => "rpcar2",
16178            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
16179            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
16180            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
16181            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
16182            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
16183            // bit-identical to the decode path — measurement corpus ONLY, never auto).
16184            Ok("rpsc") => "rpsc",
16185            Ok("rpms") => "rpms",
16186            Ok("rpmsc") => "rpmsc",
16187            Ok("rpks") => "rpks",
16188            Ok("rpksc") => "rpksc",
16189            _ => "auto",
16190        });
16191        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
16192        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
16193        // shapes qualify; anything else falls back to the register variants.
16194        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16195        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
16196        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
16197        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
16198        // forced MEMRA_MMVQ_BV values still work).
16199        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16200        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16201        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16202        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16203        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16204        let sms = *SMS.get_or_init(|| {
16205            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16206            self.gpu
16207                .ctx
16208                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16209                .unwrap_or(82)
16210        });
16211        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
16212        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
16213        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
16214        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
16215        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
16216        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
16217        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
16218        // AUTO RULE = the measured winners table (differs from NVFP4's!):
16219        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
16220        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
16221        //     r2 1258us) — kernels kept behind the force seam for the corpus;
16222        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
16223        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
16224        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
16225        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
16226        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
16227        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
16228        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
16229        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
16230        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16231        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
16232        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
16233        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16234        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16235            Ok("base") => "base",
16236            Ok("r2") => "r2",
16237            Ok("r2w8") => "r2w8",
16238            _ => "auto",
16239        });
16240        let variant: &'static str = if qtype == QT_Q4_0 {
16241            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
16242            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
16243            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
16244            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16245            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16246                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
16247                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
16248                // + syncs cost more than the stalls, bank-pad made no difference);
16249                // register load-ahead flat (nvcc already reorders). The b-tier limiter
16250                // is still unidentified — see the jsonl row.
16251                Ok("base") => "base",
16252                Ok("r2") => "r2",
16253                Ok("ms") => "ms",
16254                Ok("sm") => "sm",
16255                Ok("la") => "la",
16256                _ => "auto",
16257            });
16258            let v = if q40 != "auto" {
16259                q40
16260            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16261                "r2"
16262            } else {
16263                "base"
16264            };
16265            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
16266            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
16267            // and the limiter is the per-column activation load chain (long_scoreboard
16268            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
16269            if rp {
16270                match v {
16271                    "ms" => "r2ms_rp",
16272                    "sm" => "r2sm_rp",
16273                    "la" => "r2la_rp",
16274                    "r2" => "r2_rp",
16275                    _ => "rp",
16276                }
16277            } else if matches!(v, "ms" | "sm" | "la") {
16278                "r2"
16279            } else {
16280                v
16281            }
16282        } else if qtype != QT_NVFP4 && !kq_r2 {
16283            "base"
16284        } else if kq_r2 && rp {
16285            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
16286            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
16287            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
16288            "rp"
16289        } else if kq_r2 {
16290            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
16291            // mcols != 4 forced r2w8 falls to unbounded r2.
16292            if kq_bv != "auto" {
16293                if kq_bv == "r2w8" && mcols != 4 {
16294                    "r2"
16295                } else {
16296                    kq_bv
16297                }
16298            } else if bv != "auto" {
16299                match bv {
16300                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16301                    "r2w8" | "rpr2w8" => {
16302                        if mcols != 4 {
16303                            "r2"
16304                        } else {
16305                            "r2w8"
16306                        }
16307                    }
16308                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
16309                }
16310            } else {
16311                let blocks = (out_f + 7) / 8;
16312                let waves = blocks as f64 / (7 * sms as usize) as f64;
16313                let filled = blocks >= 4 * sms as usize;
16314                let use_r2 = if qtype == QT_Q4_K {
16315                    filled
16316                } else {
16317                    waves >= 2.0
16318                };
16319                if use_r2 { "r2" } else { "base" }
16320            }
16321        } else if bv != "auto" {
16322            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
16323            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
16324            // unsupported (shape, mcols) combos fall back to pf/r2.
16325            // On rp buffers, forced legacy names map to their rp twins (layout law).
16326            let v = if bv == "r2w8" && mcols == 2 {
16327                "r2"
16328            } else if bv == "ca" && (!ca_ok || mcols == 8) {
16329                "pf"
16330            } else if bv == "car2" && (!ca_ok || mcols == 8) {
16331                "r2"
16332            } else if bv == "pfr2" && mcols == 8 {
16333                "r2"
16334            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16335                "rpr2"
16336            }
16337            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
16338            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16339                if mcols == 8 { "rpr2w8" } else { "rpr2" }
16340            } else if bv == "rpcar2" && mcols == 2 {
16341                "rpca"
16342            }
16343            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
16344            // (rpms has no smem and no alignment need — always valid on rp buffers).
16345            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16346                "rpr2"
16347            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16348                "rpr2"
16349            } else {
16350                bv
16351            };
16352            if rp {
16353                match v {
16354                    "base" | "pf" | "ca" | "rp" => "rp",
16355                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16356                    "r2w8" | "rpr2w8" => {
16357                        if mcols == 2 {
16358                            "rpr2"
16359                        } else {
16360                            "rpr2w8"
16361                        }
16362                    }
16363                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
16364                }
16365            } else {
16366                v
16367            }
16368        } else if mcols == 8 {
16369            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
16370            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
16371            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
16372            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
16373            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
16374            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
16375            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
16376            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
16377            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
16378            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
16379            if rp {
16380                if sc_ok { "rpsc" } else { "rpr2w8" }
16381            } else {
16382                "r2w8"
16383            }
16384        } else if mcols >= 4 {
16385            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
16386            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
16387            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
16388            let blocks = (out_f + 7) / 8;
16389            let r7 = 7 * sms as usize;
16390            let r8 = 8 * sms as usize;
16391            let waves = blocks as f64 / r7 as f64;
16392            let filled = blocks >= 4 * sms as usize;
16393            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
16394            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
16395            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
16396            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16397                // the extra residency drops the INTEGER wave count -> the straggler wave a
16398                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
16399                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
16400                if rp { "rpr2w8" } else { "r2w8" }
16401            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16402                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
16403                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
16404                if rp { "rpr2" } else { "r2" }
16405            } else {
16406                // fractional straggler-wave window with no crossing, or grid too small to fill
16407                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
16408                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
16409                if rp { "rp" } else { "pf" }
16410            }
16411        } else if in_f >= 6144 {
16412            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
16413            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
16414            // stays.
16415            if rp { "rpr2" } else { "r2" }
16416        } else if rp {
16417            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
16418            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
16419            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
16420            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
16421            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16422            if sc_ok && waves >= 0.9 && waves <= 1.1 {
16423                "rpsc"
16424            } else {
16425                "rp"
16426            }
16427        } else {
16428            "base"
16429        };
16430        variant
16431    }
16432
16433    pub fn qmatvec_mmvq_batched(
16434        &self,
16435        bytes: &CudaSlice<u8>,
16436        aq: &CudaSlice<i8>,
16437        ad: &CudaSlice<f32>,
16438        m: usize,
16439        in_f: usize,
16440        out_f: usize,
16441        qtype: i32,
16442        row_bytes: usize,
16443        mcols: usize,
16444        scale: f32,
16445        rp: bool,
16446    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16447        const ROWS_PER_BLOCK: u32 = 4;
16448        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
16449        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
16450        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
16451        // weight keeps its rp-layout kernel family regardless of the override.
16452        let forced: Option<&'static str> = {
16453            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16454            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16455                .as_deref()
16456                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16457        };
16458        let variant = match forced {
16459            Some(v) if !rp || v.contains("rp") => v,
16460            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16461        };
16462        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16463            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16464        })?;
16465        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
16466        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
16467        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
16468        let variant = if mcols == 16 {
16469            if rp { "rp" } else { "base" }
16470        } else {
16471            variant
16472        };
16473        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
16474        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
16475        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
16476        // per-(token,row) chain (columns c >= m never execute in either form) ->
16477        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
16478        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
16479        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16480        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16481        if b567
16482            && qtype == QT_NVFP4
16483            && rp
16484            && mcols == 8
16485            && (5..=7).contains(&m)
16486            && matches!(variant, "rpsc" | "rpr2w8")
16487        {
16488            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16489            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
16490            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16491            let cfg = LaunchConfig {
16492                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16493                block_dim: (32, ROWS_PER_BLOCK, 1),
16494                shared_mem_bytes: 0,
16495            };
16496            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16497            let __s_b = self.gpu.stream();
16498            let mut b = __s_b.launch_builder(&f);
16499            b.arg(bytes)
16500                .arg(aq)
16501                .arg(ad)
16502                .arg(&mut y)
16503                .arg(&inf)
16504                .arg(&outf)
16505                .arg(&mi)
16506                .arg(&rb);
16507            unsafe {
16508                b.launch(cfg)?;
16509            }
16510            if scale != 1.0 {
16511                self.scale_inplace(&mut y, scale, m * out_f)?;
16512            }
16513            return Ok(y);
16514        }
16515        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16516            "base" => (base_name.into(), ROWS_PER_BLOCK),
16517            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16518            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16519            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16520            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
16521            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
16522            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
16523            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16524            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16525            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16526            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16527            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16528            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16529            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16530            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
16531        };
16532        debug_assert!(
16533            !rp || name.contains("_rp"),
16534            "rp weight dispatched to a GGUF-layout kernel"
16535        );
16536        let f = self.func(&name);
16537        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16538        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
16539        let smem = if name.contains("_r2sm_rp") {
16540            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16541        } else {
16542            0
16543        };
16544        let cfg = LaunchConfig {
16545            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16546            block_dim: (32, ROWS_PER_BLOCK, 1),
16547            shared_mem_bytes: smem,
16548        };
16549        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16550        let __s_b = self.gpu.stream();
16551        let mut b = __s_b.launch_builder(&f);
16552        b.arg(bytes)
16553            .arg(aq)
16554            .arg(ad)
16555            .arg(&mut y)
16556            .arg(&inf)
16557            .arg(&outf)
16558            .arg(&mi)
16559            .arg(&rb);
16560        unsafe {
16561            b.launch(cfg)?;
16562        }
16563        if scale != 1.0 {
16564            self.scale_inplace(&mut y, scale, m * out_f)?;
16565        }
16566        Ok(y)
16567    }
16568
16569    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
16570    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
16571    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
16572    pub fn qmatvec_batched_raw(
16573        &self,
16574        bytes: &CudaSlice<u8>,
16575        x: &CudaSlice<f32>,
16576        m: usize,
16577        in_f: usize,
16578        out_f: usize,
16579        qtype: i32,
16580        row_bytes: usize,
16581        mcols: usize,
16582        rp: bool,
16583    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16584        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16585        self.qmatvec_mmvq_batched(
16586            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16587        )
16588    }
16589
16590    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
16591    pub fn qmatvec_nvfp4_batched_raw(
16592        &self,
16593        bytes: &CudaSlice<u8>,
16594        x: &CudaSlice<f32>,
16595        m: usize,
16596        in_f: usize,
16597        out_f: usize,
16598        row_bytes: usize,
16599        mcols: usize,
16600        rp: bool,
16601    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16602        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16603    }
16604
16605    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
16606    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
16607    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
16608    fn try_fp4_gemm(
16609        &self,
16610        w: &crate::model::GpuTensor,
16611        x: &CudaSlice<f32>,
16612        m: usize,
16613        in_f: usize,
16614        out_f: usize,
16615    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16616        use crate::model::GpuTensor;
16617        if cfg!(memra_portable_cuda) {
16618            return Ok(None);
16619        }
16620        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which cu/qmatvec_gemm.cu:1234 omits on a
16621        // portable build (the mxf4 block-scale MMA is sm_120a-only). Refuse at the door.
16622        if std::env::var("MEMRA_FP4").is_ok() {
16623            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16624        }
16625        if std::env::var("MEMRA_FP4").is_err() {
16626            return Ok(None);
16627        }
16628        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
16629        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
16630        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
16631        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
16632        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
16633        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
16634        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
16635        // for the common no-macro-scale case.
16636        #[cfg(memra_cutlass)]
16637        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16638            if let GpuTensor::Quant {
16639                bytes,
16640                qtype,
16641                scale,
16642                row_bytes,
16643                cutlass,
16644                ..
16645            } = w
16646            {
16647                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16648                    if let Some(cw) = cutlass {
16649                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16650                        let y = self.cutlass_fp4_gemm(
16651                            &cw.b_packed,
16652                            &cw.sfb_swizzled,
16653                            x,
16654                            *scale,
16655                            m,
16656                            out_f,
16657                            in_f,
16658                        )?;
16659                        return Ok(Some(y));
16660                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16661                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16662                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16663                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16664                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16665                        let (b_packed, sfb_sw) =
16666                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16667                        let y =
16668                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16669                        return Ok(Some(y));
16670                    }
16671                }
16672            }
16673        }
16674        if let GpuTensor::Quant {
16675            bytes,
16676            qtype,
16677            row_bytes,
16678            scale,
16679            rp,
16680            ..
16681        } = w
16682        {
16683            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16684            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16685            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16686                let y =
16687                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16688                return Ok(Some(y));
16689            }
16690        }
16691        Ok(None)
16692    }
16693
16694    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16695    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16696    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16697    pub fn rms_norm_f16out(
16698        &self,
16699        x: &CudaSlice<f32>,
16700        w: &CudaSlice<f32>,
16701        dst: &mut CudaSlice<f32>,
16702        dst16: &mut CudaSlice<u8>,
16703        ncols: usize,
16704        nrows: usize,
16705        eps: f32,
16706    ) -> Result<(), Box<dyn std::error::Error>> {
16707        let f = self.func("rms_norm_f16out_f32");
16708        let cfg = LaunchConfig {
16709            grid_dim: (nrows as u32, 1, 1),
16710            block_dim: (rms_block(), 1, 1),
16711            shared_mem_bytes: 0,
16712        };
16713        let (nc, e) = (ncols as i32, eps);
16714        let __s_b = self.gpu.stream();
16715        let mut b = __s_b.launch_builder(&f);
16716        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16717        unsafe {
16718            b.launch(cfg)?;
16719        }
16720        Ok(())
16721    }
16722
16723    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16724    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16725    #[allow(clippy::too_many_arguments)]
16726    pub fn add_rms_norm_f16out(
16727        &self,
16728        a: &CudaSlice<f32>,
16729        b: &CudaSlice<f32>,
16730        w: &CudaSlice<f32>,
16731        res: &mut CudaSlice<f32>,
16732        dst: &mut CudaSlice<f32>,
16733        dst16: &mut CudaSlice<u8>,
16734        ncols: usize,
16735        nrows: usize,
16736        eps: f32,
16737    ) -> Result<(), Box<dyn std::error::Error>> {
16738        let f = self.func("add_rms_norm_f16out_f32");
16739        let cfg = LaunchConfig {
16740            grid_dim: (nrows as u32, 1, 1),
16741            block_dim: (rms_block(), 1, 1),
16742            shared_mem_bytes: 0,
16743        };
16744        let (nc, e) = (ncols as i32, eps);
16745        let __s_lb = self.gpu.stream();
16746        let mut lb = __s_lb.launch_builder(&f);
16747        lb.arg(a)
16748            .arg(b)
16749            .arg(w)
16750            .arg(res)
16751            .arg(dst)
16752            .arg(dst16)
16753            .arg(&nc)
16754            .arg(&e);
16755        unsafe {
16756            lb.launch(cfg)?;
16757        }
16758        Ok(())
16759    }
16760
16761    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
16762    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
16763    pub fn matmul_group_xh(
16764        &self,
16765        ws: &[&crate::model::GpuTensor],
16766        x: &CudaSlice<f32>,
16767        xh: &CudaSlice<u8>,
16768        m: usize,
16769    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16770        let mut out = Vec::with_capacity(ws.len());
16771        let in_f = ws[0].in_features();
16772        for w in ws {
16773            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16774                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16775                    out.push(y);
16776                    continue;
16777                }
16778            }
16779            out.push(self.matmul(w, x, m)?);
16780        }
16781        Ok(out)
16782    }
16783
16784    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
16785    /// GDN steps). Layouts [T, H].
16786    pub fn gdn_pad_mask(
16787        &self,
16788        beta: &mut CudaSlice<f32>,
16789        g_log: &mut CudaSlice<f32>,
16790        len_d: &CudaSlice<i32>,
16791        h: usize,
16792        t: usize,
16793    ) -> Result<(), Box<dyn std::error::Error>> {
16794        let f = self.func("gdn_pad_mask_f32");
16795        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16796        let (hi, ti) = (h as i32, t as i32);
16797        let __s_b = self.gpu.stream();
16798        let mut b = __s_b.launch_builder(&f);
16799        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16800        unsafe {
16801            b.launch(cfg)?;
16802        }
16803        Ok(())
16804    }
16805
16806    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
16807    /// gather for the padded prime graph's h_seed/hlast.
16808    pub fn row_gather_dev(
16809        &self,
16810        src: &CudaSlice<f32>,
16811        dst: &mut CudaSlice<f32>,
16812        len_d: &CudaSlice<i32>,
16813        ncols: usize,
16814    ) -> Result<(), Box<dyn std::error::Error>> {
16815        let f = self.func("row_gather_dev_f32");
16816        let cfg = LaunchConfig::for_num_elems(ncols as u32);
16817        let nc = ncols as i32;
16818        let __s_b = self.gpu.stream();
16819        let mut b = __s_b.launch_builder(&f);
16820        b.arg(src).arg(dst).arg(len_d).arg(&nc);
16821        unsafe {
16822            b.launch(cfg)?;
16823        }
16824        Ok(())
16825    }
16826
16827    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
16828    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
16829    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
16830    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
16831    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
16832    /// different in_f) falls back to its own `matmul` — behavior unchanged.
16833    pub fn matmul_group(
16834        &self,
16835        ws: &[&crate::model::GpuTensor],
16836        x: &CudaSlice<f32>,
16837        m: usize,
16838    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16839        use crate::model::GpuTensor;
16840        let mut out = Vec::with_capacity(ws.len());
16841        let any_mirror = ws
16842            .iter()
16843            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16844        if m >= 16 && any_mirror && !self.verify_exact_on() {
16845            let in_f = ws[0].in_features();
16846            let xh = self.f16_act(x, m * in_f, in_f)?;
16847            for w in ws {
16848                if w.in_features() == in_f {
16849                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16850                        out.push(y);
16851                        continue;
16852                    }
16853                }
16854                out.push(self.matmul(w, x, m)?);
16855            }
16856            return Ok(out);
16857        }
16858        for w in ws {
16859            out.push(self.matmul(w, x, m)?);
16860        }
16861        Ok(out)
16862    }
16863
16864    /// Cross-request grouped matmul (task #13): run ONE projection group over the
16865    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
16866    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
16867    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
16868    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
16869    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
16870    pub fn matmul_group_multi(
16871        &self,
16872        ws: &[&crate::model::GpuTensor],
16873        xs: &[&CudaSlice<f32>],
16874        ms: &[usize],
16875    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16876        assert_eq!(xs.len(), ms.len());
16877        let in_f = ws[0].in_features();
16878        let total: usize = ms.iter().sum();
16879        let mut xcat = self.uninit(total * in_f)?;
16880        let mut off = 0usize;
16881        for (x, &m) in xs.iter().zip(ms) {
16882            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16883            off += m;
16884        }
16885        let ys = self.matmul_group(ws, &xcat, total)?;
16886        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16887        for (w, y) in ws.iter().zip(ys) {
16888            let out_f = w.out_features();
16889            let mut off = 0usize;
16890            for (s, &m) in ms.iter().enumerate() {
16891                let mut ys_s = self.uninit(m * out_f)?;
16892                let src = y.slice(off * out_f..(off + m) * out_f);
16893                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16894                out[s].push(ys_s);
16895                off += m;
16896            }
16897        }
16898        Ok(out)
16899    }
16900
16901    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
16902    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
16903    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
16904    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
16905    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
16906    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
16907    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
16908    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
16909    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
16910    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16911        use crate::model::GpuTensor;
16912        if !legacy_quant_gemm_allowed(
16913            cfg!(memra_portable_cuda),
16914            cfg!(memra_hopper_mma),
16915            std::env::var_os("MEMRA_NO_GEMM").is_some(),
16916        ) {
16917            return false;
16918        }
16919        match w {
16920            GpuTensor::Quant { qtype, .. } => {
16921                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16922                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16923            }
16924            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16925        }
16926    }
16927
16928    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
16929    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
16930    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
16931    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
16932    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
16933    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
16934    pub fn qmatvec_gemm(
16935        &self,
16936        w: &crate::model::GpuTensor,
16937        aq: &CudaSlice<i8>,
16938        ad: &CudaSlice<f32>,
16939        m: usize,
16940    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16941        use crate::model::GpuTensor;
16942        let in_f = w.in_features();
16943        let out_f = w.out_features();
16944        let (bytes, qtype, row_bytes, scale, rp) = match w {
16945            GpuTensor::Quant {
16946                bytes,
16947                qtype,
16948                row_bytes,
16949                scale,
16950                rp,
16951                ..
16952            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16953            _ => unreachable!("gemm_supports guaranteed Quant"),
16954        };
16955        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
16956        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
16957        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
16958        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
16959        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
16960        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16961            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16962                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16963                if scale != 1.0 {
16964                    self.scale_inplace(&mut y, scale, m * out_f)?;
16965                }
16966                return Ok(y);
16967            }
16968        }
16969        let name = match qtype {
16970            QT_Q8_0 => "qmatvec_gemm_q8_0",
16971            QT_Q4_K => "qmatvec_gemm_q4_K",
16972            QT_Q4_0 => {
16973                if rp {
16974                    "qmatvec_gemm_q4_0_rp"
16975                } else {
16976                    "qmatvec_gemm_q4_0"
16977                }
16978            }
16979            QT_Q5_K => "qmatvec_gemm_q5_K",
16980            QT_Q6_K => "qmatvec_gemm_q6_K",
16981            QT_NVFP4 => {
16982                if rp {
16983                    "qmatvec_gemm_nvfp4_rp"
16984                } else {
16985                    "qmatvec_gemm_nvfp4"
16986                }
16987            }
16988            _ => unreachable!(),
16989        };
16990        let f = self.func(name);
16991        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16992        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
16993        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
16994        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
16995        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16996        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16997        let k1_tile = if is_k1 {
16998            k1_launch_override().unwrap_or((128, 128, 8))
16999        } else {
17000            (128, 128, 8)
17001        };
17002        let (bm, bn): (u32, u32) = if is_k1 {
17003            (k1_tile.0, k1_tile.1)
17004        } else {
17005            (64, 256)
17006        };
17007        let warps: u32 = if is_k1 {
17008            k1_tile.2
17009        } else {
17010            match qtype {
17011                QT_NVFP4 => 8,
17012                _ => 4,
17013            }
17014        };
17015        let cfg = LaunchConfig {
17016            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17017            block_dim: (32, warps, 1),
17018            shared_mem_bytes: 0,
17019        };
17020        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17021        let __s_b = self.gpu.stream();
17022        let mut b = __s_b.launch_builder(&f);
17023        b.arg(bytes)
17024            .arg(aq)
17025            .arg(ad)
17026            .arg(&mut y)
17027            .arg(&inf)
17028            .arg(&outf)
17029            .arg(&mi)
17030            .arg(&rb);
17031        unsafe {
17032            b.launch(cfg)?;
17033        }
17034        if scale != 1.0 {
17035            self.scale_inplace(&mut y, scale, m * out_f)?;
17036        }
17037        Ok(y)
17038    }
17039
17040    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
17041    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
17042    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
17043    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
17044    pub fn qmatvec_gemm_raw(
17045        &self,
17046        bytes: &CudaSlice<u8>,
17047        x: &CudaSlice<f32>,
17048        m: usize,
17049        in_f: usize,
17050        out_f: usize,
17051        qtype: i32,
17052        row_bytes: usize,
17053    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17054        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17055        let name = match qtype {
17056            QT_Q8_0 => "qmatvec_gemm_q8_0",
17057            QT_Q4_K => "qmatvec_gemm_q4_K",
17058            QT_Q4_0 => "qmatvec_gemm_q4_0",
17059            QT_Q5_K => "qmatvec_gemm_q5_K",
17060            QT_Q6_K => "qmatvec_gemm_q6_K",
17061            QT_NVFP4 => "qmatvec_gemm_nvfp4",
17062            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17063            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17064        };
17065        let f = self.func(name);
17066        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17067        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
17068        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
17069        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
17070        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
17071        let k1_tile = if is_k1 {
17072            k1_launch_override().unwrap_or((128, 128, 8))
17073        } else {
17074            (128, 128, 8)
17075        };
17076        let (bm, bn): (u32, u32) = if is_k1 {
17077            (k1_tile.0, k1_tile.1)
17078        } else {
17079            (64, 256)
17080        };
17081        let warps: u32 = if is_k1 {
17082            k1_tile.2
17083        } else {
17084            match qtype {
17085                QT_NVFP4 | QT_NVFP4_RP => 8,
17086                _ => 4,
17087            }
17088        };
17089        let cfg = LaunchConfig {
17090            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17091            block_dim: (32, warps, 1),
17092            shared_mem_bytes: 0,
17093        };
17094        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17095        let __s_b = self.gpu.stream();
17096        let mut b = __s_b.launch_builder(&f);
17097        b.arg(bytes)
17098            .arg(&aq)
17099            .arg(&ad)
17100            .arg(&mut y)
17101            .arg(&inf)
17102            .arg(&outf)
17103            .arg(&mi)
17104            .arg(&rb);
17105        unsafe {
17106            b.launch(cfg)?;
17107        }
17108        Ok(y)
17109    }
17110
17111    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
17112    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
17113    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
17114    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
17115    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
17116    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
17117    pub fn qmatvec_gemm_q8_0_wgmma_raw(
17118        &self,
17119        rp4: &CudaSlice<u8>,
17120        aq: &CudaSlice<i8>,
17121        ad: &CudaSlice<f32>,
17122        m: usize,
17123        in_f: usize,
17124        out_f: usize,
17125    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17126        assert!(
17127            out_f % 64 == 0 && in_f % 32 == 0,
17128            "wgmma GEMM needs out_f%64==0, in_f%32==0"
17129        );
17130        let f = self.func("qmatvec_gemm_q8_0_wgmma");
17131        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
17132        let cfg = LaunchConfig {
17133            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17134            block_dim: (128, 1, 1),
17135            shared_mem_bytes: 0,
17136        };
17137        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17138        let __s_b = self.gpu.stream();
17139        let mut b = __s_b.launch_builder(&f);
17140        b.arg(rp4)
17141            .arg(aq)
17142            .arg(ad)
17143            .arg(&mut y)
17144            .arg(&inf)
17145            .arg(&outf)
17146            .arg(&mi);
17147        unsafe {
17148            b.launch(cfg)?;
17149        }
17150        Ok(y)
17151    }
17152
17153    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
17154    pub fn scale_inplace(
17155        &self,
17156        y: &mut CudaSlice<f32>,
17157        s: f32,
17158        n: usize,
17159    ) -> Result<(), Box<dyn std::error::Error>> {
17160        let f = self.func("scale_f32");
17161        let cfg = LaunchConfig::for_num_elems(n as u32);
17162        let (sf, ni) = (s, n as i32);
17163        let __s_b = self.gpu.stream();
17164        let mut b = __s_b.launch_builder(&f);
17165        b.arg(y).arg(&sf).arg(&ni);
17166        unsafe {
17167            b.launch(cfg)?;
17168        }
17169        Ok(())
17170    }
17171
17172    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
17173    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
17174    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
17175    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
17176    pub fn bf16_to_f32(
17177        &self,
17178        data: &cudarc::driver::CudaView<'_, u8>,
17179        n: usize,
17180    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17181        let mut out = self.alloc_uninit::<f32>(n)?;
17182        let f = self.func("bf16_to_f32");
17183        let cfg = LaunchConfig::for_num_elems(n as u32);
17184        let ni = n as i32;
17185        let __s_b = self.gpu.stream();
17186        let mut b = __s_b.launch_builder(&f);
17187        b.arg(data).arg(&mut out).arg(&ni);
17188        unsafe {
17189            b.launch(cfg)?;
17190        }
17191        Ok(out)
17192    }
17193
17194    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
17195    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
17196    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
17197    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
17198    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
17199    /// calls, the spec-verify contract) vs plain linear.
17200    fn linear_bf16_chunked(
17201        &self,
17202        x: &CudaSlice<f32>,
17203        data: &CudaSlice<u8>,
17204        m: usize,
17205        in_f: usize,
17206        out_f: usize,
17207        exact: bool,
17208        canonical_chunk_rows: Option<usize>,
17209    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17210        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
17211        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
17212        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
17213        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17214        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17215        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17216        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17217        let started = timing.then(std::time::Instant::now);
17218        let result =
17219            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17220        if let Some(started) = started {
17221            use std::sync::atomic::Ordering;
17222            self.stream().synchronize()?;
17223            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17224                + started.elapsed().as_nanos() as u64;
17225            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17226                + (in_f * out_f * 2) as u64;
17227            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17228            if calls % 1024 == 0 {
17229                eprintln!(
17230                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17231                     weight_gb={:.2}",
17232                    ns as f64 / 1.0e6,
17233                    ns as f64 / calls as f64 / 1.0e3,
17234                    wb as f64 / 1.0e9,
17235                );
17236            }
17237        }
17238        result
17239    }
17240
17241    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
17242    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
17243    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
17244    /// numeric-class doors (DEV_ROUTES precedent).
17245    pub(crate) fn bf16_mmv_on() -> bool {
17246        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17247        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17248    }
17249
17250    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
17251    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
17252    fn matvec_bf16(
17253        &self,
17254        data: &CudaSlice<u8>,
17255        x: &CudaSlice<f32>,
17256        in_f: usize,
17257        out_f: usize,
17258    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17259        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17260            return Err(format!(
17261                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17262                data.len(),
17263                x.len()
17264            )
17265            .into());
17266        }
17267        let mut y = self.alloc_uninit::<f32>(out_f)?;
17268        let f = self.func("matvec_bf16_f32acc");
17269        let cfg = LaunchConfig {
17270            grid_dim: (out_f as u32, 1, 1),
17271            block_dim: (mmv_block(), 1, 1),
17272            shared_mem_bytes: 0,
17273        };
17274        let ini = in_f as i32;
17275        let __s_bld = self.gpu.stream();
17276        let mut bld = __s_bld.launch_builder(&f);
17277        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17278        unsafe {
17279            bld.launch(cfg)?;
17280        }
17281        Ok(y)
17282    }
17283
17284    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
17285    /// launches, a position upload, and the rope launch; the position is read directly from
17286    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
17287    #[allow(clippy::too_many_arguments)]
17288    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
17289    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
17290    /// Bit-identical to the split kernels; requires head_dim == 128 and
17291    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
17292    #[allow(clippy::too_many_arguments)]
17293    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
17294    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
17295    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
17296    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
17297    #[allow(clippy::too_many_arguments)]
17298    pub fn qk_norm_rope_append_inc_dcw_rows(
17299        &self,
17300        q_raw_t: &CudaSlice<f32>,
17301        k_raw_t: &CudaSlice<f32>,
17302        v_raw_t: &CudaSlice<f32>,
17303        qw: &CudaSlice<f32>,
17304        kw: &CudaSlice<f32>,
17305        q_out_t: &mut CudaSlice<f32>,
17306        k_out_t: &mut CudaSlice<f32>,
17307        tab: &CudaSlice<u64>,
17308        pos_t: &CudaSlice<i32>,
17309        same_session: bool,
17310        t: usize,
17311        kv_dim_k: usize,
17312        kv_dim_v: usize,
17313        k_tok_bytes: usize,
17314        v_tok_bytes: usize,
17315        head_dim: usize,
17316        n_dims: usize,
17317        nh_q: usize,
17318        nh_k: usize,
17319        eps: f32,
17320        freq_base: f32,
17321        freq_scale: f32,
17322        ff: Option<&CudaSlice<f32>>,
17323    ) -> Result<(), Box<dyn std::error::Error>> {
17324        if head_dim != 128
17325            || kv_dim_v != kv_dim_k
17326            || kv_dim_k != nh_k * head_dim
17327            || t == 0
17328            || t > 32
17329            || tab.len() < t * 6
17330            || pos_t.len() < t
17331            || q_raw_t.len() < t * nh_q * head_dim
17332            || k_raw_t.len() < t * nh_k * head_dim
17333            || v_raw_t.len() < t * kv_dim_v
17334            || q_out_t.len() < t * nh_q * head_dim
17335            || k_out_t.len() < t * nh_k * head_dim
17336        {
17337            return Err(format!(
17338                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17339                 nh_q={nh_q} nh_k={nh_k}"
17340            )
17341            .into());
17342        }
17343        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17344        let same_t: i32 = if same_session { t as i32 } else { 0 };
17345        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17346        let cfg = LaunchConfig {
17347            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17348            block_dim: (128, 1, 1),
17349            shared_mem_bytes: 0,
17350        };
17351        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17352        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17353        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17354        let null: u64 = 0;
17355        let __s_b = self.gpu.stream();
17356        let mut b = __s_b.launch_builder(&f);
17357        b.arg(q_raw_t)
17358            .arg(k_raw_t)
17359            .arg(v_raw_t)
17360            .arg(qw)
17361            .arg(kw)
17362            .arg(q_out_t)
17363            .arg(k_out_t)
17364            .arg(tab)
17365            .arg(pos_t)
17366            .arg(&same_t)
17367            .arg(&kvk)
17368            .arg(&kvv)
17369            .arg(&ktb)
17370            .arg(&vtb)
17371            .arg(&hd)
17372            .arg(&nd)
17373            .arg(&nq)
17374            .arg(&nk)
17375            .arg(&eps)
17376            .arg(&theta_scale)
17377            .arg(&freq_scale);
17378        match ff {
17379            Some(freqs) => {
17380                b.arg(freqs);
17381            }
17382            None => {
17383                b.arg(&null);
17384            }
17385        }
17386        unsafe {
17387            b.launch(cfg)?;
17388        }
17389        Ok(())
17390    }
17391
17392    pub fn qk_norm_rope_append_inc_dcw(
17393        &self,
17394        q_raw: &CudaSlice<f32>,
17395        k_raw: &CudaSlice<f32>,
17396        v_raw: &CudaSlice<f32>,
17397        qw: &CudaSlice<f32>,
17398        kw: &CudaSlice<f32>,
17399        q_out: &mut CudaSlice<f32>,
17400        k_out: &mut CudaSlice<f32>,
17401        pos: &CudaSlice<i32>,
17402        k_plane: &mut CudaSlice<u8>,
17403        v_plane: &mut CudaSlice<u8>,
17404        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
17405        // (single) writer, exactly like the split append+inc pair it replaces.
17406        len_dev: &CudaSlice<i32>,
17407        base_dev: Option<&CudaSlice<i32>>,
17408        done_ctr: &mut CudaSlice<u32>,
17409        kv_dim_k: usize,
17410        kv_dim_v: usize,
17411        k_tok_bytes: usize,
17412        v_tok_bytes: usize,
17413        head_dim: usize,
17414        n_dims: usize,
17415        nh_q: usize,
17416        nh_k: usize,
17417        eps: f32,
17418        freq_base: f32,
17419        freq_scale: f32,
17420        ff: Option<&CudaSlice<f32>>,
17421    ) -> Result<(), Box<dyn std::error::Error>> {
17422        if head_dim != 128
17423            || kv_dim_v != kv_dim_k
17424            || kv_dim_k != nh_k * head_dim
17425            || q_raw.len() < nh_q * head_dim
17426            || k_raw.len() < nh_k * head_dim
17427            || v_raw.len() < kv_dim_v
17428            || q_out.len() < nh_q * head_dim
17429            || k_out.len() < nh_k * head_dim
17430            || pos.is_empty()
17431            || done_ctr.is_empty()
17432        {
17433            return Err(format!(
17434                "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}"
17435            )
17436            .into());
17437        }
17438        let f = self.func("qk_norm_rope_append_inc_dcw");
17439        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17440        let cfg = LaunchConfig {
17441            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17442            block_dim: (128, 1, 1),
17443            shared_mem_bytes: 0,
17444        };
17445        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17446        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17447        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17448        let null: u64 = 0;
17449        let __s_b = self.gpu.stream();
17450        let mut b = __s_b.launch_builder(&f);
17451        b.arg(q_raw)
17452            .arg(k_raw)
17453            .arg(v_raw)
17454            .arg(qw)
17455            .arg(kw)
17456            .arg(q_out)
17457            .arg(k_out)
17458            .arg(pos)
17459            .arg(&mut *k_plane)
17460            .arg(&mut *v_plane)
17461            .arg(len_dev);
17462        match base_dev {
17463            Some(base) => {
17464                b.arg(base);
17465            }
17466            None => {
17467                b.arg(&null);
17468            }
17469        }
17470        b.arg(&mut *done_ctr)
17471            .arg(&kvk)
17472            .arg(&kvv)
17473            .arg(&ktb)
17474            .arg(&vtb)
17475            .arg(&hd)
17476            .arg(&nd)
17477            .arg(&nq)
17478            .arg(&eps)
17479            .arg(&theta_scale)
17480            .arg(&freq_scale);
17481        match ff {
17482            Some(freqs) => {
17483                b.arg(freqs);
17484            }
17485            None => {
17486                b.arg(&null);
17487            }
17488        }
17489        unsafe {
17490            b.launch(cfg)?;
17491        }
17492        Ok(())
17493    }
17494
17495    pub fn qk_norm_rope_into(
17496        &self,
17497        q_raw: &CudaSlice<f32>,
17498        k_raw: &CudaSlice<f32>,
17499        qw: &CudaSlice<f32>,
17500        kw: &CudaSlice<f32>,
17501        q_out: &mut CudaSlice<f32>,
17502        k_out: &mut CudaSlice<f32>,
17503        pos: &CudaSlice<i32>,
17504        head_dim: usize,
17505        n_dims: usize,
17506        nh_q: usize,
17507        nh_k: usize,
17508        eps: f32,
17509        freq_base: f32,
17510        freq_scale: f32,
17511        ff: Option<&CudaSlice<f32>>,
17512    ) -> Result<(), Box<dyn std::error::Error>> {
17513        if head_dim > 512
17514            || q_raw.len() < nh_q * head_dim
17515            || k_raw.len() < nh_k * head_dim
17516            || q_out.len() < nh_q * head_dim
17517            || k_out.len() < nh_k * head_dim
17518            || qw.len() < head_dim
17519            || kw.len() < head_dim
17520            || pos.is_empty()
17521        {
17522            return Err(format!(
17523                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17524            )
17525            .into());
17526        }
17527        let f = self.func("qk_norm_rope_f32");
17528        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17529        let cfg = LaunchConfig {
17530            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17531            block_dim: (128, 1, 1),
17532            shared_mem_bytes: 0,
17533        };
17534        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17535        let __s_b = self.gpu.stream();
17536        let mut b = __s_b.launch_builder(&f);
17537        b.arg(q_raw)
17538            .arg(k_raw)
17539            .arg(qw)
17540            .arg(kw)
17541            .arg(q_out)
17542            .arg(k_out)
17543            .arg(pos)
17544            .arg(&hd)
17545            .arg(&nd)
17546            .arg(&nq)
17547            .arg(&eps)
17548            .arg(&theta_scale)
17549            .arg(&freq_scale);
17550        match ff {
17551            Some(ffv) => {
17552                b.arg(ffv);
17553                unsafe {
17554                    b.launch(cfg)?;
17555                }
17556            }
17557            None => {
17558                let null: u64 = 0;
17559                b.arg(&null);
17560                unsafe {
17561                    b.launch(cfg)?;
17562                }
17563            }
17564        }
17565        Ok(())
17566    }
17567
17568    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
17569    /// launch computes a rank's whole O partial from its four canonical column blocks.
17570    #[allow(clippy::too_many_arguments)]
17571    pub fn matvec_f32_b4_into(
17572        &self,
17573        w: [&CudaSlice<f32>; 4],
17574        x: &CudaSlice<f32>,
17575        y: &mut CudaSlice<f32>,
17576        block_cols: usize,
17577        out_f: usize,
17578    ) -> Result<(), Box<dyn std::error::Error>> {
17579        if block_cols % 4 != 0
17580            || x.len() < 4 * block_cols
17581            || y.len() < out_f
17582            || w.iter().any(|w| w.len() != out_f * block_cols)
17583        {
17584            return Err(format!(
17585                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17586                x.len()
17587            )
17588            .into());
17589        }
17590        let f = self.func("matvec_f32_b4");
17591        let cfg = LaunchConfig {
17592            grid_dim: (out_f as u32, 1, 1),
17593            block_dim: (128, 1, 1),
17594            shared_mem_bytes: 0,
17595        };
17596        let (bc, of) = (block_cols as i32, out_f as i32);
17597        let __s_b = self.gpu.stream();
17598        let mut b = __s_b.launch_builder(&f);
17599        b.arg(w[0])
17600            .arg(w[1])
17601            .arg(w[2])
17602            .arg(w[3])
17603            .arg(x)
17604            .arg(y)
17605            .arg(&bc)
17606            .arg(&of);
17607        unsafe {
17608            b.launch(cfg)?;
17609        }
17610        Ok(())
17611    }
17612
17613    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
17614    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
17615    pub fn axpy_rows_seq_into(
17616        &self,
17617        x: &CudaSlice<f32>,
17618        w: &CudaSlice<f32>,
17619        y: &mut CudaSlice<f32>,
17620        width: usize,
17621        n_rows: usize,
17622    ) -> Result<(), Box<dyn std::error::Error>> {
17623        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17624            return Err(format!(
17625                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17626                x.len(),
17627                w.len(),
17628                y.len()
17629            )
17630            .into());
17631        }
17632        let f = self.func("axpy_rows_seq_f32");
17633        let cfg = LaunchConfig::for_num_elems(width as u32);
17634        let (wi, nr) = (width as i32, n_rows as i32);
17635        let __s_b = self.gpu.stream();
17636        let mut b = __s_b.launch_builder(&f);
17637        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17638        unsafe {
17639            b.launch(cfg)?;
17640        }
17641        Ok(())
17642    }
17643
17644    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
17645    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
17646    /// exact sequential FP chain of the base kernel over that window.
17647    #[allow(clippy::too_many_arguments)]
17648    pub fn axpy_rows_seq_md_off_into(
17649        &self,
17650        x: &CudaSlice<f32>,
17651        w_route: &CudaSlice<f32>,
17652        md: &CudaSlice<f32>,
17653        sel: &CudaSlice<i32>,
17654        y: &mut CudaSlice<f32>,
17655        width: usize,
17656        n_rows: usize,
17657        row0: usize,
17658    ) -> Result<(), Box<dyn std::error::Error>> {
17659        if x.len() < (row0 + n_rows) * width
17660            || w_route.len() < row0 + n_rows
17661            || sel.len() < row0 + n_rows
17662            || y.len() < width
17663        {
17664            return Err(format!(
17665                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17666                 rows={n_rows} row0={row0}",
17667                x.len(),
17668                w_route.len(),
17669                sel.len(),
17670                y.len()
17671            )
17672            .into());
17673        }
17674        let f = self.func("axpy_rows_seq_md_off_f32");
17675        let cfg = LaunchConfig::for_num_elems(width as u32);
17676        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17677        let __s_b = self.gpu.stream();
17678        let mut b = __s_b.launch_builder(&f);
17679        b.arg(x)
17680            .arg(w_route)
17681            .arg(md)
17682            .arg(sel)
17683            .arg(y)
17684            .arg(&wi)
17685            .arg(&nr)
17686            .arg(&r0);
17687        unsafe {
17688            b.launch(cfg)?;
17689        }
17690        Ok(())
17691    }
17692
17693    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
17694    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
17695    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
17696    /// outputs are bit-equal to its own t=1 launch.
17697    #[allow(clippy::too_many_arguments)]
17698    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17699        &self,
17700        gate_bank: &CudaSlice<u8>,
17701        up_bank: &CudaSlice<u8>,
17702        sel: &CudaSlice<i32>,
17703        aq: &CudaSlice<i8>,
17704        ad: &CudaSlice<f32>,
17705        yg: &mut CudaSlice<f32>,
17706        yu: &mut CudaSlice<f32>,
17707        n_sel: usize,
17708        n_sel_col: usize,
17709        in_f: usize,
17710        out_f: usize,
17711        row_bytes: usize,
17712        expert_stride: usize,
17713        act_row_stride: usize,
17714        ad_row_stride: usize,
17715    ) -> Result<(), Box<dyn std::error::Error>> {
17716        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17717        if yg.len() < n_sel * out_f
17718            || yu.len() < n_sel * out_f
17719            || sel.len() < n_sel
17720            || n_sel_col == 0
17721            || n_sel % n_sel_col != 0
17722        {
17723            return Err("NVFP4 gu tcol geometry".into());
17724        }
17725        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17726        let cfg = LaunchConfig {
17727            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17728            block_dim: (128, 1, 1),
17729            shared_mem_bytes: 0,
17730        };
17731        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17732        let (rb, es) = (row_bytes as i64, expert_stride as i64);
17733        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17734        let __s_b = self.gpu.stream();
17735        let mut b = __s_b.launch_builder(&f);
17736        b.arg(gate_bank)
17737            .arg(up_bank)
17738            .arg(sel)
17739            .arg(aq)
17740            .arg(ad)
17741            .arg(yg)
17742            .arg(yu)
17743            .arg(&inf)
17744            .arg(&outf)
17745            .arg(&ns)
17746            .arg(&rb)
17747            .arg(&es)
17748            .arg(&ars)
17749            .arg(&adrs)
17750            .arg(&nsc);
17751        unsafe {
17752            b.launch(cfg)?;
17753        }
17754        Ok(())
17755    }
17756
17757    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
17758    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
17759    #[allow(clippy::too_many_arguments)]
17760    pub fn axpy_rows_seq_md_into(
17761        &self,
17762        x: &CudaSlice<f32>,
17763        w_route: &CudaSlice<f32>,
17764        md: &CudaSlice<f32>,
17765        sel: &CudaSlice<i32>,
17766        y: &mut CudaSlice<f32>,
17767        width: usize,
17768        n_rows: usize,
17769    ) -> Result<(), Box<dyn std::error::Error>> {
17770        if x.len() < n_rows * width
17771            || w_route.len() < n_rows
17772            || sel.len() < n_rows
17773            || y.len() < width
17774        {
17775            return Err(format!(
17776                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17777                x.len(),
17778                w_route.len(),
17779                sel.len(),
17780                y.len()
17781            )
17782            .into());
17783        }
17784        let f = self.func("axpy_rows_seq_md_f32");
17785        let cfg = LaunchConfig::for_num_elems(width as u32);
17786        let (wi, nr) = (width as i32, n_rows as i32);
17787        let __s_b = self.gpu.stream();
17788        let mut b = __s_b.launch_builder(&f);
17789        b.arg(x)
17790            .arg(w_route)
17791            .arg(md)
17792            .arg(sel)
17793            .arg(y)
17794            .arg(&wi)
17795            .arg(&nr);
17796        unsafe {
17797            b.launch(cfg)?;
17798        }
17799        Ok(())
17800    }
17801
17802    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
17803    #[allow(clippy::too_many_arguments)]
17804    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
17805    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
17806    /// land column-major-of-rows: yq[c*out_q + row] etc.
17807    #[allow(clippy::too_many_arguments)]
17808    pub fn matvec_bf16_qkvg_tcol_into(
17809        &self,
17810        wq: &CudaSlice<u8>,
17811        wk: &CudaSlice<u8>,
17812        wv: &CudaSlice<u8>,
17813        wg: &CudaSlice<u8>,
17814        x_t: &CudaSlice<f32>,
17815        yq: &mut CudaSlice<f32>,
17816        yk: &mut CudaSlice<f32>,
17817        yv: &mut CudaSlice<f32>,
17818        yg: &mut CudaSlice<f32>,
17819        in_f: usize,
17820        out_q: usize,
17821        out_kv: usize,
17822        out_g: usize,
17823        t: usize,
17824    ) -> Result<(), Box<dyn std::error::Error>> {
17825        if t == 0
17826            || t > 8
17827            || in_f % 8 != 0
17828            || x_t.len() < t * in_f
17829            || yq.len() < t * out_q
17830            || yk.len() < t * out_kv
17831            || yv.len() < t * out_kv
17832            || (out_g > 0 && yg.len() < t * out_g)
17833        {
17834            return Err("matvec_bf16_qkvg_tcol geometry".into());
17835        }
17836        let grid = out_q + 2 * out_kv + out_g;
17837        let cfg = LaunchConfig {
17838            grid_dim: (grid as u32, 1, 1),
17839            block_dim: (mmv_block(), 1, 1),
17840            shared_mem_bytes: 0,
17841        };
17842        let (ini, oq, okv, og, ti) = (
17843            in_f as i32,
17844            out_q as i32,
17845            out_kv as i32,
17846            out_g as i32,
17847            t as i32,
17848        );
17849        let __s_b = self.gpu.stream();
17850        // Compile-time-T twins keep the accumulators in registers (bit-identical chain).
17851        if let Some(name) = match t {
17852            2 => Some("matvec_bf16_qkvg_tcol_t2"),
17853            4 => Some("matvec_bf16_qkvg_tcol_t4"),
17854            8 => Some("matvec_bf16_qkvg_tcol_t8"),
17855            _ => None,
17856        } {
17857            let f = self.func(name);
17858            let mut b = __s_b.launch_builder(&f);
17859            b.arg(wq)
17860                .arg(wk)
17861                .arg(wv)
17862                .arg(wg)
17863                .arg(x_t)
17864                .arg(yq)
17865                .arg(yk)
17866                .arg(yv)
17867                .arg(yg)
17868                .arg(&ini)
17869                .arg(&oq)
17870                .arg(&okv)
17871                .arg(&og);
17872            unsafe {
17873                b.launch(cfg)?;
17874            }
17875            return Ok(());
17876        }
17877        let f = self.func("matvec_bf16_qkvg_tcol");
17878        let mut b = __s_b.launch_builder(&f);
17879        b.arg(wq)
17880            .arg(wk)
17881            .arg(wv)
17882            .arg(wg)
17883            .arg(x_t)
17884            .arg(yq)
17885            .arg(yk)
17886            .arg(yv)
17887            .arg(yg)
17888            .arg(&ini)
17889            .arg(&oq)
17890            .arg(&okv)
17891            .arg(&og)
17892            .arg(&ti);
17893        unsafe {
17894            b.launch(cfg)?;
17895        }
17896        Ok(())
17897    }
17898
17899    pub fn matvec_bf16_qkvg_into(
17900        &self,
17901        wq: &CudaSlice<u8>,
17902        wk: &CudaSlice<u8>,
17903        wv: &CudaSlice<u8>,
17904        wg: &CudaSlice<u8>,
17905        x: &CudaSlice<f32>,
17906        yq: &mut CudaSlice<f32>,
17907        yk: &mut CudaSlice<f32>,
17908        yv: &mut CudaSlice<f32>,
17909        yg: &mut CudaSlice<f32>,
17910        in_f: usize,
17911        out_q: usize,
17912        out_kv: usize,
17913        out_g: usize,
17914    ) -> Result<(), Box<dyn std::error::Error>> {
17915        if in_f % 8 != 0
17916            || wq.len() != out_q * in_f * 2
17917            || wk.len() != out_kv * in_f * 2
17918            || wv.len() != out_kv * in_f * 2
17919            || wg.len() < out_g * in_f * 2
17920            || x.len() < in_f
17921            || yq.len() < out_q
17922            || yk.len() < out_kv
17923            || yv.len() < out_kv
17924            || (out_g > 0 && yg.len() < out_g)
17925        {
17926            return Err(format!(
17927                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17928            )
17929            .into());
17930        }
17931        let f = self.func("matvec_bf16_qkvg");
17932        let cfg = LaunchConfig {
17933            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17934            block_dim: (mmv_block(), 1, 1),
17935            shared_mem_bytes: 0,
17936        };
17937        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17938        let __s_b = self.gpu.stream();
17939        let mut b = __s_b.launch_builder(&f);
17940        b.arg(wq)
17941            .arg(wk)
17942            .arg(wv)
17943            .arg(wg)
17944            .arg(x)
17945            .arg(yq)
17946            .arg(yk)
17947            .arg(yv)
17948            .arg(yg)
17949            .arg(&inf)
17950            .arg(&oq)
17951            .arg(&okv)
17952            .arg(&og);
17953        unsafe {
17954            b.launch(cfg)?;
17955        }
17956        Ok(())
17957    }
17958
17959    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
17960    pub fn matvec_bf16_b4_into(
17961        &self,
17962        w: [&CudaSlice<u8>; 4],
17963        x: &CudaSlice<f32>,
17964        y: &mut CudaSlice<f32>,
17965        block_cols: usize,
17966        out_f: usize,
17967    ) -> Result<(), Box<dyn std::error::Error>> {
17968        if block_cols % 8 != 0
17969            || x.len() < 4 * block_cols
17970            || y.len() < out_f
17971            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17972        {
17973            return Err(format!(
17974                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17975                x.len()
17976            )
17977            .into());
17978        }
17979        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
17980        // bit-identical per row (the second row's stream hides the first's reduce tail).
17981        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17982        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17983        let f = self.func(if x2 {
17984            "matvec_bf16_b4_x2"
17985        } else {
17986            "matvec_bf16_b4"
17987        });
17988        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17989        let cfg = LaunchConfig {
17990            grid_dim: (grid as u32, 1, 1),
17991            block_dim: (mmv_block(), 1, 1),
17992            shared_mem_bytes: 0,
17993        };
17994        let (bc, of) = (block_cols as i32, out_f as i32);
17995        let __s_b = self.gpu.stream();
17996        let mut b = __s_b.launch_builder(&f);
17997        b.arg(w[0])
17998            .arg(w[1])
17999            .arg(w[2])
18000            .arg(w[3])
18001            .arg(x)
18002            .arg(y)
18003            .arg(&bc)
18004            .arg(&of);
18005        unsafe {
18006            b.launch(cfg)?;
18007        }
18008        Ok(())
18009    }
18010
18011    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
18012    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
18013    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
18014    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
18015    /// t=1 program).
18016    pub fn matvec_bf16_b4_tcol_into(
18017        &self,
18018        w: [&CudaSlice<u8>; 4],
18019        x_t: &CudaSlice<f32>,
18020        y_t: &mut CudaSlice<f32>,
18021        block_cols: usize,
18022        out_f: usize,
18023        t: usize,
18024    ) -> Result<(), Box<dyn std::error::Error>> {
18025        if block_cols % 8 != 0
18026            || t == 0
18027            || t > 8
18028            || x_t.len() < t * 4 * block_cols
18029            || y_t.len() < t * out_f
18030            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18031        {
18032            return Err(format!(
18033                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18034                x_t.len()
18035            )
18036            .into());
18037        }
18038        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18039            return Err(
18040                "b4 tcol verify is qualified against the plain b4 kernel only \
18041                        (MEMRA_B4_X2=1 is a different t=1 program)"
18042                    .into(),
18043            );
18044        }
18045        // Compile-time-T twins for the walk widths: the runtime-t inner loop spills the
18046        // per-column accumulators to local memory (283us vs 33 at t=8). Same FP chain
18047        // per (block, column) — bit-identical.
18048        let cfg = LaunchConfig {
18049            grid_dim: (out_f as u32, 1, 1),
18050            block_dim: (mmv_block(), 1, 1),
18051            shared_mem_bytes: 0,
18052        };
18053        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18054        let __s_b = self.gpu.stream();
18055        if let Some(name) = match t {
18056            2 => Some("matvec_bf16_b4_tcol_t2"),
18057            4 => Some("matvec_bf16_b4_tcol_t4"),
18058            8 => Some("matvec_bf16_b4_tcol_t8"),
18059            _ => None,
18060        } {
18061            let f = self.func(name);
18062            let mut b = __s_b.launch_builder(&f);
18063            b.arg(w[0])
18064                .arg(w[1])
18065                .arg(w[2])
18066                .arg(w[3])
18067                .arg(x_t)
18068                .arg(y_t)
18069                .arg(&bc)
18070                .arg(&of);
18071            unsafe {
18072                b.launch(cfg)?;
18073            }
18074            return Ok(());
18075        }
18076        let f = self.func("matvec_bf16_b4_tcol");
18077        let mut b = __s_b.launch_builder(&f);
18078        b.arg(w[0])
18079            .arg(w[1])
18080            .arg(w[2])
18081            .arg(w[3])
18082            .arg(x_t)
18083            .arg(y_t)
18084            .arg(&bc)
18085            .arg(&of)
18086            .arg(&ti);
18087        unsafe {
18088            b.launch(cfg)?;
18089        }
18090        Ok(())
18091    }
18092
18093    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
18094    pub fn matvec_bf16_into(
18095        &self,
18096        data: &CudaSlice<u8>,
18097        x: &CudaSlice<f32>,
18098        y: &mut CudaSlice<f32>,
18099        in_f: usize,
18100        out_f: usize,
18101    ) -> Result<(), Box<dyn std::error::Error>> {
18102        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18103            return Err(format!(
18104                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18105                data.len(),
18106                x.len(),
18107                y.len()
18108            )
18109            .into());
18110        }
18111        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
18112        // block, exact f32acc per-row program — cures the 1-iteration latency
18113        // starvation (shexp down measured 420GB/s at in_f=1280).
18114        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18115        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18116            && in_f <= 2048;
18117        if x4 {
18118            let f = self.func("matvec_bf16_f32acc_x4");
18119            let cfg = LaunchConfig {
18120                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18121                block_dim: (mmv_block(), 1, 1),
18122                shared_mem_bytes: 0,
18123            };
18124            let (ini, outi) = (in_f as i32, out_f as i32);
18125            let __s_b = self.gpu.stream();
18126            let mut b = __s_b.launch_builder(&f);
18127            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18128            unsafe {
18129                b.launch(cfg)?;
18130            }
18131            return Ok(());
18132        }
18133        let f = self.func("matvec_bf16_f32acc");
18134        let cfg = LaunchConfig {
18135            grid_dim: (out_f as u32, 1, 1),
18136            block_dim: (mmv_block(), 1, 1),
18137            shared_mem_bytes: 0,
18138        };
18139        let ini = in_f as i32;
18140        let __s_b = self.gpu.stream();
18141        let mut b = __s_b.launch_builder(&f);
18142        b.arg(data).arg(x).arg(y).arg(&ini);
18143        unsafe {
18144            b.launch(cfg)?;
18145        }
18146        Ok(())
18147    }
18148
18149    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
18150    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
18151    pub fn matvec_bf16_view_into(
18152        &self,
18153        data: &cudarc::driver::CudaView<'_, u8>,
18154        x: &CudaSlice<f32>,
18155        y: &mut CudaSlice<f32>,
18156        in_f: usize,
18157        out_f: usize,
18158    ) -> Result<(), Box<dyn std::error::Error>> {
18159        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18160            return Err(format!(
18161                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18162                data.len(),
18163                x.len(),
18164                y.len()
18165            )
18166            .into());
18167        }
18168        let f = self.func("matvec_bf16_f32acc");
18169        let cfg = LaunchConfig {
18170            grid_dim: (out_f as u32, 1, 1),
18171            block_dim: (mmv_block(), 1, 1),
18172            shared_mem_bytes: 0,
18173        };
18174        let ini = in_f as i32;
18175        let __s_b = self.gpu.stream();
18176        let mut b = __s_b.launch_builder(&f);
18177        b.arg(data).arg(x).arg(y).arg(&ini);
18178        unsafe {
18179            b.launch(cfg)?;
18180        }
18181        Ok(())
18182    }
18183
18184    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
18185    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
18186    pub fn matvec_bf16_raw_out(
18187        &self,
18188        w: &CudaSlice<u8>,
18189        x: &CudaSlice<f32>,
18190        y_raw: u64,
18191        in_f: usize,
18192        out_f: usize,
18193    ) -> Result<(), Box<dyn std::error::Error>> {
18194        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18195            return Err("matvec_bf16_raw_out geometry".into());
18196        }
18197        let f = self.func("matvec_bf16_f32acc");
18198        let cfg = LaunchConfig {
18199            grid_dim: (out_f as u32, 1, 1),
18200            block_dim: (mmv_block(), 1, 1),
18201            shared_mem_bytes: 0,
18202        };
18203        let ini = in_f as i32;
18204        let __s_b = self.gpu.stream();
18205        let mut b = __s_b.launch_builder(&f);
18206        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18207        unsafe {
18208            b.launch(cfg)?;
18209        }
18210        Ok(())
18211    }
18212
18213    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
18214    /// UVA pointers so the caller passes persistent-static rows without holding locks).
18215    /// Exact per-element sequence of the split add + add_scaled_rows pair.
18216    pub fn add3_raw(
18217        &self,
18218        a: &CudaSlice<f32>,
18219        b: &CudaSlice<f32>,
18220        sh_raw: u64,
18221        scale_raw: u64,
18222        dst: &mut CudaSlice<f32>,
18223        n: usize,
18224    ) -> Result<(), Box<dyn std::error::Error>> {
18225        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18226            return Err("add3_raw geometry".into());
18227        }
18228        let f = self.func("add3_f32");
18229        let cfg = LaunchConfig {
18230            grid_dim: ((n as u32).div_ceil(256), 1, 1),
18231            block_dim: (256, 1, 1),
18232            shared_mem_bytes: 0,
18233        };
18234        let ni = n as i32;
18235        let __s_b = self.gpu.stream();
18236        let mut bld = __s_b.launch_builder(&f);
18237        bld.arg(a)
18238            .arg(b)
18239            .arg(&sh_raw)
18240            .arg(&scale_raw)
18241            .arg(dst)
18242            .arg(&ni);
18243        unsafe {
18244            bld.launch(cfg)?;
18245        }
18246        Ok(())
18247    }
18248
18249    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
18250    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
18251    pub fn matvec_bf16_down_addscale_into(
18252        &self,
18253        w: &CudaSlice<u8>,
18254        x: &CudaSlice<f32>,
18255        scale: &CudaSlice<f32>,
18256        dst: &mut CudaSlice<f32>,
18257        in_f: usize,
18258        out_f: usize,
18259    ) -> Result<(), Box<dyn std::error::Error>> {
18260        if w.len() != in_f * out_f * 2
18261            || x.len() < in_f
18262            || in_f % 8 != 0
18263            || dst.len() < out_f
18264            || scale.is_empty()
18265        {
18266            return Err("matvec_bf16_down_addscale geometry".into());
18267        }
18268        let f = self.func("matvec_bf16_down_addscale");
18269        let cfg = LaunchConfig {
18270            grid_dim: (out_f as u32, 1, 1),
18271            block_dim: (mmv_block(), 1, 1),
18272            shared_mem_bytes: 0,
18273        };
18274        let ini = in_f as i32;
18275        let __s_b = self.gpu.stream();
18276        let mut b = __s_b.launch_builder(&f);
18277        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18278        unsafe {
18279            b.launch(cfg)?;
18280        }
18281        Ok(())
18282    }
18283
18284    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
18285    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
18286    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
18287    #[allow(clippy::too_many_arguments)]
18288    pub fn matvec_bf16_dual_silu_rows_into(
18289        &self,
18290        wg: &CudaSlice<u8>,
18291        wu: &CudaSlice<u8>,
18292        x: &CudaSlice<f32>,
18293        act: &mut CudaSlice<f32>,
18294        in_f: usize,
18295        out_f: usize,
18296        limit: Option<f32>,
18297        t: usize,
18298    ) -> Result<(), Box<dyn std::error::Error>> {
18299        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18300            return Err("matvec_bf16_dual_silu_rows geometry".into());
18301        }
18302        let f = self.func("matvec_bf16_dual_silu_rows");
18303        let cfg = LaunchConfig {
18304            grid_dim: (out_f as u32, t as u32, 1),
18305            block_dim: (mmv_block(), 1, 1),
18306            shared_mem_bytes: 0,
18307        };
18308        let (ini, outi) = (in_f as i32, out_f as i32);
18309        let lim = limit.unwrap_or(0.0);
18310        let __s_b = self.gpu.stream();
18311        let mut b = __s_b.launch_builder(&f);
18312        b.arg(wg)
18313            .arg(wu)
18314            .arg(x)
18315            .arg(&mut *act)
18316            .arg(&ini)
18317            .arg(&outi)
18318            .arg(&lim);
18319        unsafe {
18320            b.launch(cfg)?;
18321        }
18322        Ok(())
18323    }
18324
18325    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
18326    pub fn matvec_bf16_rows_into(
18327        &self,
18328        w: &CudaSlice<u8>,
18329        x: &CudaSlice<f32>,
18330        y: &mut CudaSlice<f32>,
18331        in_f: usize,
18332        out_f: usize,
18333        t: usize,
18334    ) -> Result<(), Box<dyn std::error::Error>> {
18335        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18336            return Err("matvec_bf16_rows geometry".into());
18337        }
18338        let f = self.func("matvec_bf16_f32acc_x4_rows");
18339        let cfg = LaunchConfig {
18340            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18341            block_dim: (mmv_block(), 1, 1),
18342            shared_mem_bytes: 0,
18343        };
18344        let (ini, outi) = (in_f as i32, out_f as i32);
18345        let __s_b = self.gpu.stream();
18346        let mut b = __s_b.launch_builder(&f);
18347        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18348        unsafe {
18349            b.launch(cfg)?;
18350        }
18351        Ok(())
18352    }
18353
18354    pub fn matvec_bf16_dual_silu_into(
18355        &self,
18356        wg: &CudaSlice<u8>,
18357        wu: &CudaSlice<u8>,
18358        x: &CudaSlice<f32>,
18359        act: &mut CudaSlice<f32>,
18360        in_f: usize,
18361        out_f: usize,
18362        limit: Option<f32>,
18363    ) -> Result<(), Box<dyn std::error::Error>> {
18364        if wg.len() != in_f * out_f * 2
18365            || wu.len() != in_f * out_f * 2
18366            || x.len() < in_f
18367            || in_f % 8 != 0
18368            || act.len() < out_f
18369        {
18370            return Err("matvec_bf16_dual_silu geometry".into());
18371        }
18372        let f = self.func("matvec_bf16_dual_silu");
18373        let cfg = LaunchConfig {
18374            grid_dim: (out_f as u32, 1, 1),
18375            block_dim: (mmv_block(), 1, 1),
18376            shared_mem_bytes: 0,
18377        };
18378        let (ini, outi) = (in_f as i32, out_f as i32);
18379        let lim = limit.unwrap_or(0.0);
18380        let __s_b = self.gpu.stream();
18381        let mut b = __s_b.launch_builder(&f);
18382        b.arg(wg)
18383            .arg(wu)
18384            .arg(x)
18385            .arg(act)
18386            .arg(&ini)
18387            .arg(&outi)
18388            .arg(&lim);
18389        unsafe {
18390            b.launch(cfg)?;
18391        }
18392        Ok(())
18393    }
18394
18395    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
18396    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
18397    #[allow(clippy::too_many_arguments)]
18398    pub fn matvec_bf16_dual_view_into(
18399        &self,
18400        wg: &cudarc::driver::CudaView<'_, u8>,
18401        wu: &cudarc::driver::CudaView<'_, u8>,
18402        x: &CudaSlice<f32>,
18403        yg: &mut CudaSlice<f32>,
18404        yu: &mut CudaSlice<f32>,
18405        in_f: usize,
18406        out_f: usize,
18407    ) -> Result<(), Box<dyn std::error::Error>> {
18408        if wg.len() != in_f * out_f * 2
18409            || wu.len() != in_f * out_f * 2
18410            || x.len() < in_f
18411            || in_f % 8 != 0
18412            || yg.len() < out_f
18413            || yu.len() < out_f
18414        {
18415            return Err(format!(
18416                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18417                wg.len(),
18418                wu.len(),
18419                x.len()
18420            )
18421            .into());
18422        }
18423        let f = self.func("matvec_bf16_dual");
18424        let cfg = LaunchConfig {
18425            grid_dim: ((2 * out_f) as u32, 1, 1),
18426            block_dim: (mmv_block(), 1, 1),
18427            shared_mem_bytes: 0,
18428        };
18429        let (ini, outi) = (in_f as i32, out_f as i32);
18430        let __s_b = self.gpu.stream();
18431        let mut b = __s_b.launch_builder(&f);
18432        b.arg(wg)
18433            .arg(wu)
18434            .arg(x)
18435            .arg(yg)
18436            .arg(yu)
18437            .arg(&ini)
18438            .arg(&outi);
18439        unsafe {
18440            b.launch(cfg)?;
18441        }
18442        Ok(())
18443    }
18444
18445    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
18446    #[allow(clippy::too_many_arguments)]
18447    pub fn matvec_bf16_dual_into(
18448        &self,
18449        wg: &CudaSlice<u8>,
18450        wu: &CudaSlice<u8>,
18451        x: &CudaSlice<f32>,
18452        yg: &mut CudaSlice<f32>,
18453        yu: &mut CudaSlice<f32>,
18454        in_f: usize,
18455        out_f: usize,
18456    ) -> Result<(), Box<dyn std::error::Error>> {
18457        if wg.len() != in_f * out_f * 2
18458            || wu.len() != in_f * out_f * 2
18459            || x.len() < in_f
18460            || in_f % 8 != 0
18461            || yg.len() < out_f
18462            || yu.len() < out_f
18463        {
18464            return Err(format!(
18465                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18466                wg.len(),
18467                wu.len(),
18468                x.len()
18469            )
18470            .into());
18471        }
18472        let f = self.func("matvec_bf16_dual");
18473        let cfg = LaunchConfig {
18474            grid_dim: ((2 * out_f) as u32, 1, 1),
18475            block_dim: (mmv_block(), 1, 1),
18476            shared_mem_bytes: 0,
18477        };
18478        let (ini, outi) = (in_f as i32, out_f as i32);
18479        let __s_b = self.gpu.stream();
18480        let mut b = __s_b.launch_builder(&f);
18481        b.arg(wg)
18482            .arg(wu)
18483            .arg(x)
18484            .arg(yg)
18485            .arg(yu)
18486            .arg(&ini)
18487            .arg(&outi);
18488        unsafe {
18489            b.launch(cfg)?;
18490        }
18491        Ok(())
18492    }
18493
18494    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
18495    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
18496    pub(crate) fn matvec_bf16_dual(
18497        &self,
18498        wg: &CudaSlice<u8>,
18499        wu: &CudaSlice<u8>,
18500        x: &CudaSlice<f32>,
18501        in_f: usize,
18502        out_f: usize,
18503    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18504        if wg.len() != in_f * out_f * 2
18505            || wu.len() != in_f * out_f * 2
18506            || x.len() < in_f
18507            || in_f % 8 != 0
18508        {
18509            return Err(format!(
18510                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18511                wg.len(),
18512                wu.len(),
18513                x.len()
18514            )
18515            .into());
18516        }
18517        let mut yg = self.alloc_uninit::<f32>(out_f)?;
18518        let mut yu = self.alloc_uninit::<f32>(out_f)?;
18519        let f = self.func("matvec_bf16_dual");
18520        let cfg = LaunchConfig {
18521            grid_dim: ((2 * out_f) as u32, 1, 1),
18522            block_dim: (mmv_block(), 1, 1),
18523            shared_mem_bytes: 0,
18524        };
18525        let (ini, outi) = (in_f as i32, out_f as i32);
18526        let __s_b = self.gpu.stream();
18527        let mut b = __s_b.launch_builder(&f);
18528        b.arg(wg)
18529            .arg(wu)
18530            .arg(x)
18531            .arg(&mut yg)
18532            .arg(&mut yu)
18533            .arg(&ini)
18534            .arg(&outi);
18535        unsafe {
18536            b.launch(cfg)?;
18537        }
18538        Ok((yg, yu))
18539    }
18540
18541    #[allow(clippy::too_many_arguments)]
18542    fn linear_bf16_chunked_inner(
18543        &self,
18544        x: &CudaSlice<f32>,
18545        data: &CudaSlice<u8>,
18546        m: usize,
18547        in_f: usize,
18548        out_f: usize,
18549        exact: bool,
18550        canonical_chunk_rows: Option<usize>,
18551    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18552        const CHUNK_BYTES: usize = 256 << 20;
18553        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
18554        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
18555        if m == 1
18556            && !exact
18557            && canonical_chunk_rows.is_none()
18558            && in_f % 8 == 0
18559            && Self::bf16_mmv_on()
18560        {
18561            return self.matvec_bf16(data, x, in_f, out_f);
18562        }
18563        let row_bytes = in_f
18564            .checked_mul(std::mem::size_of::<f32>())
18565            .ok_or("BF16 chunk row byte count overflow")?;
18566        if row_bytes == 0 || out_f == 0 {
18567            return Err("BF16 chunk dimensions must be nonzero".into());
18568        }
18569        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18570        let chunk_rows = match canonical_chunk_rows {
18571            Some(rows) if rows == 0 => {
18572                return Err("canonical BF16 chunk rows must be nonzero".into());
18573            }
18574            Some(rows) if rows > max_chunk_rows => {
18575                return Err(format!(
18576                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18577                )
18578                .into());
18579            }
18580            Some(rows) if out_f % rows != 0 => {
18581                return Err(format!(
18582                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18583                )
18584                .into());
18585            }
18586            Some(rows) => rows,
18587            None => max_chunk_rows,
18588        };
18589        if chunk_rows >= out_f {
18590            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18591            return if exact {
18592                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
18593            } else {
18594                self.linear(x, &wf32, m, in_f, out_f)
18595            };
18596        }
18597        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18598        let mut r0 = 0usize;
18599        while r0 < out_f {
18600            let rows = chunk_rows.min(out_f - r0);
18601            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
18602            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
18603            let yc = if exact {
18604                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
18605            } else {
18606                self.linear(x, &wf32, m, in_f, rows)?
18607            };
18608            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
18609            for mi in 0..m {
18610                let src = yc.slice(mi * rows..(mi + 1) * rows);
18611                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
18612                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
18613            }
18614            r0 += rows;
18615        }
18616        Ok(y)
18617    }
18618
18619    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
18620    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
18621    /// chunked BF16 numerical program instead of re-encoding the weight.
18622    pub fn linear_bf16_resident(
18623        &self,
18624        x: &CudaSlice<f32>,
18625        data: &CudaSlice<u8>,
18626        m: usize,
18627        in_f: usize,
18628        out_f: usize,
18629    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18630        if data.len() != in_f * out_f * 2 {
18631            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18632        }
18633        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
18634    }
18635
18636    /// Execute a resident BF16 projection as fixed-width output-row chunks.
18637    ///
18638    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
18639    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
18640    /// model topology rather than the active rank count.
18641    pub fn linear_bf16_resident_canonical_rows(
18642        &self,
18643        x: &CudaSlice<f32>,
18644        data: &CudaSlice<u8>,
18645        m: usize,
18646        in_f: usize,
18647        out_f: usize,
18648        canonical_chunk_rows: usize,
18649    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18650        if data.len() != in_f * out_f * 2 {
18651            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18652        }
18653        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18654    }
18655
18656    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
18657    ///
18658    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
18659    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
18660    pub fn linear_f32_resident_canonical_rows(
18661        &self,
18662        x: &CudaSlice<f32>,
18663        data: &CudaSlice<f32>,
18664        m: usize,
18665        in_f: usize,
18666        out_f: usize,
18667        canonical_chunk_rows: usize,
18668    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18669        self.linear_f32_resident_canonical_rows_inner(
18670            x,
18671            data,
18672            m,
18673            in_f,
18674            out_f,
18675            canonical_chunk_rows,
18676            false,
18677        )
18678    }
18679
18680    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
18681    ///
18682    /// The projection shapes and values are identical to
18683    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
18684    /// changes, replacing one device copy per token with one placement kernel per output chunk.
18685    pub fn linear_f32_resident_canonical_rows_strided(
18686        &self,
18687        x: &CudaSlice<f32>,
18688        data: &CudaSlice<f32>,
18689        m: usize,
18690        in_f: usize,
18691        out_f: usize,
18692        canonical_chunk_rows: usize,
18693    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18694        self.linear_f32_resident_canonical_rows_inner(
18695            x,
18696            data,
18697            m,
18698            in_f,
18699            out_f,
18700            canonical_chunk_rows,
18701            true,
18702        )
18703    }
18704
18705    fn linear_f32_resident_canonical_rows_inner(
18706        &self,
18707        x: &CudaSlice<f32>,
18708        data: &CudaSlice<f32>,
18709        m: usize,
18710        in_f: usize,
18711        out_f: usize,
18712        canonical_chunk_rows: usize,
18713        strided_output: bool,
18714    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18715        if data.len() != in_f * out_f {
18716            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18717        }
18718        if canonical_chunk_rows == 0
18719            || canonical_chunk_rows > out_f
18720            || out_f % canonical_chunk_rows != 0
18721        {
18722            return Err(format!(
18723                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18724            )
18725            .into());
18726        }
18727        if canonical_chunk_rows == out_f {
18728            return self.linear(x, data, m, in_f, out_f);
18729        }
18730
18731        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18732        let input = x.slice(0..x.len());
18733        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18734            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18735            if m == 1 {
18736                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18737                self.linear_device_into(
18738                    &input,
18739                    &weights,
18740                    &mut destination,
18741                    1,
18742                    in_f,
18743                    canonical_chunk_rows,
18744                )?;
18745                continue;
18746            }
18747            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18748            if strided_output {
18749                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18750            } else {
18751                for token in 0..m {
18752                    let source = chunk
18753                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18754                    let mut destination =
18755                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18756                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18757                }
18758            }
18759        }
18760        Ok(y)
18761    }
18762
18763    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
18764    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
18765    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
18766    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
18767    pub fn linear_f32_resident_canonical_rows_t1_into(
18768        &self,
18769        x: &CudaSlice<f32>,
18770        data: &CudaSlice<f32>,
18771        y: &mut CudaSlice<f32>,
18772        in_f: usize,
18773        out_f: usize,
18774        canonical_chunk_rows: usize,
18775    ) -> Result<(), Box<dyn std::error::Error>> {
18776        if data.len() != in_f * out_f {
18777            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18778        }
18779        if y.len() != out_f || x.len() != in_f {
18780            return Err(format!(
18781                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18782                x.len(),
18783                y.len()
18784            )
18785            .into());
18786        }
18787        if canonical_chunk_rows == 0
18788            || canonical_chunk_rows > out_f
18789            || out_f % canonical_chunk_rows != 0
18790        {
18791            return Err(format!(
18792                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18793            )
18794            .into());
18795        }
18796        let input = x.slice(0..x.len());
18797        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18798            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18799            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18800            self.linear_device_into(
18801                &input,
18802                &weights,
18803                &mut destination,
18804                1,
18805                in_f,
18806                canonical_chunk_rows,
18807            )?;
18808        }
18809        Ok(())
18810    }
18811
18812    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
18813    /// without the allocation, for workspace-resident operands.
18814    pub fn linear_t1_into(
18815        &self,
18816        x: &cudarc::driver::CudaView<'_, f32>,
18817        w: &cudarc::driver::CudaView<'_, f32>,
18818        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18819        in_f: usize,
18820        out_f: usize,
18821    ) -> Result<(), Box<dyn std::error::Error>> {
18822        self.linear_device_into(x, w, y, 1, in_f, out_f)
18823    }
18824
18825    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
18826    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
18827    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
18828    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
18829    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
18830    /// router/shexp sites and matmul_decode_exact's Float arm.
18831    pub fn linear_decode_exact(
18832        &self,
18833        x: &CudaSlice<f32>,
18834        w: &CudaSlice<f32>,
18835        m_tokens: usize,
18836        in_f: usize,
18837        out_f: usize,
18838    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18839        if m_tokens == 1 {
18840            return self.linear(x, w, 1, in_f, out_f);
18841        }
18842        let xv = self.view(x, m_tokens * in_f);
18843        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18844        for t in 0..m_tokens {
18845            let row = xv.slice(t * in_f..(t + 1) * in_f);
18846            let mut xr = self.alloc_uninit::<f32>(in_f)?;
18847            self.copy_view_into(&mut xr, 0, &row, in_f)?;
18848            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18849            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18850        }
18851        Ok(y)
18852    }
18853
18854    pub fn linear(
18855        &self,
18856        x: &CudaSlice<f32>,
18857        w: &CudaSlice<f32>,
18858        m_tokens: usize,
18859        in_f: usize,
18860        out_f: usize,
18861    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18862        self.linear_device(x, w, m_tokens, in_f, out_f)
18863    }
18864
18865    fn linear_device<I>(
18866        &self,
18867        x: &I,
18868        w: &I,
18869        m_tokens: usize,
18870        in_f: usize,
18871        out_f: usize,
18872    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18873    where
18874        I: cudarc::driver::DevicePtr<f32>,
18875    {
18876        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
18877        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
18878        Ok(c)
18879    }
18880
18881    fn linear_device_into<I, O>(
18882        &self,
18883        x: &I,
18884        w: &I,
18885        c: &mut O,
18886        m_tokens: usize,
18887        in_f: usize,
18888        out_f: usize,
18889    ) -> Result<(), Box<dyn std::error::Error>>
18890    where
18891        I: cudarc::driver::DevicePtr<f32>,
18892        O: cudarc::driver::DevicePtrMut<f32>,
18893    {
18894        use cudarc::cublaslt::{Matmul, MatmulConfig};
18895        let cfg = MatmulConfig {
18896            transa: true,
18897            transb: false,
18898            transc: false,
18899            m: out_f as u64,
18900            n: m_tokens as u64,
18901            k: in_f as u64,
18902            alpha: 1.0,
18903            lda: in_f as i64,
18904            ldb: in_f as i64,
18905            beta: 0.0,
18906            ldc: out_f as i64,
18907            stride_a: None,
18908            stride_b: None,
18909            stride_c: None,
18910            stride_bias: None,
18911            batch_size: None,
18912        };
18913        let blas = self.gpu.blas();
18914        unsafe {
18915            blas.matmul(cfg, w, x, c, None, None)?;
18916        }
18917        Ok(())
18918    }
18919
18920    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
18921    ///
18922    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
18923    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
18924    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
18925    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
18926    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
18927    /// launch error mid-request.
18928    pub fn sdpa_naive(
18929        &self,
18930        q: &CudaSlice<f32>,
18931        k: &CudaSlice<f32>,
18932        v: &CudaSlice<f32>,
18933        o: &mut CudaSlice<f32>,
18934        head_dim: usize,
18935        n_head: usize,
18936        n_head_kv: usize,
18937        t: usize,
18938        t_kv: usize,
18939        scale: f32,
18940        causal: bool,
18941    ) -> Result<(), Box<dyn std::error::Error>> {
18942        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
18943            return self.sdpa_naive_gmem(
18944                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18945            );
18946        }
18947        let f = self.func("sdpa_naive_f32");
18948        let cfg = LaunchConfig {
18949            grid_dim: (n_head as u32, t as u32, 1),
18950            block_dim: (128, 1, 1),
18951            shared_mem_bytes: (t_kv * 4) as u32,
18952        };
18953        let (hd, nh, nhkv, ti, tkvi, cz) = (
18954            head_dim as i32,
18955            n_head as i32,
18956            n_head_kv as i32,
18957            t as i32,
18958            t_kv as i32,
18959            causal as i32,
18960        );
18961        let __s_b = self.gpu.stream();
18962        let mut b = __s_b.launch_builder(&f);
18963        b.arg(q)
18964            .arg(k)
18965            .arg(v)
18966            .arg(o)
18967            .arg(&hd)
18968            .arg(&nh)
18969            .arg(&nhkv)
18970            .arg(&ti)
18971            .arg(&tkvi)
18972            .arg(&scale)
18973            .arg(&cz);
18974        unsafe {
18975            b.launch(cfg)?;
18976        }
18977        Ok(())
18978    }
18979
18980    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
18981    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
18982    /// of dynamic shared memory: identical loop structure and reduction order, so the output
18983    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
18984    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
18985    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
18986    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
18987    /// T==T_kv caller cannot silently allocate tens of GB.
18988    #[allow(clippy::too_many_arguments)]
18989    pub fn sdpa_naive_gmem(
18990        &self,
18991        q: &CudaSlice<f32>,
18992        k: &CudaSlice<f32>,
18993        v: &CudaSlice<f32>,
18994        o: &mut CudaSlice<f32>,
18995        head_dim: usize,
18996        n_head: usize,
18997        n_head_kv: usize,
18998        t: usize,
18999        t_kv: usize,
19000        scale: f32,
19001        causal: bool,
19002    ) -> Result<(), Box<dyn std::error::Error>> {
19003        let ws_len = n_head
19004            .checked_mul(t)
19005            .and_then(|x| x.checked_mul(t_kv))
19006            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19007        let ws_bytes = ws_len
19008            .checked_mul(std::mem::size_of::<f32>())
19009            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19010        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19011            return Err(format!(
19012                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19013                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19014                 needs a tiled/flash kernel, not the naive oracle"
19015            )
19016            .into());
19017        }
19018        let mut scores = self.uninit(ws_len)?;
19019        let f = self.func("sdpa_naive_gmem_f32");
19020        let cfg = LaunchConfig {
19021            grid_dim: (n_head as u32, t as u32, 1),
19022            block_dim: (128, 1, 1),
19023            shared_mem_bytes: 0,
19024        };
19025        let (hd, nh, nhkv, ti, tkvi, cz) = (
19026            head_dim as i32,
19027            n_head as i32,
19028            n_head_kv as i32,
19029            t as i32,
19030            t_kv as i32,
19031            causal as i32,
19032        );
19033        let __s_b = self.gpu.stream();
19034        let mut b = __s_b.launch_builder(&f);
19035        b.arg(q)
19036            .arg(k)
19037            .arg(v)
19038            .arg(o)
19039            .arg(&mut scores)
19040            .arg(&hd)
19041            .arg(&nh)
19042            .arg(&nhkv)
19043            .arg(&ti)
19044            .arg(&tkvi)
19045            .arg(&scale)
19046            .arg(&cz);
19047        unsafe {
19048            b.launch(cfg)?;
19049        }
19050        Ok(())
19051    }
19052
19053    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
19054    /// bidirectional image islands. `span_id` labels each absolute kv position
19055    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
19056    /// reproducing the reference's non-causal image batch. window 0 = no window.
19057    #[allow(clippy::too_many_arguments)]
19058    pub fn sdpa_naive_island(
19059        &self,
19060        q: &CudaSlice<f32>,
19061        k: &CudaSlice<f32>,
19062        v: &CudaSlice<f32>,
19063        o: &mut CudaSlice<f32>,
19064        span_id: &CudaSlice<i32>,
19065        head_dim: usize,
19066        n_head: usize,
19067        n_head_kv: usize,
19068        t: usize,
19069        t_kv: usize,
19070        scale: f32,
19071        window: usize,
19072    ) -> Result<(), Box<dyn std::error::Error>> {
19073        let f = self.func("sdpa_naive_island_f32");
19074        let cfg = LaunchConfig {
19075            grid_dim: (n_head as u32, t as u32, 1),
19076            block_dim: (128, 1, 1),
19077            shared_mem_bytes: (t_kv * 4) as u32,
19078        };
19079        let (hd, nh, nhkv, ti, tkvi, wi) = (
19080            head_dim as i32,
19081            n_head as i32,
19082            n_head_kv as i32,
19083            t as i32,
19084            t_kv as i32,
19085            window as i32,
19086        );
19087        let __s_b = self.gpu.stream();
19088        let mut b = __s_b.launch_builder(&f);
19089        b.arg(q)
19090            .arg(k)
19091            .arg(v)
19092            .arg(o)
19093            .arg(span_id)
19094            .arg(&hd)
19095            .arg(&nh)
19096            .arg(&nhkv)
19097            .arg(&ti)
19098            .arg(&tkvi)
19099            .arg(&scale)
19100            .arg(&wi);
19101        unsafe {
19102            b.launch(cfg)?;
19103        }
19104        Ok(())
19105    }
19106
19107    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
19108    #[allow(clippy::too_many_arguments)]
19109    pub fn sdpa_naive_w(
19110        &self,
19111        q: &CudaSlice<f32>,
19112        k: &CudaSlice<f32>,
19113        v: &CudaSlice<f32>,
19114        o: &mut CudaSlice<f32>,
19115        head_dim: usize,
19116        n_head: usize,
19117        n_head_kv: usize,
19118        t: usize,
19119        t_kv: usize,
19120        scale: f32,
19121        causal: bool,
19122        window: usize,
19123    ) -> Result<(), Box<dyn std::error::Error>> {
19124        let f = self.func("sdpa_naive_w_f32");
19125        let cfg = LaunchConfig {
19126            grid_dim: (n_head as u32, t as u32, 1),
19127            block_dim: (128, 1, 1),
19128            shared_mem_bytes: (t_kv * 4) as u32,
19129        };
19130        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19131            head_dim as i32,
19132            n_head as i32,
19133            n_head_kv as i32,
19134            t as i32,
19135            t_kv as i32,
19136            causal as i32,
19137            window as i32,
19138        );
19139        let __s_b = self.gpu.stream();
19140        let mut b = __s_b.launch_builder(&f);
19141        b.arg(q)
19142            .arg(k)
19143            .arg(v)
19144            .arg(o)
19145            .arg(&hd)
19146            .arg(&nh)
19147            .arg(&nhkv)
19148            .arg(&ti)
19149            .arg(&tkvi)
19150            .arg(&scale)
19151            .arg(&cz)
19152            .arg(&wi);
19153        unsafe {
19154            b.launch(cfg)?;
19155        }
19156        Ok(())
19157    }
19158
19159    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
19160    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
19161    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
19162    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
19163    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
19164    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
19165    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
19166    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
19167    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
19168    #[allow(clippy::too_many_arguments)]
19169    pub fn sdpa_naive_w_lo(
19170        &self,
19171        q: &CudaSlice<f32>,
19172        k: &CudaSlice<f32>,
19173        v: &CudaSlice<f32>,
19174        o: &mut CudaSlice<f32>,
19175        head_dim: usize,
19176        n_head: usize,
19177        n_head_kv: usize,
19178        t: usize,
19179        t_kv: usize,
19180        scale: f32,
19181        causal: bool,
19182        window: usize,
19183    ) -> Result<(), Box<dyn std::error::Error>> {
19184        let kv_lo = if window > 0 {
19185            (t_kv - t + 1).saturating_sub(window)
19186        } else {
19187            0
19188        };
19189        let smem = (t_kv - kv_lo) * 4;
19190        if smem > 48 * 1024 {
19191            return Err(format!(
19192                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19193                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19194                 a window this wide needs the multi-pass long-ctx kernel"
19195            )
19196            .into());
19197        }
19198        let f = self.func("sdpa_naive_w_lo_f32");
19199        let cfg = LaunchConfig {
19200            grid_dim: (n_head as u32, t as u32, 1),
19201            block_dim: (128, 1, 1),
19202            shared_mem_bytes: smem as u32,
19203        };
19204        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
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            kv_lo as i32,
19213        );
19214        let __s_b = self.gpu.stream();
19215        let mut b = __s_b.launch_builder(&f);
19216        b.arg(q)
19217            .arg(k)
19218            .arg(v)
19219            .arg(o)
19220            .arg(&hd)
19221            .arg(&nh)
19222            .arg(&nhkv)
19223            .arg(&ti)
19224            .arg(&tkvi)
19225            .arg(&scale)
19226            .arg(&cz)
19227            .arg(&wi)
19228            .arg(&lo);
19229        unsafe {
19230            b.launch(cfg)?;
19231        }
19232        Ok(())
19233    }
19234
19235    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
19236    pub fn sdpa_naive_view(
19237        &self,
19238        q: &CudaSlice<f32>,
19239        k: &cudarc::driver::CudaView<f32>,
19240        v: &cudarc::driver::CudaView<f32>,
19241        o: &mut CudaSlice<f32>,
19242        head_dim: usize,
19243        n_head: usize,
19244        n_head_kv: usize,
19245        t: usize,
19246        t_kv: usize,
19247        scale: f32,
19248        causal: bool,
19249    ) -> Result<(), Box<dyn std::error::Error>> {
19250        let f = self.func("sdpa_naive_f32");
19251        let cfg = LaunchConfig {
19252            grid_dim: (n_head as u32, t as u32, 1),
19253            block_dim: (128, 1, 1),
19254            shared_mem_bytes: (t_kv * 4) as u32,
19255        };
19256        let (hd, nh, nhkv, ti, tkvi, cz) = (
19257            head_dim as i32,
19258            n_head as i32,
19259            n_head_kv as i32,
19260            t as i32,
19261            t_kv as i32,
19262            causal as i32,
19263        );
19264        let __s_b = self.gpu.stream();
19265        let mut b = __s_b.launch_builder(&f);
19266        b.arg(q)
19267            .arg(k)
19268            .arg(v)
19269            .arg(o)
19270            .arg(&hd)
19271            .arg(&nh)
19272            .arg(&nhkv)
19273            .arg(&ti)
19274            .arg(&tkvi)
19275            .arg(&scale)
19276            .arg(&cz);
19277        unsafe {
19278            b.launch(cfg)?;
19279        }
19280        Ok(())
19281    }
19282
19283    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
19284    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
19285    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
19286    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
19287    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
19288    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
19289    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
19290    #[allow(clippy::too_many_arguments)]
19291    pub fn fa_dequant_kv_view_f32(
19292        &self,
19293        k: &cudarc::driver::CudaView<u8>,
19294        v: &cudarc::driver::CudaView<u8>,
19295        kf: &mut CudaSlice<f32>,
19296        vf: &mut CudaSlice<f32>,
19297        kv_dim_k: usize,
19298        kv_dim_v: usize,
19299        t_kv: usize,
19300        k_tok_bytes: usize,
19301        v_tok_bytes: usize,
19302        g: bool,
19303    ) -> Result<(), Box<dyn std::error::Error>> {
19304        let f = if g {
19305            self.func_g("fa_dequant_kv_ws_f32")
19306        } else {
19307            self.func("fa_dequant_kv_ws_f32")
19308        };
19309        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19310        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19311        let cfg = LaunchConfig {
19312            grid_dim: (nblk.max(1), 1, 1),
19313            block_dim: (256, 1, 1),
19314            shared_mem_bytes: 0,
19315        };
19316        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19317        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19318        let __s_b = self.gpu.stream();
19319        let mut b = __s_b.launch_builder(&f);
19320        b.arg(k)
19321            .arg(v)
19322            .arg(&mut *kf)
19323            .arg(&mut *vf)
19324            .arg(&kdk)
19325            .arg(&kdv)
19326            .arg(&tkvi)
19327            .arg(&ktb)
19328            .arg(&vtb);
19329        unsafe {
19330            b.launch(cfg)?;
19331        }
19332        Ok(())
19333    }
19334
19335    #[allow(clippy::too_many_arguments)]
19336    pub fn sdpa_naive_quantized_view(
19337        &self,
19338        q: &CudaSlice<f32>,
19339        k: &cudarc::driver::CudaView<u8>,
19340        v: &cudarc::driver::CudaView<u8>,
19341        o: &mut CudaSlice<f32>,
19342        head_dim: usize,
19343        n_head: usize,
19344        n_head_kv: usize,
19345        t: usize,
19346        t_kv: usize,
19347        scale: f32,
19348        causal: bool,
19349        k_tok_bytes: usize,
19350        v_tok_bytes: usize,
19351    ) -> Result<(), Box<dyn std::error::Error>> {
19352        let kv_dim = n_head_kv * head_dim;
19353        let mut kf = self.uninit(t_kv * kv_dim)?;
19354        let mut vf = self.uninit(t_kv * kv_dim)?;
19355        let f = self.func("fa_dequant_kv_ws_f32");
19356        let total = (2 * t_kv * kv_dim) as u64;
19357        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19358        let cfg = LaunchConfig {
19359            grid_dim: (nblk.max(1), 1, 1),
19360            block_dim: (256, 1, 1),
19361            shared_mem_bytes: 0,
19362        };
19363        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19364        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19365        let __s_b = self.gpu.stream();
19366        let mut b = __s_b.launch_builder(&f);
19367        b.arg(k)
19368            .arg(v)
19369            .arg(&mut kf)
19370            .arg(&mut vf)
19371            .arg(&kv_dim_i)
19372            .arg(&kv_dim_i)
19373            .arg(&t_kv_i)
19374            .arg(&k_tok_bytes_i)
19375            .arg(&v_tok_bytes_i);
19376        unsafe { b.launch(cfg)? };
19377        self.sdpa_naive(
19378            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19379        )
19380    }
19381
19382    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
19383    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
19384    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
19385    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
19386    /// unwindowed function above and produces bit-identical output at window == 0.
19387    ///
19388    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
19389    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
19390    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
19391    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
19392    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
19393    #[allow(clippy::too_many_arguments)]
19394    pub fn sdpa_naive_w_quantized_view(
19395        &self,
19396        q: &CudaSlice<f32>,
19397        k: &cudarc::driver::CudaView<u8>,
19398        v: &cudarc::driver::CudaView<u8>,
19399        o: &mut CudaSlice<f32>,
19400        head_dim: usize,
19401        n_head: usize,
19402        n_head_kv: usize,
19403        t: usize,
19404        t_kv: usize,
19405        scale: f32,
19406        causal: bool,
19407        window: usize,
19408        k_tok_bytes: usize,
19409        v_tok_bytes: usize,
19410    ) -> Result<(), Box<dyn std::error::Error>> {
19411        let kv_dim = n_head_kv * head_dim;
19412        let mut kf = self.uninit(t_kv * kv_dim)?;
19413        let mut vf = self.uninit(t_kv * kv_dim)?;
19414        let f = self.func("fa_dequant_kv_ws_f32");
19415        let total = (2 * t_kv * kv_dim) as u64;
19416        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19417        let cfg = LaunchConfig {
19418            grid_dim: (nblk.max(1), 1, 1),
19419            block_dim: (256, 1, 1),
19420            shared_mem_bytes: 0,
19421        };
19422        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19423        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19424        let __s_b = self.gpu.stream();
19425        let mut b = __s_b.launch_builder(&f);
19426        b.arg(k)
19427            .arg(v)
19428            .arg(&mut kf)
19429            .arg(&mut vf)
19430            .arg(&kv_dim_i)
19431            .arg(&kv_dim_i)
19432            .arg(&t_kv_i)
19433            .arg(&k_tok_bytes_i)
19434            .arg(&v_tok_bytes_i);
19435        unsafe { b.launch(cfg)? };
19436        self.sdpa_naive_w(
19437            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19438        )
19439    }
19440
19441    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
19442    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
19443    /// Q/K/V/O [head_dim, n_head(_kv), T].
19444    pub fn fa_prefill(
19445        &self,
19446        q: &CudaSlice<f32>,
19447        k: &CudaSlice<f32>,
19448        v: &CudaSlice<f32>,
19449        o: &mut CudaSlice<f32>,
19450        head_dim: usize,
19451        n_head: usize,
19452        n_head_kv: usize,
19453        t: usize,
19454        t_kv: usize,
19455        scale: f32,
19456        causal: bool,
19457    ) -> Result<(), Box<dyn std::error::Error>> {
19458        if portable_mma_gated() {
19459            return self.sdpa_naive(
19460                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19461            );
19462        }
19463        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
19464        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
19465        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
19466        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
19467        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
19468        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
19469        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
19470        let fa3_on = head_dim == 256
19471            && causal
19472            && t == t_kv
19473            && match std::env::var("MEMRA_FA3").as_deref() {
19474                Ok("0") => false,
19475                // The force arm consults the arch now: the bf16 stage below calls
19476                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
19477                // a portable build. Refuse at the switch, not at the lookup.
19478                Ok("1") => {
19479                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19480                    true
19481                }
19482                _ => cfg!(memra_hopper_mma),
19483            };
19484        if fa3_on {
19485            let n = t * n_head * head_dim;
19486            let nkv = t * n_head_kv * head_dim;
19487            let mut q16 = self.alloc_u8_uninit(n * 2)?;
19488            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19489            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19490            self.f32_to_bf16_into(q, &mut q16, n)?;
19491            self.f32_to_bf16_into(k, &mut k16, nkv)?;
19492            self.f32_to_bf16_into(v, &mut v16, nkv)?;
19493            let rc = {
19494                use cudarc::driver::{DevicePtr, DevicePtrMut};
19495                let stream = self.gpu.stream();
19496                let (qp, _g1) = q16.device_ptr(&stream);
19497                let (kp, _g2) = k16.device_ptr(&stream);
19498                let (vp, _g3) = v16.device_ptr(&stream);
19499                let (op, _g4) = o.device_ptr_mut(&stream);
19500                unsafe {
19501                    memra_fa3_prefill(
19502                        qp as *const core::ffi::c_void,
19503                        kp as *const core::ffi::c_void,
19504                        vp as *const core::ffi::c_void,
19505                        op as *mut f32,
19506                        t as i32,
19507                        n_head as i32,
19508                        n_head_kv as i32,
19509                        head_dim as i32,
19510                        scale,
19511                        stream.cu_stream() as *mut core::ffi::c_void,
19512                    )
19513                }
19514            };
19515            if rc != 0 {
19516                return Err(format!("memra_fa3_prefill rc={rc}").into());
19517            }
19518            return Ok(());
19519        }
19520        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
19521        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
19522        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
19523        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
19524        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19525        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19526        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19527            const BLOCK_Q: usize = 64;
19528            const BKX: usize = 32;
19529            let f = self.func("fa_prefill_bf16_p1");
19530            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19531                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19532            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19533            f.set_attribute(
19534                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19535                shmem as i32,
19536            )?;
19537            let cfg = LaunchConfig {
19538                grid_dim: (
19539                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19540                    n_head as u32,
19541                    1,
19542                ),
19543                block_dim: (32, 4, 1),
19544                shared_mem_bytes: shmem,
19545            };
19546            let (hd, nh, nhkv, ti, tkvi, cz) = (
19547                head_dim as i32,
19548                n_head as i32,
19549                n_head_kv as i32,
19550                t as i32,
19551                t_kv as i32,
19552                causal as i32,
19553            );
19554            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19555            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19556            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19557            let __s_b = self.gpu.stream();
19558            let mut b = __s_b.launch_builder(&f);
19559            b.arg(&qb)
19560                .arg(&kb)
19561                .arg(&vb)
19562                .arg(o)
19563                .arg(&hd)
19564                .arg(&nh)
19565                .arg(&nhkv)
19566                .arg(&ti)
19567                .arg(&tkvi)
19568                .arg(&scale)
19569                .arg(&cz);
19570            unsafe {
19571                b.launch(cfg)?;
19572            }
19573            return Ok(());
19574        }
19575        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
19576        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
19577        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
19578        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
19579        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
19580        const BK: usize = 32;
19581        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
19582        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
19583        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19584        let (block_q, warps, w2_sfx): (usize, u32, &str) =
19585            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19586        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
19587        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
19588        // other head_dims to sdpa_naive before reaching here.
19589        let hd_sfx = fa_hd_suffix(head_dim)?;
19590        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19591        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
19592        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
19593        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
19594        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
19595        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
19596        let (kb16, vb16) = if bf16kv {
19597            let n = t_kv * n_head_kv * head_dim;
19598            let mut kb = self.alloc_u8_uninit(n * 2)?;
19599            let mut vb = self.alloc_u8_uninit(n * 2)?;
19600            let fcv = self.func("f32_to_bf16_bulk");
19601            let ni = n as i64;
19602            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
19603            let __s_b = self.gpu.stream();
19604            let mut b = __s_b.launch_builder(&fcv);
19605            b.arg(k).arg(&mut kb).arg(&ni);
19606            unsafe {
19607                b.launch(cfgc)?;
19608            }
19609            let __s_b = self.gpu.stream();
19610            let mut b = __s_b.launch_builder(&fcv);
19611            b.arg(v).arg(&mut vb).arg(&ni);
19612            unsafe {
19613                b.launch(cfgc)?;
19614            }
19615            (Some(kb), Some(vb))
19616        } else {
19617            (None, None)
19618        };
19619        let f = self.func(&if bf16kv {
19620            format!("fa_prefill_bf16kv_pp{hd_sfx}")
19621        } else {
19622            format!(
19623                "fa_prefill_f32{}{}{hd_sfx}",
19624                if floor { "" } else { "_pp" },
19625                if floor { "" } else { w2_sfx }
19626            )
19627        });
19628        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
19629        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
19630        let kv_stages = if bf16kv { 2 } else { 1 };
19631        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19632            + 4 * (block_q * BK + 2 * block_q)) as u32;
19633        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19634        f.set_attribute(
19635            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19636            shmem as i32,
19637        )?;
19638        let cfg = LaunchConfig {
19639            grid_dim: (
19640                (t as u32 + block_q as u32 - 1) / block_q as u32,
19641                n_head as u32,
19642                1,
19643            ),
19644            block_dim: (32, warps, 1),
19645            shared_mem_bytes: shmem,
19646        };
19647        let (hd, nh, nhkv, ti, tkvi, cz) = (
19648            head_dim as i32,
19649            n_head as i32,
19650            n_head_kv as i32,
19651            t as i32,
19652            t_kv as i32,
19653            causal as i32,
19654        );
19655        let __s_b = self.gpu.stream();
19656        let mut b = __s_b.launch_builder(&f);
19657        b.arg(q);
19658        match (&kb16, &vb16) {
19659            (Some(kb), Some(vb)) => {
19660                b.arg(kb).arg(vb);
19661            }
19662            _ => {
19663                b.arg(k).arg(v);
19664            }
19665        }
19666        b.arg(o)
19667            .arg(&hd)
19668            .arg(&nh)
19669            .arg(&nhkv)
19670            .arg(&ti)
19671            .arg(&tkvi)
19672            .arg(&scale)
19673            .arg(&cz);
19674        unsafe {
19675            b.launch(cfg)?;
19676        }
19677        Ok(())
19678    }
19679
19680    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
19681    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
19682    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
19683    #[allow(clippy::too_many_arguments)]
19684    pub fn fa_prefill_w(
19685        &self,
19686        q: &CudaSlice<f32>,
19687        k: &CudaSlice<f32>,
19688        v: &CudaSlice<f32>,
19689        o: &mut CudaSlice<f32>,
19690        head_dim: usize,
19691        n_head: usize,
19692        n_head_kv: usize,
19693        t: usize,
19694        t_kv: usize,
19695        scale: f32,
19696        causal: bool,
19697        window: usize,
19698    ) -> Result<(), Box<dyn std::error::Error>> {
19699        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
19700        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
19701        if portable_mma_gated() {
19702            return self.sdpa_naive_w(
19703                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19704            );
19705        }
19706        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
19707        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
19708        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
19709        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19710        let faw_f32 =
19711            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
19712        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19713        self.fa_prefill_w_arm(
19714            q,
19715            k,
19716            v,
19717            o,
19718            head_dim,
19719            n_head,
19720            n_head_kv,
19721            t,
19722            t_kv,
19723            scale,
19724            causal,
19725            window,
19726            floor || faw_f32,
19727            floor,
19728        )
19729    }
19730
19731    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
19732    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
19733    #[allow(clippy::too_many_arguments)]
19734    pub fn fa_prefill_w_pre(
19735        &self,
19736        qb: &CudaSlice<u8>,
19737        kb: &CudaSlice<u8>,
19738        vb: &CudaSlice<u8>,
19739        o: &mut CudaSlice<f32>,
19740        head_dim: usize,
19741        n_head: usize,
19742        n_head_kv: usize,
19743        t: usize,
19744        t_kv: usize,
19745        scale: f32,
19746        causal: bool,
19747        window: usize,
19748        v_f16: bool,
19749    ) -> Result<(), Box<dyn std::error::Error>> {
19750        const BLOCK_Q: usize = 64;
19751        const BK: usize = 32;
19752        debug_assert_eq!(head_dim, 256);
19753        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19754        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19755        if hp {
19756            const BLOCK_QH: usize = 32;
19757            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
19758            // else re-encode through the pooled scratch (stream-ordered reuse).
19759            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19760            let vh: &CudaSlice<u8> = if v_f16 {
19761                vb
19762            } else {
19763                let n = t_kv * n_head_kv * head_dim;
19764                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19765                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19766                }
19767                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19768                vguard.as_ref().unwrap()
19769            };
19770            let f = self.func("fa_prefill_w_bf16_p1h2");
19771            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19772            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19773            f.set_attribute(
19774                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19775                shmem as i32,
19776            )?;
19777            let cfg = LaunchConfig {
19778                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19779                block_dim: (32, 4, 1),
19780                shared_mem_bytes: shmem,
19781            };
19782            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19783                head_dim as i32,
19784                n_head as i32,
19785                n_head_kv as i32,
19786                t as i32,
19787                t_kv as i32,
19788                causal as i32,
19789                window as i32,
19790            );
19791            let __s_b = self.gpu.stream();
19792            let mut b = __s_b.launch_builder(&f);
19793            b.arg(qb)
19794                .arg(kb)
19795                .arg(vh)
19796                .arg(o)
19797                .arg(&hd)
19798                .arg(&nh)
19799                .arg(&nhkv)
19800                .arg(&ti)
19801                .arg(&tkvi)
19802                .arg(&scale)
19803                .arg(&cz)
19804                .arg(&wi);
19805            unsafe {
19806                b.launch(cfg)?;
19807            }
19808            return Ok(());
19809        }
19810        let f = self.func("fa_prefill_w_bf16_p1");
19811        let shmem =
19812            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19813        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19814        f.set_attribute(
19815            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19816            shmem as i32,
19817        )?;
19818        let cfg = LaunchConfig {
19819            grid_dim: (
19820                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19821                n_head as u32,
19822                1,
19823            ),
19824            block_dim: (32, 4, 1),
19825            shared_mem_bytes: shmem,
19826        };
19827        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19828            head_dim as i32,
19829            n_head as i32,
19830            n_head_kv as i32,
19831            t as i32,
19832            t_kv as i32,
19833            causal as i32,
19834            window as i32,
19835        );
19836        let __s_b = self.gpu.stream();
19837        let mut b = __s_b.launch_builder(&f);
19838        b.arg(qb)
19839            .arg(kb)
19840            .arg(vb)
19841            .arg(o)
19842            .arg(&hd)
19843            .arg(&nh)
19844            .arg(&nhkv)
19845            .arg(&ti)
19846            .arg(&tkvi)
19847            .arg(&scale)
19848            .arg(&cz)
19849            .arg(&wi);
19850        unsafe {
19851            b.launch(cfg)?;
19852        }
19853        Ok(())
19854    }
19855
19856    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
19857    #[allow(clippy::too_many_arguments)]
19858    pub fn fa_prefill_w_arm(
19859        &self,
19860        q: &CudaSlice<f32>,
19861        k: &CudaSlice<f32>,
19862        v: &CudaSlice<f32>,
19863        o: &mut CudaSlice<f32>,
19864        head_dim: usize,
19865        n_head: usize,
19866        n_head_kv: usize,
19867        t: usize,
19868        t_kv: usize,
19869        scale: f32,
19870        causal: bool,
19871        window: usize,
19872        f32_stage: bool,
19873        floor: bool,
19874    ) -> Result<(), Box<dyn std::error::Error>> {
19875        const BLOCK_Q: usize = 64;
19876        const BK: usize = 32;
19877        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19878        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
19879        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
19880        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
19881        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19882        let p1 = !floor
19883            && !f32_stage
19884            && *P1_ON.get_or_init(|| {
19885                std::env::var("MEMRA_FAW_P1")
19886                    .map(|v| v != "0")
19887                    .unwrap_or(true)
19888            });
19889        let hp =
19890            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19891        if hp {
19892            const BLOCK_QH: usize = 32;
19893            let f = self.func("fa_prefill_w_bf16_p1h2");
19894            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19895            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19896            f.set_attribute(
19897                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19898                shmem as i32,
19899            )?;
19900            let cfg = LaunchConfig {
19901                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19902                block_dim: (32, 4, 1),
19903                shared_mem_bytes: shmem,
19904            };
19905            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19906                head_dim as i32,
19907                n_head as i32,
19908                n_head_kv as i32,
19909                t as i32,
19910                t_kv as i32,
19911                causal as i32,
19912                window as i32,
19913            );
19914            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19915            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19916            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19917            let __s_b = self.gpu.stream();
19918            let mut b = __s_b.launch_builder(&f);
19919            b.arg(&qb)
19920                .arg(&kb)
19921                .arg(&vh)
19922                .arg(o)
19923                .arg(&hd)
19924                .arg(&nh)
19925                .arg(&nhkv)
19926                .arg(&ti)
19927                .arg(&tkvi)
19928                .arg(&scale)
19929                .arg(&cz)
19930                .arg(&wi);
19931            unsafe {
19932                b.launch(cfg)?;
19933            }
19934            return Ok(());
19935        }
19936        if p1 {
19937            let f = self.func("fa_prefill_w_bf16_p1");
19938            let shmem =
19939                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19940            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19941            f.set_attribute(
19942                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19943                shmem as i32,
19944            )?;
19945            let cfg = LaunchConfig {
19946                grid_dim: (
19947                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19948                    n_head as u32,
19949                    1,
19950                ),
19951                block_dim: (32, 4, 1),
19952                shared_mem_bytes: shmem,
19953            };
19954            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19955                head_dim as i32,
19956                n_head as i32,
19957                n_head_kv as i32,
19958                t as i32,
19959                t_kv as i32,
19960                causal as i32,
19961                window as i32,
19962            );
19963            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19964            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19965            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19966            let __s_b = self.gpu.stream();
19967            let mut b = __s_b.launch_builder(&f);
19968            b.arg(&qb)
19969                .arg(&kb)
19970                .arg(&vb)
19971                .arg(o)
19972                .arg(&hd)
19973                .arg(&nh)
19974                .arg(&nhkv)
19975                .arg(&ti)
19976                .arg(&tkvi)
19977                .arg(&scale)
19978                .arg(&cz)
19979                .arg(&wi);
19980            unsafe {
19981                b.launch(cfg)?;
19982            }
19983            return Ok(());
19984        }
19985        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
19986        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
19987        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19988        let g4 = !floor
19989            && !f32_stage
19990            && n_head_kv == 1
19991            && n_head % 4 == 0
19992            && *G4_ON.get_or_init(|| {
19993                std::env::var("MEMRA_FAW_G4")
19994                    .map(|v| v != "0")
19995                    .unwrap_or(true)
19996            });
19997        if g4 {
19998            const SP_M: usize = 16;
19999            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
20000            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
20001            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20002            let o2 = *O2_ON.get_or_init(|| {
20003                std::env::var("MEMRA_FAW_O2")
20004                    .map(|v| v != "0")
20005                    .unwrap_or(true)
20006            });
20007            let f = self.func(if o2 {
20008                "fa_prefill_w_bf16_g4o2"
20009            } else {
20010                "fa_prefill_w_bf16_g4"
20011            });
20012            let shmem = if o2 {
20013                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20014            } else {
20015                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20016                    as u32
20017            };
20018            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20019            f.set_attribute(
20020                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20021                shmem as i32,
20022            )?;
20023            let cfg = LaunchConfig {
20024                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20025                block_dim: (32, 4, 1),
20026                shared_mem_bytes: shmem,
20027            };
20028            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20029                head_dim as i32,
20030                n_head as i32,
20031                n_head_kv as i32,
20032                t as i32,
20033                t_kv as i32,
20034                causal as i32,
20035                window as i32,
20036            );
20037            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20038            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20039            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20040            let __s_b = self.gpu.stream();
20041            let mut b = __s_b.launch_builder(&f);
20042            b.arg(&qb)
20043                .arg(&kb)
20044                .arg(&vb)
20045                .arg(o)
20046                .arg(&hd)
20047                .arg(&nh)
20048                .arg(&nhkv)
20049                .arg(&ti)
20050                .arg(&tkvi)
20051                .arg(&scale)
20052                .arg(&cz)
20053                .arg(&wi);
20054            unsafe {
20055                b.launch(cfg)?;
20056            }
20057            return Ok(());
20058        }
20059        let f = self.func(if floor {
20060            "fa_prefill_w_f32"
20061        } else if f32_stage {
20062            "fa_prefill_w_f32_pp"
20063        } else {
20064            "fa_prefill_w_bf16_pp"
20065        });
20066        let shmem =
20067            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20068        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20069        f.set_attribute(
20070            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20071            shmem as i32,
20072        )?;
20073        let cfg = LaunchConfig {
20074            grid_dim: (
20075                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20076                n_head as u32,
20077                1,
20078            ),
20079            block_dim: (32, 4, 1),
20080            shared_mem_bytes: shmem,
20081        };
20082        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20083            head_dim as i32,
20084            n_head as i32,
20085            n_head_kv as i32,
20086            t as i32,
20087            t_kv as i32,
20088            causal as i32,
20089            window as i32,
20090        );
20091        if f32_stage {
20092            let __s_b = self.gpu.stream();
20093            let mut b = __s_b.launch_builder(&f);
20094            b.arg(q)
20095                .arg(k)
20096                .arg(v)
20097                .arg(o)
20098                .arg(&hd)
20099                .arg(&nh)
20100                .arg(&nhkv)
20101                .arg(&ti)
20102                .arg(&tkvi)
20103                .arg(&scale)
20104                .arg(&cz)
20105                .arg(&wi);
20106            unsafe {
20107                b.launch(cfg)?;
20108            }
20109        } else {
20110            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20111            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20112            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20113            let __s_b = self.gpu.stream();
20114            let mut b = __s_b.launch_builder(&f);
20115            b.arg(&qb)
20116                .arg(&kb)
20117                .arg(&vb)
20118                .arg(o)
20119                .arg(&hd)
20120                .arg(&nh)
20121                .arg(&nhkv)
20122                .arg(&ti)
20123                .arg(&tkvi)
20124                .arg(&scale)
20125                .arg(&cz)
20126                .arg(&wi);
20127            unsafe {
20128                b.launch(cfg)?;
20129            }
20130        }
20131        Ok(())
20132    }
20133
20134    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
20135    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
20136    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
20137    #[allow(clippy::too_many_arguments)]
20138    pub fn fa_prefill_hd512(
20139        &self,
20140        q: &CudaSlice<f32>,
20141        k: &CudaSlice<f32>,
20142        v: &CudaSlice<f32>,
20143        o: &mut CudaSlice<f32>,
20144        head_dim: usize,
20145        n_head: usize,
20146        n_head_kv: usize,
20147        t: usize,
20148        t_kv: usize,
20149        scale: f32,
20150        causal: bool,
20151    ) -> Result<(), Box<dyn std::error::Error>> {
20152        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
20153        if portable_mma_gated() {
20154            return self.sdpa_naive(
20155                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20156            );
20157        }
20158        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
20159        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
20160        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
20161        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
20162        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
20163        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20164        let f32_stage =
20165            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20166        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
20167        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
20168        // Own numeric config (partial-sum order) — battery-gated.
20169        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20170        let sp = !f32_stage
20171            && *SP_ON.get_or_init(|| {
20172                std::env::var("MEMRA_FA512_SP")
20173                    .map(|v| v != "0")
20174                    .unwrap_or(true)
20175            });
20176        self.fa_prefill_hd512_arm(
20177            q,
20178            k,
20179            v,
20180            o,
20181            head_dim,
20182            n_head,
20183            n_head_kv,
20184            t,
20185            t_kv,
20186            scale,
20187            causal,
20188            f32_stage,
20189            sp,
20190            sp && fa_f16pv_on(),
20191        )
20192    }
20193
20194    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
20195    #[allow(clippy::too_many_arguments)]
20196    pub fn fa_prefill_hd512_pre(
20197        &self,
20198        qb: &CudaSlice<u8>,
20199        kb: &CudaSlice<u8>,
20200        vb: &CudaSlice<u8>,
20201        o: &mut CudaSlice<f32>,
20202        head_dim: usize,
20203        n_head: usize,
20204        n_head_kv: usize,
20205        t: usize,
20206        t_kv: usize,
20207        scale: f32,
20208        causal: bool,
20209        v_f16: bool,
20210    ) -> Result<(), Box<dyn std::error::Error>> {
20211        debug_assert_eq!(head_dim, 512);
20212        const SP_M: usize = 16;
20213        const BKS: usize = 32;
20214        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
20215        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
20216        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
20217        let f16pv = fa_f16pv_on();
20218        let nw = if f16pv { fa512_wide_warps() } else { 2 };
20219        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20220        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20221        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20222        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20223            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
20224            let n = t_kv * n_head_kv * head_dim;
20225            let need = n * 2;
20226            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20227                *vguard = Some(self.alloc_uninit::<u8>(need)?);
20228            }
20229            let dst = vguard.as_mut().unwrap();
20230            self.bf16_to_f16_into(vb, n, dst)?;
20231            vguard.as_ref().unwrap()
20232        } else {
20233            vb
20234        };
20235        let f = self.func(if hp {
20236            "fa_prefill_bf16_hd512_sp16h2"
20237        } else {
20238            match (f16pv, nw) {
20239                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20240                (true, _) => "fa_prefill_bf16_hd512_sp16",
20241                _ => "fa_prefill_bf16_hd512_sp",
20242            }
20243        });
20244        let (nwarp, npart) = if hp {
20245            (4usize, 4usize)
20246        } else if nw > 2 {
20247            (nw, nw)
20248        } else {
20249            (2, 1)
20250        };
20251        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
20252        let shmem = if hp {
20253            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20254                as u32
20255        } else {
20256            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20257                + 4 * (npart * SP_M * BKS + SP_M)) as u32
20258        };
20259        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20260        f.set_attribute(
20261            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20262            shmem as i32,
20263        )?;
20264        let grid_y = if hp {
20265            (n_head / 2) as u32
20266        } else {
20267            n_head as u32
20268        };
20269        let cfg = LaunchConfig {
20270            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20271            block_dim: (32, nwarp as u32, 1),
20272            shared_mem_bytes: shmem,
20273        };
20274        let (hd, nh, nhkv, ti, tkvi, cz) = (
20275            head_dim as i32,
20276            n_head as i32,
20277            n_head_kv as i32,
20278            t as i32,
20279            t_kv as i32,
20280            causal as i32,
20281        );
20282        let __s_b = self.gpu.stream();
20283        let mut b = __s_b.launch_builder(&f);
20284        b.arg(qb)
20285            .arg(kb)
20286            .arg(vref)
20287            .arg(o)
20288            .arg(&hd)
20289            .arg(&nh)
20290            .arg(&nhkv)
20291            .arg(&ti)
20292            .arg(&tkvi)
20293            .arg(&scale)
20294            .arg(&cz);
20295        unsafe {
20296            b.launch(cfg)?;
20297        }
20298        Ok(())
20299    }
20300
20301    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
20302    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
20303    #[allow(clippy::too_many_arguments)]
20304    pub fn fa_prefill_hd512_arm(
20305        &self,
20306        q: &CudaSlice<f32>,
20307        k: &CudaSlice<f32>,
20308        v: &CudaSlice<f32>,
20309        o: &mut CudaSlice<f32>,
20310        head_dim: usize,
20311        n_head: usize,
20312        n_head_kv: usize,
20313        t: usize,
20314        t_kv: usize,
20315        scale: f32,
20316        causal: bool,
20317        f32_stage: bool,
20318        sp: bool,
20319        f16pv: bool,
20320    ) -> Result<(), Box<dyn std::error::Error>> {
20321        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20322        if sp && !f32_stage {
20323            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
20324            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
20325            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
20326            const SP_M: usize = 16;
20327            const BKS: usize = 32;
20328            let nw = if f16pv { fa512_wide_warps() } else { 2 };
20329            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20330            let f = self.func(if hp {
20331                "fa_prefill_bf16_hd512_sp16h2"
20332            } else {
20333                match (f16pv, nw) {
20334                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20335                    (true, _) => "fa_prefill_bf16_hd512_sp16",
20336                    _ => "fa_prefill_bf16_hd512_sp",
20337                }
20338            });
20339            let (nwarp, npart) = if hp {
20340                (4usize, 4usize)
20341            } else if nw > 2 {
20342                (nw, nw)
20343            } else {
20344                (2, 1)
20345            };
20346            let shmem = if hp {
20347                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20348                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20349            } else {
20350                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20351                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
20352            };
20353            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20354            f.set_attribute(
20355                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20356                shmem as i32,
20357            )?;
20358            let grid_y = if hp {
20359                (n_head / 2) as u32
20360            } else {
20361                n_head as u32
20362            };
20363            let cfg = LaunchConfig {
20364                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20365                block_dim: (32, nwarp as u32, 1),
20366                shared_mem_bytes: shmem,
20367            };
20368            let (hd, nh, nhkv, ti, tkvi, cz) = (
20369                head_dim as i32,
20370                n_head as i32,
20371                n_head_kv as i32,
20372                t as i32,
20373                t_kv as i32,
20374                causal as i32,
20375            );
20376            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20377            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20378            let vb = if f16pv {
20379                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20380            } else {
20381                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20382            };
20383            let __s_b = self.gpu.stream();
20384            let mut b = __s_b.launch_builder(&f);
20385            b.arg(&qb)
20386                .arg(&kb)
20387                .arg(&vb)
20388                .arg(o)
20389                .arg(&hd)
20390                .arg(&nh)
20391                .arg(&nhkv)
20392                .arg(&ti)
20393                .arg(&tkvi)
20394                .arg(&scale)
20395                .arg(&cz);
20396            unsafe {
20397                b.launch(cfg)?;
20398            }
20399            return Ok(());
20400        }
20401        const BLOCK_Q: usize = 32;
20402        const BK: usize = 32;
20403        const HALF: usize = 256;
20404        let f = self.func(if f32_stage {
20405            "fa_prefill_f32_hd512"
20406        } else {
20407            "fa_prefill_bf16_hd512"
20408        });
20409        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
20410        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20411            + 4 * BLOCK_Q) as u32;
20412        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20413        f.set_attribute(
20414            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20415            shmem as i32,
20416        )?;
20417        let cfg = LaunchConfig {
20418            grid_dim: (
20419                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20420                n_head as u32,
20421                2,
20422            ),
20423            block_dim: (32, 2, 1),
20424            shared_mem_bytes: shmem,
20425        };
20426        let (hd, nh, nhkv, ti, tkvi, cz) = (
20427            head_dim as i32,
20428            n_head as i32,
20429            n_head_kv as i32,
20430            t as i32,
20431            t_kv as i32,
20432            causal as i32,
20433        );
20434        if f32_stage {
20435            let __s_b = self.gpu.stream();
20436            let mut b = __s_b.launch_builder(&f);
20437            b.arg(q)
20438                .arg(k)
20439                .arg(v)
20440                .arg(o)
20441                .arg(&hd)
20442                .arg(&nh)
20443                .arg(&nhkv)
20444                .arg(&ti)
20445                .arg(&tkvi)
20446                .arg(&scale)
20447                .arg(&cz);
20448            unsafe {
20449                b.launch(cfg)?;
20450            }
20451        } else {
20452            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20453            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20454            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20455            let __s_b = self.gpu.stream();
20456            let mut b = __s_b.launch_builder(&f);
20457            b.arg(&qb)
20458                .arg(&kb)
20459                .arg(&vb)
20460                .arg(o)
20461                .arg(&hd)
20462                .arg(&nh)
20463                .arg(&nhkv)
20464                .arg(&ti)
20465                .arg(&tkvi)
20466                .arg(&scale)
20467                .arg(&cz);
20468            unsafe {
20469                b.launch(cfg)?;
20470            }
20471        }
20472        Ok(())
20473    }
20474
20475    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
20476    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
20477    /// separate f32_to_bf16 the FA entries would run).
20478    #[allow(clippy::too_many_arguments)]
20479    pub fn rope_neox2_bf16e(
20480        &self,
20481        q: &mut CudaSlice<f32>,
20482        k: &mut CudaSlice<f32>,
20483        qb: &mut CudaSlice<u8>,
20484        kb: &mut CudaSlice<u8>,
20485        pos: &CudaSlice<i32>,
20486        head_dim: usize,
20487        n_dims: usize,
20488        nh_q: usize,
20489        nh_k: usize,
20490        n_tokens: usize,
20491        base: f32,
20492        freq_scale: f32,
20493        ff: Option<&CudaSlice<f32>>,
20494    ) -> Result<(), Box<dyn std::error::Error>> {
20495        let f = self.func("rope_neox2_bf16e_f32");
20496        let rows = ((nh_q + nh_k) * n_tokens) as u32;
20497        let cfg = LaunchConfig {
20498            grid_dim: (rows, 1, 1),
20499            block_dim: ((head_dim / 2) as u32, 1, 1),
20500            shared_mem_bytes: 0,
20501        };
20502        let theta_scale = base.powf(-2.0 / n_dims as f32);
20503        let (hd, nd, nhq, nhk, nt) = (
20504            head_dim as i32,
20505            n_dims as i32,
20506            nh_q as i32,
20507            nh_k as i32,
20508            n_tokens as i32,
20509        );
20510        let __s_b = self.gpu.stream();
20511        let mut b = __s_b.launch_builder(&f);
20512        match ff {
20513            Some(t) => {
20514                b.arg(&mut *q)
20515                    .arg(&mut *k)
20516                    .arg(&mut *qb)
20517                    .arg(&mut *kb)
20518                    .arg(pos)
20519                    .arg(&hd)
20520                    .arg(&nd)
20521                    .arg(&nhq)
20522                    .arg(&nhk)
20523                    .arg(&nt)
20524                    .arg(&theta_scale)
20525                    .arg(&freq_scale)
20526                    .arg(t);
20527                unsafe {
20528                    b.launch(cfg)?;
20529                }
20530            }
20531            None => {
20532                let null: u64 = 0;
20533                b.arg(&mut *q)
20534                    .arg(&mut *k)
20535                    .arg(&mut *qb)
20536                    .arg(&mut *kb)
20537                    .arg(pos)
20538                    .arg(&hd)
20539                    .arg(&nd)
20540                    .arg(&nhq)
20541                    .arg(&nhk)
20542                    .arg(&nt)
20543                    .arg(&theta_scale)
20544                    .arg(&freq_scale)
20545                    .arg(&null);
20546                unsafe {
20547                    b.launch(cfg)?;
20548                }
20549            }
20550        }
20551        Ok(())
20552    }
20553
20554    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
20555    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
20556    pub fn f32_to_bf16(
20557        &self,
20558        x: &CudaSlice<f32>,
20559        n: usize,
20560    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20561        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
20562        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20563        let f = self.func("f32_to_bf16_flat");
20564        let n_i = n as i64;
20565        let cfg = LaunchConfig {
20566            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20567            block_dim: (256, 1, 1),
20568            shared_mem_bytes: 0,
20569        };
20570        let __s_b = self.gpu.stream();
20571        let mut b = __s_b.launch_builder(&f);
20572        b.arg(x).arg(&mut y).arg(&n_i);
20573        unsafe {
20574            b.launch(cfg)?;
20575        }
20576        Ok(y)
20577    }
20578
20579    pub fn f32_to_f16(
20580        &self,
20581        x: &CudaSlice<f32>,
20582        n: usize,
20583    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20584        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20585        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20586        let f = self.func("f32_to_f16_flat");
20587        let n_i = n as i64;
20588        let cfg = LaunchConfig {
20589            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20590            block_dim: (256, 1, 1),
20591            shared_mem_bytes: 0,
20592        };
20593        let __s_b = self.gpu.stream();
20594        let mut b = __s_b.launch_builder(&f);
20595        b.arg(x).arg(&mut y).arg(&n_i);
20596        unsafe {
20597            b.launch(cfg)?;
20598        }
20599        Ok(y)
20600    }
20601
20602    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
20603    pub fn bf16_to_f16(
20604        &self,
20605        xb: &CudaSlice<u8>,
20606        n: usize,
20607    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20608        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20609        self.bf16_to_f16_into(xb, n, &mut y)?;
20610        Ok(y)
20611    }
20612
20613    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
20614    pub fn bf16_to_f16_into(
20615        &self,
20616        xb: &CudaSlice<u8>,
20617        n: usize,
20618        y: &mut CudaSlice<u8>,
20619    ) -> Result<(), Box<dyn std::error::Error>> {
20620        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
20621        assert!(y.len() >= n * 2);
20622        let f = self.func("bf16_to_f16_flat");
20623        let n2 = (n / 2) as i64;
20624        let cfg = LaunchConfig {
20625            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
20626            block_dim: (256, 1, 1),
20627            shared_mem_bytes: 0,
20628        };
20629        let __s_b = self.gpu.stream();
20630        let mut b = __s_b.launch_builder(&f);
20631        b.arg(xb).arg(y).arg(&n2);
20632        unsafe {
20633            b.launch(cfg)?;
20634        }
20635        Ok(())
20636    }
20637
20638    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
20639    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
20640    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
20641    /// head_dim in {256, 128}, bf16kv lane on.
20642    #[allow(clippy::too_many_arguments)]
20643    pub fn fa_prefill_vl8(
20644        &self,
20645        seqs: &[FaSeqVl],
20646        head_dim: usize,
20647        n_head: usize,
20648        n_head_kv: usize,
20649        scale: f32,
20650    ) -> Result<(), Box<dyn std::error::Error>> {
20651        const BK: usize = 32;
20652        let b = seqs.len();
20653        assert!(b >= 1 && b <= 8);
20654        let mut packed = [FaSeqVl::default(); 8];
20655        packed[..b].copy_from_slice(seqs);
20656        let v = FaVl8(packed);
20657        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20658        let ept = (n_head_kv * head_dim) as i32;
20659        {
20660            let f = self.func("fa_mirror_vl");
20661            let max_n = (max_t as i64) * ept as i64;
20662            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
20663            for which in 0..2i32 {
20664                let cfg = LaunchConfig {
20665                    grid_dim: (blocks, 1, b as u32),
20666                    block_dim: (256, 1, 1),
20667                    shared_mem_bytes: 0,
20668                };
20669                let __s_lb = self.gpu.stream();
20670                let mut lb = __s_lb.launch_builder(&f);
20671                lb.arg(&v).arg(&ept).arg(&which);
20672                unsafe {
20673                    lb.launch(cfg)?;
20674                }
20675            }
20676        }
20677        let hd_sfx = fa_hd_suffix(head_dim)?;
20678        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
20679        let block_q = 64usize;
20680        let kv_stages = 2usize;
20681        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20682            + 4 * (block_q * BK + 2 * block_q)) as u32;
20683        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20684        f.set_attribute(
20685            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20686            shmem as i32,
20687        )?;
20688        let cfg = LaunchConfig {
20689            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
20690            block_dim: (32, 4, 1),
20691            shared_mem_bytes: shmem,
20692        };
20693        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20694        let __s_lb = self.gpu.stream();
20695        let mut lb = __s_lb.launch_builder(&f);
20696        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
20697        unsafe {
20698            lb.launch(cfg)?;
20699        }
20700        Ok(())
20701    }
20702
20703    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
20704    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
20705    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
20706    #[allow(clippy::too_many_arguments)]
20707    pub fn attn_pre_vl8(
20708        &self,
20709        seqs: &[AttnPreVl],
20710        wq: &CudaSlice<f32>,
20711        wk: &CudaSlice<f32>,
20712        head_dim: usize,
20713        rope_dims: usize,
20714        n_head: usize,
20715        n_head_kv: usize,
20716        eps: f32,
20717        freq_base: f32,
20718        freq_scale: f32,
20719        kv_dim_k: usize,
20720        kv_dim_v: usize,
20721        k_tok_bytes: usize,
20722        v_tok_bytes: usize,
20723    ) -> Result<(), Box<dyn std::error::Error>> {
20724        let b = seqs.len();
20725        assert!(b >= 1 && b <= 8);
20726        let mut packed = [AttnPreVl::default(); 8];
20727        packed[..b].copy_from_slice(seqs);
20728        let v = AttnPreVl8(packed);
20729        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20730        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20731        {
20732            let f = self.func("q_gate_split_vl");
20733            let n = max_t * (n_head * head_dim) as u32;
20734            let cfg = LaunchConfig {
20735                grid_dim: (n.div_ceil(256), 1, b as u32),
20736                block_dim: (256, 1, 1),
20737                shared_mem_bytes: 0,
20738            };
20739            let __s_lb = self.gpu.stream();
20740            let mut lb = __s_lb.launch_builder(&f);
20741            lb.arg(&v).arg(&hd).arg(&nh);
20742            unsafe {
20743                lb.launch(cfg)?;
20744            }
20745        }
20746        {
20747            let f = self.func("attn_rms_vl");
20748            let cfg = LaunchConfig {
20749                grid_dim: (max_t * n_head as u32, 2, b as u32),
20750                block_dim: (rms_block(), 1, 1),
20751                shared_mem_bytes: 0,
20752            };
20753            let __s_lb = self.gpu.stream();
20754            let mut lb = __s_lb.launch_builder(&f);
20755            lb.arg(&v)
20756                .arg(wq)
20757                .arg(wk)
20758                .arg(&hd)
20759                .arg(&nh)
20760                .arg(&nhkv)
20761                .arg(&eps);
20762            unsafe {
20763                lb.launch(cfg)?;
20764            }
20765        }
20766        {
20767            let f = self.func("attn_rope_vl");
20768            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20769            let nd = rope_dims as i32;
20770            let cfg = LaunchConfig {
20771                grid_dim: (max_t * n_head as u32, 2, b as u32),
20772                block_dim: ((head_dim / 2) as u32, 1, 1),
20773                shared_mem_bytes: 0,
20774            };
20775            let __s_lb = self.gpu.stream();
20776            let mut lb = __s_lb.launch_builder(&f);
20777            lb.arg(&v)
20778                .arg(&hd)
20779                .arg(&nd)
20780                .arg(&nh)
20781                .arg(&nhkv)
20782                .arg(&theta_scale)
20783                .arg(&freq_scale);
20784            unsafe {
20785                lb.launch(cfg)?;
20786            }
20787        }
20788        {
20789            let f = self.func("append_kv_vl");
20790            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20791            let cfg = LaunchConfig {
20792                grid_dim: (nblk, max_t, b as u32),
20793                block_dim: (32, 1, 1),
20794                shared_mem_bytes: 0,
20795            };
20796            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20797            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20798            let __s_lb = self.gpu.stream();
20799            let mut lb = __s_lb.launch_builder(&f);
20800            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20801            unsafe {
20802                lb.launch(cfg)?;
20803            }
20804        }
20805        Ok(())
20806    }
20807
20808    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
20809    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
20810    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
20811    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
20812    pub fn fa_prefill_view(
20813        &self,
20814        q: &CudaSlice<f32>,
20815        k: &cudarc::driver::CudaView<u8>,
20816        v: &cudarc::driver::CudaView<u8>,
20817        o: &mut CudaSlice<f32>,
20818        head_dim: usize,
20819        n_head: usize,
20820        n_head_kv: usize,
20821        t: usize,
20822        t_kv: usize,
20823        scale: f32,
20824        causal: bool,
20825        k_tok_bytes: usize,
20826        v_tok_bytes: usize,
20827        g: bool,
20828    ) -> Result<(), Box<dyn std::error::Error>> {
20829        if portable_mma_gated() {
20830            return self.sdpa_naive_quantized_view(
20831                q,
20832                k,
20833                v,
20834                o,
20835                head_dim,
20836                n_head,
20837                n_head_kv,
20838                t,
20839                t_kv,
20840                scale,
20841                causal,
20842                k_tok_bytes,
20843                v_tok_bytes,
20844            );
20845        }
20846        const BLOCK_Q: usize = 64;
20847        const BK: usize = 32;
20848        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
20849        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
20850        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20851        let f = if g {
20852            self.func_g(&name)
20853        } else {
20854            self.func(&name)
20855        };
20856        let shmem =
20857            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20858        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20859        f.set_attribute(
20860            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20861            shmem as i32,
20862        )?;
20863        let cfg = LaunchConfig {
20864            grid_dim: (
20865                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20866                n_head as u32,
20867                1,
20868            ),
20869            block_dim: (32, 4, 1),
20870            shared_mem_bytes: shmem,
20871        };
20872        let (hd, nh, nhkv, ti, tkvi, cz) = (
20873            head_dim as i32,
20874            n_head as i32,
20875            n_head_kv as i32,
20876            t as i32,
20877            t_kv as i32,
20878            causal as i32,
20879        );
20880        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20881        let __s_b = self.gpu.stream();
20882        let mut b = __s_b.launch_builder(&f);
20883        b.arg(q)
20884            .arg(k)
20885            .arg(v)
20886            .arg(o)
20887            .arg(&hd)
20888            .arg(&nh)
20889            .arg(&nhkv)
20890            .arg(&ti)
20891            .arg(&tkvi)
20892            .arg(&scale)
20893            .arg(&cz)
20894            .arg(&ktb)
20895            .arg(&vtb);
20896        unsafe {
20897            b.launch(cfg)?;
20898        }
20899        Ok(())
20900    }
20901
20902    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
20903    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
20904    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
20905    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
20906    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
20907    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
20908    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
20909    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
20910    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
20911    #[allow(clippy::too_many_arguments)]
20912    pub fn fa_prefill_view_ws(
20913        &self,
20914        q: &CudaSlice<f32>,
20915        k: &cudarc::driver::CudaView<u8>,
20916        v: &cudarc::driver::CudaView<u8>,
20917        o: &mut CudaSlice<f32>,
20918        head_dim: usize,
20919        n_head: usize,
20920        n_head_kv: usize,
20921        t: usize,
20922        t_kv: usize,
20923        scale: f32,
20924        causal: bool,
20925        k_tok_bytes: usize,
20926        v_tok_bytes: usize,
20927        g: bool,
20928    ) -> Result<(), Box<dyn std::error::Error>> {
20929        if portable_mma_gated() {
20930            return self.sdpa_naive_quantized_view(
20931                q,
20932                k,
20933                v,
20934                o,
20935                head_dim,
20936                n_head,
20937                n_head_kv,
20938                t,
20939                t_kv,
20940                scale,
20941                causal,
20942                k_tok_bytes,
20943                v_tok_bytes,
20944            );
20945        }
20946        const BLOCK_Q: usize = 64;
20947        const BK: usize = 32;
20948        let kv_dim_k = n_head_kv * head_dim;
20949        let kv_dim_v = n_head_kv * head_dim;
20950        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20951        let v_ws_bytes = t_kv * kv_dim_v * 2;
20952        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
20953        let mut guard = self.prime_deqw_ws.lock().unwrap();
20954        let need_grow = match guard.as_ref() {
20955            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20956            None => true,
20957        };
20958        if need_grow {
20959            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20960            let (ck, cv) = guard
20961                .as_ref()
20962                .map(|(a, b)| (a.len(), b.len()))
20963                .unwrap_or((0, 0));
20964            *guard = Some((
20965                self.alloc_u8(grow(ck, k_ws_bytes))?,
20966                self.alloc_u8(grow(cv, v_ws_bytes))?,
20967            ));
20968        }
20969        let (kw, vw) = guard.as_mut().unwrap();
20970        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
20971        {
20972            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
20973            let f = if g {
20974                self.func_g("fa_dequant_kv_ws_bf16")
20975            } else {
20976                self.func("fa_dequant_kv_ws_bf16")
20977            };
20978            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20979            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20980            let cfg = LaunchConfig {
20981                grid_dim: (nblk.max(1), 1, 1),
20982                block_dim: (256, 1, 1),
20983                shared_mem_bytes: 0,
20984            };
20985            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20986            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20987            let __s_b = self.gpu.stream();
20988            let mut b = __s_b.launch_builder(&f);
20989            b.arg(k)
20990                .arg(v)
20991                .arg(&mut *kw)
20992                .arg(&mut *vw)
20993                .arg(&kdk)
20994                .arg(&kdv)
20995                .arg(&tkvi)
20996                .arg(&ktb)
20997                .arg(&vtb);
20998            unsafe {
20999                b.launch(cfg)?;
21000            }
21001        }
21002        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
21003        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
21004        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
21005        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
21006        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
21007        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
21008        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
21009        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21010            .map(|v| v != "0")
21011            .unwrap_or(true);
21012        {
21013            let hd_sfx = fa_hd_suffix(head_dim)?;
21014            let f = self.func(&format!(
21015                "fa_prefill_qw{}{hd_sfx}",
21016                if db { "_db" } else { "" }
21017            ));
21018            let shmem = if db {
21019                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
21020                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21021            } else {
21022                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21023            };
21024            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21025            f.set_attribute(
21026                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21027                shmem as i32,
21028            )?;
21029            let cfg = LaunchConfig {
21030                grid_dim: (
21031                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21032                    n_head as u32,
21033                    1,
21034                ),
21035                block_dim: (32, 4, 1),
21036                shared_mem_bytes: shmem,
21037            };
21038            let (hd, nh, nhkv, ti, tkvi, cz) = (
21039                head_dim as i32,
21040                n_head as i32,
21041                n_head_kv as i32,
21042                t as i32,
21043                t_kv as i32,
21044                causal as i32,
21045            );
21046            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21047            let __s_b = self.gpu.stream();
21048            let mut b = __s_b.launch_builder(&f);
21049            b.arg(q)
21050                .arg(&*kw)
21051                .arg(&*vw)
21052                .arg(o)
21053                .arg(&hd)
21054                .arg(&nh)
21055                .arg(&nhkv)
21056                .arg(&ti)
21057                .arg(&tkvi)
21058                .arg(&scale)
21059                .arg(&cz)
21060                .arg(&kdk)
21061                .arg(&kdv);
21062            unsafe {
21063                b.launch(cfg)?;
21064            }
21065        }
21066        Ok(())
21067    }
21068
21069    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
21070    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
21071    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
21072    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
21073    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
21074    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
21075    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
21076    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
21077    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
21078    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
21079    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
21080    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
21081    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
21082    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
21083    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
21084    #[allow(clippy::too_many_arguments)]
21085    pub fn fa_prefill_view_ws_w_hd128(
21086        &self,
21087        q: &CudaSlice<f32>,
21088        k: &cudarc::driver::CudaView<u8>,
21089        v: &cudarc::driver::CudaView<u8>,
21090        o: &mut CudaSlice<f32>,
21091        head_dim: usize,
21092        n_head: usize,
21093        n_head_kv: usize,
21094        t: usize,
21095        t_kv: usize,
21096        scale: f32,
21097        causal: bool,
21098        window: usize,
21099        k_tok_bytes: usize,
21100        v_tok_bytes: usize,
21101    ) -> Result<(), Box<dyn std::error::Error>> {
21102        assert_eq!(
21103            head_dim, 128,
21104            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21105        );
21106        if portable_mma_gated() {
21107            return self.sdpa_naive_w_quantized_view(
21108                q,
21109                k,
21110                v,
21111                o,
21112                head_dim,
21113                n_head,
21114                n_head_kv,
21115                t,
21116                t_kv,
21117                scale,
21118                causal,
21119                window,
21120                k_tok_bytes,
21121                v_tok_bytes,
21122            );
21123        }
21124        const BLOCK_Q: usize = 64;
21125        const BK: usize = 32;
21126        let kv_dim_k = n_head_kv * head_dim;
21127        let kv_dim_v = n_head_kv * head_dim;
21128        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
21129        let v_ws_bytes = t_kv * kv_dim_v * 2;
21130        let mut guard = self.prime_deqw_ws.lock().unwrap();
21131        let need_grow = match guard.as_ref() {
21132            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21133            None => true,
21134        };
21135        if need_grow {
21136            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21137            let (ck, cv) = guard
21138                .as_ref()
21139                .map(|(a, b)| (a.len(), b.len()))
21140                .unwrap_or((0, 0));
21141            *guard = Some((
21142                self.alloc_u8(grow(ck, k_ws_bytes))?,
21143                self.alloc_u8(grow(cv, v_ws_bytes))?,
21144            ));
21145        }
21146        let (kw, vw) = guard.as_mut().unwrap();
21147        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
21148        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
21149        {
21150            let f = self.func("fa_dequant_kv_ws_bf16");
21151            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21152            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21153            let cfg = LaunchConfig {
21154                grid_dim: (nblk.max(1), 1, 1),
21155                block_dim: (256, 1, 1),
21156                shared_mem_bytes: 0,
21157            };
21158            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21159            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21160            let __s_b = self.gpu.stream();
21161            let mut b = __s_b.launch_builder(&f);
21162            b.arg(k)
21163                .arg(v)
21164                .arg(&mut *kw)
21165                .arg(&mut *vw)
21166                .arg(&kdk)
21167                .arg(&kdv)
21168                .arg(&tkvi)
21169                .arg(&ktb)
21170                .arg(&vtb);
21171            unsafe {
21172                b.launch(cfg)?;
21173            }
21174        }
21175        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
21176        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21177            .map(|v| v != "0")
21178            .unwrap_or(true);
21179        {
21180            let f = self.func(if db {
21181                "fa_prefill_qw_db_w_hd128"
21182            } else {
21183                "fa_prefill_qw_w_hd128"
21184            });
21185            let shmem = if db {
21186                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21187            } else {
21188                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21189            };
21190            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21191            f.set_attribute(
21192                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21193                shmem as i32,
21194            )?;
21195            let cfg = LaunchConfig {
21196                grid_dim: (
21197                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21198                    n_head as u32,
21199                    1,
21200                ),
21201                block_dim: (32, 4, 1),
21202                shared_mem_bytes: shmem,
21203            };
21204            let (hd, nh, nhkv, ti, tkvi, cz) = (
21205                head_dim as i32,
21206                n_head as i32,
21207                n_head_kv as i32,
21208                t as i32,
21209                t_kv as i32,
21210                causal as i32,
21211            );
21212            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21213            let __s_b = self.gpu.stream();
21214            let mut b = __s_b.launch_builder(&f);
21215            b.arg(q)
21216                .arg(&*kw)
21217                .arg(&*vw)
21218                .arg(o)
21219                .arg(&hd)
21220                .arg(&nh)
21221                .arg(&nhkv)
21222                .arg(&ti)
21223                .arg(&tkvi)
21224                .arg(&scale)
21225                .arg(&cz)
21226                .arg(&kdk)
21227                .arg(&kdv)
21228                .arg(&wnd);
21229            unsafe {
21230                b.launch(cfg)?;
21231            }
21232        }
21233        Ok(())
21234    }
21235
21236    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
21237    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
21238    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
21239    pub fn fa_decode(
21240        &self,
21241        q: &CudaSlice<f32>,
21242        k: &cudarc::driver::CudaView<u8>,
21243        v: &cudarc::driver::CudaView<u8>,
21244        o: &mut CudaSlice<f32>,
21245        head_dim: usize,
21246        n_head: usize,
21247        n_head_kv: usize,
21248        t_kv: usize,
21249        scale: f32,
21250        k_tok_bytes: usize,
21251        v_tok_bytes: usize,
21252    ) -> Result<(), Box<dyn std::error::Error>> {
21253        self.fa_decode_kvmod(
21254            q,
21255            k,
21256            v,
21257            o,
21258            head_dim,
21259            n_head,
21260            n_head_kv,
21261            t_kv,
21262            scale,
21263            k_tok_bytes,
21264            v_tok_bytes,
21265            false,
21266        )
21267    }
21268
21269    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
21270    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
21271    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
21272    #[allow(clippy::too_many_arguments)]
21273    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
21274    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
21275    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
21276    #[allow(clippy::too_many_arguments)]
21277    #[allow(clippy::too_many_arguments)]
21278    fn fa_decode_scalar_unified(
21279        &self,
21280        q: &cudarc::driver::CudaView<f32>,
21281        k: &cudarc::driver::CudaView<u8>,
21282        v: &cudarc::driver::CudaView<u8>,
21283        o: &mut cudarc::driver::CudaViewMut<f32>,
21284        head_dim: usize,
21285        n_head: usize,
21286        n_head_kv: usize,
21287        t_kv_host: usize,
21288        t_kv_dev: Option<&CudaSlice<i32>>,
21289        scale: f32,
21290        n_splits: usize,
21291        split_keys: usize,
21292        k_tok_bytes: usize,
21293        v_tok_bytes: usize,
21294        g: bool,
21295        part_o: &mut CudaSlice<f32>,
21296        part_m: &mut CudaSlice<f32>,
21297        part_l: &mut CudaSlice<f32>,
21298        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21299    ) -> Result<(), Box<dyn std::error::Error>> {
21300        let f = if g {
21301            self.func_g("fa_decode_f32")
21302        } else {
21303            self.fa_func("fa_decode_f32", head_dim)
21304        };
21305        let cfg = LaunchConfig {
21306            grid_dim: (n_head as u32, n_splits as u32, 1),
21307            block_dim: (head_dim as u32, 1, 1),
21308            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21309        };
21310        let (hd, nh, nhkv, nsp) = (
21311            head_dim as i32,
21312            n_head as i32,
21313            n_head_kv as i32,
21314            n_splits as i32,
21315        );
21316        let (ktb, vtb, tkvi, ski) = (
21317            k_tok_bytes as i64,
21318            v_tok_bytes as i64,
21319            t_kv_host as i32,
21320            split_keys as i32,
21321        );
21322        let __s_b = self.gpu.stream();
21323        let mut b = __s_b.launch_builder(&f);
21324        match t_kv_dev {
21325            Some(d) => {
21326                b.arg(q)
21327                    .arg(k)
21328                    .arg(v)
21329                    .arg(&mut *part_o)
21330                    .arg(&mut *part_m)
21331                    .arg(&mut *part_l)
21332                    .arg(&hd)
21333                    .arg(&nh)
21334                    .arg(&nhkv)
21335                    .arg(&tkvi)
21336                    .arg(d)
21337                    .arg(&scale)
21338                    .arg(&nsp)
21339                    .arg(&ski)
21340                    .arg(&ktb)
21341                    .arg(&vtb);
21342                unsafe {
21343                    b.launch(cfg)?;
21344                }
21345            }
21346            None => {
21347                let null: u64 = 0;
21348                b.arg(q)
21349                    .arg(k)
21350                    .arg(v)
21351                    .arg(&mut *part_o)
21352                    .arg(&mut *part_m)
21353                    .arg(&mut *part_l)
21354                    .arg(&hd)
21355                    .arg(&nh)
21356                    .arg(&nhkv)
21357                    .arg(&tkvi)
21358                    .arg(&null)
21359                    .arg(&scale)
21360                    .arg(&nsp)
21361                    .arg(&ski)
21362                    .arg(&ktb)
21363                    .arg(&vtb);
21364                unsafe {
21365                    b.launch(cfg)?;
21366                }
21367            }
21368        }
21369        let cfg2 = LaunchConfig {
21370            grid_dim: (n_head as u32, 1, 1),
21371            block_dim: (head_dim as u32, 1, 1),
21372            shared_mem_bytes: 0,
21373        };
21374        if let Some((oq, od)) = q8_out {
21375            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
21376            let fc = if g {
21377                self.func_g("fa_decode_combine_q8_1")
21378            } else {
21379                self.fa_func("fa_decode_combine_q8_1", head_dim)
21380            };
21381            let __s_b2 = self.gpu.stream();
21382            let mut b2 = __s_b2.launch_builder(&fc);
21383            b2.arg(&*part_o)
21384                .arg(&*part_m)
21385                .arg(&*part_l)
21386                .arg(oq)
21387                .arg(od)
21388                .arg(&hd)
21389                .arg(&nh)
21390                .arg(&nsp);
21391            unsafe {
21392                b2.launch(cfg2)?;
21393            }
21394            return Ok(());
21395        }
21396        let fc = if g {
21397            self.func_g("fa_decode_combine_f32")
21398        } else {
21399            self.fa_func("fa_decode_combine_f32", head_dim)
21400        };
21401        let __s_b2 = self.gpu.stream();
21402        let mut b2 = __s_b2.launch_builder(&fc);
21403        b2.arg(&*part_o)
21404            .arg(&*part_m)
21405            .arg(&*part_l)
21406            .arg(o)
21407            .arg(&hd)
21408            .arg(&nh)
21409            .arg(&nsp);
21410        unsafe {
21411            b2.launch(cfg2)?;
21412        }
21413        Ok(())
21414    }
21415
21416    pub fn fa_decode_kvmod(
21417        &self,
21418        q: &CudaSlice<f32>,
21419        k: &cudarc::driver::CudaView<u8>,
21420        v: &cudarc::driver::CudaView<u8>,
21421        o: &mut CudaSlice<f32>,
21422        head_dim: usize,
21423        n_head: usize,
21424        n_head_kv: usize,
21425        t_kv: usize,
21426        scale: f32,
21427        k_tok_bytes: usize,
21428        v_tok_bytes: usize,
21429        g: bool,
21430    ) -> Result<(), Box<dyn std::error::Error>> {
21431        let q_view = q.as_view();
21432        let mut o_view = o.as_view_mut();
21433        self.fa_decode_kvmod_view(
21434            &q_view,
21435            k,
21436            v,
21437            &mut o_view,
21438            head_dim,
21439            n_head,
21440            n_head_kv,
21441            t_kv,
21442            scale,
21443            k_tok_bytes,
21444            v_tok_bytes,
21445            g,
21446        )
21447    }
21448
21449    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
21450    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
21451    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
21452    /// per-session KV view and FA launch.
21453    #[allow(clippy::too_many_arguments)]
21454    pub fn fa_decode_kvmod_view(
21455        &self,
21456        q: &cudarc::driver::CudaView<f32>,
21457        k: &cudarc::driver::CudaView<u8>,
21458        v: &cudarc::driver::CudaView<u8>,
21459        o: &mut cudarc::driver::CudaViewMut<f32>,
21460        head_dim: usize,
21461        n_head: usize,
21462        n_head_kv: usize,
21463        t_kv: usize,
21464        scale: f32,
21465        k_tok_bytes: usize,
21466        v_tok_bytes: usize,
21467        g: bool,
21468    ) -> Result<(), Box<dyn std::error::Error>> {
21469        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
21470        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
21471        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
21472        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
21473        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
21474        //
21475        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
21476        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
21477        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
21478        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
21479        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
21480        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
21481        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
21482        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
21483        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
21484        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
21485        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
21486        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
21487        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
21488        // fall to the exact scalar there instead of the broken register arm.
21489        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21490        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
21491        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
21492        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
21493        if g && head_dim == 256 && !fa_v4_at(t_kv) {
21494            fa_vec = false;
21495        }
21496        let sp = fa_split_keys(t_kv, n_head_kv);
21497        let n_splits = if fa_vec {
21498            ((t_kv + sp - 1) / sp).max(1)
21499        } else {
21500            ((t_kv + 255) / 256).max(1)
21501        };
21502        let o_len = n_head * n_splits * head_dim;
21503        let ml_len = n_head * n_splits;
21504        let mut part_guard = self.fa_part_pool.lock().unwrap();
21505        if part_guard
21506            .as_ref()
21507            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21508            .unwrap_or(true)
21509        {
21510            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21511            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21512            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21513            // later live allocations land at those addresses, and the next graph REPLAY writes
21514            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21515            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21516            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21517            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21518            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21519            // (total retired < final size).
21520            let old = part_guard.take();
21521            let (co, cm) = old
21522                .as_ref()
21523                .map(|pp| (pp.0.len(), pp.1.len()))
21524                .unwrap_or((0, 0));
21525            if let Some(old) = old {
21526                self.fa_part_retired.lock().unwrap().push(old);
21527            }
21528            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21529                eprintln!(
21530                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21531                    co, o_len, cm, ml_len
21532                );
21533            }
21534            *part_guard = Some((
21535                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21536                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21537                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21538            ));
21539        }
21540        let pg = part_guard.as_mut().unwrap();
21541        self.gpu
21542            .stream()
21543            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21544        self.gpu
21545            .stream()
21546            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21547        self.gpu
21548            .stream()
21549            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21550        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21551        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21552        let (hd, nh, nhkv, tkvi, nsp) = (
21553            head_dim as i32,
21554            n_head as i32,
21555            n_head_kv as i32,
21556            t_kv as i32,
21557            n_splits as i32,
21558        );
21559        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21560        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
21561        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
21562        // silently truncating the accumulator.
21563        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
21564        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
21565        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
21566        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
21567        // 178.4 -> 173.7 when 512 rode vec unconditionally).
21568        let fa512_min = fa512_min_tkv();
21569        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
21570        // g-module keeps the v4 pick (its class is not the depth-decay class).
21571        let deep = fa_vec
21572            && head_dim == 256
21573            && fa_v4_at(t_kv)
21574            && !g
21575            && fa_deep_at(t_kv)
21576            && !matches!(fa_v4_mode(), "noB3" | "stage");
21577        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21578            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
21579            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
21580            let gqa = (n_head / n_head_kv).max(1) as u32;
21581            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21582            (
21583                fv,
21584                LaunchConfig {
21585                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21586                    block_dim: (32, gqa, 1),
21587                    shared_mem_bytes: 0,
21588                },
21589            )
21590        } else if fa_vec && head_dim <= 256 {
21591            let gqa = (n_head / n_head_kv).max(1) as u32;
21592            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
21593            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
21594            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
21595            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
21596            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
21597            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
21598            // dequant each tile ONCE per block.
21599            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
21600            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
21601            // there by 12x — latency, not bandwidth, rules small KV).
21602            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21603            let smem_tkv = *SMEM_TKV.get_or_init(|| {
21604                std::env::var("MEMRA_FA_SMEM_TKV")
21605                    .ok()
21606                    .and_then(|v| v.parse().ok())
21607                    .unwrap_or_else(|| {
21608                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21609                    })
21610            });
21611            if fa_v4_at(t_kv) && head_dim == 256 {
21612                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
21613                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
21614                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
21615                let v4name = match fa_v4_mode() {
21616                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
21617                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
21618                    _ if deep => "fa_decode_vec_q_v4_deep",
21619                    _ => "fa_decode_vec_q_v4",
21620                };
21621                let fv = if g {
21622                    self.func_g(v4name)
21623                } else {
21624                    self.func(v4name)
21625                };
21626                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
21627                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
21628                let shmem = (if deep { 12160 } else { 11520 }
21629                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
21630                use cudarc::driver::sys::CUfunction_attribute_enum as A;
21631                fv.set_attribute(
21632                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21633                    shmem as i32,
21634                )?;
21635                (
21636                    fv,
21637                    LaunchConfig {
21638                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21639                        block_dim: (32, gqa, 1),
21640                        shared_mem_bytes: shmem,
21641                    },
21642                )
21643            } else if fa_v3_active(head_dim) {
21644                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
21645                // smem = sV only (half of v2's).
21646                let fv = if g {
21647                    self.func_g("fa_decode_vec_q_v3")
21648                } else {
21649                    self.func("fa_decode_vec_q_v3")
21650                };
21651                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
21652                (
21653                    fv,
21654                    LaunchConfig {
21655                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21656                        block_dim: (32, gqa, 1),
21657                        shared_mem_bytes: shmem,
21658                    },
21659                )
21660            } else if fa_v2_on() {
21661                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
21662                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
21663                // partials; same 32KB sK+sV tile as the smem twin.
21664                let fv = if g {
21665                    self.func_g("fa_decode_vec_q_v2")
21666                } else {
21667                    self.func("fa_decode_vec_q_v2")
21668                };
21669                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
21670                (
21671                    fv,
21672                    LaunchConfig {
21673                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21674                        block_dim: (32, gqa, 1),
21675                        shared_mem_bytes: shmem,
21676                    },
21677                )
21678            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
21679            {
21680                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
21681                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
21682                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
21683                let fv = if g {
21684                    self.func_g("fa_decode_vec_q_smem")
21685                } else {
21686                    self.func("fa_decode_vec_q_smem")
21687                };
21688                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
21689                use cudarc::driver::sys::CUfunction_attribute_enum as A;
21690                fv.set_attribute(
21691                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21692                    shmem as i32,
21693                )?;
21694                (
21695                    fv,
21696                    LaunchConfig {
21697                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21698                        block_dim: (32, gqa, 1),
21699                        shared_mem_bytes: shmem,
21700                    },
21701                )
21702            } else {
21703                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
21704                // dequant, zero dynamic shared memory.
21705                let fv = if g {
21706                    self.func_g("fa_decode_vec_q")
21707                } else {
21708                    self.func("fa_decode_vec_q")
21709                };
21710                (
21711                    fv,
21712                    LaunchConfig {
21713                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21714                        block_dim: (32, gqa, 1),
21715                        shared_mem_bytes: 0,
21716                    },
21717                )
21718            }
21719        } else {
21720            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
21721            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
21722            return self.fa_decode_scalar_unified(
21723                q,
21724                k,
21725                v,
21726                o,
21727                head_dim,
21728                n_head,
21729                n_head_kv,
21730                t_kv,
21731                None,
21732                scale,
21733                n_splits,
21734                if fa_vec { sp } else { 256 },
21735                k_tok_bytes,
21736                v_tok_bytes,
21737                g,
21738                part_o,
21739                part_m,
21740                part_l,
21741                None,
21742            );
21743        };
21744        let __s_b = self.gpu.stream();
21745        let mut b = __s_b.launch_builder(&f);
21746        b.arg(q)
21747            .arg(k)
21748            .arg(v)
21749            .arg(&mut *part_o)
21750            .arg(&mut *part_m)
21751            .arg(&mut *part_l)
21752            .arg(&hd)
21753            .arg(&nh)
21754            .arg(&nhkv)
21755            .arg(&tkvi)
21756            .arg(&scale)
21757            .arg(&nsp)
21758            .arg(&ktb)
21759            .arg(&vtb);
21760        unsafe {
21761            b.launch(cfg)?;
21762        }
21763        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
21764        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
21765        let (fc, cfg2) = (
21766            if g {
21767                self.func_g("fa_decode_combine_f32")
21768            } else {
21769                self.fa_func("fa_decode_combine_f32", head_dim)
21770            },
21771            LaunchConfig {
21772                grid_dim: (n_head as u32, 1, 1),
21773                block_dim: (head_dim as u32, 1, 1),
21774                shared_mem_bytes: 0,
21775            },
21776        );
21777        let __s_b2 = self.gpu.stream();
21778        let mut b2 = __s_b2.launch_builder(&fc);
21779        b2.arg(&*part_o)
21780            .arg(&*part_m)
21781            .arg(&*part_l)
21782            .arg(o)
21783            .arg(&hd)
21784            .arg(&nh)
21785            .arg(&nsp);
21786        unsafe {
21787            b2.launch(cfg2)?;
21788        }
21789        Ok(())
21790    }
21791
21792    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
21793    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
21794    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
21795    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
21796    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
21797    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
21798    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
21799    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
21800    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
21801    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
21802    #[allow(clippy::too_many_arguments)]
21803    pub fn fa_decode_batch_seqs_v4(
21804        &self,
21805        q: &CudaSlice<f32>,
21806        kv_ptrs: &cudarc::driver::CudaView<u64>,
21807        pos_seq: &CudaSlice<i32>,
21808        o: &mut CudaSlice<f32>,
21809        head_dim: usize,
21810        n_head: usize,
21811        n_head_kv: usize,
21812        b_n: usize,
21813        t_kv_max: usize,
21814        scale: f32,
21815        split_keys: usize,
21816        k_tok_bytes: usize,
21817        v_tok_bytes: usize,
21818    ) -> Result<(), Box<dyn std::error::Error>> {
21819        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21820        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21821        let o_len = b_n * n_head * n_splits_max * head_dim;
21822        let ml_len = b_n * n_head * n_splits_max;
21823        let mut part_guard = self.fa_part_pool.lock().unwrap();
21824        if part_guard
21825            .as_ref()
21826            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21827            .unwrap_or(true)
21828        {
21829            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21830            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21831            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21832            // later live allocations land at those addresses, and the next graph REPLAY writes
21833            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21834            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21835            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21836            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21837            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21838            // (total retired < final size).
21839            let old = part_guard.take();
21840            let (co, cm) = old
21841                .as_ref()
21842                .map(|pp| (pp.0.len(), pp.1.len()))
21843                .unwrap_or((0, 0));
21844            if let Some(old) = old {
21845                self.fa_part_retired.lock().unwrap().push(old);
21846            }
21847            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21848                eprintln!(
21849                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21850                    co, o_len, cm, ml_len
21851                );
21852            }
21853            *part_guard = Some((
21854                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21855                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21856                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21857            ));
21858        }
21859        let pg = part_guard.as_mut().unwrap();
21860        self.gpu
21861            .stream()
21862            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21863        self.gpu
21864            .stream()
21865            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21866        self.gpu
21867            .stream()
21868            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21869        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21870        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21871        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21872        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21873        let gqa = (n_head / n_head_kv).max(1) as u32;
21874        let f = self.func("fa_decode_vec_q_seqs_v4");
21875        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
21876        let shmem = (11520 + 32 * head_dim * 2) as u32;
21877        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21878        f.set_attribute(
21879            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21880            shmem as i32,
21881        )?;
21882        let cfg = LaunchConfig {
21883            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21884            block_dim: (32, gqa, 1),
21885            shared_mem_bytes: shmem,
21886        };
21887        {
21888            let __s_b = self.gpu.stream();
21889            let mut b = __s_b.launch_builder(&f);
21890            b.arg(q)
21891                .arg(kv_ptrs)
21892                .arg(pos_seq)
21893                .arg(&mut *part_o)
21894                .arg(&mut *part_m)
21895                .arg(&mut *part_l)
21896                .arg(&hd)
21897                .arg(&nh)
21898                .arg(&nhkv)
21899                .arg(&scale)
21900                .arg(&nspm)
21901                .arg(&spk)
21902                .arg(&ktb)
21903                .arg(&vtb);
21904            unsafe {
21905                b.launch(cfg)?;
21906            }
21907        }
21908        let fc = self.func("fa_decode_combine_seqs");
21909        let cfg2 = LaunchConfig {
21910            grid_dim: (n_head as u32, b_n as u32, 1),
21911            block_dim: (head_dim as u32, 1, 1),
21912            shared_mem_bytes: 0,
21913        };
21914        let __s_b2 = self.gpu.stream();
21915        let mut b2 = __s_b2.launch_builder(&fc);
21916        b2.arg(&*part_o)
21917            .arg(&*part_m)
21918            .arg(&*part_l)
21919            .arg(o)
21920            .arg(&hd)
21921            .arg(&nh)
21922            .arg(pos_seq)
21923            .arg(&nspm)
21924            .arg(&spk);
21925        unsafe {
21926            b2.launch(cfg2)?;
21927        }
21928        Ok(())
21929    }
21930
21931    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
21932    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
21933    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
21934    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
21935    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
21936    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
21937    #[allow(clippy::too_many_arguments)]
21938    pub fn append_kv_quantized_seqs(
21939        &self,
21940        k_rows: &CudaSlice<f32>,
21941        v_rows: &CudaSlice<f32>,
21942        kv_ptrs: &cudarc::driver::CudaView<u64>,
21943        pos_seq: &CudaSlice<i32>,
21944        b_n: usize,
21945        kv_dim_k: usize,
21946        kv_dim_v: usize,
21947        k_tok_bytes: usize,
21948        v_tok_bytes: usize,
21949    ) -> Result<(), Box<dyn std::error::Error>> {
21950        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21951        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21952        let cfg = LaunchConfig {
21953            grid_dim: (nblk, b_n as u32, 1),
21954            block_dim: (32, 1, 1),
21955            shared_mem_bytes: 0,
21956        };
21957        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21958        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21959        let __s_b = self.gpu.stream();
21960        let mut b = __s_b.launch_builder(&f);
21961        b.arg(k_rows)
21962            .arg(v_rows)
21963            .arg(kv_ptrs)
21964            .arg(pos_seq)
21965            .arg(&kdk)
21966            .arg(&kdv)
21967            .arg(&ktb)
21968            .arg(&vtb);
21969        unsafe {
21970            b.launch(cfg)?;
21971        }
21972        Ok(())
21973    }
21974
21975    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
21976    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
21977    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
21978    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
21979    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
21980    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21981        std::env::var("MEMRA_NO_FA_VEC").is_err()
21982            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21983            && base_len + 1 >= fa_vec_min_tkv()
21984            && head_dim <= 256
21985            && head_dim % 32 == 0
21986    }
21987
21988    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
21989    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
21990    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
21991    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
21992    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
21993    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
21994    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
21995    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
21996    #[allow(clippy::too_many_arguments)]
21997    pub fn fa_decode_rows(
21998        &self,
21999        q: &CudaSlice<f32>,
22000        k: &cudarc::driver::CudaView<u8>,
22001        v: &cudarc::driver::CudaView<u8>,
22002        o: &mut CudaSlice<f32>,
22003        head_dim: usize,
22004        n_head: usize,
22005        n_head_kv: usize,
22006        base_len: usize,
22007        t: usize,
22008        scale: f32,
22009        k_tok_bytes: usize,
22010        v_tok_bytes: usize,
22011        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
22012        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
22013        // keep the host arg. None is a bug for hd512 (asserted below).
22014        base_dev: Option<(&CudaSlice<i32>, i32)>,
22015        // K and V planes hold the same values (gemma globals, wv:=wk): pick
22016        // the _kv twin — V plane never read, value rides the q8_0 key dq.
22017        kv_shared: bool,
22018        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
22019        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
22020        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
22021        g: bool,
22022        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
22023        // (hd512 path) — the standalone quantize launch folds away.
22024        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22025    ) -> Result<(), Box<dyn std::error::Error>> {
22026        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22027        let t_kv_max = base_len + t; // LAST row's key bound
22028        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
22029        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
22030        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
22031        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
22032        // (parity law), so the partition is freely tunable — verify and decode move together.
22033        if head_dim == 512 {
22034            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22035            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
22036            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
22037            let v = *SP512.get_or_init(|| {
22038                std::env::var("MEMRA_FA_SP512")
22039                    .ok()
22040                    .and_then(|x| x.parse().ok())
22041                    .unwrap_or(0)
22042            });
22043            sp = if v >= 8 {
22044                v
22045            } else {
22046                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22047            };
22048        }
22049        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22050        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22051        let gqa = (n_head / n_head_kv).max(1) as u32;
22052        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
22053        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
22054        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
22055        // the different partition changes the combine's FP order (greedy tie flips at depth;
22056        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
22057        // consecutive rows by their OWN ladder value and launch once per group — each row then
22058        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
22059        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
22060        // sp override is t_kv-independent by construction).
22061        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
22062        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22063            groups.push((0, t, sp));
22064        } else {
22065            let mut r0 = 0usize;
22066            while r0 < t {
22067                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22068                let mut r1 = r0 + 1;
22069                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22070                    r1 += 1;
22071                }
22072                groups.push((r0, r1 - r0, sp_g));
22073                r0 = r1;
22074            }
22075        }
22076        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
22077        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
22078        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
22079        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22080        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22081            std::env::var("MEMRA_FA_SMEM_TKV")
22082                .ok()
22083                .and_then(|v| v.parse().ok())
22084                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22085        });
22086        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22087        let v3 = fa_v3_active(head_dim);
22088        let smem_rows =
22089            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22090        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
22091        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
22092        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
22093        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
22094        let _ = kv_shared;
22095        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
22096        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
22097        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22098        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
22099        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
22100        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
22101        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
22102        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
22103        // (kv_head, split) stages its tile once and loops the rows over it — kills the
22104        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
22105        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
22106        // shared by every hd512 caller through this wrapper (decode+verify flip together;
22107        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
22108        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
22109        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
22110        // not unpack-bound; jsonl 2026-07-14.
22111        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22112        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
22113        let tb512 = head_dim == 512
22114            && sp <= 32
22115            && n_head / n_head_kv.max(1) <= 16
22116            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22117        let fname = if tb512 {
22118            "fa_decode_vec_q_rows_v4_512_tb"
22119        } else if i2 {
22120            "fa_decode_vec_q_rows_dpl16_i2"
22121        } else if head_dim == 512 {
22122            "fa_decode_vec_q_rows_dpl16"
22123        }
22124        // gemma globals (parity law)
22125        else if v4 {
22126            "fa_decode_vec_q_rows_v4"
22127        } else if v3 {
22128            "fa_decode_vec_q_rows_v3"
22129        } else if fa_v2_on() {
22130            "fa_decode_vec_q_rows_v2"
22131        } else if smem_rows {
22132            "fa_decode_vec_q_rows_smem"
22133        } else {
22134            "fa_decode_vec_q_rows"
22135        };
22136        let f = if head_dim == 512 {
22137            self.fa_func(fname, head_dim)
22138        } else if g {
22139            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
22140            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
22141            // g-module rows against decode's g-module v4 — different programs, short-VG
22142            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
22143            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
22144            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
22145            // dq macros are format-aware.
22146            self.func_g(if smem_rows {
22147                "fa_decode_vec_q_rows"
22148            } else {
22149                fname
22150            })
22151        } else {
22152            self.func(fname)
22153        };
22154        let shmem = if tb512 {
22155            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
22156            let gk = Self::gkv_on();
22157            let sh =
22158                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22159            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22160            f.set_attribute(
22161                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22162                sh as i32,
22163            )?;
22164            sh
22165        } else if v4 || v3 || smem_rows || fa_v2_on() {
22166            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
22167            let sh = (if v4 {
22168                11520 + 32 * head_dim * if g { 1 } else { 2 }
22169            } else if v3 {
22170                32 * head_dim * 2
22171            } else {
22172                2 * 32 * head_dim * 2
22173            }) as u32;
22174            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22175            f.set_attribute(
22176                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22177                sh as i32,
22178            )?;
22179            sh
22180        } else {
22181            0
22182        };
22183        // Per-GROUP launches (single group in the common case — identical to the pre-fix
22184        // single launch there): each group gets its own partials (the rows kernel indexes
22185        // partials by its LOCAL grid.z row) and q/o row-offset views.
22186        for &(r0, t_g, sp_g) in &groups {
22187            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22188            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22189            let base_i = (base_len + r0) as i32;
22190            let o_len = t_g * n_head * n_splits_g * head_dim;
22191            let ml_len = t_g * n_head * n_splits_g;
22192            let mut part_guard = self.fa_part_pool.lock().unwrap();
22193            if part_guard
22194                .as_ref()
22195                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22196                .unwrap_or(true)
22197            {
22198                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22199                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22200                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22201                // later live allocations land at those addresses, and the next graph REPLAY writes
22202                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22203                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22204                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22205                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22206                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22207                // (total retired < final size).
22208                let old = part_guard.take();
22209                let (co, cm) = old
22210                    .as_ref()
22211                    .map(|pp| (pp.0.len(), pp.1.len()))
22212                    .unwrap_or((0, 0));
22213                if let Some(old) = old {
22214                    self.fa_part_retired.lock().unwrap().push(old);
22215                }
22216                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22217                    eprintln!(
22218                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22219                        co, o_len, cm, ml_len
22220                    );
22221                }
22222                *part_guard = Some((
22223                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22224                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22225                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22226                ));
22227            }
22228            let pg = part_guard.as_mut().unwrap();
22229            self.gpu
22230                .stream()
22231                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22232            self.gpu
22233                .stream()
22234                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22235            self.gpu
22236                .stream()
22237                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22238            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22239            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22240            let qv = self.view(q, t * n_head * head_dim);
22241            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22242            let cfg = LaunchConfig {
22243                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22244                block_dim: (32, gqa, 1),
22245                shared_mem_bytes: shmem,
22246            };
22247            {
22248                let __s_b = self.gpu.stream();
22249                let mut b = __s_b.launch_builder(&f);
22250                if tb512 {
22251                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
22252                    let (bd, plus) =
22253                        base_dev.expect("hd512 rows twin requires a device base counter");
22254                    let plus_g = plus + r0 as i32;
22255                    let nr = t_g as i32;
22256                    if Self::pdl_on() && Self::pdl_wb_on() {
22257                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
22258                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22259                        let s = &self.gpu.stream();
22260                        let (pq, _b0) = q_g.device_ptr(s);
22261                        let (pk, _b1) = k.device_ptr(s);
22262                        let (pv, _b2) = v.device_ptr(s);
22263                        let (po, _b3) = part_o.device_ptr_mut(s);
22264                        let (pm, _b4) = part_m.device_ptr_mut(s);
22265                        let (pl, _b5) = part_l.device_ptr_mut(s);
22266                        let (pb, _b6) = bd.device_ptr(s);
22267                        let mut ps = [
22268                            &pq as *const _ as *mut std::ffi::c_void,
22269                            &pk as *const _ as *mut _,
22270                            &pv as *const _ as *mut _,
22271                            &po as *const _ as *mut _,
22272                            &pm as *const _ as *mut _,
22273                            &pl as *const _ as *mut _,
22274                            &hd as *const _ as *mut _,
22275                            &nh as *const _ as *mut _,
22276                            &nhkv as *const _ as *mut _,
22277                            &pb as *const _ as *mut _,
22278                            &plus_g as *const _ as *mut _,
22279                            &scale as *const _ as *mut _,
22280                            &nspm as *const _ as *mut _,
22281                            &spk as *const _ as *mut _,
22282                            &ktb as *const _ as *mut _,
22283                            &vtb as *const _ as *mut _,
22284                            &nr as *const _ as *mut _,
22285                        ];
22286                        unsafe {
22287                            self.launch_pdl_flash(
22288                                Self::gkv_on(),
22289                                "fa_decode_vec_q_rows_v4_512_tb",
22290                                (n_head_kv as u32, n_splits_g as u32, 1),
22291                                (32, gqa, 1),
22292                                shmem,
22293                                &mut ps,
22294                            )?;
22295                        }
22296                    } else {
22297                        let cfg_tb = LaunchConfig {
22298                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22299                            block_dim: (32, gqa, 1),
22300                            shared_mem_bytes: shmem,
22301                        };
22302                        b.arg(&q_g)
22303                            .arg(k)
22304                            .arg(v)
22305                            .arg(&mut *part_o)
22306                            .arg(&mut *part_m)
22307                            .arg(&mut *part_l)
22308                            .arg(&hd)
22309                            .arg(&nh)
22310                            .arg(&nhkv)
22311                            .arg(bd)
22312                            .arg(&plus_g)
22313                            .arg(&scale)
22314                            .arg(&nspm)
22315                            .arg(&spk)
22316                            .arg(&ktb)
22317                            .arg(&vtb)
22318                            .arg(&nr);
22319                        unsafe {
22320                            b.launch(cfg_tb)?;
22321                        }
22322                    }
22323                } else if head_dim == 512 {
22324                    let (bd, plus) =
22325                        base_dev.expect("hd512 rows twin requires a device base counter");
22326                    let plus_g = plus + r0 as i32;
22327                    b.arg(&q_g)
22328                        .arg(k)
22329                        .arg(v)
22330                        .arg(&mut *part_o)
22331                        .arg(&mut *part_m)
22332                        .arg(&mut *part_l)
22333                        .arg(&hd)
22334                        .arg(&nh)
22335                        .arg(&nhkv)
22336                        .arg(bd)
22337                        .arg(&plus_g)
22338                        .arg(&scale)
22339                        .arg(&nspm)
22340                        .arg(&spk)
22341                        .arg(&ktb)
22342                        .arg(&vtb);
22343                    unsafe {
22344                        b.launch(cfg)?;
22345                    }
22346                } else {
22347                    b.arg(&q_g)
22348                        .arg(k)
22349                        .arg(v)
22350                        .arg(&mut *part_o)
22351                        .arg(&mut *part_m)
22352                        .arg(&mut *part_l)
22353                        .arg(&hd)
22354                        .arg(&nh)
22355                        .arg(&nhkv)
22356                        .arg(&base_i)
22357                        .arg(&scale)
22358                        .arg(&nspm)
22359                        .arg(&spk)
22360                        .arg(&ktb)
22361                        .arg(&vtb);
22362                    unsafe {
22363                        b.launch(cfg)?;
22364                    }
22365                }
22366            }
22367            let cfg2 = LaunchConfig {
22368                grid_dim: (n_head as u32, t_g as u32, 1),
22369                block_dim: (head_dim as u32, 1, 1),
22370                shared_mem_bytes: 0,
22371            };
22372            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22373            if head_dim == 512 {
22374                // device-len combine (shared by verify/eager/graph — parity by symbol): the
22375                // per-row n_splits derives from the SAME counter the rows kernel read.
22376                let (bd, plus) = base_dev.unwrap();
22377                let plus_g = plus + r0 as i32;
22378                if let Some((oq, od)) = q8_out.as_mut() {
22379                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
22380                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22381                    if Self::pdl_on() && Self::pdl_wb_on() {
22382                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
22383                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22384                        let s = &self.gpu.stream();
22385                        let (po, _g0) = part_o.device_ptr(s);
22386                        let (pm, _g1) = part_m.device_ptr(s);
22387                        let (pl, _g2) = part_l.device_ptr(s);
22388                        let (pq, _g3) = oq.device_ptr_mut(s);
22389                        let (pd, _g4) = od.device_ptr_mut(s);
22390                        let (pb, _g5) = bd.device_ptr(s);
22391                        let mut ps = [
22392                            &po as *const _ as *mut std::ffi::c_void,
22393                            &pm as *const _ as *mut _,
22394                            &pl as *const _ as *mut _,
22395                            &pq as *const _ as *mut _,
22396                            &pd as *const _ as *mut _,
22397                            &hd as *const _ as *mut _,
22398                            &nh as *const _ as *mut _,
22399                            &pb as *const _ as *mut _,
22400                            &plus_g as *const _ as *mut _,
22401                            &nspm as *const _ as *mut _,
22402                            &spk as *const _ as *mut _,
22403                        ];
22404                        unsafe {
22405                            self.launch_pdl_flash(
22406                                Self::gkv_on(),
22407                                "fa_decode_combine_rows_dc_q8_1",
22408                                cfg2.grid_dim,
22409                                cfg2.block_dim,
22410                                0,
22411                                &mut ps,
22412                            )?;
22413                        }
22414                        continue;
22415                    }
22416                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22417                    let __s_b2 = self.gpu.stream();
22418                    let mut b2 = __s_b2.launch_builder(&fc);
22419                    b2.arg(&*part_o)
22420                        .arg(&*part_m)
22421                        .arg(&*part_l)
22422                        .arg(&mut **oq)
22423                        .arg(&mut **od)
22424                        .arg(&hd)
22425                        .arg(&nh)
22426                        .arg(bd)
22427                        .arg(&plus_g)
22428                        .arg(&nspm)
22429                        .arg(&spk);
22430                    unsafe {
22431                        b2.launch(cfg2)?;
22432                    }
22433                    continue;
22434                }
22435                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22436                let __s_b2 = self.gpu.stream();
22437                let mut b2 = __s_b2.launch_builder(&fc);
22438                b2.arg(&*part_o)
22439                    .arg(&*part_m)
22440                    .arg(&*part_l)
22441                    .arg(&mut o_g)
22442                    .arg(&hd)
22443                    .arg(&nh)
22444                    .arg(bd)
22445                    .arg(&plus_g)
22446                    .arg(&nspm)
22447                    .arg(&spk);
22448                unsafe {
22449                    b2.launch(cfg2)?;
22450                }
22451            } else {
22452                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
22453                // leave the caller's pair unwritten (consumer would read garbage).
22454                assert!(
22455                    q8_out.is_none(),
22456                    "rows q8 emit requires the hd512 dc combine"
22457                );
22458                let fc = self.func("fa_decode_combine_rows");
22459                let __s_b2 = self.gpu.stream();
22460                let mut b2 = __s_b2.launch_builder(&fc);
22461                b2.arg(&*part_o)
22462                    .arg(&*part_m)
22463                    .arg(&*part_l)
22464                    .arg(&mut o_g)
22465                    .arg(&hd)
22466                    .arg(&nh)
22467                    .arg(&base_i)
22468                    .arg(&nspm)
22469                    .arg(&spk);
22470                unsafe {
22471                    b2.launch(cfg2)?;
22472                }
22473            }
22474        }
22475        Ok(())
22476    }
22477
22478    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
22479    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
22480    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
22481    #[allow(clippy::too_many_arguments)]
22482    pub fn fa_decode_rows_w(
22483        &self,
22484        q: &CudaSlice<f32>,
22485        k: &cudarc::driver::CudaView<u8>,
22486        v: &cudarc::driver::CudaView<u8>,
22487        o: &mut CudaSlice<f32>,
22488        head_dim: usize,
22489        n_head: usize,
22490        n_head_kv: usize,
22491        base_dev: &CudaSlice<i32>,
22492        base_plus: i32,
22493        t: usize,
22494        scale: f32,
22495        window: usize,
22496        k_tok_bytes: usize,
22497        v_tok_bytes: usize,
22498        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22499    ) -> Result<(), Box<dyn std::error::Error>> {
22500        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
22501        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
22502        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
22503        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
22504        debug_assert!(head_dim == 256);
22505        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
22506        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
22507        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
22508        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
22509        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
22510        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
22511        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
22512        let sp = {
22513            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22514            let v = *SPW.get_or_init(|| {
22515                std::env::var("MEMRA_FA_SPW")
22516                    .ok()
22517                    .and_then(|x| x.parse().ok())
22518                    .unwrap_or(0)
22519            });
22520            if v >= 8 {
22521                v
22522            } else {
22523                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22524            }
22525        };
22526        let n_splits_max = (window + sp - 1) / sp;
22527        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22528        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22529        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22530        let gqa = (n_head / n_head_kv).max(1) as u32;
22531        let o_len = t * n_head * n_splits_max * head_dim;
22532        let ml_len = t * n_head * n_splits_max;
22533        let mut part_guard = self.fa_part_pool.lock().unwrap();
22534        if part_guard
22535            .as_ref()
22536            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22537            .unwrap_or(true)
22538        {
22539            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22540            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22541            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22542            // later live allocations land at those addresses, and the next graph REPLAY writes
22543            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22544            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22545            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22546            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22547            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22548            // (total retired < final size).
22549            let old = part_guard.take();
22550            let (co, cm) = old
22551                .as_ref()
22552                .map(|pp| (pp.0.len(), pp.1.len()))
22553                .unwrap_or((0, 0));
22554            if let Some(old) = old {
22555                self.fa_part_retired.lock().unwrap().push(old);
22556            }
22557            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22558                eprintln!(
22559                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22560                    co, o_len, cm, ml_len
22561                );
22562            }
22563            *part_guard = Some((
22564                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22565                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22566                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22567            ));
22568        }
22569        let pg = part_guard.as_mut().unwrap();
22570        self.gpu
22571            .stream()
22572            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22573        self.gpu
22574            .stream()
22575            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22576        self.gpu
22577            .stream()
22578            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22579        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22580        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
22581        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
22582        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
22583        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
22584        // floor (deep-ctx broadcast win); register twin between.
22585        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22586        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22587            std::env::var("MEMRA_FA_SMEM_TKV")
22588                .ok()
22589                .and_then(|v| v.parse().ok())
22590                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22591        });
22592        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
22593        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
22594        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
22595        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
22596        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
22597        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22598        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
22599        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
22600        // per (lane, format-module) keeps parity structural; the old register-i2 detour
22601        // (-33%) is retired.
22602        let wg = Self::wkv_on();
22603        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
22604        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
22605        let sp2 =
22606            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
22607        if sp2 {
22608            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22609            if Self::pdl_on() && Self::pdl_wb_on() {
22610                // wave-B2b: flavor mirrors wg.
22611                use cudarc::driver::{DevicePtr, DevicePtrMut};
22612                let s = &self.gpu.stream();
22613                let (pq, _b0) = q.device_ptr(s);
22614                let (pk, _b1) = k.device_ptr(s);
22615                let (pv, _b2) = v.device_ptr(s);
22616                let (po, _b3) = part_o.device_ptr_mut(s);
22617                let (pm, _b4) = part_m.device_ptr_mut(s);
22618                let (pl, _b5) = part_l.device_ptr_mut(s);
22619                let (pb, _b6) = base_dev.device_ptr(s);
22620                let mut ps = [
22621                    &pq as *const _ as *mut std::ffi::c_void,
22622                    &pk as *const _ as *mut _,
22623                    &pv as *const _ as *mut _,
22624                    &po as *const _ as *mut _,
22625                    &pm as *const _ as *mut _,
22626                    &pl as *const _ as *mut _,
22627                    &hd as *const _ as *mut _,
22628                    &nh as *const _ as *mut _,
22629                    &nhkv as *const _ as *mut _,
22630                    &pb as *const _ as *mut _,
22631                    &base_plus as *const _ as *mut _,
22632                    &scale as *const _ as *mut _,
22633                    &nspm as *const _ as *mut _,
22634                    &spk as *const _ as *mut _,
22635                    &ktb as *const _ as *mut _,
22636                    &vtb as *const _ as *mut _,
22637                    &wini as *const _ as *mut _,
22638                ];
22639                unsafe {
22640                    self.launch_pdl_flash(
22641                        wg,
22642                        "fa_decode_vec_q_rows_v4_w_sp",
22643                        (n_head_kv as u32, n_splits_max as u32, t as u32),
22644                        (32, gqa + 1, 1),
22645                        sh,
22646                        &mut ps,
22647                    )?;
22648                }
22649            } else {
22650                let f = if wg {
22651                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
22652                } else {
22653                    self.func("fa_decode_vec_q_rows_v4_w_sp")
22654                };
22655                f.set_attribute(
22656                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22657                    sh as i32,
22658                )?;
22659                let cfg = LaunchConfig {
22660                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22661                    block_dim: (32, gqa + 1, 1),
22662                    shared_mem_bytes: sh,
22663                };
22664                let __s_b = self.gpu.stream();
22665                let mut b = __s_b.launch_builder(&f);
22666                b.arg(q)
22667                    .arg(k)
22668                    .arg(v)
22669                    .arg(&mut *part_o)
22670                    .arg(&mut *part_m)
22671                    .arg(&mut *part_l)
22672                    .arg(&hd)
22673                    .arg(&nh)
22674                    .arg(&nhkv)
22675                    .arg(base_dev)
22676                    .arg(&base_plus)
22677                    .arg(&scale)
22678                    .arg(&nspm)
22679                    .arg(&spk)
22680                    .arg(&ktb)
22681                    .arg(&vtb)
22682                    .arg(&wini);
22683                unsafe {
22684                    b.launch(cfg)?;
22685                }
22686            }
22687        } else {
22688            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
22689                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
22690                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22691                use cudarc::driver::{DevicePtr, DevicePtrMut};
22692                let s = &self.gpu.stream();
22693                let (pq, _b0) = q.device_ptr(s);
22694                let (pk, _b1) = k.device_ptr(s);
22695                let (pv, _b2) = v.device_ptr(s);
22696                let (po, _b3) = part_o.device_ptr_mut(s);
22697                let (pm, _b4) = part_m.device_ptr_mut(s);
22698                let (pl, _b5) = part_l.device_ptr_mut(s);
22699                let (pb, _b6) = base_dev.device_ptr(s);
22700                let mut ps = [
22701                    &pq as *const _ as *mut std::ffi::c_void,
22702                    &pk as *const _ as *mut _,
22703                    &pv as *const _ as *mut _,
22704                    &po as *const _ as *mut _,
22705                    &pm as *const _ as *mut _,
22706                    &pl as *const _ as *mut _,
22707                    &hd as *const _ as *mut _,
22708                    &nh as *const _ as *mut _,
22709                    &nhkv as *const _ as *mut _,
22710                    &pb as *const _ as *mut _,
22711                    &base_plus as *const _ as *mut _,
22712                    &scale as *const _ as *mut _,
22713                    &nspm as *const _ as *mut _,
22714                    &spk as *const _ as *mut _,
22715                    &ktb as *const _ as *mut _,
22716                    &vtb as *const _ as *mut _,
22717                    &wini as *const _ as *mut _,
22718                ];
22719                unsafe {
22720                    self.launch_pdl_flash(
22721                        wg,
22722                        "fa_decode_vec_q_rows_v4_w",
22723                        (n_head_kv as u32, n_splits_max as u32, t as u32),
22724                        (32, gqa, 1),
22725                        sh,
22726                        &mut ps,
22727                    )?;
22728                }
22729            } else {
22730                let pick = |name: &str| {
22731                    if wg {
22732                        self.func_g(name)
22733                    } else {
22734                        self.func(name)
22735                    }
22736                };
22737                let (f, sh) = if fa_v4_at(window) {
22738                    let f = pick("fa_decode_vec_q_rows_v4_w");
22739                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22740                } else if smem_tkv > 0 && window >= smem_tkv {
22741                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
22742                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
22743                    (
22744                        pick("fa_decode_vec_q_rows_smem_w"),
22745                        (2 * 32 * head_dim * 2) as u32,
22746                    )
22747                } else {
22748                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22749                };
22750                f.set_attribute(
22751                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22752                    sh as i32,
22753                )?;
22754                let cfg = LaunchConfig {
22755                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22756                    block_dim: (32, gqa, 1),
22757                    shared_mem_bytes: sh,
22758                };
22759                let __s_b = self.gpu.stream();
22760                let mut b = __s_b.launch_builder(&f);
22761                b.arg(q)
22762                    .arg(k)
22763                    .arg(v)
22764                    .arg(&mut *part_o)
22765                    .arg(&mut *part_m)
22766                    .arg(&mut *part_l)
22767                    .arg(&hd)
22768                    .arg(&nh)
22769                    .arg(&nhkv)
22770                    .arg(base_dev)
22771                    .arg(&base_plus)
22772                    .arg(&scale)
22773                    .arg(&nspm)
22774                    .arg(&spk)
22775                    .arg(&ktb)
22776                    .arg(&vtb)
22777                    .arg(&wini);
22778                unsafe {
22779                    b.launch(cfg)?;
22780                }
22781            }
22782        }
22783        let cfg2 = LaunchConfig {
22784            grid_dim: (n_head as u32, t as u32, 1),
22785            block_dim: (head_dim as u32, 1, 1),
22786            shared_mem_bytes: 0,
22787        };
22788        if let Some((oq, od)) = q8_out {
22789            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
22790            // consumes the pair directly; the standalone quantize launch folds away.
22791            if Self::pdl_on() && Self::pdl_wb_on() {
22792                // wave-B2: flavor mirrors the builder's wg choice.
22793                use cudarc::driver::{DevicePtr, DevicePtrMut};
22794                let s = &self.gpu.stream();
22795                let (po, _g0) = part_o.device_ptr(s);
22796                let (pm, _g1) = part_m.device_ptr(s);
22797                let (pl, _g2) = part_l.device_ptr(s);
22798                let (pq, _g3) = oq.device_ptr_mut(s);
22799                let (pd, _g4) = od.device_ptr_mut(s);
22800                let mut ps = [
22801                    &po as *const _ as *mut std::ffi::c_void,
22802                    &pm as *const _ as *mut _,
22803                    &pl as *const _ as *mut _,
22804                    &pq as *const _ as *mut _,
22805                    &pd as *const _ as *mut _,
22806                    &hd as *const _ as *mut _,
22807                    &nh as *const _ as *mut _,
22808                    &nspm as *const _ as *mut _,
22809                    &spk as *const _ as *mut _,
22810                    &wini as *const _ as *mut _,
22811                ];
22812                unsafe {
22813                    self.launch_pdl_flash(
22814                        wg,
22815                        "fa_decode_combine_rows_w_q8_1",
22816                        cfg2.grid_dim,
22817                        cfg2.block_dim,
22818                        0,
22819                        &mut ps,
22820                    )?;
22821                }
22822                return Ok(());
22823            }
22824            let fc = if wg {
22825                self.func_g("fa_decode_combine_rows_w_q8_1")
22826            } else {
22827                self.func("fa_decode_combine_rows_w_q8_1")
22828            };
22829            let __s_b2 = self.gpu.stream();
22830            let mut b2 = __s_b2.launch_builder(&fc);
22831            b2.arg(&*part_o)
22832                .arg(&*part_m)
22833                .arg(&*part_l)
22834                .arg(oq)
22835                .arg(od)
22836                .arg(&hd)
22837                .arg(&nh)
22838                .arg(&nspm)
22839                .arg(&spk)
22840                .arg(&wini);
22841            unsafe {
22842                b2.launch(cfg2)?;
22843            }
22844            return Ok(());
22845        }
22846        let fc = if wg {
22847            self.func_g("fa_decode_combine_rows_w")
22848        } else {
22849            self.func("fa_decode_combine_rows_w")
22850        };
22851        let __s_b2 = self.gpu.stream();
22852        let mut b2 = __s_b2.launch_builder(&fc);
22853        b2.arg(&*part_o)
22854            .arg(&*part_m)
22855            .arg(&*part_l)
22856            .arg(o)
22857            .arg(&hd)
22858            .arg(&nh)
22859            .arg(&nspm)
22860            .arg(&spk)
22861            .arg(&wini);
22862        unsafe {
22863            b2.launch(cfg2)?;
22864        }
22865        Ok(())
22866    }
22867
22868    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
22869    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
22870    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
22871    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
22872    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
22873    #[allow(clippy::too_many_arguments)]
22874    pub fn fa_decode_rows_dc(
22875        &self,
22876        q: &CudaSlice<f32>,
22877        k: &cudarc::driver::CudaView<u8>,
22878        v: &cudarc::driver::CudaView<u8>,
22879        o: &mut CudaSlice<f32>,
22880        head_dim: usize,
22881        n_head: usize,
22882        n_head_kv: usize,
22883        base_dev: &CudaSlice<i32>,
22884        t_kv_upper: usize,
22885        t: usize,
22886        scale: f32,
22887        k_tok_bytes: usize,
22888        v_tok_bytes: usize,
22889        base_plus: i32,
22890        g: bool,
22891    ) -> Result<(), Box<dyn std::error::Error>> {
22892        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22893        assert!(
22894            v4 || fa_v3_active(head_dim),
22895            "stream fa rows requires the v3 or v4 lane"
22896        );
22897        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22898        if v4 {
22899            let sp = fa_split_keys(t_kv_upper, n_head_kv);
22900            let n_splits_max = (t_kv_upper + sp - 1) / sp;
22901            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22902            let (nspm, spk) = (n_splits_max as i32, sp as i32);
22903            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22904            let gqa = (n_head / n_head_kv).max(1) as u32;
22905            let o_len = t * n_head * n_splits_max * head_dim;
22906            let ml_len = t * n_head * n_splits_max;
22907            let mut part_guard = self.fa_part_pool.lock().unwrap();
22908            if part_guard
22909                .as_ref()
22910                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22911                .unwrap_or(true)
22912            {
22913                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22914                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22915                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22916                // later live allocations land at those addresses, and the next graph REPLAY writes
22917                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22918                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22919                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22920                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22921                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22922                // (total retired < final size).
22923                let old = part_guard.take();
22924                let (co, cm) = old
22925                    .as_ref()
22926                    .map(|pp| (pp.0.len(), pp.1.len()))
22927                    .unwrap_or((0, 0));
22928                if let Some(old) = old {
22929                    self.fa_part_retired.lock().unwrap().push(old);
22930                }
22931                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22932                    eprintln!(
22933                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22934                        co, o_len, cm, ml_len
22935                    );
22936                }
22937                *part_guard = Some((
22938                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22939                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22940                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22941                ));
22942            }
22943            let pg = part_guard.as_mut().unwrap();
22944            self.gpu
22945                .stream()
22946                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22947            self.gpu
22948                .stream()
22949                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22950            self.gpu
22951                .stream()
22952                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22953            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22954            let f = if g {
22955                self.func_g("fa_decode_vec_q_rows_v4_dc")
22956            } else {
22957                self.func("fa_decode_vec_q_rows_v4_dc")
22958            };
22959            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22960            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22961            f.set_attribute(
22962                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22963                sh as i32,
22964            )?;
22965            let cfg = LaunchConfig {
22966                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22967                block_dim: (32, gqa, 1),
22968                shared_mem_bytes: sh,
22969            };
22970            let __s_b = self.gpu.stream();
22971            let mut b = __s_b.launch_builder(&f);
22972            b.arg(q)
22973                .arg(k)
22974                .arg(v)
22975                .arg(&mut *part_o)
22976                .arg(&mut *part_m)
22977                .arg(&mut *part_l)
22978                .arg(&hd)
22979                .arg(&nh)
22980                .arg(&nhkv)
22981                .arg(base_dev)
22982                .arg(&base_plus)
22983                .arg(&scale)
22984                .arg(&nspm)
22985                .arg(&spk)
22986                .arg(&ktb)
22987                .arg(&vtb);
22988            unsafe {
22989                b.launch(cfg)?;
22990            }
22991            let fc = self.func("fa_decode_combine_rows_dc");
22992            let cfg2 = LaunchConfig {
22993                grid_dim: (n_head as u32, t as u32, 1),
22994                block_dim: (head_dim as u32, 1, 1),
22995                shared_mem_bytes: 0,
22996            };
22997            let __s_b2 = self.gpu.stream();
22998            let mut b2 = __s_b2.launch_builder(&fc);
22999            b2.arg(&*part_o)
23000                .arg(&*part_m)
23001                .arg(&*part_l)
23002                .arg(o)
23003                .arg(&hd)
23004                .arg(&nh)
23005                .arg(base_dev)
23006                .arg(&base_plus)
23007                .arg(&nspm)
23008                .arg(&spk);
23009            unsafe {
23010                b2.launch(cfg2)?;
23011            }
23012            return Ok(());
23013        }
23014        let sp = fa_split_keys(t_kv_upper, n_head_kv);
23015        let n_splits_max = (t_kv_upper + sp - 1) / sp;
23016        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23017        let (nspm, spk) = (n_splits_max as i32, sp as i32);
23018        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23019        let gqa = (n_head / n_head_kv).max(1) as u32;
23020        let o_len = t * n_head * n_splits_max * head_dim;
23021        let ml_len = t * n_head * n_splits_max;
23022        let mut part_guard = self.fa_part_pool.lock().unwrap();
23023        if part_guard
23024            .as_ref()
23025            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23026            .unwrap_or(true)
23027        {
23028            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23029            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23030            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23031            // later live allocations land at those addresses, and the next graph REPLAY writes
23032            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23033            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23034            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23035            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23036            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23037            // (total retired < final size).
23038            let old = part_guard.take();
23039            let (co, cm) = old
23040                .as_ref()
23041                .map(|pp| (pp.0.len(), pp.1.len()))
23042                .unwrap_or((0, 0));
23043            if let Some(old) = old {
23044                self.fa_part_retired.lock().unwrap().push(old);
23045            }
23046            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23047                eprintln!(
23048                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23049                    co, o_len, cm, ml_len
23050                );
23051            }
23052            *part_guard = Some((
23053                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23054                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23055                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23056            ));
23057        }
23058        let pg = part_guard.as_mut().unwrap();
23059        self.gpu
23060            .stream()
23061            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23062        self.gpu
23063            .stream()
23064            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23065        self.gpu
23066            .stream()
23067            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23068        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23069        let f = self.func("fa_decode_vec_q_rows_v3_dc");
23070        let sh = (32 * head_dim * 2) as u32;
23071        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23072        f.set_attribute(
23073            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23074            sh as i32,
23075        )?;
23076        let cfg = LaunchConfig {
23077            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23078            block_dim: (32, gqa, 1),
23079            shared_mem_bytes: sh,
23080        };
23081        let __s_b = self.gpu.stream();
23082        let mut b = __s_b.launch_builder(&f);
23083        b.arg(q)
23084            .arg(k)
23085            .arg(v)
23086            .arg(&mut *part_o)
23087            .arg(&mut *part_m)
23088            .arg(&mut *part_l)
23089            .arg(&hd)
23090            .arg(&nh)
23091            .arg(&nhkv)
23092            .arg(base_dev)
23093            .arg(&scale)
23094            .arg(&nspm)
23095            .arg(&spk)
23096            .arg(&ktb)
23097            .arg(&vtb);
23098        unsafe {
23099            b.launch(cfg)?;
23100        }
23101        let fc = self.func("fa_decode_combine_rows_dc");
23102        let cfg2 = LaunchConfig {
23103            grid_dim: (n_head as u32, t as u32, 1),
23104            block_dim: (head_dim as u32, 1, 1),
23105            shared_mem_bytes: 0,
23106        };
23107        let plus0 = 0i32;
23108        let __s_b2 = self.gpu.stream();
23109        let mut b2 = __s_b2.launch_builder(&fc);
23110        b2.arg(&*part_o)
23111            .arg(&*part_m)
23112            .arg(&*part_l)
23113            .arg(o)
23114            .arg(&hd)
23115            .arg(&nh)
23116            .arg(base_dev)
23117            .arg(&plus0)
23118            .arg(&nspm)
23119            .arg(&spk);
23120        unsafe {
23121            b2.launch(cfg2)?;
23122        }
23123        Ok(())
23124    }
23125
23126    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
23127    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
23128    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
23129    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
23130    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
23131    ///
23132    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
23133    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
23134    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
23135    /// grouping (different but mathematically-equal log-sum-exp merge).
23136    pub fn fa_decode_dc(
23137        &self,
23138        q: &CudaSlice<f32>,
23139        k: &cudarc::driver::CudaView<u8>,
23140        v: &cudarc::driver::CudaView<u8>,
23141        o: &mut CudaSlice<f32>,
23142        head_dim: usize,
23143        n_head: usize,
23144        n_head_kv: usize,
23145        t_kv_dev: &CudaSlice<i32>,
23146        bucket_max: usize,
23147        scale: f32,
23148        k_tok_bytes: usize,
23149        v_tok_bytes: usize,
23150        g: bool,
23151    ) -> Result<(), Box<dyn std::error::Error>> {
23152        self.fa_decode_dc_q8(
23153            q,
23154            k,
23155            v,
23156            o,
23157            head_dim,
23158            n_head,
23159            n_head_kv,
23160            t_kv_dev,
23161            bucket_max,
23162            scale,
23163            k_tok_bytes,
23164            v_tok_bytes,
23165            g,
23166            None,
23167        )
23168    }
23169
23170    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
23171    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
23172    #[allow(clippy::too_many_arguments)]
23173    pub fn fa_decode_dc_q8(
23174        &self,
23175        q: &CudaSlice<f32>,
23176        k: &cudarc::driver::CudaView<u8>,
23177        v: &cudarc::driver::CudaView<u8>,
23178        o: &mut CudaSlice<f32>,
23179        head_dim: usize,
23180        n_head: usize,
23181        n_head_kv: usize,
23182        t_kv_dev: &CudaSlice<i32>,
23183        bucket_max: usize,
23184        scale: f32,
23185        k_tok_bytes: usize,
23186        v_tok_bytes: usize,
23187        g: bool,
23188        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23189    ) -> Result<(), Box<dyn std::error::Error>> {
23190        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
23191        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
23192        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
23193        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
23194        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
23195        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
23196        // 2026-07-12).
23197        let mut fa_vec =
23198            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23199        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23200            fa_vec = false;
23201        } // mirror kvmod/geom
23202        let sp = fa_split_keys(bucket_max, n_head_kv);
23203        let n_splits = if fa_vec {
23204            ((bucket_max + sp - 1) / sp).max(1)
23205        } else {
23206            ((bucket_max + 255) / 256).max(1)
23207        };
23208        let o_len = n_head * n_splits * head_dim;
23209        let ml_len = n_head * n_splits;
23210        let mut part_guard = self.fa_part_pool.lock().unwrap();
23211        if part_guard
23212            .as_ref()
23213            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23214            .unwrap_or(true)
23215        {
23216            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23217            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23218            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23219            // later live allocations land at those addresses, and the next graph REPLAY writes
23220            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23221            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23222            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23223            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23224            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23225            // (total retired < final size).
23226            let old = part_guard.take();
23227            let (co, cm) = old
23228                .as_ref()
23229                .map(|pp| (pp.0.len(), pp.1.len()))
23230                .unwrap_or((0, 0));
23231            if let Some(old) = old {
23232                self.fa_part_retired.lock().unwrap().push(old);
23233            }
23234            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23235                eprintln!(
23236                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23237                    co, o_len, cm, ml_len
23238                );
23239            }
23240            *part_guard = Some((
23241                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23242                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23243                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23244            ));
23245        }
23246        let pg = part_guard.as_mut().unwrap();
23247        self.gpu
23248            .stream()
23249            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23250        self.gpu
23251            .stream()
23252            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23253        self.gpu
23254            .stream()
23255            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23256        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23257        let (hd, nh, nhkv, nsp) = (
23258            head_dim as i32,
23259            n_head as i32,
23260            n_head_kv as i32,
23261            n_splits as i32,
23262        );
23263        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23264        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23265        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
23266        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
23267        let deep = fa_vec
23268            && head_dim == 256
23269            && fa_v4_at(bucket_max)
23270            && !g
23271            && fa_deep_at(bucket_max)
23272            && !matches!(fa_v4_mode(), "noB3" | "stage");
23273        let (f, cfg) = if fa_vec
23274            && head_dim == 512
23275            && bucket_max >= {
23276                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23277                *FA512_MIN_DC.get_or_init(|| {
23278                    std::env::var("MEMRA_FA512_MIN")
23279                        .ok()
23280                        .and_then(|v| v.parse().ok())
23281                        .unwrap_or(512)
23282                })
23283            } {
23284            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
23285            let gqa = (n_head / n_head_kv).max(1) as u32;
23286            (
23287                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23288                LaunchConfig {
23289                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23290                    block_dim: (32, gqa, 1),
23291                    shared_mem_bytes: 0,
23292                },
23293            )
23294        } else if fa_vec && head_dim == 512 {
23295            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
23296            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
23297            let q_view = q.as_view();
23298            let mut o_view = o.as_view_mut();
23299            return self.fa_decode_scalar_unified(
23300                &q_view,
23301                k,
23302                v,
23303                &mut o_view,
23304                head_dim,
23305                n_head,
23306                n_head_kv,
23307                0,
23308                Some(t_kv_dev),
23309                scale,
23310                n_splits,
23311                sp,
23312                k_tok_bytes,
23313                v_tok_bytes,
23314                g,
23315                &mut *part_o,
23316                &mut *part_m,
23317                &mut *part_l,
23318                q8_out,
23319            );
23320        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23321            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
23322            // incl the g-module route + raw-e4m3 sV sizing.
23323            let gqa = (n_head / n_head_kv).max(1) as u32;
23324            let fv = if g {
23325                self.func_g("fa_decode_vec_q_v4_dc")
23326            } else if deep {
23327                self.func("fa_decode_vec_q_v4_deep_dc")
23328            } else {
23329                self.func("fa_decode_vec_q_v4_dc")
23330            };
23331            let shmem =
23332                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23333            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23334            fv.set_attribute(
23335                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23336                shmem as i32,
23337            )?;
23338            (
23339                fv,
23340                LaunchConfig {
23341                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23342                    block_dim: (32, gqa, 1),
23343                    shared_mem_bytes: shmem,
23344                },
23345            )
23346        } else if fa_vec && fa_v3_active(head_dim) {
23347            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
23348            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
23349            let gqa = (n_head / n_head_kv).max(1) as u32;
23350            let fv = if g {
23351                self.func_g("fa_decode_vec_q_v3_dc")
23352            } else {
23353                self.func("fa_decode_vec_q_v3_dc")
23354            };
23355            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
23356            (
23357                fv,
23358                LaunchConfig {
23359                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23360                    block_dim: (32, gqa, 1),
23361                    shared_mem_bytes: shmem,
23362                },
23363            )
23364        } else if fa_vec && fa_v2_on() {
23365            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
23366            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
23367            // a numeric config; eager, rows-verify and graph all switch together).
23368            let gqa = (n_head / n_head_kv).max(1) as u32;
23369            let fv = if g {
23370                self.func_g("fa_decode_vec_q_v2_dc")
23371            } else {
23372                self.func("fa_decode_vec_q_v2_dc")
23373            };
23374            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
23375            (
23376                fv,
23377                LaunchConfig {
23378                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23379                    block_dim: (32, gqa, 1),
23380                    shared_mem_bytes: shmem,
23381                },
23382            )
23383        } else if fa_vec {
23384            let gqa = (n_head / n_head_kv).max(1) as u32;
23385            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
23386            let fv = if g {
23387                self.func_g("fa_decode_vec_q_dc")
23388            } else {
23389                self.func("fa_decode_vec_q_dc")
23390            };
23391            (
23392                fv,
23393                LaunchConfig {
23394                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23395                    block_dim: (32, gqa, 1),
23396                    shared_mem_bytes: 0,
23397                },
23398            )
23399        } else {
23400            let q_view = q.as_view();
23401            let mut o_view = o.as_view_mut();
23402            return self.fa_decode_scalar_unified(
23403                &q_view,
23404                k,
23405                v,
23406                &mut o_view,
23407                head_dim,
23408                n_head,
23409                n_head_kv,
23410                0,
23411                Some(t_kv_dev),
23412                scale,
23413                n_splits,
23414                if fa_vec { sp } else { 256 },
23415                k_tok_bytes,
23416                v_tok_bytes,
23417                g,
23418                &mut *part_o,
23419                &mut *part_m,
23420                &mut *part_l,
23421                q8_out,
23422            );
23423        };
23424        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
23425        let __s_b = self.gpu.stream();
23426        let mut b = __s_b.launch_builder(&f);
23427        b.arg(q)
23428            .arg(k)
23429            .arg(v)
23430            .arg(&mut *part_o)
23431            .arg(&mut *part_m)
23432            .arg(&mut *part_l)
23433            .arg(&hd)
23434            .arg(&nh)
23435            .arg(&nhkv)
23436            .arg(t_kv_dev)
23437            .arg(&scale)
23438            .arg(&nsp)
23439            .arg(&ski)
23440            .arg(&ktb)
23441            .arg(&vtb);
23442        unsafe {
23443            b.launch(cfg)?;
23444        }
23445        let cfg2 = LaunchConfig {
23446            grid_dim: (n_head as u32, 1, 1),
23447            block_dim: (head_dim as u32, 1, 1),
23448            shared_mem_bytes: 0,
23449        };
23450        if let Some((oq, od)) = q8_out {
23451            let fc = if g {
23452                self.func_g("fa_decode_combine_q8_1")
23453            } else {
23454                self.fa_func("fa_decode_combine_q8_1", head_dim)
23455            };
23456            let __s_b2 = self.gpu.stream();
23457            let mut b2 = __s_b2.launch_builder(&fc);
23458            b2.arg(&*part_o)
23459                .arg(&*part_m)
23460                .arg(&*part_l)
23461                .arg(oq)
23462                .arg(od)
23463                .arg(&hd)
23464                .arg(&nh)
23465                .arg(&nsp);
23466            unsafe {
23467                b2.launch(cfg2)?;
23468            }
23469            return Ok(());
23470        }
23471        let fc = if g {
23472            self.func_g("fa_decode_combine_f32")
23473        } else {
23474            self.fa_func("fa_decode_combine_f32", head_dim)
23475        };
23476        let __s_b2 = self.gpu.stream();
23477        let mut b2 = __s_b2.launch_builder(&fc);
23478        b2.arg(&*part_o)
23479            .arg(&*part_m)
23480            .arg(&*part_l)
23481            .arg(o)
23482            .arg(&hd)
23483            .arg(&nh)
23484            .arg(&nsp);
23485        unsafe {
23486            b2.launch(cfg2)?;
23487        }
23488        Ok(())
23489    }
23490
23491    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
23492    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
23493    /// at equal rows.
23494    #[allow(clippy::too_many_arguments)]
23495    pub fn append_kv_quantized_dcw(
23496        &self,
23497        k_row: &CudaSlice<f32>,
23498        v_row: &CudaSlice<f32>,
23499        kc: &mut CudaSlice<u8>,
23500        vc: &mut CudaSlice<u8>,
23501        len_dev: &CudaSlice<i32>,
23502        base_dev: Option<&CudaSlice<i32>>,
23503        kv_dim_k: usize,
23504        kv_dim_v: usize,
23505        k_tok_bytes: usize,
23506        v_tok_bytes: usize,
23507    ) -> Result<(), Box<dyn std::error::Error>> {
23508        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23509        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23510        let cfg = LaunchConfig {
23511            grid_dim: (nblk, 1, 1),
23512            block_dim: (32, 1, 1),
23513            shared_mem_bytes: 0,
23514        };
23515        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23516        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23517        let null: u64 = 0;
23518        let __s_b = self.gpu.stream();
23519        let mut b = __s_b.launch_builder(&f);
23520        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23521        match base_dev {
23522            Some(base) => {
23523                b.arg(base);
23524            }
23525            None => {
23526                b.arg(&null);
23527            }
23528        }
23529        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23530        unsafe {
23531            b.launch(cfg)?;
23532        }
23533        Ok(())
23534    }
23535
23536    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
23537    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23538        let f = self.func("inc_i32");
23539        let cfg = LaunchConfig {
23540            grid_dim: (1, 1, 1),
23541            block_dim: (1, 1, 1),
23542            shared_mem_bytes: 0,
23543        };
23544        let __s_b = self.gpu.stream();
23545        let mut b = __s_b.launch_builder(&f);
23546        b.arg(counter);
23547        unsafe {
23548            b.launch(cfg)?;
23549        }
23550        Ok(())
23551    }
23552
23553    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
23554    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
23555    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
23556    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
23557    /// kernel class on this lane); callers keep eager below the vec floor and for any other
23558    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
23559    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
23560    /// alive across bucket growth.
23561    #[allow(clippy::too_many_arguments)]
23562    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
23563    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
23564    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
23565    fn fa_part_pool_grow(
23566        &self,
23567        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
23568        o_len: usize,
23569        ml_len: usize,
23570    ) -> Result<(), Box<dyn std::error::Error>> {
23571        if part_guard
23572            .as_ref()
23573            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23574            .unwrap_or(true)
23575        {
23576            let old = part_guard.take();
23577            let (co, cm) = old
23578                .as_ref()
23579                .map(|pp| (pp.0.len(), pp.1.len()))
23580                .unwrap_or((0, 0));
23581            if let Some(old) = old {
23582                self.fa_part_retired.lock().unwrap().push(old);
23583            }
23584            *part_guard = Some((
23585                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23586                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23587                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23588            ));
23589        }
23590        Ok(())
23591    }
23592
23593    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
23594    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
23595    pub fn fa_dcw_pool_ensure(
23596        &self,
23597        head_dim: usize,
23598        n_head: usize,
23599        n_head_kv: usize,
23600        bucket_max: usize,
23601    ) -> Result<(), Box<dyn std::error::Error>> {
23602        let sp = fa_split_keys(bucket_max, n_head_kv);
23603        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23604        let o_len = n_head * n_splits * head_dim;
23605        let ml_len = n_head * n_splits;
23606        let mut part_guard = self.fa_part_pool.lock().unwrap();
23607        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
23608    }
23609
23610    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
23611    /// appended; one launch walks the KV stream once with two query rows (per-row causal
23612    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
23613    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
23614    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
23615    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
23616    /// outputs (the head gate fuses into the combine as in the t=1 path).
23617    #[allow(clippy::too_many_arguments)]
23618    pub fn fa_decode_dcw2(
23619        &self,
23620        q2: &CudaSlice<f32>,
23621        k_ring: &cudarc::driver::CudaView<u8>,
23622        v_ring: &cudarc::driver::CudaView<u8>,
23623        o2: &mut CudaSlice<f32>,
23624        head_dim: usize,
23625        n_head: usize,
23626        n_head_kv: usize,
23627        len_dev: &CudaSlice<i32>,
23628        base_dev: Option<&CudaSlice<i32>>,
23629        window: usize,
23630        bucket_max: usize,
23631        scale: f32,
23632        k_tok_bytes: usize,
23633        v_tok_bytes: usize,
23634        gate2: &CudaSlice<f32>,
23635    ) -> Result<(), Box<dyn std::error::Error>> {
23636        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23637        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23638            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
23639        }
23640        let sp = fa_split_keys(bucket_max, n_head_kv);
23641        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23642        // Partials for BOTH rows: row-major halves.
23643        let o_len = 2 * n_head * n_splits * head_dim;
23644        let ml_len = 2 * n_head * n_splits;
23645        let mut part_guard = self.fa_part_pool.lock().unwrap();
23646        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23647        let pg = part_guard.as_mut().unwrap();
23648        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23649        let (hd, nh, nhkv, nsp) = (
23650            head_dim as i32,
23651            n_head as i32,
23652            n_head_kv as i32,
23653            n_splits as i32,
23654        );
23655        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23656        let (ski, win) = (sp as i32, window as i32);
23657        let gqa = (n_head / n_head_kv).max(1) as u32;
23658        let smem = (32 * head_dim * 2) as u32;
23659        let f = self.func("fa_decode_vec_q_v3_dcw2");
23660        let cfg = LaunchConfig {
23661            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23662            block_dim: (32, gqa, 1),
23663            shared_mem_bytes: smem,
23664        };
23665        let null: u64 = 0;
23666        {
23667            let __s_b = self.gpu.stream();
23668            let mut b = __s_b.launch_builder(&f);
23669            b.arg(q2)
23670                .arg(k_ring)
23671                .arg(v_ring)
23672                .arg(&mut *part_o)
23673                .arg(&mut *part_m)
23674                .arg(&mut *part_l)
23675                .arg(&hd)
23676                .arg(&nh)
23677                .arg(&nhkv)
23678                .arg(len_dev);
23679            match base_dev {
23680                Some(base) => {
23681                    b.arg(base);
23682                }
23683                None => {
23684                    b.arg(&null);
23685                }
23686            }
23687            b.arg(&win)
23688                .arg(&scale)
23689                .arg(&nsp)
23690                .arg(&ski)
23691                .arg(&ktb)
23692                .arg(&vtb);
23693            unsafe {
23694                b.launch(cfg)?;
23695            }
23696        }
23697        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
23698        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
23699        // one launch covers both rows with the exact t=1 program per (row, head).
23700        let fc = {
23701            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23702            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23703                self.func("fa_decode_combine_gate_f32_s")
23704            } else {
23705                self.func("fa_decode_combine_gate_f32")
23706            }
23707        };
23708        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23709        let nh2 = (2 * n_head) as i32;
23710        let cfg2 = LaunchConfig {
23711            grid_dim: ((2 * n_head) as u32, 1, 1),
23712            block_dim: (head_dim as u32, 1, 1),
23713            shared_mem_bytes: if combine_shared {
23714                (2 * n_splits * 4) as u32
23715            } else {
23716                0
23717            },
23718        };
23719        let __s_b2 = self.gpu.stream();
23720        let mut b2 = __s_b2.launch_builder(&fc);
23721        b2.arg(&*part_o)
23722            .arg(&*part_m)
23723            .arg(&*part_l)
23724            .arg(gate2)
23725            .arg(o2)
23726            .arg(&hd)
23727            .arg(&nh2)
23728            .arg(&nsp);
23729        unsafe {
23730            b2.launch(cfg2)?;
23731        }
23732        Ok(())
23733    }
23734
23735    /// T-ROW dcw decode attention over a per-row session table (the per-session
23736    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
23737    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
23738    /// program verbatim with that row's ring/len/base and its own split geometry, so each
23739    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
23740    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
23741    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
23742    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
23743    #[allow(clippy::too_many_arguments)]
23744    pub fn fa_decode_dcw_rows(
23745        &self,
23746        q_rows: &CudaSlice<f32>,
23747        tab: &CudaSlice<u64>,
23748        o_rows: &mut CudaSlice<f32>,
23749        t: usize,
23750        head_dim: usize,
23751        n_head: usize,
23752        n_head_kv: usize,
23753        window: usize,
23754        max_ns: usize,
23755        scale: f32,
23756        k_tok_bytes: usize,
23757        v_tok_bytes: usize,
23758        gate_rows: &CudaSlice<f32>,
23759    ) -> Result<(), Box<dyn std::error::Error>> {
23760        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
23761            || head_dim > 256
23762            || head_dim % 32 != 0
23763            || !fa_v3_on()
23764        {
23765            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
23766        }
23767        if fa_sm_count() < 128
23768            || std::env::var("MEMRA_FA_SPLIT").is_ok()
23769            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
23770            || std::env::var("MEMRA_FA_SP16").is_ok()
23771        {
23772            return Err(
23773                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
23774                 (or a <128-SM rig) keep the per-row path"
23775                    .into(),
23776            );
23777        }
23778        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
23779            return Err("fa_decode_dcw_rows geometry".into());
23780        }
23781        let o_len = t * n_head * max_ns * head_dim;
23782        let ml_len = t * n_head * max_ns;
23783        let mut part_guard = self.fa_part_pool.lock().unwrap();
23784        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23785        let pg = part_guard.as_mut().unwrap();
23786        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23787        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23788        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23789        let (win, mns) = (window as i32, max_ns as i32);
23790        let gqa = (n_head / n_head_kv).max(1) as u32;
23791        let smem = (32 * head_dim * 2) as u32;
23792        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
23793        let cfg = LaunchConfig {
23794            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
23795            block_dim: (32, gqa, 1),
23796            shared_mem_bytes: smem,
23797        };
23798        {
23799            let __s_b = self.gpu.stream();
23800            let mut b = __s_b.launch_builder(&f);
23801            b.arg(q_rows)
23802                .arg(tab)
23803                .arg(&mut *part_o)
23804                .arg(&mut *part_m)
23805                .arg(&mut *part_l)
23806                .arg(&hd)
23807                .arg(&nh)
23808                .arg(&nhkv)
23809                .arg(&win)
23810                .arg(&scale)
23811                .arg(&mns)
23812                .arg(&ktb)
23813                .arg(&vtb);
23814            unsafe {
23815                b.launch(cfg)?;
23816            }
23817        }
23818        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
23819        // row r head h reads its own partial bank; splits past a row's ns_eff carry
23820        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
23821        let fc = {
23822            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23823            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23824                self.func("fa_decode_combine_gate_f32_s")
23825            } else {
23826                self.func("fa_decode_combine_gate_f32")
23827            }
23828        };
23829        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23830        let nht = (t * n_head) as i32;
23831        let cfg2 = LaunchConfig {
23832            grid_dim: ((t * n_head) as u32, 1, 1),
23833            block_dim: (head_dim as u32, 1, 1),
23834            shared_mem_bytes: if combine_shared {
23835                (2 * max_ns * 4) as u32
23836            } else {
23837                0
23838            },
23839        };
23840        let __s_b2 = self.gpu.stream();
23841        let mut b2 = __s_b2.launch_builder(&fc);
23842        b2.arg(&*part_o)
23843            .arg(&*part_m)
23844            .arg(&*part_l)
23845            .arg(gate_rows)
23846            .arg(o_rows)
23847            .arg(&hd)
23848            .arg(&nht)
23849            .arg(&mns);
23850        unsafe {
23851            b2.launch(cfg2)?;
23852        }
23853        Ok(())
23854    }
23855
23856    pub fn fa_decode_dcw(
23857        &self,
23858        q: &CudaSlice<f32>,
23859        k_ring: &cudarc::driver::CudaView<u8>,
23860        v_ring: &cudarc::driver::CudaView<u8>,
23861        o: &mut CudaSlice<f32>,
23862        head_dim: usize,
23863        n_head: usize,
23864        n_head_kv: usize,
23865        len_dev: &CudaSlice<i32>,
23866        base_dev: Option<&CudaSlice<i32>>,
23867        window: usize,
23868        bucket_max: usize,
23869        scale: f32,
23870        k_tok_bytes: usize,
23871        v_tok_bytes: usize,
23872        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
23873        // one launch saved); `o` then receives the GATED output and the caller skips its
23874        // attn_head_gate call.
23875        fused_gate: Option<&CudaSlice<f32>>,
23876    ) -> Result<(), Box<dyn std::error::Error>> {
23877        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23878        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23879            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"
23880                .into());
23881        }
23882        let sp = fa_split_keys(bucket_max, n_head_kv);
23883        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23884        let o_len = n_head * n_splits * head_dim;
23885        let ml_len = n_head * n_splits;
23886        let mut part_guard = self.fa_part_pool.lock().unwrap();
23887        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23888        let pg = part_guard.as_mut().unwrap();
23889        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
23890        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
23891        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
23892        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
23893        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23894        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
23895        // finds the attention children BY their three-memset signature and updates the
23896        // memset widths per bucket — capturing without them silently kills retargeting
23897        // (battery-v8 token drift, 2026-08-21).
23898        let memset_on = *MEMSET_ON
23899            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23900            || crate::tp::token_graph_building();
23901        if memset_on {
23902            self.gpu
23903                .stream()
23904                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23905            self.gpu
23906                .stream()
23907                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23908            self.gpu
23909                .stream()
23910                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23911        }
23912        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23913        let (hd, nh, nhkv, nsp) = (
23914            head_dim as i32,
23915            n_head as i32,
23916            n_head_kv as i32,
23917            n_splits as i32,
23918        );
23919        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23920        let (ski, win) = (sp as i32, window as i32);
23921        let gqa = (n_head / n_head_kv).max(1) as u32;
23922        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
23923        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
23924        // see fa_dec_v3_walk_u). Same launch geometry.
23925        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23926        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23927        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23928            Ok("2") => 2,
23929            Ok("1") => 1,
23930            _ => 0,
23931        });
23932        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
23933        // permission-blocked in this container and the module params are not exposed, so this
23934        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
23935        // prints cumulative cycle shares every 430 launches.
23936        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23937        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23938        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23939            std::sync::Mutex::new(None);
23940        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
23941        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
23942        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
23943        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23944        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23945            && (n_head / n_head_kv) % 2 == 0
23946            && (n_head / n_head_kv) >= 2;
23947        let f = if fprof {
23948            self.func("fa_decode_vec_q_v3_dcw_prof")
23949        } else if hs2 {
23950            self.func("fa_decode_vec_q_v3_dcw_hs2")
23951        } else if hoist == 2 {
23952            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
23953            self.func("fa_decode_vec_q_v3_dcw_hc")
23954        } else if hoist == 1 {
23955            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
23956            self.func("fa_decode_vec_q_v3_dcw_h")
23957        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23958            self.func("fa_decode_vec_q_v3_dcw_u8")
23959        } else {
23960            self.func("fa_decode_vec_q_v3_dcw")
23961        };
23962        let cfg = LaunchConfig {
23963            grid_dim: if hs2 {
23964                ((2 * n_head_kv) as u32, n_splits as u32, 1)
23965            } else {
23966                (n_head_kv as u32, n_splits as u32, 1)
23967            },
23968            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23969            shared_mem_bytes: smem,
23970        };
23971        let null: u64 = 0;
23972        let __s_b = self.gpu.stream();
23973        let mut b = __s_b.launch_builder(&f);
23974        b.arg(q)
23975            .arg(k_ring)
23976            .arg(v_ring)
23977            .arg(&mut *part_o)
23978            .arg(&mut *part_m)
23979            .arg(&mut *part_l)
23980            .arg(&hd)
23981            .arg(&nh)
23982            .arg(&nhkv)
23983            .arg(len_dev);
23984        match base_dev {
23985            Some(base) => {
23986                b.arg(base);
23987            }
23988            None => {
23989                b.arg(&null);
23990            }
23991        }
23992        b.arg(&win)
23993            .arg(&scale)
23994            .arg(&nsp)
23995            .arg(&ski)
23996            .arg(&ktb)
23997            .arg(&vtb);
23998        if fprof {
23999            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24000            if guard
24001                .as_ref()
24002                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24003            {
24004                *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24005            }
24006            let (_, buf) = guard.as_mut().expect("armed above");
24007            b.arg(&*buf);
24008            unsafe {
24009                b.launch(cfg)?;
24010            }
24011            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24012            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24013            if n % 430 == 0 {
24014                self.stream().synchronize()?;
24015                let h = self.dtoh_u64(buf)?;
24016                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24017                let tot: u64 = h[..6].iter().sum();
24018                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24019                for (i, name) in phases.iter().enumerate() {
24020                    let pct = if tot > 0 {
24021                        h[i] as f64 / tot as f64 * 100.0
24022                    } else {
24023                        0.0
24024                    };
24025                    line.push_str(&format!(" {name}={pct:.1}%"));
24026                }
24027                if h[6] > 0 {
24028                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24029                }
24030                eprintln!("{line}");
24031            }
24032        } else {
24033            unsafe {
24034                b.launch(cfg)?;
24035            }
24036        }
24037        let mut combine_shared = false;
24038        let fc = if fused_gate.is_some() {
24039            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
24040            // n_splits-deep dependent global load chain every thread used to walk twice).
24041            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24042            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24043                combine_shared = true;
24044                self.func("fa_decode_combine_gate_f32_s")
24045            } else {
24046                self.func("fa_decode_combine_gate_f32")
24047            }
24048        } else {
24049            self.fa_func("fa_decode_combine_f32", head_dim)
24050        };
24051        let cfg2 = LaunchConfig {
24052            grid_dim: (n_head as u32, 1, 1),
24053            block_dim: (head_dim as u32, 1, 1),
24054            shared_mem_bytes: if combine_shared {
24055                (2 * n_splits * 4) as u32
24056            } else {
24057                0
24058            },
24059        };
24060        let __s_b2 = self.gpu.stream();
24061        let mut b2 = __s_b2.launch_builder(&fc);
24062        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24063        if let Some(gate_row) = fused_gate {
24064            b2.arg(gate_row);
24065        }
24066        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24067        unsafe {
24068            b2.launch(cfg2)?;
24069        }
24070        Ok(())
24071    }
24072
24073    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
24074    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
24075    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
24076    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
24077    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
24078    pub fn fa_geom_eager(
24079        &self,
24080        t_kv: usize,
24081        head_dim: usize,
24082        n_head_kv: usize,
24083        g: bool,
24084    ) -> (bool, usize) {
24085        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
24086        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
24087        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
24088        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24089        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
24090        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
24091        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
24092        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
24093        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
24094        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24095        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24096        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
24097        // family; everything else falls to the g-module scalar.
24098        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
24099        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
24100        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
24101        if g && head_dim == 256 && !fa_v4_at(t_kv) {
24102            fa_vec = false;
24103        }
24104        let sp = fa_split_keys(t_kv, n_head_kv);
24105        let n_splits = if fa_vec {
24106            ((t_kv + sp - 1) / sp).max(1)
24107        } else {
24108            ((t_kv + 255) / 256).max(1)
24109        };
24110        (fa_vec, n_splits)
24111    }
24112
24113    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
24114    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
24115    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
24116    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
24117    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
24118    pub fn fa_bucket_key(
24119        &self,
24120        t_kv: usize,
24121        head_dim: usize,
24122        n_head_kv: usize,
24123        g: bool,
24124    ) -> (bool, usize) {
24125        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24126    }
24127
24128    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
24129    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
24130    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
24131    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
24132    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
24133    /// device data) — every per-step varying scalar must come from a device counter. Returns the
24134    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
24135    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
24136    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
24137    /// replays (transients returning to the pool get reused by unrelated work and corrupt
24138    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
24139    pub fn capture_graph_retained<F>(
24140        &self,
24141        step: F,
24142    ) -> Result<
24143        (
24144            cudarc::driver::CudaGraph,
24145            Vec<Box<dyn std::any::Any + Send>>,
24146        ),
24147        Box<dyn std::error::Error>,
24148    >
24149    where
24150        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24151    {
24152        use cudarc::driver::sys::CUgraphInstantiate_flags;
24153        self.capture_graph_retained_flags(
24154            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24155            step,
24156        )
24157    }
24158
24159    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
24160    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
24161    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
24162    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
24163    pub fn capture_graph_retained_flags<F>(
24164        &self,
24165        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24166        mut step: F,
24167    ) -> Result<
24168        (
24169            cudarc::driver::CudaGraph,
24170            Vec<Box<dyn std::any::Any + Send>>,
24171        ),
24172        Box<dyn std::error::Error>,
24173    >
24174    where
24175        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24176    {
24177        use cudarc::driver::sys::CUstreamCaptureMode;
24178        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
24179        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
24180        // while the capture region is open become dead copy NODES replayed every launch
24181        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
24182        // warmup runs allocate the same transient sequence at the same pool addresses, so
24183        // retaining the warmup clones preserves the draft-graph fix without polluting the
24184        // captured graph.
24185        self.capture_keep.lock().unwrap().clear();
24186        let was_tracking = self.gpu.ctx.is_event_tracking();
24187        if was_tracking {
24188            unsafe {
24189                self.gpu.ctx.disable_event_tracking();
24190            }
24191        }
24192        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24193            self.capture_keep_on
24194                .store(true, std::sync::atomic::Ordering::Relaxed);
24195            let w = (|| {
24196                step(self)?;
24197                step(self)
24198            })();
24199            self.capture_keep_on
24200                .store(false, std::sync::atomic::Ordering::Relaxed);
24201            w?;
24202            self.gpu.stream().synchronize()?;
24203            self.gpu
24204                .stream()
24205                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24206            let r = step(self);
24207            let g = self.gpu.stream().end_capture(flags);
24208            r?;
24209            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24210            graph.upload()?;
24211            Ok(graph)
24212        };
24213        let result = run();
24214        self.capture_keep_on
24215            .store(false, std::sync::atomic::Ordering::Relaxed);
24216        if was_tracking {
24217            unsafe {
24218                self.gpu.ctx.enable_event_tracking();
24219            }
24220        }
24221        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24222        Ok((result?, keeper))
24223    }
24224
24225    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
24226    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
24227    /// alloc-free with persistent operands, and their bodies carry device side effects
24228    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
24229    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
24230    pub fn capture_graph_retained_nowarm<F>(
24231        &self,
24232        mut step: F,
24233    ) -> Result<
24234        (
24235            cudarc::driver::CudaGraph,
24236            Vec<Box<dyn std::any::Any + Send>>,
24237        ),
24238        Box<dyn std::error::Error>,
24239    >
24240    where
24241        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24242    {
24243        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24244        let was_tracking = self.gpu.ctx.is_event_tracking();
24245        if was_tracking {
24246            unsafe {
24247                self.gpu.ctx.disable_event_tracking();
24248            }
24249        }
24250        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24251            self.gpu.stream().synchronize()?;
24252            self.gpu
24253                .stream()
24254                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24255            let r = step(self);
24256            let g = self.gpu.stream().end_capture(
24257                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24258            );
24259            r?;
24260            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24261            graph.upload()?;
24262            Ok(graph)
24263        };
24264        let result = run();
24265        if was_tracking {
24266            unsafe {
24267                self.gpu.ctx.enable_event_tracking();
24268            }
24269        }
24270        Ok((result?, Vec::new()))
24271    }
24272
24273    pub fn capture_graph<F>(
24274        &self,
24275        mut step: F,
24276    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24277    where
24278        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24279    {
24280        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24281        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
24282        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
24283        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
24284        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
24285        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
24286        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
24287        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
24288        let was_tracking = self.gpu.ctx.is_event_tracking();
24289        if was_tracking {
24290            unsafe {
24291                self.gpu.ctx.disable_event_tracking();
24292            }
24293        }
24294        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
24295        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
24296        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
24297        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
24298        // measure that scan's real cost on the generic path. Diagnostic door only; the
24299        // default stays AUTO_FREE until a measured A/B justifies moving it.
24300        let iflag = {
24301            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24302            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24303                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
24304                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
24305                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24306                Ok("priority") => {
24307                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24308                }
24309                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24310            })
24311        };
24312        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
24313        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
24314        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
24315        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
24316        // eager step executions and are node-count-invariant. Printing the split bounds the
24317        // refactor's ceiling instead of assuming it.
24318        let ct = {
24319            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24320            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24321        };
24322        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
24323        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
24324        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
24325        // chased, and node-count-invariant, so no capture-body refactor could touch it.
24326        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
24327        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
24328        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
24329        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
24330        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
24331        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
24332        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
24333        // grow and never frees, resident counters/scratch, cache set in place), and the
24334        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
24335        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
24336        // settling and pool mapping. Arbitrated adversarially, not by taste:
24337        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
24338        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
24339        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
24340        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
24341        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
24342        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
24343        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
24344        let warmups = {
24345            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24346            *W.get_or_init(|| {
24347                std::env::var("MEMRA_GRAPH_WARMUPS")
24348                    .ok()
24349                    .and_then(|v| v.parse().ok())
24350                    .filter(|n| *n >= 1)
24351                    .unwrap_or(1)
24352            })
24353        };
24354        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24355            let t_w = std::time::Instant::now();
24356            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
24357            for _ in 0..warmups {
24358                step(self)?;
24359            }
24360            self.gpu.stream().synchronize()?;
24361            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24362            // capture the third run.
24363            let t_c = std::time::Instant::now();
24364            self.gpu
24365                .stream()
24366                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24367            // If the body errors mid-capture, end the capture before propagating so the stream isn't
24368            // left in a capturing state.
24369            let r = step(self);
24370            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24371            let t_i = std::time::Instant::now();
24372            let g = self.gpu.stream().end_capture(iflag);
24373            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24374            r?;
24375            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24376            let t_u = std::time::Instant::now();
24377            graph.upload()?;
24378            if ct {
24379                println!(
24380                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
24381                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
24382                    t_u.elapsed().as_secs_f64() * 1e3
24383                );
24384            }
24385            Ok(graph)
24386        };
24387        let result = run();
24388        if was_tracking {
24389            unsafe {
24390                self.gpu.ctx.enable_event_tracking();
24391            }
24392        }
24393        result
24394    }
24395
24396    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
24397    pub fn gdn_scan_s128_view(
24398        &self,
24399        q: &CudaSlice<f32>,
24400        k: &CudaSlice<f32>,
24401        v: &CudaSlice<f32>,
24402        g: &CudaSlice<f32>,
24403        beta: &CudaSlice<f32>,
24404        state_in: &cudarc::driver::CudaView<f32>,
24405        state_out: &mut cudarc::driver::CudaViewMut<f32>,
24406        o: &mut CudaSlice<f32>,
24407        n_head: usize,
24408        t: usize,
24409        scale: f32,
24410    ) -> Result<(), Box<dyn std::error::Error>> {
24411        let f = self.func("gdn_scan_s128");
24412        const S_V: u32 = 128;
24413        const WARP: u32 = 32;
24414        const COLS: u32 = 4;
24415        let cfg = LaunchConfig {
24416            grid_dim: (n_head as u32, 1, S_V / COLS),
24417            block_dim: (WARP, COLS, 1),
24418            shared_mem_bytes: 0,
24419        };
24420        let (h, ti) = (n_head as i32, t as i32);
24421        let __s_b = self.gpu.stream();
24422        let mut b = __s_b.launch_builder(&f);
24423        b.arg(q)
24424            .arg(k)
24425            .arg(v)
24426            .arg(g)
24427            .arg(beta)
24428            .arg(state_in)
24429            .arg(state_out)
24430            .arg(o)
24431            .arg(&h)
24432            .arg(&ti)
24433            .arg(&scale);
24434        unsafe {
24435            b.launch(cfg)?;
24436        }
24437        Ok(())
24438    }
24439
24440    /// conv1d where the input is a CudaView (resident conv state assembled in place).
24441    pub fn ssm_conv1d_view(
24442        &self,
24443        x: &cudarc::driver::CudaView<f32>,
24444        w: &CudaSlice<f32>,
24445        y: &mut CudaSlice<f32>,
24446        conv_dim: usize,
24447        t: usize,
24448        d_conv: usize,
24449        silu: bool,
24450    ) -> Result<(), Box<dyn std::error::Error>> {
24451        let f = self.func("ssm_conv1d_silu_f32");
24452        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
24453        let cfg = LaunchConfig {
24454            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24455            block_dim: (256, 1, 1),
24456            shared_mem_bytes: 0,
24457        };
24458        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24459        let __s_b = self.gpu.stream();
24460        let mut b = __s_b.launch_builder(&f);
24461        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24462        unsafe {
24463            b.launch(cfg)?;
24464        }
24465        Ok(())
24466    }
24467
24468    /// Depthwise causal conv1d + optional SiLU.
24469    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
24470    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
24471    /// FUSED prefill conv (token-major input, zero left-state): replaces
24472    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
24473    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
24474    pub fn ssm_conv1d_tm(
24475        &self,
24476        qkv_tm: &CudaSlice<f32>,
24477        w: &CudaSlice<f32>,
24478        y: &mut CudaSlice<f32>,
24479        conv_dim: usize,
24480        t: usize,
24481        d_conv: usize,
24482    ) -> Result<(), Box<dyn std::error::Error>> {
24483        let f = self.func("ssm_conv1d_tm_f32");
24484        let cfg = LaunchConfig {
24485            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24486            block_dim: (256, 1, 1),
24487            shared_mem_bytes: 0,
24488        };
24489        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24490        let __s_b = self.gpu.stream();
24491        let mut b = __s_b.launch_builder(&f);
24492        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24493        unsafe {
24494            b.launch(cfg)?;
24495        }
24496        Ok(())
24497    }
24498
24499    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
24500    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
24501    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
24502    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
24503    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
24504    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
24505    /// columns; the final ring == what T sequential decode ring rolls leave).
24506    pub fn ssm_conv1d_tm_state(
24507        &self,
24508        qkv_tm: &CudaSlice<f32>,
24509        conv_state: &mut CudaSlice<f32>,
24510        w: &CudaSlice<f32>,
24511        y: &mut CudaSlice<f32>,
24512        conv_dim: usize,
24513        t: usize,
24514        d_conv: usize,
24515    ) -> Result<(), Box<dyn std::error::Error>> {
24516        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24517    }
24518
24519    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
24520    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
24521    #[allow(clippy::too_many_arguments)]
24522    pub fn ssm_conv1d_tm_state_pad(
24523        &self,
24524        qkv_tm: &CudaSlice<f32>,
24525        conv_state: &mut CudaSlice<f32>,
24526        w: &CudaSlice<f32>,
24527        y: &mut CudaSlice<f32>,
24528        conv_dim: usize,
24529        t: usize,
24530        d_conv: usize,
24531        pad_len: Option<&CudaSlice<i32>>,
24532    ) -> Result<(), Box<dyn std::error::Error>> {
24533        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24534        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
24535        // the window kernel both read the pre-roll ring; the roll launches after both) — but
24536        // cloning first keeps the ordering trivially correct under any future stream split.
24537        let ring_old = if t < d_conv - 1 {
24538            Some(self.clone_dtod(conv_state)?)
24539        } else {
24540            None
24541        };
24542        {
24543            let f = self.func("ssm_conv1d_tm_state_f32");
24544            let cfg = LaunchConfig {
24545                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24546                block_dim: (256, 1, 1),
24547                shared_mem_bytes: 0,
24548            };
24549            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24550            let __s_b = self.gpu.stream();
24551            let mut b = __s_b.launch_builder(&f);
24552            b.arg(qkv_tm)
24553                .arg(&*conv_state)
24554                .arg(w)
24555                .arg(y)
24556                .arg(&cd)
24557                .arg(&ti)
24558                .arg(&dc);
24559            unsafe {
24560                b.launch(cfg)?;
24561            }
24562        }
24563        match (ring_old, pad_len) {
24564            (None, Some(len_d)) => {
24565                let f = self.func("ssm_conv_ring_update_dev_f32");
24566                let n = conv_dim * (d_conv - 1);
24567                let cfg = LaunchConfig::for_num_elems(n as u32);
24568                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24569                let __s_b = self.gpu.stream();
24570                let mut b = __s_b.launch_builder(&f);
24571                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24572                unsafe {
24573                    b.launch(cfg)?;
24574                }
24575            }
24576            (None, None) => {
24577                let f = self.func("ssm_conv_ring_update_f32");
24578                let n = conv_dim * (d_conv - 1);
24579                let cfg = LaunchConfig::for_num_elems(n as u32);
24580                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24581                let __s_b = self.gpu.stream();
24582                let mut b = __s_b.launch_builder(&f);
24583                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24584                unsafe {
24585                    b.launch(cfg)?;
24586                }
24587            }
24588            (Some(old), _) => {
24589                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
24590            }
24591        }
24592        Ok(())
24593    }
24594
24595    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
24596    pub fn ssm_conv1d_tm_state_pad_v(
24597        &self,
24598        qkv_tm: &cudarc::driver::CudaView<f32>,
24599        conv_state: &mut CudaSlice<f32>,
24600        w: &CudaSlice<f32>,
24601        y: &mut CudaSlice<f32>,
24602        conv_dim: usize,
24603        t: usize,
24604        d_conv: usize,
24605        pad_len: Option<&CudaSlice<i32>>,
24606    ) -> Result<(), Box<dyn std::error::Error>> {
24607        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24608        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
24609        // the window kernel both read the pre-roll ring; the roll launches after both) — but
24610        // cloning first keeps the ordering trivially correct under any future stream split.
24611        let ring_old = if t < d_conv - 1 {
24612            Some(self.clone_dtod(conv_state)?)
24613        } else {
24614            None
24615        };
24616        {
24617            let f = self.func("ssm_conv1d_tm_state_f32");
24618            let cfg = LaunchConfig {
24619                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24620                block_dim: (256, 1, 1),
24621                shared_mem_bytes: 0,
24622            };
24623            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24624            let __s_b = self.gpu.stream();
24625            let mut b = __s_b.launch_builder(&f);
24626            b.arg(qkv_tm)
24627                .arg(&*conv_state)
24628                .arg(w)
24629                .arg(y)
24630                .arg(&cd)
24631                .arg(&ti)
24632                .arg(&dc);
24633            unsafe {
24634                b.launch(cfg)?;
24635            }
24636        }
24637        match (ring_old, pad_len) {
24638            (None, Some(len_d)) => {
24639                let f = self.func("ssm_conv_ring_update_dev_f32");
24640                let n = conv_dim * (d_conv - 1);
24641                let cfg = LaunchConfig::for_num_elems(n as u32);
24642                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24643                let __s_b = self.gpu.stream();
24644                let mut b = __s_b.launch_builder(&f);
24645                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24646                unsafe {
24647                    b.launch(cfg)?;
24648                }
24649            }
24650            (None, None) => {
24651                let f = self.func("ssm_conv_ring_update_f32");
24652                let n = conv_dim * (d_conv - 1);
24653                let cfg = LaunchConfig::for_num_elems(n as u32);
24654                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24655                let __s_b = self.gpu.stream();
24656                let mut b = __s_b.launch_builder(&f);
24657                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24658                unsafe {
24659                    b.launch(cfg)?;
24660                }
24661            }
24662            (Some(_), _) => unreachable!(
24663                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
24664            ),
24665        }
24666        Ok(())
24667    }
24668
24669    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
24670    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
24671    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
24672    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
24673    pub fn ssm_conv_ring_rebuild(
24674        &self,
24675        qkv_tm: &CudaSlice<f32>,
24676        ring_old: &CudaSlice<f32>,
24677        conv_state: &mut CudaSlice<f32>,
24678        conv_dim: usize,
24679        tc: usize,
24680        d_conv: usize,
24681    ) -> Result<(), Box<dyn std::error::Error>> {
24682        let f = self.func("ssm_conv_ring_rebuild_f32");
24683        let n = conv_dim * (d_conv - 1);
24684        let cfg = LaunchConfig::for_num_elems(n as u32);
24685        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
24686        let __s_b = self.gpu.stream();
24687        let mut b = __s_b.launch_builder(&f);
24688        b.arg(qkv_tm)
24689            .arg(ring_old)
24690            .arg(conv_state)
24691            .arg(&cd)
24692            .arg(&ti)
24693            .arg(&dc);
24694        unsafe {
24695            b.launch(cfg)?;
24696        }
24697        Ok(())
24698    }
24699
24700    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
24701    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
24702    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
24703    /// the argmax + run-spec gates are the authority.
24704    #[allow(clippy::too_many_arguments)]
24705    pub fn gdn_prep_decode(
24706        &self,
24707        conv_out: &CudaSlice<f32>,
24708        beta_raw: &CudaSlice<f32>,
24709        alpha: &CudaSlice<f32>,
24710        dt_bias: &CudaSlice<f32>,
24711        a: &CudaSlice<f32>,
24712        q_l2: &mut CudaSlice<f32>,
24713        k_l2: &mut CudaSlice<f32>,
24714        v_g: &mut CudaSlice<f32>,
24715        beta: &mut CudaSlice<f32>,
24716        g_log: &mut CudaSlice<f32>,
24717        d_state: usize,
24718        num_v: usize,
24719        num_k: usize,
24720        key_dim: usize,
24721        eps: f32,
24722    ) -> Result<(), Box<dyn std::error::Error>> {
24723        let f = self.func("gdn_prep_decode_f32");
24724        let cfg = LaunchConfig {
24725            grid_dim: (num_v as u32, 1, 1),
24726            block_dim: (32, 4, 1),
24727            shared_mem_bytes: 0,
24728        };
24729        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24730        let __s_b = self.gpu.stream();
24731        let mut b = __s_b.launch_builder(&f);
24732        b.arg(conv_out)
24733            .arg(beta_raw)
24734            .arg(alpha)
24735            .arg(dt_bias)
24736            .arg(a)
24737            .arg(q_l2)
24738            .arg(k_l2)
24739            .arg(v_g)
24740            .arg(beta)
24741            .arg(g_log)
24742            .arg(&ds)
24743            .arg(&nv)
24744            .arg(&nk)
24745            .arg(&kd)
24746            .arg(&eps);
24747        unsafe {
24748            b.launch(cfg)?;
24749        }
24750        Ok(())
24751    }
24752
24753    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
24754    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
24755    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
24756    #[allow(clippy::too_many_arguments)]
24757    pub fn ssm_conv1d_gdn(
24758        &self,
24759        qkv_tm: &CudaSlice<f32>,
24760        w: &CudaSlice<f32>,
24761        q_g: &mut CudaSlice<f32>,
24762        k_g: &mut CudaSlice<f32>,
24763        v_g: &mut CudaSlice<f32>,
24764        conv_dim: usize,
24765        t: usize,
24766        d_conv: usize,
24767        d_state: usize,
24768        num_v: usize,
24769        num_k: usize,
24770        key_dim: usize,
24771    ) -> Result<(), Box<dyn std::error::Error>> {
24772        let f = self.func("ssm_conv1d_gdn_f32");
24773        let cfg = LaunchConfig {
24774            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24775            block_dim: (256, 1, 1),
24776            shared_mem_bytes: 0,
24777        };
24778        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24779        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24780        let __s_b = self.gpu.stream();
24781        let mut b = __s_b.launch_builder(&f);
24782        b.arg(qkv_tm)
24783            .arg(w)
24784            .arg(q_g)
24785            .arg(k_g)
24786            .arg(v_g)
24787            .arg(&cd)
24788            .arg(&ti)
24789            .arg(&dc)
24790            .arg(&ds)
24791            .arg(&nv)
24792            .arg(&nk)
24793            .arg(&kd);
24794        unsafe {
24795            b.launch(cfg)?;
24796        }
24797        Ok(())
24798    }
24799
24800    pub fn ssm_conv1d(
24801        &self,
24802        x: &CudaSlice<f32>,
24803        w: &CudaSlice<f32>,
24804        y: &mut CudaSlice<f32>,
24805        conv_dim: usize,
24806        t: usize,
24807        d_conv: usize,
24808        silu: bool,
24809    ) -> Result<(), Box<dyn std::error::Error>> {
24810        let f = self.func("ssm_conv1d_silu_f32");
24811        let cfg = LaunchConfig {
24812            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24813            block_dim: (256, 1, 1),
24814            shared_mem_bytes: 0,
24815        };
24816        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24817        let __s_b = self.gpu.stream();
24818        let mut b = __s_b.launch_builder(&f);
24819        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24820        unsafe {
24821            b.launch(cfg)?;
24822        }
24823        Ok(())
24824    }
24825
24826    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
24827    /// o:[128,H,T]. Single sequence.
24828    pub fn gdn_scan_s128(
24829        &self,
24830        q: &CudaSlice<f32>,
24831        k: &CudaSlice<f32>,
24832        v: &CudaSlice<f32>,
24833        g: &CudaSlice<f32>,
24834        beta: &CudaSlice<f32>,
24835        state_in: &CudaSlice<f32>,
24836        state_out: &mut CudaSlice<f32>,
24837        o: &mut CudaSlice<f32>,
24838        n_head: usize,
24839        t: usize,
24840        scale: f32,
24841    ) -> Result<(), Box<dyn std::error::Error>> {
24842        let f = self.func("gdn_scan_s128");
24843        const S_V: u32 = 128;
24844        const WARP: u32 = 32;
24845        const COLS_PER_BLOCK: u32 = 4;
24846        let cfg = LaunchConfig {
24847            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
24848            block_dim: (WARP, COLS_PER_BLOCK, 1),
24849            shared_mem_bytes: 0,
24850        };
24851        let (h, ti) = (n_head as i32, t as i32);
24852        let __s_b = self.gpu.stream();
24853        let mut b = __s_b.launch_builder(&f);
24854        b.arg(q)
24855            .arg(k)
24856            .arg(v)
24857            .arg(g)
24858            .arg(beta)
24859            .arg(state_in)
24860            .arg(state_out)
24861            .arg(o)
24862            .arg(&h)
24863            .arg(&ti)
24864            .arg(&scale);
24865        unsafe {
24866            b.launch(cfg)?;
24867        }
24868        Ok(())
24869    }
24870
24871    // ==== B2' batched decode state ops (decode_batch.rs) ====
24872    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
24873    // Bodies are the single-seq kernels per sequence — bit-identical per row.
24874
24875    #[allow(clippy::too_many_arguments)]
24876    pub fn ssm_conv1d_fused_decode_b(
24877        &self,
24878        qkv_cols: &CudaSlice<f32>,
24879        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24880        w: &CudaSlice<f32>,
24881        conv_outs: &mut CudaSlice<f32>,
24882        conv_dim: usize,
24883        d_conv: usize,
24884        b_n: usize,
24885    ) -> Result<(), Box<dyn std::error::Error>> {
24886        let f = self.func("ssm_conv1d_fused_decode_b_f32");
24887        let cfg = LaunchConfig {
24888            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24889            block_dim: (256, 1, 1),
24890            shared_mem_bytes: 0,
24891        };
24892        let (cd, dc) = (conv_dim as i32, d_conv as i32);
24893        let __s_b = self.gpu.stream();
24894        let mut b = __s_b.launch_builder(&f);
24895        b.arg(qkv_cols)
24896            .arg(conv_state_ptrs)
24897            .arg(w)
24898            .arg(conv_outs)
24899            .arg(&cd)
24900            .arg(&dc);
24901        unsafe {
24902            b.launch(cfg)?;
24903        }
24904        Ok(())
24905    }
24906
24907    #[allow(clippy::too_many_arguments)]
24908    pub fn gdn_prep_decode_b(
24909        &self,
24910        conv_outs: &CudaSlice<f32>,
24911        beta_raws: &CudaSlice<f32>,
24912        alphas: &CudaSlice<f32>,
24913        dt_bias: &CudaSlice<f32>,
24914        a: &CudaSlice<f32>,
24915        q_l2: &mut CudaSlice<f32>,
24916        k_l2: &mut CudaSlice<f32>,
24917        v_g: &mut CudaSlice<f32>,
24918        beta: &mut CudaSlice<f32>,
24919        g_log: &mut CudaSlice<f32>,
24920        d_state: usize,
24921        num_v: usize,
24922        num_k: usize,
24923        key_dim: usize,
24924        eps: f32,
24925        conv_dim: usize,
24926        b_n: usize,
24927    ) -> Result<(), Box<dyn std::error::Error>> {
24928        let f = self.func("gdn_prep_decode_b_f32");
24929        let cfg = LaunchConfig {
24930            grid_dim: (num_v as u32, 1, b_n as u32),
24931            block_dim: (32, 4, 1),
24932            shared_mem_bytes: 0,
24933        };
24934        let (ds, nv, nk, kd, cd) = (
24935            d_state as i32,
24936            num_v as i32,
24937            num_k as i32,
24938            key_dim as i32,
24939            conv_dim as i32,
24940        );
24941        let __s_b = self.gpu.stream();
24942        let mut b = __s_b.launch_builder(&f);
24943        b.arg(conv_outs)
24944            .arg(beta_raws)
24945            .arg(alphas)
24946            .arg(dt_bias)
24947            .arg(a)
24948            .arg(q_l2)
24949            .arg(k_l2)
24950            .arg(v_g)
24951            .arg(beta)
24952            .arg(g_log)
24953            .arg(&ds)
24954            .arg(&nv)
24955            .arg(&nk)
24956            .arg(&kd)
24957            .arg(&eps)
24958            .arg(&cd);
24959        unsafe {
24960            b.launch(cfg)?;
24961        }
24962        Ok(())
24963    }
24964
24965    #[allow(clippy::too_many_arguments)]
24966    pub fn gdn_scan_s128_batched(
24967        &self,
24968        q: &CudaSlice<f32>,
24969        k: &CudaSlice<f32>,
24970        v: &CudaSlice<f32>,
24971        g: &CudaSlice<f32>,
24972        beta: &CudaSlice<f32>,
24973        state_in_ptrs: &cudarc::driver::CudaView<u64>,
24974        state_out_ptrs: &cudarc::driver::CudaView<u64>,
24975        o: &mut CudaSlice<f32>,
24976        n_head: usize,
24977        b_n: usize,
24978        scale: f32,
24979    ) -> Result<(), Box<dyn std::error::Error>> {
24980        let f = self.func("gdn_scan_s128_b");
24981        const S_V: u32 = 128;
24982        const WARP: u32 = 32;
24983        const COLS_PER_BLOCK: u32 = 4;
24984        let cfg = LaunchConfig {
24985            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24986            block_dim: (WARP, COLS_PER_BLOCK, 1),
24987            shared_mem_bytes: 0,
24988        };
24989        let h = n_head as i32;
24990        let __s_b = self.gpu.stream();
24991        let mut b = __s_b.launch_builder(&f);
24992        b.arg(q)
24993            .arg(k)
24994            .arg(v)
24995            .arg(g)
24996            .arg(beta)
24997            .arg(state_in_ptrs)
24998            .arg(state_out_ptrs)
24999            .arg(o)
25000            .arg(&h)
25001            .arg(&scale);
25002        unsafe {
25003            b.launch(cfg)?;
25004        }
25005        Ok(())
25006    }
25007
25008    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
25009    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
25010    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
25011    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
25012    /// numeric class; only the pointer arithmetic moved host-side.
25013    #[allow(clippy::too_many_arguments)]
25014    pub fn ssm_conv1d_fused_decode_b_view(
25015        &self,
25016        qkv_cols: &cudarc::driver::CudaView<f32>,
25017        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25018        w: &CudaSlice<f32>,
25019        conv_outs: &mut CudaSlice<f32>,
25020        conv_dim: usize,
25021        d_conv: usize,
25022        b_n: usize,
25023    ) -> Result<(), Box<dyn std::error::Error>> {
25024        let f = self.func("ssm_conv1d_fused_decode_b_f32");
25025        let cfg = LaunchConfig {
25026            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25027            block_dim: (256, 1, 1),
25028            shared_mem_bytes: 0,
25029        };
25030        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25031        let __s_b = self.gpu.stream();
25032        let mut b = __s_b.launch_builder(&f);
25033        b.arg(qkv_cols)
25034            .arg(conv_state_ptrs)
25035            .arg(w)
25036            .arg(conv_outs)
25037            .arg(&cd)
25038            .arg(&dc);
25039        unsafe {
25040            b.launch(cfg)?;
25041        }
25042        Ok(())
25043    }
25044
25045    #[allow(clippy::too_many_arguments)]
25046    pub fn gdn_prep_decode_b_view(
25047        &self,
25048        conv_outs: &CudaSlice<f32>,
25049        beta_raws: &cudarc::driver::CudaView<f32>,
25050        alphas: &cudarc::driver::CudaView<f32>,
25051        dt_bias: &CudaSlice<f32>,
25052        a: &CudaSlice<f32>,
25053        q_l2: &mut CudaSlice<f32>,
25054        k_l2: &mut CudaSlice<f32>,
25055        v_g: &mut CudaSlice<f32>,
25056        beta: &mut CudaSlice<f32>,
25057        g_log: &mut CudaSlice<f32>,
25058        d_state: usize,
25059        num_v: usize,
25060        num_k: usize,
25061        key_dim: usize,
25062        eps: f32,
25063        conv_dim: usize,
25064        b_n: usize,
25065    ) -> Result<(), Box<dyn std::error::Error>> {
25066        let f = self.func("gdn_prep_decode_b_f32");
25067        let cfg = LaunchConfig {
25068            grid_dim: (num_v as u32, 1, b_n as u32),
25069            block_dim: (32, 4, 1),
25070            shared_mem_bytes: 0,
25071        };
25072        let (ds, nv, nk, kd, cd) = (
25073            d_state as i32,
25074            num_v as i32,
25075            num_k as i32,
25076            key_dim as i32,
25077            conv_dim as i32,
25078        );
25079        let __s_b = self.gpu.stream();
25080        let mut b = __s_b.launch_builder(&f);
25081        b.arg(conv_outs)
25082            .arg(beta_raws)
25083            .arg(alphas)
25084            .arg(dt_bias)
25085            .arg(a)
25086            .arg(q_l2)
25087            .arg(k_l2)
25088            .arg(v_g)
25089            .arg(beta)
25090            .arg(g_log)
25091            .arg(&ds)
25092            .arg(&nv)
25093            .arg(&nk)
25094            .arg(&kd)
25095            .arg(&eps)
25096            .arg(&cd);
25097        unsafe {
25098            b.launch(cfg)?;
25099        }
25100        Ok(())
25101    }
25102
25103    #[allow(clippy::too_many_arguments)]
25104    pub fn gdn_scan_s128_batched_view(
25105        &self,
25106        q: &CudaSlice<f32>,
25107        k: &CudaSlice<f32>,
25108        v: &CudaSlice<f32>,
25109        g: &CudaSlice<f32>,
25110        beta: &CudaSlice<f32>,
25111        state_in_ptrs: &cudarc::driver::CudaView<u64>,
25112        state_out_ptrs: &cudarc::driver::CudaView<u64>,
25113        o: &mut cudarc::driver::CudaViewMut<f32>,
25114        n_head: usize,
25115        b_n: usize,
25116        scale: f32,
25117    ) -> Result<(), Box<dyn std::error::Error>> {
25118        let f = self.func("gdn_scan_s128_b");
25119        const S_V: u32 = 128;
25120        const WARP: u32 = 32;
25121        const COLS_PER_BLOCK: u32 = 4;
25122        let cfg = LaunchConfig {
25123            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25124            block_dim: (WARP, COLS_PER_BLOCK, 1),
25125            shared_mem_bytes: 0,
25126        };
25127        let h = n_head as i32;
25128        let __s_b = self.gpu.stream();
25129        let mut b = __s_b.launch_builder(&f);
25130        b.arg(q)
25131            .arg(k)
25132            .arg(v)
25133            .arg(g)
25134            .arg(beta)
25135            .arg(state_in_ptrs)
25136            .arg(state_out_ptrs)
25137            .arg(o)
25138            .arg(&h)
25139            .arg(&scale);
25140        unsafe {
25141            b.launch(cfg)?;
25142        }
25143        Ok(())
25144    }
25145
25146    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
25147    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
25148    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
25149    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
25150    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
25151    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
25152    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
25153    /// identity law); prime_cache/forward/forward_last are the only callers.
25154    pub fn gdn_chunked_enabled() -> bool {
25155        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25156        *E.get_or_init(|| {
25157            std::env::var("MEMRA_GDN_CHUNKED")
25158                .map(|v| v != "0")
25159                .unwrap_or(true)
25160        })
25161    }
25162
25163    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
25164    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
25165    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
25166    /// of 32 in [32, 128] (kernel row mappings require it).
25167    pub fn gdn_chunk_size() -> usize {
25168        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25169        *C.get_or_init(|| {
25170            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25171                .ok()
25172                .and_then(|v| v.parse().ok())
25173                .unwrap_or(32);
25174            c.clamp(32, 128) / 32 * 32
25175        })
25176    }
25177
25178    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
25179    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
25180    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
25181    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
25182    #[allow(clippy::too_many_arguments)]
25183    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
25184    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
25185    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25186    #[allow(clippy::too_many_arguments)]
25187    pub fn gdn_chunk_k123(
25188        &self,
25189        q: &CudaSlice<f32>,
25190        k: &CudaSlice<f32>,
25191        v: &CudaSlice<f32>,
25192        g: &CudaSlice<f32>,
25193        beta: &CudaSlice<f32>,
25194        wb16: Option<&mut CudaSlice<u8>>,
25195        n_head: usize,
25196        t: usize,
25197        c: usize,
25198        hk: usize,
25199        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25200    ) -> Result<
25201        (
25202            CudaSlice<f32>,
25203            CudaSlice<f32>,
25204            CudaSlice<f32>,
25205            CudaSlice<f32>,
25206        ),
25207        Box<dyn std::error::Error>,
25208    > {
25209        const D: usize = 128;
25210        let h = n_head;
25211        let nc = (t + c - 1) / c;
25212        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25213        let mut gcum = self.uninit(t * h)?;
25214        let mut a = self.uninit(nc * h * c * c)?;
25215        let mut p = self.uninit(nc * h * c * c)?;
25216        let mut u = self.uninit(nc * h * c * D)?;
25217        let mut w = self.uninit(nc * h * c * D)?;
25218        {
25219            // K1
25220            let f = self.func("gdn_chunk_cumgate_f32");
25221            let cfg = LaunchConfig {
25222                grid_dim: (nc as u32, h as u32, 1),
25223                block_dim: (32, 1, 1),
25224                shared_mem_bytes: 0,
25225            };
25226            let __s_b = self.gpu.stream();
25227            let mut b = __s_b.launch_builder(&f);
25228            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25229            unsafe {
25230                b.launch(cfg)?;
25231            }
25232        }
25233        if let Some((qb, kb, pb)) = k2w {
25234            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
25235            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
25236            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25237            let f = self.func("gdn_k2_wgmma");
25238            let cfg = LaunchConfig {
25239                grid_dim: (nc as u32, h as u32, 1),
25240                block_dim: (128, 1, 1),
25241                shared_mem_bytes: 0,
25242            };
25243            let hki = hk as i32;
25244            let __s_b = self.gpu.stream();
25245            let mut b = __s_b.launch_builder(&f);
25246            b.arg(qb)
25247                .arg(kb)
25248                .arg(&gcum)
25249                .arg(beta)
25250                .arg(&mut a)
25251                .arg(&mut *pb)
25252                .arg(&hi)
25253                .arg(&ti)
25254                .arg(&ci)
25255                .arg(&hki);
25256            unsafe {
25257                b.launch(cfg)?;
25258            }
25259        } else if c <= 64 && !portable_mma_gated() {
25260            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
25261            let f = self.func("gdn_chunk_attn_f32");
25262            f.set_attribute(
25263                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25264                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25265            )?;
25266            let jt = ((c + 31) / 32) as u32;
25267            let cfg = LaunchConfig {
25268                grid_dim: (nc as u32, h as u32, jt),
25269                block_dim: (256, 1, 1),
25270                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25271            };
25272            let hki = hk as i32;
25273            let __s_b = self.gpu.stream();
25274            let mut b = __s_b.launch_builder(&f);
25275            b.arg(q)
25276                .arg(k)
25277                .arg(&gcum)
25278                .arg(beta)
25279                .arg(&mut a)
25280                .arg(&mut p)
25281                .arg(&hi)
25282                .arg(&ti)
25283                .arg(&ci)
25284                .arg(&hki);
25285            unsafe {
25286                b.launch(cfg)?;
25287            }
25288        } else {
25289            // K2 generic (C = 128, or the portable target's low-smem fallback)
25290            assert!(
25291                hk == h,
25292                "generic K2 is broadcast-only (de-broadcast rides C==32)"
25293            );
25294            let f = self.func("gdn_chunk_attn_g_f32");
25295            let cfg = LaunchConfig {
25296                grid_dim: (nc as u32, h as u32, 1),
25297                block_dim: (32, 8, 1),
25298                shared_mem_bytes: 0,
25299            };
25300            let __s_b = self.gpu.stream();
25301            let mut b = __s_b.launch_builder(&f);
25302            b.arg(q)
25303                .arg(k)
25304                .arg(&gcum)
25305                .arg(beta)
25306                .arg(&mut a)
25307                .arg(&mut p)
25308                .arg(&hi)
25309                .arg(&ti)
25310                .arg(&ci);
25311            unsafe {
25312                b.launch(cfg)?;
25313            }
25314        }
25315        {
25316            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
25317            let cfg = LaunchConfig {
25318                grid_dim: (nc as u32, h as u32, 1),
25319                block_dim: (256, 1, 1),
25320                shared_mem_bytes: 0,
25321            };
25322            match c {
25323                32 | 64 => {
25324                    let f = self.func(if c == 32 {
25325                        "gdn_chunk_solve32_f32"
25326                    } else {
25327                        "gdn_chunk_solve64_f32"
25328                    });
25329                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
25330                    let wb: u64 = match wb16 {
25331                        Some(d) => self.addr_u8(d),
25332                        None => 0,
25333                    };
25334                    let hki = hk as i32;
25335                    let __s_b = self.gpu.stream();
25336                    let mut b = __s_b.launch_builder(&f);
25337                    b.arg(v)
25338                        .arg(k)
25339                        .arg(&a)
25340                        .arg(&gcum)
25341                        .arg(&mut u)
25342                        .arg(&mut w)
25343                        .arg(&wb)
25344                        .arg(&hi)
25345                        .arg(&ti)
25346                        .arg(&hki);
25347                    unsafe {
25348                        b.launch(cfg)?;
25349                    }
25350                }
25351                _ => {
25352                    assert!(hk == h, "generic K3 is broadcast-only");
25353                    let f = self.func("gdn_chunk_solve_f32");
25354                    let __s_b = self.gpu.stream();
25355                    let mut b = __s_b.launch_builder(&f);
25356                    b.arg(v)
25357                        .arg(k)
25358                        .arg(&a)
25359                        .arg(&gcum)
25360                        .arg(&mut u)
25361                        .arg(&mut w)
25362                        .arg(&hi)
25363                        .arg(&ti)
25364                        .arg(&ci);
25365                    unsafe {
25366                        b.launch(cfg)?;
25367                    }
25368                }
25369            }
25370        }
25371        Ok((gcum, p, u, w))
25372    }
25373
25374    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
25375    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
25376    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
25377    pub fn gdn_db_on() -> bool {
25378        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25379    }
25380
25381    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
25382    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
25383    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
25384    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
25385    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
25386    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
25387    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
25388    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
25389    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25390        !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    }
25398
25399    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
25400    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
25401    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
25402    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
25403    /// force would silently produce garbage. Required since the sm_120a mma default
25404    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
25405    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25406        cfg!(memra_hopper_mma)
25407            && self.gdn_mma_enabled(c)
25408            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25409    }
25410
25411    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
25412    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
25413    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
25414    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
25415    #[allow(clippy::too_many_arguments)]
25416    pub fn ssm_conv1d_gdn_state_pad(
25417        &self,
25418        qkv_tm: &cudarc::driver::CudaView<f32>,
25419        conv_state: &mut CudaSlice<f32>,
25420        w: &CudaSlice<f32>,
25421        q_g: &mut CudaSlice<f32>,
25422        k_g: &mut CudaSlice<f32>,
25423        v_g: &mut CudaSlice<f32>,
25424        conv_dim: usize,
25425        t: usize,
25426        d_conv: usize,
25427        d_state: usize,
25428        num_v: usize,
25429        num_k: usize,
25430        key_dim: usize,
25431        hk: usize,
25432        pad_len: Option<&CudaSlice<i32>>,
25433    ) -> Result<(), Box<dyn std::error::Error>> {
25434        assert!(
25435            t >= d_conv - 1,
25436            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25437        );
25438        {
25439            let f = self.func("ssm_conv1d_gdn_state_f32");
25440            let cfg = LaunchConfig {
25441                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25442                block_dim: (256, 1, 1),
25443                shared_mem_bytes: 0,
25444            };
25445            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25446            let (ds, nv, nk, kd, hki) = (
25447                d_state as i32,
25448                num_v as i32,
25449                num_k as i32,
25450                key_dim as i32,
25451                hk as i32,
25452            );
25453            let __s_b = self.gpu.stream();
25454            let mut b = __s_b.launch_builder(&f);
25455            b.arg(qkv_tm)
25456                .arg(&*conv_state)
25457                .arg(w)
25458                .arg(q_g)
25459                .arg(k_g)
25460                .arg(v_g)
25461                .arg(&cd)
25462                .arg(&ti)
25463                .arg(&dc)
25464                .arg(&ds)
25465                .arg(&nv)
25466                .arg(&nk)
25467                .arg(&kd)
25468                .arg(&hki);
25469            unsafe {
25470                b.launch(cfg)?;
25471            }
25472        }
25473        match pad_len {
25474            Some(len_d) => {
25475                let f = self.func("ssm_conv_ring_update_dev_f32");
25476                let n = conv_dim * (d_conv - 1);
25477                let cfg = LaunchConfig::for_num_elems(n as u32);
25478                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25479                let __s_b = self.gpu.stream();
25480                let mut b = __s_b.launch_builder(&f);
25481                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25482                unsafe {
25483                    b.launch(cfg)?;
25484                }
25485            }
25486            None => {
25487                let f = self.func("ssm_conv_ring_update_f32");
25488                let n = conv_dim * (d_conv - 1);
25489                let cfg = LaunchConfig::for_num_elems(n as u32);
25490                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25491                let __s_b = self.gpu.stream();
25492                let mut b = __s_b.launch_builder(&f);
25493                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25494                unsafe {
25495                    b.launch(cfg)?;
25496                }
25497            }
25498        }
25499        Ok(())
25500    }
25501
25502    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
25503    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
25504    /// K2/K3 can write them.
25505    pub fn gdn_chunk_alloc(
25506        &self,
25507        n_head: usize,
25508        t: usize,
25509        c: usize,
25510        hk: usize,
25511    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25512        const D: usize = 128;
25513        assert!(
25514            c == 32,
25515            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25516        );
25517        let h = n_head;
25518        let nc = (t + c - 1) / c;
25519        Ok(GdnChunkBufs {
25520            gcum: self.uninit(t * h)?,
25521            a: self.uninit(nc * h * c * c)?,
25522            p: self.uninit(nc * h * c * c)?,
25523            u: self.uninit(nc * h * c * D)?,
25524            w: self.uninit(nc * h * c * D)?,
25525            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25526            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25527            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25528            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25529            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25530            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25531            o: self.uninit(D * h * t)?,
25532            t,
25533            nc,
25534        })
25535    }
25536
25537    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
25538    pub fn f32_to_bf16_v(
25539        &self,
25540        x: &cudarc::driver::CudaView<f32>,
25541        dst: &mut CudaSlice<u8>,
25542        n: usize,
25543    ) -> Result<(), Box<dyn std::error::Error>> {
25544        let f = self.func("f32_to_bf16_bulk");
25545        let ni = n as i64;
25546        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25547        let __s_b = self.gpu.stream();
25548        let mut b = __s_b.launch_builder(&f);
25549        b.arg(x).arg(dst).arg(&ni);
25550        unsafe {
25551            b.launch(cfg)?;
25552        }
25553        Ok(())
25554    }
25555
25556    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
25557    pub fn f32_to_bf16_into(
25558        &self,
25559        x: &CudaSlice<f32>,
25560        dst: &mut CudaSlice<u8>,
25561        n: usize,
25562    ) -> Result<(), Box<dyn std::error::Error>> {
25563        let f = self.func("f32_to_bf16_bulk");
25564        let ni = n as i64;
25565        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25566        let __s_b = self.gpu.stream();
25567        let mut b = __s_b.launch_builder(&f);
25568        b.arg(x).arg(dst).arg(&ni);
25569        unsafe {
25570            b.launch(cfg)?;
25571        }
25572        Ok(())
25573    }
25574
25575    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
25576    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
25577    pub fn gdn_chunk_k123_vl8(
25578        &self,
25579        seqs: &[GdnSeqVl],
25580        n_head: usize,
25581        hk: usize,
25582        wq: Option<&GdnWVl8>,
25583    ) -> Result<(), Box<dyn std::error::Error>> {
25584        let b = seqs.len();
25585        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
25586        let mut packed = [GdnSeqVl::default(); 8];
25587        packed[..b].copy_from_slice(seqs);
25588        let v = GdnVl8(packed);
25589        let (hi, ci) = (n_head as i32, 32i32);
25590        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25591        {
25592            let f = self.func("gdn_chunk_cumgate_vl");
25593            let cfg = LaunchConfig {
25594                grid_dim: (max_nc, n_head as u32, b as u32),
25595                block_dim: (32, 1, 1),
25596                shared_mem_bytes: 0,
25597            };
25598            let __s_lb = self.gpu.stream();
25599            let mut lb = __s_lb.launch_builder(&f);
25600            lb.arg(&v).arg(&hi).arg(&ci);
25601            unsafe {
25602                lb.launch(cfg)?;
25603            }
25604        }
25605        let hki = hk as i32;
25606        if let Some(w) = wq {
25607            // K2-wgmma vl twin (writes A + pre-masked Pb16)
25608            let f = self.func("gdn_k2_wgmma_vl");
25609            let cfg = LaunchConfig {
25610                grid_dim: (max_nc, n_head as u32, b as u32),
25611                block_dim: (128, 1, 1),
25612                shared_mem_bytes: 0,
25613            };
25614            let __s_lb = self.gpu.stream();
25615            let mut lb = __s_lb.launch_builder(&f);
25616            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
25617            unsafe {
25618                lb.launch(cfg)?;
25619            }
25620        } else {
25621            let f = self.func("gdn_chunk_attn_vl");
25622            f.set_attribute(
25623                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25624                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25625            )?;
25626            let cfg = LaunchConfig {
25627                grid_dim: (max_nc, n_head as u32, b as u32),
25628                block_dim: (256, 1, 1),
25629                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25630            };
25631            let __s_lb = self.gpu.stream();
25632            let mut lb = __s_lb.launch_builder(&f);
25633            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25634            unsafe {
25635                lb.launch(cfg)?;
25636            }
25637        }
25638        {
25639            let f = self.func("gdn_chunk_solve32_vl");
25640            let cfg = LaunchConfig {
25641                grid_dim: (max_nc, n_head as u32, b as u32),
25642                block_dim: (256, 1, 1),
25643                shared_mem_bytes: 0,
25644            };
25645            let __s_lb = self.gpu.stream();
25646            let mut lb = __s_lb.launch_builder(&f);
25647            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25648            unsafe {
25649                lb.launch(cfg)?;
25650            }
25651        }
25652        Ok(())
25653    }
25654
25655    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
25656    /// fused gate-prep, 5 launches for every sequence (per-element math identical
25657    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
25658    #[allow(clippy::too_many_arguments)]
25659    pub fn gdn_prep_vl8(
25660        &self,
25661        seqs: &[GdnPrepVl],
25662        conv_w: &CudaSlice<f32>,
25663        dt_bias: &CudaSlice<f32>,
25664        a: &CudaSlice<f32>,
25665        conv_dim: usize,
25666        d_conv: usize,
25667        d_state: usize,
25668        num_v: usize,
25669        num_k: usize,
25670        key_dim: usize,
25671        hk: usize,
25672        eps: f32,
25673    ) -> Result<(), Box<dyn std::error::Error>> {
25674        let b = seqs.len();
25675        assert!(b >= 1 && b <= 8);
25676        let mut packed = [GdnPrepVl::default(); 8];
25677        packed[..b].copy_from_slice(seqs);
25678        let v = GdnPrepVl8(packed);
25679        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25680        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
25681        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
25682        assert!(
25683            conv_fuse || hk == num_v,
25684            "de-broadcast requires the fused conv"
25685        );
25686        if conv_fuse {
25687            let f = self.func("ssm_conv1d_gdn_state_vl");
25688            let cfg = LaunchConfig {
25689                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25690                block_dim: (256, 1, 1),
25691                shared_mem_bytes: 0,
25692            };
25693            let (dsi, nvi, nki, kdi, hki) = (
25694                d_state as i32,
25695                num_v as i32,
25696                num_k as i32,
25697                key_dim as i32,
25698                hk as i32,
25699            );
25700            let __s_lb = self.gpu.stream();
25701            let mut lb = __s_lb.launch_builder(&f);
25702            lb.arg(&v)
25703                .arg(conv_w)
25704                .arg(&cdi)
25705                .arg(&dci)
25706                .arg(&dsi)
25707                .arg(&nvi)
25708                .arg(&nki)
25709                .arg(&kdi)
25710                .arg(&hki);
25711            unsafe {
25712                lb.launch(cfg)?;
25713            }
25714        } else {
25715            let f = self.func("ssm_conv1d_tm_state_vl");
25716            let cfg = LaunchConfig {
25717                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25718                block_dim: (256, 1, 1),
25719                shared_mem_bytes: 0,
25720            };
25721            let __s_lb = self.gpu.stream();
25722            let mut lb = __s_lb.launch_builder(&f);
25723            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
25724            unsafe {
25725                lb.launch(cfg)?;
25726            }
25727        }
25728        {
25729            let f = self.func("ssm_conv_ring_update_vl");
25730            let n = (conv_dim * (d_conv - 1)) as u32;
25731            let cfg = LaunchConfig {
25732                grid_dim: (n.div_ceil(256), 1, b as u32),
25733                block_dim: (256, 1, 1),
25734                shared_mem_bytes: 0,
25735            };
25736            let __s_lb = self.gpu.stream();
25737            let mut lb = __s_lb.launch_builder(&f);
25738            lb.arg(&v).arg(&cdi).arg(&dci);
25739            unsafe {
25740                lb.launch(cfg)?;
25741            }
25742        }
25743        if !conv_fuse {
25744            let f = self.func("qkv_to_gdn_repack_vl");
25745            let n = max_t * (num_v * d_state) as u32;
25746            let cfg = LaunchConfig {
25747                grid_dim: (n.div_ceil(256), 1, b as u32),
25748                block_dim: (256, 1, 1),
25749                shared_mem_bytes: 0,
25750            };
25751            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25752            let __s_lb = self.gpu.stream();
25753            let mut lb = __s_lb.launch_builder(&f);
25754            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
25755            unsafe {
25756                lb.launch(cfg)?;
25757            }
25758        }
25759        if Self::l2_v2_on(d_state) {
25760            let f = self.func("gdn_l2_v2_vl");
25761            let cfg = LaunchConfig {
25762                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
25763                block_dim: (256, 1, 1),
25764                shared_mem_bytes: 0,
25765            };
25766            let (dsi, nvi) = (d_state as i32, hk as i32);
25767            let __s_lb = self.gpu.stream();
25768            let mut lb = __s_lb.launch_builder(&f);
25769            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25770            unsafe {
25771                lb.launch(cfg)?;
25772            }
25773        } else {
25774            let f = self.func("gdn_l2_vl");
25775            let cfg = LaunchConfig {
25776                grid_dim: (max_t * hk as u32, 2, b as u32),
25777                block_dim: (256, 1, 1),
25778                shared_mem_bytes: 0,
25779            };
25780            let (dsi, nvi) = (d_state as i32, hk as i32);
25781            let __s_lb = self.gpu.stream();
25782            let mut lb = __s_lb.launch_builder(&f);
25783            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25784            unsafe {
25785                lb.launch(cfg)?;
25786            }
25787        }
25788        {
25789            let f = self.func("gdn_gate_prep_vl");
25790            let n = max_t * num_v as u32;
25791            let cfg = LaunchConfig {
25792                grid_dim: (n.div_ceil(256), 1, b as u32),
25793                block_dim: (256, 1, 1),
25794                shared_mem_bytes: 0,
25795            };
25796            let nvi = num_v as i32;
25797            let __s_lb = self.gpu.stream();
25798            let mut lb = __s_lb.launch_builder(&f);
25799            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
25800            unsafe {
25801                lb.launch(cfg)?;
25802            }
25803        }
25804        Ok(())
25805    }
25806
25807    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
25808    pub fn gdn_mirror_vl8(
25809        &self,
25810        seqs: &[GdnSeqVl],
25811        n_head: usize,
25812        which: i32,
25813        hk: usize,
25814    ) -> Result<(), Box<dyn std::error::Error>> {
25815        let b = seqs.len();
25816        assert!(b >= 1 && b <= 8);
25817        let mut packed = [GdnSeqVl::default(); 8];
25818        packed[..b].copy_from_slice(seqs);
25819        let v = GdnVl8(packed);
25820        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
25821        let max_n = seqs
25822            .iter()
25823            .map(|s| {
25824                if which == 0 {
25825                    s.t as i64 * ept as i64
25826                } else {
25827                    s.nc as i64 * ept as i64 * 32
25828                }
25829            })
25830            .max()
25831            .unwrap();
25832        let f = self.func("gdn_mirror_vl");
25833        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25834        let cfg = LaunchConfig {
25835            grid_dim: (blocks, 1, b as u32),
25836            block_dim: (256, 1, 1),
25837            shared_mem_bytes: 0,
25838        };
25839        let __s_lb = self.gpu.stream();
25840        let mut lb = __s_lb.launch_builder(&f);
25841        lb.arg(&v).arg(&ept).arg(&which);
25842        unsafe {
25843            lb.launch(cfg)?;
25844        }
25845        Ok(())
25846    }
25847
25848    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
25849    pub fn gdn_tail_vl8(
25850        &self,
25851        seqs: &[GdnPrepVl],
25852        norm_w: &CudaSlice<f32>,
25853        d_state: usize,
25854        num_v: usize,
25855        eps: f32,
25856    ) -> Result<(), Box<dyn std::error::Error>> {
25857        let b = seqs.len();
25858        assert!(b >= 1 && b <= 8);
25859        let mut packed = [GdnPrepVl::default(); 8];
25860        packed[..b].copy_from_slice(seqs);
25861        let v = GdnPrepVl8(packed);
25862        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25863        let f = self.func("gated_rmsnorm_f16out_vl");
25864        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25865        let cfg = LaunchConfig {
25866            grid_dim: (max_t * num_v as u32, 1, b as u32),
25867            block_dim: (128, 1, 1),
25868            shared_mem_bytes: 0,
25869        };
25870        let (dsi, nvi) = (d_state as i32, num_v as i32);
25871        let __s_lb = self.gpu.stream();
25872        let mut lb = __s_lb.launch_builder(&f);
25873        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25874        unsafe {
25875            lb.launch(cfg)?;
25876        }
25877        Ok(())
25878    }
25879
25880    /// Raw device address helpers for the varlen by-value arg struct (single-stream
25881    /// launches; every buffer outlives the call — the f16 FFI discipline).
25882    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25883        use cudarc::driver::DevicePtr;
25884        let s = self.gpu.stream();
25885        let (p, _g) = x.device_ptr(&s);
25886        p as u64
25887    }
25888    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25889        use cudarc::driver::DevicePtrMut;
25890        let s = self.gpu.stream();
25891        let (p, _g) = x.device_ptr_mut(&s);
25892        p as u64
25893    }
25894    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25895        use cudarc::driver::DevicePtr;
25896        let s = self.gpu.stream();
25897        let (p, _g) = x.device_ptr(&s);
25898        p as u64
25899    }
25900    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25901        use cudarc::driver::DevicePtr;
25902        let s = self.gpu.stream();
25903        let (p, _g) = x.device_ptr(&s);
25904        p as u64
25905    }
25906
25907    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
25908    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
25909    /// launches, so this is strictly bit-gateable against them).
25910    pub fn gdn_chunk_vl8(
25911        &self,
25912        seqs: &[GdnSeqVl],
25913        n_head: usize,
25914        scale: f32,
25915        hk: usize,
25916        wq: Option<&GdnWVl8>,
25917    ) -> Result<(), Box<dyn std::error::Error>> {
25918        const NSPLIT: u32 = 4;
25919        let b = seqs.len();
25920        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25921        let mut packed = [GdnSeqVl::default(); 8];
25922        packed[..b].copy_from_slice(seqs);
25923        let v = GdnVl8(packed);
25924        let (hi, ci) = (n_head as i32, 32i32);
25925        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25926        let hki = hk as i32;
25927        if let Some(w) = wq {
25928            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
25929            let f = self.func("gdn_k45_wgmma_vl");
25930            let cfg = LaunchConfig {
25931                grid_dim: (n_head as u32, NSPLIT, b as u32),
25932                block_dim: (256, 1, 1),
25933                shared_mem_bytes: 0,
25934            };
25935            let __s_lb = self.gpu.stream();
25936            let mut lb = __s_lb.launch_builder(&f);
25937            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25938            unsafe {
25939                lb.launch(cfg)?;
25940            }
25941            let _ = max_nc;
25942            return Ok(());
25943        }
25944        {
25945            let f = self.func("gdn_chunk_state_mma_vl");
25946            let cfg = LaunchConfig {
25947                grid_dim: (n_head as u32, NSPLIT, b as u32),
25948                block_dim: (256, 1, 1),
25949                shared_mem_bytes: 0,
25950            };
25951            let __s_lb = self.gpu.stream();
25952            let mut lb = __s_lb.launch_builder(&f);
25953            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25954            unsafe {
25955                lb.launch(cfg)?;
25956            }
25957        }
25958        {
25959            let f = self.func("gdn_chunk_output_mma_vl");
25960            let cfg = LaunchConfig {
25961                grid_dim: (max_nc, n_head as u32, b as u32),
25962                block_dim: (256, 1, 1),
25963                shared_mem_bytes: 0,
25964            };
25965            let __s_lb = self.gpu.stream();
25966            let mut lb = __s_lb.launch_builder(&f);
25967            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25968            unsafe {
25969                lb.launch(cfg)?;
25970            }
25971        }
25972        Ok(())
25973    }
25974    pub fn gdn_scan_chunked(
25975        &self,
25976        q: &CudaSlice<f32>,
25977        k: &CudaSlice<f32>,
25978        v: &CudaSlice<f32>,
25979        g: &CudaSlice<f32>,
25980        beta: &CudaSlice<f32>,
25981        kb16_pre: Option<&CudaSlice<u8>>,
25982        qb16_pre: Option<&CudaSlice<u8>>,
25983        state_in: &CudaSlice<f32>,
25984        state_out: &mut CudaSlice<f32>,
25985        o: &mut CudaSlice<f32>,
25986        n_head: usize,
25987        t: usize,
25988        scale: f32,
25989        c: usize,
25990        hk: usize,
25991    ) -> Result<(), Box<dyn std::error::Error>> {
25992        const D: usize = 128;
25993        const NSPLIT: u32 = 4;
25994        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25995        let h = n_head;
25996        let nc = (t + c - 1) / c;
25997        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25998        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
25999        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
26000        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
26001        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
26002        let gdn_mma_pre = !portable_mma_gated()
26003            && c == 32
26004            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26005                Ok("1") => true,
26006                Ok("0") => false,
26007                _ => gdn_mma_default_on(),
26008            };
26009        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26010            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26011        } else {
26012            None
26013        };
26014        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
26015        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
26016        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
26017        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
26018        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26019            && gdn_mma_pre
26020            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26021        let nk = t * hk * D;
26022        let mut kb16_local: Option<CudaSlice<u8>> = None;
26023        if gdn_mma_pre && kb16_pre.is_none() {
26024            let mut kb = self.alloc_u8_uninit(nk * 2)?;
26025            let f = self.func("f32_to_bf16_bulk");
26026            let n2 = nk as i64;
26027            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26028            let __s_b = self.gpu.stream();
26029            let mut b = __s_b.launch_builder(&f);
26030            b.arg(k).arg(&mut kb).arg(&n2);
26031            unsafe {
26032                b.launch(cfg2)?;
26033            }
26034            kb16_local = Some(kb);
26035        }
26036        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26037        if let Some(kb) = kb16_pre {
26038            assert!(kb.len() >= nk * 2, "kb16_pre too small");
26039        }
26040        let mut qb16: Option<CudaSlice<u8>> = None;
26041        let mut pb16: Option<CudaSlice<u8>> = None;
26042        if gdn_wgmma_pre {
26043            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
26044            // the standalone bulk cvt only serves callers without the prep mirror.
26045            if qb16_pre.is_none() {
26046                let mut qb = self.alloc_u8_uninit(nk * 2)?;
26047                let f = self.func("f32_to_bf16_bulk");
26048                let n2 = nk as i64;
26049                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26050                let __s_b = self.gpu.stream();
26051                let mut b = __s_b.launch_builder(&f);
26052                b.arg(q).arg(&mut qb).arg(&n2);
26053                unsafe {
26054                    b.launch(cfg2)?;
26055                }
26056                qb16 = Some(qb);
26057            } else if let Some(qb) = qb16_pre {
26058                assert!(qb.len() >= nk * 2, "qb16_pre too small");
26059            }
26060            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26061        }
26062        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26063        let k2w = if gdn_wgmma_pre {
26064            Some((
26065                *qb16_ref0.as_ref().unwrap(),
26066                *kb16_ref0.as_ref().unwrap(),
26067                pb16.as_mut().unwrap(),
26068            ))
26069        } else {
26070            None
26071        };
26072        let (gcum, p, u, w) =
26073            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26074        let _ = &w;
26075        let mut y = self.uninit(nc * h * c * D)?;
26076        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
26077        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
26078        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
26079        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
26080        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
26081        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
26082        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
26083        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
26084        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
26085        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
26086        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
26087        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
26088        // sites must agree or the pre-work arms while the scan takes the scalar route.
26089        let gdn_mma = !portable_mma_gated()
26090            && c == 32
26091            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26092                Ok("1") => true,
26093                Ok("0") => false,
26094                _ => gdn_mma_default_on(),
26095            };
26096        if gdn_mma {
26097            let wb16 = wb16_pre
26098                .take()
26099                .expect("mma path pre-allocates wb16 (K3 store fold)");
26100            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26101            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
26102            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
26103            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
26104            // pass runs inside the persistent-M kernel; Y and Ssnap are never
26105            // materialized. New numeric class (gk folds into k^T instead of ys) —
26106            // explicit opt-in until the state-carry battery promotes it. Env read per
26107            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
26108            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
26109            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
26110            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
26111            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
26112            if gdn_wgmma_pre {
26113                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
26114                let qb16 = qb16_ref0.unwrap();
26115                let pb16 = pb16.as_ref().unwrap();
26116                {
26117                    let f = self.func("gdn_k45_wgmma");
26118                    let cfg = LaunchConfig {
26119                        grid_dim: (h as u32, 4, 1),
26120                        block_dim: (256, 1, 1),
26121                        shared_mem_bytes: 0,
26122                    };
26123                    let hki = hk as i32;
26124                    let __s_b = self.gpu.stream();
26125                    let mut b = __s_b.launch_builder(&f);
26126                    b.arg(kb16_ref)
26127                        .arg(&gcum)
26128                        .arg(beta)
26129                        .arg(&u)
26130                        .arg(&wb16)
26131                        .arg(qb16)
26132                        .arg(pb16)
26133                        .arg(o)
26134                        .arg(&scale)
26135                        .arg(state_in)
26136                        .arg(&mut *state_out)
26137                        .arg(&hi)
26138                        .arg(&ti)
26139                        .arg(&ci)
26140                        .arg(&hki);
26141                    unsafe {
26142                        b.launch(cfg)?;
26143                    }
26144                }
26145                return Ok(());
26146            }
26147            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
26148            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
26149            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
26150            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26151            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26152            {
26153                let f = self.func("gdn_chunk_state_mma");
26154                let cfg = LaunchConfig {
26155                    grid_dim: (h as u32, NSPLIT, 1),
26156                    block_dim: (256, 1, 1),
26157                    shared_mem_bytes: 0,
26158                };
26159                let hki = hk as i32;
26160                let __s_b = self.gpu.stream();
26161                let mut b = __s_b.launch_builder(&f);
26162                b.arg(kb16_ref)
26163                    .arg(&gcum)
26164                    .arg(beta)
26165                    .arg(&u)
26166                    .arg(&wb16)
26167                    .arg(&mut y16)
26168                    .arg(&mut ssnap16)
26169                    .arg(state_in)
26170                    .arg(&mut *state_out)
26171                    .arg(&hi)
26172                    .arg(&ti)
26173                    .arg(&ci)
26174                    .arg(&hki);
26175                unsafe {
26176                    b.launch(cfg)?;
26177                }
26178            }
26179            {
26180                // K5-mma (bf16 St/Y consumers)
26181                let f = self.func("gdn_chunk_output_mma");
26182                let jt = ((c + 31) / 32) as u32;
26183                let cfg = LaunchConfig {
26184                    grid_dim: (nc as u32, h as u32, jt),
26185                    block_dim: (256, 1, 1),
26186                    shared_mem_bytes: 0,
26187                };
26188                let hki = hk as i32;
26189                let __s_b = self.gpu.stream();
26190                let mut b = __s_b.launch_builder(&f);
26191                b.arg(q)
26192                    .arg(&gcum)
26193                    .arg(&p)
26194                    .arg(&y16)
26195                    .arg(&ssnap16)
26196                    .arg(o)
26197                    .arg(&hi)
26198                    .arg(&ti)
26199                    .arg(&ci)
26200                    .arg(&scale)
26201                    .arg(&hki);
26202                unsafe {
26203                    b.launch(cfg)?;
26204                }
26205            }
26206            return Ok(());
26207        }
26208        {
26209            // K4 (sequential over chunks inside; blocks col-partition the state)
26210            let f = self.func("gdn_chunk_state_f32");
26211            let cfg = LaunchConfig {
26212                grid_dim: (h as u32, NSPLIT, 1),
26213                block_dim: (256, 1, 1),
26214                shared_mem_bytes: 0,
26215            };
26216            let __s_b = self.gpu.stream();
26217            let mut b = __s_b.launch_builder(&f);
26218            b.arg(k)
26219                .arg(&gcum)
26220                .arg(beta)
26221                .arg(&u)
26222                .arg(&w)
26223                .arg(&mut y)
26224                .arg(&mut ssnap)
26225                .arg(state_in)
26226                .arg(&mut *state_out)
26227                .arg(&hi)
26228                .arg(&ti)
26229                .arg(&ci);
26230            unsafe {
26231                b.launch(cfg)?;
26232            }
26233        }
26234        {
26235            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
26236            let f = self.func("gdn_chunk_output_f32");
26237            let jt = ((c + 31) / 32) as u32;
26238            let cfg = LaunchConfig {
26239                grid_dim: (nc as u32, h as u32, jt),
26240                block_dim: (256, 1, 1),
26241                shared_mem_bytes: 0,
26242            };
26243            let __s_b = self.gpu.stream();
26244            let mut b = __s_b.launch_builder(&f);
26245            b.arg(q)
26246                .arg(&gcum)
26247                .arg(&p)
26248                .arg(&y)
26249                .arg(&ssnap)
26250                .arg(o)
26251                .arg(&hi)
26252                .arg(&ti)
26253                .arg(&ci)
26254                .arg(&scale);
26255            unsafe {
26256                b.launch(cfg)?;
26257            }
26258        }
26259        Ok(())
26260    }
26261
26262    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
26263    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
26264    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
26265    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
26266    ///
26267    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
26268    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
26269    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
26270    #[allow(clippy::too_many_arguments)]
26271    #[allow(clippy::too_many_arguments)]
26272    pub fn gdn_scan_prefill(
26273        &self,
26274        q: &CudaSlice<f32>,
26275        k: &CudaSlice<f32>,
26276        v: &CudaSlice<f32>,
26277        g: &CudaSlice<f32>,
26278        beta: &CudaSlice<f32>,
26279        kb16_pre: Option<&CudaSlice<u8>>,
26280        qb16_pre: Option<&CudaSlice<u8>>,
26281        state_in: &CudaSlice<f32>,
26282        state_out: &mut CudaSlice<f32>,
26283        o: &mut CudaSlice<f32>,
26284        n_head: usize,
26285        t: usize,
26286        scale: f32,
26287        hk: usize,
26288    ) -> Result<(), Box<dyn std::error::Error>> {
26289        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26290            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26291            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26292        }
26293        if Self::gdn_chunked_enabled() && t >= 16 {
26294            self.gdn_scan_chunked(
26295                q,
26296                k,
26297                v,
26298                g,
26299                beta,
26300                kb16_pre,
26301                qb16_pre,
26302                state_in,
26303                state_out,
26304                o,
26305                n_head,
26306                t,
26307                scale,
26308                Self::gdn_chunk_size(),
26309                hk,
26310            )
26311        } else {
26312            assert!(
26313                hk == n_head,
26314                "s128 scan is broadcast-only (prep guarantees by predicate)"
26315            );
26316            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26317        }
26318    }
26319
26320    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
26321    #[allow(clippy::too_many_arguments)]
26322    fn gdn_scan_diff(
26323        &self,
26324        q: &CudaSlice<f32>,
26325        k: &CudaSlice<f32>,
26326        v: &CudaSlice<f32>,
26327        g: &CudaSlice<f32>,
26328        beta: &CudaSlice<f32>,
26329        state_in: &CudaSlice<f32>,
26330        state_out: &mut CudaSlice<f32>,
26331        o: &mut CudaSlice<f32>,
26332        n_head: usize,
26333        t: usize,
26334        scale: f32,
26335    ) -> Result<(), Box<dyn std::error::Error>> {
26336        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26337        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26338        let mut o_c = self.uninit(o.len())?;
26339        let mut st_c = self.uninit(state_out.len())?;
26340        self.gdn_scan_chunked(
26341            q,
26342            k,
26343            v,
26344            g,
26345            beta,
26346            None,
26347            None,
26348            state_in,
26349            &mut st_c,
26350            &mut o_c,
26351            n_head,
26352            t,
26353            scale,
26354            Self::gdn_chunk_size(),
26355            n_head,
26356        )?;
26357        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26358        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26359        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26360        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26361            let mut max_abs = 0f32;
26362            let mut max_rel = 0f32;
26363            let mut sum_rel = 0f64;
26364            for (x, y) in a.iter().zip(b) {
26365                let ad = (x - y).abs();
26366                let rel = ad / x.abs().max(y.abs()).max(1e-3);
26367                if ad > max_abs {
26368                    max_abs = ad;
26369                }
26370                if rel > max_rel {
26371                    max_rel = rel;
26372                }
26373                sum_rel += rel as f64;
26374            }
26375            (max_abs, max_rel, sum_rel / a.len() as f64)
26376        };
26377        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26378        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26379        println!(
26380            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26381                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26382            Self::gdn_chunk_size()
26383        );
26384        Ok(())
26385    }
26386
26387    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
26388    pub fn gdn_glog(
26389        &self,
26390        alpha: &CudaSlice<f32>,
26391        dt_bias: &CudaSlice<f32>,
26392        a: &CudaSlice<f32>,
26393        g_log: &mut CudaSlice<f32>,
26394        n_head: usize,
26395        t: usize,
26396    ) -> Result<(), Box<dyn std::error::Error>> {
26397        let f = self.func("gdn_glog_f32");
26398        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26399        let (h, ti) = (n_head as i32, t as i32);
26400        let __s_b = self.gpu.stream();
26401        let mut b = __s_b.launch_builder(&f);
26402        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26403        unsafe {
26404            b.launch(cfg)?;
26405        }
26406        Ok(())
26407    }
26408
26409    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
26410    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
26411    pub fn sigmoid_v(
26412        &self,
26413        x: &cudarc::driver::CudaView<f32>,
26414        y: &mut CudaSlice<f32>,
26415        n: usize,
26416    ) -> Result<(), Box<dyn std::error::Error>> {
26417        let f = self.func("sigmoid_f32");
26418        let cfg = LaunchConfig::for_num_elems(n as u32);
26419        let ni = n as i32;
26420        let __s_b = self.gpu.stream();
26421        let mut b = __s_b.launch_builder(&f);
26422        b.arg(x).arg(y).arg(&ni);
26423        unsafe {
26424            b.launch(cfg)?;
26425        }
26426        Ok(())
26427    }
26428
26429    pub fn gdn_glog_v(
26430        &self,
26431        alpha: &cudarc::driver::CudaView<f32>,
26432        dt_bias: &CudaSlice<f32>,
26433        a: &CudaSlice<f32>,
26434        g_log: &mut CudaSlice<f32>,
26435        n_head: usize,
26436        t: usize,
26437    ) -> Result<(), Box<dyn std::error::Error>> {
26438        let f = self.func("gdn_glog_f32");
26439        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26440        let (h, ti) = (n_head as i32, t as i32);
26441        let __s_b = self.gpu.stream();
26442        let mut b = __s_b.launch_builder(&f);
26443        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26444        unsafe {
26445            b.launch(cfg)?;
26446        }
26447        Ok(())
26448    }
26449
26450    pub fn sigmoid(
26451        &self,
26452        x: &CudaSlice<f32>,
26453        y: &mut CudaSlice<f32>,
26454        n: usize,
26455    ) -> Result<(), Box<dyn std::error::Error>> {
26456        let f = self.func("sigmoid_f32");
26457        let cfg = LaunchConfig::for_num_elems(n as u32);
26458        let ni = n as i32;
26459        let __s_b = self.gpu.stream();
26460        let mut b = __s_b.launch_builder(&f);
26461        b.arg(x).arg(y).arg(&ni);
26462        unsafe {
26463            b.launch(cfg)?;
26464        }
26465        Ok(())
26466    }
26467
26468    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
26469    /// (replaces sigmoid + mul + convert). Bit-identical class.
26470    pub fn sig_mul_f16out(
26471        &self,
26472        a: &CudaSlice<f32>,
26473        g: &CudaSlice<f32>,
26474        dst: &mut CudaSlice<f32>,
26475        dst16: &mut CudaSlice<u8>,
26476        n: usize,
26477    ) -> Result<(), Box<dyn std::error::Error>> {
26478        let f = self.func("sig_mul_f16out_f32");
26479        let cfg = LaunchConfig::for_num_elems(n as u32);
26480        let ni = n as i32;
26481        let __s_b = self.gpu.stream();
26482        let mut b = __s_b.launch_builder(&f);
26483        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26484        unsafe {
26485            b.launch(cfg)?;
26486        }
26487        Ok(())
26488    }
26489
26490    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
26491    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
26492    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
26493    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
26494    ///
26495    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
26496    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
26497    /// applies the wrong number of distinct gate values.
26498    #[allow(clippy::too_many_arguments)]
26499    pub fn attn_head_gate(
26500        &self,
26501        a: &CudaSlice<f32>,
26502        g: &CudaSlice<f32>,
26503        dst: &mut CudaSlice<f32>,
26504        dst16: Option<&mut CudaSlice<u8>>,
26505        head_dim: usize,
26506        n_head: usize,
26507        t: usize,
26508    ) -> Result<(), Box<dyn std::error::Error>> {
26509        let f = self.func("attn_head_gate_f32");
26510        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26511        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26512        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
26513        let d16: u64 = match dst16 {
26514            Some(d) => self.addr_u8(d),
26515            None => 0,
26516        };
26517        let __s_b = self.gpu.stream();
26518        let mut b = __s_b.launch_builder(&f);
26519        b.arg(a)
26520            .arg(g)
26521            .arg(dst)
26522            .arg(&d16)
26523            .arg(&hd)
26524            .arg(&nh)
26525            .arg(&ti);
26526        unsafe {
26527            b.launch(cfg)?;
26528        }
26529        Ok(())
26530    }
26531
26532    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
26533    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
26534    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
26535    ///
26536    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
26537    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
26538    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
26539    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
26540    #[allow(clippy::too_many_arguments)]
26541    pub fn swiglu_clamped_mul_scaled(
26542        &self,
26543        gate: &CudaSlice<f32>,
26544        up: &CudaSlice<f32>,
26545        gs: f32,
26546        us: f32,
26547        limit: f32,
26548        dst: &mut CudaSlice<f32>,
26549        n: usize,
26550    ) -> Result<(), Box<dyn std::error::Error>> {
26551        debug_assert!(
26552            limit > 1e-6,
26553            "swiglu_clamped needs a live limit; use silu_mul_scaled"
26554        );
26555        let f = self.func("swiglu_clamped_mul_scaled_f32");
26556        let cfg = LaunchConfig::for_num_elems(n as u32);
26557        let ni = n as i32;
26558        let __s_b = self.gpu.stream();
26559        let mut b = __s_b.launch_builder(&f);
26560        b.arg(gate)
26561            .arg(up)
26562            .arg(&gs)
26563            .arg(&us)
26564            .arg(&limit)
26565            .arg(dst)
26566            .arg(&ni);
26567        unsafe {
26568            b.launch(cfg)?;
26569        }
26570        Ok(())
26571    }
26572
26573    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
26574    pub fn gated_rmsnorm(
26575        &self,
26576        o: &CudaSlice<f32>,
26577        w: &CudaSlice<f32>,
26578        z: &CudaSlice<f32>,
26579        dst: &mut CudaSlice<f32>,
26580        ncols: usize,
26581        nrows: usize,
26582        eps: f32,
26583    ) -> Result<(), Box<dyn std::error::Error>> {
26584        let f = self.func("gated_rmsnorm_f32");
26585        let cfg = LaunchConfig {
26586            grid_dim: (nrows as u32, 1, 1),
26587            block_dim: (128, 1, 1),
26588            shared_mem_bytes: 0,
26589        };
26590        let (nc, e) = (ncols as i32, eps);
26591        let __s_b = self.gpu.stream();
26592        let mut b = __s_b.launch_builder(&f);
26593        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26594        unsafe {
26595            b.launch(cfg)?;
26596        }
26597        Ok(())
26598    }
26599
26600    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
26601    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
26602    pub fn gated_rmsnorm_f16out(
26603        &self,
26604        o: &CudaSlice<f32>,
26605        w: &CudaSlice<f32>,
26606        z: &CudaSlice<f32>,
26607        dst: &mut CudaSlice<f32>,
26608        dst16: &mut CudaSlice<u8>,
26609        ncols: usize,
26610        nrows: usize,
26611        eps: f32,
26612    ) -> Result<(), Box<dyn std::error::Error>> {
26613        let f = self.func("gated_rmsnorm_f16out_f32");
26614        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
26615        let cfg = LaunchConfig {
26616            grid_dim: (nrows as u32, 1, 1),
26617            block_dim: (128, 1, 1),
26618            shared_mem_bytes: 0,
26619        };
26620        let (nc, e) = (ncols as i32, eps);
26621        let __s_b = self.gpu.stream();
26622        let mut b = __s_b.launch_builder(&f);
26623        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26624        unsafe {
26625            b.launch(cfg)?;
26626        }
26627        Ok(())
26628    }
26629
26630    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
26631    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
26632    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
26633    #[allow(clippy::too_many_arguments)]
26634    pub fn add_rms_norm_zq8(
26635        &self,
26636        a: &CudaSlice<f32>,
26637        b_in: &CudaSlice<f32>,
26638        w: &CudaSlice<f32>,
26639        res: &mut CudaSlice<f32>,
26640        z: &mut CudaSlice<f32>,
26641        ncols: usize,
26642        nrows: usize,
26643        eps: f32,
26644    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26645        assert!(ncols % 32 == 0);
26646        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
26647        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26648        let f = self.func("add_rms_norm_zq8");
26649        let cfg = LaunchConfig {
26650            grid_dim: (nrows as u32, 1, 1),
26651            block_dim: (1024, 1, 1),
26652            shared_mem_bytes: 0,
26653        };
26654        let (nc, ep) = (ncols as i32, eps);
26655        let __s_b = self.gpu.stream();
26656        let mut b = __s_b.launch_builder(&f);
26657        b.arg(a)
26658            .arg(b_in)
26659            .arg(w)
26660            .arg(res)
26661            .arg(z)
26662            .arg(&mut q)
26663            .arg(&mut d)
26664            .arg(&nc)
26665            .arg(&ep);
26666        unsafe {
26667            b.launch(cfg)?;
26668        }
26669        Ok((q, d))
26670    }
26671
26672    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
26673    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
26674    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
26675    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
26676    pub fn gated_rmsnorm_zv(
26677        &self,
26678        o: &CudaSlice<f32>,
26679        w: &CudaSlice<f32>,
26680        z: &cudarc::driver::CudaView<f32>,
26681        dst: &mut CudaSlice<f32>,
26682        ncols: usize,
26683        nrows: usize,
26684        eps: f32,
26685    ) -> Result<(), Box<dyn std::error::Error>> {
26686        let f = self.func("gated_rmsnorm_f32");
26687        let cfg = LaunchConfig {
26688            grid_dim: (nrows as u32, 1, 1),
26689            block_dim: (128, 1, 1),
26690            shared_mem_bytes: 0,
26691        };
26692        let (nc, e) = (ncols as i32, eps);
26693        let __s_b = self.gpu.stream();
26694        let mut b = __s_b.launch_builder(&f);
26695        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26696        unsafe {
26697            b.launch(cfg)?;
26698        }
26699        Ok(())
26700    }
26701
26702    pub fn gated_rmsnorm_f16out_zv(
26703        &self,
26704        o: &CudaSlice<f32>,
26705        w: &CudaSlice<f32>,
26706        z: &cudarc::driver::CudaView<f32>,
26707        dst: &mut CudaSlice<f32>,
26708        dst16: &mut CudaSlice<u8>,
26709        ncols: usize,
26710        nrows: usize,
26711        eps: f32,
26712    ) -> Result<(), Box<dyn std::error::Error>> {
26713        let f = self.func("gated_rmsnorm_f16out_f32");
26714        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
26715        let cfg = LaunchConfig {
26716            grid_dim: (nrows as u32, 1, 1),
26717            block_dim: (128, 1, 1),
26718            shared_mem_bytes: 0,
26719        };
26720        let (nc, e) = (ncols as i32, eps);
26721        let __s_b = self.gpu.stream();
26722        let mut b = __s_b.launch_builder(&f);
26723        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26724        unsafe {
26725            b.launch(cfg)?;
26726        }
26727        Ok(())
26728    }
26729
26730    pub fn gated_rmsnorm_q8_1(
26731        &self,
26732        o: &CudaSlice<f32>,
26733        w: &CudaSlice<f32>,
26734        z: &CudaSlice<f32>,
26735        ncols: usize,
26736        nrows: usize,
26737        eps: f32,
26738    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26739        assert!(ncols % 32 == 0);
26740        let f = self.func("gated_rmsnorm_q8_1");
26741        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
26742        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26743        let cfg = LaunchConfig {
26744            grid_dim: (nrows as u32, 1, 1),
26745            block_dim: (128, 1, 1),
26746            shared_mem_bytes: 0,
26747        };
26748        let (nc, ep) = (ncols as i32, eps);
26749        let __s_b = self.gpu.stream();
26750        let mut b = __s_b.launch_builder(&f);
26751        b.arg(o)
26752            .arg(w)
26753            .arg(z)
26754            .arg(&mut out_q)
26755            .arg(&mut out_d)
26756            .arg(&nc)
26757            .arg(&ep);
26758        unsafe {
26759            b.launch(cfg)?;
26760        }
26761        Ok((out_q, out_d))
26762    }
26763
26764    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
26765    pub fn transpose(
26766        &self,
26767        inp: &CudaSlice<f32>,
26768        rows: usize,
26769        cols: usize,
26770    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26771        let f = self.func("transpose_f32");
26772        let mut out = self.zeros(rows * cols)?;
26773        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
26774        let (r, c) = (rows as i32, cols as i32);
26775        let __s_b = self.gpu.stream();
26776        let mut b = __s_b.launch_builder(&f);
26777        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
26778        unsafe {
26779            b.launch(cfg)?;
26780        }
26781        Ok(out)
26782    }
26783
26784    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
26785    pub fn repeat_heads(
26786        &self,
26787        inp: &CudaSlice<f32>,
26788        out: &mut CudaSlice<f32>,
26789        head_dim: usize,
26790        n_in: usize,
26791        n_out: usize,
26792        t: usize,
26793    ) -> Result<(), Box<dyn std::error::Error>> {
26794        let f = self.func("repeat_heads_f32");
26795        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
26796        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
26797        let __s_b = self.gpu.stream();
26798        let mut b = __s_b.launch_builder(&f);
26799        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
26800        unsafe {
26801            b.launch(cfg)?;
26802        }
26803        Ok(())
26804    }
26805
26806    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
26807    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
26808    ///
26809    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
26810    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
26811    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
26812    pub fn q_gate_split(
26813        &self,
26814        qf: &CudaSlice<f32>,
26815        q_out: &mut CudaSlice<f32>,
26816        gate_out: &mut CudaSlice<f32>,
26817        head_dim: usize,
26818        n_head: usize,
26819        t: usize,
26820    ) -> Result<(), Box<dyn std::error::Error>> {
26821        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
26822        let out_need = head_dim * n_head * t;
26823        if q_out.len() < out_need || gate_out.len() < out_need {
26824            return Err(format!(
26825                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
26826                q_out.len(),
26827                gate_out.len()
26828            )
26829            .into());
26830        }
26831        let f = self.func("q_gate_split_f32");
26832        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26833        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26834        let __s_b = self.gpu.stream();
26835        let mut b = __s_b.launch_builder(&f);
26836        b.arg(qf)
26837            .arg(q_out)
26838            .arg(gate_out)
26839            .arg(&hd)
26840            .arg(&nh)
26841            .arg(&ti);
26842        unsafe {
26843            b.launch(cfg)?;
26844        }
26845        Ok(())
26846    }
26847
26848    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
26849    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
26850    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
26851    pub fn qkv_to_gdn_repack(
26852        &self,
26853        conv_out: &CudaSlice<f32>,
26854        q_g: &mut CudaSlice<f32>,
26855        k_g: &mut CudaSlice<f32>,
26856        v_g: &mut CudaSlice<f32>,
26857        d_state: usize,
26858        num_v: usize,
26859        num_k: usize,
26860        key_dim: usize,
26861        t: usize,
26862    ) -> Result<(), Box<dyn std::error::Error>> {
26863        let f = self.func("qkv_to_gdn_repack_f32");
26864        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26865        let (ds, nv, nk, kd, ti) = (
26866            d_state as i32,
26867            num_v as i32,
26868            num_k as i32,
26869            key_dim as i32,
26870            t as i32,
26871        );
26872        let __s_b = self.gpu.stream();
26873        let mut b = __s_b.launch_builder(&f);
26874        b.arg(conv_out)
26875            .arg(q_g)
26876            .arg(k_g)
26877            .arg(v_g)
26878            .arg(&ds)
26879            .arg(&nv)
26880            .arg(&nk)
26881            .arg(&kd)
26882            .arg(&ti);
26883        unsafe {
26884            b.launch(cfg)?;
26885        }
26886        Ok(())
26887    }
26888
26889    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
26890    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
26891    pub fn conv_left_pad(
26892        &self,
26893        src: &CudaSlice<f32>,
26894        dst: &mut CudaSlice<f32>,
26895        conv_dim: usize,
26896        t: usize,
26897        pad: usize,
26898    ) -> Result<(), Box<dyn std::error::Error>> {
26899        let f = self.func("conv_left_pad_f32");
26900        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26901        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26902        let __s_b = self.gpu.stream();
26903        let mut b = __s_b.launch_builder(&f);
26904        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26905        unsafe {
26906            b.launch(cfg)?;
26907        }
26908        Ok(())
26909    }
26910
26911    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
26912    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
26913    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
26914    pub fn conv_assemble_and_roll(
26915        &self,
26916        qkv_col: &CudaSlice<f32>,
26917        conv_state: &mut CudaSlice<f32>,
26918        conv_in: &mut CudaSlice<f32>,
26919        conv_dim: usize,
26920        pad: usize,
26921    ) -> Result<(), Box<dyn std::error::Error>> {
26922        let f = self.func("conv_assemble_and_roll_f32");
26923        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26924        let (cd, p) = (conv_dim as i32, pad as i32);
26925        let __s_b = self.gpu.stream();
26926        let mut b = __s_b.launch_builder(&f);
26927        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26928        unsafe {
26929            b.launch(cfg)?;
26930        }
26931        Ok(())
26932    }
26933
26934    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
26935    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
26936    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
26937    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
26938    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
26939    pub fn ssm_conv1d_fused_decode(
26940        &self,
26941        qkv_col: &CudaSlice<f32>,
26942        conv_state: &mut CudaSlice<f32>,
26943        w: &CudaSlice<f32>,
26944        conv_out: &mut CudaSlice<f32>,
26945        conv_dim: usize,
26946        d_conv: usize,
26947    ) -> Result<(), Box<dyn std::error::Error>> {
26948        let f = self.func("ssm_conv1d_fused_decode_f32");
26949        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26950        let (cd, dc) = (conv_dim as i32, d_conv as i32);
26951        let __s_b = self.gpu.stream();
26952        let mut b = __s_b.launch_builder(&f);
26953        b.arg(qkv_col)
26954            .arg(conv_state)
26955            .arg(w)
26956            .arg(conv_out)
26957            .arg(&cd)
26958            .arg(&dc);
26959        unsafe {
26960            b.launch(cfg)?;
26961        }
26962        Ok(())
26963    }
26964
26965    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
26966    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
26967    pub fn slice_range(
26968        &self,
26969        src: &CudaSlice<f32>,
26970        start: usize,
26971        len: usize,
26972    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26973        let host = self.gpu.stream().clone_dtoh(src)?;
26974        self.gpu.stream().synchronize()?;
26975        Ok(self.htod(&host[start..start + len])?)
26976    }
26977}
26978
26979#[cfg(test)]
26980mod target_dispatch_tests {
26981    use super::legacy_quant_gemm_allowed;
26982
26983    #[test]
26984    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26985        // sm_120a native lane
26986        assert!(legacy_quant_gemm_allowed(false, false, false));
26987        assert!(!legacy_quant_gemm_allowed(false, false, true));
26988        // pure portable lane (sm_89): gated
26989        assert!(!legacy_quant_gemm_allowed(true, false, false));
26990        assert!(!legacy_quant_gemm_allowed(true, false, true));
26991        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
26992        assert!(legacy_quant_gemm_allowed(true, true, false));
26993        assert!(!legacy_quant_gemm_allowed(true, true, true));
26994    }
26995
26996    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26997    #[test]
26998    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26999        assert!(!legacy_quant_gemm_allowed(
27000            cfg!(memra_portable_cuda),
27001            cfg!(memra_hopper_mma),
27002            false
27003        ));
27004    }
27005
27006    #[cfg(memra_hopper_mma)]
27007    #[test]
27008    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27009        assert!(legacy_quant_gemm_allowed(
27010            cfg!(memra_portable_cuda),
27011            cfg!(memra_hopper_mma),
27012            false
27013        ));
27014        assert!(super::portable_mma_gated() == false);
27015    }
27016}
27017
27018/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
27019/// inherent methods (inherent methods win name resolution, so no recursion).
27020impl memra_kv::KvDev for Engine {
27021    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27022        Engine::zeros(self, n)
27023    }
27024    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27025        Engine::uninit(self, n)
27026    }
27027    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27028        Engine::alloc_u8(self, n)
27029    }
27030    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27031        Engine::htod_i32(self, v)
27032    }
27033    fn clone_dtod(
27034        &self,
27035        src: &CudaSlice<f32>,
27036    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27037        Engine::clone_dtod(self, src)
27038    }
27039    fn copy_into(
27040        &self,
27041        dst: &mut CudaSlice<f32>,
27042        off: usize,
27043        src: &CudaSlice<f32>,
27044        len: usize,
27045    ) -> Result<(), Box<dyn std::error::Error>> {
27046        Engine::copy_into(self, dst, off, src, len)
27047    }
27048    fn set_i32_one(
27049        &self,
27050        d: &mut CudaSlice<i32>,
27051        v: i32,
27052    ) -> Result<(), Box<dyn std::error::Error>> {
27053        Engine::set_i32_one(self, d, v)
27054    }
27055}
27056
27057#[cfg(test)]
27058mod fused_gate_bounds_tests {
27059    use super::*;
27060
27061    /// The fused `[q|gate]` split's read-site guard, on the device.
27062    ///
27063    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
27064    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
27065    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
27066    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
27067    /// `FusedQGateExtent` before the launch.
27068    ///
27069    /// Catch demonstration for this test (guard temporarily removed, then restored):
27070    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
27071    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
27072    /// the call returns `Err`. Receipt in the lane report.
27073    #[test]
27074    #[ignore = "requires a CUDA GPU"]
27075    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27076        let e = Engine::new(0).unwrap();
27077        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27078        let fused = 2 * head_dim * n_head * t;
27079        let out_n = head_dim * n_head * t;
27080
27081        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
27082        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27083        let mut q = e.uninit(out_n).unwrap();
27084        let mut gate = e.uninit(out_n).unwrap();
27085        let err = e
27086            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27087            .expect_err("half-width wq must be refused, not read past")
27088            .to_string();
27089        assert!(err.contains("NO fused gate"), "{err}");
27090        assert!(err.contains(&format!("{fused}")), "{err}");
27091
27092        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
27093        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
27094        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27095        let wide = e.htod(&host).unwrap();
27096        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27097            .expect("full-width wq splits");
27098        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27099        for tok in 0..t {
27100            for hh in 0..n_head {
27101                for d in 0..head_dim {
27102                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27103                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27104                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27105                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27106                }
27107            }
27108        }
27109
27110        // undersized destinations are refused too (the other half of the extent contract)
27111        let mut small = e.uninit(out_n - 1).unwrap();
27112        assert!(
27113            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27114                .is_err()
27115        );
27116    }
27117}
27118
27119/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
27120/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
27121/// any launch, so the refusal is testable without a device.
27122#[cfg(test)]
27123mod fused_rope_width_tests {
27124    use super::Engine;
27125
27126    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
27127    /// safetensors route derives the same), which is why the fusion is legal there today.
27128    #[test]
27129    fn full_width_is_accepted() {
27130        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27131        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27132        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27133    }
27134
27135    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
27136    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
27137    ///
27138    /// ```text
27139    /// attention.key_length     512   rope.dimension_count     512   (global class)
27140    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
27141    /// ```
27142    ///
27143    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
27144    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
27145    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
27146    /// instead of a silently over-rotated head.
27147    #[test]
27148    fn gemma4_official_artifact_widths_pass() {
27149        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27150        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27151    }
27152
27153    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
27154    /// with no `n_dims`, silently rotating the pass-through band.
27155    #[test]
27156    fn partial_rotary_is_refused_with_the_geometry_named() {
27157        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
27158        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27159            .expect_err("partial rotary must refuse");
27160        let msg = err.to_string();
27161        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27162        assert!(msg.contains("n_rot 64"), "{msg}");
27163        assert!(msg.contains("head_dim 256"), "{msg}");
27164        assert!(
27165            msg.contains("64..256"),
27166            "names the band it would corrupt: {msg}"
27167        );
27168        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
27169        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27170        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
27171        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27172    }
27173}