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_STEP_TP_W8=1: q8_0 mirror of the step TP attention projections for DECODE.
590///
591/// NUMERIC-CLASS door, same class and acceptance as `MEMRA_STEP_TP_QKV_FUSED` /
592/// `MEMRA_BF16_MMV` / `MEMRA_SEL_GU_WPR`: the per-row arithmetic becomes an int8 dp4a dot
593/// with per-32 scales instead of a bf16xf32 fma chain, so a bit-tape cannot apply and the
594/// acceptance is the argmax gate plus the boot battery. Motivation is measured, not assumed
595/// (`decode-kernel-census`, 2026-08-25): the fused qkv shape runs 23.0 us in bf16 at
596/// 1.83 TB/s and 14.0 us in q8_0 at 1.60, and o_proj 24.2 -> 11.7 us — together
597/// ~-1.0 ms of a 13.16 ms token. Default OFF.
598pub(crate) fn step_tp_w8_on() -> bool {
599    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
600    *ON.get_or_init(|| std::env::var("MEMRA_STEP_TP_W8").as_deref() == Ok("1"))
601}
602
603/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
604/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
605/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
606/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
607pub(crate) fn sig_expf_dev_on() -> bool {
608    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
609    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
610}
611
612pub(crate) fn topk_fast_on() -> bool {
613    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
614    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
615}
616
617/// Select the sigmoid-router kernel without ever sending a shape wider than the fast
618/// kernels' fixed eight-pick scratch. The generic and dexp kernels support the full host
619/// contract; both `_fast` twins index `[warp][8]` storage and would write out of bounds for
620/// `n_used > 8` (Hermes `0d220d8c9a3eb634`).
621fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
622    match (sig_expf, fast && n_used <= 8) {
623        (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
624        (true, false) => "moe_router_sigmoid_topk_f32_dexp",
625        (false, true) => "moe_router_sigmoid_topk_f32_fast",
626        (false, false) => "moe_router_sigmoid_topk_f32",
627    }
628}
629
630#[cfg(test)]
631mod sigmoid_topk_dispatch_tests {
632    #[test]
633    fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
634        use super::sigmoid_topk_kernel;
635
636        assert_eq!(
637            sigmoid_topk_kernel(false, true, 8),
638            "moe_router_sigmoid_topk_f32_fast"
639        );
640        assert_eq!(
641            sigmoid_topk_kernel(true, true, 8),
642            "moe_router_sigmoid_topk_f32_dexp_fast"
643        );
644        assert_eq!(
645            sigmoid_topk_kernel(false, true, 9),
646            "moe_router_sigmoid_topk_f32"
647        );
648        assert_eq!(
649            sigmoid_topk_kernel(true, true, 9),
650            "moe_router_sigmoid_topk_f32_dexp"
651        );
652    }
653}
654
655pub(crate) fn rms_block() -> u32 {
656    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
657    *V.get_or_init(|| {
658        std::env::var("MEMRA_RMS_BLOCK")
659            .ok()
660            .and_then(|v| v.parse().ok())
661            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
662    })
663}
664
665pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
666    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
667    if let Some(forced) = *S.get_or_init(|| {
668        std::env::var("MEMRA_FA_SPLIT")
669            .ok()
670            .and_then(|v| v.parse().ok())
671            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
672    }) {
673        return forced;
674    }
675    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
676    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
677    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
678    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
679    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
680    //
681    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
682    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
683    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
684    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
685    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
686    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
687    // rig-divergence law: this branch is measured on 188 SMs only).
688    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
689    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
690    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
691    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
692    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
693        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
694    {
695        return if t_kv <= 8192 {
696            16
697        } else if t_kv <= 16384 {
698            64
699        } else {
700            128
701        };
702    }
703    let big_rig = fa_sm_count() >= 128;
704    if big_rig {
705        let _ = n_head_kv;
706        if t_kv <= 2048 {
707            // MEMRA_FA_SP_SHORT=N: the SHORT rung only (the SWA layers' capped t_kv lands
708            // here on step37: 33 of 45 layers at t_kv=512). At 16 the tile loop runs
709            // HALF-EMPTY (FA_DEC_TILE=32 -> nt=16 per split), so the V staging pass moves a
710            // half tile per iteration and the combine carries 2x the partials; 32 makes each
711            // split exactly one full tile. A global MEMRA_FA_SPLIT cannot isolate this — it
712            // moves the deep-ctx rung too, where more splits measured worse.
713            // NUMERIC-CLASS door (key partition -> different per-split partials/combine):
714            // new tape + battery, exactly like every other split-ladder change.
715            static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
716            if let Some(sp) = *SHORT.get_or_init(|| {
717                std::env::var("MEMRA_FA_SP_SHORT")
718                    .ok()
719                    .and_then(|v| v.parse().ok())
720                    .filter(|&s: &usize| s >= 8 && s % 8 == 0)
721            }) {
722                return sp;
723            }
724            16
725        } else if t_kv <= 16384 {
726            64
727        } else {
728            128
729        }
730    } else if n_head_kv <= 4 {
731        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
732        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
733        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
734        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
735        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
736        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
737        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
738        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
739        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
740        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
741        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
742        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
743        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
744        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
745        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
746        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
747        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
748        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
749        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
750        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
751        if t_kv <= 512 {
752            8
753        } else if t_kv <= 16384 {
754            64
755        } else {
756            128
757        }
758    } else {
759        if t_kv <= 8192 {
760            32
761        } else if t_kv <= 16384 {
762            64
763        } else {
764            128
765        }
766    }
767}
768
769/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
770/// same attribute Engine::batched_variant reads).
771pub(crate) fn fa_sm_count() -> i32 {
772    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
773    *N.get_or_init(|| {
774        cudarc::driver::result::init().ok();
775        cudarc::driver::result::device::get(0)
776            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
777                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
778            .unwrap_or(82)
779    })
780}
781
782/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
783/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
784/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
785fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
786    match head_dim {
787        256 => Ok(""),
788        128 => Ok("_hd128"),
789        d => Err(format!(
790            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
791                          callers must gate to sdpa_naive"
792        )
793        .into()),
794    }
795}
796
797/// Quant type codes matching qmatvec.cu QType enum.
798pub const QT_Q8_0: i32 = 0;
799pub const QT_Q4_K: i32 = 1;
800pub const QT_Q6_K: i32 = 2;
801pub const QT_Q5_K: i32 = 3;
802pub const QT_Q3_K: i32 = 4;
803pub const QT_IQ4_XS: i32 = 5;
804pub const QT_IQ3_S: i32 = 6;
805pub const QT_NVFP4: i32 = 7;
806/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
807/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
808/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
809/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
810/// — ONE weight copy total, no Q8_0 re-encode duplicate.
811pub const QT_F8_E4M3: i32 = 10;
812/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
813/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
814pub const QT_NVFP4_RP: i32 = 9;
815/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
816pub const QT_F32: i32 = 8;
817pub const QT_BF16: i32 = 11;
818pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
819/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
820/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
821/// dp4a/MMQ implementation exists.
822pub const QT_Q2_K: i32 = 13;
823/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
824/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
825/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
826/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
827/// scalar `scale` field is 1.0 by the layout contract.
828///
829/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
830/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
831/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
832/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
833/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
834/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
835/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
836/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
837/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
838pub const QT_F8_E4M3_BLK: i32 = 14;
839
840/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
841pub struct Engine {
842    pub gpu: memra_runtime::Gpu,
843    module: Arc<CudaModule>,
844    hybrid: Arc<CudaModule>,
845    qmatvec: Arc<CudaModule>,
846    flash: Arc<CudaModule>,
847    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
848    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
849    /// Lazy: loaded on first global-format use; None until then.
850    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
851    gemm: Arc<CudaModule>,
852    router: Arc<CudaModule>,
853    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
854    sample: Arc<CudaModule>,
855    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
856    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
857    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
858    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
859    /// MEMRA_STEP_TP_W8, hybrid half: q8_0 mirrors of bf16 GEMV weights that do NOT live in a
860    /// TP resident bank (the LM head, the shared expert, the dense-FFN layers), keyed by the
861    /// bf16 slab's device pointer and built on first decode use. The mirror is 1.0625 B/w
862    /// against bf16's 2, and the raw slab stays resident, so prefill keeps its arithmetic.
863    w8_mirrors: Mutex<std::collections::HashMap<u64, CudaSlice<u8>>>,
864    /// Per-`in_f` q8_1 activation scratch for those mirrors (allocating per call would cost
865    /// more than the door saves).
866    w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
867    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
868    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
869    /// the single largest block. The cache still owns every address for its full lifetime.
870    moe_cache_layout: Mutex<Option<Vec<usize>>>,
871    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
872    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
873    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
874    /// verify between replays) reuse their addresses and the replay reads/writes live memory
875    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
876    capture_keep_on: std::sync::atomic::AtomicBool,
877    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
878    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
879    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
880    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
881    verify_exact: std::sync::atomic::AtomicBool,
882    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
883    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
884    pub copy_stream: Arc<CudaStream>,
885    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
886    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
887    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
888    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
889    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
890    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
891    #[cfg(memra_cutlass)]
892    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
893    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
894    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
895    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
896    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
897    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
898    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
899    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
900    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
901    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
902    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
903    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
904    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
905    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
906    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
907    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
908    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
909    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
910    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
911    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
912    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
913    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
914    /// before capture under the generate_graph tracking-off window so it carries no events).
915    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
916    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
917    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
918    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
919    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
920    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
921    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
922    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
923    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
924    router_stage: Mutex<Option<PinnedStage>>,
925}
926
927/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
928/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
929/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
930/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
931/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
932/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
933/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
934/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
935/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
936fn fa_v2_on() -> bool {
937    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
938    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
939    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
940    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
941    // + graph bit-identity green on all three models.
942    std::env::var("MEMRA_FA_V2")
943        .map(|v| v != "0")
944        .unwrap_or(true)
945}
946
947/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
948/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
949/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
950/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
951/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
952/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
953/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
954pub(crate) fn fa_v3_on() -> bool {
955    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
956    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
957    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
958    std::env::var("MEMRA_FA_V3")
959        .map(|v| v != "0")
960        .unwrap_or(true)
961}
962
963/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
964/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
965/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
966/// predicate so the twins can never diverge.
967fn fa_v4_mode() -> &'static str {
968    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
969    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
970}
971fn fa_v4_on() -> bool {
972    fa_v4_mode() != "0"
973} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
974/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
975/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
976/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
977/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
978/// stays kernel-family-identical to decode at the same t_kv.
979/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
980/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
981pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
982    std::sync::atomic::AtomicUsize::new(1024);
983pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
984    std::sync::atomic::AtomicUsize::new(usize::MAX);
985pub fn fa_v4_at_pub(t_kv: usize) -> bool {
986    fa_v4_at(t_kv)
987}
988fn fa_v4_at(t_kv: usize) -> bool {
989    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
990    let mx = *M.get_or_init(|| {
991        std::env::var("MEMRA_FA_V4_MAX")
992            .ok()
993            .and_then(|v| v.parse().ok())
994            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
995    });
996    fa_v4_on() && t_kv < mx
997}
998/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
999/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
1000/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
1001/// (same split partition, same softmax/accumulation order, same partials/combine) and only
1002/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
1003/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
1004/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
1005/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
1006/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
1007/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
1008/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
1009/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
1010/// within one process (the v2/v3 pattern).
1011pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1012fn fa_deep_at(t_kv: usize) -> bool {
1013    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1014        return false;
1015    }
1016    let min = std::env::var("MEMRA_FA_DEEP_MIN")
1017        .ok()
1018        .and_then(|v| v.parse().ok())
1019        .unwrap_or(FA_DEEP_MIN_DEFAULT);
1020    t_kv >= min
1021}
1022/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
1023pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1024    fa_deep_at(t_kv)
1025}
1026
1027fn fa_v3_active(head_dim: usize) -> bool {
1028    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
1029    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
1030    fa_v3_on()
1031        && head_dim % 128 == 0
1032        && kv_cache_formats() == ("q8_0", "q5_1")
1033        && !Engine::kv_fp8_on()
1034}
1035
1036/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
1037/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
1038/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
1039/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
1040/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
1041/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
1042/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
1043pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1044    std::env::var("MEMRA_NO_FA_VEC").is_err()
1045        && t_kv >= fa_vec_min_tkv()
1046        && head_dim == 256
1047        && fa_v4_at(t_kv)
1048        && !matches!(fa_v4_mode(), "noB3" | "stage")
1049        && !Engine::kv_fp8_on()
1050}
1051/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
1052pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1053    fa_split_keys(t_kv, n_head_kv)
1054}
1055
1056/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
1057/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
1058/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
1059/// so we allocate through `result::malloc_host` with flags=0 directly.
1060struct PinnedStage {
1061    ptr: *mut u8,
1062    cap: usize,
1063}
1064unsafe impl Send for PinnedStage {}
1065impl PinnedStage {
1066    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1067        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1068        Ok(PinnedStage { ptr, cap })
1069    }
1070}
1071impl Drop for PinnedStage {
1072    fn drop(&mut self) {
1073        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1074    }
1075}
1076
1077/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
1078/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
1079pub const ARGMAX_NB: usize = 256;
1080
1081/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
1082pub(crate) use memra_fa3_vl as fa3_vl_raw;
1083
1084unsafe extern "C" {
1085    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
1086    fn memra_fa3_prefill(
1087        q16: *const core::ffi::c_void,
1088        k16: *const core::ffi::c_void,
1089        v16: *const core::ffi::c_void,
1090        o: *mut f32,
1091        t: i32,
1092        h: i32,
1093        hkv: i32,
1094        d: i32,
1095        scale: f32,
1096        stream: *mut core::ffi::c_void,
1097    ) -> i32;
1098    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
1099    pub(crate) fn memra_fa3_vl(
1100        q16s: *const *const core::ffi::c_void,
1101        k16s: *const *const core::ffi::c_void,
1102        v16s: *const *const core::ffi::c_void,
1103        os: *const *mut f32,
1104        ts: *const i32,
1105        b: i32,
1106        h: i32,
1107        hkv: i32,
1108        d: i32,
1109        scale: f32,
1110        stream: *mut core::ffi::c_void,
1111    ) -> i32;
1112}
1113
1114/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1115/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1116/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1117/// (slots are never re-allocated), so passing raw values is stable across the launch.
1118#[repr(C)]
1119#[derive(Clone, Copy)]
1120pub struct WPtr8(pub [u64; 8]);
1121unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1122
1123/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1124/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1125/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1126/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1127#[repr(C)]
1128#[derive(Clone, Copy, Default)]
1129pub struct GdnSeqVl {
1130    pub kb16: u64,
1131    pub gcum: u64,
1132    pub beta: u64,
1133    pub u: u64,
1134    pub wb16: u64,
1135    pub y: u64,
1136    pub ssnap: u64,
1137    pub state_in: u64,
1138    pub state_out: u64,
1139    pub q: u64,
1140    pub p: u64,
1141    pub o: u64,
1142    pub k: u64,
1143    pub v: u64,
1144    pub g: u64,
1145    pub a: u64,
1146    pub w: u64,
1147    pub t: i32,
1148    pub nc: i32,
1149}
1150unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1151#[repr(C)]
1152#[derive(Clone, Copy)]
1153pub struct GdnVl8(pub [GdnSeqVl; 8]);
1154unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1155
1156/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1157/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1158#[repr(C)]
1159#[derive(Clone, Copy, Default)]
1160pub struct GdnWVl {
1161    pub qb16: u64,
1162    pub pb16: u64,
1163}
1164unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1165#[repr(C)]
1166#[derive(Clone, Copy)]
1167pub struct GdnWVl8(pub [GdnWVl; 8]);
1168unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1169
1170/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1171#[repr(C)]
1172#[derive(Clone, Copy, Default)]
1173pub struct GdnPrepVl {
1174    pub qkv: u64,
1175    pub conv_state: u64,
1176    pub conv_out: u64,
1177    pub q_g: u64,
1178    pub k_g: u64,
1179    pub v_g: u64,
1180    pub q_l2: u64,
1181    pub k_l2: u64,
1182    pub beta_raw: u64,
1183    pub alpha: u64,
1184    pub beta: u64,
1185    pub g_log: u64,
1186    pub o: u64,
1187    pub z: u64,
1188    pub gn: u64,
1189    pub gn16: u64,
1190    pub kb16: u64,
1191    pub qb16: u64,
1192    pub t: i32,
1193    pub pad: i32,
1194}
1195unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1196#[repr(C)]
1197#[derive(Clone, Copy)]
1198pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1199unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1200
1201/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1202#[repr(C)]
1203#[derive(Clone, Copy, Default)]
1204pub struct FaSeqVl {
1205    pub q: u64,
1206    pub k16: u64,
1207    pub v16: u64,
1208    pub o: u64,
1209    pub kf: u64,
1210    pub vf: u64,
1211    pub t: i32,
1212    pub pad: i32,
1213}
1214unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1215#[repr(C)]
1216#[derive(Clone, Copy)]
1217pub struct FaVl8(pub [FaSeqVl; 8]);
1218unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1219
1220/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1221#[repr(C)]
1222#[derive(Clone, Copy, Default)]
1223pub struct AttnPreVl {
1224    pub qf: u64,
1225    pub kf: u64,
1226    pub vf: u64,
1227    pub q: u64,
1228    pub gate: u64,
1229    pub qn: u64,
1230    pub kn: u64,
1231    pub kc: u64,
1232    pub vc: u64,
1233    pub t: i32,
1234    pub pad: i32,
1235}
1236unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1237#[repr(C)]
1238#[derive(Clone, Copy)]
1239pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1240unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1241
1242/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1243/// varlen K1-K5 chain fills them).
1244pub struct GdnChunkBufs {
1245    pub gcum: CudaSlice<f32>,
1246    pub a: CudaSlice<f32>,
1247    pub p: CudaSlice<f32>,
1248    pub u: CudaSlice<f32>,
1249    pub w: CudaSlice<f32>,
1250    pub kb16: CudaSlice<u8>,
1251    pub wb16: CudaSlice<u8>,
1252    pub y16: CudaSlice<u8>,
1253    pub ssnap16: CudaSlice<u8>,
1254    pub qb16: CudaSlice<u8>,
1255    pub pb16: CudaSlice<u8>,
1256    pub o: CudaSlice<f32>,
1257    pub t: usize,
1258    pub nc: usize,
1259}
1260
1261/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1262#[repr(C)]
1263#[derive(Clone, Copy)]
1264pub struct F32x8(pub [f32; 8]);
1265unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1266
1267/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1268/// process. Bench binaries read it right after the call to print gen-only throughput without the
1269/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1270pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1271
1272/// RAII guard from `Engine::exact_scope`: restores the pre-scope `verify_exact` value on
1273/// drop, so error propagation (`?`) can never leave the engine latched in the
1274/// decode-exact matmul program (hermes finding, fixed 2026-08-23). Holds the flag, not
1275/// the Engine, so the restoration contract is unit-testable without a GPU.
1276#[must_use = "dropping immediately ends the exact scope"]
1277pub struct ExactScope<'a> {
1278    flag: &'a std::sync::atomic::AtomicBool,
1279    prev: bool,
1280}
1281
1282impl<'a> ExactScope<'a> {
1283    pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
1284        let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
1285        flag.store(on, std::sync::atomic::Ordering::Relaxed);
1286        ExactScope { flag, prev }
1287    }
1288}
1289
1290impl Drop for ExactScope<'_> {
1291    fn drop(&mut self) {
1292        self.flag
1293            .store(self.prev, std::sync::atomic::Ordering::Relaxed);
1294    }
1295}
1296
1297#[cfg(test)]
1298mod exact_scope_tests {
1299    use std::sync::atomic::{AtomicBool, Ordering};
1300
1301    #[test]
1302    fn error_path_restores_verify_exact() {
1303        // TOOTH (hermes finding, fixed 2026-08-23): dspark_spec_session_burst called
1304        // set_verify_exact(true)/(false) manually with `?`s in between — any error left
1305        // the engine latched in the decode-exact matmul program for every later request.
1306        // The RAII scope must restore across an error propagation.
1307        let flag = AtomicBool::new(false);
1308        let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
1309            let _scope = super::ExactScope::set(flag, true);
1310            assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
1311            Err("draft forward failed")? // the `?` exit the manual pair leaked on
1312        };
1313        assert!(failing(&flag).is_err());
1314        assert!(
1315            !flag.load(Ordering::Relaxed),
1316            "error propagation must restore the pre-scope value"
1317        );
1318        // Nested/previous-value contract: a scope entered while already ON restores ON.
1319        let flag = AtomicBool::new(true);
1320        {
1321            let _scope = super::ExactScope::set(&flag, true);
1322        }
1323        assert!(flag.load(Ordering::Relaxed));
1324        // Early drop ends the scope exactly where the manual `false` used to sit.
1325        let flag = AtomicBool::new(false);
1326        let scope = super::ExactScope::set(&flag, true);
1327        drop(scope);
1328        assert!(!flag.load(Ordering::Relaxed));
1329    }
1330}
1331
1332impl Engine {
1333    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1334        let gpu = memra_runtime::Gpu::new(ordinal)?;
1335        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1336        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1337        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1338        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1339            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1340            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1341                .and_then(|d| unsafe {
1342                    Ok((
1343                        cudarc::driver::result::device::get_attribute(
1344                            d,
1345                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1346                        )?,
1347                        cudarc::driver::result::device::get_attribute(
1348                            d,
1349                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1350                        )?,
1351                    ))
1352                })
1353                .unwrap_or((0, 0));
1354            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1355            let ok = matches!(
1356                (built, maj, min),
1357                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1358            );
1359            if !ok {
1360                return Err(format!(
1361                    "memra was built for sm_{built} but device {ordinal} reports compute \
1362                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1363                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1364                )
1365                .into());
1366            }
1367        }
1368        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1369        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1370        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1371        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1372        unsafe {
1373            use cudarc::driver::sys;
1374            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1375            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1376            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1377                let mut thresh: u64 = u64::MAX;
1378                let _ = sys::cuMemPoolSetAttribute(
1379                    pool,
1380                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1381                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1382                );
1383            }
1384        }
1385        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1386        let hybrid = gpu
1387            .ctx
1388            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1389        let qmatvec = gpu
1390            .ctx
1391            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1392        let flash = gpu
1393            .ctx
1394            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1395        let gemm = gpu
1396            .ctx
1397            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1398        let router = gpu
1399            .ctx
1400            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1401        let sample = gpu
1402            .ctx
1403            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1404        let copy_stream = gpu.ctx.new_stream()?;
1405        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1406        // cudarc is in multi-stream mode (main stream +
1407        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1408        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1409        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1410        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
1411        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1412        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1413        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1414        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1415        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1416        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1417        // implicit event tracking.
1418        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1419        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1420        if std::env::var("MEMRA_EVT")
1421            .map(|v| v == "1")
1422            .unwrap_or(false)
1423        {
1424            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1425        } else {
1426            unsafe {
1427                gpu.ctx.disable_event_tracking();
1428            }
1429        }
1430        Ok(Self {
1431            gpu,
1432            module,
1433            hybrid,
1434            qmatvec,
1435            flash,
1436            flash_g: std::sync::OnceLock::new(),
1437            gemm,
1438            router,
1439            sample,
1440            moe_cache: Mutex::new(None),
1441            w8_mirrors: Mutex::new(std::collections::HashMap::new()),
1442            w8_act: Mutex::new(std::collections::HashMap::new()),
1443            moe_cache_layout: Mutex::new(None),
1444            copy_stream,
1445            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1446            verify_exact: std::sync::atomic::AtomicBool::new(false),
1447            capture_keep: Mutex::new(Vec::new()),
1448            argmax_partials: Mutex::new(None),
1449            prime_deqw_ws: Mutex::new(None),
1450            router_stage: Mutex::new(None),
1451            fp8_scratch: Mutex::new(None),
1452            fa_vf16_scratch: Mutex::new(None),
1453            fa_part_pool: Mutex::new(None),
1454            fa_part_retired: Mutex::new(Vec::new()),
1455            fn_cache: Mutex::new(Default::default()),
1456            f16_scratch: Mutex::new(None),
1457            #[cfg(memra_cutlass)]
1458            cutlass_scratch: Mutex::new(None),
1459        })
1460    }
1461
1462    pub fn ctx(&self) -> &Arc<CudaContext> {
1463        &self.gpu.ctx
1464    }
1465
1466    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1467    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1468    ///
1469    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1470    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1471    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1472    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1473    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1474    ///
1475    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1476    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1477    /// under-count headroom does not belong in a gate that queues real work, but the honest
1478    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1479    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1480    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1481    ///
1482    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1483    pub fn pool_cached_bytes(&self) -> usize {
1484        let (reserved, used) = self.pool_reserved_used();
1485        reserved.saturating_sub(used)
1486    }
1487
1488    /// Bytes the driver's per-device CUDA GRAPH memory pool currently holds RESERVED
1489    /// (`cuDeviceGetGraphMemAttribute` RESERVED_MEM_CURRENT) — the backing store of every
1490    /// captured alloc node, which on this engine means the dspark verify-graph pool
1491    /// (decode/step graphs bake pre-allocated buffers and own no alloc nodes). This memory
1492    /// is DISTINCT from the async pool above: `mem_get_info`'s `free` already excludes it,
1493    /// it is never released back (the vgraph pool has no eviction by design), and it GROWS
1494    /// as new (segment, vt)/(vt, rung, hi) keys capture — the growth is what
1495    /// `dspark_vg_admission_debt` charges at admission. Returns 0 if the attribute cannot
1496    /// be queried (never a false headroom claim, matching `pool_cached_bytes`).
1497    pub fn device_graph_mem_reserved(&self) -> usize {
1498        use cudarc::driver::sys as cus;
1499        let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1500            return 0;
1501        };
1502        let mut bytes: u64 = 0;
1503        let rc = unsafe {
1504            cus::cuDeviceGetGraphMemAttribute(
1505                dev,
1506                cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1507                &mut bytes as *mut u64 as *mut std::ffi::c_void,
1508            )
1509        };
1510        if rc == cus::cudaError_enum::CUDA_SUCCESS {
1511            bytes as usize
1512        } else {
1513            0
1514        }
1515    }
1516
1517    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1518    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1519    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1520    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1521    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1522    /// (0, 0) if the pool cannot be queried.
1523    /// Release every CACHED (freed-but-retained) block of the default async mempool
1524    /// back to the driver (deploy-headroom lane, 2026-08-27). The boot-time
1525    /// RELEASE_THRESHOLD=u64::MAX pin keeps freed blocks cached for graph-launch speed,
1526    /// which is right for steady serving and wrong at a blue/green overlap: a green
1527    /// PROCESS cannot use blue's cached pool. cuMemPoolTrimTo(0) frees only unused
1528    /// blocks — live allocations are untouched; later allocs re-map once. Returns the
1529    /// bytes released (reserved delta), 0 if the pool cannot be queried.
1530    pub fn pool_trim_to_zero(&self) -> usize {
1531        use cudarc::driver::sys;
1532        let (before, _) = self.pool_reserved_used();
1533        unsafe {
1534            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1535            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1536                != sys::CUresult::CUDA_SUCCESS
1537            {
1538                return 0;
1539            }
1540            let _ = sys::cuMemPoolTrimTo(pool, 0);
1541        }
1542        let (after, _) = self.pool_reserved_used();
1543        before.saturating_sub(after)
1544    }
1545
1546    pub fn pool_reserved_used(&self) -> (usize, usize) {
1547        use cudarc::driver::sys;
1548        unsafe {
1549            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1550            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1551                != sys::CUresult::CUDA_SUCCESS
1552            {
1553                return (0, 0);
1554            }
1555            let (mut reserved, mut used) = (0u64, 0u64);
1556            if sys::cuMemPoolGetAttribute(
1557                pool,
1558                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1559                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1560            ) != sys::CUresult::CUDA_SUCCESS
1561            {
1562                return (0, 0);
1563            }
1564            if sys::cuMemPoolGetAttribute(
1565                pool,
1566                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1567                &mut used as *mut u64 as *mut core::ffi::c_void,
1568            ) != sys::CUresult::CUDA_SUCCESS
1569            {
1570                return (0, 0);
1571            }
1572            (reserved as usize, used as usize)
1573        }
1574    }
1575
1576    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1577    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1578    pub fn stream(&self) -> Arc<CudaStream> {
1579        self.gpu.stream()
1580    }
1581    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1582    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1583    pub fn gkv_on() -> bool {
1584        memra_kv::gkv_on()
1585    }
1586
1587    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1588    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1589    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1590    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1591    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1592    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1593    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1594    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1595    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1596    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1597    /// ON for both — no acceptance cost measured.
1598    pub fn wkv_on() -> bool {
1599        memra_kv::wkv_on()
1600    }
1601
1602    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1603    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1604    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1605    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1606    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1607    pub fn kv_fp8_on() -> bool {
1608        memra_kv::kv_fp8_on()
1609    }
1610
1611    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1612    /// when the fp8-globals arm is on; everything else from the default flash module.
1613    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1614        if head_dim == 512 && Self::gkv_on() {
1615            self.func_g(name)
1616        } else {
1617            self.func(name)
1618        }
1619    }
1620
1621    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1622    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1623    /// per-format fatbins; fall back to the base modules for those.
1624    fn func_g(&self, name: &str) -> CudaFunction {
1625        let m = self.flash_g.get_or_init(|| {
1626            self.gpu
1627                .ctx
1628                .load_module(cudarc::nvrtc::Ptx::from_binary(
1629                    FLASH_FATBIN_KF8VF8.to_vec(),
1630                ))
1631                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1632        });
1633        let key = format!("g:{name}");
1634        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1635            return f.clone();
1636        }
1637        let f = match m.load_function(name) {
1638            Ok(f) => f,
1639            Err(_) => self.func(name),
1640        };
1641        self.fn_cache.lock().unwrap().insert(key, f.clone());
1642        f
1643    }
1644
1645    fn func(&self, name: &str) -> CudaFunction {
1646        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1647        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1648        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1649            return f.clone();
1650        }
1651        let f = self
1652            .module
1653            .load_function(name)
1654            .or_else(|_| self.hybrid.load_function(name))
1655            .or_else(|_| self.qmatvec.load_function(name))
1656            .or_else(|_| self.flash.load_function(name))
1657            .or_else(|_| self.gemm.load_function(name))
1658            .or_else(|_| self.router.load_function(name))
1659            .or_else(|_| self.sample.load_function(name))
1660            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1661        self.fn_cache
1662            .lock()
1663            .unwrap()
1664            .insert(name.to_string(), f.clone());
1665        f
1666    }
1667
1668    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1669    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1670    pub fn scatter_trim_logits(
1671        &self,
1672        src: &CudaSlice<f32>,
1673        d2t: &CudaSlice<u32>,
1674        dst: &mut CudaSlice<f32>,
1675        d_vocab: usize,
1676        n_vocab: usize,
1677    ) -> Result<(), Box<dyn std::error::Error>> {
1678        let f1 = self.func("scatter_trim_logits_f32");
1679        let f2 = self.func("scatter_trim_logits_pass2_f32");
1680        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1681        let cfg1 = LaunchConfig {
1682            grid_dim: (256, 1, 1),
1683            block_dim: (256, 1, 1),
1684            shared_mem_bytes: 0,
1685        };
1686        let __s_b1 = self.gpu.stream();
1687        let mut b1 = __s_b1.launch_builder(&f1);
1688        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1689        unsafe {
1690            b1.launch(cfg1)?;
1691        }
1692        let cfg2 = LaunchConfig {
1693            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1694            block_dim: (256, 1, 1),
1695            shared_mem_bytes: 0,
1696        };
1697        let __s_b2 = self.gpu.stream();
1698        let mut b2 = __s_b2.launch_builder(&f2);
1699        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1700        unsafe {
1701            b2.launch(cfg2)?;
1702        }
1703        Ok(())
1704    }
1705
1706    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1707    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1708
1709    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1710    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1711    #[allow(clippy::too_many_arguments)]
1712    pub fn filter_stats(
1713        &self,
1714        x: &CudaSlice<f32>,
1715        row_stride: usize,
1716        rows: &CudaSlice<i32>,
1717        out_th: &mut CudaSlice<f32>,
1718        out_z: &mut CudaSlice<f32>,
1719        out_max: &mut CudaSlice<f32>,
1720        n: usize,
1721        nrow: usize,
1722        temp: f32,
1723        top_k: i32,
1724        top_p: f32,
1725        min_p: f32,
1726    ) -> Result<(), Box<dyn std::error::Error>> {
1727        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1728        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1729        // L2-resident, so the extra passes are near-free while the per-thread selection list
1730        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1731        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1732        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1733        //
1734        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1735        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1736        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1737        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1738        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1739        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1740        //
1741        // DETERMINISTIC KEYING (hermes finding, fixed 2026-08-23): the old admission
1742        // `16*nrow <= sm_count` fell back to the single-block program PER CALL when a tick
1743        // carried too many rows — and the two programs are NOT bit-identical (measured
1744        // ~1e-7 rel on the renorm mass: different f32 partial-sum order), so a request's
1745        // sampling threshold arithmetic depended on how many rows shared its serve tick.
1746        // Coop is now THE program on every coop-capable device: rows are CHUNKED to the
1747        // co-residency cap (sm_count/16 rows per cooperative launch) and each row's
1748        // arithmetic uses only its own 16 slices + its own ws region, so the per-row bits
1749        // are independent of batch width by construction — the kernel-check
1750        // FILTER-COOP-CHUNK cell pins exactly that. The single-block program remains only
1751        // behind the deployment-keyed seams: MEMRA_FILTER_COOP=0, or a device with
1752        // sm_count < 16 (fixed per device class, never per call).
1753        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1754        let coop_on =
1755            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1756        if coop_on && self.sm_count() >= 16 {
1757            let cap = self.sm_count() as usize / 16;
1758            let mut done = 0usize;
1759            while done < nrow {
1760                let chunk = cap.min(nrow - done);
1761                self.filter_stats_coop_chunk(
1762                    x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1763                    top_p, min_p,
1764                )?;
1765                done += chunk;
1766            }
1767            return Ok(());
1768        }
1769        self.filter_stats_plain_program(
1770            x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1771        )
1772    }
1773
1774    /// One cooperative `filter_stats` launch over rows `row0..row0+chunk` (pub so the
1775    /// kernel-check FILTER-COOP-CHUNK cell can pin batch-width independence directly).
1776    /// The kernel indexes `rows`/outputs by blockIdx.y, so the chunk is expressed as
1777    /// sub-views at `row0` — per-row arithmetic is untouched by the offset.
1778    #[allow(clippy::too_many_arguments)]
1779    pub fn filter_stats_coop_chunk(
1780        &self,
1781        x: &CudaSlice<f32>,
1782        row_stride: usize,
1783        rows: &CudaSlice<i32>,
1784        row0: usize,
1785        out_th: &mut CudaSlice<f32>,
1786        out_z: &mut CudaSlice<f32>,
1787        out_max: &mut CudaSlice<f32>,
1788        n: usize,
1789        chunk: usize,
1790        temp: f32,
1791        top_k: i32,
1792        top_p: f32,
1793        min_p: f32,
1794    ) -> Result<(), Box<dyn std::error::Error>> {
1795        let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1796        let f = self.func("filter_stats_coop_f32");
1797        let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1798        let cfg = LaunchConfig {
1799            grid_dim: (16, chunk as u32, 1),
1800            block_dim: (512, 1, 1),
1801            shared_mem_bytes: 0,
1802        };
1803        let rows_v = rows.slice(row0..row0 + chunk);
1804        let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1805        let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1806        let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1807        let __s_b = self.gpu.stream();
1808        let mut b = __s_b.launch_builder(&f);
1809        b.arg(x)
1810            .arg(&rs)
1811            .arg(&rows_v)
1812            .arg(&mut th_v)
1813            .arg(&mut z_v)
1814            .arg(&mut mx_v)
1815            .arg(&mut ws)
1816            .arg(&ni)
1817            .arg(&nr)
1818            .arg(&temp)
1819            .arg(&top_k)
1820            .arg(&top_p)
1821            .arg(&min_p);
1822        unsafe {
1823            b.launch_cooperative(cfg)?;
1824        }
1825        Ok(())
1826    }
1827
1828    /// The single-block-per-row `filter_stats` program (the pre-coop form; the
1829    /// MEMRA_FILTER_COOP=0 rollback and the occupancy fallback). Gate-callable twin of
1830    /// `filter_stats_coop_program`.
1831    #[allow(clippy::too_many_arguments)]
1832    pub fn filter_stats_plain_program(
1833        &self,
1834        x: &CudaSlice<f32>,
1835        row_stride: usize,
1836        rows: &CudaSlice<i32>,
1837        out_th: &mut CudaSlice<f32>,
1838        out_z: &mut CudaSlice<f32>,
1839        out_max: &mut CudaSlice<f32>,
1840        n: usize,
1841        nrow: usize,
1842        temp: f32,
1843        top_k: i32,
1844        top_p: f32,
1845        min_p: f32,
1846    ) -> Result<(), Box<dyn std::error::Error>> {
1847        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1848        let f = self.func("filter_stats_f32");
1849        let cfg = LaunchConfig {
1850            grid_dim: (nrow as u32, 1, 1),
1851            block_dim: (1024, 1, 1),
1852            shared_mem_bytes: 0,
1853        };
1854        let __s_b = self.gpu.stream();
1855        let mut b = __s_b.launch_builder(&f);
1856        b.arg(x)
1857            .arg(&rs)
1858            .arg(rows)
1859            .arg(&mut *out_th)
1860            .arg(&mut *out_z)
1861            .arg(&mut *out_max)
1862            .arg(&ni)
1863            .arg(&nr)
1864            .arg(&temp)
1865            .arg(&top_k)
1866            .arg(&top_p)
1867            .arg(&min_p);
1868        unsafe {
1869            b.launch(cfg)?;
1870        }
1871        Ok(())
1872    }
1873
1874    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1875    #[allow(clippy::too_many_arguments)]
1876    pub fn softmax_gather_filtered(
1877        &self,
1878        x: &CudaSlice<f32>,
1879        row_stride: usize,
1880        ids: &CudaSlice<u32>,
1881        rows: &CudaSlice<i32>,
1882        th: &CudaSlice<f32>,
1883        z: &CudaSlice<f32>,
1884        out: &mut CudaSlice<f32>,
1885        n: usize,
1886        npair: usize,
1887        temp: f32,
1888    ) -> Result<(), Box<dyn std::error::Error>> {
1889        let f = self.func("softmax_gather_filtered_f32");
1890        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1891        let cfg = LaunchConfig {
1892            grid_dim: (npair as u32, 1, 1),
1893            block_dim: (256, 1, 1),
1894            shared_mem_bytes: 0,
1895        };
1896        let __s_b = self.gpu.stream();
1897        let mut b = __s_b.launch_builder(&f);
1898        b.arg(x)
1899            .arg(&rs)
1900            .arg(ids)
1901            .arg(rows)
1902            .arg(th)
1903            .arg(z)
1904            .arg(&mut *out)
1905            .arg(&ni)
1906            .arg(&np)
1907            .arg(&temp);
1908        unsafe {
1909            b.launch(cfg)?;
1910        }
1911        Ok(())
1912    }
1913
1914    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1915    #[allow(clippy::too_many_arguments)]
1916    pub fn residual_sample_filtered(
1917        &self,
1918        p: &CudaSlice<f32>,
1919        q: Option<&CudaSlice<f32>>,
1920        n: usize,
1921        temp: f32,
1922        seed: u64,
1923        stream_pos: u32,
1924        p_stats: (f32, f32, f32),
1925        q_stats: (f32, f32, f32),
1926        out_tok: &mut CudaSlice<u32>,
1927    ) -> Result<(), Box<dyn std::error::Error>> {
1928        let f = self.func("residual_sample_filtered_f32");
1929        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1930        let has_q: i32 = q.is_some() as i32;
1931        let qbuf = q.unwrap_or(p);
1932        let (pm, pth, pz) = p_stats;
1933        let (qm, qth, qz) = q_stats;
1934        let cfg = LaunchConfig {
1935            grid_dim: (1, 1, 1),
1936            block_dim: (1024, 1, 1),
1937            shared_mem_bytes: 0,
1938        };
1939        let __s_b = self.gpu.stream();
1940        let mut b = __s_b.launch_builder(&f);
1941        b.arg(p)
1942            .arg(qbuf)
1943            .arg(&has_q)
1944            .arg(&ni)
1945            .arg(&temp)
1946            .arg(&slo)
1947            .arg(&shi)
1948            .arg(&stream_pos)
1949            .arg(&pm)
1950            .arg(&pth)
1951            .arg(&pz)
1952            .arg(&qm)
1953            .arg(&qth)
1954            .arg(&qz)
1955            .arg(&mut *out_tok);
1956        unsafe {
1957            b.launch(cfg)?;
1958        }
1959        Ok(())
1960    }
1961
1962    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1963    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1964    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1965    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1966    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1967    #[allow(clippy::too_many_arguments)]
1968    pub fn residual_sample_sparse_q(
1969        &self,
1970        p: &CudaSlice<f32>,
1971        cand_ids: &CudaSlice<u32>,
1972        q_probs: &CudaSlice<f32>,
1973        n_cand: usize,
1974        n: usize,
1975        temp: f32,
1976        seed: u64,
1977        stream_pos: u32,
1978        p_stats: (f32, f32, f32),
1979        out_tok: &mut CudaSlice<u32>,
1980    ) -> Result<(), Box<dyn std::error::Error>> {
1981        assert!(
1982            n_cand >= 1 && n_cand <= 32,
1983            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1984        );
1985        let f = self.func("residual_sample_sparse_q_f32");
1986        let (ni, nc) = (n as i32, n_cand as i32);
1987        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1988        let (pm, pth, pz) = p_stats;
1989        let cfg = LaunchConfig {
1990            grid_dim: (1, 1, 1),
1991            block_dim: (1024, 1, 1),
1992            shared_mem_bytes: 0,
1993        };
1994        let __s_b = self.gpu.stream();
1995        let mut b = __s_b.launch_builder(&f);
1996        b.arg(p)
1997            .arg(cand_ids)
1998            .arg(q_probs)
1999            .arg(&nc)
2000            .arg(&ni)
2001            .arg(&temp)
2002            .arg(&slo)
2003            .arg(&shi)
2004            .arg(&stream_pos)
2005            .arg(&pm)
2006            .arg(&pth)
2007            .arg(&pz)
2008            .arg(&mut *out_tok);
2009        unsafe {
2010            b.launch(cfg)?;
2011        }
2012        Ok(())
2013    }
2014
2015    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
2016    #[allow(clippy::too_many_arguments)]
2017    pub fn gumbel_perturb_filtered(
2018        &self,
2019        x: &CudaSlice<f32>,
2020        y: &mut CudaSlice<f32>,
2021        n: usize,
2022        seed: u64,
2023        stream_pos: u32,
2024        temp: f32,
2025        row_max: f32,
2026        th: f32,
2027    ) -> Result<(), Box<dyn std::error::Error>> {
2028        let f = self.func("gumbel_perturb_filtered_f32");
2029        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2030        let cfg = LaunchConfig {
2031            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2032            block_dim: (256, 1, 1),
2033            shared_mem_bytes: 0,
2034        };
2035        let __s_b = self.gpu.stream();
2036        let mut b = __s_b.launch_builder(&f);
2037        b.arg(x)
2038            .arg(&mut *y)
2039            .arg(&ni)
2040            .arg(&slo)
2041            .arg(&shi)
2042            .arg(&stream_pos)
2043            .arg(&temp)
2044            .arg(&row_max)
2045            .arg(&th);
2046        unsafe {
2047            b.launch(cfg)?;
2048        }
2049        Ok(())
2050    }
2051
2052    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
2053    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
2054    /// filtered rejection sampling exact for the penalized target.
2055    #[allow(clippy::too_many_arguments)]
2056    pub fn penalize_logits(
2057        &self,
2058        x: &mut CudaSlice<f32>,
2059        hist: &CudaSlice<u32>,
2060        n_hist: usize,
2061        rep: f32,
2062        freq: f32,
2063        present: f32,
2064        n: usize,
2065    ) -> Result<(), Box<dyn std::error::Error>> {
2066        if n_hist == 0 {
2067            return Ok(());
2068        }
2069        let f = self.func("penalize_logits_f32");
2070        let (nh, ni) = (n_hist as i32, n as i32);
2071        let cfg = LaunchConfig {
2072            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2073            block_dim: (128, 1, 1),
2074            shared_mem_bytes: 0,
2075        };
2076        let __s_b = self.gpu.stream();
2077        let mut b = __s_b.launch_builder(&f);
2078        b.arg(&mut *x)
2079            .arg(hist)
2080            .arg(&nh)
2081            .arg(&rep)
2082            .arg(&freq)
2083            .arg(&present)
2084            .arg(&ni);
2085        unsafe {
2086            b.launch(cfg)?;
2087        }
2088        Ok(())
2089    }
2090
2091    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
2092    #[allow(clippy::too_many_arguments)]
2093    pub fn penalize_logits_rows(
2094        &self,
2095        x: &mut CudaSlice<f32>,
2096        hist: &CudaSlice<u32>,
2097        n_hist: usize,
2098        rep: f32,
2099        freq: f32,
2100        present: f32,
2101        n: usize,
2102        nrow: usize,
2103    ) -> Result<(), Box<dyn std::error::Error>> {
2104        if n_hist == 0 || nrow == 0 {
2105            return Ok(());
2106        }
2107        let f = self.func("penalize_logits_rows_f32");
2108        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2109        let cfg = LaunchConfig {
2110            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2111            block_dim: (128, 1, 1),
2112            shared_mem_bytes: 0,
2113        };
2114        let __s_b = self.gpu.stream();
2115        let mut b = __s_b.launch_builder(&f);
2116        b.arg(&mut *x)
2117            .arg(hist)
2118            .arg(&nh)
2119            .arg(&rep)
2120            .arg(&freq)
2121            .arg(&present)
2122            .arg(&ni)
2123            .arg(&nr);
2124        unsafe {
2125            b.launch(cfg)?;
2126        }
2127        Ok(())
2128    }
2129
2130    /// Heterogeneous serving-batch penalties over host-maintained sparse window counts.
2131    /// `offsets[r]..offsets[r+1]` indexes the unique positive-count `(id,count)` entries for logits row
2132    /// `rows[r]`; each row may carry independent repetition/frequency/presence coefficients.
2133    /// One thread owns one distinct logit, so the kernel needs neither atomics nor the
2134    /// history-squared dedup scan used by the speculative raw-history oracle.
2135    #[allow(clippy::too_many_arguments)]
2136    pub fn penalize_logits_sparse_rows(
2137        &self,
2138        x: &mut CudaSlice<f32>,
2139        ids: &[u32],
2140        counts: &[u32],
2141        offsets: &[i32],
2142        rows: &[i32],
2143        reps: &[f32],
2144        freqs: &[f32],
2145        presents: &[f32],
2146        n: usize,
2147    ) -> Result<(), Box<dyn std::error::Error>> {
2148        let nrow = rows.len();
2149        if nrow == 0 {
2150            return Ok(());
2151        }
2152        let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2153        let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2154        let entry_count =
2155            i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2156        if ids.len() != counts.len()
2157            || offsets.len() != nrow + 1
2158            || reps.len() != nrow
2159            || freqs.len() != nrow
2160            || presents.len() != nrow
2161            || offsets.first().copied() != Some(0)
2162            || offsets.last().copied() != Some(entry_count)
2163        {
2164            return Err("sparse penalty row metadata shape mismatch".into());
2165        }
2166        if counts.contains(&0) {
2167            return Err("sparse penalty counts must be positive".into());
2168        }
2169        let mut max_len = 0usize;
2170        for pair in offsets.windows(2) {
2171            if pair[0] < 0 || pair[1] < pair[0] {
2172                return Err("sparse penalty offsets must be monotonic".into());
2173            }
2174            max_len = max_len.max((pair[1] - pair[0]) as usize);
2175        }
2176        if max_len == 0 {
2177            return Ok(());
2178        }
2179
2180        let mut seen = std::collections::HashSet::with_capacity(ids.len());
2181        for (r, &row) in rows.iter().enumerate() {
2182            if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2183                return Err("sparse penalty row index exceeds logits shape".into());
2184            }
2185            let begin = offsets[r] as usize;
2186            let end = offsets[r + 1] as usize;
2187            for &id in &ids[begin..end] {
2188                if id as usize >= n {
2189                    return Err("sparse penalty token id exceeds logits row".into());
2190                }
2191                if !seen.insert((row, id)) {
2192                    return Err("sparse penalty entries must be unique per logits row".into());
2193                }
2194            }
2195        }
2196
2197        // SAFETY: the checks above establish every invariant of the launch-only helper.
2198        unsafe {
2199            self.penalize_logits_sparse_rows_unchecked(
2200                x, ids, counts, offsets, rows, reps, freqs, presents, n,
2201            )
2202        }
2203    }
2204
2205    /// Launch-only form for the serving hot path, whose `HashMap`-backed producer already
2206    /// guarantees unique ids and whose rows are enumerated from the live batch.
2207    ///
2208    /// # Safety
2209    ///
2210    /// Shapes must match the safe wrapper, offsets must be monotonic and in bounds, every row
2211    /// must index `x`, and each `(row,id)` pair must occur at most once. Token ids outside the
2212    /// logits row are safe no-ops because the kernel bounds-checks them before computing `x`.
2213    #[allow(clippy::too_many_arguments)]
2214    pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2215        &self,
2216        x: &mut CudaSlice<f32>,
2217        ids: &[u32],
2218        counts: &[u32],
2219        offsets: &[i32],
2220        rows: &[i32],
2221        reps: &[f32],
2222        freqs: &[f32],
2223        presents: &[f32],
2224        n: usize,
2225    ) -> Result<(), Box<dyn std::error::Error>> {
2226        let nrow = rows.len();
2227        if nrow == 0 {
2228            return Ok(());
2229        }
2230        let max_len = offsets
2231            .windows(2)
2232            .map(|pair| (pair[1] - pair[0]) as usize)
2233            .max()
2234            .unwrap_or(0);
2235        if max_len == 0 {
2236            return Ok(());
2237        }
2238        let ids_d = self.htod_u32_v(ids)?;
2239        let counts_d = self.htod_u32_v(counts)?;
2240        let offsets_d = self.htod_i32(offsets)?;
2241        let rows_d = self.htod_i32(rows)?;
2242        let reps_d = self.htod(reps)?;
2243        let freqs_d = self.htod(freqs)?;
2244        let presents_d = self.htod(presents)?;
2245        let f = self.func("penalize_logits_sparse_rows_f32");
2246        let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2247        let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2248        let cfg = LaunchConfig {
2249            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2250            block_dim: (128, 1, 1),
2251            shared_mem_bytes: 0,
2252        };
2253        let __s_b = self.gpu.stream();
2254        let mut b = __s_b.launch_builder(&f);
2255        b.arg(&mut *x)
2256            .arg(&ids_d)
2257            .arg(&counts_d)
2258            .arg(&offsets_d)
2259            .arg(&rows_d)
2260            .arg(&reps_d)
2261            .arg(&freqs_d)
2262            .arg(&presents_d)
2263            .arg(&ni)
2264            .arg(&nr);
2265        unsafe {
2266            b.launch(cfg)?;
2267        }
2268        Ok(())
2269    }
2270
2271    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
2272    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
2273    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
2274    /// is the within-round evolving penalty state block drafting needs: verify row r's
2275    /// target is penalized by every token committed before it INCLUDING same-round
2276    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
2277    /// approximation this exists to replace on the dspark route.
2278    #[allow(clippy::too_many_arguments)]
2279    pub fn penalize_logits_rows_inc(
2280        &self,
2281        x: &mut CudaSlice<f32>,
2282        hist: &CudaSlice<u32>,
2283        n_hist0: usize,
2284        rep: f32,
2285        freq: f32,
2286        present: f32,
2287        n: usize,
2288        nrow: usize,
2289        win: usize,
2290    ) -> Result<(), Box<dyn std::error::Error>> {
2291        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2292            return Ok(());
2293        }
2294        debug_assert!(
2295            hist.len() >= n_hist0 + nrow - 1,
2296            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2297        );
2298        let f = self.func("penalize_logits_rows_inc_f32");
2299        let max_len = win.min(n_hist0 + nrow - 1).max(1);
2300        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2301        let cfg = LaunchConfig {
2302            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2303            block_dim: (128, 1, 1),
2304            shared_mem_bytes: 0,
2305        };
2306        let __s_b = self.gpu.stream();
2307        let mut b = __s_b.launch_builder(&f);
2308        b.arg(&mut *x)
2309            .arg(hist)
2310            .arg(&nh)
2311            .arg(&rep)
2312            .arg(&freq)
2313            .arg(&present)
2314            .arg(&ni)
2315            .arg(&nr)
2316            .arg(&wi);
2317        unsafe {
2318            b.launch(cfg)?;
2319        }
2320        Ok(())
2321    }
2322
2323    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
2324    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
2325    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
2326    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
2327    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
2328    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
2329    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
2330    pub fn wpf_level() -> u32 {
2331        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2332        *ON.get_or_init(|| {
2333            std::env::var("MEMRA_WPF")
2334                .ok()
2335                .and_then(|v| v.parse().ok())
2336                .unwrap_or(1)
2337        })
2338    }
2339
2340    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
2341    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
2342    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
2343    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
2344    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
2345    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
2346    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
2347    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
2348    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
2349    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
2350    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
2351    /// Prefer `exact_scope` — the RAII form — anywhere a `?` can exit the scope: a manual
2352    /// true/false pair leaves the flag LATCHED engine-wide when an error propagates
2353    /// between the two calls (hermes finding on dspark_spec_session_burst, fixed
2354    /// 2026-08-23), and every later request then runs the exact-GEMM program.
2355    pub fn set_verify_exact(&self, on: bool) {
2356        self.verify_exact
2357            .store(on, std::sync::atomic::Ordering::Relaxed);
2358    }
2359    pub(crate) fn verify_exact_on(&self) -> bool {
2360        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2361    }
2362
2363    /// RAII scope over `verify_exact`: sets the flag to `on` now and restores the
2364    /// PREVIOUS value on drop — unwind, early `return`, and every `?` exit included.
2365    /// This is the required form for any scope an error can leave (see
2366    /// `set_verify_exact`); dropping the guard early (`drop(scope)`) ends the scope
2367    /// exactly where the manual `set_verify_exact(false)` used to sit.
2368    pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2369        ExactScope::set(&self.verify_exact, on)
2370    }
2371
2372    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
2373    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
2374    pub fn qkv_append_on() -> bool {
2375        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2376        *ON.get_or_init(|| {
2377            std::env::var("MEMRA_QKV_APPEND")
2378                .map(|v| v != "0")
2379                .unwrap_or(true)
2380        })
2381    }
2382
2383    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
2384    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
2385    pub fn pdl_wb_on() -> bool {
2386        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2387        *ON.get_or_init(|| {
2388            std::env::var("MEMRA_PDL_WB")
2389                .map(|v| v != "0")
2390                .unwrap_or(true)
2391        })
2392    }
2393
2394    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
2395    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
2396    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
2397    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
2398    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
2399    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
2400    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
2401    pub fn norm_ilp_on() -> bool {
2402        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2403        *ON.get_or_init(|| {
2404            std::env::var("MEMRA_NORM_ILP")
2405                .map(|v| v != "0")
2406                .unwrap_or(true)
2407        })
2408    }
2409
2410    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
2411    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
2412    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
2413    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
2414    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
2415    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
2416    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
2417    pub fn tk_ffn_dual_on() -> bool {
2418        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2419        *ON.get_or_init(|| {
2420            std::env::var("MEMRA_TK_FFN_DUAL")
2421                .map(|v| v != "0")
2422                .unwrap_or(true)
2423        })
2424    }
2425
2426    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
2427    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
2428    /// per-model no-harm bisect knob.
2429    pub fn pdl_mmvq_on() -> bool {
2430        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2431        *ON.get_or_init(|| {
2432            std::env::var("MEMRA_PDL_MMVQ")
2433                .map(|v| v != "0")
2434                .unwrap_or(true)
2435        })
2436    }
2437
2438    pub fn pdl_on() -> bool {
2439        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2440        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2441    }
2442
2443    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
2444    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
2445    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
2446    /// on the producer before any read), bit-identical by construction.
2447    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
2448    pub fn pdl_nvfp4q8_on() -> bool {
2449        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2450        *ON.get_or_init(|| {
2451            std::env::var("MEMRA_PDL_NVFP4")
2452                .map(|v| v != "0")
2453                .unwrap_or(true)
2454        })
2455    }
2456
2457    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
2458    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
2459    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
2460    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
2461    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
2462    fn q40_mr1_on() -> bool {
2463        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2464        match *Q40MR.get_or_init(|| {
2465            std::env::var("MEMRA_Q40_MR")
2466                .ok()
2467                .and_then(|v| v.parse().ok())
2468        }) {
2469            Some(v) => v == 1,
2470            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2471        }
2472    }
2473
2474    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
2475    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
2476    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
2477    /// writes wrong bytes silently.
2478    fn pdl_func_flash(
2479        &self,
2480        g: bool,
2481        name: &'static str,
2482    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2483        use cudarc::driver::sys as cu;
2484        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
2485        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
2486        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
2487        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
2488        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
2489        // this engine's CUcontext; single-context runs behave exactly as before.
2490        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2491            std::sync::Mutex::new(None);
2492        static FNS: std::sync::Mutex<
2493            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2494        > = std::sync::Mutex::new(None);
2495        let ctx_key = self.ctx().cu_ctx() as usize;
2496        if let Some(&f) = FNS
2497            .lock()
2498            .unwrap()
2499            .get_or_insert_with(Default::default)
2500            .get(&(ctx_key, g, name))
2501        {
2502            return Ok(f as cu::CUfunction);
2503        }
2504        let module = {
2505            let mut mods = MODS.lock().unwrap();
2506            let map = mods.get_or_insert_with(Default::default);
2507            match map.get(&(ctx_key, g)) {
2508                Some(&m) => m,
2509                None => {
2510                    let m = self.pdl_load_module_in_ctx(if g {
2511                        FLASH_FATBIN_KF8VF8
2512                    } else {
2513                        FLASH_FATBIN
2514                    })?;
2515                    map.insert((ctx_key, g), m);
2516                    m
2517                }
2518            }
2519        };
2520        let cname = std::ffi::CString::new(name)?;
2521        let mut f: cu::CUfunction = std::ptr::null_mut();
2522        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2523        if r != cu::CUresult::CUDA_SUCCESS {
2524            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2525        }
2526        FNS.lock()
2527            .unwrap()
2528            .get_or_insert_with(Default::default)
2529            .insert((ctx_key, g, name), f as usize);
2530        Ok(f)
2531    }
2532
2533    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2534    /// the module to the thread's CURRENT context — a remote-stage engine must not
2535    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2536    /// current context before returning.
2537    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2538        use cudarc::driver::sys as cu;
2539        let mut prev: cu::CUcontext = std::ptr::null_mut();
2540        unsafe {
2541            cu::cuCtxGetCurrent(&mut prev).result()?;
2542        }
2543        self.ctx().bind_to_thread()?;
2544        let mut m: cu::CUmodule = std::ptr::null_mut();
2545        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2546        let restore = if prev.is_null() {
2547            cu::CUresult::CUDA_SUCCESS
2548        } else {
2549            unsafe { cu::cuCtxSetCurrent(prev) }
2550        };
2551        if r != cu::CUresult::CUDA_SUCCESS {
2552            return Err(format!("pdl module load: {r:?}").into());
2553        }
2554        if restore != cu::CUresult::CUDA_SUCCESS {
2555            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2556        }
2557        Ok(m as usize)
2558    }
2559
2560    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
2561    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
2562    pub fn raw_kernel_function(
2563        &self,
2564        name: &'static str,
2565    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2566        self.pdl_func(name)
2567    }
2568
2569    fn pdl_func(
2570        &self,
2571        name: &'static str,
2572    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2573        use cudarc::driver::sys as cu;
2574        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2575        // are context-scoped; key everything by this engine's CUcontext).
2576        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2577            std::sync::Mutex::new(None);
2578        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2579        // duplicate module, loaded lazily on the first kernels-module miss.
2580        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2581            std::sync::Mutex::new(None);
2582        static FNS: std::sync::Mutex<
2583            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2584        > = std::sync::Mutex::new(None);
2585        let ctx_key = self.ctx().cu_ctx() as usize;
2586        if let Some(&f) = FNS
2587            .lock()
2588            .unwrap()
2589            .get_or_insert_with(Default::default)
2590            .get(&(ctx_key, name))
2591        {
2592            return Ok(f as cu::CUfunction);
2593        }
2594        let module = {
2595            let mut mods = MODULES.lock().unwrap();
2596            let map = mods.get_or_insert_with(Default::default);
2597            match map.get(&ctx_key) {
2598                Some(&m) => m,
2599                None => {
2600                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2601                    map.insert(ctx_key, m);
2602                    m
2603                }
2604            }
2605        };
2606        let cname = std::ffi::CString::new(name)?;
2607        let mut f: cu::CUfunction = std::ptr::null_mut();
2608        let mut r =
2609            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2610        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2611            let qmodule = {
2612                let mut mods = QMODULES.lock().unwrap();
2613                let map = mods.get_or_insert_with(Default::default);
2614                match map.get(&ctx_key) {
2615                    Some(&m) => m,
2616                    None => {
2617                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2618                        map.insert(ctx_key, m);
2619                        m
2620                    }
2621                }
2622            };
2623            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2624        }
2625        if r != cu::CUresult::CUDA_SUCCESS {
2626            return Err(format!("pdl_func {name}: {r:?}").into());
2627        }
2628        FNS.lock()
2629            .unwrap()
2630            .get_or_insert_with(Default::default)
2631            .insert((ctx_key, name), f as usize);
2632        Ok(f)
2633    }
2634
2635    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2636    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2637    ///
2638    /// # Safety
2639    /// `params` must match the kernel's exact parameter list (order, types, count) —
2640    /// a mismatch corrupts the launch silently.
2641    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2642    /// builder path's fa_func/func_g choice exactly).
2643    ///
2644    /// # Safety
2645    /// Same contract as `launch_pdl`.
2646    unsafe fn launch_pdl_flash(
2647        &self,
2648        g: bool,
2649        name: &'static str,
2650        grid: (u32, u32, u32),
2651        block: (u32, u32, u32),
2652        smem: u32,
2653        params: &mut [*mut std::ffi::c_void],
2654    ) -> Result<(), Box<dyn std::error::Error>> {
2655        use cudarc::driver::sys as cu;
2656        let f = self.pdl_func_flash(g, name)?;
2657        if smem > 0 {
2658            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2659            let r =
2660                unsafe {
2661                    cu::cuFuncSetAttribute(f,
2662                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2663                smem as i32)
2664                };
2665            if r != cu::CUresult::CUDA_SUCCESS {
2666                return Err(format!("pdl smem attr {name}: {r:?}").into());
2667            }
2668        }
2669        let mut attr = cu::CUlaunchAttribute {
2670            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2671            pad: [0; 4],
2672            value: cu::CUlaunchAttributeValue {
2673                programmaticStreamSerializationAllowed: 1,
2674            },
2675        };
2676        let cfg = cu::CUlaunchConfig {
2677            gridDimX: grid.0,
2678            gridDimY: grid.1,
2679            gridDimZ: grid.2,
2680            blockDimX: block.0,
2681            blockDimY: block.1,
2682            blockDimZ: block.2,
2683            sharedMemBytes: smem,
2684            hStream: self.gpu.stream().cu_stream(),
2685            attrs: &mut attr,
2686            numAttrs: 1,
2687        };
2688        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2689        if r != cu::CUresult::CUDA_SUCCESS {
2690            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2691        }
2692        Ok(())
2693    }
2694
2695    unsafe fn launch_pdl(
2696        &self,
2697        name: &'static str,
2698        grid: (u32, u32, u32),
2699        block: (u32, u32, u32),
2700        params: &mut [*mut std::ffi::c_void],
2701    ) -> Result<(), Box<dyn std::error::Error>> {
2702        use cudarc::driver::sys as cu;
2703        let f = self.pdl_func(name)?;
2704        let mut attr = cu::CUlaunchAttribute {
2705            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2706            pad: [0; 4],
2707            value: cu::CUlaunchAttributeValue {
2708                programmaticStreamSerializationAllowed: 1,
2709            },
2710        };
2711        let cfg = cu::CUlaunchConfig {
2712            gridDimX: grid.0,
2713            gridDimY: grid.1,
2714            gridDimZ: grid.2,
2715            blockDimX: block.0,
2716            blockDimY: block.1,
2717            blockDimZ: block.2,
2718            sharedMemBytes: 0,
2719            hStream: self.gpu.stream().cu_stream(),
2720            attrs: &mut attr,
2721            numAttrs: 1,
2722        };
2723        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2724        if r != cu::CUresult::CUDA_SUCCESS {
2725            return Err(format!("launch_pdl {name}: {r:?}").into());
2726        }
2727        Ok(())
2728    }
2729
2730    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2731    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2732    pub fn prefetch_weight_l2(
2733        &self,
2734        w: &crate::model::GpuTensor,
2735    ) -> Result<(), Box<dyn std::error::Error>> {
2736        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2737            let p = rp4.as_ref().unwrap_or(bytes);
2738            self.prefetch_l2(p, p.len())?;
2739        }
2740        Ok(())
2741    }
2742
2743    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2744    /// by the DEVICE token id at tok[idx] into f32.
2745    pub fn gather_row_bf16(
2746        &self,
2747        table: &CudaSlice<u8>,
2748        tok: &CudaSlice<u32>,
2749        idx: usize,
2750        dst: &mut CudaSlice<f32>,
2751        ncols: usize,
2752    ) -> Result<(), Box<dyn std::error::Error>> {
2753        let f = self.func("gather_row_bf16_f32");
2754        let cfg = LaunchConfig {
2755            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2756            block_dim: (256, 1, 1),
2757            shared_mem_bytes: 0,
2758        };
2759        let (nc, ix) = (ncols as i32, idx as i32);
2760        let __s_b = self.gpu.stream();
2761        let mut b = __s_b.launch_builder(&f);
2762        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2763        unsafe {
2764            b.launch(cfg)?;
2765        }
2766        Ok(())
2767    }
2768
2769    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2770    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2771    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2772    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2773    /// finish(1).
2774    #[allow(clippy::too_many_arguments)]
2775    pub fn dflash2_dynconv(
2776        &self,
2777        x: &CudaSlice<f32>,
2778        dyn_: &CudaSlice<f32>,
2779        base: &CudaSlice<f32>,
2780        out: &mut CudaSlice<f32>,
2781        rows: usize,
2782        hidden: usize,
2783        group_size: usize,
2784        ksize: usize,
2785        half: usize,
2786    ) -> Result<(), Box<dyn std::error::Error>> {
2787        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2788        let f = self.func("dflash2_dynconv_f32");
2789        let n = rows * hidden;
2790        let cfg = LaunchConfig {
2791            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2792            block_dim: (256, 1, 1),
2793            shared_mem_bytes: 0,
2794        };
2795        let (ri, hi, gi, ki, hf) = (
2796            rows as i32,
2797            hidden as i32,
2798            group_size as i32,
2799            ksize as i32,
2800            half as i32,
2801        );
2802        let __s_b = self.gpu.stream();
2803        let mut b = __s_b.launch_builder(&f);
2804        b.arg(x)
2805            .arg(dyn_)
2806            .arg(base)
2807            .arg(out)
2808            .arg(&ri)
2809            .arg(&hi)
2810            .arg(&gi)
2811            .arg(&ki)
2812            .arg(&hf);
2813        unsafe {
2814            b.launch(cfg)?;
2815        }
2816        Ok(())
2817    }
2818
2819    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2820    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2821    /// value-descending, ties to the lower index.
2822    pub fn topk_rows(
2823        &self,
2824        logits: &CudaSlice<f32>,
2825        n_rows: usize,
2826        n_cols: usize,
2827        k: usize,
2828    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2829        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2830        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2831        let f = self.func("topk_rows_f32");
2832        let nth = 256usize;
2833        let mut vals = self.uninit(n_rows * k)?;
2834        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2835        let cfg = LaunchConfig {
2836            grid_dim: (n_rows as u32, 1, 1),
2837            block_dim: (nth as u32, 1, 1),
2838            shared_mem_bytes: (nth * k * 8) as u32,
2839        };
2840        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2841        let __s_b = self.gpu.stream();
2842        let mut b = __s_b.launch_builder(&f);
2843        b.arg(logits)
2844            .arg(&nr)
2845            .arg(&nc)
2846            .arg(&ki)
2847            .arg(&mut vals)
2848            .arg(&mut idxs);
2849        unsafe {
2850            b.launch(cfg)?;
2851        }
2852        Ok((vals, idxs))
2853    }
2854
2855    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2856    pub fn add_row_inplace(
2857        &self,
2858        logits: &mut CudaSlice<f32>,
2859        bias: &CudaSlice<f32>,
2860        n: usize,
2861        row_off: usize,
2862    ) -> Result<(), Box<dyn std::error::Error>> {
2863        let f = self.func("add_row_inplace_f32");
2864        let cfg = LaunchConfig {
2865            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2866            block_dim: (256, 1, 1),
2867            shared_mem_bytes: 0,
2868        };
2869        let (ni, off) = (n as i32, row_off as i64);
2870        let __s_b = self.gpu.stream();
2871        let mut b = __s_b.launch_builder(&f);
2872        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2873        unsafe {
2874            b.launch(cfg)?;
2875        }
2876        Ok(())
2877    }
2878
2879    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2880    pub fn prefetch_l2(
2881        &self,
2882        p: &CudaSlice<u8>,
2883        n: usize,
2884    ) -> Result<(), Box<dyn std::error::Error>> {
2885        let f = self.func("prefetch_l2_bytes");
2886        let lines = n.div_ceil(128);
2887        let ni = n as i64;
2888        let cfg = LaunchConfig {
2889            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2890            block_dim: (256, 1, 1),
2891            shared_mem_bytes: 0,
2892        };
2893        let __s_b = self.gpu.stream();
2894        let mut b = __s_b.launch_builder(&f);
2895        b.arg(p).arg(&ni);
2896        unsafe {
2897            b.launch(cfg)?;
2898        }
2899        Ok(())
2900    }
2901
2902    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2903    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2904    pub fn router_gemv(
2905        &self,
2906        w: &CudaSlice<f32>,
2907        x: &CudaSlice<f32>,
2908        n_embd: usize,
2909        n_experts: usize,
2910        t: usize,
2911    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2912        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2913        // stream differs) — too small to justify a numeric config change; deleted.
2914        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2915        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2916        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2917        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2918            Ok("0") => false,
2919            Ok(_) => true,
2920            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2921        };
2922        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2923        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2924        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2925        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2926        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2927        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2928        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2929        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2930        // (perf-only, bits equal).
2931        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2932        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2933    }
2934
2935    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2936    /// force both forms; `batch` requires `w8`).
2937    pub fn router_gemv_form(
2938        &self,
2939        w: &CudaSlice<f32>,
2940        x: &CudaSlice<f32>,
2941        n_embd: usize,
2942        n_experts: usize,
2943        t: usize,
2944        w8: bool,
2945        batch: bool,
2946    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2947        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2948        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2949        let f = if batch {
2950            self.func("router_gemv_f32_w8_batch")
2951        } else if w8 {
2952            self.func("router_gemv_f32_w8")
2953        } else {
2954            self.func("router_gemv_f32")
2955        };
2956        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2957        let cfg = if batch {
2958            LaunchConfig {
2959                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2960                block_dim: (32, 8, 1),
2961                shared_mem_bytes: 0,
2962            }
2963        } else {
2964            LaunchConfig {
2965                grid_dim: (n_experts as u32, t as u32, 1),
2966                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2967                shared_mem_bytes: 0,
2968            }
2969        };
2970        let __s_b = self.gpu.stream();
2971        let mut b = __s_b.launch_builder(&f);
2972        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2973        unsafe {
2974            b.launch(cfg)?;
2975        }
2976        Ok(y)
2977    }
2978
2979    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
2980    /// buffer — token-graph alloc-free.
2981    pub fn router_gemv_into(
2982        &self,
2983        w: &CudaSlice<f32>,
2984        x: &CudaSlice<f32>,
2985        y: &mut CudaSlice<f32>,
2986        n_embd: usize,
2987        n_experts: usize,
2988        t: usize,
2989    ) -> Result<(), Box<dyn std::error::Error>> {
2990        if y.len() < t * n_experts {
2991            return Err("router_gemv_into output too small".into());
2992        }
2993        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2994            Ok("0") => false,
2995            Ok(_) => true,
2996            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2997        };
2998        let f = if w8 {
2999            self.func("router_gemv_f32_w8")
3000        } else {
3001            self.func("router_gemv_f32")
3002        };
3003        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3004        let cfg = LaunchConfig {
3005            grid_dim: (n_experts as u32, t as u32, 1),
3006            block_dim: (32, if w8 { 8 } else { 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(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
3012        unsafe {
3013            b.launch(cfg)?;
3014        }
3015        Ok(())
3016    }
3017
3018    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
3019    pub fn rows_permute(
3020        &self,
3021        src: &CudaSlice<f32>,
3022        idx: &CudaSlice<i32>,
3023        nrows: usize,
3024        ncols: usize,
3025    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3026        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3027        let f = self.func("rows_permute_f32");
3028        let (nc, nr) = (ncols as i32, nrows as i32);
3029        let cfg = LaunchConfig {
3030            grid_dim: (nrows as u32, 1, 1),
3031            block_dim: (256, 1, 1),
3032            shared_mem_bytes: 0,
3033        };
3034        let __s_b = self.gpu.stream();
3035        let mut b = __s_b.launch_builder(&f);
3036        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3037        unsafe {
3038            b.launch(cfg)?;
3039        }
3040        Ok(dst)
3041    }
3042
3043    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
3044    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
3045    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
3046    /// decode chain and the small-t spec-verify chain match per row by construction.
3047    pub fn sigmoid_dot_rows(
3048        &self,
3049        x: &CudaSlice<f32>,
3050        w: &CudaSlice<f32>,
3051        n_embd: usize,
3052        t: usize,
3053    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3054        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
3055        // config; same class as MEMRA_ROUTER_V2).
3056        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3057        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3058            let gs = self.linear(x, w, t, n_embd, 1)?;
3059            let mut g = self.uninit(t)?;
3060            self.sigmoid(&gs, &mut g, t)?;
3061            return Ok(g);
3062        }
3063        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
3064        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
3065        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
3066        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
3067        // flags doctrine; this per-token form serves every t.
3068        let mut g = self.alloc_uninit::<f32>(t)?;
3069        let f = self.func("sigmoid_dot_rows_f32");
3070        let (ne, ti) = (n_embd as i32, t as i32);
3071        let cfg = LaunchConfig {
3072            grid_dim: (t as u32, 1, 1),
3073            block_dim: (32, 8, 1),
3074            shared_mem_bytes: 0,
3075        };
3076        let __s_b = self.gpu.stream();
3077        let mut b = __s_b.launch_builder(&f);
3078        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3079        unsafe {
3080            b.launch(cfg)?;
3081        }
3082        Ok(g)
3083    }
3084
3085    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
3086    pub fn sigmoid_dot_rows_into(
3087        &self,
3088        x: &CudaSlice<f32>,
3089        w: &CudaSlice<f32>,
3090        g: &mut CudaSlice<f32>,
3091        n_embd: usize,
3092        t: usize,
3093    ) -> Result<(), Box<dyn std::error::Error>> {
3094        if g.len() < t {
3095            return Err("sigmoid_dot_rows_into output too small".into());
3096        }
3097        let f = self.func("sigmoid_dot_rows_f32");
3098        let (ne, ti) = (n_embd as i32, t as i32);
3099        let cfg = LaunchConfig {
3100            grid_dim: (t as u32, 1, 1),
3101            block_dim: (32, 8, 1),
3102            shared_mem_bytes: 0,
3103        };
3104        let __s_b = self.gpu.stream();
3105        let mut b = __s_b.launch_builder(&f);
3106        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3107        unsafe {
3108            b.launch(cfg)?;
3109        }
3110        Ok(())
3111    }
3112
3113    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
3114    pub fn spec_rollback_stream(
3115        &self,
3116        len_ptrs: &CudaSlice<u64>,
3117        pos_start: &CudaSlice<i32>,
3118        acc: &CudaSlice<u32>,
3119        base: usize,
3120        n_rows: usize,
3121    ) -> Result<(), Box<dyn std::error::Error>> {
3122        let f = self.func("spec_rollback_stream");
3123        let (b, nr) = (base as i32, n_rows as i32);
3124        let cfg = LaunchConfig {
3125            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3126            block_dim: (64, 1, 1),
3127            shared_mem_bytes: 0,
3128        };
3129        let __s_bl = self.gpu.stream();
3130        let mut bl = __s_bl.launch_builder(&f);
3131        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3132        unsafe {
3133            bl.launch(cfg)?;
3134        }
3135        Ok(())
3136    }
3137
3138    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
3139    pub fn plain_tok_ring(
3140        &self,
3141        vam: &CudaSlice<u32>,
3142        pos_start: &CudaSlice<i32>,
3143        base: usize,
3144        ring: &mut CudaSlice<u32>,
3145    ) -> Result<(), Box<dyn std::error::Error>> {
3146        let f = self.func("plain_tok_ring");
3147        let (b, cap) = (base as i32, ring.len() as i32);
3148        let cfg = LaunchConfig {
3149            grid_dim: (1, 1, 1),
3150            block_dim: (32, 1, 1),
3151            shared_mem_bytes: 0,
3152        };
3153        let __s_bl = self.gpu.stream();
3154        let mut bl = __s_bl.launch_builder(&f);
3155        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3156        unsafe {
3157            bl.launch(cfg)?;
3158        }
3159        Ok(())
3160    }
3161
3162    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
3163    pub fn spec_ring_commit(
3164        &self,
3165        vtok: &CudaSlice<u32>,
3166        acc: &CudaSlice<u32>,
3167        brk: &CudaSlice<u32>,
3168        ring: &mut CudaSlice<u32>,
3169        pend: &mut CudaSlice<u32>,
3170    ) -> Result<(), Box<dyn std::error::Error>> {
3171        let f = self.func("spec_ring_commit");
3172        let cfg = LaunchConfig {
3173            grid_dim: (1, 1, 1),
3174            block_dim: (32, 1, 1),
3175            shared_mem_bytes: 0,
3176        };
3177        let __s_b = self.gpu.stream();
3178        let mut b = __s_b.launch_builder(&f);
3179        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3180        unsafe {
3181            b.launch(cfg)?;
3182        }
3183        Ok(())
3184    }
3185    pub fn i32_copy_add(
3186        &self,
3187        src: &CudaSlice<i32>,
3188        dst: &mut CudaSlice<i32>,
3189        delta: i32,
3190    ) -> Result<(), Box<dyn std::error::Error>> {
3191        let f = self.func("i32_copy_add");
3192        let cfg = LaunchConfig {
3193            grid_dim: (1, 1, 1),
3194            block_dim: (32, 1, 1),
3195            shared_mem_bytes: 0,
3196        };
3197        let __s_b = self.gpu.stream();
3198        let mut b = __s_b.launch_builder(&f);
3199        b.arg(src).arg(dst).arg(&delta);
3200        unsafe {
3201            b.launch(cfg)?;
3202        }
3203        Ok(())
3204    }
3205    pub fn u32_copy(
3206        &self,
3207        src: &CudaSlice<u32>,
3208        dst: &mut CudaSlice<u32>,
3209    ) -> Result<(), Box<dyn std::error::Error>> {
3210        let f = self.func("u32_copy");
3211        let cfg = LaunchConfig {
3212            grid_dim: (1, 1, 1),
3213            block_dim: (32, 1, 1),
3214            shared_mem_bytes: 0,
3215        };
3216        let __s_b = self.gpu.stream();
3217        let mut b = __s_b.launch_builder(&f);
3218        b.arg(src).arg(dst);
3219        unsafe {
3220            b.launch(cfg)?;
3221        }
3222        Ok(())
3223    }
3224
3225    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
3226    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
3227    /// caps acceptance exactly like drafting fewer tokens).
3228    pub fn spec_adapt_k(
3229        &self,
3230        acc: &CudaSlice<u32>,
3231        brk: &mut CudaSlice<u32>,
3232        floor: usize,
3233        cap: usize,
3234    ) -> Result<(), Box<dyn std::error::Error>> {
3235        let f = self.func("spec_adapt_k");
3236        let (fl, cp) = (floor as i32, cap as i32);
3237        let cfg = LaunchConfig {
3238            grid_dim: (1, 1, 1),
3239            block_dim: (32, 1, 1),
3240            shared_mem_bytes: 0,
3241        };
3242        let __s_b = self.gpu.stream();
3243        let mut b = __s_b.launch_builder(&f);
3244        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3245        unsafe {
3246            b.launch(cfg)?;
3247        }
3248        Ok(())
3249    }
3250
3251    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
3252    pub fn spec_accept_greedy_dc(
3253        &self,
3254        preds: &CudaSlice<u32>,
3255        vtok: &CudaSlice<u32>,
3256        last_pred: &CudaSlice<u32>,
3257        brk: &CudaSlice<u32>,
3258        out: &mut CudaSlice<u32>,
3259    ) -> Result<(), Box<dyn std::error::Error>> {
3260        let f = self.func("spec_accept_greedy_dc");
3261        let cfg = LaunchConfig {
3262            grid_dim: (1, 1, 1),
3263            block_dim: (32, 1, 1),
3264            shared_mem_bytes: 0,
3265        };
3266        let __s_b = self.gpu.stream();
3267        let mut b = __s_b.launch_builder(&f);
3268        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3269        unsafe {
3270            b.launch(cfg)?;
3271        }
3272        Ok(())
3273    }
3274
3275    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
3276    pub fn pos_iota(
3277        &self,
3278        pos0: &CudaSlice<i32>,
3279        out: &mut CudaSlice<i32>,
3280        t: usize,
3281    ) -> Result<(), Box<dyn std::error::Error>> {
3282        let f = self.func("pos_iota_i32");
3283        let ti = t as i32;
3284        let cfg = LaunchConfig {
3285            grid_dim: (1, 1, 1),
3286            block_dim: (t.max(1) as u32, 1, 1),
3287            shared_mem_bytes: 0,
3288        };
3289        let __s_b = self.gpu.stream();
3290        let mut b = __s_b.launch_builder(&f);
3291        b.arg(pos0).arg(out).arg(&ti);
3292        unsafe {
3293            b.launch(cfg)?;
3294        }
3295        Ok(())
3296    }
3297    #[allow(clippy::too_many_arguments)]
3298    pub fn append_kv_quantized_rows_dc(
3299        &self,
3300        k_rows: &CudaSlice<f32>,
3301        v_rows: &CudaSlice<f32>,
3302        kc: &mut CudaSlice<u8>,
3303        vc: &mut CudaSlice<u8>,
3304        t0_dev: &CudaSlice<i32>,
3305        t: usize,
3306        kv_dim_k: usize,
3307        kv_dim_v: usize,
3308        k_tok_bytes: usize,
3309        v_tok_bytes: usize,
3310        g: bool,
3311    ) -> Result<(), Box<dyn std::error::Error>> {
3312        let f = if g {
3313            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3314        } else {
3315            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3316        };
3317        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3318        let cfg = LaunchConfig {
3319            grid_dim: (nblk, t as u32, 1),
3320            block_dim: (32, 1, 1),
3321            shared_mem_bytes: 0,
3322        };
3323        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3324        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3325        let __s_b = self.gpu.stream();
3326        let mut b = __s_b.launch_builder(&f);
3327        b.arg(k_rows)
3328            .arg(v_rows)
3329            .arg(kc)
3330            .arg(vc)
3331            .arg(t0_dev)
3332            .arg(&kdk)
3333            .arg(&kdv)
3334            .arg(&ktb)
3335            .arg(&vtb);
3336        unsafe {
3337            b.launch(cfg)?;
3338        }
3339        Ok(())
3340    }
3341
3342    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
3343    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
3344    #[allow(clippy::too_many_arguments)]
3345    pub fn append_kv_quantized_row_dc_inc(
3346        &self,
3347        k_row: &CudaSlice<f32>,
3348        v_row: &CudaSlice<f32>,
3349        kc: &mut CudaSlice<u8>,
3350        vc: &mut CudaSlice<u8>,
3351        t0_dev: &mut CudaSlice<i32>,
3352        kv_dim_k: usize,
3353        kv_dim_v: usize,
3354        k_tok_bytes: usize,
3355        v_tok_bytes: usize,
3356        g: bool,
3357    ) -> Result<(), Box<dyn std::error::Error>> {
3358        let f = if g {
3359            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3360        } else {
3361            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3362        };
3363        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3364        let cfg = LaunchConfig {
3365            grid_dim: (1, 1, 1),
3366            block_dim: (nthreads, 1, 1),
3367            shared_mem_bytes: 0,
3368        };
3369        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3370        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3371        let __s_b = self.gpu.stream();
3372        let mut b = __s_b.launch_builder(&f);
3373        b.arg(k_row)
3374            .arg(v_row)
3375            .arg(kc)
3376            .arg(vc)
3377            .arg(t0_dev)
3378            .arg(&kdk)
3379            .arg(&kdv)
3380            .arg(&ktb)
3381            .arg(&vtb);
3382        unsafe {
3383            b.launch(cfg)?;
3384        }
3385        Ok(())
3386    }
3387
3388    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
3389    pub fn pack_tok_p(
3390        &self,
3391        tok: &CudaSlice<u32>,
3392        p: &CudaSlice<f32>,
3393        out: &mut CudaSlice<u32>,
3394        slot: usize,
3395    ) -> Result<(), Box<dyn std::error::Error>> {
3396        let f = self.func("pack_tok_p");
3397        let sl = slot as i32;
3398        let cfg = LaunchConfig {
3399            grid_dim: (1, 1, 1),
3400            block_dim: (32, 1, 1),
3401            shared_mem_bytes: 0,
3402        };
3403        let __s_b = self.gpu.stream();
3404        let mut b = __s_b.launch_builder(&f);
3405        b.arg(tok).arg(p).arg(out).arg(&sl);
3406        unsafe {
3407            b.launch(cfg)?;
3408        }
3409        Ok(())
3410    }
3411    pub fn tok_map_u32(
3412        &self,
3413        tok: &mut CudaSlice<u32>,
3414        map: &CudaSlice<u32>,
3415    ) -> Result<(), Box<dyn std::error::Error>> {
3416        let f = self.func("tok_map_u32");
3417        let cfg = LaunchConfig {
3418            grid_dim: (1, 1, 1),
3419            block_dim: (32, 1, 1),
3420            shared_mem_bytes: 0,
3421        };
3422        let __s_b = self.gpu.stream();
3423        let mut b = __s_b.launch_builder(&f);
3424        b.arg(tok).arg(map);
3425        unsafe {
3426            b.launch(cfg)?;
3427        }
3428        Ok(())
3429    }
3430
3431    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
3432    #[allow(clippy::too_many_arguments)]
3433    pub fn spec_assemble_verify(
3434        &self,
3435        tokp: &CudaSlice<u32>,
3436        pend: &CudaSlice<u32>,
3437        d2t: Option<&CudaSlice<u32>>,
3438        vtok: &mut CudaSlice<u32>,
3439        brk: &mut CudaSlice<u32>,
3440        p_min: f32,
3441        k: usize,
3442        pmin0: bool,
3443    ) -> Result<(), Box<dyn std::error::Error>> {
3444        let f = self.func("spec_assemble_verify");
3445        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3446        let cfg = LaunchConfig {
3447            grid_dim: (1, 1, 1),
3448            block_dim: (32, 1, 1),
3449            shared_mem_bytes: 0,
3450        };
3451        let __s_b = self.gpu.stream();
3452        let mut b = __s_b.launch_builder(&f);
3453        match d2t {
3454            Some(m) => {
3455                b.arg(tokp)
3456                    .arg(pend)
3457                    .arg(m)
3458                    .arg(vtok)
3459                    .arg(brk)
3460                    .arg(&p_min)
3461                    .arg(&ki)
3462                    .arg(&pm);
3463                unsafe {
3464                    b.launch(cfg)?;
3465                }
3466            }
3467            None => {
3468                let null: u64 = 0;
3469                b.arg(tokp)
3470                    .arg(pend)
3471                    .arg(&null)
3472                    .arg(vtok)
3473                    .arg(brk)
3474                    .arg(&p_min)
3475                    .arg(&ki)
3476                    .arg(&pm);
3477                unsafe {
3478                    b.launch(cfg)?;
3479                }
3480            }
3481        }
3482        Ok(())
3483    }
3484
3485    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
3486    #[allow(clippy::too_many_arguments)]
3487    pub fn ssm_conv_ring_rebuild_dc(
3488        &self,
3489        qkv_tm: &CudaSlice<f32>,
3490        ring_old: &CudaSlice<f32>,
3491        conv_state: &mut CudaSlice<f32>,
3492        conv_dim: usize,
3493        acc: &CudaSlice<u32>,
3494        base: usize,
3495        t_v: usize,
3496        d_conv: usize,
3497    ) -> Result<(), Box<dyn std::error::Error>> {
3498        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3499        let n = conv_dim * (d_conv - 1);
3500        let cfg = LaunchConfig::for_num_elems(n as u32);
3501        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3502        let __s_b = self.gpu.stream();
3503        let mut b = __s_b.launch_builder(&f);
3504        b.arg(qkv_tm)
3505            .arg(ring_old)
3506            .arg(conv_state)
3507            .arg(&cd)
3508            .arg(acc)
3509            .arg(&b0)
3510            .arg(&tv)
3511            .arg(&dc);
3512        unsafe {
3513            b.launch(cfg)?;
3514        }
3515        Ok(())
3516    }
3517    #[allow(clippy::too_many_arguments)]
3518    pub fn gdn_scan_s128_dc(
3519        &self,
3520        q: &CudaSlice<f32>,
3521        k: &CudaSlice<f32>,
3522        v: &CudaSlice<f32>,
3523        g: &CudaSlice<f32>,
3524        beta: &CudaSlice<f32>,
3525        state_in: &CudaSlice<f32>,
3526        state_out: &mut CudaSlice<f32>,
3527        o: &mut CudaSlice<f32>,
3528        n_head: usize,
3529        acc: &CudaSlice<u32>,
3530        base: usize,
3531        t_v: usize,
3532        scale: f32,
3533    ) -> Result<(), Box<dyn std::error::Error>> {
3534        let f = self.func("gdn_scan_s128_dc");
3535        const S_V: u32 = 128;
3536        const WARP: u32 = 32;
3537        const COLS_PER_BLOCK: u32 = 4;
3538        let cfg = LaunchConfig {
3539            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3540            block_dim: (WARP, COLS_PER_BLOCK, 1),
3541            shared_mem_bytes: 0,
3542        };
3543        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3544        let __s_b = self.gpu.stream();
3545        let mut b = __s_b.launch_builder(&f);
3546        b.arg(q)
3547            .arg(k)
3548            .arg(v)
3549            .arg(g)
3550            .arg(beta)
3551            .arg(state_in)
3552            .arg(state_out)
3553            .arg(o)
3554            .arg(&h)
3555            .arg(acc)
3556            .arg(&b0)
3557            .arg(&tv)
3558            .arg(&scale);
3559        unsafe {
3560            b.launch(cfg)?;
3561        }
3562        Ok(())
3563    }
3564
3565    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
3566    pub fn spec_rollback_kv(
3567        &self,
3568        len_ptrs: &CudaSlice<u64>,
3569        saved: &CudaSlice<i32>,
3570        acc: &CudaSlice<u32>,
3571        base: usize,
3572        n_layer: usize,
3573    ) -> Result<(), Box<dyn std::error::Error>> {
3574        let f = self.func("spec_rollback_kv");
3575        let (b, nl) = (base as i32, n_layer as i32);
3576        let cfg = LaunchConfig {
3577            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3578            block_dim: (64, 1, 1),
3579            shared_mem_bytes: 0,
3580        };
3581        let __s_bl = self.gpu.stream();
3582        let mut bl = __s_bl.launch_builder(&f);
3583        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3584        unsafe {
3585            bl.launch(cfg)?;
3586        }
3587        Ok(())
3588    }
3589
3590    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3591    pub fn spec_fork_valid(
3592        &self,
3593        acc: &CudaSlice<u32>,
3594        optimistic_pending: u32,
3595        valid: &mut CudaSlice<u32>,
3596    ) -> Result<(), Box<dyn std::error::Error>> {
3597        let f = self.func("spec_fork_valid");
3598        let cfg = LaunchConfig {
3599            grid_dim: (1, 1, 1),
3600            block_dim: (1, 1, 1),
3601            shared_mem_bytes: 0,
3602        };
3603        let __s_bl = self.gpu.stream();
3604        let mut bl = __s_bl.launch_builder(&f);
3605        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3606        unsafe {
3607            bl.launch(cfg)?;
3608        }
3609        Ok(())
3610    }
3611
3612    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3613    pub fn spec_fork_reconcile_kv(
3614        &self,
3615        len_ptrs: &CudaSlice<u64>,
3616        saved: &CudaSlice<i32>,
3617        acc: &CudaSlice<u32>,
3618        valid: &CudaSlice<u32>,
3619        base: usize,
3620        n_layer: usize,
3621    ) -> Result<(), Box<dyn std::error::Error>> {
3622        let f = self.func("spec_fork_reconcile_kv");
3623        let (b, nl) = (base as i32, n_layer as i32);
3624        let cfg = LaunchConfig {
3625            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3626            block_dim: (64, 1, 1),
3627            shared_mem_bytes: 0,
3628        };
3629        let __s_bl = self.gpu.stream();
3630        let mut bl = __s_bl.launch_builder(&f);
3631        bl.arg(len_ptrs)
3632            .arg(saved)
3633            .arg(acc)
3634            .arg(valid)
3635            .arg(&b)
3636            .arg(&nl);
3637        unsafe {
3638            bl.launch(cfg)?;
3639        }
3640        Ok(())
3641    }
3642
3643    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3644    pub fn spec_fork_restore_f32(
3645        &self,
3646        snapshot: &CudaSlice<f32>,
3647        state: &mut CudaSlice<f32>,
3648        valid: &CudaSlice<u32>,
3649    ) -> Result<(), Box<dyn std::error::Error>> {
3650        assert_eq!(
3651            snapshot.len(),
3652            state.len(),
3653            "fork recurrent snapshot shape mismatch"
3654        );
3655        let f = self.func("spec_fork_restore_f32");
3656        let n = state.len() as i32;
3657        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3658        let cfg = LaunchConfig {
3659            grid_dim: (blocks, 1, 1),
3660            block_dim: (256, 1, 1),
3661            shared_mem_bytes: 0,
3662        };
3663        let __s_bl = self.gpu.stream();
3664        let mut bl = __s_bl.launch_builder(&f);
3665        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3666        unsafe {
3667            bl.launch(cfg)?;
3668        }
3669        Ok(())
3670    }
3671
3672    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3673    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3674    pub fn spec_seed_gather(
3675        &self,
3676        vx: &CudaSlice<f32>,
3677        fill_prev: &CudaSlice<f32>,
3678        acc: &CudaSlice<u32>,
3679        h_seed: &mut CudaSlice<f32>,
3680        base: usize,
3681        n_embd: usize,
3682    ) -> Result<(), Box<dyn std::error::Error>> {
3683        let f = self.func("spec_seed_gather");
3684        let (b, ne) = (base as i32, n_embd as i32);
3685        let cfg = LaunchConfig {
3686            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3687            block_dim: (256, 1, 1),
3688            shared_mem_bytes: 0,
3689        };
3690        let __s_bl = self.gpu.stream();
3691        let mut bl = __s_bl.launch_builder(&f);
3692        bl.arg(vx)
3693            .arg(fill_prev)
3694            .arg(acc)
3695            .arg(h_seed)
3696            .arg(&b)
3697            .arg(&ne);
3698        unsafe {
3699            bl.launch(cfg)?;
3700        }
3701        Ok(())
3702    }
3703
3704    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3705    pub fn spec_accept_greedy(
3706        &self,
3707        preds: &CudaSlice<u32>,
3708        draft: &CudaSlice<u32>,
3709        last_pred: u32,
3710        base: usize,
3711        k_round: usize,
3712        out: &mut CudaSlice<u32>,
3713    ) -> Result<(), Box<dyn std::error::Error>> {
3714        let f = self.func("spec_accept_greedy");
3715        let (b, k) = (base as i32, k_round as i32);
3716        let cfg = LaunchConfig {
3717            grid_dim: (1, 1, 1),
3718            block_dim: (32, 1, 1),
3719            shared_mem_bytes: 0,
3720        };
3721        let __s_bl = self.gpu.stream();
3722        let mut bl = __s_bl.launch_builder(&f);
3723        bl.arg(preds)
3724            .arg(draft)
3725            .arg(&last_pred)
3726            .arg(&b)
3727            .arg(&k)
3728            .arg(out);
3729        unsafe {
3730            bl.launch(cfg)?;
3731        }
3732        Ok(())
3733    }
3734
3735    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3736    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3737    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3738
3739    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3740    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3741    pub fn gumbel_perturb(
3742        &self,
3743        x: &CudaSlice<f32>,
3744        y: &mut CudaSlice<f32>,
3745        n: usize,
3746        seed: u64,
3747        stream_pos: u32,
3748        temp: f32,
3749    ) -> Result<(), Box<dyn std::error::Error>> {
3750        let f = self.func("gumbel_perturb_f32");
3751        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3752        let cfg = LaunchConfig {
3753            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3754            block_dim: (256, 1, 1),
3755            shared_mem_bytes: 0,
3756        };
3757        let __s_b = self.gpu.stream();
3758        let mut b = __s_b.launch_builder(&f);
3759        b.arg(x)
3760            .arg(&mut *y)
3761            .arg(&ni)
3762            .arg(&slo)
3763            .arg(&shi)
3764            .arg(&stream_pos)
3765            .arg(&temp);
3766        unsafe {
3767            b.launch(cfg)?;
3768        }
3769        Ok(())
3770    }
3771
3772    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3773    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3774    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3775    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3776    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3777    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3778    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3779    pub fn mask_logits_col(
3780        &self,
3781        logits: &mut CudaSlice<f32>,
3782        mask: &CudaSlice<u32>,
3783        col: usize,
3784        n: usize,
3785        mask_words: usize,
3786    ) -> Result<(), Box<dyn std::error::Error>> {
3787        let f = self.func("mask_logits_f32");
3788        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3789        let cfg = LaunchConfig {
3790            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3791            block_dim: (256, 1, 1),
3792            shared_mem_bytes: 0,
3793        };
3794        let __s_b = self.gpu.stream();
3795        let mut b = __s_b.launch_builder(&f);
3796        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3797        unsafe {
3798            b.launch(cfg)?;
3799        }
3800        Ok(())
3801    }
3802
3803    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3804    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3805    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3806    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3807    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3808    /// pointer-invariance IS the serving isolation contract for sampled rows.
3809    pub fn gumbel_perturb_col(
3810        &self,
3811        x: &CudaSlice<f32>,
3812        col: usize,
3813        y: &mut CudaSlice<f32>,
3814        n: usize,
3815        seed: u64,
3816        stream_pos: u32,
3817        temp: f32,
3818    ) -> Result<(), Box<dyn std::error::Error>> {
3819        let f = self.func("gumbel_perturb_f32");
3820        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3821        let col_view = x.slice(col * n..(col + 1) * n);
3822        let cfg = LaunchConfig {
3823            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3824            block_dim: (256, 1, 1),
3825            shared_mem_bytes: 0,
3826        };
3827        let __s_b = self.gpu.stream();
3828        let mut b = __s_b.launch_builder(&f);
3829        b.arg(&col_view)
3830            .arg(&mut *y)
3831            .arg(&ni)
3832            .arg(&slo)
3833            .arg(&shi)
3834            .arg(&stream_pos)
3835            .arg(&temp);
3836        unsafe {
3837            b.launch(cfg)?;
3838        }
3839        Ok(())
3840    }
3841
3842    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3843    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3844    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3845    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3846    /// the serving isolation contract for sampled rows).
3847    #[allow(clippy::too_many_arguments)]
3848    pub fn gumbel_perturb_filtered_col(
3849        &self,
3850        x: &CudaSlice<f32>,
3851        col: usize,
3852        y: &mut CudaSlice<f32>,
3853        n: usize,
3854        seed: u64,
3855        stream_pos: u32,
3856        temp: f32,
3857        stat_max: &CudaSlice<f32>,
3858        stat_th: &CudaSlice<f32>,
3859        stat_idx: usize,
3860    ) -> Result<(), Box<dyn std::error::Error>> {
3861        let f = self.func("gumbel_perturb_filtered_col_f32");
3862        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3863        let (ci, si) = (col as i32, stat_idx as i32);
3864        let cfg = LaunchConfig {
3865            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3866            block_dim: (256, 1, 1),
3867            shared_mem_bytes: 0,
3868        };
3869        let __s_b = self.gpu.stream();
3870        let mut b = __s_b.launch_builder(&f);
3871        b.arg(x)
3872            .arg(&ci)
3873            .arg(&mut *y)
3874            .arg(&ni)
3875            .arg(&slo)
3876            .arg(&shi)
3877            .arg(&stream_pos)
3878            .arg(&temp)
3879            .arg(stat_max)
3880            .arg(stat_th)
3881            .arg(&si);
3882        unsafe {
3883            b.launch(cfg)?;
3884        }
3885        Ok(())
3886    }
3887
3888    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3889    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3890    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3891    /// reads it (counter is data, not state — graph-replay-safe).
3892    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3893        let f = self.func("memra_sctr_inc");
3894        let cfg = LaunchConfig {
3895            grid_dim: (1, 1, 1),
3896            block_dim: (1, 1, 1),
3897            shared_mem_bytes: 0,
3898        };
3899        let __s_b = self.gpu.stream();
3900        let mut b = __s_b.launch_builder(&f);
3901        b.arg(&mut *ctr);
3902        unsafe {
3903            b.launch(cfg)?;
3904        }
3905        Ok(())
3906    }
3907
3908    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3909    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3910    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3911    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3912    pub fn gumbel_perturb_ctr(
3913        &self,
3914        x: &CudaSlice<f32>,
3915        y: &mut CudaSlice<f32>,
3916        n: usize,
3917        seed: u64,
3918        ctr: &CudaSlice<u32>,
3919        temp: f32,
3920    ) -> Result<(), Box<dyn std::error::Error>> {
3921        let f = self.func("gumbel_perturb_ctr_f32");
3922        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3923        let cfg = LaunchConfig {
3924            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3925            block_dim: (256, 1, 1),
3926            shared_mem_bytes: 0,
3927        };
3928        let __s_b = self.gpu.stream();
3929        let mut b = __s_b.launch_builder(&f);
3930        b.arg(x)
3931            .arg(&mut *y)
3932            .arg(&ni)
3933            .arg(&slo)
3934            .arg(&shi)
3935            .arg(ctr)
3936            .arg(&temp);
3937        unsafe {
3938            b.launch(cfg)?;
3939        }
3940        Ok(())
3941    }
3942
3943    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3944    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3945    /// (smallest-index tie-break — matches the argmax-gate contract).
3946    pub fn softmax_gather(
3947        &self,
3948        x: &CudaSlice<f32>,
3949        row_stride: usize,
3950        ids: &CudaSlice<u32>,
3951        rows: &CudaSlice<i32>,
3952        out: &mut CudaSlice<f32>,
3953        n: usize,
3954        npair: usize,
3955        temp: f32,
3956    ) -> Result<(), Box<dyn std::error::Error>> {
3957        let f = self.func("softmax_gather_f32");
3958        let (ni, rs) = (n as i32, row_stride as i64);
3959        let np = npair as i32;
3960        let cfg = LaunchConfig {
3961            grid_dim: (npair as u32, 1, 1),
3962            block_dim: (256, 1, 1),
3963            shared_mem_bytes: 0,
3964        };
3965        let __s_b = self.gpu.stream();
3966        let mut b = __s_b.launch_builder(&f);
3967        b.arg(x)
3968            .arg(&rs)
3969            .arg(ids)
3970            .arg(rows)
3971            .arg(&mut *out)
3972            .arg(&ni)
3973            .arg(&np)
3974            .arg(&temp);
3975        unsafe {
3976            b.launch(cfg)?;
3977        }
3978        Ok(())
3979    }
3980
3981    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3982    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3983    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3984    pub fn residual_sample(
3985        &self,
3986        p: &CudaSlice<f32>,
3987        q: Option<&CudaSlice<f32>>,
3988        n: usize,
3989        temp: f32,
3990        seed: u64,
3991        stream_pos: u32,
3992        out_tok: &mut CudaSlice<u32>,
3993    ) -> Result<(), Box<dyn std::error::Error>> {
3994        let f = self.func("residual_sample_f32");
3995        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3996        let nth = 1024u32;
3997        let cfg = LaunchConfig {
3998            grid_dim: (1, 1, 1),
3999            block_dim: (nth, 1, 1),
4000            shared_mem_bytes: 0,
4001        };
4002        let has_q: i32 = q.is_some() as i32;
4003        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
4004        let __s_b = self.gpu.stream();
4005        let mut b = __s_b.launch_builder(&f);
4006        b.arg(p)
4007            .arg(qbuf)
4008            .arg(&has_q)
4009            .arg(&ni)
4010            .arg(&temp)
4011            .arg(&slo)
4012            .arg(&shi)
4013            .arg(&stream_pos)
4014            .arg(&mut *out_tok);
4015        unsafe {
4016            b.launch(cfg)?;
4017        }
4018        Ok(())
4019    }
4020
4021    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
4022    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
4023    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
4024    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
4025    pub fn with_moe_cache<R>(
4026        &self,
4027        max_block_bytes: usize,
4028        f: impl FnOnce(
4029            &mut crate::moe_cache::MoeSlotCache,
4030            &Engine,
4031        ) -> Result<R, Box<dyn std::error::Error>>,
4032    ) -> Result<R, Box<dyn std::error::Error>> {
4033        let mut guard = self.moe_cache.lock().unwrap();
4034        if guard.is_none() {
4035            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4036        }
4037        let cache = guard.as_mut().unwrap();
4038        f(cache, self)
4039    }
4040
4041    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
4042    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
4043    pub fn freeze_moe_cache(&self) {
4044        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4045            cache.freeze();
4046        }
4047    }
4048
4049    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
4050    /// Never constructs a cache.
4051    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4052        self.moe_cache
4053            .lock()
4054            .unwrap()
4055            .as_ref()
4056            .map(crate::moe_cache::MoeSlotCache::export_residency)
4057    }
4058
4059    pub(crate) fn moe_cache_frozen(&self) -> bool {
4060        self.moe_cache
4061            .lock()
4062            .unwrap()
4063            .as_ref()
4064            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4065    }
4066
4067    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
4068    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
4069    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
4070    /// while leaving the profiling warmup's established batched behavior untouched.
4071    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
4072    /// tokenwise arm anyway.)
4073    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4074        crate::cpu_experts::configured()
4075            && self.moe_cache_frozen()
4076            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4077    }
4078
4079    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
4080    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4081        assert!(
4082            self.moe_cache.lock().unwrap().is_none(),
4083            "MoE cache layout configured after cache construction"
4084        );
4085        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4086    }
4087
4088    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4089        self.moe_cache_layout.lock().unwrap().clone()
4090    }
4091
4092    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
4093    pub fn moe_cache_enabled() -> bool {
4094        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4095    }
4096
4097    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
4098    /// Returns None if the cache was never built (disabled or no MoE forward ran).
4099    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4100        let guard = self.moe_cache.lock().unwrap();
4101        guard
4102            .as_ref()
4103            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4104    }
4105
4106    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
4107    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
4108    /// callers compare a before/after snapshot around a decode window.
4109    pub fn cpu_expert_stats(
4110        &self,
4111    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4112        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4113    }
4114
4115    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
4116    /// the backend tail that resident-GPU expert work did not hide.
4117    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4118        crate::cpu_experts::predictor_stats()
4119    }
4120
4121    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4122        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4123    }
4124
4125    /// CPU-routed expert selections grouped by how many of their three projections were already
4126    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
4127    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4128        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4129    }
4130
4131    /// Positioned-read proof-backend counters:
4132    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
4133    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4134        let guard = self.moe_cache.lock().unwrap();
4135        guard
4136            .as_ref()
4137            .and_then(|cache| cache.pread_stats())
4138            .map(|stats| {
4139                (
4140                    stats.reads,
4141                    stats.bytes,
4142                    stats.read_errors,
4143                    stats.short_reads,
4144                    stats.fallbacks,
4145                    stats.buffer_waits,
4146                    stats.ring_full,
4147                )
4148            })
4149    }
4150
4151    /// Spill configuration values that warned and substituted their documented defaults.
4152    pub fn spill_config_fallbacks(&self) -> u64 {
4153        crate::spill_pread::config_fallbacks()
4154    }
4155
4156    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
4157    pub fn moe_cache_reset_counters(&self) {
4158        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4159            c.reset_counters();
4160        }
4161    }
4162
4163    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4164        Ok(self.gpu.stream().clone_htod(v)?)
4165    }
4166
4167    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
4168    /// past the final q4_0 block through their aligned window — the bytes never reach a
4169    /// result (funnelshift discards them) but must be mapped memory.
4170    pub fn htod_bytes_padded(
4171        &self,
4172        v: &[u8],
4173        pad: usize,
4174    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4175        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4176        {
4177            let mut view = d.slice_mut(0..v.len());
4178            self.gpu.stream().memcpy_htod(v, &mut view)?;
4179        }
4180        Ok(d)
4181    }
4182
4183    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
4184    pub fn copy_into(
4185        &self,
4186        dst: &mut CudaSlice<f32>,
4187        off: usize,
4188        src: &CudaSlice<f32>,
4189        len: usize,
4190    ) -> Result<(), Box<dyn std::error::Error>> {
4191        let mut view = dst.slice_mut(off..off + len);
4192        self.gpu
4193            .stream()
4194            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4195        Ok(())
4196    }
4197
4198    /// D2D copy with an offset on BOTH sides. `copy_into` always reads the source from 0,
4199    /// which cannot express "copy the TAIL of this buffer" — the shape a sliding-window draft
4200    /// KV export needs (lane/dspark-draft-plane-20260827).
4201    pub fn copy_range_into(
4202        &self,
4203        dst: &mut CudaSlice<f32>,
4204        dst_off: usize,
4205        src: &CudaSlice<f32>,
4206        src_off: usize,
4207        len: usize,
4208    ) -> Result<(), Box<dyn std::error::Error>> {
4209        let mut view = dst.slice_mut(dst_off..dst_off + len);
4210        self.gpu
4211            .stream()
4212            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
4213        Ok(())
4214    }
4215
4216    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
4217    /// u8 twin of copy_into (D2D byte-range copy at an offset).
4218    pub fn copy_u8_into(
4219        &self,
4220        dst: &mut CudaSlice<u8>,
4221        off: usize,
4222        src: &CudaSlice<u8>,
4223        len: usize,
4224    ) -> Result<(), Box<dyn std::error::Error>> {
4225        let mut view = dst.slice_mut(off..off + len);
4226        self.gpu
4227            .stream()
4228            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4229        Ok(())
4230    }
4231
4232    /// D2D byte-range copy with explicit source and destination offsets.
4233    pub fn copy_u8_range_into(
4234        &self,
4235        dst: &mut CudaSlice<u8>,
4236        dst_off: usize,
4237        src: &CudaSlice<u8>,
4238        src_off: usize,
4239        len: usize,
4240    ) -> Result<(), Box<dyn std::error::Error>> {
4241        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
4242        self.gpu
4243            .stream()
4244            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4245        Ok(())
4246    }
4247
4248    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
4249    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
4250    /// keeping the audited attention range contiguous without changing its absolute start.
4251    pub fn prepare_kv_append(
4252        &self,
4253        kv: &mut crate::cache::KvLayer,
4254        retain_from: usize,
4255        append_rows: usize,
4256    ) -> Result<usize, Box<dyn std::error::Error>> {
4257        let Some(plan) = kv
4258            .ring
4259            .as_ref()
4260            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4261            .transpose()?
4262        else {
4263            return Ok(kv.len);
4264        };
4265        match plan {
4266            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4267            crate::cache::KvRingAppend::Rebase {
4268                src_row,
4269                keep_rows,
4270                new_base,
4271                write_row,
4272            } => {
4273                if keep_rows > 0 {
4274                    let k_len = keep_rows * kv.k_tok_bytes;
4275                    let v_len = keep_rows * kv.v_tok_bytes;
4276                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4277                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4278                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4279                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4280                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4281                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4282                }
4283                kv.ring.as_mut().unwrap().apply_rebase(new_base);
4284                Ok(write_row)
4285            }
4286        }
4287    }
4288
4289    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
4290    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
4291    pub fn htod_u8_into(
4292        &self,
4293        dst: &mut CudaSlice<u8>,
4294        off: usize,
4295        src: &[u8],
4296    ) -> Result<(), Box<dyn std::error::Error>> {
4297        let mut view = dst.slice_mut(off..off + src.len());
4298        self.gpu.stream().memcpy_htod(src, &mut view)?;
4299        Ok(())
4300    }
4301
4302    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4303        b.slice(0..len)
4304    }
4305
4306    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
4307    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
4308    pub fn view_u8_range<'a>(
4309        &self,
4310        b: &'a CudaSlice<u8>,
4311        start: usize,
4312        end: usize,
4313    ) -> cudarc::driver::CudaView<'a, u8> {
4314        b.slice(start..end)
4315    }
4316    pub fn view_u8<'a>(
4317        &self,
4318        b: &'a CudaSlice<u8>,
4319        len: usize,
4320    ) -> cudarc::driver::CudaView<'a, u8> {
4321        b.slice(0..len)
4322    }
4323
4324    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
4325    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
4326    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
4327    pub fn append_kv_quantized(
4328        &self,
4329        k_row: &CudaSlice<f32>,
4330        v_row: &CudaSlice<f32>,
4331        kc: &mut CudaSlice<u8>,
4332        vc: &mut CudaSlice<u8>,
4333        t: usize,
4334        kv_dim_k: usize,
4335        kv_dim_v: usize,
4336        k_tok_bytes: usize,
4337        v_tok_bytes: usize,
4338        g: bool,
4339    ) -> Result<(), Box<dyn std::error::Error>> {
4340        let f = if g {
4341            self.func_g("append_quantize_kv_q8_0_q5_1")
4342        } else {
4343            self.func("append_quantize_kv_q8_0_q5_1")
4344        };
4345        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4346        let cfg = LaunchConfig {
4347            grid_dim: (nblk, 1, 1),
4348            block_dim: (32, 1, 1),
4349            shared_mem_bytes: 0,
4350        };
4351        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4352        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4353        let __s_b = self.gpu.stream();
4354        let mut b = __s_b.launch_builder(&f);
4355        b.arg(k_row)
4356            .arg(v_row)
4357            .arg(kc)
4358            .arg(vc)
4359            .arg(&ti)
4360            .arg(&kdk)
4361            .arg(&kdv)
4362            .arg(&ktb)
4363            .arg(&vtb);
4364        unsafe {
4365            b.launch(cfg)?;
4366        }
4367        Ok(())
4368    }
4369
4370    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
4371    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
4372    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
4373    pub fn append_kv_quantized_dc(
4374        &self,
4375        k_row: &CudaSlice<f32>,
4376        v_row: &CudaSlice<f32>,
4377        kc: &mut CudaSlice<u8>,
4378        vc: &mut CudaSlice<u8>,
4379        t_dev: &CudaSlice<i32>,
4380        kv_dim_k: usize,
4381        kv_dim_v: usize,
4382        k_tok_bytes: usize,
4383        v_tok_bytes: usize,
4384        g: bool,
4385    ) -> Result<(), Box<dyn std::error::Error>> {
4386        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4387        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4388        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4389        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
4390        if Self::pdl_on() && Self::pdl_wb_on() {
4391            use cudarc::driver::{DevicePtr, DevicePtrMut};
4392            let s = &self.gpu.stream();
4393            let (pk, _g0) = k_row.device_ptr(s);
4394            let (pv, _g1) = v_row.device_ptr(s);
4395            let (pkc, _g2) = kc.device_ptr_mut(s);
4396            let (pvc, _g3) = vc.device_ptr_mut(s);
4397            let (pt, _g4) = t_dev.device_ptr(s);
4398            let mut ps = [
4399                &pk as *const _ as *mut std::ffi::c_void,
4400                &pv as *const _ as *mut _,
4401                &pkc as *const _ as *mut _,
4402                &pvc as *const _ as *mut _,
4403                &pt as *const _ as *mut _,
4404                &kdk as *const _ as *mut _,
4405                &kdv as *const _ as *mut _,
4406                &ktb as *const _ as *mut _,
4407                &vtb as *const _ as *mut _,
4408            ];
4409            unsafe {
4410                self.launch_pdl_flash(
4411                    g,
4412                    "append_quantize_kv_q8_0_q5_1_dc",
4413                    (nblk, 1, 1),
4414                    (32, 1, 1),
4415                    0,
4416                    &mut ps,
4417                )?;
4418            }
4419            return Ok(());
4420        }
4421        let f = if g {
4422            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4423        } else {
4424            self.func("append_quantize_kv_q8_0_q5_1_dc")
4425        };
4426        let cfg = LaunchConfig {
4427            grid_dim: (nblk, 1, 1),
4428            block_dim: (32, 1, 1),
4429            shared_mem_bytes: 0,
4430        };
4431        let __s_b = self.gpu.stream();
4432        let mut b = __s_b.launch_builder(&f);
4433        b.arg(k_row)
4434            .arg(v_row)
4435            .arg(kc)
4436            .arg(vc)
4437            .arg(t_dev)
4438            .arg(&kdk)
4439            .arg(&kdv)
4440            .arg(&ktb)
4441            .arg(&vtb);
4442        unsafe {
4443            b.launch(cfg)?;
4444        }
4445        Ok(())
4446    }
4447
4448    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
4449    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
4450    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
4451    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
4452    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
4453    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
4454    #[allow(clippy::too_many_arguments)]
4455    pub fn append_kv_quantized_rows(
4456        &self,
4457        k_rows: &CudaSlice<f32>,
4458        v_rows: &CudaSlice<f32>,
4459        kc: &mut CudaSlice<u8>,
4460        vc: &mut CudaSlice<u8>,
4461        t0: usize,
4462        t: usize,
4463        kv_dim_k: usize,
4464        kv_dim_v: usize,
4465        k_tok_bytes: usize,
4466        v_tok_bytes: usize,
4467        g: bool,
4468    ) -> Result<(), Box<dyn std::error::Error>> {
4469        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4470            for i in 0..t {
4471                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4472                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4473                self.append_kv_quantized_view(
4474                    &k_row,
4475                    &v_row,
4476                    kc,
4477                    vc,
4478                    t0 + i,
4479                    kv_dim_k,
4480                    kv_dim_v,
4481                    k_tok_bytes,
4482                    v_tok_bytes,
4483                    g,
4484                )?;
4485            }
4486            return Ok(());
4487        }
4488        let f = if g {
4489            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4490        } else {
4491            self.func("append_quantize_kv_q8_0_q5_1_rows")
4492        };
4493        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4494        let cfg = LaunchConfig {
4495            grid_dim: (nblk, t as u32, 1),
4496            block_dim: (32, 1, 1),
4497            shared_mem_bytes: 0,
4498        };
4499        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4500        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4501        let __s_b = self.gpu.stream();
4502        let mut b = __s_b.launch_builder(&f);
4503        b.arg(k_rows)
4504            .arg(v_rows)
4505            .arg(kc)
4506            .arg(vc)
4507            .arg(&t0i)
4508            .arg(&kdk)
4509            .arg(&kdv)
4510            .arg(&ktb)
4511            .arg(&vtb);
4512        unsafe {
4513            b.launch(cfg)?;
4514        }
4515        Ok(())
4516    }
4517
4518    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
4519    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
4520    /// later, inside a captured graph) without a host round-trip.
4521    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4522        let f = self.func("inc_i32");
4523        let cfg = LaunchConfig {
4524            grid_dim: (1, 1, 1),
4525            block_dim: (1, 1, 1),
4526            shared_mem_bytes: 0,
4527        };
4528        let __s_b = self.gpu.stream();
4529        let mut b = __s_b.launch_builder(&f);
4530        b.arg(p);
4531        unsafe {
4532            b.launch(cfg)?;
4533        }
4534        Ok(())
4535    }
4536
4537    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
4538    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
4539    pub fn append_kv_quantized_view(
4540        &self,
4541        k_row: &cudarc::driver::CudaView<f32>,
4542        v_row: &cudarc::driver::CudaView<f32>,
4543        kc: &mut CudaSlice<u8>,
4544        vc: &mut CudaSlice<u8>,
4545        t: usize,
4546        kv_dim_k: usize,
4547        kv_dim_v: usize,
4548        k_tok_bytes: usize,
4549        v_tok_bytes: usize,
4550        g: bool,
4551    ) -> Result<(), Box<dyn std::error::Error>> {
4552        let stream = self.gpu.stream();
4553        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4554        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4555        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4556        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4557        let f = if g {
4558            self.func_g("append_quantize_kv_q8_0_q5_1")
4559        } else {
4560            self.func("append_quantize_kv_q8_0_q5_1")
4561        };
4562        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4563        let cfg = LaunchConfig {
4564            grid_dim: (nblk, 1, 1),
4565            block_dim: (32, 1, 1),
4566            shared_mem_bytes: 0,
4567        };
4568        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4569        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4570        let mut b = stream.launch_builder(&f);
4571        b.arg(k_row)
4572            .arg(v_row)
4573            .arg(kc)
4574            .arg(vc)
4575            .arg(&ti)
4576            .arg(&kdk)
4577            .arg(&kdv)
4578            .arg(&ktb)
4579            .arg(&vtb);
4580        unsafe {
4581            b.launch(cfg)?;
4582        }
4583        Ok(())
4584    }
4585
4586    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
4587    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
4588    pub fn copy_view_into(
4589        &self,
4590        dst: &mut CudaSlice<f32>,
4591        off: usize,
4592        src: &cudarc::driver::CudaView<f32>,
4593        len: usize,
4594    ) -> Result<(), Box<dyn std::error::Error>> {
4595        let mut view = dst.slice_mut(off..off + len);
4596        self.gpu
4597            .stream()
4598            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4599        Ok(())
4600    }
4601
4602    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
4603    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
4604    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
4605    pub fn clone_dtod(
4606        &self,
4607        src: &CudaSlice<f32>,
4608    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4609        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4610        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4611        Ok(dst)
4612    }
4613
4614    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4615    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4616    pub fn dtod_copy_view(
4617        &self,
4618        src: &cudarc::driver::CudaView<f32>,
4619        dst: &mut CudaSlice<f32>,
4620    ) -> Result<(), Box<dyn std::error::Error>> {
4621        self.gpu.stream().memcpy_dtod(src, dst)?;
4622        Ok(())
4623    }
4624
4625    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4626    pub fn dtod_copy_view_i8(
4627        &self,
4628        src: &cudarc::driver::CudaView<i8>,
4629        dst: &mut CudaSlice<i8>,
4630    ) -> Result<(), Box<dyn std::error::Error>> {
4631        self.gpu.stream().memcpy_dtod(src, dst)?;
4632        Ok(())
4633    }
4634
4635    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4636    pub fn dtod_copy_into(
4637        &self,
4638        src: &CudaSlice<f32>,
4639        dst: &mut CudaSlice<f32>,
4640        offset: usize,
4641    ) -> Result<(), Box<dyn std::error::Error>> {
4642        let n = src.len();
4643        let mut dv = dst.slice_mut(offset..offset + n);
4644        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4645        Ok(())
4646    }
4647
4648    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4649    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4650    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4651    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4652    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4653    pub fn copy_batch_uniform_f32(
4654        &self,
4655        table: &CudaSlice<u64>,
4656        n: usize,
4657        words: usize,
4658    ) -> Result<(), Box<dyn std::error::Error>> {
4659        if n == 0 || words == 0 {
4660            return Ok(());
4661        }
4662        debug_assert!(
4663            table.len() >= 2 * n,
4664            "pointer table must hold n srcs + n dsts"
4665        );
4666        let f = self.func("copy_batch_uniform_f32");
4667        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4668        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4669        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4670        let (ni, wi) = (n as i32, words as i32);
4671        let cfg = LaunchConfig {
4672            grid_dim: (chunks, n as u32, 1),
4673            block_dim: (256, 1, 1),
4674            shared_mem_bytes: 0,
4675        };
4676        let __s = self.gpu.stream();
4677        let mut b = __s.launch_builder(&f);
4678        b.arg(table).arg(&ni).arg(&wi);
4679        unsafe {
4680            b.launch(cfg)?;
4681        }
4682        Ok(())
4683    }
4684
4685    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4686    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4687    pub fn htod_u64_into(
4688        &self,
4689        v: &[u64],
4690        dst: &mut CudaSlice<u64>,
4691    ) -> Result<(), Box<dyn std::error::Error>> {
4692        let mut view = dst.slice_mut(0..v.len());
4693        self.gpu.stream().memcpy_htod(v, &mut view)?;
4694        Ok(())
4695    }
4696
4697    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4698    /// device pointer-table entry at run time, so a captured graph follows the gdn
4699    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4700    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4701    pub fn copy_indirect_src_f32(
4702        &self,
4703        src_entry: &cudarc::driver::CudaView<u64>,
4704        dst: &mut CudaSlice<f32>,
4705        dst_off: usize,
4706        words: usize,
4707    ) -> Result<(), Box<dyn std::error::Error>> {
4708        let f = self.func("copy_indirect_src_f32");
4709        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4710        let wi = words as i32;
4711        let cfg = LaunchConfig {
4712            grid_dim: (chunks, 1, 1),
4713            block_dim: (256, 1, 1),
4714            shared_mem_bytes: 0,
4715        };
4716        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4717        let __s = self.gpu.stream();
4718        let mut b = __s.launch_builder(&f);
4719        b.arg(src_entry).arg(&mut dv).arg(&wi);
4720        unsafe {
4721            b.launch(cfg)?;
4722        }
4723        Ok(())
4724    }
4725
4726    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4727    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4728        self.alloc_uninit::<i8>(n)
4729    }
4730
4731    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4732    pub fn qmatvec(
4733        &self,
4734        w: &CudaSlice<u8>,
4735        x: &CudaSlice<f32>,
4736        m: usize,
4737        in_f: usize,
4738        out_f: usize,
4739        qtype: i32,
4740        row_bytes: usize,
4741    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4742        let f = self.func("qmatvec_f32");
4743        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4744        let cfg = LaunchConfig {
4745            grid_dim: (out_f as u32, m as u32, 1),
4746            block_dim: (256, 1, 1),
4747            shared_mem_bytes: 0,
4748        };
4749        let (inf, outf, mi, qt, rb) =
4750            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4751        let __s_b = self.gpu.stream();
4752        let mut b = __s_b.launch_builder(&f);
4753        b.arg(w)
4754            .arg(x)
4755            .arg(&mut y)
4756            .arg(&inf)
4757            .arg(&outf)
4758            .arg(&mi)
4759            .arg(&qt)
4760            .arg(&rb);
4761        unsafe {
4762            b.launch(cfg)?;
4763        }
4764        Ok(y)
4765    }
4766
4767    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4768    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4769        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4770        self.keep_if_capturing(&s);
4771        Ok(s)
4772    }
4773
4774    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4775    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4776    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4777    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4778        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4779        self.keep_if_capturing(&s);
4780        Ok(s)
4781    }
4782
4783    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4784    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4785    pub fn memset_zeros_view(
4786        &self,
4787        dst: &mut cudarc::driver::CudaViewMut<f32>,
4788    ) -> Result<(), Box<dyn std::error::Error>> {
4789        self.gpu.stream().memset_zeros(dst)?;
4790        Ok(())
4791    }
4792
4793    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4794    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4795    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4796    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4797    /// stream would require an event).
4798    pub fn stage_expert(
4799        &self,
4800        host_bytes: &[u8],
4801        scratch: &mut CudaSlice<u8>,
4802        off: usize,
4803    ) -> Result<(), Box<dyn std::error::Error>> {
4804        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4805        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4806        Ok(())
4807    }
4808
4809    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4810    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4811    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4812    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4813    /// One CTA per token row, 256 threads (one per expert).
4814    pub fn moe_router_topk(
4815        &self,
4816        logits: &CudaSlice<f32>,
4817        t: usize,
4818        n_expert: usize,
4819        n_used: usize,
4820    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4821        let f = self.func("moe_router_topk_f32");
4822        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4823        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4824        let cfg = LaunchConfig {
4825            grid_dim: (t as u32, 1, 1),
4826            block_dim: (n_expert as u32, 1, 1),
4827            shared_mem_bytes: 0,
4828        };
4829        let (ne, nu) = (n_expert as i32, n_used as i32);
4830        let __s_b = self.gpu.stream();
4831        let mut b = __s_b.launch_builder(&f);
4832        b.arg(logits)
4833            .arg(&mut sel_idx)
4834            .arg(&mut sel_w)
4835            .arg(&ne)
4836            .arg(&nu);
4837        unsafe {
4838            b.launch(cfg)?;
4839        }
4840        Ok((sel_idx, sel_w))
4841    }
4842
4843    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4844    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4845    pub fn moe_router_topk_scaled(
4846        &self,
4847        logits: &CudaSlice<f32>,
4848        t: usize,
4849        n_expert: usize,
4850        n_used: usize,
4851        ex_scale: &CudaSlice<f32>,
4852    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4853        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4854        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4855        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4856        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4857        let f = self.func("moe_router_topk_scaled_f32");
4858        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4859        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4860        let cfg = LaunchConfig {
4861            grid_dim: (t as u32, 1, 1),
4862            block_dim: (n_expert as u32, 1, 1),
4863            shared_mem_bytes: 0,
4864        };
4865        let (ne, nu) = (n_expert as i32, n_used as i32);
4866        let __s_b = self.gpu.stream();
4867        let mut b = __s_b.launch_builder(&f);
4868        b.arg(logits)
4869            .arg(&mut sel_idx)
4870            .arg(&mut sel_w)
4871            .arg(&ne)
4872            .arg(&nu)
4873            .arg(ex_scale);
4874        unsafe {
4875            b.launch(cfg)?;
4876        }
4877        Ok((sel_idx, sel_w))
4878    }
4879
4880    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4881    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4882    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4883    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4884    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4885    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4886    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4887    pub fn moe_router_topk_host(
4888        &self,
4889        logits: &CudaSlice<f32>,
4890        t: usize,
4891        n_expert: usize,
4892        n_used: usize,
4893    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4894        let f = self.func("moe_router_topk_f32");
4895        let n = t * n_used;
4896        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4897        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4898        let cfg = LaunchConfig {
4899            grid_dim: (t as u32, 1, 1),
4900            block_dim: (n_expert as u32, 1, 1),
4901            shared_mem_bytes: 0,
4902        };
4903        let (ne, nu) = (n_expert as i32, n_used as i32);
4904        let __s_b = self.gpu.stream();
4905        let mut b = __s_b.launch_builder(&f);
4906        b.arg(logits)
4907            .arg(&mut sel_idx)
4908            .arg(&mut sel_w)
4909            .arg(&ne)
4910            .arg(&nu);
4911        unsafe {
4912            b.launch(cfg)?;
4913        }
4914        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4915        let bytes = n * 8;
4916        let mut guard = self.router_stage.lock().unwrap();
4917        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4918            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4919        }
4920        let stage = guard.as_mut().unwrap();
4921        let (si, sw) = unsafe {
4922            (
4923                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4924                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4925            )
4926        };
4927        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4928        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4929        self.gpu.stream().synchronize()?; // ONE sync for both
4930        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4931    }
4932
4933    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4934    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4935    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4936    #[allow(clippy::too_many_arguments)]
4937    pub fn moe_router_sigmoid_topk(
4938        &self,
4939        logits: &CudaSlice<f32>,
4940        t: usize,
4941        n_expert: usize,
4942        n_used: usize,
4943        active_count: usize,
4944        correction_bias: &CudaSlice<f32>,
4945        active: &CudaSlice<u8>,
4946        scaling_factor: f32,
4947        route_norm: bool,
4948    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4949        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4950        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4951            return Err(format!(
4952                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4953            )
4954            .into());
4955        }
4956        if logits.len() < t * n_expert
4957            || correction_bias.len() != n_expert
4958            || active.len() != n_expert
4959        {
4960            return Err(format!(
4961                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4962                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4963            ).into());
4964        }
4965        let f = self.func(crate::sigmoid_topk_kernel(
4966            crate::sig_expf_dev_on(),
4967            crate::topk_fast_on(),
4968            n_used,
4969        ));
4970        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4971        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4972        let threads = n_expert.div_ceil(32) * 32;
4973        let cfg = LaunchConfig {
4974            grid_dim: (t as u32, 1, 1),
4975            block_dim: (threads as u32, 1, 1),
4976            shared_mem_bytes: 0,
4977        };
4978        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4979        let __s_b = self.gpu.stream();
4980        let mut b = __s_b.launch_builder(&f);
4981        b.arg(logits)
4982            .arg(correction_bias)
4983            .arg(active)
4984            .arg(&mut sel_idx)
4985            .arg(&mut sel_w)
4986            .arg(&ne)
4987            .arg(&nu)
4988            .arg(&scaling_factor)
4989            .arg(&rn);
4990        unsafe {
4991            b.launch(cfg)?;
4992        }
4993        Ok((sel_idx, sel_w))
4994    }
4995
4996    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
4997    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
4998    #[allow(clippy::too_many_arguments)]
4999    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
5000    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
5001    /// the model engine can wait on it with a same-device stream memop.
5002    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
5003        if ptr == 0 {
5004            return Err("ring_flag_raw: unarmed flag".into());
5005        }
5006        let f = self.func("memra_ring_flag");
5007        let cfg = LaunchConfig {
5008            grid_dim: (1, 1, 1),
5009            block_dim: (32, 1, 1),
5010            shared_mem_bytes: 0,
5011        };
5012        let __s_b = self.gpu.stream();
5013        let mut b = __s_b.launch_builder(&f);
5014        b.arg(&ptr).arg(&value);
5015        unsafe {
5016            b.launch(cfg)?;
5017        }
5018        Ok(())
5019    }
5020
5021    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
5022    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
5023    pub fn moe_sel_w_mirror(
5024        &self,
5025        sel_src: &CudaSlice<i32>,
5026        w_src: &CudaSlice<f32>,
5027        sel_dst: &mut CudaSlice<i32>,
5028        w_dst: &mut CudaSlice<f32>,
5029        n: usize,
5030    ) -> Result<(), Box<dyn std::error::Error>> {
5031        if n == 0
5032            || n > 32
5033            || sel_src.len() < n
5034            || w_src.len() < n
5035            || sel_dst.len() < n
5036            || w_dst.len() < n
5037        {
5038            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
5039        }
5040        let f = self.func("moe_sel_w_mirror");
5041        let cfg = LaunchConfig {
5042            grid_dim: (1, 1, 1),
5043            block_dim: (32, 1, 1),
5044            shared_mem_bytes: 0,
5045        };
5046        let ni = n as i32;
5047        let __s_b = self.gpu.stream();
5048        let mut b = __s_b.launch_builder(&f);
5049        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5050        unsafe {
5051            b.launch(cfg)?;
5052        }
5053        Ok(())
5054    }
5055
5056    pub fn moe_router_sigmoid_topk_into(
5057        &self,
5058        logits: &CudaSlice<f32>,
5059        t: usize,
5060        n_expert: usize,
5061        n_used: usize,
5062        active_count: usize,
5063        correction_bias: &CudaSlice<f32>,
5064        active: &CudaSlice<u8>,
5065        scaling_factor: f32,
5066        route_norm: bool,
5067        sel_idx: &mut CudaSlice<i32>,
5068        sel_w: &mut CudaSlice<f32>,
5069    ) -> Result<(), Box<dyn std::error::Error>> {
5070        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5071        if n_expert == 0
5072            || n_expert > 1024
5073            || n_used == 0
5074            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
5075            || n_used > n_expert
5076            || logits.len() < t * n_expert
5077            || correction_bias.len() != n_expert
5078            || active.len() != n_expert
5079            || sel_idx.len() < t * n_used
5080            || sel_w.len() < t * n_used
5081        {
5082            return Err("sigmoid router _into geometry mismatch".into());
5083        }
5084        let f = self.func(crate::sigmoid_topk_kernel(
5085            crate::sig_expf_dev_on(),
5086            crate::topk_fast_on(),
5087            n_used,
5088        ));
5089        let threads = n_expert.div_ceil(32) * 32;
5090        let cfg = LaunchConfig {
5091            grid_dim: (t as u32, 1, 1),
5092            block_dim: (threads as u32, 1, 1),
5093            shared_mem_bytes: 0,
5094        };
5095        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5096        let __s_b = self.gpu.stream();
5097        let mut b = __s_b.launch_builder(&f);
5098        b.arg(logits)
5099            .arg(correction_bias)
5100            .arg(active)
5101            .arg(&mut *sel_idx)
5102            .arg(&mut *sel_w)
5103            .arg(&ne)
5104            .arg(&nu)
5105            .arg(&scaling_factor)
5106            .arg(&rn);
5107        unsafe {
5108            b.launch(cfg)?;
5109        }
5110        Ok(())
5111    }
5112
5113    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
5114    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
5115    #[allow(clippy::too_many_arguments)]
5116    pub fn moe_router_sigmoid_topk_host(
5117        &self,
5118        logits: &CudaSlice<f32>,
5119        t: usize,
5120        n_expert: usize,
5121        n_used: usize,
5122        active_count: usize,
5123        correction_bias: &CudaSlice<f32>,
5124        active: &CudaSlice<u8>,
5125        scaling_factor: f32,
5126        route_norm: bool,
5127    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5128        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5129            logits,
5130            t,
5131            n_expert,
5132            n_used,
5133            active_count,
5134            correction_bias,
5135            active,
5136            scaling_factor,
5137            route_norm,
5138        )?;
5139        let n = t * n_used;
5140        let bytes = n * 8;
5141        let mut guard = self.router_stage.lock().unwrap();
5142        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5143            *guard = Some(PinnedStage::new(bytes.max(4096))?);
5144        }
5145        let stage = guard.as_mut().unwrap();
5146        let (si, sw) = unsafe {
5147            (
5148                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5149                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5150            )
5151        };
5152        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5153        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5154        self.gpu.stream().synchronize()?;
5155        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5156    }
5157
5158    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
5159    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
5160    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
5161    pub fn stage_expert_async(
5162        &self,
5163        host_bytes: &[u8],
5164        scratch: &mut CudaSlice<u8>,
5165        off: usize,
5166    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5167        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5168        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5169        Ok(self.copy_stream.record_event(None)?)
5170    }
5171
5172    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
5173    pub fn compute_wait(
5174        &self,
5175        ev: &cudarc::driver::CudaEvent,
5176    ) -> Result<(), Box<dyn std::error::Error>> {
5177        self.gpu.stream().wait(ev)?;
5178        Ok(())
5179    }
5180
5181    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
5182    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
5183    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
5184    /// CudaView base+offset pointer is honored by the launch arg.
5185    pub fn qmatvec_view(
5186        &self,
5187        w: &CudaSlice<u8>,
5188        range: std::ops::Range<usize>,
5189        x: &cudarc::driver::CudaView<f32>,
5190        m: usize,
5191        in_f: usize,
5192        out_f: usize,
5193        qtype: i32,
5194        row_bytes: usize,
5195    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5196        let f = self.func("qmatvec_f32");
5197        let wv = w.slice(range); // CudaView<u8>, offset honored
5198        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
5199        let cfg = LaunchConfig {
5200            grid_dim: (out_f as u32, m as u32, 1),
5201            block_dim: (256, 1, 1),
5202            shared_mem_bytes: 0,
5203        };
5204        let (inf, outf, mi, qt, rb) =
5205            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5206        let __s_b = self.gpu.stream();
5207        let mut b = __s_b.launch_builder(&f);
5208        b.arg(&wv)
5209            .arg(x)
5210            .arg(&mut y)
5211            .arg(&inf)
5212            .arg(&outf)
5213            .arg(&mi)
5214            .arg(&qt)
5215            .arg(&rb);
5216        unsafe {
5217            b.launch(cfg)?;
5218        }
5219        Ok(y)
5220    }
5221
5222    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
5223    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
5224    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
5225    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
5226    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
5227    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
5228    #[allow(clippy::too_many_arguments)]
5229    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
5230    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
5231    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
5232    pub fn moe_gate_up_silu8_q8(
5233        &self,
5234        gp: WPtr8,
5235        up: WPtr8,
5236        aq: &CudaSlice<i8>,
5237        ad: &CudaSlice<f32>,
5238        in_f: usize,
5239        n_ff: usize,
5240        n_used: usize,
5241        qt_g: i32,
5242        qt_u: i32,
5243        rb_g: usize,
5244        rb_u: usize,
5245    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5246        let f = self.func("moe_gate_up_silu8_q8");
5247        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5248        let cfg = LaunchConfig {
5249            grid_dim: (n_ff as u32, n_used as u32, 1),
5250            block_dim: (32, 1, 1),
5251            shared_mem_bytes: 0,
5252        };
5253        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5254        let __s_b = self.gpu.stream();
5255        let mut b = __s_b.launch_builder(&f);
5256        b.arg(&gp)
5257            .arg(&up)
5258            .arg(aq)
5259            .arg(ad)
5260            .arg(&mut act)
5261            .arg(&inf)
5262            .arg(&nff)
5263            .arg(&qt_g)
5264            .arg(&qt_u)
5265            .arg(&rbg)
5266            .arg(&rbu);
5267        unsafe {
5268            b.launch(cfg)?;
5269        }
5270        Ok(act)
5271    }
5272
5273    #[allow(clippy::too_many_arguments)]
5274    pub fn moe_down8_fma_q8(
5275        &self,
5276        dp: WPtr8,
5277        w: F32x8,
5278        aq2: &CudaSlice<i8>,
5279        ad2: &CudaSlice<f32>,
5280        dst: &mut cudarc::driver::CudaViewMut<f32>,
5281        in_f: usize,
5282        out_f: usize,
5283        n_used: usize,
5284        qt: i32,
5285        rb: usize,
5286    ) -> Result<(), Box<dyn std::error::Error>> {
5287        let f = self.func("moe_down8_fma_q8");
5288        let cfg = LaunchConfig {
5289            grid_dim: (out_f as u32, 1, 1),
5290            block_dim: (32, 1, 1),
5291            shared_mem_bytes: 0,
5292        };
5293        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5294        let __s_b = self.gpu.stream();
5295        let mut b = __s_b.launch_builder(&f);
5296        b.arg(&dp)
5297            .arg(&w)
5298            .arg(aq2)
5299            .arg(ad2)
5300            .arg(dst)
5301            .arg(&inf)
5302            .arg(&outf)
5303            .arg(&nu)
5304            .arg(&qt)
5305            .arg(&rbi);
5306        unsafe {
5307            b.launch(cfg)?;
5308        }
5309        Ok(())
5310    }
5311
5312    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
5313    pub fn qmatvec_expert_q8(
5314        &self,
5315        w: &CudaSlice<u8>,
5316        range: std::ops::Range<usize>,
5317        aq: &CudaSlice<i8>,
5318        ad: &CudaSlice<f32>,
5319        m: usize,
5320        in_f: usize,
5321        out_f: usize,
5322        qtype: i32,
5323        row_bytes: usize,
5324    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5325        let f = self.func("qmatvec_expert_q8");
5326        let wv = w.slice(range);
5327        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5328        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
5329        let cfg = LaunchConfig {
5330            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5331            block_dim: (32, ROWS, 1),
5332            shared_mem_bytes: 0,
5333        };
5334        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5335        let __s_b = self.gpu.stream();
5336        let mut b = __s_b.launch_builder(&f);
5337        b.arg(&wv)
5338            .arg(aq)
5339            .arg(ad)
5340            .arg(&mut y)
5341            .arg(&inf)
5342            .arg(&outf)
5343            .arg(&mi)
5344            .arg(&qtype)
5345            .arg(&rbi);
5346        unsafe {
5347            b.launch(cfg)?;
5348        }
5349        Ok(y)
5350    }
5351
5352    pub fn moe_gate_up_silu8(
5353        &self,
5354        gp: WPtr8,
5355        up: WPtr8,
5356        x: &cudarc::driver::CudaView<f32>,
5357        in_f: usize,
5358        n_ff: usize,
5359        n_used: usize,
5360        qt_g: i32,
5361        qt_u: i32,
5362        rb_g: usize,
5363        rb_u: usize,
5364    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5365        let f = self.func("moe_gate_up_silu8_f32");
5366        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
5367        let cfg = LaunchConfig {
5368            grid_dim: (n_ff as u32, n_used as u32, 1),
5369            block_dim: (256, 1, 1),
5370            shared_mem_bytes: 0,
5371        };
5372        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5373        let __s_b = self.gpu.stream();
5374        let mut b = __s_b.launch_builder(&f);
5375        b.arg(&gp)
5376            .arg(&up)
5377            .arg(x)
5378            .arg(&mut act)
5379            .arg(&inf)
5380            .arg(&nff)
5381            .arg(&qt_g)
5382            .arg(&qt_u)
5383            .arg(&rbg)
5384            .arg(&rbu);
5385        unsafe {
5386            b.launch(cfg)?;
5387        }
5388        Ok(act)
5389    }
5390
5391    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
5392    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
5393    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
5394    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
5395    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
5396    #[allow(clippy::too_many_arguments)]
5397    pub fn moe_down8_fma_into(
5398        &self,
5399        dp: WPtr8,
5400        w: F32x8,
5401        act: &CudaSlice<f32>,
5402        dst: &mut cudarc::driver::CudaViewMut<f32>,
5403        in_f: usize,
5404        out_f: usize,
5405        n_used: usize,
5406        qt: i32,
5407        rb: usize,
5408    ) -> Result<(), Box<dyn std::error::Error>> {
5409        let f = self.func("moe_down8_fma_f32");
5410        let cfg = LaunchConfig {
5411            grid_dim: (out_f as u32, 1, 1),
5412            block_dim: (256, 1, 1),
5413            shared_mem_bytes: 0,
5414        };
5415        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5416        let __s_b = self.gpu.stream();
5417        let mut b = __s_b.launch_builder(&f);
5418        b.arg(&dp)
5419            .arg(&w)
5420            .arg(act)
5421            .arg(dst)
5422            .arg(&inf)
5423            .arg(&outf)
5424            .arg(&nu)
5425            .arg(&qt)
5426            .arg(&rbv);
5427        unsafe {
5428            b.launch(cfg)?;
5429        }
5430        Ok(())
5431    }
5432
5433    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
5434    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
5435    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
5436    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
5437    #[allow(clippy::too_many_arguments)]
5438    /// dp4a q8 twin of the _dev pair (resident-experts arc).
5439    ///
5440    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
5441    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
5442    /// down's FMA chain stays slot-ordered serial). Seams:
5443    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
5444    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
5445    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
5446    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
5447    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
5448    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
5449    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
5450    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
5451    ///                       only) | w8h2 (h2 x slot-parallel)
5452    #[allow(clippy::too_many_arguments)]
5453    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
5454    #[allow(clippy::too_many_arguments)]
5455    pub fn moe_pairs_matvec_q8(
5456        &self,
5457        table: &CudaSlice<u64>,
5458        proj: i32,
5459        pair_tok: &CudaSlice<i32>,
5460        pair_ex: &CudaSlice<i32>,
5461        aq: &CudaSlice<i8>,
5462        ad: &CudaSlice<f32>,
5463        in_f: usize,
5464        out_f: usize,
5465        n_expert: usize,
5466        n_pairs: usize,
5467        qtype: i32,
5468        row_bytes: usize,
5469    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5470        let f = self.func("moe_pairs_matvec_q8");
5471        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5472        const ROWS: u32 = 4;
5473        let cfg = LaunchConfig {
5474            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5475            block_dim: (32, ROWS, 1),
5476            shared_mem_bytes: 0,
5477        };
5478        let (inf, outf, ne, np, rbi) = (
5479            in_f as i32,
5480            out_f as i32,
5481            n_expert as i32,
5482            n_pairs as i32,
5483            row_bytes as i64,
5484        );
5485        let __s_b = self.gpu.stream();
5486        let mut b = __s_b.launch_builder(&f);
5487        b.arg(table)
5488            .arg(&proj)
5489            .arg(pair_tok)
5490            .arg(pair_ex)
5491            .arg(aq)
5492            .arg(ad)
5493            .arg(&mut y)
5494            .arg(&inf)
5495            .arg(&outf)
5496            .arg(&ne)
5497            .arg(&np)
5498            .arg(&qtype)
5499            .arg(&rbi);
5500        unsafe {
5501            b.launch(cfg)?;
5502        }
5503        Ok(y)
5504    }
5505
5506    /// Expert-major pair matvec (weight-reuse across each expert's token group).
5507    #[allow(clippy::too_many_arguments)]
5508    pub fn moe_pairs_matvec_q8_em(
5509        &self,
5510        table: &CudaSlice<u64>,
5511        proj: i32,
5512        ex_ids: &CudaSlice<i32>,
5513        ex_off: &CudaSlice<i32>,
5514        ex_pairs: &CudaSlice<i32>,
5515        pair_tok: &CudaSlice<i32>,
5516        aq: &CudaSlice<i8>,
5517        ad: &CudaSlice<f32>,
5518        in_f: usize,
5519        out_f: usize,
5520        n_expert: usize,
5521        n_active: usize,
5522        n_pairs: usize,
5523        qtype: i32,
5524        row_bytes: usize,
5525    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5526        let f = self.func("moe_pairs_matvec_q8_em");
5527        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5528        const ROWS: u32 = 4;
5529        let cfg = LaunchConfig {
5530            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5531            block_dim: (32, ROWS, 1),
5532            shared_mem_bytes: 0,
5533        };
5534        let (inf, outf, ne, na, rbi) = (
5535            in_f as i32,
5536            out_f as i32,
5537            n_expert as i32,
5538            n_active as i32,
5539            row_bytes as i64,
5540        );
5541        let __s_b = self.gpu.stream();
5542        let mut b = __s_b.launch_builder(&f);
5543        b.arg(table)
5544            .arg(&proj)
5545            .arg(ex_ids)
5546            .arg(ex_off)
5547            .arg(ex_pairs)
5548            .arg(pair_tok)
5549            .arg(aq)
5550            .arg(ad)
5551            .arg(&mut y)
5552            .arg(&inf)
5553            .arg(&outf)
5554            .arg(&ne)
5555            .arg(&na)
5556            .arg(&qtype)
5557            .arg(&rbi);
5558        unsafe {
5559            b.launch(cfg)?;
5560        }
5561        Ok(y)
5562    }
5563
5564    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
5565    // weight group once per (row,group) then dp4a's across the expert's token group.
5566    #[allow(clippy::too_many_arguments)]
5567    pub fn moe_pairs_matvec_q8_dec(
5568        &self,
5569        table: &CudaSlice<u64>,
5570        proj: i32,
5571        ex_ids: &CudaSlice<i32>,
5572        ex_off: &CudaSlice<i32>,
5573        ex_pairs: &CudaSlice<i32>,
5574        pair_tok: &CudaSlice<i32>,
5575        aq: &CudaSlice<i8>,
5576        ad: &CudaSlice<f32>,
5577        in_f: usize,
5578        out_f: usize,
5579        n_expert: usize,
5580        n_active: usize,
5581        n_pairs: usize,
5582        qtype: i32,
5583        row_bytes: usize,
5584    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5585        let f = self.func("moe_pairs_matvec_q8_dec");
5586        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5587        const ROWS: u32 = 4;
5588        let cfg = LaunchConfig {
5589            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5590            block_dim: (32, ROWS, 1),
5591            shared_mem_bytes: 0,
5592        };
5593        let (inf, outf, ne, na, rbi) = (
5594            in_f as i32,
5595            out_f as i32,
5596            n_expert as i32,
5597            n_active as i32,
5598            row_bytes as i64,
5599        );
5600        let __s_b = self.gpu.stream();
5601        let mut b = __s_b.launch_builder(&f);
5602        b.arg(table)
5603            .arg(&proj)
5604            .arg(ex_ids)
5605            .arg(ex_off)
5606            .arg(ex_pairs)
5607            .arg(pair_tok)
5608            .arg(aq)
5609            .arg(ad)
5610            .arg(&mut y)
5611            .arg(&inf)
5612            .arg(&outf)
5613            .arg(&ne)
5614            .arg(&na)
5615            .arg(&qtype)
5616            .arg(&rbi);
5617        unsafe {
5618            b.launch(cfg)?;
5619        }
5620        Ok(y)
5621    }
5622
5623    pub fn moe_pairs_gelu_mul(
5624        &self,
5625        gate: &CudaSlice<f32>,
5626        up: &CudaSlice<f32>,
5627        n: usize,
5628    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5629        let f = self.func("moe_pairs_gelu_mul");
5630        let mut act = self.alloc_uninit::<f32>(n)?;
5631        let cfg = LaunchConfig::for_num_elems(n as u32);
5632        let nl = n as i64;
5633        let __s_b = self.gpu.stream();
5634        let mut b = __s_b.launch_builder(&f);
5635        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5636        unsafe {
5637            b.launch(cfg)?;
5638        }
5639        Ok(act)
5640    }
5641
5642    pub fn moe_pairs_silu_mul(
5643        &self,
5644        gate: &CudaSlice<f32>,
5645        up: &CudaSlice<f32>,
5646        n: usize,
5647    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5648        let f = self.func("moe_pairs_silu_mul");
5649        let mut act = self.alloc_uninit::<f32>(n)?;
5650        let cfg = LaunchConfig::for_num_elems(n as u32);
5651        let nl = n as i64;
5652        let __s_b = self.gpu.stream();
5653        let mut b = __s_b.launch_builder(&f);
5654        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5655        unsafe {
5656            b.launch(cfg)?;
5657        }
5658        Ok(act)
5659    }
5660
5661    #[allow(clippy::too_many_arguments)]
5662    pub fn moe_pairs_scatter(
5663        &self,
5664        y_down: &CudaSlice<f32>,
5665        pair_w: &CudaSlice<f32>,
5666        tok_pair_off: &CudaSlice<i32>,
5667        tok_pair_ids: &CudaSlice<i32>,
5668        moe_out: &mut CudaSlice<f32>,
5669        t: usize,
5670        n_embd: usize,
5671    ) -> Result<(), Box<dyn std::error::Error>> {
5672        let f = self.func("moe_pairs_scatter");
5673        let cfg = LaunchConfig {
5674            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5675            block_dim: (256, 1, 1),
5676            shared_mem_bytes: 0,
5677        };
5678        let ne = n_embd as i32;
5679        let __s_b = self.gpu.stream();
5680        let mut b = __s_b.launch_builder(&f);
5681        b.arg(y_down)
5682            .arg(pair_w)
5683            .arg(tok_pair_off)
5684            .arg(tok_pair_ids)
5685            .arg(moe_out)
5686            .arg(&ne);
5687        unsafe {
5688            b.launch(cfg)?;
5689        }
5690        Ok(())
5691    }
5692
5693    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
5694    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
5695    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
5696    #[allow(clippy::too_many_arguments)]
5697    pub fn moe_gate_up_gelu8_dev_q8(
5698        &self,
5699        table: &CudaSlice<u64>,
5700        sel: &cudarc::driver::CudaView<i32>,
5701        aq: &CudaSlice<i8>,
5702        ad: &CudaSlice<f32>,
5703        in_f: usize,
5704        n_ff: usize,
5705        n_used: usize,
5706        n_expert: usize,
5707        qt_g: i32,
5708        qt_u: i32,
5709        rb_g: usize,
5710        rb_u: usize,
5711    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5712        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5713        let (inf, nff, ne, rbg, rbu) = (
5714            in_f as i32,
5715            n_ff as i32,
5716            n_expert as i32,
5717            rb_g as i64,
5718            rb_u as i64,
5719        );
5720        let f = self.func("moe_gate_up_gelu8_dev_q8");
5721        let cfg = LaunchConfig {
5722            grid_dim: (n_ff as u32, n_used as u32, 1),
5723            block_dim: (32, 1, 1),
5724            shared_mem_bytes: 0,
5725        };
5726        let __s_b = self.gpu.stream();
5727        let mut b = __s_b.launch_builder(&f);
5728        b.arg(table)
5729            .arg(sel)
5730            .arg(aq)
5731            .arg(ad)
5732            .arg(&mut act)
5733            .arg(&inf)
5734            .arg(&nff)
5735            .arg(&ne)
5736            .arg(&qt_g)
5737            .arg(&qt_u)
5738            .arg(&rbg)
5739            .arg(&rbu);
5740        unsafe {
5741            b.launch(cfg)?;
5742        }
5743        Ok(act)
5744    }
5745
5746    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5747    #[allow(clippy::too_many_arguments)]
5748    pub fn moe_gate_up_gelu8_dev_q8_rows(
5749        &self,
5750        table: &CudaSlice<u64>,
5751        sel: &CudaSlice<i32>,
5752        aq: &CudaSlice<i8>,
5753        ad: &CudaSlice<f32>,
5754        t: usize,
5755        in_f: usize,
5756        n_ff: usize,
5757        n_used: usize,
5758        n_expert: usize,
5759        qt_g: i32,
5760        qt_u: i32,
5761        rb_g: usize,
5762        rb_u: usize,
5763    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5764        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5765        let (inf, nff, ne, rbg, rbu, nu) = (
5766            in_f as i32,
5767            n_ff as i32,
5768            n_expert as i32,
5769            rb_g as i64,
5770            rb_u as i64,
5771            n_used as i32,
5772        );
5773        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5774        let cfg = LaunchConfig {
5775            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5776            block_dim: (32, 1, 1),
5777            shared_mem_bytes: 0,
5778        };
5779        let __s_b = self.gpu.stream();
5780        let mut b = __s_b.launch_builder(&f);
5781        b.arg(table)
5782            .arg(sel)
5783            .arg(aq)
5784            .arg(ad)
5785            .arg(&mut act)
5786            .arg(&inf)
5787            .arg(&nff)
5788            .arg(&ne)
5789            .arg(&qt_g)
5790            .arg(&qt_u)
5791            .arg(&rbg)
5792            .arg(&rbu)
5793            .arg(&nu);
5794        unsafe {
5795            b.launch(cfg)?;
5796        }
5797        Ok(act)
5798    }
5799
5800    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5801    #[allow(clippy::too_many_arguments)]
5802    pub fn moe_gate_up_gelu8_dev_q8_csr(
5803        &self,
5804        table: &CudaSlice<u64>,
5805        sel: &CudaSlice<i32>,
5806        aq: &CudaSlice<i8>,
5807        ad: &CudaSlice<f32>,
5808        n_pairs: usize,
5809        in_f: usize,
5810        n_ff: usize,
5811        n_used: usize,
5812        n_expert: usize,
5813        qt_g: i32,
5814        qt_u: i32,
5815        rb_g: usize,
5816        rb_u: usize,
5817    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5818        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5819        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5820            in_f as i32,
5821            n_ff as i32,
5822            n_expert as i32,
5823            rb_g as i64,
5824            rb_u as i64,
5825            n_used as i32,
5826            n_pairs as i32,
5827        );
5828        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5829        let cfg = LaunchConfig {
5830            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5831            block_dim: (32, 1, 1),
5832            shared_mem_bytes: 0,
5833        };
5834        let __s_b = self.gpu.stream();
5835        let mut b = __s_b.launch_builder(&f);
5836        b.arg(table)
5837            .arg(sel)
5838            .arg(aq)
5839            .arg(ad)
5840            .arg(&mut act)
5841            .arg(&inf)
5842            .arg(&nff)
5843            .arg(&ne)
5844            .arg(&qt_g)
5845            .arg(&qt_u)
5846            .arg(&rbg)
5847            .arg(&rbu)
5848            .arg(&nu)
5849            .arg(&npi);
5850        unsafe {
5851            b.launch(cfg)?;
5852        }
5853        Ok(act)
5854    }
5855
5856    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5857    #[allow(clippy::too_many_arguments)]
5858    pub fn moe_down8_fma_dev_q8_rows_g(
5859        &self,
5860        table: &CudaSlice<u64>,
5861        sel: &CudaSlice<i32>,
5862        w: &CudaSlice<f32>,
5863        aq2: &CudaSlice<i8>,
5864        ad2: &CudaSlice<f32>,
5865        dst: &mut CudaSlice<f32>,
5866        t: usize,
5867        in_f: usize,
5868        out_f: usize,
5869        n_used: usize,
5870        n_expert: usize,
5871        qt: i32,
5872        rb: usize,
5873    ) -> Result<(), Box<dyn std::error::Error>> {
5874        let (inf, outf, nu, ne, rbi) = (
5875            in_f as i32,
5876            out_f as i32,
5877            n_used as i32,
5878            n_expert as i32,
5879            rb as i64,
5880        );
5881        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5882        // eight warps, then replay the original slot-ordered FMA chain. Every
5883        // other shape retains the generic one-warp rows kernel.
5884        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5885        let f = self.func(if step_b1_w8 {
5886            "moe_down8_fma_dev_q8_rows_w8"
5887        } else {
5888            "moe_down8_fma_dev_q8_rows_g"
5889        });
5890        let cfg = LaunchConfig {
5891            grid_dim: (out_f as u32, 1, t as u32),
5892            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5893            shared_mem_bytes: 0,
5894        };
5895        let __s_b = self.gpu.stream();
5896        let mut b = __s_b.launch_builder(&f);
5897        b.arg(table)
5898            .arg(sel)
5899            .arg(w)
5900            .arg(aq2)
5901            .arg(ad2)
5902            .arg(dst)
5903            .arg(&inf)
5904            .arg(&outf)
5905            .arg(&nu)
5906            .arg(&ne)
5907            .arg(&qt)
5908            .arg(&rbi);
5909        unsafe {
5910            b.launch(cfg)?;
5911        }
5912        Ok(())
5913    }
5914
5915    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5916    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5917    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5918    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5919        let (out_f, in_f) = (2048usize, 2816usize);
5920        let nblk = in_f / 32;
5921        let mut seed = 0x9E3779B97F4A7C15u64;
5922        let mut rng = move || {
5923            seed = seed
5924                .wrapping_mul(6364136223846793005)
5925                .wrapping_add(1442695040888963407);
5926            (seed >> 33) as u8
5927        };
5928        let mut w = vec![0u8; out_f * nblk * 18];
5929        for b in w.iter_mut() {
5930            *b = rng();
5931        }
5932        for r in 0..out_f {
5933            for g in 0..nblk {
5934                let off = (r * nblk + g) * 18;
5935                w[off] = 0x00;
5936                w[off + 1] = 0x2C; // sane half d
5937            }
5938        }
5939        let qplane = out_f * nblk * 16;
5940        let mut wrp = vec![0u8; w.len()];
5941        for r in 0..out_f {
5942            for g in 0..nblk {
5943                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5944                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5945                    .copy_from_slice(&src[0..2]);
5946                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5947            }
5948        }
5949        let w_d = self.htod_bytes(&w)?;
5950        let wrp_d = self.htod_bytes(&wrp)?;
5951        let mut aq = vec![0i8; m * in_f];
5952        for v in aq.iter_mut() {
5953            *v = rng() as i8;
5954        }
5955        let aq_d = self.htod_i8(&aq)?;
5956        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5957        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5958        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5959        const RPB: u32 = 4;
5960        let cfg = LaunchConfig {
5961            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5962            block_dim: (32, RPB, 1),
5963            shared_mem_bytes: 0,
5964        };
5965        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5966        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5967        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5968        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5969        {
5970            let __s_b = self.gpu.stream();
5971            let mut b = __s_b.launch_builder(&fb);
5972            b.arg(&w_d)
5973                .arg(&aq_d)
5974                .arg(&ad_d)
5975                .arg(&mut y0)
5976                .arg(&inf)
5977                .arg(&outf)
5978                .arg(&mi)
5979                .arg(&rb);
5980            unsafe {
5981                b.launch(cfg)?;
5982            }
5983            let __s_b = self.gpu.stream();
5984            let mut b = __s_b.launch_builder(&fr);
5985            b.arg(&wrp_d)
5986                .arg(&aq_d)
5987                .arg(&ad_d)
5988                .arg(&mut y1)
5989                .arg(&inf)
5990                .arg(&outf)
5991                .arg(&mi)
5992                .arg(&qp);
5993            unsafe {
5994                b.launch(cfg)?;
5995            }
5996        }
5997        self.gpu.stream().synchronize()?;
5998        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5999        let nd = h0
6000            .iter()
6001            .zip(&h1)
6002            .filter(|(a, b)| a.to_bits() != b.to_bits())
6003            .count();
6004        if nd != 0 {
6005            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
6006        }
6007        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
6008            self.gpu.stream().synchronize()?;
6009            let t0 = std::time::Instant::now();
6010            for _ in 0..500 {
6011                if rp {
6012                    let __s_b = self.gpu.stream();
6013                    let mut b = __s_b.launch_builder(&fr);
6014                    b.arg(&wrp_d)
6015                        .arg(&aq_d)
6016                        .arg(&ad_d)
6017                        .arg(&mut y1)
6018                        .arg(&inf)
6019                        .arg(&outf)
6020                        .arg(&mi)
6021                        .arg(&qp);
6022                    unsafe {
6023                        b.launch(cfg)?;
6024                    }
6025                } else {
6026                    let __s_b = self.gpu.stream();
6027                    let mut b = __s_b.launch_builder(&fb);
6028                    b.arg(&w_d)
6029                        .arg(&aq_d)
6030                        .arg(&ad_d)
6031                        .arg(&mut y0)
6032                        .arg(&inf)
6033                        .arg(&outf)
6034                        .arg(&mi)
6035                        .arg(&rb);
6036                    unsafe {
6037                        b.launch(cfg)?;
6038                    }
6039                }
6040            }
6041            self.gpu.stream().synchronize()?;
6042            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6043        };
6044        let _ = time(false)?;
6045        let _ = time(true)?; // warm
6046        Ok((time(false)?, time(true)?))
6047    }
6048
6049    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
6050    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
6051    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
6052    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
6053    pub fn build_q4_rp4(
6054        &self,
6055        t: &mut crate::model::GpuTensor,
6056    ) -> Result<(), Box<dyn std::error::Error>> {
6057        use crate::model::GpuTensor;
6058        let GpuTensor::Quant {
6059            bytes,
6060            qtype,
6061            row_bytes,
6062            ne,
6063            rp4,
6064            ..
6065        } = t
6066        else {
6067            return Ok(());
6068        };
6069        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6070            return Ok(());
6071        }
6072        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6073        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6074            return Ok(());
6075        }
6076        let nblk = in_f / 32;
6077        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6078        let f = self.func("q4_0_split_rp_build");
6079        let n = (out_f * nblk) as i32;
6080        let cfg = LaunchConfig {
6081            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6082            block_dim: (256, 1, 1),
6083            shared_mem_bytes: 0,
6084        };
6085        let (of, nb) = (out_f as i32, nblk as i32);
6086        let _ = n;
6087        let __s_b = self.gpu.stream();
6088        let mut b = __s_b.launch_builder(&f);
6089        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6090        unsafe {
6091            b.launch(cfg)?;
6092        }
6093        *rp4 = Some(dst);
6094        Ok(())
6095    }
6096
6097    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
6098    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
6099    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
6100    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
6101    pub fn build_q8_rp4(
6102        &self,
6103        t: &mut crate::model::GpuTensor,
6104    ) -> Result<(), Box<dyn std::error::Error>> {
6105        use crate::model::GpuTensor;
6106        let GpuTensor::Quant {
6107            bytes,
6108            qtype,
6109            row_bytes,
6110            ne,
6111            rp4,
6112            ..
6113        } = t
6114        else {
6115            return Ok(());
6116        };
6117        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6118            return Ok(());
6119        }
6120        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6121        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6122            return Ok(());
6123        }
6124        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6125        Ok(())
6126    }
6127
6128    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
6129    /// mirror without a GpuTensor (same kernel the loader path above uses).
6130    pub fn build_q8_rp4_raw(
6131        &self,
6132        bytes: &CudaSlice<u8>,
6133        in_f: usize,
6134        out_f: usize,
6135    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6136        assert!(in_f % 32 == 0);
6137        let nblk = in_f / 32;
6138        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6139        let f = self.func("q8_0_split_rp_build");
6140        let cfg = LaunchConfig {
6141            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6142            block_dim: (256, 1, 1),
6143            shared_mem_bytes: 0,
6144        };
6145        let (of, nb) = (out_f as i32, nblk as i32);
6146        let __s_b = self.gpu.stream();
6147        let mut b = __s_b.launch_builder(&f);
6148        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6149        unsafe {
6150            b.launch(cfg)?;
6151        }
6152        Ok(dst)
6153    }
6154
6155    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
6156    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
6157    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
6158    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
6159    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
6160    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
6161    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
6162    pub fn build_q4k_rp4(
6163        &self,
6164        t: &mut crate::model::GpuTensor,
6165    ) -> Result<(), Box<dyn std::error::Error>> {
6166        use crate::model::GpuTensor;
6167        let GpuTensor::Quant {
6168            bytes,
6169            qtype,
6170            row_bytes,
6171            ne,
6172            rp4,
6173            ..
6174        } = t
6175        else {
6176            return Ok(());
6177        };
6178        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6179            return Ok(());
6180        }
6181        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6182        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6183            return Ok(());
6184        }
6185        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6186        Ok(())
6187    }
6188
6189    pub fn build_q6k_rp4(
6190        &self,
6191        t: &mut crate::model::GpuTensor,
6192    ) -> Result<(), Box<dyn std::error::Error>> {
6193        use crate::model::GpuTensor;
6194        let GpuTensor::Quant {
6195            bytes,
6196            qtype,
6197            row_bytes,
6198            ne,
6199            rp4,
6200            ..
6201        } = t
6202        else {
6203            return Ok(());
6204        };
6205        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6206            return Ok(());
6207        }
6208        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6209        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6210            return Ok(());
6211        }
6212        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6213        Ok(())
6214    }
6215
6216    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
6217    pub fn build_kq_rp4_raw(
6218        &self,
6219        bytes: &CudaSlice<u8>,
6220        in_f: usize,
6221        out_f: usize,
6222        qtype: i32,
6223    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6224        assert!(in_f % 256 == 0);
6225        let nsbk = in_f / 256;
6226        let (sb_bytes, kname) = match qtype {
6227            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6228            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6229            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6230        };
6231        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6232        let f = self.func(kname);
6233        let cfg = LaunchConfig {
6234            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6235            block_dim: (256, 1, 1),
6236            shared_mem_bytes: 0,
6237        };
6238        let (of, nb) = (out_f as i32, nsbk as i32);
6239        let __s_b = self.gpu.stream();
6240        let mut b = __s_b.launch_builder(&f);
6241        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6242        unsafe {
6243            b.launch(cfg)?;
6244        }
6245        Ok(dst)
6246    }
6247
6248    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
6249    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
6250    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
6251    pub fn kqrp_enabled() -> bool {
6252        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6253        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6254            Ok("0") => false,
6255            Ok(_) => true,
6256            Err(_) => cfg!(memra_hopper_mma),
6257        })
6258    }
6259
6260    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
6261    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
6262    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
6263    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
6264    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
6265    pub fn build_q4_rp_swap(
6266        &self,
6267        t: &mut crate::model::GpuTensor,
6268    ) -> Result<bool, Box<dyn std::error::Error>> {
6269        use crate::model::GpuTensor;
6270        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
6271        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
6272        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
6273        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
6274        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
6275        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
6276        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
6277        // this fn's OWN builder serves may ever be swapped; everything else refuses
6278        // here, regardless of walk ordering.
6279        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6280            return Ok(false);
6281        }
6282        self.build_q4_rp4(t)?;
6283        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
6284        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6285            return Ok(false);
6286        };
6287        match rp4.take() {
6288            Some(split) => {
6289                *bytes = split; // the GGUF-layout buffer drops here
6290                *rp = true;
6291                Ok(true)
6292            }
6293            None => Ok(false),
6294        }
6295    }
6296
6297    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
6298    pub fn q4rp_enabled() -> bool {
6299        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6300        *ON.get_or_init(|| {
6301            std::env::var("MEMRA_Q4RP")
6302                .map(|v| v != "0")
6303                .unwrap_or(true)
6304        })
6305    }
6306
6307    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
6308    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
6309    pub fn copy_rows_strided(
6310        &self,
6311        src: &CudaSlice<f32>,
6312        dst: &mut CudaSlice<f32>,
6313        row_elems: usize,
6314        n_rows: usize,
6315        src_stride: usize,
6316        src_off: usize,
6317    ) -> Result<(), Box<dyn std::error::Error>> {
6318        let f = self.func("copy_rows_strided_f32");
6319        let cfg = LaunchConfig {
6320            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6321            block_dim: (256, 1, 1),
6322            shared_mem_bytes: 0,
6323        };
6324        let (re, nr) = (row_elems as i32, n_rows as i32);
6325        let (st, off) = (src_stride as i64, src_off as i64);
6326        let __s_b = self.gpu.stream();
6327        let mut b = __s_b.launch_builder(&f);
6328        b.arg(src)
6329            .arg(&mut *dst)
6330            .arg(&re)
6331            .arg(&nr)
6332            .arg(&st)
6333            .arg(&off);
6334        unsafe {
6335            b.launch(cfg)?;
6336        }
6337        Ok(())
6338    }
6339
6340    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
6341    ///
6342    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
6343    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
6344    /// one peer copy per token.
6345    pub fn place_rows_strided(
6346        &self,
6347        src: &CudaSlice<f32>,
6348        dst: &mut CudaSlice<f32>,
6349        row_elems: usize,
6350        n_rows: usize,
6351        dst_stride: usize,
6352        dst_off: usize,
6353    ) -> Result<(), Box<dyn std::error::Error>> {
6354        if row_elems == 0 || n_rows == 0 {
6355            return Err("strided row placement requires nonzero rows and row width".into());
6356        }
6357        let src_len = n_rows
6358            .checked_mul(row_elems)
6359            .ok_or("strided row placement source size overflow")?;
6360        let dst_len = n_rows
6361            .checked_sub(1)
6362            .and_then(|rows| rows.checked_mul(dst_stride))
6363            .and_then(|base| base.checked_add(dst_off))
6364            .and_then(|base| base.checked_add(row_elems))
6365            .ok_or("strided row placement destination size overflow")?;
6366        let row_end = dst_off
6367            .checked_add(row_elems)
6368            .ok_or("strided row placement row size overflow")?;
6369        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6370            return Err(format!(
6371                "strided row placement geometry mismatch: src={} need_src={src_len} \
6372                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6373                 dst_stride={dst_stride} dst_off={dst_off}",
6374                src.len(),
6375                dst.len(),
6376            )
6377            .into());
6378        }
6379        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6380            return Err("strided row placement exceeds CUDA kernel geometry".into());
6381        }
6382        let f = self.func("place_rows_strided_f32");
6383        let cfg = LaunchConfig {
6384            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6385            block_dim: (256, 1, 1),
6386            shared_mem_bytes: 0,
6387        };
6388        let (re, nr) = (row_elems as i32, n_rows as i32);
6389        let (st, off) = (dst_stride as i64, dst_off as i64);
6390        let __s_b = self.gpu.stream();
6391        let mut b = __s_b.launch_builder(&f);
6392        b.arg(src)
6393            .arg(&mut *dst)
6394            .arg(&re)
6395            .arg(&nr)
6396            .arg(&st)
6397            .arg(&off);
6398        unsafe {
6399            b.launch(cfg)?;
6400        }
6401        Ok(())
6402    }
6403
6404    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
6405    pub fn u32_set_k(
6406        &self,
6407        dst: &mut CudaSlice<u32>,
6408        v: u32,
6409        idx: usize,
6410    ) -> Result<(), Box<dyn std::error::Error>> {
6411        let f = self.func("u32_set_k");
6412        let cfg = LaunchConfig {
6413            grid_dim: (1, 1, 1),
6414            block_dim: (1, 1, 1),
6415            shared_mem_bytes: 0,
6416        };
6417        let ii = idx as i32;
6418        let __s_b = self.gpu.stream();
6419        let mut b = __s_b.launch_builder(&f);
6420        b.arg(dst).arg(&v).arg(&ii);
6421        unsafe {
6422            b.launch(cfg)?;
6423        }
6424        Ok(())
6425    }
6426
6427    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
6428    pub fn i32_add_k(
6429        &self,
6430        d: &mut CudaSlice<i32>,
6431        v: i32,
6432    ) -> Result<(), Box<dyn std::error::Error>> {
6433        let f = self.func("i32_add_k");
6434        let cfg = LaunchConfig {
6435            grid_dim: (1, 1, 1),
6436            block_dim: (32, 1, 1),
6437            shared_mem_bytes: 0,
6438        };
6439        let __s_b = self.gpu.stream();
6440        let mut b = __s_b.launch_builder(&f);
6441        b.arg(d).arg(&v);
6442        unsafe {
6443            b.launch(cfg)?;
6444        }
6445        Ok(())
6446    }
6447
6448    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
6449    pub fn i32_iota_from(
6450        &self,
6451        ctr: &CudaSlice<i32>,
6452        dst: &mut CudaSlice<i32>,
6453        n: usize,
6454    ) -> Result<(), Box<dyn std::error::Error>> {
6455        let f = self.func("i32_iota_from");
6456        let cfg = LaunchConfig::for_num_elems(n as u32);
6457        let ni = n as i32;
6458        let __s_b = self.gpu.stream();
6459        let mut b = __s_b.launch_builder(&f);
6460        b.arg(ctr).arg(dst).arg(&ni);
6461        unsafe {
6462            b.launch(cfg)?;
6463        }
6464        Ok(())
6465    }
6466
6467    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
6468    pub fn u32_map_k(
6469        &self,
6470        buf: &mut CudaSlice<u32>,
6471        map: &CudaSlice<u32>,
6472        idx: usize,
6473    ) -> Result<(), Box<dyn std::error::Error>> {
6474        let f = self.func("u32_map_k");
6475        let cfg = LaunchConfig {
6476            grid_dim: (1, 1, 1),
6477            block_dim: (1, 1, 1),
6478            shared_mem_bytes: 0,
6479        };
6480        let ii = idx as i32;
6481        let __s_b = self.gpu.stream();
6482        let mut b = __s_b.launch_builder(&f);
6483        b.arg(buf).arg(map).arg(&ii);
6484        unsafe {
6485            b.launch(cfg)?;
6486        }
6487        Ok(())
6488    }
6489
6490    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
6491    #[allow(clippy::too_many_arguments)]
6492    pub fn u32_pack2(
6493        &self,
6494        a: &CudaSlice<u32>,
6495        off_a: usize,
6496        n1: usize,
6497        b_in: &CudaSlice<u32>,
6498        n2: usize,
6499        out: &mut CudaSlice<u32>,
6500    ) -> Result<(), Box<dyn std::error::Error>> {
6501        let f = self.func("u32_pack2");
6502        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6503        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6504        let __s_b = self.gpu.stream();
6505        let mut b = __s_b.launch_builder(&f);
6506        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6507        unsafe {
6508            b.launch(cfg)?;
6509        }
6510        Ok(())
6511    }
6512
6513    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
6514    pub fn moe_w_exscale(
6515        &self,
6516        w: &mut CudaSlice<f32>,
6517        sel: &CudaSlice<i32>,
6518        s: &CudaSlice<f32>,
6519        n: usize,
6520    ) -> Result<(), Box<dyn std::error::Error>> {
6521        let f = self.func("moe_w_exscale");
6522        let cfg = LaunchConfig::for_num_elems(n as u32);
6523        let ni = n as i32;
6524        let __s_b = self.gpu.stream();
6525        let mut b = __s_b.launch_builder(&f);
6526        b.arg(w).arg(sel).arg(s).arg(&ni);
6527        unsafe {
6528            b.launch(cfg)?;
6529        }
6530        Ok(())
6531    }
6532
6533    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
6534    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
6535    pub fn moe_w_scale_by_expert(
6536        &self,
6537        w: &mut CudaSlice<f32>,
6538        sel: &CudaSlice<i32>,
6539        macros: &CudaSlice<f32>,
6540        n_expert: usize,
6541        n: usize,
6542    ) -> Result<(), Box<dyn std::error::Error>> {
6543        let f = self.func("moe_w_scale_by_expert");
6544        let cfg = LaunchConfig {
6545            grid_dim: (n.div_ceil(64) as u32, 1, 1),
6546            block_dim: (64, 1, 1),
6547            shared_mem_bytes: 0,
6548        };
6549        let (ne, nn) = (n_expert as i32, n as i32);
6550        let __s_b = self.gpu.stream();
6551        let mut b = __s_b.launch_builder(&f);
6552        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6553        unsafe {
6554            b.launch(cfg)?;
6555        }
6556        Ok(())
6557    }
6558
6559    pub fn moe_gate_up_silu8_dev_q8(
6560        &self,
6561        table: &CudaSlice<u64>,
6562        sel: &cudarc::driver::CudaView<i32>,
6563        aq: &CudaSlice<i8>,
6564        ad: &CudaSlice<f32>,
6565        in_f: usize,
6566        n_ff: usize,
6567        n_used: usize,
6568        n_expert: usize,
6569        qt_g: i32,
6570        qt_u: i32,
6571        rb_g: usize,
6572        rb_u: usize,
6573        macros: &CudaSlice<f32>,
6574    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6575        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6576        let (mode, wpb) = GU.get_or_init(|| {
6577            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6578            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6579                .ok()
6580                .and_then(|v| v.parse().ok())
6581                .unwrap_or(4u32)
6582                .clamp(1, 16);
6583            (mode, wpb)
6584        });
6585        let (mode, wpb) = (mode.as_str(), *wpb);
6586        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6587        let (inf, nff, ne, rbg, rbu) = (
6588            in_f as i32,
6589            n_ff as i32,
6590            n_expert as i32,
6591            rb_g as i64,
6592            rb_u as i64,
6593        );
6594        let (f, cfg) = match mode {
6595            "1" | "2" | "4" => {
6596                let rpw: u32 = mode.parse().unwrap();
6597                let f = self.func(match rpw {
6598                    1 => "moe_gate_up_silu8_dev_q8_r1",
6599                    2 => "moe_gate_up_silu8_dev_q8_r2",
6600                    _ => "moe_gate_up_silu8_dev_q8_r4",
6601                });
6602                let rows_per_block = (rpw * wpb) as usize;
6603                let gx = n_ff.div_ceil(rows_per_block) as u32;
6604                (
6605                    f,
6606                    LaunchConfig {
6607                        grid_dim: (gx, n_used as u32, 1),
6608                        block_dim: (32, wpb, 1),
6609                        shared_mem_bytes: 0,
6610                    },
6611                )
6612            }
6613            "j8" if n_used <= 32 => (
6614                self.func("moe_gate_up_silu8_dev_q8_j8"),
6615                LaunchConfig {
6616                    grid_dim: (n_ff as u32, 1, 1),
6617                    block_dim: (32, n_used as u32, 1),
6618                    shared_mem_bytes: 0,
6619                },
6620            ),
6621            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
6622            "vsm2" => {
6623                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6624                let sh = (rb_g + rb_u) as u32;
6625                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6626                f.set_attribute(
6627                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6628                    sh as i32,
6629                )?;
6630                (
6631                    f,
6632                    LaunchConfig {
6633                        grid_dim: (n_ff as u32, n_used as u32, 1),
6634                        block_dim: (32, 1, 1),
6635                        shared_mem_bytes: sh,
6636                    },
6637                )
6638            }
6639            "vsm" => {
6640                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6641                let sh = (rb_g + rb_u) as u32;
6642                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6643                f.set_attribute(
6644                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6645                    sh as i32,
6646                )?;
6647                (
6648                    f,
6649                    LaunchConfig {
6650                        grid_dim: (n_ff as u32, n_used as u32, 1),
6651                        block_dim: (32, 1, 1),
6652                        shared_mem_bytes: sh,
6653                    },
6654                )
6655            }
6656            "sg" => (
6657                self.func("moe_gate_up_silu8_dev_q8_sg"),
6658                LaunchConfig {
6659                    grid_dim: (n_ff as u32, n_used as u32, 1),
6660                    block_dim: (32, 1, 1),
6661                    shared_mem_bytes: 0,
6662                },
6663            ),
6664            "j8sg" if n_used <= 32 => (
6665                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6666                LaunchConfig {
6667                    grid_dim: (n_ff as u32, 1, 1),
6668                    block_dim: (32, n_used as u32, 1),
6669                    shared_mem_bytes: 0,
6670                },
6671            ),
6672            "u64" if in_f == 2048 => (
6673                self.func("moe_gate_up_silu8_dev_q8_u64"),
6674                LaunchConfig {
6675                    grid_dim: (n_ff as u32, n_used as u32, 1),
6676                    block_dim: (32, 1, 1),
6677                    shared_mem_bytes: 0,
6678                },
6679            ),
6680            "gs4" if in_f == 2048 => (
6681                self.func("moe_gate_up_silu8_dev_q8_gs4"),
6682                LaunchConfig {
6683                    grid_dim: (n_ff as u32, n_used as u32, 1),
6684                    block_dim: (32, 4, 1),
6685                    shared_mem_bytes: 0,
6686                },
6687            ),
6688            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
6689            "v" | "" => (
6690                self.func("moe_gate_up_silu8_dev_q8_v"),
6691                LaunchConfig {
6692                    grid_dim: (n_ff as u32, n_used as u32, 1),
6693                    block_dim: (32, 1, 1),
6694                    shared_mem_bytes: 0,
6695                },
6696            ),
6697            "s2" => (
6698                self.func("moe_gate_up_silu8_dev_q8_s2"),
6699                LaunchConfig {
6700                    grid_dim: (n_ff as u32, n_used as u32, 1),
6701                    block_dim: (32, 2, 1),
6702                    shared_mem_bytes: 0,
6703                },
6704            ),
6705            "s2z" => {
6706                let rz = wpb.min(16); // s2z smem tile is [16][2]
6707                (
6708                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
6709                    LaunchConfig {
6710                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6711                        block_dim: (32, 2, rz),
6712                        shared_mem_bytes: 0,
6713                    },
6714                )
6715            }
6716            _ => (
6717                self.func("moe_gate_up_silu8_dev_q8"),
6718                LaunchConfig {
6719                    grid_dim: (n_ff as u32, n_used as u32, 1),
6720                    block_dim: (32, 1, 1),
6721                    shared_mem_bytes: 0,
6722                },
6723            ),
6724        };
6725        let __s_b = self.gpu.stream();
6726        let mut b = __s_b.launch_builder(&f);
6727        b.arg(table)
6728            .arg(sel)
6729            .arg(aq)
6730            .arg(ad)
6731            .arg(&mut act)
6732            .arg(&inf)
6733            .arg(&nff)
6734            .arg(&ne)
6735            .arg(&qt_g)
6736            .arg(&qt_u)
6737            .arg(&rbg)
6738            .arg(&rbu)
6739            .arg(macros);
6740        unsafe {
6741            b.launch(cfg)?;
6742        }
6743        Ok(act)
6744    }
6745
6746    #[allow(clippy::too_many_arguments)]
6747    pub fn moe_down8_fma_dev_q8(
6748        &self,
6749        table: &CudaSlice<u64>,
6750        sel: &cudarc::driver::CudaView<i32>,
6751        w: &cudarc::driver::CudaView<f32>,
6752        aq2: &CudaSlice<i8>,
6753        ad2: &CudaSlice<f32>,
6754        dst: &mut cudarc::driver::CudaViewMut<f32>,
6755        in_f: usize,
6756        out_f: usize,
6757        n_used: usize,
6758        n_expert: usize,
6759        qt: i32,
6760        rb: usize,
6761    ) -> Result<(), Box<dyn std::error::Error>> {
6762        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6763        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6764        let (inf, outf, nu, ne, rbi) = (
6765            in_f as i32,
6766            out_f as i32,
6767            n_used as i32,
6768            n_expert as i32,
6769            rb as i64,
6770        );
6771        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
6772        // the h2 twins are nsb==16 (in_f==512) shape-gated.
6773        let (f, cfg) = match mode.as_str() {
6774            m @ ("1" | "2" | "4") if n_used <= 8 => {
6775                let rpw: usize = m.parse().unwrap();
6776                let f = self.func(match rpw {
6777                    1 => "moe_down8_fma_dev_q8_w8r1",
6778                    2 => "moe_down8_fma_dev_q8_w8r2",
6779                    _ => "moe_down8_fma_dev_q8_w8r4",
6780                });
6781                (
6782                    f,
6783                    LaunchConfig {
6784                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6785                        block_dim: (32, n_used as u32, 1),
6786                        shared_mem_bytes: 0,
6787                    },
6788                )
6789            }
6790            "h2" if in_f == 512 => (
6791                self.func("moe_down8_fma_dev_q8_h2"),
6792                LaunchConfig {
6793                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6794                    block_dim: (32, 1, 1),
6795                    shared_mem_bytes: 0,
6796                },
6797            ),
6798            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6799            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6800            "" if in_f == 704 && n_used <= 8 => (
6801                self.func("moe_down8_fma_dev_q8_w8r2"),
6802                LaunchConfig {
6803                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6804                    block_dim: (32, n_used as u32, 1),
6805                    shared_mem_bytes: 0,
6806                },
6807            ),
6808            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6809            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6810            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6811            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6812                self.func("moe_down8_fma_dev_q8_w8h2v"),
6813                LaunchConfig {
6814                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6815                    block_dim: (32, n_used as u32, 1),
6816                    shared_mem_bytes: 0,
6817                },
6818            ),
6819            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6820                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6821                LaunchConfig {
6822                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6823                    block_dim: (32, n_used as u32, 1),
6824                    shared_mem_bytes: 0,
6825                },
6826            ),
6827            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6828                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6829                LaunchConfig {
6830                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6831                    block_dim: (32, n_used as u32, 1),
6832                    shared_mem_bytes: 0,
6833                },
6834            ),
6835            "w8h2" if in_f == 512 && n_used <= 8 => (
6836                self.func("moe_down8_fma_dev_q8_w8h2"),
6837                LaunchConfig {
6838                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6839                    block_dim: (32, n_used as u32, 1),
6840                    shared_mem_bytes: 0,
6841                },
6842            ),
6843            _ => (
6844                self.func("moe_down8_fma_dev_q8"),
6845                LaunchConfig {
6846                    grid_dim: (out_f as u32, 1, 1),
6847                    block_dim: (32, 1, 1),
6848                    shared_mem_bytes: 0,
6849                },
6850            ),
6851        };
6852        let __s_b = self.gpu.stream();
6853        let mut b = __s_b.launch_builder(&f);
6854        b.arg(table)
6855            .arg(sel)
6856            .arg(w)
6857            .arg(aq2)
6858            .arg(ad2)
6859            .arg(dst)
6860            .arg(&inf)
6861            .arg(&outf)
6862            .arg(&nu)
6863            .arg(&ne)
6864            .arg(&qt)
6865            .arg(&rbi);
6866        unsafe {
6867            b.launch(cfg)?;
6868        }
6869        Ok(())
6870    }
6871
6872    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6873    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6874    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6875    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6876    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6877    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6878    #[allow(clippy::too_many_arguments)]
6879    pub fn moe_gate_up_silu8_dev_q8_rows(
6880        &self,
6881        table: &CudaSlice<u64>,
6882        sel: &CudaSlice<i32>,
6883        aq: &CudaSlice<i8>,
6884        ad: &CudaSlice<f32>,
6885        t: usize,
6886        in_f: usize,
6887        n_ff: usize,
6888        n_used: usize,
6889        n_expert: usize,
6890        qt_g: i32,
6891        qt_u: i32,
6892        rb_g: usize,
6893        rb_u: usize,
6894        macros: &CudaSlice<f32>,
6895    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6896        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6897        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6898        let cfg = LaunchConfig {
6899            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6900            block_dim: (32, 1, 1),
6901            shared_mem_bytes: 0,
6902        };
6903        let (inf, nff, ne, nu, rbg, rbu) = (
6904            in_f as i32,
6905            n_ff as i32,
6906            n_expert as i32,
6907            n_used as i32,
6908            rb_g as i64,
6909            rb_u as i64,
6910        );
6911        let __s_b = self.gpu.stream();
6912        let mut b = __s_b.launch_builder(&f);
6913        b.arg(table)
6914            .arg(sel)
6915            .arg(aq)
6916            .arg(ad)
6917            .arg(&mut act)
6918            .arg(&inf)
6919            .arg(&nff)
6920            .arg(&ne)
6921            .arg(&qt_g)
6922            .arg(&qt_u)
6923            .arg(&rbg)
6924            .arg(&rbu)
6925            .arg(&nu)
6926            .arg(macros);
6927        unsafe {
6928            b.launch(cfg)?;
6929        }
6930        Ok(act)
6931    }
6932
6933    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6934    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6935    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6936    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6937    #[allow(clippy::too_many_arguments)]
6938    pub fn moe_down8_fma_dev_q8_rows(
6939        &self,
6940        table: &CudaSlice<u64>,
6941        sel: &CudaSlice<i32>,
6942        w: &CudaSlice<f32>,
6943        aq2: &CudaSlice<i8>,
6944        ad2: &CudaSlice<f32>,
6945        dst: &mut CudaSlice<f32>,
6946        t: usize,
6947        in_f: usize,
6948        out_f: usize,
6949        n_used: usize,
6950        n_expert: usize,
6951        qt: i32,
6952        rb: usize,
6953    ) -> Result<(), Box<dyn std::error::Error>> {
6954        assert!(
6955            in_f == 512 && n_used <= 8,
6956            "down rows twin is w8h2v shape-gated"
6957        );
6958        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6959        let cfg = LaunchConfig {
6960            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6961            block_dim: (32, n_used as u32, 1),
6962            shared_mem_bytes: 0,
6963        };
6964        let (inf, outf, nu, ne, rbi) = (
6965            in_f as i32,
6966            out_f as i32,
6967            n_used as i32,
6968            n_expert as i32,
6969            rb as i64,
6970        );
6971        let __s_b = self.gpu.stream();
6972        let mut b = __s_b.launch_builder(&f);
6973        b.arg(table)
6974            .arg(sel)
6975            .arg(w)
6976            .arg(aq2)
6977            .arg(ad2)
6978            .arg(dst)
6979            .arg(&inf)
6980            .arg(&outf)
6981            .arg(&nu)
6982            .arg(&ne)
6983            .arg(&qt)
6984            .arg(&rbi);
6985        unsafe {
6986            b.launch(cfg)?;
6987        }
6988        Ok(())
6989    }
6990
6991    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6992    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6993    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6994    #[allow(clippy::too_many_arguments)]
6995    pub fn moe_gate_up_silu8_dev_q8_csr(
6996        &self,
6997        table: &CudaSlice<u64>,
6998        sel: &CudaSlice<i32>,
6999        aq: &CudaSlice<i8>,
7000        ad: &CudaSlice<f32>,
7001        n_pairs: usize,
7002        in_f: usize,
7003        n_ff: usize,
7004        n_used: usize,
7005        n_expert: usize,
7006        qt_g: i32,
7007        qt_u: i32,
7008        rb_g: usize,
7009        rb_u: usize,
7010    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7011        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
7012        // host gate guarantees qt_g == qt_u within a supported class.
7013        let f = if qt_g == crate::QT_NVFP4 {
7014            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
7015        } else {
7016            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
7017        };
7018        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7019        let cfg = LaunchConfig {
7020            grid_dim: (n_ff as u32, n_pairs as u32, 1),
7021            block_dim: (32, 1, 1),
7022            shared_mem_bytes: 0,
7023        };
7024        let (inf, nff, ne, nu, npi, rbg, rbu) = (
7025            in_f as i32,
7026            n_ff as i32,
7027            n_expert as i32,
7028            n_used as i32,
7029            n_pairs as i32,
7030            rb_g as i64,
7031            rb_u as i64,
7032        );
7033        let __s_b = self.gpu.stream();
7034        let mut b = __s_b.launch_builder(&f);
7035        b.arg(table)
7036            .arg(sel)
7037            .arg(aq)
7038            .arg(ad)
7039            .arg(&mut act)
7040            .arg(&inf)
7041            .arg(&nff)
7042            .arg(&ne)
7043            .arg(&qt_g)
7044            .arg(&qt_u)
7045            .arg(&rbg)
7046            .arg(&rbu)
7047            .arg(&nu)
7048            .arg(&npi);
7049        unsafe {
7050            b.launch(cfg)?;
7051        }
7052        Ok(act)
7053    }
7054
7055    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
7056    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
7057    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
7058    #[allow(clippy::too_many_arguments)]
7059    pub fn moe_down8_fma_dev_q8_variant(
7060        &self,
7061        variant: &str,
7062        table: &CudaSlice<u64>,
7063        sel: &cudarc::driver::CudaView<i32>,
7064        w: &cudarc::driver::CudaView<f32>,
7065        aq2: &CudaSlice<i8>,
7066        ad2: &CudaSlice<f32>,
7067        dst: &mut cudarc::driver::CudaViewMut<f32>,
7068        in_f: usize,
7069        out_f: usize,
7070        n_used: usize,
7071        n_expert: usize,
7072        qt: i32,
7073        rb: usize,
7074    ) -> Result<(), Box<dyn std::error::Error>> {
7075        let (inf, outf, nu, ne, rbi) = (
7076            in_f as i32,
7077            out_f as i32,
7078            n_used as i32,
7079            n_expert as i32,
7080            rb as i64,
7081        );
7082        let (f, cfg) = match variant {
7083            "w8h2" | "w8h2v" => (
7084                self.func(if variant == "w8h2" {
7085                    "moe_down8_fma_dev_q8_w8h2"
7086                } else {
7087                    "moe_down8_fma_dev_q8_w8h2v"
7088                }),
7089                LaunchConfig {
7090                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7091                    block_dim: (32, n_used as u32, 1),
7092                    shared_mem_bytes: 0,
7093                },
7094            ),
7095            "w8h2r2" | "w8h2r2v" => (
7096                self.func(if variant == "w8h2r2" {
7097                    "moe_down8_fma_dev_q8_w8h2r2"
7098                } else {
7099                    "moe_down8_fma_dev_q8_w8h2r2v"
7100                }),
7101                LaunchConfig {
7102                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7103                    block_dim: (32, n_used as u32, 1),
7104                    shared_mem_bytes: 0,
7105                },
7106            ),
7107            _ => (
7108                self.func("moe_down8_fma_dev_q8"),
7109                LaunchConfig {
7110                    grid_dim: (out_f as u32, 1, 1),
7111                    block_dim: (32, 1, 1),
7112                    shared_mem_bytes: 0,
7113                },
7114            ),
7115        };
7116        let __s_b = self.gpu.stream();
7117        let mut b = __s_b.launch_builder(&f);
7118        b.arg(table)
7119            .arg(sel)
7120            .arg(w)
7121            .arg(aq2)
7122            .arg(ad2)
7123            .arg(dst)
7124            .arg(&inf)
7125            .arg(&outf)
7126            .arg(&nu)
7127            .arg(&ne)
7128            .arg(&qt)
7129            .arg(&rbi);
7130        unsafe {
7131            b.launch(cfg)?;
7132        }
7133        Ok(())
7134    }
7135
7136    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
7137    #[allow(clippy::too_many_arguments)]
7138    pub fn moe_gate_up_silu8_dev_q8_variant(
7139        &self,
7140        variant: &str,
7141        table: &CudaSlice<u64>,
7142        sel: &cudarc::driver::CudaView<i32>,
7143        aq: &CudaSlice<i8>,
7144        ad: &CudaSlice<f32>,
7145        in_f: usize,
7146        n_ff: usize,
7147        n_used: usize,
7148        n_expert: usize,
7149        qt_g: i32,
7150        qt_u: i32,
7151        rb_g: usize,
7152        rb_u: usize,
7153    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7154        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7155        let (inf, nff, ne, rbg, rbu) = (
7156            in_f as i32,
7157            n_ff as i32,
7158            n_expert as i32,
7159            rb_g as i64,
7160            rb_u as i64,
7161        );
7162        let f = self.func(if variant == "v" {
7163            "moe_gate_up_silu8_dev_q8_v"
7164        } else {
7165            "moe_gate_up_silu8_dev_q8"
7166        });
7167        let cfg = LaunchConfig {
7168            grid_dim: (n_ff as u32, n_used as u32, 1),
7169            block_dim: (32, 1, 1),
7170            shared_mem_bytes: 0,
7171        };
7172        let __s_b = self.gpu.stream();
7173        let mut b = __s_b.launch_builder(&f);
7174        b.arg(table)
7175            .arg(sel)
7176            .arg(aq)
7177            .arg(ad)
7178            .arg(&mut act)
7179            .arg(&inf)
7180            .arg(&nff)
7181            .arg(&ne)
7182            .arg(&qt_g)
7183            .arg(&qt_u)
7184            .arg(&rbg)
7185            .arg(&rbu);
7186        unsafe {
7187            b.launch(cfg)?;
7188        }
7189        Ok(act)
7190    }
7191
7192    pub fn moe_gate_up_silu8_dev(
7193        &self,
7194        table: &CudaSlice<u64>,
7195        sel: &cudarc::driver::CudaView<i32>,
7196        x: &cudarc::driver::CudaView<f32>,
7197        in_f: usize,
7198        n_ff: usize,
7199        n_used: usize,
7200        n_expert: usize,
7201        qt_g: i32,
7202        qt_u: i32,
7203        rb_g: usize,
7204        rb_u: usize,
7205        macros: &CudaSlice<f32>,
7206    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7207        let f = self.func("moe_gate_up_silu8_dev");
7208        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
7209        let cfg = LaunchConfig {
7210            grid_dim: (n_ff as u32, n_used as u32, 1),
7211            block_dim: (256, 1, 1),
7212            shared_mem_bytes: 0,
7213        };
7214        let (inf, nff, ne, rbg, rbu) = (
7215            in_f as i32,
7216            n_ff as i32,
7217            n_expert as i32,
7218            rb_g as i64,
7219            rb_u as i64,
7220        );
7221        let __s_b = self.gpu.stream();
7222        let mut b = __s_b.launch_builder(&f);
7223        b.arg(table)
7224            .arg(sel)
7225            .arg(x)
7226            .arg(&mut act)
7227            .arg(&inf)
7228            .arg(&nff)
7229            .arg(&ne)
7230            .arg(&qt_g)
7231            .arg(&qt_u)
7232            .arg(&rbg)
7233            .arg(&rbu)
7234            .arg(macros);
7235        unsafe {
7236            b.launch(cfg)?;
7237        }
7238        Ok(act)
7239    }
7240
7241    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
7242    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
7243    #[allow(clippy::too_many_arguments)]
7244    pub fn moe_down8_fma_dev(
7245        &self,
7246        table: &CudaSlice<u64>,
7247        sel: &cudarc::driver::CudaView<i32>,
7248        w: &cudarc::driver::CudaView<f32>,
7249        act: &CudaSlice<f32>,
7250        dst: &mut cudarc::driver::CudaViewMut<f32>,
7251        in_f: usize,
7252        out_f: usize,
7253        n_used: usize,
7254        n_expert: usize,
7255        qt: i32,
7256        rb: usize,
7257    ) -> Result<(), Box<dyn std::error::Error>> {
7258        let f = self.func("moe_down8_fma_dev");
7259        let cfg = LaunchConfig {
7260            grid_dim: (out_f as u32, 1, 1),
7261            block_dim: (256, 1, 1),
7262            shared_mem_bytes: 0,
7263        };
7264        let (inf, outf, nu, ne, rbv) = (
7265            in_f as i32,
7266            out_f as i32,
7267            n_used as i32,
7268            n_expert as i32,
7269            rb as i64,
7270        );
7271        let __s_b = self.gpu.stream();
7272        let mut b = __s_b.launch_builder(&f);
7273        b.arg(table)
7274            .arg(sel)
7275            .arg(w)
7276            .arg(act)
7277            .arg(dst)
7278            .arg(&inf)
7279            .arg(&outf)
7280            .arg(&nu)
7281            .arg(&ne)
7282            .arg(&qt)
7283            .arg(&rbv);
7284        unsafe {
7285            b.launch(cfg)?;
7286        }
7287        Ok(())
7288    }
7289
7290    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
7291    pub fn axpy_into(
7292        &self,
7293        src: &CudaSlice<f32>,
7294        alpha: f32,
7295        dst: &mut cudarc::driver::CudaViewMut<f32>,
7296        n: usize,
7297    ) -> Result<(), Box<dyn std::error::Error>> {
7298        let f = self.func("axpy_f32");
7299        let cfg = LaunchConfig::for_num_elems(n as u32);
7300        let (a, ni) = (alpha, n as i32);
7301        let __s_b = self.gpu.stream();
7302        let mut b = __s_b.launch_builder(&f);
7303        b.arg(src).arg(dst).arg(&a).arg(&ni);
7304        unsafe {
7305            b.launch(cfg)?;
7306        }
7307        Ok(())
7308    }
7309
7310    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
7311    pub fn axpy_host_into(
7312        &self,
7313        src: &cudarc::driver::CudaView<'_, f32>,
7314        alpha: f32,
7315        dst: &mut cudarc::driver::CudaViewMut<f32>,
7316        n: usize,
7317    ) -> Result<(), Box<dyn std::error::Error>> {
7318        let f = self.func("axpy_host_f32");
7319        let cfg = LaunchConfig::for_num_elems(n as u32);
7320        let (a, ni) = (alpha, n as i32);
7321        let __s_b = self.gpu.stream();
7322        let mut b = __s_b.launch_builder(&f);
7323        b.arg(src).arg(dst).arg(&a).arg(&ni);
7324        unsafe {
7325            b.launch(cfg)?;
7326        }
7327        Ok(())
7328    }
7329
7330    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
7331    pub fn add_scaled_rows(
7332        &self,
7333        src: &CudaSlice<f32>,
7334        scale: &CudaSlice<f32>,
7335        dst: &mut CudaSlice<f32>,
7336        ncols: usize,
7337        nrows: usize,
7338    ) -> Result<(), Box<dyn std::error::Error>> {
7339        let f = self.func("add_scaled_rows_f32");
7340        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7341        let (nc, nr) = (ncols as i32, nrows as i32);
7342        let __s_b = self.gpu.stream();
7343        let mut b = __s_b.launch_builder(&f);
7344        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7345        unsafe {
7346            b.launch(cfg)?;
7347        }
7348        Ok(())
7349    }
7350
7351    // ======== A2 GROUPED MoE PREFILL KERNELS ========
7352
7353    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
7354    pub fn gather_rows(
7355        &self,
7356        src: &CudaSlice<f32>,
7357        idx: &CudaSlice<i32>,
7358        dst: &mut CudaSlice<f32>,
7359        ncols: usize,
7360        m_e: usize,
7361    ) -> Result<(), Box<dyn std::error::Error>> {
7362        let f = self.func("gather_rows_f32");
7363        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7364        let (nc, me) = (ncols as i32, m_e as i32);
7365        let __s_b = self.gpu.stream();
7366        let mut b = __s_b.launch_builder(&f);
7367        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7368        unsafe {
7369            b.launch(cfg)?;
7370        }
7371        Ok(())
7372    }
7373
7374    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
7375    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
7376    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
7377    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
7378    pub fn scatter_slot(
7379        &self,
7380        src: &CudaSlice<f32>,
7381        tok_idx: &CudaSlice<i32>,
7382        slot_idx: &CudaSlice<i32>,
7383        weight: &CudaSlice<f32>,
7384        dst: &mut CudaSlice<f32>,
7385        wbuf: &mut CudaSlice<f32>,
7386        ncols: usize,
7387        n_used: usize,
7388        m_e: usize,
7389    ) -> Result<(), Box<dyn std::error::Error>> {
7390        let f = self.func("scatter_add_slot_f32");
7391        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7392        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7393        let __s_b = self.gpu.stream();
7394        let mut b = __s_b.launch_builder(&f);
7395        b.arg(src)
7396            .arg(tok_idx)
7397            .arg(slot_idx)
7398            .arg(weight)
7399            .arg(dst)
7400            .arg(wbuf)
7401            .arg(&nc)
7402            .arg(&nu)
7403            .arg(&me);
7404        unsafe {
7405            b.launch(cfg)?;
7406        }
7407        Ok(())
7408    }
7409
7410    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
7411    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
7412    /// Uses FMA for bit-identity with the sequential axpy path.
7413    pub fn reduce_slots(
7414        &self,
7415        slots: &CudaSlice<f32>,
7416        wbuf: &CudaSlice<f32>,
7417        dst: &mut CudaSlice<f32>,
7418        ncols: usize,
7419        n_used: usize,
7420        t: usize,
7421    ) -> Result<(), Box<dyn std::error::Error>> {
7422        let f = self.func("reduce_slots_f32");
7423        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7424        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7425        let __s_b = self.gpu.stream();
7426        let mut b = __s_b.launch_builder(&f);
7427        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7428        unsafe {
7429            b.launch(cfg)?;
7430        }
7431        Ok(())
7432    }
7433
7434    /// Canonical slot-order reduction with separately rounded multiply and add.
7435    ///
7436    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
7437    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
7438    pub fn reduce_slots_host(
7439        &self,
7440        slots: &CudaSlice<f32>,
7441        wbuf: &CudaSlice<f32>,
7442        dst: &mut CudaSlice<f32>,
7443        ncols: usize,
7444        n_used: usize,
7445        t: usize,
7446    ) -> Result<(), Box<dyn std::error::Error>> {
7447        let f = self.func("reduce_slots_host_f32");
7448        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7449        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7450        let __s_b = self.gpu.stream();
7451        let mut b = __s_b.launch_builder(&f);
7452        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7453        unsafe {
7454            b.launch(cfg)?;
7455        }
7456        Ok(())
7457    }
7458
7459    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
7460    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
7461    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
7462    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
7463    /// GPU time, ~half of it redundant re-quantization of the same row.
7464    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
7465    pub fn quantize_q8_1_view(
7466        &self,
7467        x: &cudarc::driver::CudaView<f32>,
7468        m: usize,
7469        in_f: usize,
7470    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7471        let f = self.func("quantize_q8_1");
7472        let nblk = in_f / 32;
7473        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7474        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7475        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7476        let (inf, mi) = (in_f as i32, m as i32);
7477        let __s_b = self.gpu.stream();
7478        let mut b = __s_b.launch_builder(&f);
7479        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7480        unsafe {
7481            b.launch(cfg)?;
7482        }
7483        Ok((q, d))
7484    }
7485
7486    pub fn quantize_q8_1(
7487        &self,
7488        x: &CudaSlice<f32>,
7489        m: usize,
7490        in_f: usize,
7491    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7492        let nblk = in_f / 32;
7493        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
7494        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
7495        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
7496        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7497        let (inf, mi) = (in_f as i32, m as i32);
7498        if Self::pdl_on() && Self::pdl_wb_on() {
7499            {
7500                use cudarc::driver::{DevicePtr, DevicePtrMut};
7501                let s = &self.gpu.stream();
7502                let (px, _g0) = x.device_ptr(s);
7503                let (pq, _g1) = q.device_ptr_mut(s);
7504                let (pd, _g2) = d.device_ptr_mut(s);
7505                let mut ps = [
7506                    &px as *const _ as *mut std::ffi::c_void,
7507                    &pq as *const _ as *mut _,
7508                    &pd as *const _ as *mut _,
7509                    &inf as *const _ as *mut _,
7510                    &mi as *const _ as *mut _,
7511                ];
7512                unsafe {
7513                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7514                }
7515            }
7516            return Ok((q, d));
7517        }
7518        let f = self.func("quantize_q8_1");
7519        let __s_b = self.gpu.stream();
7520        let mut b = __s_b.launch_builder(&f);
7521        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7522        unsafe {
7523            b.launch(cfg)?;
7524        }
7525        Ok((q, d))
7526    }
7527
7528    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
7529    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
7530    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
7531    pub fn quantize_fp4_act(
7532        &self,
7533        x: &CudaSlice<f32>,
7534        m: usize,
7535        in_f: usize,
7536    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7537        let f = self.func("quantize_fp4_act");
7538        let nb16 = in_f / 16;
7539        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
7540        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
7541        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
7542        let (inf, mi) = (in_f as i32, m as i32);
7543        let __s_b = self.gpu.stream();
7544        let mut b = __s_b.launch_builder(&f);
7545        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7546        unsafe {
7547            b.launch(cfg)?;
7548        }
7549        Ok((aq4, ad4))
7550    }
7551
7552    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
7553    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
7554    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
7555    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
7556    pub fn qmatvec_gemm_nvfp4_fp4(
7557        &self,
7558        bytes: &CudaSlice<u8>,
7559        x: &CudaSlice<f32>,
7560        m: usize,
7561        in_f: usize,
7562        out_f: usize,
7563        row_bytes: usize,
7564        scale: f32,
7565    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7566        assert!(
7567            in_f % 64 == 0,
7568            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7569        );
7570        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7571        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7572        if scale != 1.0 {
7573            self.scale_inplace(&mut y, scale, m * out_f)?;
7574        }
7575        Ok(y)
7576    }
7577
7578    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
7579    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
7580    fn fp4_gemm_launch(
7581        &self,
7582        bytes: &CudaSlice<u8>,
7583        aq4: &CudaSlice<u32>,
7584        ad4: &CudaSlice<u8>,
7585        m: usize,
7586        in_f: usize,
7587        out_f: usize,
7588        row_bytes: usize,
7589    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7590        let f = self.func("qmatvec_gemm_nvfp4_fp4");
7591        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7592        const BM: u32 = 64;
7593        const BN: u32 = 256;
7594        let cfg = LaunchConfig {
7595            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7596            block_dim: (32, 4, 1),
7597            shared_mem_bytes: 0,
7598        };
7599        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7600        let __s_b = self.gpu.stream();
7601        let mut b = __s_b.launch_builder(&f);
7602        b.arg(bytes)
7603            .arg(aq4)
7604            .arg(ad4)
7605            .arg(&mut y)
7606            .arg(&inf)
7607            .arg(&outf)
7608            .arg(&mi)
7609            .arg(&rb);
7610        unsafe {
7611            b.launch(cfg)?;
7612        }
7613        Ok(y)
7614    }
7615
7616    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
7617    pub fn qmatvec_gemm_nvfp4_fp4_raw(
7618        &self,
7619        bytes: &CudaSlice<u8>,
7620        x: &CudaSlice<f32>,
7621        m: usize,
7622        in_f: usize,
7623        out_f: usize,
7624        row_bytes: usize,
7625    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7626        assert!(
7627            in_f % 64 == 0,
7628            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7629        );
7630        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7631        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7632    }
7633
7634    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
7635    pub fn qmatvec_q8_0_fast(
7636        &self,
7637        w: &CudaSlice<u8>,
7638        x: &CudaSlice<f32>,
7639        m: usize,
7640        in_f: usize,
7641        out_f: usize,
7642        row_bytes: usize,
7643    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7644        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7645        let f = self.func("qmatvec_q8_0_dp4a");
7646        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7647        let cfg = LaunchConfig {
7648            grid_dim: (out_f as u32, m as u32, 1),
7649            block_dim: (128, 1, 1),
7650            shared_mem_bytes: 0,
7651        };
7652        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7653        let __s_b = self.gpu.stream();
7654        let mut b = __s_b.launch_builder(&f);
7655        b.arg(w)
7656            .arg(&aq)
7657            .arg(&ad)
7658            .arg(&mut y)
7659            .arg(&inf)
7660            .arg(&outf)
7661            .arg(&mi)
7662            .arg(&rb);
7663        unsafe {
7664            b.launch(cfg)?;
7665        }
7666        Ok(y)
7667    }
7668
7669    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7670    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7671    pub fn qmatvec_q4_K_fast(
7672        &self,
7673        w: &CudaSlice<u8>,
7674        x: &CudaSlice<f32>,
7675        m: usize,
7676        in_f: usize,
7677        out_f: usize,
7678        row_bytes: usize,
7679    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7680        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7681        let f = self.func("qmatvec_q4_K_dp4a");
7682        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7683        let cfg = LaunchConfig {
7684            grid_dim: (out_f as u32, m as u32, 1),
7685            block_dim: (128, 1, 1),
7686            shared_mem_bytes: 0,
7687        };
7688        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7689        let __s_b = self.gpu.stream();
7690        let mut b = __s_b.launch_builder(&f);
7691        b.arg(w)
7692            .arg(&aq)
7693            .arg(&ad)
7694            .arg(&mut y)
7695            .arg(&inf)
7696            .arg(&outf)
7697            .arg(&mi)
7698            .arg(&rb);
7699        unsafe {
7700            b.launch(cfg)?;
7701        }
7702        Ok(y)
7703    }
7704
7705    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7706    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7707    pub fn qmatvec_q6_K_fast(
7708        &self,
7709        w: &CudaSlice<u8>,
7710        x: &CudaSlice<f32>,
7711        m: usize,
7712        in_f: usize,
7713        out_f: usize,
7714        row_bytes: usize,
7715    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7716        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7717        let f = self.func("qmatvec_q6_K_dp4a");
7718        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7719        let cfg = LaunchConfig {
7720            grid_dim: (out_f as u32, m as u32, 1),
7721            block_dim: (128, 1, 1),
7722            shared_mem_bytes: 0,
7723        };
7724        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7725        let __s_b = self.gpu.stream();
7726        let mut b = __s_b.launch_builder(&f);
7727        b.arg(w)
7728            .arg(&aq)
7729            .arg(&ad)
7730            .arg(&mut y)
7731            .arg(&inf)
7732            .arg(&outf)
7733            .arg(&mi)
7734            .arg(&rb);
7735        unsafe {
7736            b.launch(cfg)?;
7737        }
7738        Ok(y)
7739    }
7740
7741    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7742    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7743    pub fn qmatvec_q5_K_fast(
7744        &self,
7745        w: &CudaSlice<u8>,
7746        x: &CudaSlice<f32>,
7747        m: usize,
7748        in_f: usize,
7749        out_f: usize,
7750        row_bytes: usize,
7751    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7752        self.qmatvec_dp4a_named(
7753            "qmatvec_q5_K_dp4a",
7754            &w.slice(0..w.len()),
7755            x,
7756            m,
7757            in_f,
7758            out_f,
7759            row_bytes,
7760        )
7761    }
7762    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7763    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7764    pub fn qmatvec_q3_K_fast(
7765        &self,
7766        w: &CudaSlice<u8>,
7767        x: &CudaSlice<f32>,
7768        m: usize,
7769        in_f: usize,
7770        out_f: usize,
7771        row_bytes: usize,
7772    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7773        self.qmatvec_dp4a_named(
7774            "qmatvec_q3_K_dp4a",
7775            &w.slice(0..w.len()),
7776            x,
7777            m,
7778            in_f,
7779            out_f,
7780            row_bytes,
7781        )
7782    }
7783    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
7784    pub fn qmatvec_nvfp4_fast_rp(
7785        &self,
7786        w: &CudaSlice<u8>,
7787        x: &CudaSlice<f32>,
7788        m: usize,
7789        in_f: usize,
7790        out_f: usize,
7791        row_bytes: usize,
7792    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7793        assert!(
7794            in_f % 64 == 0,
7795            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7796        );
7797        self.qmatvec_dp4a_named(
7798            "qmatvec_nvfp4_dp4a_rp",
7799            &w.slice(0..w.len()),
7800            x,
7801            m,
7802            in_f,
7803            out_f,
7804            row_bytes,
7805        )
7806    }
7807    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
7808    pub fn qmatvec_nvfp4_fast(
7809        &self,
7810        w: &cudarc::driver::CudaView<'_, u8>,
7811        x: &CudaSlice<f32>,
7812        m: usize,
7813        in_f: usize,
7814        out_f: usize,
7815        row_bytes: usize,
7816    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7817        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
7818        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
7819        assert!(
7820            in_f % 64 == 0,
7821            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7822        );
7823        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7824    }
7825    /// v2-layout twin of `qmatvec_nvfp4_fast` for the slot-major expert banks
7826    /// (MEMRA_NVFP4_BANK_V2) — bit-identical per row, coalesced reads.
7827    pub fn qmatvec_nvfp4_fast_v2(
7828        &self,
7829        w: &cudarc::driver::CudaView<'_, u8>,
7830        x: &CudaSlice<f32>,
7831        m: usize,
7832        in_f: usize,
7833        out_f: usize,
7834        row_bytes: usize,
7835    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7836        assert!(
7837            in_f % 64 == 0,
7838            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7839        );
7840        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7841    }
7842    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
7843    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7844    pub fn qmatvec_iq4_XS_fast(
7845        &self,
7846        w: &CudaSlice<u8>,
7847        x: &CudaSlice<f32>,
7848        m: usize,
7849        in_f: usize,
7850        out_f: usize,
7851        row_bytes: usize,
7852    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7853        self.qmatvec_dp4a_named(
7854            "qmatvec_iq4_XS_dp4a",
7855            &w.slice(0..w.len()),
7856            x,
7857            m,
7858            in_f,
7859            out_f,
7860            row_bytes,
7861        )
7862    }
7863
7864    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
7865    fn qmatvec_dp4a_named(
7866        &self,
7867        name: &str,
7868        w: &cudarc::driver::CudaView<'_, u8>,
7869        x: &CudaSlice<f32>,
7870        m: usize,
7871        in_f: usize,
7872        out_f: usize,
7873        row_bytes: usize,
7874    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7875        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7876        let f = self.func(name);
7877        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7878        let cfg = LaunchConfig {
7879            grid_dim: (out_f as u32, m as u32, 1),
7880            block_dim: (128, 1, 1),
7881            shared_mem_bytes: 0,
7882        };
7883        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7884        let __s_b = self.gpu.stream();
7885        let mut b = __s_b.launch_builder(&f);
7886        b.arg(w)
7887            .arg(&aq)
7888            .arg(&ad)
7889            .arg(&mut y)
7890            .arg(&inf)
7891            .arg(&outf)
7892            .arg(&mi)
7893            .arg(&rb);
7894        unsafe {
7895            b.launch(cfg)?;
7896        }
7897        Ok(y)
7898    }
7899
7900    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
7901    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
7902    /// its output); this entry exists so a routed-expert program can quantize one activation
7903    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
7904    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
7905    #[allow(clippy::too_many_arguments)]
7906    pub fn qmatvec_nvfp4_fast_prequant_into(
7907        &self,
7908        w: &CudaSlice<u8>,
7909        aq: &CudaSlice<i8>,
7910        ad: &CudaSlice<f32>,
7911        y: &mut CudaSlice<f32>,
7912        m: usize,
7913        in_f: usize,
7914        out_f: usize,
7915        row_bytes: usize,
7916    ) -> Result<(), Box<dyn std::error::Error>> {
7917        assert!(
7918            in_f % 64 == 0,
7919            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7920        );
7921        if y.len() < m * out_f {
7922            return Err(format!(
7923                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7924                y.len()
7925            )
7926            .into());
7927        }
7928        let f = self.func("qmatvec_nvfp4_dp4a");
7929        let cfg = LaunchConfig {
7930            grid_dim: (out_f as u32, m as u32, 1),
7931            block_dim: (128, 1, 1),
7932            shared_mem_bytes: 0,
7933        };
7934        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7935        let __s_b = self.gpu.stream();
7936        let mut b = __s_b.launch_builder(&f);
7937        b.arg(w)
7938            .arg(aq)
7939            .arg(ad)
7940            .arg(y)
7941            .arg(&inf)
7942            .arg(&outf)
7943            .arg(&mi)
7944            .arg(&rb);
7945        unsafe {
7946            b.launch(cfg)?;
7947        }
7948        Ok(())
7949    }
7950
7951    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
7952    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
7953    #[allow(clippy::too_many_arguments)]
7954    pub fn matvec_f32_qkv_into(
7955        &self,
7956        wq: &CudaSlice<f32>,
7957        wk: &CudaSlice<f32>,
7958        wv: &CudaSlice<f32>,
7959        wg: &CudaSlice<f32>,
7960        x: &CudaSlice<f32>,
7961        yq: &mut CudaSlice<f32>,
7962        yk: &mut CudaSlice<f32>,
7963        yv: &mut CudaSlice<f32>,
7964        yg: &mut CudaSlice<f32>,
7965        in_f: usize,
7966        out_q: usize,
7967        out_kv: usize,
7968        out_g: usize,
7969    ) -> Result<(), Box<dyn std::error::Error>> {
7970        if in_f % 4 != 0
7971            || wq.len() != out_q * in_f
7972            || wk.len() != out_kv * in_f
7973            || wv.len() != out_kv * in_f
7974            || wg.len() < out_g * in_f
7975            || x.len() < in_f
7976            || yq.len() < out_q
7977            || yk.len() < out_kv
7978            || yv.len() < out_kv
7979            || (out_g > 0 && yg.len() < out_g)
7980        {
7981            return Err(format!(
7982                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7983                 wq={} wk={} wv={} wg={}",
7984                wq.len(),
7985                wk.len(),
7986                wv.len(),
7987                wg.len()
7988            )
7989            .into());
7990        }
7991        let f = self.func("matvec_f32_qkv");
7992        let cfg = LaunchConfig {
7993            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7994            block_dim: (128, 1, 1),
7995            shared_mem_bytes: 0,
7996        };
7997        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7998        let __s_b = self.gpu.stream();
7999        let mut b = __s_b.launch_builder(&f);
8000        b.arg(wq)
8001            .arg(wk)
8002            .arg(wv)
8003            .arg(wg)
8004            .arg(x)
8005            .arg(yq)
8006            .arg(yk)
8007            .arg(yv)
8008            .arg(yg)
8009            .arg(&inf)
8010            .arg(&oq)
8011            .arg(&okv)
8012            .arg(&og);
8013        unsafe {
8014            b.launch(cfg)?;
8015        }
8016        Ok(())
8017    }
8018
8019    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
8020    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
8021    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
8022    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
8023    /// kernel — the batching only removes host launch latency.
8024    #[allow(clippy::too_many_arguments)]
8025    /// FUSION #2a: gate+up sweeps in one launch (v2 banks only; identical geometry both
8026    /// banks, caller-guarded). Per-row bit-identical to two qmatvec_nvfp4_sel_into calls.
8027    #[allow(clippy::too_many_arguments)]
8028    pub fn qmatvec_nvfp4_sel_gu_into(
8029        &self,
8030        gate_bank: &CudaSlice<u8>,
8031        up_bank: &CudaSlice<u8>,
8032        sel: &CudaSlice<i32>,
8033        aq: &CudaSlice<i8>,
8034        ad: &CudaSlice<f32>,
8035        yg: &mut CudaSlice<f32>,
8036        yu: &mut CudaSlice<f32>,
8037        n_sel: usize,
8038        in_f: usize,
8039        out_f: usize,
8040        row_bytes: usize,
8041        expert_stride: usize,
8042    ) -> Result<(), Box<dyn std::error::Error>> {
8043        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8044        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8045            return Err("NVFP4 gu sel geometry".into());
8046        }
8047        // MEMRA_SEL_GU_RPW=2|4: multirow twin (activation group read once, reused across
8048        // RPW rows' gate+up dots) — bit-identical per row, one block per RPW rows.
8049        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
8050        let rpw = *RPW.get_or_init(|| {
8051            std::env::var("MEMRA_SEL_GU_RPW")
8052                .ok()
8053                .and_then(|v| v.parse().ok())
8054                .filter(|r| *r == 2 || *r == 4)
8055                .unwrap_or(1)
8056        });
8057        let rpw = if out_f % rpw == 0 { rpw } else { 1 };
8058        // MEMRA_SEL_GU_WPR=1: warp-per-row (NUMERIC-CLASS — per-row reduction order changes;
8059        // acceptance is the argmax gate + battery, the QKV_FUSED/BF16_MMV class).
8060        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8061        let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
8062        let f = self.func(match (wpr, rpw) {
8063            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
8064            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
8065            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
8066            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
8067        });
8068        let cfg = LaunchConfig {
8069            grid_dim: if wpr {
8070                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
8071            } else if rpw == 1 {
8072                ((2 * out_f) as u32, n_sel as u32, 1)
8073            } else {
8074                ((out_f / rpw) as u32, n_sel as u32, 1)
8075            },
8076            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
8077            shared_mem_bytes: 0,
8078        };
8079        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8080        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8081        let (ars, adrs) = (0i64, 0i64);
8082        let __s_b = self.gpu.stream();
8083        let mut b = __s_b.launch_builder(&f);
8084        b.arg(gate_bank)
8085            .arg(up_bank)
8086            .arg(sel)
8087            .arg(aq)
8088            .arg(ad)
8089            .arg(yg)
8090            .arg(yu)
8091            .arg(&inf)
8092            .arg(&outf)
8093            .arg(&ns)
8094            .arg(&rb)
8095            .arg(&es)
8096            .arg(&ars)
8097            .arg(&adrs);
8098        unsafe {
8099            b.launch(cfg)?;
8100        }
8101        Ok(())
8102    }
8103
8104    /// MEMRA_SEL_DOWN8=1: the DOWN sweep and the route-weight combine in ONE launch
8105    /// (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm ported to the
8106    /// NVFP4 banks). Block = (32, n_sel): one warp per slot instead of one warp per
8107    /// (row, slot), and the n_sel x out_f partial buffer disappears. Bit-identical to
8108    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce
8109    /// tree, same slot-ordered chain. Requires the v2 banks and nsb <= 32 (the fit-block
8110    /// class the reduce identity is argued at).
8111    #[allow(clippy::too_many_arguments)]
8112    pub fn qmatvec_nvfp4_sel_down8_into(
8113        &self,
8114        bank: &CudaSlice<u8>,
8115        sel: &CudaSlice<i32>,
8116        aq: &CudaSlice<i8>,
8117        ad: &CudaSlice<f32>,
8118        route_w: &CudaSlice<f32>,
8119        md: &CudaSlice<f32>,
8120        dst: &mut CudaSlice<f32>,
8121        n_sel: usize,
8122        in_f: usize,
8123        out_f: usize,
8124        row_bytes: usize,
8125        expert_stride: usize,
8126        act_row_stride: usize,
8127        ad_row_stride: usize,
8128    ) -> Result<(), Box<dyn std::error::Error>> {
8129        if in_f % 64 != 0
8130            || n_sel == 0
8131            || n_sel > 8
8132            || (in_f >> 5) > 32
8133            || dst.len() < out_f
8134            || sel.len() < n_sel
8135            || route_w.len() < n_sel
8136        {
8137            return Err(format!(
8138                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
8139                dst.len()
8140            )
8141            .into());
8142        }
8143        if !crate::tp::nvfp4_bank_v2_on() {
8144            return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
8145        }
8146        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
8147        let cfg = LaunchConfig {
8148            grid_dim: (out_f as u32, 1, 1),
8149            block_dim: (32, n_sel as u32, 1),
8150            shared_mem_bytes: 0,
8151        };
8152        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8153        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8154        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8155        let __s_b = self.gpu.stream();
8156        let mut b = __s_b.launch_builder(&f);
8157        b.arg(bank)
8158            .arg(sel)
8159            .arg(aq)
8160            .arg(ad)
8161            .arg(route_w)
8162            .arg(md)
8163            .arg(dst)
8164            .arg(&inf)
8165            .arg(&outf)
8166            .arg(&ns)
8167            .arg(&rb)
8168            .arg(&es)
8169            .arg(&ars)
8170            .arg(&adrs);
8171        unsafe {
8172            b.launch(cfg)?;
8173        }
8174        Ok(())
8175    }
8176
8177    /// T-ROW twin of the down8 fusion (spec verify + batched serving MoE): one block per
8178    /// (output row, token row) = the exact t=1 down8 program per token — bit-identical
8179    /// per row to its own down8/axpy pair at any t.
8180    #[allow(clippy::too_many_arguments)]
8181    pub fn qmatvec_nvfp4_sel_down8_rows_into(
8182        &self,
8183        bank: &CudaSlice<u8>,
8184        sel: &CudaSlice<i32>,
8185        aq: &CudaSlice<i8>,
8186        ad: &CudaSlice<f32>,
8187        route_w: &CudaSlice<f32>,
8188        md: &CudaSlice<f32>,
8189        dst: &mut CudaSlice<f32>,
8190        t: usize,
8191        n_sel_col: usize,
8192        in_f: usize,
8193        out_f: usize,
8194        row_bytes: usize,
8195        expert_stride: usize,
8196        act_row_stride: usize,
8197        ad_row_stride: usize,
8198    ) -> Result<(), Box<dyn std::error::Error>> {
8199        let n_sel = t * n_sel_col;
8200        if in_f % 64 != 0
8201            || n_sel_col == 0
8202            || n_sel_col > 8
8203            || t == 0
8204            || t > 64
8205            || (in_f >> 5) > 32
8206            || dst.len() < t * out_f
8207            || sel.len() < n_sel
8208            || route_w.len() < n_sel
8209        {
8210            return Err(format!(
8211                "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
8212                dst.len()
8213            )
8214            .into());
8215        }
8216        if !crate::tp::nvfp4_bank_v2_on() {
8217            return Err(
8218                "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8219            );
8220        }
8221        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8222        let cfg = LaunchConfig {
8223            grid_dim: (out_f as u32, t as u32, 1),
8224            block_dim: (32, n_sel_col as u32, 1),
8225            shared_mem_bytes: 0,
8226        };
8227        let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8228        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8229        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8230        let __s_b = self.gpu.stream();
8231        let mut b = __s_b.launch_builder(&f);
8232        b.arg(bank)
8233            .arg(sel)
8234            .arg(aq)
8235            .arg(ad)
8236            .arg(route_w)
8237            .arg(md)
8238            .arg(dst)
8239            .arg(&inf)
8240            .arg(&outf)
8241            .arg(&nsc)
8242            .arg(&rb)
8243            .arg(&es)
8244            .arg(&ars)
8245            .arg(&adrs);
8246        unsafe {
8247            b.launch(cfg)?;
8248        }
8249        Ok(())
8250    }
8251
8252    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
8253    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
8254    #[allow(clippy::too_many_arguments)]
8255    pub fn qmatvec_nvfp4_sel_gu_ep_into(
8256        &self,
8257        gate_bank: &CudaSlice<u8>,
8258        up_bank: &CudaSlice<u8>,
8259        sel: &CudaSlice<i32>,
8260        aq: &CudaSlice<i8>,
8261        ad: &CudaSlice<f32>,
8262        yg: &mut CudaSlice<f32>,
8263        yu: &mut CudaSlice<f32>,
8264        n_sel: usize,
8265        in_f: usize,
8266        out_f: usize,
8267        row_bytes: usize,
8268        expert_stride: usize,
8269        owner: usize,
8270    ) -> Result<(), Box<dyn std::error::Error>> {
8271        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8272        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8273            return Err("NVFP4 gu ep geometry".into());
8274        }
8275        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8276        let cfg = LaunchConfig {
8277            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8278            block_dim: (128, 1, 1),
8279            shared_mem_bytes: 0,
8280        };
8281        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8282        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8283        let (ars, adrs) = (0i64, 0i64);
8284        let __s_b = self.gpu.stream();
8285        let mut b = __s_b.launch_builder(&f);
8286        b.arg(gate_bank)
8287            .arg(up_bank)
8288            .arg(sel)
8289            .arg(aq)
8290            .arg(ad)
8291            .arg(yg)
8292            .arg(yu)
8293            .arg(&inf)
8294            .arg(&outf)
8295            .arg(&ns)
8296            .arg(&rb)
8297            .arg(&es)
8298            .arg(&ars)
8299            .arg(&adrs)
8300            .arg(&own);
8301        unsafe {
8302            b.launch(cfg)?;
8303        }
8304        Ok(())
8305    }
8306
8307    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
8308    #[allow(clippy::too_many_arguments)]
8309    pub fn silu_mul_scaled_q8_1_sel_ep_into(
8310        &self,
8311        gate: &CudaSlice<f32>,
8312        up: &CudaSlice<f32>,
8313        gmac: &CudaSlice<f32>,
8314        umac: &CudaSlice<f32>,
8315        sel: &CudaSlice<i32>,
8316        limit: Option<f32>,
8317        out_q: &mut CudaSlice<i8>,
8318        out_d: &mut CudaSlice<f32>,
8319        n_per: usize,
8320        n_sel: usize,
8321        owner: usize,
8322    ) -> Result<(), Box<dyn std::error::Error>> {
8323        if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8324            return Err("NVFP4 silu ep geometry".into());
8325        }
8326        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8327        let warps = n_sel * n_per / 32;
8328        let cfg = LaunchConfig {
8329            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8330            block_dim: (128, 1, 1),
8331            shared_mem_bytes: 0,
8332        };
8333        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8334        let (lim, has) = match limit {
8335            Some(l) => (l, 1i32),
8336            None => (0.0f32, 0i32),
8337        };
8338        let __s_b = self.gpu.stream();
8339        let mut b = __s_b.launch_builder(&f);
8340        b.arg(gate)
8341            .arg(up)
8342            .arg(gmac)
8343            .arg(umac)
8344            .arg(sel)
8345            .arg(&lim)
8346            .arg(&has)
8347            .arg(out_q)
8348            .arg(out_d)
8349            .arg(&np)
8350            .arg(&ns)
8351            .arg(&own);
8352        unsafe {
8353            b.launch(cfg)?;
8354        }
8355        Ok(())
8356    }
8357
8358    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
8359    #[allow(clippy::too_many_arguments)]
8360    pub fn qmatvec_nvfp4_sel_down8_ep_into(
8361        &self,
8362        bank: &CudaSlice<u8>,
8363        sel: &CudaSlice<i32>,
8364        aq: &CudaSlice<i8>,
8365        ad: &CudaSlice<f32>,
8366        route_w: &CudaSlice<f32>,
8367        md: &CudaSlice<f32>,
8368        dst: &mut CudaSlice<f32>,
8369        n_sel: usize,
8370        in_f: usize,
8371        out_f: usize,
8372        row_bytes: usize,
8373        expert_stride: usize,
8374        act_row_stride: usize,
8375        ad_row_stride: usize,
8376        owner: usize,
8377    ) -> Result<(), Box<dyn std::error::Error>> {
8378        if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8379            return Err("NVFP4 down8 ep geometry".into());
8380        }
8381        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8382        let cfg = LaunchConfig {
8383            grid_dim: (out_f as u32, 1, 1),
8384            block_dim: (32, n_sel as u32, 1),
8385            shared_mem_bytes: 0,
8386        };
8387        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8388        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8389        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8390        let __s_b = self.gpu.stream();
8391        let mut b = __s_b.launch_builder(&f);
8392        b.arg(bank)
8393            .arg(sel)
8394            .arg(aq)
8395            .arg(ad)
8396            .arg(route_w)
8397            .arg(md)
8398            .arg(dst)
8399            .arg(&inf)
8400            .arg(&outf)
8401            .arg(&ns)
8402            .arg(&rb)
8403            .arg(&es)
8404            .arg(&ars)
8405            .arg(&adrs)
8406            .arg(&own);
8407        unsafe {
8408            b.launch(cfg)?;
8409        }
8410        Ok(())
8411    }
8412
8413    pub fn qmatvec_nvfp4_sel_into(
8414        &self,
8415        bank: &CudaSlice<u8>,
8416        sel: &CudaSlice<i32>,
8417        aq: &CudaSlice<i8>,
8418        ad: &CudaSlice<f32>,
8419        y: &mut CudaSlice<f32>,
8420        n_sel: usize,
8421        in_f: usize,
8422        out_f: usize,
8423        row_bytes: usize,
8424        expert_stride: usize,
8425        act_row_stride: usize,
8426        ad_row_stride: usize,
8427    ) -> Result<(), Box<dyn std::error::Error>> {
8428        assert!(
8429            in_f % 64 == 0,
8430            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8431        );
8432        if y.len() < n_sel * out_f || sel.len() < n_sel {
8433            return Err(format!(
8434                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8435                y.len(),
8436                sel.len()
8437            )
8438            .into());
8439        }
8440        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
8441        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
8442        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
8443        // sequential-rows variant was flat). Default stays the single-row form.
8444        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
8445        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
8446        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8447        let mode = *MR.get_or_init(|| {
8448            if crate::tp::nvfp4_bank_v2_on() {
8449                3
8450            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8451                2
8452            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8453                1
8454            } else {
8455                0
8456            }
8457        });
8458        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8459        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
8460        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
8461        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
8462        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8463        let v2s = mode == 3
8464            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8465            && row_bytes % 16 == 0
8466            && in_f <= 4096;
8467        let f = match (mode, v2s) {
8468            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8469            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8470            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8471            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8472            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8473        };
8474        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
8475        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
8476        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
8477        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
8478        let nsb = in_f >> 5;
8479        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8480            32
8481        } else if mode == 1 {
8482            512
8483        } else {
8484            128
8485        };
8486        let cfg = LaunchConfig {
8487            grid_dim: (
8488                if v2s {
8489                    (out_f as u32).div_ceil(8)
8490                } else {
8491                    match mode {
8492                        2 => (out_f as u32).div_ceil(16),
8493                        1 => (out_f as u32).div_ceil(4),
8494                        _ => out_f as u32,
8495                    }
8496                },
8497                n_sel as u32,
8498                1,
8499            ),
8500            block_dim: (fit_block, 1, 1),
8501            shared_mem_bytes: 0,
8502        };
8503        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8504        let (rb, es, ars, adrs) = (
8505            row_bytes as i64,
8506            expert_stride as i64,
8507            act_row_stride as i64,
8508            ad_row_stride as i64,
8509        );
8510        let __s_b = self.gpu.stream();
8511        let mut b = __s_b.launch_builder(&f);
8512        b.arg(bank)
8513            .arg(sel)
8514            .arg(aq)
8515            .arg(ad)
8516            .arg(y)
8517            .arg(&inf)
8518            .arg(&outf)
8519            .arg(&ns)
8520            .arg(&rb)
8521            .arg(&es)
8522            .arg(&ars)
8523            .arg(&adrs);
8524        unsafe {
8525            b.launch(cfg)?;
8526        }
8527        Ok(())
8528    }
8529
8530    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
8531    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
8532    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
8533    /// takes the plain SiLU kernel.
8534    #[allow(clippy::too_many_arguments)]
8535    pub fn silu_mul_scaled_q8_1_sel_into(
8536        &self,
8537        gate: &CudaSlice<f32>,
8538        up: &CudaSlice<f32>,
8539        gmac: &CudaSlice<f32>,
8540        umac: &CudaSlice<f32>,
8541        sel: &CudaSlice<i32>,
8542        limit: Option<f32>,
8543        out_q: &mut CudaSlice<i8>,
8544        out_d: &mut CudaSlice<f32>,
8545        n_per: usize,
8546        n_sel: usize,
8547    ) -> Result<(), Box<dyn std::error::Error>> {
8548        let n = n_per * n_sel;
8549        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8550            return Err(format!(
8551                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8552                out_q.len(),
8553                out_d.len()
8554            )
8555            .into());
8556        }
8557        if let Some(limit) = limit {
8558            if limit <= 1e-6 {
8559                return Err(format!(
8560                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8561                )
8562                .into());
8563            }
8564            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8565            let cfg = LaunchConfig::for_num_elems(n as u32);
8566            let (np, ns) = (n_per as i32, n_sel as i32);
8567            let __s_b = self.gpu.stream();
8568            let mut b = __s_b.launch_builder(&f);
8569            b.arg(gate)
8570                .arg(up)
8571                .arg(gmac)
8572                .arg(umac)
8573                .arg(sel)
8574                .arg(&limit)
8575                .arg(out_q)
8576                .arg(out_d)
8577                .arg(&np)
8578                .arg(&ns);
8579            unsafe {
8580                b.launch(cfg)?;
8581            }
8582            return Ok(());
8583        }
8584        let f = self.func("silu_mul_scaled_q8_1_sel");
8585        let cfg = LaunchConfig::for_num_elems(n as u32);
8586        let (np, ns) = (n_per as i32, n_sel as i32);
8587        let __s_b = self.gpu.stream();
8588        let mut b = __s_b.launch_builder(&f);
8589        b.arg(gate)
8590            .arg(up)
8591            .arg(gmac)
8592            .arg(umac)
8593            .arg(sel)
8594            .arg(out_q)
8595            .arg(out_d)
8596            .arg(&np)
8597            .arg(&ns);
8598        unsafe {
8599            b.launch(cfg)?;
8600        }
8601        Ok(())
8602    }
8603
8604    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8605        Ok(self.gpu.stream().clone_htod(v)?)
8606    }
8607    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8608        Ok(self.gpu.stream().clone_htod(v)?)
8609    }
8610    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
8611    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8612        Ok(self.gpu.stream().clone_htod(v)?)
8613    }
8614    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8615        Ok(self.gpu.stream().clone_htod(v)?)
8616    }
8617    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
8618    pub fn dtoh_view(
8619        &self,
8620        d: &cudarc::driver::CudaView<f32>,
8621    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8622        let v = self.gpu.stream().clone_dtoh(d)?;
8623        self.gpu.stream().synchronize()?;
8624        Ok(v)
8625    }
8626    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8627        let v = self.gpu.stream().clone_dtoh(d)?;
8628        self.gpu.stream().synchronize()?;
8629        Ok(v)
8630    }
8631    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
8632    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
8633    /// issuing them together avoids a second stream synchronization in every trunk layer.
8634    pub fn dtoh_pair(
8635        &self,
8636        a: &CudaSlice<f32>,
8637        b: &CudaSlice<f32>,
8638    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8639        let av = self.gpu.stream().clone_dtoh(a)?;
8640        let bv = self.gpu.stream().clone_dtoh(b)?;
8641        self.gpu.stream().synchronize()?;
8642        Ok((av, bv))
8643    }
8644    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
8645    /// cross a shape-sensitive host boundary.
8646    pub fn dtoh_pair_views(
8647        &self,
8648        a: &cudarc::driver::CudaView<f32>,
8649        b: &cudarc::driver::CudaView<f32>,
8650    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8651        let av = self.gpu.stream().clone_dtoh(a)?;
8652        let bv = self.gpu.stream().clone_dtoh(b)?;
8653        self.gpu.stream().synchronize()?;
8654        Ok((av, bv))
8655    }
8656    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
8657    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8658        let v = self.gpu.stream().clone_dtoh(d)?;
8659        self.gpu.stream().synchronize()?;
8660        Ok(v)
8661    }
8662    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
8663    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8664        let v = self.gpu.stream().clone_dtoh(d)?;
8665        self.gpu.stream().synchronize()?;
8666        Ok(v)
8667    }
8668    pub fn dtoh_u8_view(
8669        &self,
8670        d: &cudarc::driver::CudaView<u8>,
8671    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8672        let v = self.gpu.stream().clone_dtoh(d)?;
8673        self.gpu.stream().synchronize()?;
8674        Ok(v)
8675    }
8676    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8677        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8678        self.keep_if_capturing(&s);
8679        Ok(s)
8680    }
8681
8682    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
8683    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
8684    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
8685    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
8686    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
8687    /// back (or kept resident for graph replay). Returns the device token buffer.
8688    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
8689    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
8690    pub fn prob_of_token_device(
8691        &self,
8692        logits: &CudaSlice<f32>,
8693        tok: &CudaSlice<u32>,
8694        n_vocab: usize,
8695    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8696        let nb = ARGMAX_NB;
8697        let mut part = self.alloc_uninit::<f32>(nb)?;
8698        let mut p = self.alloc_uninit::<f32>(1)?;
8699        let f1 = self.func("prob_of_token_partial_f32");
8700        let cfg1 = LaunchConfig {
8701            grid_dim: (nb as u32, 1, 1),
8702            block_dim: (256, 1, 1),
8703            shared_mem_bytes: 0,
8704        };
8705        let nv = n_vocab as i32;
8706        let __s_b1 = self.gpu.stream();
8707        let mut b1 = __s_b1.launch_builder(&f1);
8708        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8709        unsafe {
8710            b1.launch(cfg1)?;
8711        }
8712        let f2 = self.func("prob_of_token_final_f32");
8713        let cfg2 = LaunchConfig {
8714            grid_dim: (1, 1, 1),
8715            block_dim: (256, 1, 1),
8716            shared_mem_bytes: 0,
8717        };
8718        let nbi = nb as i32;
8719        let __s_b2 = self.gpu.stream();
8720        let mut b2 = __s_b2.launch_builder(&f2);
8721        b2.arg(&part).arg(&mut p).arg(&nbi);
8722        unsafe {
8723            b2.launch(cfg2)?;
8724        }
8725        Ok(p)
8726    }
8727
8728    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
8729    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
8730    /// where the host reads the p-min confidence between replays. Same kernels, same math.
8731    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
8732    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
8733    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
8734    pub fn prob_of_token_device_col(
8735        &self,
8736        logits: &CudaSlice<f32>,
8737        tok_all: &CudaSlice<u32>,
8738        tok_idx: usize,
8739        p_out: &mut CudaSlice<f32>,
8740        p_idx: usize,
8741        n_vocab: usize,
8742    ) -> Result<(), Box<dyn std::error::Error>> {
8743        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8744        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8745        let nb = ARGMAX_NB;
8746        let mut part = self.alloc_uninit::<f32>(nb)?;
8747        let f1 = self.func("prob_of_token_partial_f32");
8748        let cfg1 = LaunchConfig {
8749            grid_dim: (nb as u32, 1, 1),
8750            block_dim: (256, 1, 1),
8751            shared_mem_bytes: 0,
8752        };
8753        let nv = n_vocab as i32;
8754        let __s_b1 = self.gpu.stream();
8755        let mut b1 = __s_b1.launch_builder(&f1);
8756        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8757        unsafe {
8758            b1.launch(cfg1)?;
8759        }
8760        let f2 = self.func("prob_of_token_final_f32");
8761        let cfg2 = LaunchConfig {
8762            grid_dim: (1, 1, 1),
8763            block_dim: (256, 1, 1),
8764            shared_mem_bytes: 0,
8765        };
8766        let nbi = nb as i32;
8767        let __s_b2 = self.gpu.stream();
8768        let mut b2 = __s_b2.launch_builder(&f2);
8769        b2.arg(&part).arg(&mut p_v).arg(&nbi);
8770        unsafe {
8771            b2.launch(cfg2)?;
8772        }
8773        Ok(())
8774    }
8775
8776    pub fn prob_of_token_device_into(
8777        &self,
8778        logits: &CudaSlice<f32>,
8779        tok: &CudaSlice<u32>,
8780        p_out: &mut CudaSlice<f32>,
8781        n_vocab: usize,
8782    ) -> Result<(), Box<dyn std::error::Error>> {
8783        let nb = ARGMAX_NB;
8784        let mut part = self.alloc_uninit::<f32>(nb)?;
8785        let f1 = self.func("prob_of_token_partial_f32");
8786        let cfg1 = LaunchConfig {
8787            grid_dim: (nb as u32, 1, 1),
8788            block_dim: (256, 1, 1),
8789            shared_mem_bytes: 0,
8790        };
8791        let nv = n_vocab as i32;
8792        let __s_b1 = self.gpu.stream();
8793        let mut b1 = __s_b1.launch_builder(&f1);
8794        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8795        unsafe {
8796            b1.launch(cfg1)?;
8797        }
8798        let f2 = self.func("prob_of_token_final_f32");
8799        let cfg2 = LaunchConfig {
8800            grid_dim: (1, 1, 1),
8801            block_dim: (256, 1, 1),
8802            shared_mem_bytes: 0,
8803        };
8804        let nbi = nb as i32;
8805        let __s_b2 = self.gpu.stream();
8806        let mut b2 = __s_b2.launch_builder(&f2);
8807        b2.arg(&part).arg(p_out).arg(&nbi);
8808        unsafe {
8809            b2.launch(cfg2)?;
8810        }
8811        Ok(())
8812    }
8813
8814    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
8815    /// (graph-constant params, device-varying index). Capture-safe.
8816    pub fn u32_hist_append(
8817        &self,
8818        tok: &CudaSlice<u32>,
8819        hist: &mut CudaSlice<u32>,
8820        idx: &mut CudaSlice<i32>,
8821    ) -> Result<(), Box<dyn std::error::Error>> {
8822        let f = self.func("u32_hist_append");
8823        let cfg = LaunchConfig {
8824            grid_dim: (1, 1, 1),
8825            block_dim: (32, 1, 1),
8826            shared_mem_bytes: 0,
8827        };
8828        let __s_b = self.gpu.stream();
8829        let mut b = __s_b.launch_builder(&f);
8830        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8831        unsafe {
8832            b.launch(cfg)?;
8833        }
8834        Ok(())
8835    }
8836
8837    pub fn argmax_token_device(
8838        &self,
8839        logits: &CudaSlice<f32>,
8840        n_vocab: usize,
8841    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8842        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8843        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8844        Ok(tok)
8845    }
8846    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
8847    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
8848    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
8849    /// pointer is baked once and the token id never round-trips to host inside steady state. The
8850    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
8851    /// captured passes bake fixed addresses.
8852    pub fn argmax_token_device_into(
8853        &self,
8854        logits: &CudaSlice<f32>,
8855        tok: &mut CudaSlice<u32>,
8856        n_vocab: usize,
8857    ) -> Result<(), Box<dyn std::error::Error>> {
8858        let nb = ARGMAX_NB;
8859        let f1 = self.func("argmax_partial_f32");
8860        let f2 = self.func("argmax_final_f32");
8861        let mut guard = self.argmax_partials.lock().unwrap();
8862        if guard.is_none() {
8863            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8864            // buffers carry no cudarc events (illegal inside capture).
8865            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8866            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8867            *guard = Some((pv, pi));
8868        }
8869        let (part_v, part_i) = guard.as_mut().unwrap();
8870        let nv = n_vocab as i32;
8871        let nbi = nb as i32;
8872        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8873        let cfg1 = LaunchConfig {
8874            grid_dim: (nb as u32, 1, 1),
8875            block_dim: (256, 1, 1),
8876            shared_mem_bytes: 0,
8877        };
8878        let __s_b1 = self.gpu.stream();
8879        let mut b1 = __s_b1.launch_builder(&f1);
8880        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8881        unsafe {
8882            b1.launch(cfg1)?;
8883        }
8884        // pass 2: one block reduces NB partials -> token_out[0].
8885        let cfg2 = LaunchConfig {
8886            grid_dim: (1, 1, 1),
8887            block_dim: (256, 1, 1),
8888            shared_mem_bytes: 0,
8889        };
8890        let __s_b2 = self.gpu.stream();
8891        let mut b2 = __s_b2.launch_builder(&f2);
8892        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8893        unsafe {
8894            b2.launch(cfg2)?;
8895        }
8896        Ok(())
8897    }
8898    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8899    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8900    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8901    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8902    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8903    pub fn argmax_token_device_col(
8904        &self,
8905        logits: &CudaSlice<f32>,
8906        col: usize,
8907        n_vocab: usize,
8908        toks: &mut CudaSlice<u32>,
8909        out_idx: usize,
8910    ) -> Result<(), Box<dyn std::error::Error>> {
8911        let nb = ARGMAX_NB;
8912        let f1 = self.func("argmax_partial_f32");
8913        let f2 = self.func("argmax_final_f32");
8914        let mut guard = self.argmax_partials.lock().unwrap();
8915        if guard.is_none() {
8916            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8917            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8918            *guard = Some((pv, pi));
8919        }
8920        let (part_v, part_i) = guard.as_mut().unwrap();
8921        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8922        let nv = n_vocab as i32;
8923        let nbi = nb as i32;
8924        let cfg1 = LaunchConfig {
8925            grid_dim: (nb as u32, 1, 1),
8926            block_dim: (256, 1, 1),
8927            shared_mem_bytes: 0,
8928        };
8929        let __s_b1 = self.gpu.stream();
8930        let mut b1 = __s_b1.launch_builder(&f1);
8931        b1.arg(&col_view)
8932            .arg(&mut *part_v)
8933            .arg(&mut *part_i)
8934            .arg(&nv);
8935        unsafe {
8936            b1.launch(cfg1)?;
8937        }
8938        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8939        let cfg2 = LaunchConfig {
8940            grid_dim: (1, 1, 1),
8941            block_dim: (256, 1, 1),
8942            shared_mem_bytes: 0,
8943        };
8944        let __s_b2 = self.gpu.stream();
8945        let mut b2 = __s_b2.launch_builder(&f2);
8946        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8947        unsafe {
8948            b2.launch(cfg2)?;
8949        }
8950        Ok(())
8951    }
8952    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8953    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8954        Ok(self.gpu.stream().clone_htod(v)?)
8955    }
8956    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8957        let v = self.gpu.stream().clone_dtoh(d)?;
8958        self.gpu.stream().synchronize()?;
8959        Ok(v)
8960    }
8961
8962    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8963        let v = self.gpu.stream().clone_dtoh(d)?;
8964        self.gpu.stream().synchronize()?;
8965        Ok(v)
8966    }
8967    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8968    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8969    /// contents change every step, the address must not, so a captured graph can read it).
8970    pub fn htod_u32_into(
8971        &self,
8972        dst: &mut CudaSlice<u32>,
8973        src: &[u32],
8974    ) -> Result<(), Box<dyn std::error::Error>> {
8975        let mut view = dst.slice_mut(0..src.len());
8976        self.gpu.stream().memcpy_htod(src, &mut view)?;
8977        Ok(())
8978    }
8979
8980    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8981    /// table without changing the device address its reconcile kernel consumes.
8982    pub fn htod_i32_into(
8983        &self,
8984        dst: &mut CudaSlice<i32>,
8985        src: &[i32],
8986    ) -> Result<(), Box<dyn std::error::Error>> {
8987        let mut view = dst.slice_mut(0..src.len());
8988        self.gpu.stream().memcpy_htod(src, &mut view)?;
8989        Ok(())
8990    }
8991
8992    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8993        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8994        self.keep_if_capturing(&s);
8995        Ok(s)
8996    }
8997    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8998    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8999    pub fn embed_gather_device_into(
9000        &self,
9001        embd: &CudaSlice<u8>,
9002        token_d: &CudaSlice<u32>,
9003        x_out: &mut CudaSlice<f32>,
9004        n_embd: usize,
9005        qtype: i32,
9006        row_bytes: usize,
9007    ) -> Result<(), Box<dyn std::error::Error>> {
9008        let f = self.func("embed_gather_u32");
9009        let cfg = LaunchConfig {
9010            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9011            block_dim: (256, 1, 1),
9012            shared_mem_bytes: 0,
9013        };
9014        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9015        let __s_b = self.gpu.stream();
9016        let mut b = __s_b.launch_builder(&f);
9017        b.arg(embd)
9018            .arg(token_d)
9019            .arg(x_out)
9020            .arg(&ne)
9021            .arg(&qt)
9022            .arg(&rb);
9023        unsafe {
9024            b.launch(cfg)?;
9025        }
9026        Ok(())
9027    }
9028    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
9029    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
9030        let v = self.gpu.stream().clone_dtoh(d)?;
9031        self.gpu.stream().synchronize()?;
9032        Ok(v[0])
9033    }
9034    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
9035    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
9036    /// the counter value after the throwaway capture warmups corrupt it.
9037    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
9038    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
9039    /// copy (fine at stream-idle boundaries, poison mid-round).
9040    pub fn i32_set_k(
9041        &self,
9042        dst: &mut CudaSlice<i32>,
9043        v: i32,
9044    ) -> Result<(), Box<dyn std::error::Error>> {
9045        let f = self.func("i32_set_k");
9046        let cfg = LaunchConfig {
9047            grid_dim: (1, 1, 1),
9048            block_dim: (1, 1, 1),
9049            shared_mem_bytes: 0,
9050        };
9051        let idx = 0i32;
9052        let __s_b = self.gpu.stream();
9053        let mut b = __s_b.launch_builder(&f);
9054        b.arg(dst).arg(&v).arg(&idx);
9055        unsafe {
9056            b.launch(cfg)?;
9057        }
9058        Ok(())
9059    }
9060
9061    pub fn set_i32_one(
9062        &self,
9063        d: &mut CudaSlice<i32>,
9064        v: i32,
9065    ) -> Result<(), Box<dyn std::error::Error>> {
9066        self.gpu.stream().memcpy_htod(&[v], d)?;
9067        Ok(())
9068    }
9069    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
9070    /// during priming / capture-state restore.
9071    pub fn set_u32_one(
9072        &self,
9073        d: &mut CudaSlice<u32>,
9074        v: u32,
9075    ) -> Result<(), Box<dyn std::error::Error>> {
9076        self.gpu.stream().memcpy_htod(&[v], d)?;
9077        Ok(())
9078    }
9079    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
9080    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9081        let v = self.gpu.stream().clone_dtoh(d)?;
9082        self.gpu.stream().synchronize()?;
9083        Ok(v[0])
9084    }
9085    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
9086    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9087        Ok(self.gpu.stream().clone_htod(bytes)?)
9088    }
9089    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
9090    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
9091    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
9092    pub fn embed_gather_device(
9093        &self,
9094        embd: &CudaSlice<u8>,
9095        token_d: &CudaSlice<u32>,
9096        n_embd: usize,
9097        qtype: i32,
9098        row_bytes: usize,
9099    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9100        let f = self.func("embed_gather_u32");
9101        let mut x = self.alloc_uninit::<f32>(n_embd)?;
9102        let cfg = LaunchConfig {
9103            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9104            block_dim: (256, 1, 1),
9105            shared_mem_bytes: 0,
9106        };
9107        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9108        let __s_b = self.gpu.stream();
9109        let mut b = __s_b.launch_builder(&f);
9110        b.arg(embd)
9111            .arg(token_d)
9112            .arg(&mut x)
9113            .arg(&ne)
9114            .arg(&qt)
9115            .arg(&rb);
9116        unsafe {
9117            b.launch(cfg)?;
9118        }
9119        Ok(x)
9120    }
9121
9122    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
9123    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
9124    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
9125    pub fn embed_gather_device_t(
9126        &self,
9127        embd: &CudaSlice<u8>,
9128        tokens: &[u32],
9129        n_embd: usize,
9130        qtype: i32,
9131        row_bytes: usize,
9132    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9133        let t = tokens.len();
9134        let tok_d = self.gpu.stream().clone_htod(tokens)?;
9135        let f = self.func("embed_gather_u32_t");
9136        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9137        let cfg = LaunchConfig {
9138            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9139            block_dim: (256, 1, 1),
9140            shared_mem_bytes: 0,
9141        };
9142        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9143        let __s_b = self.gpu.stream();
9144        let mut b = __s_b.launch_builder(&f);
9145        b.arg(embd)
9146            .arg(&tok_d)
9147            .arg(&mut x)
9148            .arg(&ne)
9149            .arg(&qt)
9150            .arg(&rb)
9151            .arg(&ti);
9152        unsafe {
9153            b.launch(cfg)?;
9154        }
9155        Ok(x)
9156    }
9157
9158    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
9159    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
9160    /// as embed_gather_device_t — bit-identical rows.
9161    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
9162    pub fn embed_gather_device_tv(
9163        &self,
9164        embd: &CudaSlice<u8>,
9165        tok_v: &cudarc::driver::CudaView<u32>,
9166        t: usize,
9167        n_embd: usize,
9168        qtype: i32,
9169        row_bytes: usize,
9170    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9171        let f = self.func("embed_gather_u32_t");
9172        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9173        let cfg = LaunchConfig {
9174            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9175            block_dim: (256, 1, 1),
9176            shared_mem_bytes: 0,
9177        };
9178        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9179        let __s_b = self.gpu.stream();
9180        let mut b = __s_b.launch_builder(&f);
9181        b.arg(embd)
9182            .arg(tok_v)
9183            .arg(&mut x)
9184            .arg(&ne)
9185            .arg(&qt)
9186            .arg(&rb)
9187            .arg(&ti);
9188        unsafe {
9189            b.launch(cfg)?;
9190        }
9191        Ok(x)
9192    }
9193
9194    pub fn embed_gather_device_td(
9195        &self,
9196        embd: &CudaSlice<u8>,
9197        tok_d: &CudaSlice<u32>,
9198        t: usize,
9199        n_embd: usize,
9200        qtype: i32,
9201        row_bytes: usize,
9202    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9203        let f = self.func("embed_gather_u32_t");
9204        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9205        let cfg = LaunchConfig {
9206            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9207            block_dim: (256, 1, 1),
9208            shared_mem_bytes: 0,
9209        };
9210        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9211        let __s_b = self.gpu.stream();
9212        let mut b = __s_b.launch_builder(&f);
9213        b.arg(embd)
9214            .arg(tok_d)
9215            .arg(&mut x)
9216            .arg(&ne)
9217            .arg(&qt)
9218            .arg(&rb)
9219            .arg(&ti);
9220        unsafe {
9221            b.launch(cfg)?;
9222        }
9223        Ok(x)
9224    }
9225
9226    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
9227    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
9228    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
9229    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
9230    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
9231    #[inline]
9232    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
9233    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9234        if self
9235            .capture_keep_on
9236            .load(std::sync::atomic::Ordering::Relaxed)
9237        {
9238            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9239        }
9240    }
9241
9242    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9243        &self,
9244        n: usize,
9245    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9246        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9247        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
9248        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
9249        // not cover engine-internal buffers). Debug-only: massive launch overhead.
9250        {
9251            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9252            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9253                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
9254                use cudarc::driver::DevicePtrMut;
9255                let n_bytes = s.len() * std::mem::size_of::<T>();
9256                let stream = self.gpu.stream();
9257                let (p_, _g) = s.device_ptr_mut(&stream);
9258                unsafe {
9259                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9260                        .result()?;
9261                }
9262            }
9263        }
9264        self.keep_if_capturing(&s);
9265        Ok(s)
9266    }
9267
9268    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
9269    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
9270    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
9271    /// consumers alloc through this (m=1 decode arms).
9272    pub fn uninit_q8_pair(
9273        &self,
9274        n: usize,
9275    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9276        Ok((
9277            self.alloc_uninit::<i8>(n)?,
9278            self.alloc_uninit::<f32>(n / 32)?,
9279        ))
9280    }
9281
9282    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9283        self.alloc_uninit::<f32>(n)
9284    }
9285
9286    /// i8 uninitialized scratch (same contract as `uninit`).
9287    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9288        self.alloc_uninit::<i8>(n)
9289    }
9290
9291    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
9292    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
9293    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
9294    #[allow(clippy::too_many_arguments)]
9295    pub fn rms_norm3(
9296        &self,
9297        x: &CudaSlice<f32>,
9298        w0: &CudaSlice<f32>,
9299        w1: &CudaSlice<f32>,
9300        w2: &CudaSlice<f32>,
9301        d0: &mut CudaSlice<f32>,
9302        d1: &mut CudaSlice<f32>,
9303        d2: &mut CudaSlice<f32>,
9304        ncols: usize,
9305        nrows: usize,
9306        eps: f32,
9307    ) -> Result<(), Box<dyn std::error::Error>> {
9308        let f = self.func("rms_norm3_f32");
9309        let cfg = LaunchConfig {
9310            grid_dim: (nrows as u32, 1, 1),
9311            block_dim: (rms_block(), 1, 1),
9312            shared_mem_bytes: 0,
9313        };
9314        let (nc, e) = (ncols as i32, eps);
9315        let __s_b = self.gpu.stream();
9316        let mut b = __s_b.launch_builder(&f);
9317        b.arg(x)
9318            .arg(w0)
9319            .arg(w1)
9320            .arg(w2)
9321            .arg(d0)
9322            .arg(d1)
9323            .arg(d2)
9324            .arg(&nc)
9325            .arg(&e);
9326        unsafe {
9327            b.launch(cfg)?;
9328        }
9329        Ok(())
9330    }
9331
9332    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
9333    #[allow(clippy::too_many_arguments)]
9334    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
9335    /// piggybacks on the same conditions.
9336    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9337        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9338        *WARP_ON.get_or_init(|| {
9339            std::env::var("MEMRA_QKVNORM_W")
9340                .map(|v| v != "0")
9341                .unwrap_or(true)
9342        }) && ncols % 4 == 0
9343            && rows >= 64
9344    }
9345
9346    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
9347    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
9348    #[allow(clippy::too_many_arguments)]
9349    pub fn rms_norm_qkv_w4b(
9350        &self,
9351        q: &CudaSlice<f32>,
9352        k: &CudaSlice<f32>,
9353        v: &CudaSlice<f32>,
9354        wq: &CudaSlice<f32>,
9355        wk: &CudaSlice<f32>,
9356        wv: &CudaSlice<f32>,
9357        dq: &mut CudaSlice<f32>,
9358        dk: &mut CudaSlice<f32>,
9359        dv: &mut CudaSlice<f32>,
9360        dvb: &mut CudaSlice<u8>,
9361        ncols: usize,
9362        rq: usize,
9363        rk: usize,
9364        eps: f32,
9365        vf16: bool,
9366    ) -> Result<(), Box<dyn std::error::Error>> {
9367        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9368        let f = self.func("rms_norm_qkv_w4b_f32");
9369        let rows = (rq + 2 * rk) as u32;
9370        let cfg = LaunchConfig {
9371            grid_dim: (rows.div_ceil(8), 1, 1),
9372            block_dim: (256, 1, 1),
9373            shared_mem_bytes: 0,
9374        };
9375        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9376        let vf = vf16 as i32;
9377        let __s_b = self.gpu.stream();
9378        let mut b = __s_b.launch_builder(&f);
9379        b.arg(q)
9380            .arg(k)
9381            .arg(v)
9382            .arg(wq)
9383            .arg(wk)
9384            .arg(wv)
9385            .arg(dq)
9386            .arg(dk)
9387            .arg(dv)
9388            .arg(&mut *dvb)
9389            .arg(&nc)
9390            .arg(&rqi)
9391            .arg(&rki)
9392            .arg(&rvi)
9393            .arg(&e)
9394            .arg(&vf);
9395        unsafe {
9396            b.launch(cfg)?;
9397        }
9398        Ok(())
9399    }
9400
9401    pub fn rms_norm_qkv(
9402        &self,
9403        q: &CudaSlice<f32>,
9404        k: &CudaSlice<f32>,
9405        v: &CudaSlice<f32>,
9406        wq: &CudaSlice<f32>,
9407        wk: &CudaSlice<f32>,
9408        wv: &CudaSlice<f32>,
9409        dq: &mut CudaSlice<f32>,
9410        dk: &mut CudaSlice<f32>,
9411        dv: &mut CudaSlice<f32>,
9412        ncols: usize,
9413        rq: usize,
9414        rk: usize,
9415        eps: f32,
9416    ) -> Result<(), Box<dyn std::error::Error>> {
9417        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
9418        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
9419        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
9420        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9421        let warp_on = *WARP_ON.get_or_init(|| {
9422            std::env::var("MEMRA_QKVNORM_W")
9423                .map(|v| v != "0")
9424                .unwrap_or(true)
9425        });
9426        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
9427        // replay numerics are untouched on every model; only prefill depth takes the new config.
9428        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9429            let f = self.func("rms_norm_qkv_w4_f32");
9430            let rows = (rq + 2 * rk) as u32;
9431            let cfg = LaunchConfig {
9432                grid_dim: (rows.div_ceil(8), 1, 1),
9433                block_dim: (256, 1, 1),
9434                shared_mem_bytes: 0,
9435            };
9436            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9437            let __s_b = self.gpu.stream();
9438            let mut b = __s_b.launch_builder(&f);
9439            b.arg(q)
9440                .arg(k)
9441                .arg(v)
9442                .arg(wq)
9443                .arg(wk)
9444                .arg(wv)
9445                .arg(dq)
9446                .arg(dk)
9447                .arg(dv)
9448                .arg(&nc)
9449                .arg(&rqi)
9450                .arg(&rki)
9451                .arg(&rvi)
9452                .arg(&e);
9453            unsafe {
9454                b.launch(cfg)?;
9455            }
9456            return Ok(());
9457        }
9458        let f = self.func("rms_norm_qkv_f32");
9459        let grid = (rq + 2 * rk) as u32;
9460        let cfg = LaunchConfig {
9461            grid_dim: (grid, 1, 1),
9462            block_dim: (rms_block(), 1, 1),
9463            shared_mem_bytes: 0,
9464        };
9465        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9466        let __s_b = self.gpu.stream();
9467        let mut b = __s_b.launch_builder(&f);
9468        b.arg(q)
9469            .arg(k)
9470            .arg(v)
9471            .arg(wq)
9472            .arg(wk)
9473            .arg(wv)
9474            .arg(dq)
9475            .arg(dk)
9476            .arg(dv)
9477            .arg(&nc)
9478            .arg(&rqi)
9479            .arg(&rki)
9480            .arg(&e);
9481        unsafe {
9482            b.launch(cfg)?;
9483        }
9484        Ok(())
9485    }
9486
9487    /// gemma4 fused pair of rms_norms over two different inputs (same width).
9488    #[allow(clippy::too_many_arguments)]
9489    pub fn rms_norm2x(
9490        &self,
9491        a: &CudaSlice<f32>,
9492        bb: &CudaSlice<f32>,
9493        wa: &CudaSlice<f32>,
9494        wb: &CudaSlice<f32>,
9495        da: &mut CudaSlice<f32>,
9496        db: &mut CudaSlice<f32>,
9497        ncols: usize,
9498        nrows: usize,
9499        eps: f32,
9500    ) -> Result<(), Box<dyn std::error::Error>> {
9501        let f = self.func("rms_norm2x_f32");
9502        let cfg = LaunchConfig {
9503            grid_dim: (2 * nrows as u32, 1, 1),
9504            block_dim: (rms_block(), 1, 1),
9505            shared_mem_bytes: 0,
9506        };
9507        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9508        let __s_b = self.gpu.stream();
9509        let mut b = __s_b.launch_builder(&f);
9510        b.arg(a)
9511            .arg(bb)
9512            .arg(wa)
9513            .arg(wb)
9514            .arg(da)
9515            .arg(db)
9516            .arg(&nc)
9517            .arg(&nr)
9518            .arg(&e);
9519        unsafe {
9520            b.launch(cfg)?;
9521        }
9522        Ok(())
9523    }
9524
9525    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
9526    pub fn softcap(
9527        &self,
9528        y: &mut CudaSlice<f32>,
9529        cap: f32,
9530        n: usize,
9531    ) -> Result<(), Box<dyn std::error::Error>> {
9532        let f = self.func("softcap_f32");
9533        let cfg = LaunchConfig::for_num_elems(n as u32);
9534        let ni = n as i32;
9535        let __s_b = self.gpu.stream();
9536        let mut b = __s_b.launch_builder(&f);
9537        b.arg(y).arg(&cap).arg(&ni);
9538        unsafe {
9539            b.launch(cfg)?;
9540        }
9541        Ok(())
9542    }
9543
9544    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
9545    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
9546    pub fn mask_ids_rows(
9547        &self,
9548        y: &mut CudaSlice<f32>,
9549        ids: &CudaSlice<i32>,
9550        n_ids: usize,
9551        n_vocab: usize,
9552        t: usize,
9553    ) -> Result<(), Box<dyn std::error::Error>> {
9554        let f = self.func("mask_ids_rows_f32");
9555        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9556        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9557        let __s_b = self.gpu.stream();
9558        let mut b = __s_b.launch_builder(&f);
9559        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9560        unsafe {
9561            b.launch(cfg)?;
9562        }
9563        Ok(())
9564    }
9565
9566    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
9567    #[allow(clippy::too_many_arguments)]
9568    pub fn add_scale_rms_norm(
9569        &self,
9570        a: &CudaSlice<f32>,
9571        b_in: &CudaSlice<f32>,
9572        c: f32,
9573        w: &CudaSlice<f32>,
9574        res: &mut CudaSlice<f32>,
9575        dst: &mut CudaSlice<f32>,
9576        ncols: usize,
9577        nrows: usize,
9578        eps: f32,
9579    ) -> Result<(), Box<dyn std::error::Error>> {
9580        let f = self.func("add_scale_rms_norm_f32");
9581        let cfg = LaunchConfig {
9582            grid_dim: (nrows as u32, 1, 1),
9583            block_dim: (rms_block(), 1, 1),
9584            shared_mem_bytes: 0,
9585        };
9586        let (nc, e2) = (ncols as i32, eps);
9587        let __s_b = self.gpu.stream();
9588        let mut b = __s_b.launch_builder(&f);
9589        b.arg(a)
9590            .arg(b_in)
9591            .arg(&c)
9592            .arg(w)
9593            .arg(res)
9594            .arg(dst)
9595            .arg(&nc)
9596            .arg(&e2);
9597        unsafe {
9598            b.launch(cfg)?;
9599        }
9600        Ok(())
9601    }
9602
9603    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
9604    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
9605    #[allow(clippy::too_many_arguments)]
9606    pub fn add_scale_rms_norm_q8_1(
9607        &self,
9608        a: &CudaSlice<f32>,
9609        b_in: &CudaSlice<f32>,
9610        c: f32,
9611        w: &CudaSlice<f32>,
9612        res: &mut CudaSlice<f32>,
9613        ncols: usize,
9614        nrows: usize,
9615        eps: f32,
9616    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9617        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9618        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9619        let (nc, e2) = (ncols as i32, eps);
9620        if Self::pdl_on() && Self::pdl_wb_on() {
9621            {
9622                use cudarc::driver::{DevicePtr, DevicePtrMut};
9623                let s = &self.gpu.stream();
9624                let (pa, _g0) = a.device_ptr(s);
9625                let (pb, _g1) = b_in.device_ptr(s);
9626                let (pw, _g2) = w.device_ptr(s);
9627                let (pr, _g3) = res.device_ptr_mut(s);
9628                let (pq, _g4) = out_q.device_ptr_mut(s);
9629                let (pd, _g5) = out_d.device_ptr_mut(s);
9630                let mut ps = [
9631                    &pa as *const _ as *mut std::ffi::c_void,
9632                    &pb as *const _ as *mut _,
9633                    &c as *const _ as *mut _,
9634                    &pw as *const _ as *mut _,
9635                    &pr as *const _ as *mut _,
9636                    &pq as *const _ as *mut _,
9637                    &pd as *const _ as *mut _,
9638                    &nc as *const _ as *mut _,
9639                    &e2 as *const _ as *mut _,
9640                ];
9641                unsafe {
9642                    self.launch_pdl(
9643                        "add_scale_rms_norm_q8_1",
9644                        (nrows as u32, 1, 1),
9645                        (rms_block(), 1, 1),
9646                        &mut ps,
9647                    )?;
9648                }
9649            }
9650            return Ok((out_q, out_d));
9651        }
9652        let f = self.func("add_scale_rms_norm_q8_1");
9653        let cfg = LaunchConfig {
9654            grid_dim: (nrows as u32, 1, 1),
9655            block_dim: (rms_block(), 1, 1),
9656            shared_mem_bytes: 0,
9657        };
9658        let __s_b = self.gpu.stream();
9659        let mut b = __s_b.launch_builder(&f);
9660        b.arg(a)
9661            .arg(b_in)
9662            .arg(&c)
9663            .arg(w)
9664            .arg(res)
9665            .arg(&mut out_q)
9666            .arg(&mut out_d)
9667            .arg(&nc)
9668            .arg(&e2);
9669        unsafe {
9670            b.launch(cfg)?;
9671        }
9672        Ok((out_q, out_d))
9673    }
9674
9675    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
9676    #[allow(clippy::too_many_arguments)]
9677    pub fn add_scale_rms_norm_q8_1_into(
9678        &self,
9679        a: &CudaSlice<f32>,
9680        b_in: &CudaSlice<f32>,
9681        c: f32,
9682        w: &CudaSlice<f32>,
9683        res: &mut CudaSlice<f32>,
9684        ncols: usize,
9685        nrows: usize,
9686        eps: f32,
9687        out_q: &mut CudaSlice<i8>,
9688        out_d: &mut CudaSlice<f32>,
9689    ) -> Result<(), Box<dyn std::error::Error>> {
9690        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9691        let (nc, e2) = (ncols as i32, eps);
9692        if Self::pdl_on() && Self::pdl_wb_on() {
9693            use cudarc::driver::{DevicePtr, DevicePtrMut};
9694            let s = &self.gpu.stream();
9695            let (pa, _g0) = a.device_ptr(s);
9696            let (pb, _g1) = b_in.device_ptr(s);
9697            let (pw, _g2) = w.device_ptr(s);
9698            let (pr, _g3) = res.device_ptr_mut(s);
9699            let (pq, _g4) = out_q.device_ptr_mut(s);
9700            let (pd, _g5) = out_d.device_ptr_mut(s);
9701            let mut ps = [
9702                &pa as *const _ as *mut std::ffi::c_void,
9703                &pb as *const _ as *mut _,
9704                &c as *const _ as *mut _,
9705                &pw as *const _ as *mut _,
9706                &pr as *const _ as *mut _,
9707                &pq as *const _ as *mut _,
9708                &pd as *const _ as *mut _,
9709                &nc as *const _ as *mut _,
9710                &e2 as *const _ as *mut _,
9711            ];
9712            unsafe {
9713                self.launch_pdl(
9714                    "add_scale_rms_norm_q8_1",
9715                    (nrows as u32, 1, 1),
9716                    (rms_block(), 1, 1),
9717                    &mut ps,
9718                )?;
9719            }
9720            return Ok(());
9721        }
9722        let f = self.func("add_scale_rms_norm_q8_1");
9723        let cfg = LaunchConfig {
9724            grid_dim: (nrows as u32, 1, 1),
9725            block_dim: (rms_block(), 1, 1),
9726            shared_mem_bytes: 0,
9727        };
9728        let __s_b = self.gpu.stream();
9729        let mut b = __s_b.launch_builder(&f);
9730        b.arg(a)
9731            .arg(b_in)
9732            .arg(&c)
9733            .arg(w)
9734            .arg(res)
9735            .arg(&mut *out_q)
9736            .arg(&mut *out_d)
9737            .arg(&nc)
9738            .arg(&e2);
9739        unsafe {
9740            b.launch(cfg)?;
9741        }
9742        Ok(())
9743    }
9744
9745    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
9746    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
9747    #[allow(clippy::too_many_arguments)]
9748    pub fn rms_pre_add_scale_rms_norm_q8_1(
9749        &self,
9750        a: &CudaSlice<f32>,
9751        wa: &CudaSlice<f32>,
9752        b_in: &CudaSlice<f32>,
9753        c: f32,
9754        w: &CudaSlice<f32>,
9755        res: &mut CudaSlice<f32>,
9756        ncols: usize,
9757        nrows: usize,
9758        eps: f32,
9759    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9760        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9761        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9762        let (nc, e2) = (ncols as i32, eps);
9763        if Self::pdl_on() {
9764            {
9765                use cudarc::driver::{DevicePtr, DevicePtrMut};
9766                let s = &self.gpu.stream();
9767                let (pa, _g0) = a.device_ptr(s);
9768                let (pwa, _g1) = wa.device_ptr(s);
9769                let (pb, _g2) = b_in.device_ptr(s);
9770                let (pw, _g3) = w.device_ptr(s);
9771                let (pr, _g4) = res.device_ptr_mut(s);
9772                let (pq, _g5) = out_q.device_ptr_mut(s);
9773                let (pd, _g6) = out_d.device_ptr_mut(s);
9774                let mut ps = [
9775                    &pa as *const _ as *mut std::ffi::c_void,
9776                    &pwa as *const _ as *mut _,
9777                    &pb as *const _ as *mut _,
9778                    &c as *const _ as *mut _,
9779                    &pw as *const _ as *mut _,
9780                    &pr as *const _ as *mut _,
9781                    &pq as *const _ as *mut _,
9782                    &pd as *const _ as *mut _,
9783                    &nc as *const _ as *mut _,
9784                    &e2 as *const _ as *mut _,
9785                ];
9786                unsafe {
9787                    self.launch_pdl(
9788                        "rms_pre_add_scale_rms_norm_q8_1",
9789                        (nrows as u32, 1, 1),
9790                        (rms_block(), 1, 1),
9791                        &mut ps,
9792                    )?;
9793                }
9794            }
9795            return Ok((out_q, out_d));
9796        }
9797        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9798        let cfg = LaunchConfig {
9799            grid_dim: (nrows as u32, 1, 1),
9800            block_dim: (rms_block(), 1, 1),
9801            shared_mem_bytes: 0,
9802        };
9803        let __s_b = self.gpu.stream();
9804        let mut b = __s_b.launch_builder(&f);
9805        b.arg(a)
9806            .arg(wa)
9807            .arg(b_in)
9808            .arg(&c)
9809            .arg(w)
9810            .arg(res)
9811            .arg(&mut out_q)
9812            .arg(&mut out_d)
9813            .arg(&nc)
9814            .arg(&e2);
9815        unsafe {
9816            b.launch(cfg)?;
9817        }
9818        Ok((out_q, out_d))
9819    }
9820
9821    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
9822    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
9823    pub fn gelu_tanh_mul_q8_1(
9824        &self,
9825        gate: &CudaSlice<f32>,
9826        up: &cudarc::driver::CudaView<f32>,
9827        act: &mut CudaSlice<f32>,
9828        ncols: usize,
9829        nrows: usize,
9830    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9831        debug_assert!(ncols % 128 == 0);
9832        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9833        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9834        let nc = ncols as i32;
9835        if Self::pdl_on() {
9836            {
9837                use cudarc::driver::{DevicePtr, DevicePtrMut};
9838                let s = &self.gpu.stream();
9839                let (pg, _g0) = gate.device_ptr(s);
9840                let (pu, _g1) = up.device_ptr(s);
9841                let (pact, _g2) = act.device_ptr_mut(s);
9842                let (pq, _g3) = out_q.device_ptr_mut(s);
9843                let (pd, _g4) = out_d.device_ptr_mut(s);
9844                let mut ps = [
9845                    &pg as *const _ as *mut std::ffi::c_void,
9846                    &pu as *const _ as *mut _,
9847                    &pact as *const _ as *mut _,
9848                    &pq as *const _ as *mut _,
9849                    &pd as *const _ as *mut _,
9850                    &nc as *const _ as *mut _,
9851                ];
9852                unsafe {
9853                    self.launch_pdl(
9854                        "gelu_tanh_mul_q8_1",
9855                        (nrows as u32, 1, 1),
9856                        (rms_block(), 1, 1),
9857                        &mut ps,
9858                    )?;
9859                }
9860            }
9861            return Ok((out_q, out_d));
9862        }
9863        let f = self.func("gelu_tanh_mul_q8_1");
9864        let cfg = LaunchConfig {
9865            grid_dim: (nrows as u32, 1, 1),
9866            block_dim: (rms_block(), 1, 1),
9867            shared_mem_bytes: 0,
9868        };
9869        let __s_b = self.gpu.stream();
9870        let mut b = __s_b.launch_builder(&f);
9871        b.arg(gate)
9872            .arg(up)
9873            .arg(act)
9874            .arg(&mut out_q)
9875            .arg(&mut out_d)
9876            .arg(&nc);
9877        unsafe {
9878            b.launch(cfg)?;
9879        }
9880        Ok((out_q, out_d))
9881    }
9882
9883    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9884    #[allow(clippy::too_many_arguments)]
9885    pub fn gelu_tanh_mul_q8_1_into(
9886        &self,
9887        gate: &CudaSlice<f32>,
9888        up: &cudarc::driver::CudaView<f32>,
9889        act: &mut CudaSlice<f32>,
9890        ncols: usize,
9891        nrows: usize,
9892        out_q: &mut CudaSlice<i8>,
9893        out_d: &mut CudaSlice<f32>,
9894    ) -> Result<(), Box<dyn std::error::Error>> {
9895        debug_assert!(ncols % 128 == 0);
9896        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9897        let nc = ncols as i32;
9898        if Self::pdl_on() {
9899            use cudarc::driver::{DevicePtr, DevicePtrMut};
9900            let s = &self.gpu.stream();
9901            let (pg, _g0) = gate.device_ptr(s);
9902            let (pu, _g1) = up.device_ptr(s);
9903            let (pact, _g2) = act.device_ptr_mut(s);
9904            let (pq, _g3) = out_q.device_ptr_mut(s);
9905            let (pd, _g4) = out_d.device_ptr_mut(s);
9906            let mut ps = [
9907                &pg as *const _ as *mut std::ffi::c_void,
9908                &pu as *const _ as *mut _,
9909                &pact as *const _ as *mut _,
9910                &pq as *const _ as *mut _,
9911                &pd as *const _ as *mut _,
9912                &nc as *const _ as *mut _,
9913            ];
9914            unsafe {
9915                self.launch_pdl(
9916                    "gelu_tanh_mul_q8_1",
9917                    (nrows as u32, 1, 1),
9918                    (rms_block(), 1, 1),
9919                    &mut ps,
9920                )?;
9921            }
9922            return Ok(());
9923        }
9924        let f = self.func("gelu_tanh_mul_q8_1");
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(gate)
9933            .arg(up)
9934            .arg(&mut *act)
9935            .arg(&mut *out_q)
9936            .arg(&mut *out_d)
9937            .arg(&nc);
9938        unsafe {
9939            b.launch(cfg)?;
9940        }
9941        Ok(())
9942    }
9943
9944    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9945    #[allow(clippy::too_many_arguments)]
9946    pub fn add_rms_norm3_q8z(
9947        &self,
9948        a: &CudaSlice<f32>,
9949        b_in: &CudaSlice<f32>,
9950        w0: &CudaSlice<f32>,
9951        w1: &CudaSlice<f32>,
9952        w2: &CudaSlice<f32>,
9953        res: &mut CudaSlice<f32>,
9954        out1: &mut CudaSlice<f32>,
9955        ncols: usize,
9956        nrows: usize,
9957        eps: f32,
9958    ) -> Result<
9959        (
9960            (CudaSlice<i8>, CudaSlice<f32>),
9961            (CudaSlice<i8>, CudaSlice<f32>),
9962        ),
9963        Box<dyn std::error::Error>,
9964    > {
9965        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9966        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9967        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9968        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9969        let f = self.func("add_rms_norm3_q8z_f32");
9970        let cfg = LaunchConfig {
9971            grid_dim: (nrows as u32, 1, 1),
9972            block_dim: (rms_block(), 1, 1),
9973            shared_mem_bytes: 0,
9974        };
9975        let (nc, e2) = (ncols as i32, eps);
9976        let __s_b = self.gpu.stream();
9977        let mut b = __s_b.launch_builder(&f);
9978        b.arg(a)
9979            .arg(b_in)
9980            .arg(w0)
9981            .arg(w1)
9982            .arg(w2)
9983            .arg(res)
9984            .arg(&mut q0)
9985            .arg(&mut d0)
9986            .arg(out1)
9987            .arg(&mut q2)
9988            .arg(&mut d2)
9989            .arg(&nc)
9990            .arg(&e2);
9991        unsafe {
9992            b.launch(cfg)?;
9993        }
9994        Ok(((q0, d0), (q2, d2)))
9995    }
9996
9997    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9998    #[allow(clippy::too_many_arguments)]
9999    pub fn add_rms_norm3(
10000        &self,
10001        a: &CudaSlice<f32>,
10002        b_in: &CudaSlice<f32>,
10003        w0: &CudaSlice<f32>,
10004        w1: &CudaSlice<f32>,
10005        w2: &CudaSlice<f32>,
10006        res: &mut CudaSlice<f32>,
10007        d0: &mut CudaSlice<f32>,
10008        d1: &mut CudaSlice<f32>,
10009        d2: &mut CudaSlice<f32>,
10010        ncols: usize,
10011        nrows: usize,
10012        eps: f32,
10013    ) -> Result<(), Box<dyn std::error::Error>> {
10014        let f = self.func("add_rms_norm3_f32");
10015        let cfg = LaunchConfig {
10016            grid_dim: (nrows as u32, 1, 1),
10017            block_dim: (rms_block(), 1, 1),
10018            shared_mem_bytes: 0,
10019        };
10020        let (nc, e2) = (ncols as i32, eps);
10021        let __s_b = self.gpu.stream();
10022        let mut b = __s_b.launch_builder(&f);
10023        b.arg(a)
10024            .arg(b_in)
10025            .arg(w0)
10026            .arg(w1)
10027            .arg(w2)
10028            .arg(res)
10029            .arg(d0)
10030            .arg(d1)
10031            .arg(d2)
10032            .arg(&nc)
10033            .arg(&e2);
10034        unsafe {
10035            b.launch(cfg)?;
10036        }
10037        Ok(())
10038    }
10039
10040    /// dst = (a + b) * c (residual add + layer scale, one launch).
10041    pub fn add_scale(
10042        &self,
10043        a: &CudaSlice<f32>,
10044        b_in: &CudaSlice<f32>,
10045        c: f32,
10046        dst: &mut CudaSlice<f32>,
10047        n: usize,
10048    ) -> Result<(), Box<dyn std::error::Error>> {
10049        let f = self.func("add_scale_f32");
10050        let cfg = LaunchConfig::for_num_elems(n as u32);
10051        let ni = n as i32;
10052        let __s_b = self.gpu.stream();
10053        let mut b = __s_b.launch_builder(&f);
10054        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10055        unsafe {
10056            b.launch(cfg)?;
10057        }
10058        Ok(())
10059    }
10060
10061    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
10062    pub fn layer_norm_bias(
10063        &self,
10064        x: &CudaSlice<f32>,
10065        w: &CudaSlice<f32>,
10066        b: &CudaSlice<f32>,
10067        dst: &mut CudaSlice<f32>,
10068        ncols: usize,
10069        nrows: usize,
10070        eps: f32,
10071    ) -> Result<(), Box<dyn std::error::Error>> {
10072        let f = self.func("layer_norm_bias_f32");
10073        let (nc, e) = (ncols as i32, eps);
10074        let cfg = LaunchConfig {
10075            grid_dim: (nrows as u32, 1, 1),
10076            block_dim: (256, 1, 1),
10077            shared_mem_bytes: 0,
10078        };
10079        let __s_b = self.gpu.stream();
10080        let mut lb = __s_b.launch_builder(&f);
10081        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10082        unsafe {
10083            lb.launch(cfg)?;
10084        }
10085        Ok(())
10086    }
10087
10088    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
10089    pub fn gelu_tanh(
10090        &self,
10091        x: &CudaSlice<f32>,
10092        dst: &mut CudaSlice<f32>,
10093        n: usize,
10094    ) -> Result<(), Box<dyn std::error::Error>> {
10095        let f = self.func("gelu_tanh_f32");
10096        let ni = n as i64;
10097        let cfg = LaunchConfig {
10098            grid_dim: (n.div_ceil(256) as u32, 1, 1),
10099            block_dim: (256, 1, 1),
10100            shared_mem_bytes: 0,
10101        };
10102        let __s_b = self.gpu.stream();
10103        let mut lb = __s_b.launch_builder(&f);
10104        lb.arg(x).arg(&mut *dst).arg(&ni);
10105        unsafe {
10106            lb.launch(cfg)?;
10107        }
10108        Ok(())
10109    }
10110
10111    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
10112    pub fn row_softmax(
10113        &self,
10114        x: &mut CudaSlice<f32>,
10115        ncols: usize,
10116        nrows: usize,
10117    ) -> Result<(), Box<dyn std::error::Error>> {
10118        let f = self.func("row_softmax_f32");
10119        let nc = ncols as i32;
10120        let cfg = LaunchConfig {
10121            grid_dim: (nrows as u32, 1, 1),
10122            block_dim: (256, 1, 1),
10123            shared_mem_bytes: 0,
10124        };
10125        let __s_b = self.gpu.stream();
10126        let mut lb = __s_b.launch_builder(&f);
10127        lb.arg(&mut *x).arg(&nc);
10128        unsafe {
10129            lb.launch(cfg)?;
10130        }
10131        Ok(())
10132    }
10133
10134    pub fn rms_norm(
10135        &self,
10136        x: &CudaSlice<f32>,
10137        w: &CudaSlice<f32>,
10138        dst: &mut CudaSlice<f32>,
10139        ncols: usize,
10140        nrows: usize,
10141        eps: f32,
10142    ) -> Result<(), Box<dyn std::error::Error>> {
10143        let (nc, e) = (ncols as i32, eps);
10144        let kname = if Self::norm_ilp_on() {
10145            "rms_norm_f32_v2"
10146        } else {
10147            "rms_norm_f32"
10148        };
10149        if Self::pdl_on() && Self::pdl_wb_on() {
10150            use cudarc::driver::{DevicePtr, DevicePtrMut};
10151            let s = &self.gpu.stream();
10152            let (px, _g0) = x.device_ptr(s);
10153            let (pw, _g1) = w.device_ptr(s);
10154            let (pd, _g2) = dst.device_ptr_mut(s);
10155            let mut ps = [
10156                &px as *const _ as *mut std::ffi::c_void,
10157                &pw as *const _ as *mut _,
10158                &pd as *const _ as *mut _,
10159                &nc as *const _ as *mut _,
10160                &e as *const _ as *mut _,
10161            ];
10162            unsafe {
10163                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10164            }
10165            return Ok(());
10166        }
10167        let f = self.func(kname);
10168        let cfg = LaunchConfig {
10169            grid_dim: (nrows as u32, 1, 1),
10170            block_dim: (rms_block(), 1, 1),
10171            shared_mem_bytes: 0,
10172        };
10173        let __s_b = self.gpu.stream();
10174        let mut b = __s_b.launch_builder(&f);
10175        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10176        unsafe {
10177            b.launch(cfg)?;
10178        }
10179        Ok(())
10180    }
10181
10182    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
10183    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
10184    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
10185    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
10186    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
10187    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
10188    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
10189    pub fn rms_norm_decode(
10190        &self,
10191        x: &CudaSlice<f32>,
10192        w: &CudaSlice<f32>,
10193        dst: &mut CudaSlice<f32>,
10194        ncols: usize,
10195        nrows: usize,
10196        eps: f32,
10197    ) -> Result<(), Box<dyn std::error::Error>> {
10198        let f = self.func(if Self::norm_ilp_on() {
10199            "rms_norm_f32_v2"
10200        } else {
10201            "rms_norm_f32"
10202        });
10203        let cfg = LaunchConfig {
10204            grid_dim: (nrows as u32, 1, 1),
10205            block_dim: (1024, 1, 1),
10206            shared_mem_bytes: 0,
10207        };
10208        let (nc, e) = (ncols as i32, eps);
10209        let __s_b = self.gpu.stream();
10210        let mut b = __s_b.launch_builder(&f);
10211        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10212        unsafe {
10213            b.launch(cfg)?;
10214        }
10215        Ok(())
10216    }
10217
10218    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
10219    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
10220    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
10221    pub fn rms_norm_q8_1(
10222        &self,
10223        x: &CudaSlice<f32>,
10224        w: &CudaSlice<f32>,
10225        ncols: usize,
10226        nrows: usize,
10227        eps: f32,
10228    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10229        let nblk = ncols / 32;
10230        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10231        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10232        let (nc, e) = (ncols as i32, eps);
10233        if Self::pdl_on() {
10234            {
10235                use cudarc::driver::{DevicePtr, DevicePtrMut};
10236                let s = &self.gpu.stream();
10237                let (px, _g0) = x.device_ptr(s);
10238                let (pw, _g1) = w.device_ptr(s);
10239                let (pq, _g2) = q.device_ptr_mut(s);
10240                let (pd, _g3) = d.device_ptr_mut(s);
10241                let mut ps = [
10242                    &px as *const _ as *mut std::ffi::c_void,
10243                    &pw as *const _ as *mut _,
10244                    &pq as *const _ as *mut _,
10245                    &pd as *const _ as *mut _,
10246                    &nc as *const _ as *mut _,
10247                    &e as *const _ as *mut _,
10248                ];
10249                unsafe {
10250                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10251                }
10252            }
10253            return Ok((q, d));
10254        }
10255        let f = self.func("rms_norm_q8_1");
10256        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
10257        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
10258        let cfg = LaunchConfig {
10259            grid_dim: (nrows as u32, 1, 1),
10260            block_dim: (1024, 1, 1),
10261            shared_mem_bytes: 0,
10262        };
10263        let __s_b = self.gpu.stream();
10264        let mut b = __s_b.launch_builder(&f);
10265        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10266        unsafe {
10267            b.launch(cfg)?;
10268        }
10269        Ok((q, d))
10270    }
10271
10272    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
10273    /// PDL arm), caller-owned outputs.
10274    pub fn rms_norm_q8_1_into(
10275        &self,
10276        x: &CudaSlice<f32>,
10277        w: &CudaSlice<f32>,
10278        ncols: usize,
10279        nrows: usize,
10280        eps: f32,
10281        q: &mut CudaSlice<i8>,
10282        d: &mut CudaSlice<f32>,
10283    ) -> Result<(), Box<dyn std::error::Error>> {
10284        let nblk = ncols / 32;
10285        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10286        let (nc, e) = (ncols as i32, eps);
10287        if Self::pdl_on() {
10288            use cudarc::driver::{DevicePtr, DevicePtrMut};
10289            let s = &self.gpu.stream();
10290            let (px, _g0) = x.device_ptr(s);
10291            let (pw, _g1) = w.device_ptr(s);
10292            let (pq, _g2) = q.device_ptr_mut(s);
10293            let (pd, _g3) = d.device_ptr_mut(s);
10294            let mut ps = [
10295                &px as *const _ as *mut std::ffi::c_void,
10296                &pw as *const _ as *mut _,
10297                &pq as *const _ as *mut _,
10298                &pd as *const _ as *mut _,
10299                &nc as *const _ as *mut _,
10300                &e as *const _ as *mut _,
10301            ];
10302            unsafe {
10303                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10304            }
10305            return Ok(());
10306        }
10307        let f = self.func("rms_norm_q8_1");
10308        let cfg = LaunchConfig {
10309            grid_dim: (nrows as u32, 1, 1),
10310            block_dim: (1024, 1, 1),
10311            shared_mem_bytes: 0,
10312        };
10313        let __s_b = self.gpu.stream();
10314        let mut b = __s_b.launch_builder(&f);
10315        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10316        unsafe {
10317            b.launch(cfg)?;
10318        }
10319        Ok(())
10320    }
10321
10322    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
10323    pub fn quantize_q8_1_into(
10324        &self,
10325        x: &CudaSlice<f32>,
10326        m: usize,
10327        in_f: usize,
10328        q: &mut CudaSlice<i8>,
10329        d: &mut CudaSlice<f32>,
10330    ) -> Result<(), Box<dyn std::error::Error>> {
10331        let nblk = in_f / 32;
10332        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10333        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10334        let (inf, mi) = (in_f as i32, m as i32);
10335        if Self::pdl_on() && Self::pdl_wb_on() {
10336            use cudarc::driver::{DevicePtr, DevicePtrMut};
10337            let s = &self.gpu.stream();
10338            let (px, _g0) = x.device_ptr(s);
10339            let (pq, _g1) = q.device_ptr_mut(s);
10340            let (pd, _g2) = d.device_ptr_mut(s);
10341            let mut ps = [
10342                &px as *const _ as *mut std::ffi::c_void,
10343                &pq as *const _ as *mut _,
10344                &pd as *const _ as *mut _,
10345                &inf as *const _ as *mut _,
10346                &mi as *const _ as *mut _,
10347            ];
10348            unsafe {
10349                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10350            }
10351            return Ok(());
10352        }
10353        let f = self.func("quantize_q8_1");
10354        let __s_b = self.gpu.stream();
10355        let mut b = __s_b.launch_builder(&f);
10356        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10357        unsafe {
10358            b.launch(cfg)?;
10359        }
10360        Ok(())
10361    }
10362
10363    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
10364    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
10365    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
10366    pub fn add_rms_norm_q8_1(
10367        &self,
10368        a: &CudaSlice<f32>,
10369        b_in: &CudaSlice<f32>,
10370        w: &CudaSlice<f32>,
10371        res: &mut CudaSlice<f32>,
10372        ncols: usize,
10373        nrows: usize,
10374        eps: f32,
10375    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10376        let nblk = ncols / 32;
10377        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10378        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10379        let f = self.func("add_rms_norm_q8_1");
10380        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
10381        let cfg = LaunchConfig {
10382            grid_dim: (nrows as u32, 1, 1),
10383            block_dim: (1024, 1, 1),
10384            shared_mem_bytes: 0,
10385        };
10386        let (nc, e) = (ncols as i32, eps);
10387        let __s_bld = self.gpu.stream();
10388        let mut bld = __s_bld.launch_builder(&f);
10389        bld.arg(a)
10390            .arg(b_in)
10391            .arg(w)
10392            .arg(res)
10393            .arg(&mut q)
10394            .arg(&mut d)
10395            .arg(&nc)
10396            .arg(&e);
10397        unsafe {
10398            bld.launch(cfg)?;
10399        }
10400        Ok((q, d))
10401    }
10402
10403    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
10404    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
10405    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
10406    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
10407    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
10408    #[allow(clippy::too_many_arguments)]
10409    pub fn join_add_rms_norm_raw(
10410        &self,
10411        a0_raw: u64,
10412        a1_raw: u64,
10413        x: &CudaSlice<f32>,
10414        w: &CudaSlice<f32>,
10415        res: &mut CudaSlice<f32>,
10416        dst: &mut CudaSlice<f32>,
10417        ncols: usize,
10418        eps: f32,
10419    ) -> Result<(), Box<dyn std::error::Error>> {
10420        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10421            return Err("join_add_rms_norm geometry".into());
10422        }
10423        let f = self.func("join_add_rms_norm_f32");
10424        let cfg = LaunchConfig {
10425            grid_dim: (1, 1, 1),
10426            block_dim: (rms_block(), 1, 1),
10427            shared_mem_bytes: 0,
10428        };
10429        let (nc, e) = (ncols as i32, eps);
10430        let __s_b = self.gpu.stream();
10431        let mut b = __s_b.launch_builder(&f);
10432        b.arg(&a0_raw)
10433            .arg(&a1_raw)
10434            .arg(x)
10435            .arg(w)
10436            .arg(&mut *res)
10437            .arg(&mut *dst)
10438            .arg(&nc)
10439            .arg(&e);
10440        unsafe {
10441            b.launch(cfg)?;
10442        }
10443        Ok(())
10444    }
10445
10446    pub fn add_rms_norm(
10447        &self,
10448        a: &CudaSlice<f32>,
10449        b: &CudaSlice<f32>,
10450        w: &CudaSlice<f32>,
10451        res: &mut CudaSlice<f32>,
10452        dst: &mut CudaSlice<f32>,
10453        ncols: usize,
10454        nrows: usize,
10455        eps: f32,
10456    ) -> Result<(), Box<dyn std::error::Error>> {
10457        let (nc, e) = (ncols as i32, eps);
10458        let kname = if Self::norm_ilp_on() {
10459            "add_rms_norm_f32_v2"
10460        } else {
10461            "add_rms_norm_f32"
10462        };
10463        if Self::pdl_on() && Self::pdl_wb_on() {
10464            use cudarc::driver::{DevicePtr, DevicePtrMut};
10465            let s = &self.gpu.stream();
10466            let (pa, _g0) = a.device_ptr(s);
10467            let (pb, _g1) = b.device_ptr(s);
10468            let (pw, _g2) = w.device_ptr(s);
10469            let (pr, _g3) = res.device_ptr_mut(s);
10470            let (pd, _g4) = dst.device_ptr_mut(s);
10471            let mut ps = [
10472                &pa as *const _ as *mut std::ffi::c_void,
10473                &pb as *const _ as *mut _,
10474                &pw as *const _ as *mut _,
10475                &pr as *const _ as *mut _,
10476                &pd as *const _ as *mut _,
10477                &nc as *const _ as *mut _,
10478                &e as *const _ as *mut _,
10479            ];
10480            unsafe {
10481                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10482            }
10483            return Ok(());
10484        }
10485        let f = self.func(kname);
10486        let cfg = LaunchConfig {
10487            grid_dim: (nrows as u32, 1, 1),
10488            block_dim: (rms_block(), 1, 1),
10489            shared_mem_bytes: 0,
10490        };
10491        let __s_b2 = self.gpu.stream();
10492        let mut b2 = __s_b2.launch_builder(&f);
10493        b2.arg(a)
10494            .arg(b)
10495            .arg(w)
10496            .arg(&mut *res)
10497            .arg(&mut *dst)
10498            .arg(&nc)
10499            .arg(&e);
10500        unsafe {
10501            b2.launch(cfg)?;
10502        }
10503        Ok(())
10504    }
10505
10506    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
10507    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
10508    #[allow(clippy::too_many_arguments)]
10509    pub fn rms_pre_add_rms_norm(
10510        &self,
10511        a: &CudaSlice<f32>,
10512        wa: &CudaSlice<f32>,
10513        b: &CudaSlice<f32>,
10514        w: &CudaSlice<f32>,
10515        res: &mut CudaSlice<f32>,
10516        dst: &mut CudaSlice<f32>,
10517        ncols: usize,
10518        nrows: usize,
10519        eps: f32,
10520    ) -> Result<(), Box<dyn std::error::Error>> {
10521        let f = self.func("rms_pre_add_rms_norm_f32");
10522        let cfg = LaunchConfig {
10523            grid_dim: (nrows as u32, 1, 1),
10524            block_dim: (rms_block(), 1, 1),
10525            shared_mem_bytes: 0,
10526        };
10527        let (nc, e) = (ncols as i32, eps);
10528        let __s_b2 = self.gpu.stream();
10529        let mut b2 = __s_b2.launch_builder(&f);
10530        b2.arg(a)
10531            .arg(wa)
10532            .arg(b)
10533            .arg(w)
10534            .arg(&mut *res)
10535            .arg(&mut *dst)
10536            .arg(&nc)
10537            .arg(&e);
10538        unsafe {
10539            b2.launch(cfg)?;
10540        }
10541        Ok(())
10542    }
10543
10544    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
10545    #[allow(clippy::too_many_arguments)]
10546    pub fn rms_pre_add_rms_norm_q8z(
10547        &self,
10548        a: &CudaSlice<f32>,
10549        wa: &CudaSlice<f32>,
10550        b: &CudaSlice<f32>,
10551        w: &CudaSlice<f32>,
10552        res: &mut CudaSlice<f32>,
10553        dst: &mut CudaSlice<f32>,
10554        ncols: usize,
10555        nrows: usize,
10556        eps: f32,
10557    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10558        debug_assert!(ncols % 128 == 0);
10559        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10560        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10561        let (nc, e) = (ncols as i32, eps);
10562        if Self::pdl_on() {
10563            {
10564                use cudarc::driver::{DevicePtr, DevicePtrMut};
10565                let s = &self.gpu.stream();
10566                let (pa, _g0) = a.device_ptr(s);
10567                let (pwa, _g1) = wa.device_ptr(s);
10568                let (pb, _g2) = b.device_ptr(s);
10569                let (pw, _g3) = w.device_ptr(s);
10570                let (pr, _g4) = res.device_ptr_mut(s);
10571                let (pdst, _g5) = dst.device_ptr_mut(s);
10572                let (pq, _g6) = out_q.device_ptr_mut(s);
10573                let (pd, _g7) = out_d.device_ptr_mut(s);
10574                let mut ps = [
10575                    &pa as *const _ as *mut std::ffi::c_void,
10576                    &pwa as *const _ as *mut _,
10577                    &pb as *const _ as *mut _,
10578                    &pw as *const _ as *mut _,
10579                    &pr as *const _ as *mut _,
10580                    &pdst as *const _ as *mut _,
10581                    &pq as *const _ as *mut _,
10582                    &pd as *const _ as *mut _,
10583                    &nc as *const _ as *mut _,
10584                    &e as *const _ as *mut _,
10585                ];
10586                unsafe {
10587                    self.launch_pdl(
10588                        "rms_pre_add_rms_norm_q8z_f32",
10589                        (nrows as u32, 1, 1),
10590                        (rms_block(), 1, 1),
10591                        &mut ps,
10592                    )?;
10593                }
10594            }
10595            return Ok((out_q, out_d));
10596        }
10597        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10598        let cfg = LaunchConfig {
10599            grid_dim: (nrows as u32, 1, 1),
10600            block_dim: (rms_block(), 1, 1),
10601            shared_mem_bytes: 0,
10602        };
10603        let __s_b2 = self.gpu.stream();
10604        let mut b2 = __s_b2.launch_builder(&f);
10605        b2.arg(a)
10606            .arg(wa)
10607            .arg(b)
10608            .arg(w)
10609            .arg(&mut *res)
10610            .arg(&mut *dst)
10611            .arg(&mut out_q)
10612            .arg(&mut out_d)
10613            .arg(&nc)
10614            .arg(&e);
10615        unsafe {
10616            b2.launch(cfg)?;
10617        }
10618        Ok((out_q, out_d))
10619    }
10620
10621    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
10622    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
10623    /// body must stay attribute-free (the fused2_into precedent).
10624    #[allow(clippy::too_many_arguments)]
10625    pub fn rms_pre_add_rms_norm_q8z_into(
10626        &self,
10627        a: &CudaSlice<f32>,
10628        wa: &CudaSlice<f32>,
10629        b: &CudaSlice<f32>,
10630        w: &CudaSlice<f32>,
10631        res: &mut CudaSlice<f32>,
10632        dst: &mut CudaSlice<f32>,
10633        ncols: usize,
10634        nrows: usize,
10635        eps: f32,
10636        out_q: &mut CudaSlice<i8>,
10637        out_d: &mut CudaSlice<f32>,
10638    ) -> Result<(), Box<dyn std::error::Error>> {
10639        debug_assert!(ncols % 128 == 0);
10640        let (nc, e) = (ncols as i32, eps);
10641        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10642        let cfg = LaunchConfig {
10643            grid_dim: (nrows as u32, 1, 1),
10644            block_dim: (rms_block(), 1, 1),
10645            shared_mem_bytes: 0,
10646        };
10647        let __s_b = self.gpu.stream();
10648        let mut b2 = __s_b.launch_builder(&f);
10649        b2.arg(a)
10650            .arg(wa)
10651            .arg(b)
10652            .arg(w)
10653            .arg(&mut *res)
10654            .arg(&mut *dst)
10655            .arg(&mut *out_q)
10656            .arg(&mut *out_d)
10657            .arg(&nc)
10658            .arg(&e);
10659        unsafe {
10660            b2.launch(cfg)?;
10661        }
10662        Ok(())
10663    }
10664
10665    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
10666    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
10667    #[allow(clippy::too_many_arguments)]
10668    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10669        &self,
10670        a: &CudaSlice<f32>,
10671        wa: &CudaSlice<f32>,
10672        b_in: &CudaSlice<f32>,
10673        c: f32,
10674        w: &CudaSlice<f32>,
10675        res: &mut CudaSlice<f32>,
10676        ncols: usize,
10677        nrows: usize,
10678        eps: f32,
10679        out_q: &mut CudaSlice<i8>,
10680        out_d: &mut CudaSlice<f32>,
10681    ) -> Result<(), Box<dyn std::error::Error>> {
10682        debug_assert!(ncols % 128 == 0);
10683        let (nc, e2) = (ncols as i32, eps);
10684        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10685        let cfg = LaunchConfig {
10686            grid_dim: (nrows as u32, 1, 1),
10687            block_dim: (rms_block(), 1, 1),
10688            shared_mem_bytes: 0,
10689        };
10690        let __s_b = self.gpu.stream();
10691        let mut b2 = __s_b.launch_builder(&f);
10692        b2.arg(a)
10693            .arg(wa)
10694            .arg(b_in)
10695            .arg(&c)
10696            .arg(w)
10697            .arg(&mut *res)
10698            .arg(&mut *out_q)
10699            .arg(&mut *out_d)
10700            .arg(&nc)
10701            .arg(&e2);
10702        unsafe {
10703            b2.launch(cfg)?;
10704        }
10705        Ok(())
10706    }
10707
10708    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
10709    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
10710    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
10711    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
10712    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
10713    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
10714    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
10715    pub fn g4_pnfold_on() -> bool {
10716        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10717        *ON.get_or_init(|| {
10718            std::env::var("MEMRA_G4_PNFOLD")
10719                .map(|v| v != "0")
10720                .unwrap_or(true)
10721        })
10722    }
10723
10724    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
10725    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
10726    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
10727    pub fn build_q4_out_concat3(
10728        &self,
10729        w0: &crate::model::GpuTensor,
10730        w1: &crate::model::GpuTensor,
10731        w2: &crate::model::GpuTensor,
10732    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10733        use crate::model::GpuTensor;
10734        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10735            match w {
10736                GpuTensor::Quant {
10737                    qtype,
10738                    row_bytes,
10739                    rp,
10740                    ..
10741                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10742                _ => None,
10743            }
10744        };
10745        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10746        else {
10747            return Ok(None);
10748        };
10749        if rb0 != rb1
10750            || rb0 != rb2
10751            || w0.in_features() != w1.in_features()
10752            || w0.in_features() != w2.in_features()
10753        {
10754            return Ok(None);
10755        }
10756        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10757            match w {
10758                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10759                _ => unreachable!(),
10760            }
10761        }
10762        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10763        let total = rb0 * (o0 + o1 + o2);
10764        let mut cat = self.alloc_u8(total)?;
10765        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10766        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10767        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10768        Ok(Some(GpuTensor::Quant {
10769            bytes: cat,
10770            qtype: QT_Q4_0,
10771            row_bytes: rb0,
10772            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10773            scale: 1.0,
10774            rp: false,
10775            #[cfg(memra_cutlass)]
10776            cutlass: None,
10777            fp8: None,
10778            blk: None,
10779            rp4: None,
10780            f16: None,
10781        }))
10782    }
10783
10784    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
10785    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
10786    ///
10787    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
10788    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
10789    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
10790    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
10791    ///
10792    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
10793    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
10794    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
10795    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
10796    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
10797    ///
10798    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
10799    /// width. A future partial-rotary caller fails at its first launch with the geometry named
10800    /// instead of serving quietly wrong logits.
10801    fn full_width_rope_only(
10802        kernel: &str,
10803        n_rot: usize,
10804        head_dim: usize,
10805    ) -> Result<(), Box<dyn std::error::Error>> {
10806        if n_rot == head_dim {
10807            return Ok(());
10808        }
10809        Err(format!(
10810            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10811             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10812             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10813             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10814             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10815        )
10816        .into())
10817    }
10818
10819    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
10820    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10821    /// ([`Engine::full_width_rope_only`]).
10822    #[allow(clippy::too_many_arguments)]
10823    pub fn rms_norm_qkv_rope_cat(
10824        &self,
10825        qkv: &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    ) -> Result<(), Box<dyn std::error::Error>> {
10844        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10845        let rows = rq + rk + rk;
10846        let theta_scale = base.powf(-2.0 / head_dim as f32);
10847        let (nc, rqi, rki, nhq, nhk) = (
10848            head_dim as i32,
10849            rq as i32,
10850            rk as i32,
10851            nh_q as i32,
10852            nh_k as i32,
10853        );
10854        if Self::pdl_on() {
10855            use cudarc::driver::{DevicePtr, DevicePtrMut};
10856            let s = &self.gpu.stream();
10857            let (pqkv, _g0) = qkv.device_ptr(s);
10858            let (pwq, _g1) = wq.device_ptr(s);
10859            let (pwk, _g2) = wk.device_ptr(s);
10860            let (pwv, _g3) = wv.device_ptr(s);
10861            let (pq, _g4) = q.device_ptr_mut(s);
10862            let (pk, _g5) = k.device_ptr_mut(s);
10863            let (pv, _g6) = v.device_ptr_mut(s);
10864            let (ppos, _g7) = pos.device_ptr(s);
10865            let (pff, _g8) = match ff {
10866                Some(t) => {
10867                    let (p, g) = t.device_ptr(s);
10868                    (p, Some(g))
10869                }
10870                None => (0, None),
10871            };
10872            let mut ps = [
10873                &pqkv as *const _ as *mut std::ffi::c_void,
10874                &pwq as *const _ as *mut _,
10875                &pwk as *const _ as *mut _,
10876                &pwv as *const _ as *mut _,
10877                &pq as *const _ as *mut _,
10878                &pk as *const _ as *mut _,
10879                &pv as *const _ as *mut _,
10880                &nc as *const _ as *mut _,
10881                &rqi as *const _ as *mut _,
10882                &rki as *const _ as *mut _,
10883                &ppos as *const _ as *mut _,
10884                &nhq as *const _ as *mut _,
10885                &nhk as *const _ as *mut _,
10886                &theta_scale as *const _ as *mut _,
10887                &freq_scale as *const _ as *mut _,
10888                &pff as *const _ as *mut _,
10889                &eps as *const _ as *mut _,
10890            ];
10891            unsafe {
10892                self.launch_pdl(
10893                    "rms_norm_qkv_rope_cat_f32",
10894                    (rows as u32, 1, 1),
10895                    (rms_block(), 1, 1),
10896                    &mut ps,
10897                )?;
10898            }
10899            return Ok(());
10900        }
10901        let f = self.func("rms_norm_qkv_rope_cat_f32");
10902        let cfg = LaunchConfig {
10903            grid_dim: (rows as u32, 1, 1),
10904            block_dim: (rms_block(), 1, 1),
10905            shared_mem_bytes: 0,
10906        };
10907        let __s_b = self.gpu.stream();
10908        let mut b = __s_b.launch_builder(&f);
10909        match ff {
10910            Some(t) => {
10911                b.arg(qkv)
10912                    .arg(wq)
10913                    .arg(wk)
10914                    .arg(wv)
10915                    .arg(&mut *q)
10916                    .arg(&mut *k)
10917                    .arg(&mut *v)
10918                    .arg(&nc)
10919                    .arg(&rqi)
10920                    .arg(&rki)
10921                    .arg(pos)
10922                    .arg(&nhq)
10923                    .arg(&nhk)
10924                    .arg(&theta_scale)
10925                    .arg(&freq_scale)
10926                    .arg(t)
10927                    .arg(&eps);
10928                unsafe {
10929                    b.launch(cfg)?;
10930                }
10931            }
10932            None => {
10933                let null: u64 = 0;
10934                b.arg(qkv)
10935                    .arg(wq)
10936                    .arg(wk)
10937                    .arg(wv)
10938                    .arg(&mut *q)
10939                    .arg(&mut *k)
10940                    .arg(&mut *v)
10941                    .arg(&nc)
10942                    .arg(&rqi)
10943                    .arg(&rki)
10944                    .arg(pos)
10945                    .arg(&nhq)
10946                    .arg(&nhk)
10947                    .arg(&theta_scale)
10948                    .arg(&freq_scale)
10949                    .arg(&null)
10950                    .arg(&eps);
10951                unsafe {
10952                    b.launch(cfg)?;
10953                }
10954            }
10955        }
10956        Ok(())
10957    }
10958
10959    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10960    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10961    /// ([`Engine::full_width_rope_only`]).
10962    #[allow(clippy::too_many_arguments)]
10963    pub fn rms_norm_qkv_rope(
10964        &self,
10965        q0: &CudaSlice<f32>,
10966        k0: &CudaSlice<f32>,
10967        v0: &CudaSlice<f32>,
10968        wq: &CudaSlice<f32>,
10969        wk: &CudaSlice<f32>,
10970        wv: &CudaSlice<f32>,
10971        q: &mut CudaSlice<f32>,
10972        k: &mut CudaSlice<f32>,
10973        v: &mut CudaSlice<f32>,
10974        head_dim: usize,
10975        n_rot: usize,
10976        rq: usize,
10977        rk: usize,
10978        pos: &CudaSlice<i32>,
10979        nh_q: usize,
10980        nh_k: usize,
10981        base: f32,
10982        freq_scale: f32,
10983        ff: Option<&CudaSlice<f32>>,
10984        eps: f32,
10985    ) -> Result<(), Box<dyn std::error::Error>> {
10986        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10987        let f = self.func("rms_norm_qkv_rope_f32");
10988        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10989        let cfg = LaunchConfig {
10990            grid_dim: (rows as u32, 1, 1),
10991            block_dim: (rms_block(), 1, 1),
10992            shared_mem_bytes: 0,
10993        };
10994        let theta_scale = base.powf(-2.0 / head_dim as f32);
10995        let (nc, rqi, rki, nhq, nhk) = (
10996            head_dim as i32,
10997            rq as i32,
10998            rk as i32,
10999            nh_q as i32,
11000            nh_k as i32,
11001        );
11002        let __s_b = self.gpu.stream();
11003        let mut b = __s_b.launch_builder(&f);
11004        match ff {
11005            Some(t) => {
11006                b.arg(q0)
11007                    .arg(k0)
11008                    .arg(v0)
11009                    .arg(wq)
11010                    .arg(wk)
11011                    .arg(wv)
11012                    .arg(&mut *q)
11013                    .arg(&mut *k)
11014                    .arg(&mut *v)
11015                    .arg(&nc)
11016                    .arg(&rqi)
11017                    .arg(&rki)
11018                    .arg(pos)
11019                    .arg(&nhq)
11020                    .arg(&nhk)
11021                    .arg(&theta_scale)
11022                    .arg(&freq_scale)
11023                    .arg(t)
11024                    .arg(&eps);
11025                unsafe {
11026                    b.launch(cfg)?;
11027                }
11028            }
11029            None => {
11030                let null: u64 = 0;
11031                b.arg(q0)
11032                    .arg(k0)
11033                    .arg(v0)
11034                    .arg(wq)
11035                    .arg(wk)
11036                    .arg(wv)
11037                    .arg(&mut *q)
11038                    .arg(&mut *k)
11039                    .arg(&mut *v)
11040                    .arg(&nc)
11041                    .arg(&rqi)
11042                    .arg(&rki)
11043                    .arg(pos)
11044                    .arg(&nhq)
11045                    .arg(&nhk)
11046                    .arg(&theta_scale)
11047                    .arg(&freq_scale)
11048                    .arg(&null)
11049                    .arg(&eps);
11050                unsafe {
11051                    b.launch(cfg)?;
11052                }
11053            }
11054        }
11055        Ok(())
11056    }
11057
11058    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
11059    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
11060    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
11061    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
11062    /// ([`Engine::full_width_rope_only`]).
11063    #[allow(clippy::too_many_arguments)]
11064    pub fn rms_norm_qkv_rope_append_dc(
11065        &self,
11066        q0: &CudaSlice<f32>,
11067        k0: &CudaSlice<f32>,
11068        v0: &CudaSlice<f32>,
11069        wq: &CudaSlice<f32>,
11070        wk: &CudaSlice<f32>,
11071        wv: &CudaSlice<f32>,
11072        q: &mut CudaSlice<f32>,
11073        k: &mut CudaSlice<f32>,
11074        v: &mut CudaSlice<f32>,
11075        head_dim: usize,
11076        n_rot: usize,
11077        rq: usize,
11078        rk: usize,
11079        pos: &CudaSlice<i32>,
11080        nh_q: usize,
11081        nh_k: usize,
11082        base: f32,
11083        freq_scale: f32,
11084        ff: Option<&CudaSlice<f32>>,
11085        eps: f32,
11086        kc: &mut CudaSlice<u8>,
11087        vc: &mut CudaSlice<u8>,
11088        t_dev: &CudaSlice<i32>,
11089        k_tok_bytes: usize,
11090        v_tok_bytes: usize,
11091        g: bool,
11092    ) -> Result<(), Box<dyn std::error::Error>> {
11093        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11094        let rows = rq + rk + rk;
11095        let theta_scale = base.powf(-2.0 / head_dim as f32);
11096        let (nc, rqi, rki, nhq, nhk) = (
11097            head_dim as i32,
11098            rq as i32,
11099            rk as i32,
11100            nh_q as i32,
11101            nh_k as i32,
11102        );
11103        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11104        if Self::pdl_on() && Self::pdl_wb_on() {
11105            use cudarc::driver::{DevicePtr, DevicePtrMut};
11106            let s = &self.gpu.stream();
11107            let (p0, _a0) = q0.device_ptr(s);
11108            let (p1, _a1) = k0.device_ptr(s);
11109            let (p2, _a2) = v0.device_ptr(s);
11110            let (pwq, _a3) = wq.device_ptr(s);
11111            let (pwk, _a4) = wk.device_ptr(s);
11112            let (pwv, _a5) = wv.device_ptr(s);
11113            let (pq, _a6) = q.device_ptr_mut(s);
11114            let (pk, _a7) = k.device_ptr_mut(s);
11115            let (pv, _a8) = v.device_ptr_mut(s);
11116            let (pp, _a9) = pos.device_ptr(s);
11117            let pff: u64 = match ff {
11118                Some(t) => {
11119                    let (p, _gg) = t.device_ptr(s);
11120                    p as u64
11121                }
11122                None => 0,
11123            };
11124            let (pkc, _a10) = kc.device_ptr_mut(s);
11125            let (pvc, _a11) = vc.device_ptr_mut(s);
11126            let (pt, _a12) = t_dev.device_ptr(s);
11127            let mut ps = [
11128                &p0 as *const _ as *mut std::ffi::c_void,
11129                &p1 as *const _ as *mut _,
11130                &p2 as *const _ as *mut _,
11131                &pwq as *const _ as *mut _,
11132                &pwk as *const _ as *mut _,
11133                &pwv as *const _ as *mut _,
11134                &pq as *const _ as *mut _,
11135                &pk as *const _ as *mut _,
11136                &pv as *const _ as *mut _,
11137                &nc as *const _ as *mut _,
11138                &rqi as *const _ as *mut _,
11139                &rki as *const _ as *mut _,
11140                &pp as *const _ as *mut _,
11141                &nhq as *const _ as *mut _,
11142                &nhk as *const _ as *mut _,
11143                &theta_scale as *const _ as *mut _,
11144                &freq_scale as *const _ as *mut _,
11145                &pff as *const _ as *mut _,
11146                &eps as *const _ as *mut _,
11147                &pkc as *const _ as *mut _,
11148                &pvc as *const _ as *mut _,
11149                &pt as *const _ as *mut _,
11150                &ktb as *const _ as *mut _,
11151                &vtb as *const _ as *mut _,
11152            ];
11153            unsafe {
11154                self.launch_pdl_flash(
11155                    g,
11156                    "rms_norm_qkv_rope_append_dc_f32",
11157                    (rows as u32, 1, 1),
11158                    (rms_block(), 1, 1),
11159                    0,
11160                    &mut ps,
11161                )?;
11162            }
11163            return Ok(());
11164        }
11165        let f = if g {
11166            self.func_g("rms_norm_qkv_rope_append_dc_f32")
11167        } else {
11168            self.func("rms_norm_qkv_rope_append_dc_f32")
11169        };
11170        let cfg = LaunchConfig {
11171            grid_dim: (rows as u32, 1, 1),
11172            block_dim: (rms_block(), 1, 1),
11173            shared_mem_bytes: 0,
11174        };
11175        let __s_b = self.gpu.stream();
11176        let mut b = __s_b.launch_builder(&f);
11177        match ff {
11178            Some(t) => {
11179                b.arg(q0)
11180                    .arg(k0)
11181                    .arg(v0)
11182                    .arg(wq)
11183                    .arg(wk)
11184                    .arg(wv)
11185                    .arg(&mut *q)
11186                    .arg(&mut *k)
11187                    .arg(&mut *v)
11188                    .arg(&nc)
11189                    .arg(&rqi)
11190                    .arg(&rki)
11191                    .arg(pos)
11192                    .arg(&nhq)
11193                    .arg(&nhk)
11194                    .arg(&theta_scale)
11195                    .arg(&freq_scale)
11196                    .arg(t)
11197                    .arg(&eps)
11198                    .arg(&mut *kc)
11199                    .arg(&mut *vc)
11200                    .arg(t_dev)
11201                    .arg(&ktb)
11202                    .arg(&vtb);
11203                unsafe {
11204                    b.launch(cfg)?;
11205                }
11206            }
11207            None => {
11208                let null: u64 = 0;
11209                b.arg(q0)
11210                    .arg(k0)
11211                    .arg(v0)
11212                    .arg(wq)
11213                    .arg(wk)
11214                    .arg(wv)
11215                    .arg(&mut *q)
11216                    .arg(&mut *k)
11217                    .arg(&mut *v)
11218                    .arg(&nc)
11219                    .arg(&rqi)
11220                    .arg(&rki)
11221                    .arg(pos)
11222                    .arg(&nhq)
11223                    .arg(&nhk)
11224                    .arg(&theta_scale)
11225                    .arg(&freq_scale)
11226                    .arg(&null)
11227                    .arg(&eps)
11228                    .arg(&mut *kc)
11229                    .arg(&mut *vc)
11230                    .arg(t_dev)
11231                    .arg(&ktb)
11232                    .arg(&vtb);
11233                unsafe {
11234                    b.launch(cfg)?;
11235                }
11236            }
11237        }
11238        Ok(())
11239    }
11240
11241    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
11242    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
11243    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
11244    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
11245    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
11246    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
11247    /// `head_dim` ([`Engine::full_width_rope_only`]).
11248    #[allow(clippy::too_many_arguments)]
11249    pub fn rms_norm_qkv_rope_append(
11250        &self,
11251        q0: &CudaSlice<f32>,
11252        k0: &CudaSlice<f32>,
11253        v0: &CudaSlice<f32>,
11254        wq: &CudaSlice<f32>,
11255        wk: &CudaSlice<f32>,
11256        wv: &CudaSlice<f32>,
11257        q: &mut CudaSlice<f32>,
11258        k: &mut CudaSlice<f32>,
11259        v: &mut CudaSlice<f32>,
11260        head_dim: usize,
11261        n_rot: usize,
11262        rq: usize,
11263        rk: usize,
11264        pos: &CudaSlice<i32>,
11265        nh_q: usize,
11266        nh_k: usize,
11267        base: f32,
11268        freq_scale: f32,
11269        ff: Option<&CudaSlice<f32>>,
11270        eps: f32,
11271        kc: &mut CudaSlice<u8>,
11272        vc: &mut CudaSlice<u8>,
11273        t: usize,
11274        k_tok_bytes: usize,
11275        v_tok_bytes: usize,
11276        g: bool,
11277    ) -> Result<(), Box<dyn std::error::Error>> {
11278        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11279        let rows = rq + rk + rk;
11280        let theta_scale = base.powf(-2.0 / head_dim as f32);
11281        let (nc, rqi, rki, nhq, nhk) = (
11282            head_dim as i32,
11283            rq as i32,
11284            rk as i32,
11285            nh_q as i32,
11286            nh_k as i32,
11287        );
11288        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11289        let ti = t as i32;
11290        if Self::pdl_on() && Self::pdl_wb_on() {
11291            use cudarc::driver::{DevicePtr, DevicePtrMut};
11292            let s = &self.gpu.stream();
11293            let (p0, _a0) = q0.device_ptr(s);
11294            let (p1, _a1) = k0.device_ptr(s);
11295            let (p2, _a2) = v0.device_ptr(s);
11296            let (pwq, _a3) = wq.device_ptr(s);
11297            let (pwk, _a4) = wk.device_ptr(s);
11298            let (pwv, _a5) = wv.device_ptr(s);
11299            let (pq, _a6) = q.device_ptr_mut(s);
11300            let (pk, _a7) = k.device_ptr_mut(s);
11301            let (pv, _a8) = v.device_ptr_mut(s);
11302            let (pp, _a9) = pos.device_ptr(s);
11303            let pff: u64 = match ff {
11304                Some(t) => {
11305                    let (p, _gg) = t.device_ptr(s);
11306                    p as u64
11307                }
11308                None => 0,
11309            };
11310            let (pkc, _a10) = kc.device_ptr_mut(s);
11311            let (pvc, _a11) = vc.device_ptr_mut(s);
11312            let mut ps = [
11313                &p0 as *const _ as *mut std::ffi::c_void,
11314                &p1 as *const _ as *mut _,
11315                &p2 as *const _ as *mut _,
11316                &pwq as *const _ as *mut _,
11317                &pwk as *const _ as *mut _,
11318                &pwv as *const _ as *mut _,
11319                &pq as *const _ as *mut _,
11320                &pk as *const _ as *mut _,
11321                &pv as *const _ as *mut _,
11322                &nc as *const _ as *mut _,
11323                &rqi as *const _ as *mut _,
11324                &rki as *const _ as *mut _,
11325                &pp as *const _ as *mut _,
11326                &nhq as *const _ as *mut _,
11327                &nhk as *const _ as *mut _,
11328                &theta_scale as *const _ as *mut _,
11329                &freq_scale as *const _ as *mut _,
11330                &pff as *const _ as *mut _,
11331                &eps as *const _ as *mut _,
11332                &pkc as *const _ as *mut _,
11333                &pvc as *const _ as *mut _,
11334                &ti as *const _ as *mut _,
11335                &ktb as *const _ as *mut _,
11336                &vtb as *const _ as *mut _,
11337            ];
11338            unsafe {
11339                self.launch_pdl_flash(
11340                    g,
11341                    "rms_norm_qkv_rope_append_f32",
11342                    (rows as u32, 1, 1),
11343                    (rms_block(), 1, 1),
11344                    0,
11345                    &mut ps,
11346                )?;
11347            }
11348            return Ok(());
11349        }
11350        let f = if g {
11351            self.func_g("rms_norm_qkv_rope_append_f32")
11352        } else {
11353            self.func("rms_norm_qkv_rope_append_f32")
11354        };
11355        let cfg = LaunchConfig {
11356            grid_dim: (rows as u32, 1, 1),
11357            block_dim: (rms_block(), 1, 1),
11358            shared_mem_bytes: 0,
11359        };
11360        let __s_b = self.gpu.stream();
11361        let mut b = __s_b.launch_builder(&f);
11362        let null: u64 = 0;
11363        b.arg(q0)
11364            .arg(k0)
11365            .arg(v0)
11366            .arg(wq)
11367            .arg(wk)
11368            .arg(wv)
11369            .arg(&mut *q)
11370            .arg(&mut *k)
11371            .arg(&mut *v)
11372            .arg(&nc)
11373            .arg(&rqi)
11374            .arg(&rki)
11375            .arg(pos)
11376            .arg(&nhq)
11377            .arg(&nhk)
11378            .arg(&theta_scale)
11379            .arg(&freq_scale);
11380        match ff {
11381            Some(t) => {
11382                b.arg(t);
11383            }
11384            None => {
11385                b.arg(&null);
11386            }
11387        }
11388        b.arg(&eps)
11389            .arg(&mut *kc)
11390            .arg(&mut *vc)
11391            .arg(&ti)
11392            .arg(&ktb)
11393            .arg(&vtb);
11394        unsafe {
11395            b.launch(cfg)?;
11396        }
11397        Ok(())
11398    }
11399
11400    pub fn add_q8_1(
11401        &self,
11402        a: &CudaSlice<f32>,
11403        b: &CudaSlice<f32>,
11404        res: &mut CudaSlice<f32>,
11405        ncols: usize,
11406        nrows: usize,
11407    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11408        debug_assert!(ncols % 128 == 0);
11409        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11410        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11411        let f = self.func("add_q8_1_f32");
11412        let cfg = LaunchConfig {
11413            grid_dim: (nrows as u32, 1, 1),
11414            block_dim: (rms_block(), 1, 1),
11415            shared_mem_bytes: 0,
11416        };
11417        let nc = ncols as i32;
11418        let __s_b2 = self.gpu.stream();
11419        let mut b2 = __s_b2.launch_builder(&f);
11420        b2.arg(a)
11421            .arg(b)
11422            .arg(&mut *res)
11423            .arg(&mut out_q)
11424            .arg(&mut out_d)
11425            .arg(&nc);
11426        unsafe {
11427            b2.launch(cfg)?;
11428        }
11429        Ok((out_q, out_d))
11430    }
11431
11432    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
11433    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
11434    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
11435    pub fn rms_pre_add_q8_1(
11436        &self,
11437        a: &CudaSlice<f32>,
11438        wa: &CudaSlice<f32>,
11439        b: &CudaSlice<f32>,
11440        res: &mut CudaSlice<f32>,
11441        ncols: usize,
11442        nrows: usize,
11443        eps: f32,
11444    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11445        debug_assert!(ncols % 128 == 0);
11446        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11447        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11448        let f = self.func("rms_pre_add_q8_1_f32");
11449        let cfg = LaunchConfig {
11450            grid_dim: (nrows as u32, 1, 1),
11451            block_dim: (rms_block(), 1, 1),
11452            shared_mem_bytes: 0,
11453        };
11454        let (nc, ep) = (ncols as i32, eps);
11455        let __s_b2 = self.gpu.stream();
11456        let mut b2 = __s_b2.launch_builder(&f);
11457        b2.arg(a)
11458            .arg(wa)
11459            .arg(b)
11460            .arg(&mut *res)
11461            .arg(&mut out_q)
11462            .arg(&mut out_d)
11463            .arg(&nc)
11464            .arg(&ep);
11465        unsafe {
11466            b2.launch(cfg)?;
11467        }
11468        Ok((out_q, out_d))
11469    }
11470
11471    /// L2 norm per row (head_dim), no weight.
11472    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
11473    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
11474    pub fn l2_v2_on(ncols: usize) -> bool {
11475        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11476    }
11477
11478    pub fn l2_norm_pp(
11479        &self,
11480        x: &CudaSlice<f32>,
11481        dst: &mut CudaSlice<f32>,
11482        dst16: Option<&mut CudaSlice<u8>>,
11483        ncols: usize,
11484        nrows: usize,
11485        eps: f32,
11486    ) -> Result<(), Box<dyn std::error::Error>> {
11487        if Self::l2_v2_on(ncols) {
11488            let f = self.func("l2_norm_pp_v2_f32");
11489            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
11490            let cfg = LaunchConfig {
11491                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11492                block_dim: (256, 1, 1),
11493                shared_mem_bytes: 0,
11494            };
11495            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11496            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
11497            let d16: u64 = match dst16 {
11498                Some(d) => self.addr_u8(d),
11499                None => 0,
11500            };
11501            let __s_b = self.gpu.stream();
11502            let mut b = __s_b.launch_builder(&f);
11503            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11504            unsafe {
11505                b.launch(cfg)?;
11506            }
11507            return Ok(());
11508        }
11509        self.l2_norm(x, dst, ncols, nrows, eps)
11510    }
11511
11512    pub fn l2_norm(
11513        &self,
11514        x: &CudaSlice<f32>,
11515        dst: &mut CudaSlice<f32>,
11516        ncols: usize,
11517        nrows: usize,
11518        eps: f32,
11519    ) -> Result<(), Box<dyn std::error::Error>> {
11520        let f = self.func("l2_norm_f32");
11521        let cfg = LaunchConfig {
11522            grid_dim: (nrows as u32, 1, 1),
11523            block_dim: (256, 1, 1),
11524            shared_mem_bytes: 0,
11525        };
11526        let (nc, e) = (ncols as i32, eps);
11527        let __s_b = self.gpu.stream();
11528        let mut b = __s_b.launch_builder(&f);
11529        b.arg(x).arg(dst).arg(&nc).arg(&e);
11530        unsafe {
11531            b.launch(cfg)?;
11532        }
11533        Ok(())
11534    }
11535
11536    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
11537    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
11538    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
11539    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
11540    /// propagate through gdn_scan and flip argmax on marginal logits.
11541    pub fn l2_norm_decode(
11542        &self,
11543        x: &CudaSlice<f32>,
11544        dst: &mut CudaSlice<f32>,
11545        ncols: usize,
11546        nrows: usize,
11547        eps: f32,
11548    ) -> Result<(), Box<dyn std::error::Error>> {
11549        let f = self.func("l2_norm_f32");
11550        let cfg = LaunchConfig {
11551            grid_dim: (nrows as u32, 1, 1),
11552            block_dim: (32, 1, 1),
11553            shared_mem_bytes: 0,
11554        };
11555        let (nc, e) = (ncols as i32, eps);
11556        let __s_b = self.gpu.stream();
11557        let mut b = __s_b.launch_builder(&f);
11558        b.arg(x).arg(dst).arg(&nc).arg(&e);
11559        unsafe {
11560            b.launch(cfg)?;
11561        }
11562        Ok(())
11563    }
11564
11565    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
11566    pub fn rope_neox(
11567        &self,
11568        x: &mut CudaSlice<f32>,
11569        pos: &CudaSlice<i32>,
11570        head_dim: usize,
11571        n_dims: usize,
11572        n_heads: usize,
11573        n_tokens: usize,
11574        freq_base: f32,
11575        freq_scale: f32,
11576    ) -> Result<(), Box<dyn std::error::Error>> {
11577        let f = self.func("rope_neox_f32");
11578        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11579        let grid = (n_heads * n_tokens) as u32;
11580        let cfg = LaunchConfig {
11581            grid_dim: (grid, 1, 1),
11582            block_dim: ((head_dim / 2) as u32, 1, 1),
11583            shared_mem_bytes: 0,
11584        };
11585        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11586        let __s_b = self.gpu.stream();
11587        let mut b = __s_b.launch_builder(&f);
11588        b.arg(x)
11589            .arg(pos)
11590            .arg(&hd)
11591            .arg(&nd)
11592            .arg(&nh)
11593            .arg(&theta_scale)
11594            .arg(&freq_scale);
11595        unsafe {
11596            b.launch(cfg)?;
11597        }
11598        Ok(())
11599    }
11600
11601    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
11602    pub fn rope_neox_ff(
11603        &self,
11604        x: &mut CudaSlice<f32>,
11605        pos: &CudaSlice<i32>,
11606        head_dim: usize,
11607        n_dims: usize,
11608        n_heads: usize,
11609        n_tokens: usize,
11610        freq_base: f32,
11611        freq_scale: f32,
11612        ff: &CudaSlice<f32>,
11613    ) -> Result<(), Box<dyn std::error::Error>> {
11614        let f = self.func("rope_neox_ff_f32");
11615        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11616        let grid = (n_heads * n_tokens) as u32;
11617        let cfg = LaunchConfig {
11618            grid_dim: (grid, 1, 1),
11619            block_dim: ((head_dim / 2) as u32, 1, 1),
11620            shared_mem_bytes: 0,
11621        };
11622        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11623        let __s_b = self.gpu.stream();
11624        let mut b = __s_b.launch_builder(&f);
11625        b.arg(x)
11626            .arg(pos)
11627            .arg(&hd)
11628            .arg(&nd)
11629            .arg(&nh)
11630            .arg(&theta_scale)
11631            .arg(&freq_scale)
11632            .arg(ff);
11633        unsafe {
11634            b.launch(cfg)?;
11635        }
11636        Ok(())
11637    }
11638
11639    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
11640    #[allow(clippy::too_many_arguments)]
11641    pub fn rope_neox2(
11642        &self,
11643        q: &mut CudaSlice<f32>,
11644        k: &mut CudaSlice<f32>,
11645        pos: &CudaSlice<i32>,
11646        head_dim: usize,
11647        n_dims: usize,
11648        nh_q: usize,
11649        nh_k: usize,
11650        n_tokens: usize,
11651        freq_base: f32,
11652        freq_scale: f32,
11653        ff: Option<&CudaSlice<f32>>,
11654    ) -> Result<(), Box<dyn std::error::Error>> {
11655        let f = self.func("rope_neox2_f32");
11656        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11657        let grid = ((nh_q + nh_k) * n_tokens) as u32;
11658        let cfg = LaunchConfig {
11659            grid_dim: (grid, 1, 1),
11660            block_dim: ((head_dim / 2) as u32, 1, 1),
11661            shared_mem_bytes: 0,
11662        };
11663        let (hd, nd, nq, nk, nt) = (
11664            head_dim as i32,
11665            n_dims as i32,
11666            nh_q as i32,
11667            nh_k as i32,
11668            n_tokens as i32,
11669        );
11670        let __s_b = self.gpu.stream();
11671        let mut b = __s_b.launch_builder(&f);
11672        b.arg(q)
11673            .arg(k)
11674            .arg(pos)
11675            .arg(&hd)
11676            .arg(&nd)
11677            .arg(&nq)
11678            .arg(&nk)
11679            .arg(&nt)
11680            .arg(&theta_scale)
11681            .arg(&freq_scale);
11682        match ff {
11683            Some(ffv) => {
11684                b.arg(ffv);
11685                unsafe {
11686                    b.launch(cfg)?;
11687                }
11688            }
11689            None => {
11690                let null: u64 = 0;
11691                b.arg(&null);
11692                unsafe {
11693                    b.launch(cfg)?;
11694                }
11695            }
11696        }
11697        Ok(())
11698    }
11699
11700    /// gemma4 R1: dst = GELU_tanh(gate) * up.
11701    pub fn gelu_tanh_mul(
11702        &self,
11703        gate: &CudaSlice<f32>,
11704        up: &CudaSlice<f32>,
11705        dst: &mut CudaSlice<f32>,
11706        n: usize,
11707    ) -> Result<(), Box<dyn std::error::Error>> {
11708        let f = self.func("gelu_tanh_mul_f32");
11709        let cfg = LaunchConfig::for_num_elems(n as u32);
11710        let ni = n as i32;
11711        let __s_b = self.gpu.stream();
11712        let mut b = __s_b.launch_builder(&f);
11713        b.arg(gate).arg(up).arg(dst).arg(&ni);
11714        unsafe {
11715            b.launch(cfg)?;
11716        }
11717        Ok(())
11718    }
11719
11720    pub fn silu_mul(
11721        &self,
11722        gate: &CudaSlice<f32>,
11723        up: &CudaSlice<f32>,
11724        dst: &mut CudaSlice<f32>,
11725        n: usize,
11726    ) -> Result<(), Box<dyn std::error::Error>> {
11727        let f = self.func("silu_mul_f32");
11728        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11729        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11730        let ni = n as i32;
11731        let __s_b = self.gpu.stream();
11732        let mut b = __s_b.launch_builder(&f);
11733        b.arg(gate).arg(up).arg(dst).arg(&ni);
11734        unsafe {
11735            b.launch(cfg)?;
11736        }
11737        Ok(())
11738    }
11739
11740    /// SwiGLU twin using Memra's host-matching expf transcription.
11741    pub fn silu_mul_host_expf(
11742        &self,
11743        gate: &CudaSlice<f32>,
11744        up: &CudaSlice<f32>,
11745        dst: &mut CudaSlice<f32>,
11746        n: usize,
11747    ) -> Result<(), Box<dyn std::error::Error>> {
11748        let f = self.func("silu_mul_host_expf_f32");
11749        let cfg = LaunchConfig::for_num_elems(n as u32);
11750        let ni = n as i32;
11751        let __s_b = self.gpu.stream();
11752        let mut b = __s_b.launch_builder(&f);
11753        b.arg(gate).arg(up).arg(dst).arg(&ni);
11754        unsafe {
11755            b.launch(cfg)?;
11756        }
11757        Ok(())
11758    }
11759
11760    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
11761    pub fn silu_clamped_mul_host_expf(
11762        &self,
11763        gate: &CudaSlice<f32>,
11764        up: &CudaSlice<f32>,
11765        limit: f32,
11766        dst: &mut CudaSlice<f32>,
11767        n: usize,
11768    ) -> Result<(), Box<dyn std::error::Error>> {
11769        if !limit.is_finite() || limit <= 0.0 {
11770            return Err(
11771                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11772            );
11773        }
11774        let f = self.func("silu_clamped_mul_host_expf_f32");
11775        let cfg = LaunchConfig::for_num_elems(n as u32);
11776        let ni = n as i32;
11777        let __s_b = self.gpu.stream();
11778        let mut b = __s_b.launch_builder(&f);
11779        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11780        unsafe {
11781            b.launch(cfg)?;
11782        }
11783        Ok(())
11784    }
11785
11786    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
11787    /// for the down projection — kills the standalone convert pass. Bit-identical class.
11788    pub fn silu_mul_f16out(
11789        &self,
11790        gate: &CudaSlice<f32>,
11791        up: &CudaSlice<f32>,
11792        dst: &mut CudaSlice<f32>,
11793        dst16: &mut CudaSlice<u8>,
11794        n: usize,
11795    ) -> Result<(), Box<dyn std::error::Error>> {
11796        let f = self.func("silu_mul_f16out_f32");
11797        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11798        let ni = n as i32;
11799        let __s_b = self.gpu.stream();
11800        let mut b = __s_b.launch_builder(&f);
11801        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11802        unsafe {
11803            b.launch(cfg)?;
11804        }
11805        Ok(())
11806    }
11807
11808    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
11809    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
11810    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
11811    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
11812    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
11813    /// launches per dense FFN layer (the gate+up post-matmul scales).
11814    pub fn silu_mul_scaled(
11815        &self,
11816        gate: &CudaSlice<f32>,
11817        up: &CudaSlice<f32>,
11818        gs: f32,
11819        us: f32,
11820        dst: &mut CudaSlice<f32>,
11821        n: usize,
11822    ) -> Result<(), Box<dyn std::error::Error>> {
11823        let f = self.func("silu_mul_scaled_f32");
11824        let cfg = LaunchConfig::for_num_elems(n as u32);
11825        let ni = n as i32;
11826        let (gsf, usf) = (gs, us);
11827        let __s_b = self.gpu.stream();
11828        let mut b = __s_b.launch_builder(&f);
11829        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11830        unsafe {
11831            b.launch(cfg)?;
11832        }
11833        Ok(())
11834    }
11835
11836    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
11837    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
11838    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
11839    #[allow(clippy::too_many_arguments)]
11840    pub fn swigluoai_mul_scaled(
11841        &self,
11842        gate: &CudaSlice<f32>,
11843        up: &CudaSlice<f32>,
11844        gs: f32,
11845        us: f32,
11846        alpha: f32,
11847        limit: f32,
11848        dst: &mut CudaSlice<f32>,
11849        n: usize,
11850    ) -> Result<(), Box<dyn std::error::Error>> {
11851        let f = self.func("swigluoai_mul_scaled_f32");
11852        let cfg = LaunchConfig::for_num_elems(n as u32);
11853        let ni = n as i32;
11854        let __s_b = self.gpu.stream();
11855        let mut b = __s_b.launch_builder(&f);
11856        b.arg(gate)
11857            .arg(up)
11858            .arg(&gs)
11859            .arg(&us)
11860            .arg(&alpha)
11861            .arg(&limit)
11862            .arg(dst)
11863            .arg(&ni);
11864        unsafe {
11865            b.launch(cfg)?;
11866        }
11867        Ok(())
11868    }
11869
11870    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11871    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11872    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11873    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11874    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11875    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11876    /// n must be a multiple of 32 (n_ff always is).
11877    pub fn silu_mul_scaled_q8_1(
11878        &self,
11879        gate: &CudaSlice<f32>,
11880        up: &CudaSlice<f32>,
11881        gs: f32,
11882        us: f32,
11883        n: usize,
11884    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11885        let f = self.func("silu_mul_scaled_q8_1");
11886        let nblk = n / 32;
11887        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11888        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11889        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11890        let cfg = LaunchConfig::for_num_elems(n as u32);
11891        let (gsf, usf, ni) = (gs, us, n as i32);
11892        let __s_b = self.gpu.stream();
11893        let mut b = __s_b.launch_builder(&f);
11894        b.arg(gate)
11895            .arg(up)
11896            .arg(&gsf)
11897            .arg(&usf)
11898            .arg(&mut aq)
11899            .arg(&mut ad)
11900            .arg(&ni);
11901        unsafe {
11902            b.launch(cfg)?;
11903        }
11904        Ok((aq, ad))
11905    }
11906
11907    pub fn add(
11908        &self,
11909        a: &CudaSlice<f32>,
11910        b_in: &CudaSlice<f32>,
11911        dst: &mut CudaSlice<f32>,
11912        n: usize,
11913    ) -> Result<(), Box<dyn std::error::Error>> {
11914        let f = self.func("add_f32");
11915        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11916        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11917        let ni = n as i32;
11918        let __s_bld = self.gpu.stream();
11919        let mut bld = __s_bld.launch_builder(&f);
11920        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11921        unsafe {
11922            bld.launch(cfg)?;
11923        }
11924        Ok(())
11925    }
11926
11927    pub fn mul(
11928        &self,
11929        a: &CudaSlice<f32>,
11930        b_in: &CudaSlice<f32>,
11931        dst: &mut CudaSlice<f32>,
11932        n: usize,
11933    ) -> Result<(), Box<dyn std::error::Error>> {
11934        let f = self.func("mul_f32");
11935        let cfg = LaunchConfig::for_num_elems(n as u32);
11936        let ni = n as i32;
11937        let __s_bld = self.gpu.stream();
11938        let mut bld = __s_bld.launch_builder(&f);
11939        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11940        unsafe {
11941            bld.launch(cfg)?;
11942        }
11943        Ok(())
11944    }
11945
11946    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11947    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11948    pub fn matmul(
11949        &self,
11950        w: &crate::model::GpuTensor,
11951        x: &CudaSlice<f32>,
11952        m: usize,
11953    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11954        use crate::model::GpuTensor;
11955        let in_f = w.in_features();
11956        let out_f = w.out_features();
11957        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11958        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11959        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11960        // gives nothing). Quantize the activation once here then call the GEMM.
11961        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11962        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11963        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11964        #[allow(non_snake_case)]
11965        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11966        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11967        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11968            usize::MAX
11969        } else {
11970            16usize
11971        };
11972
11973        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11974        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11975        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11976        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11977        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11978        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11979        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11980        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11981        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11982        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11983        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11984        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11985        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11986        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11987        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11988        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11989        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11990        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11991        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11992        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11993        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11994        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11995        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11996        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11997        if m >= GEMM_M_THRESHOLD {
11998            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11999                return Ok(y);
12000            }
12001            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
12002            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
12003            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
12004            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
12005            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
12006            // tile defaults differently by operand source.
12007            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
12008                return Ok(y);
12009            }
12010            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
12011            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
12012            if let Some(y) = self.try_f16_gemm(w, x, m)? {
12013                return Ok(y);
12014            }
12015        }
12016        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
12017        // m threshold the rest of this method uses:
12018        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
12019        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
12020        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
12021        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
12022        //     across every tier by construction with no batched twin needed.
12023        //
12024        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
12025        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
12026        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
12027        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
12028        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
12029        // arms is what makes sure it never gets there.
12030        if let GpuTensor::Quant { qtype, .. } = w {
12031            if *qtype == QT_F8_E4M3_BLK {
12032                if m >= GEMM_M_THRESHOLD {
12033                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
12034                        return Ok(y);
12035                    }
12036                }
12037                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12038                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12039                    return Ok(y);
12040                }
12041            }
12042        }
12043        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12044            return self.qmatvec_mmq(w, x, m);
12045        }
12046        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12047            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12048            return self.qmatvec_gemm(w, &aq, &ad, m);
12049        }
12050        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
12051        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
12052        if m >= GEMM_M_THRESHOLD {
12053            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12054                return Ok(y);
12055            }
12056        }
12057        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
12058        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
12059        // to Stage-A f32-dequant (the correctness oracle path).
12060        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12061        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
12062        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
12063        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
12064        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
12065        if m == 1 && fast {
12066            if let GpuTensor::Quant {
12067                bytes,
12068                qtype,
12069                row_bytes,
12070                rp,
12071                rp4,
12072                scale,
12073                ..
12074            } = w
12075            {
12076                if self.mmvq_supports(*qtype) {
12077                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
12078                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
12079                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
12080                    let (bytes, rp) = match rp4 {
12081                        Some(m4) => (m4, true),
12082                        None => (bytes, *rp),
12083                    };
12084                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12085                    return self.qmatvec_mmvq(
12086                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12087                    );
12088                }
12089            }
12090        }
12091        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
12092        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
12093        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
12094        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
12095        // block below. MEMRA_NO_BATCHED -> per-m path.
12096        //
12097        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
12098        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
12099        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
12100        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
12101        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
12102        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
12103        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
12104        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
12105        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
12106        if (2..=16).contains(&m)
12107            && fast
12108            && std::env::var("MEMRA_NO_BATCHED").is_err()
12109            && (m <= 4 || Self::b8_enabled())
12110        {
12111            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
12112            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
12113            // is present (rp4) — the mirror pick below then routes to the _rp family.
12114            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
12115            // because the native e4m3 row layout is already aligned and needs no mirror.
12116            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
12117            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
12118            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
12119            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
12120            let m_ok = m <= 8
12121                || matches!(w, GpuTensor::Quant { qtype, .. }
12122                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12123                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12124            if m_ok {
12125                if let GpuTensor::Quant {
12126                    bytes,
12127                    qtype,
12128                    row_bytes,
12129                    rp,
12130                    rp4,
12131                    ..
12132                } = w
12133                {
12134                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12135                        let (bytes, rp) = match rp4 {
12136                            Some(m4) => (m4, true),
12137                            None => (bytes, *rp),
12138                        };
12139                        let mcols = Self::batched_mcols(m);
12140                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12141                        let mut y = self.qmatvec_mmvq_batched(
12142                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12143                        )?;
12144                        if let GpuTensor::Quant { scale, .. } = w {
12145                            if *scale != 1.0 {
12146                                self.scale_inplace(&mut y, *scale, m * out_f)?;
12147                            }
12148                        }
12149                        return Ok(y);
12150                    }
12151                }
12152            }
12153        }
12154        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
12155        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
12156        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
12157        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
12158        // for this dtype, so the generic match below must never see it under `fast`.
12159        if fast {
12160            if let GpuTensor::Quant {
12161                bytes,
12162                qtype,
12163                row_bytes,
12164                scale,
12165                ..
12166            } = w
12167            {
12168                if *qtype == QT_F8_E4M3 {
12169                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12170                    return self.qmatvec_mmvq(
12171                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12172                    );
12173                }
12174            }
12175        }
12176        let mut y = match w {
12177            GpuTensor::Quant {
12178                bytes,
12179                qtype,
12180                row_bytes,
12181                ..
12182            } if fast && *qtype == QT_Q8_0 => {
12183                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12184            }
12185            GpuTensor::Quant {
12186                bytes,
12187                qtype,
12188                row_bytes,
12189                ..
12190            } if fast && *qtype == QT_Q4_K => {
12191                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12192            }
12193            GpuTensor::Quant {
12194                bytes,
12195                qtype,
12196                row_bytes,
12197                ..
12198            } if fast && *qtype == QT_Q6_K => {
12199                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12200            }
12201            GpuTensor::Quant {
12202                bytes,
12203                qtype,
12204                row_bytes,
12205                ..
12206            } if fast && *qtype == QT_Q5_K => {
12207                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12208            }
12209            GpuTensor::Quant {
12210                bytes,
12211                qtype,
12212                row_bytes,
12213                ..
12214            } if fast && *qtype == QT_Q3_K => {
12215                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12216            }
12217            GpuTensor::Quant {
12218                bytes,
12219                qtype,
12220                row_bytes,
12221                rp,
12222                ..
12223            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12224                if *rp {
12225                    "qmatvec_nvfp4_dp4a_rp"
12226                } else {
12227                    "qmatvec_nvfp4_dp4a"
12228                },
12229                &bytes.slice(0..bytes.len()),
12230                x,
12231                m,
12232                in_f,
12233                out_f,
12234                *row_bytes,
12235            )?,
12236            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
12237            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
12238            // anomaly (research/kat-anomaly-20260802/).
12239            GpuTensor::Quant {
12240                bytes,
12241                qtype,
12242                row_bytes,
12243                ..
12244            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12245                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12246            }
12247            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
12248            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
12249            // without first writing the matching kernel, or func() will panic
12250            // "kernel ... not in any fatbin".
12251            GpuTensor::Quant {
12252                bytes,
12253                qtype,
12254                row_bytes,
12255                rp,
12256                ..
12257            } =>
12258            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
12259            // deq(row,j) form cannot address the planes; same value/product order).
12260            {
12261                self.qmatvec(
12262                    bytes,
12263                    x,
12264                    m,
12265                    in_f,
12266                    out_f,
12267                    if *rp && *qtype == QT_NVFP4 {
12268                        QT_NVFP4_RP
12269                    } else {
12270                        *qtype
12271                    },
12272                    *row_bytes,
12273                )?
12274            }
12275            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12276            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
12277            // cuBLASLt f32 GEMV as the Float arm.
12278            GpuTensor::FloatBf16 { data, .. } => {
12279                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
12280                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
12281                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
12282                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
12283                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
12284                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
12285                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12286                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12287                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12288                    y
12289                } else {
12290                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12291                }
12292            }
12293        };
12294        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
12295        if let GpuTensor::Quant { scale, .. } = w {
12296            if *scale != 1.0 {
12297                self.scale_inplace(&mut y, *scale, m * out_f)?;
12298            }
12299        }
12300        Ok(y)
12301    }
12302
12303    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
12304    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
12305    ///
12306    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
12307    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
12308    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
12309    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
12310    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
12311    /// path must not pay an env lookup for a flag that is off.
12312    pub fn stage_a_raw_needed() -> bool {
12313        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12314        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12315    }
12316
12317    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
12318    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
12319    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12320        use crate::model::GpuTensor;
12321        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12322            return false;
12323        }
12324        match w {
12325            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
12326            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
12327            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
12328            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
12329            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
12330            // block class has no fused twin yet, so each of its projections takes its own launch.
12331            GpuTensor::Quant { qtype, .. } => {
12332                matches!(
12333                    *qtype,
12334                    QT_Q8_0
12335                        | QT_Q4_K
12336                        | QT_Q6_K
12337                        | QT_Q5_K
12338                        | QT_Q3_K
12339                        | QT_NVFP4
12340                        | QT_F8_E4M3
12341                        | QT_F8_E4M3_BLK
12342                        | QT_Q4_0
12343                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12344            }
12345            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12346        }
12347    }
12348
12349    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
12350    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
12351    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
12352    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
12353    pub fn matmul_pre(
12354        &self,
12355        w: &crate::model::GpuTensor,
12356        aq: &CudaSlice<i8>,
12357        ad: &CudaSlice<f32>,
12358        x_fallback: &CudaSlice<f32>,
12359        m: usize,
12360    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12361        use crate::model::GpuTensor;
12362        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
12363        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
12364        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
12365        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
12366        // rc=30013 dig, 2026-07-31).
12367        let x_raw_ok = x_fallback.len() >= m * w.in_features();
12368        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
12369        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
12370        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12371            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12372                return Ok(y);
12373            }
12374            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
12375            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
12376            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12377                return Ok(y);
12378            }
12379            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
12380            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12381                return Ok(y);
12382            }
12383        }
12384        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
12385        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
12386        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
12387        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
12388        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
12389        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12390            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12391                return Ok(y);
12392            }
12393        }
12394        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12395            return Ok(y);
12396        }
12397        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
12398        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
12399        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
12400        // aq/ad.
12401        if m >= 16
12402            && w.out_features() >= 128
12403            && self.mmq_supports(w)
12404            && !self.verify_exact_on()
12405            && x_raw_ok
12406        {
12407            return self.qmatvec_mmq(w, x_fallback, m);
12408        }
12409        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
12410        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
12411        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12412            if let Some(y) =
12413                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12414            {
12415                return Ok(y);
12416            }
12417        }
12418        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
12419        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
12420        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12421            return self.qmatvec_gemm(w, aq, ad, m);
12422        }
12423        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
12424        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
12425        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
12426        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
12427        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
12428        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
12429        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
12430        // which reads `m * in_f` floats out of a 0-byte allocation ->
12431        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
12432        // it poisons the context, so every LATER request in that process fails with an unrelated
12433        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
12434        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
12435        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
12436        // dense artifact and left the arm with no working truth instrument.
12437        //
12438        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
12439        // strictly better than an illegal address surfacing later at an unrelated sync point, and
12440        // an oracle that cannot run must say so rather than corrupt the context it runs in.
12441        if !self.uses_q8_1_fast(w) {
12442            if !x_raw_ok {
12443                return Err(format!(
12444                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12445                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12446                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
12447                     activation (see Engine::rms_norm_decode, which is bit-identical to \
12448                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12449                    x_fallback.len(),
12450                    m,
12451                    w.in_features(),
12452                    m * w.in_features()
12453                )
12454                .into());
12455            }
12456            return self.matmul(w, x_fallback, m);
12457        }
12458        let in_f = w.in_features();
12459        let out_f = w.out_features();
12460        let (bytes, qtype, row_bytes, scale, rp) = match w {
12461            GpuTensor::Quant {
12462                bytes,
12463                qtype,
12464                row_bytes,
12465                scale,
12466                rp,
12467                ..
12468            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12469            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12470        };
12471        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
12472        // the dp4a/oracle tails below keep the raw GGUF bytes.
12473        let (mbytes, mrp) = match w {
12474            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12475            _ => (bytes, rp),
12476        };
12477        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
12478        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
12479        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
12480        if m == 1 && self.mmvq_supports(qtype) {
12481            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12482        }
12483        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
12484        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
12485        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
12486        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
12487        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
12488        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
12489        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
12490        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
12491        // m=5..8 on the old per-m path (b8-tier-only seam).
12492        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
12493        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
12494        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
12495        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12496            && std::env::var("MEMRA_NO_BATCHED").is_err()
12497            && (m <= 4 || Self::b8_enabled())
12498            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
12499            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
12500            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
12501            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12502                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12503        {
12504            let mcols = Self::batched_mcols(m);
12505            return self.qmatvec_mmvq_batched(
12506                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12507            );
12508        }
12509        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
12510        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
12511        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
12512        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
12513        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
12514        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12515            let (b2, r2) = if qtype == QT_Q4_0 {
12516                (mbytes, mrp)
12517            } else {
12518                (bytes, rp)
12519            };
12520            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12521        }
12522        let name = match qtype {
12523            QT_Q8_0 => "qmatvec_q8_0_dp4a",
12524            QT_Q4_K => "qmatvec_q4_K_dp4a",
12525            QT_Q6_K => "qmatvec_q6_K_dp4a",
12526            QT_Q5_K => "qmatvec_q5_K_dp4a",
12527            QT_Q3_K => "qmatvec_q3_K_dp4a",
12528            QT_NVFP4 => {
12529                if rp {
12530                    "qmatvec_nvfp4_dp4a_rp"
12531                } else {
12532                    "qmatvec_nvfp4_dp4a"
12533                }
12534            }
12535            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12536            _ => unreachable!(),
12537        };
12538        let f = self.func(name);
12539        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
12540        let cfg = LaunchConfig {
12541            grid_dim: (out_f as u32, m as u32, 1),
12542            block_dim: (128, 1, 1),
12543            shared_mem_bytes: 0,
12544        };
12545        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12546        let __s_b = self.gpu.stream();
12547        let mut b = __s_b.launch_builder(&f);
12548        b.arg(bytes)
12549            .arg(aq)
12550            .arg(ad)
12551            .arg(&mut y)
12552            .arg(&inf)
12553            .arg(&outf)
12554            .arg(&mi)
12555            .arg(&rb);
12556        unsafe {
12557            b.launch(cfg)?;
12558        }
12559        if scale != 1.0 {
12560            self.scale_inplace(&mut y, scale, m * out_f)?;
12561        }
12562        Ok(y)
12563    }
12564
12565    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
12566    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
12567    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
12568    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
12569    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
12570    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
12571    /// reduce as m=1); this method just forces that path unconditionally.
12572    pub fn matmul_decode_exact(
12573        &self,
12574        w: &crate::model::GpuTensor,
12575        x: &CudaSlice<f32>,
12576        m: usize,
12577    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12578        use crate::model::GpuTensor;
12579        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
12580        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
12581        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
12582        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
12583        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
12584        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
12585        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
12586        if let GpuTensor::Float { data, .. } = w {
12587            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12588        }
12589        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
12590        // float linear (same n-independent reduction contract as the Float arm above).
12591        if let GpuTensor::FloatBf16 { data, .. } = w {
12592            let (in_f, out_f) = (w.in_features(), w.out_features());
12593            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
12594            // contract — the whole-weight f32 dequant disappears too).
12595            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12596                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12597                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12598                return Ok(y);
12599            }
12600            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12601        }
12602        if !self.uses_q8_1_fast(w) {
12603            return self.matmul(w, x, m);
12604        }
12605        let in_f = w.in_features();
12606        let out_f = w.out_features();
12607        let (bytes, qtype, row_bytes, scale, rp) = match w {
12608            GpuTensor::Quant {
12609                bytes,
12610                qtype,
12611                row_bytes,
12612                scale,
12613                rp,
12614                ..
12615            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12616            _ => return self.matmul(w, x, m),
12617        };
12618        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
12619        // which does its own mirror pick).
12620        let (bytes, rp) = match w {
12621            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12622            _ => (bytes, rp),
12623        };
12624        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12625        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
12626        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
12627        // (token,row) by construction, which is exactly what this method exists to guarantee.
12628        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12629            return Ok(y);
12630        }
12631        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
12632        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
12633        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
12634        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
12635        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
12636        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
12637        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
12638        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
12639        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12640            && std::env::var("MEMRA_NO_BATCHED").is_err()
12641            && (m <= 4 || Self::b8_enabled())
12642            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
12643            // no mirror precondition, `rp` selects the layout only.
12644            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12645                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12646        {
12647            let mcols = Self::batched_mcols(m);
12648            return self.qmatvec_mmvq_batched(
12649                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12650            );
12651        }
12652        if self.mmvq_supports(qtype) {
12653            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
12654            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
12655            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12656        }
12657        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
12658        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
12659        self.matmul_pre(w, &aq, &ad, x, m)
12660    }
12661
12662    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
12663    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
12664    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
12665    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
12666    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
12667    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
12668    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
12669    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
12670    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
12671    pub fn matmul_decode_exact_pre(
12672        &self,
12673        w: &crate::model::GpuTensor,
12674        aq: &CudaSlice<i8>,
12675        ad: &CudaSlice<f32>,
12676        m: usize,
12677    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12678        use crate::model::GpuTensor;
12679        debug_assert!(
12680            self.uses_q8_1_fast(w),
12681            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12682        );
12683        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
12684        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12685            return Ok(y);
12686        }
12687        let in_f = w.in_features();
12688        let out_f = w.out_features();
12689        let (bytes, qtype, row_bytes, scale, rp) = match w {
12690            GpuTensor::Quant {
12691                bytes,
12692                qtype,
12693                row_bytes,
12694                scale,
12695                rp,
12696                ..
12697            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12698            _ => {
12699                return Err(
12700                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12701                );
12702            }
12703        };
12704        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
12705        let (bytes, rp) = match w {
12706            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12707            _ => (bytes, rp),
12708        };
12709        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
12710        if (2..=16).contains(&m)
12711            && self.batched_supports(qtype)
12712            && self.mmvq_supports(qtype)
12713            && std::env::var("MEMRA_NO_BATCHED").is_err()
12714            && (m <= 4 || Self::b8_enabled())
12715            && (m <= 8
12716                || qtype == QT_Q4_0
12717                || qtype == QT_Q6_K
12718                || qtype == QT_F8_E4M3
12719                || qtype == QT_NVFP4
12720                || qtype == QT_Q4_K
12721                || qtype == QT_Q5_K
12722                || qtype == QT_Q8_0)
12723        {
12724            let mcols = Self::batched_mcols(m);
12725            return self.qmatvec_mmvq_batched(
12726                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12727            );
12728        }
12729        if self.mmvq_supports(qtype) {
12730            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12731        }
12732        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
12733        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
12734        let x0 = self.zeros(0)?;
12735        self.matmul_pre(w, aq, ad, &x0, m)
12736    }
12737
12738    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
12739    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
12740    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
12741    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
12742    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
12743    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
12744    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
12745    /// per-tensor path.
12746    pub fn matmul_decode_exact_dual_pre(
12747        &self,
12748        w0: &crate::model::GpuTensor,
12749        w1: &crate::model::GpuTensor,
12750        aq: &CudaSlice<i8>,
12751        ad: &CudaSlice<f32>,
12752        m: usize,
12753    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12754    {
12755        use crate::model::GpuTensor;
12756        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12757        let on = *ON.get_or_init(|| {
12758            std::env::var("MEMRA_SPEC_DUAL_T")
12759                .map(|v| v != "0")
12760                .unwrap_or(true)
12761        });
12762        if !on
12763            || !(2..=7).contains(&m)
12764            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12765            || !self.uses_q8_1_fast(w0)
12766            || !self.uses_q8_1_fast(w1)
12767        {
12768            return Ok(None);
12769        }
12770        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
12771        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
12772        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
12773        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
12774        if !self.mmvq_supports(QT_NVFP4) {
12775            return Ok(None);
12776        }
12777        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12778        if w1.in_features() != in_f || w1.out_features() != out_f {
12779            return Ok(None);
12780        }
12781        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12782            (
12783                GpuTensor::Quant {
12784                    bytes: b0,
12785                    qtype: q0,
12786                    row_bytes: rb0,
12787                    scale: s0,
12788                    rp: rp0,
12789                    rp4: None,
12790                    ..
12791                },
12792                GpuTensor::Quant {
12793                    bytes: b1,
12794                    qtype: q1,
12795                    row_bytes: rb1,
12796                    scale: s1,
12797                    rp: rp1,
12798                    rp4: None,
12799                    ..
12800                },
12801            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12802                (b0, b1, *rb0, *s0, *s1, *rp0)
12803            }
12804            _ => return Ok(None),
12805        };
12806        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
12807        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
12808        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12809        {
12810            return Ok(None);
12811        }
12812        let (y0, y1) =
12813            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12814        Ok(Some(((y0, s0), (y1, s1))))
12815    }
12816
12817    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
12818    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
12819    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
12820    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
12821    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
12822    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
12823    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
12824    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
12825    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
12826    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
12827    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
12828    pub fn matmul_decode_exact_group4_pre(
12829        &self,
12830        ws: [&crate::model::GpuTensor; 4],
12831        aq: &CudaSlice<i8>,
12832        ad: &CudaSlice<f32>,
12833        m: usize,
12834    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12835        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12836        let on = *ON.get_or_init(|| {
12837            std::env::var("MEMRA_TK_GDN_GROUP")
12838                .map(|v| v != "0")
12839                .unwrap_or(true)
12840        });
12841        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12842    }
12843
12844    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
12845    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
12846    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
12847    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
12848    pub fn matmul_decode_exact_group3_pre(
12849        &self,
12850        ws: [&crate::model::GpuTensor; 3],
12851        aq: &CudaSlice<i8>,
12852        ad: &CudaSlice<f32>,
12853        m: usize,
12854    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12855        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12856        let on = *ON.get_or_init(|| {
12857            std::env::var("MEMRA_TK_FA_GROUP")
12858                .map(|v| v != "0")
12859                .unwrap_or(true)
12860        });
12861        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12862    }
12863
12864    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
12865    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
12866    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
12867    fn matmul_decode_exact_group_pre(
12868        &self,
12869        ws: &[&crate::model::GpuTensor],
12870        aq: &CudaSlice<i8>,
12871        ad: &CudaSlice<f32>,
12872        m: usize,
12873        on: bool,
12874        tag: &'static str,
12875    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12876        use crate::model::GpuTensor;
12877        if !on
12878            || !(2..=16).contains(&m)
12879            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12880            || (m > 4 && !Self::b8_enabled())
12881            || !self.mmvq_supports(QT_NVFP4)
12882            || !self.batched_supports(QT_NVFP4)
12883        {
12884            return Ok(None);
12885        }
12886        let in_f = ws[0].in_features();
12887        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12888        for w in ws {
12889            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12890                return Ok(None);
12891            }
12892            match w {
12893                GpuTensor::Quant {
12894                    bytes,
12895                    qtype,
12896                    scale,
12897                    rp: true,
12898                    rp4: None,
12899                    ..
12900                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12901                    parts.push((bytes, w.out_features(), *scale));
12902                }
12903                _ => return Ok(None),
12904            }
12905        }
12906        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12907        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12908        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12909        let mcols = if (5..=7).contains(&m) && b567 {
12910            m
12911        } else {
12912            Self::batched_mcols(m)
12913        };
12914        let kname: &'static str = match mcols {
12915            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12916            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12917            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12918            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12919            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12920            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12921            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12922            _ => return Ok(None),
12923        };
12924        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12925        // the second door's print on the slice-D battery — key the once-set by tag.
12926        if std::env::var("MEMRA_DEBUG").is_ok() {
12927            use std::sync::Mutex;
12928            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12929            let mut seen = SEEN.lock().unwrap();
12930            if !seen.contains(&tag) {
12931                seen.push(tag);
12932                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12933            }
12934        }
12935        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12936        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12937        let total: usize = parts.iter().map(|p| p.1).sum();
12938        let three = parts.len() == 3;
12939        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12940        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12941        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12942        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12943        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12944        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12945        let cfg = LaunchConfig {
12946            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12947            block_dim: (32, ROWS_PER_BLOCK, 1),
12948            shared_mem_bytes: 0,
12949        };
12950        let (inf, mi) = (in_f as i32, m as i32);
12951        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12952        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12953        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12954        let s3 = if three { 1.0f32 } else { parts[3].2 };
12955        let w3 = if three { parts[0].0 } else { parts[3].0 };
12956        let f = self.func(kname);
12957        let __s_b = self.gpu.stream();
12958        let mut b = __s_b.launch_builder(&f);
12959        b.arg(parts[0].0)
12960            .arg(parts[1].0)
12961            .arg(parts[2].0)
12962            .arg(w3)
12963            .arg(aq)
12964            .arg(ad)
12965            .arg(&mut y0)
12966            .arg(&mut y1)
12967            .arg(&mut y2)
12968            .arg(&mut y3)
12969            .arg(&inf)
12970            .arg(&n0)
12971            .arg(&n1)
12972            .arg(&n2)
12973            .arg(&n3)
12974            .arg(&mi)
12975            .arg(&s0)
12976            .arg(&s1)
12977            .arg(&s2)
12978            .arg(&s3);
12979        unsafe {
12980            b.launch(cfg)?;
12981        }
12982        Ok(Some(if three {
12983            vec![y0, y1, y2]
12984        } else {
12985            vec![y0, y1, y2, y3]
12986        }))
12987    }
12988
12989    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12990    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12991    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12992    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12993    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12994    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12995    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12996    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12997    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12998    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12999    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
13000    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
13001    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
13002    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
13003    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
13004    pub fn matmul_decode_exact_dual(
13005        &self,
13006        w0: &crate::model::GpuTensor,
13007        w1: &crate::model::GpuTensor,
13008        x: &CudaSlice<f32>,
13009        m: usize,
13010    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13011        use crate::model::GpuTensor;
13012        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13013        let on = *ON.get_or_init(|| {
13014            std::env::var("MEMRA_SPEC_DUAL_T")
13015                .map(|v| v != "0")
13016                .unwrap_or(true)
13017        });
13018        if !on
13019            || !(2..=4).contains(&m)
13020            || std::env::var("MEMRA_NO_BATCHED").is_ok()
13021            || !self.uses_q8_1_fast(w0)
13022            || !self.uses_q8_1_fast(w1)
13023        {
13024            return Ok(None);
13025        }
13026        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
13027        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
13028        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
13029        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
13030        if !self.mmvq_supports(QT_NVFP4) {
13031            return Ok(None);
13032        }
13033        let (in_f, out_f) = (w0.in_features(), w0.out_features());
13034        if w1.in_features() != in_f || w1.out_features() != out_f {
13035            return Ok(None);
13036        }
13037        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
13038            (
13039                GpuTensor::Quant {
13040                    bytes: b0,
13041                    qtype: q0,
13042                    row_bytes: rb0,
13043                    scale: s0,
13044                    rp: rp0,
13045                    rp4: None,
13046                    ..
13047                },
13048                GpuTensor::Quant {
13049                    bytes: b1,
13050                    qtype: q1,
13051                    row_bytes: rb1,
13052                    scale: s1,
13053                    rp: rp1,
13054                    rp4: None,
13055                    ..
13056                },
13057            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13058                (b0, b1, *rb0, *s0, *s1, *rp0)
13059            }
13060            _ => return Ok(None),
13061        };
13062        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
13063        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
13064        if std::env::var("MEMRA_DEBUG").is_ok() {
13065            static ONCE: std::sync::Once = std::sync::Once::new();
13066            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13067        }
13068        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13069        let (y0, y1) =
13070            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13071        let mut y0 = y0;
13072        let mut y1 = y1;
13073        if s0 != 1.0 {
13074            self.scale_inplace(&mut y0, s0, m * out_f)?;
13075        }
13076        if s1 != 1.0 {
13077            self.scale_inplace(&mut y1, s1, m * out_f)?;
13078        }
13079        Ok(Some((y0, y1)))
13080    }
13081
13082    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
13083    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
13084    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
13085    /// twins (both buffers must be the repacked layout).
13086    #[allow(clippy::too_many_arguments)]
13087    pub fn qmatvec_batched_dual_raw(
13088        &self,
13089        b0: &CudaSlice<u8>,
13090        b1: &CudaSlice<u8>,
13091        aq: &CudaSlice<i8>,
13092        ad: &CudaSlice<f32>,
13093        m: usize,
13094        in_f: usize,
13095        out_f: usize,
13096        row_bytes: usize,
13097        rp: bool,
13098    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13099        const ROWS_PER_BLOCK: u32 = 4;
13100        let mcols = Self::batched_mcols(m);
13101        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
13102        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
13103        let tiny_rp1 = rp
13104            && mcols == 4
13105            && out_f <= 128
13106            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13107        let (name, rows_per_block) = if tiny_rp1 {
13108            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13109        } else {
13110            match (mcols, rp, m) {
13111                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13112                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13113                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13114                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13115                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13116                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13117                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13118                _ => {
13119                    return Err(
13120                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13121                    );
13122                }
13123            }
13124        };
13125        let f = self.func(name);
13126        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13127        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13128        let cfg = LaunchConfig {
13129            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13130            block_dim: (32, ROWS_PER_BLOCK, 1),
13131            shared_mem_bytes: 0,
13132        };
13133        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13134        let __s_b = self.gpu.stream();
13135        let mut b = __s_b.launch_builder(&f);
13136        b.arg(b0)
13137            .arg(b1)
13138            .arg(aq)
13139            .arg(ad)
13140            .arg(&mut y0)
13141            .arg(&mut y1)
13142            .arg(&inf)
13143            .arg(&outf)
13144            .arg(&mi)
13145            .arg(&rb);
13146        unsafe {
13147            b.launch(cfg)?;
13148        }
13149        Ok((y0, y1))
13150    }
13151
13152    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
13153    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
13154    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
13155    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
13156    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
13157    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
13158    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
13159    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
13160    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
13161    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
13162    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
13163    pub fn matmul_pre_dual_noscale(
13164        &self,
13165        w0: &crate::model::GpuTensor,
13166        w1: &crate::model::GpuTensor,
13167        aq: &CudaSlice<i8>,
13168        ad: &CudaSlice<f32>,
13169        m: usize,
13170    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13171    {
13172        use crate::model::GpuTensor;
13173        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13174            return Ok(None);
13175        }
13176        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
13177        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
13178        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
13179        // would mix dispatch families across the pair — the exact class `q8_fused_params`
13180        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
13181        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
13182        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
13183        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
13184        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
13185        if !self.mmvq_supports(QT_NVFP4) {
13186            return Ok(None);
13187        }
13188        let (in_f, out_f) = (w0.in_features(), w0.out_features());
13189        if w1.in_features() != in_f || w1.out_features() != out_f {
13190            return Ok(None);
13191        }
13192        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
13193        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
13194        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
13195        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
13196        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
13197        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
13198        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
13199        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
13200        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
13201        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
13202        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
13203        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
13204        let no_mirror =
13205            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13206        if self.q8_ffn_fuse2_on()
13207            && no_mirror(w0)
13208            && no_mirror(w1)
13209            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13210        {
13211            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13212            return Ok(Some(((y0, 1.0), (y1, 1.0))));
13213        }
13214        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
13215        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
13216        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
13217        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
13218        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
13219        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
13220        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
13221        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
13222        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
13223        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13224            let (y0, y1) =
13225                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13226            return Ok(Some(((y0, p0.3), (y1, p1.3))));
13227        }
13228        let (b0, q0, rb0, s0, rp0) = match w0 {
13229            GpuTensor::Quant {
13230                bytes,
13231                qtype,
13232                row_bytes,
13233                scale,
13234                rp,
13235                ..
13236            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13237            _ => return Ok(None),
13238        };
13239        let (b1, q1, rb1, s1, rp1) = match w1 {
13240            GpuTensor::Quant {
13241                bytes,
13242                qtype,
13243                row_bytes,
13244                scale,
13245                rp,
13246                ..
13247            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13248            _ => return Ok(None),
13249        };
13250        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13251            return Ok(None);
13252        }
13253        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13254        const RPW: u32 = 2;
13255        let rows_per_block = ROWS_PER_BLOCK * RPW;
13256        let f = self.func(if rp0 {
13257            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13258        } else {
13259            "qmatvec_nvfp4_mmvq_dual_mr2"
13260        });
13261        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13262        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13263        let cfg = LaunchConfig {
13264            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13265            block_dim: (32, ROWS_PER_BLOCK, 1),
13266            shared_mem_bytes: 0,
13267        };
13268        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13269        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
13270        // yscale args stay 1.0 here (they exist for the single-tensor callers).
13271        let one = 1.0f32;
13272        let __s_b = self.gpu.stream();
13273        let mut b = __s_b.launch_builder(&f);
13274        b.arg(b0)
13275            .arg(b1)
13276            .arg(aq)
13277            .arg(ad)
13278            .arg(&mut y0)
13279            .arg(&mut y1)
13280            .arg(&inf)
13281            .arg(&outf)
13282            .arg(&mi)
13283            .arg(&rb)
13284            .arg(&one)
13285            .arg(&one);
13286        unsafe {
13287            b.launch(cfg)?;
13288        }
13289        Ok(Some(((y0, s0), (y1, s1))))
13290    }
13291
13292    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
13293    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
13294    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
13295    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
13296    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
13297    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
13298    /// back to the three singles.
13299    #[allow(clippy::too_many_arguments)]
13300    pub fn matmul_nvfp4_fused3(
13301        &self,
13302        w0: &crate::model::GpuTensor,
13303        w1: &crate::model::GpuTensor,
13304        w2: &crate::model::GpuTensor,
13305        aq: &CudaSlice<i8>,
13306        ad: &CudaSlice<f32>,
13307        m: usize,
13308    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13309    {
13310        use crate::model::GpuTensor;
13311        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13312        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
13313        // verbatim, weight rows read once for all m columns, bit-identical per
13314        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
13315        // segments would re-read the weight per row" note described the grid.y=m lift,
13316        // which this twin deliberately is NOT.
13317        if !self.mmvq_supports(QT_NVFP4)
13318            || !self.uses_q8_1_fast(w0)
13319            || !self.uses_q8_1_fast(w1)
13320            || !self.uses_q8_1_fast(w2)
13321        {
13322            return Ok(None);
13323        }
13324        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
13325        // door — same family and bit-identity law as the fused4 delegate above.
13326        if (9..=16).contains(&m) {
13327            return Ok(
13328                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13329                    Some(mut ys) => {
13330                        let y2 = ys.pop().unwrap();
13331                        let y1 = ys.pop().unwrap();
13332                        let y0 = ys.pop().unwrap();
13333                        Some((y0, y1, y2))
13334                    }
13335                    None => None,
13336                },
13337            );
13338        }
13339        if !(1..=8).contains(&m) {
13340            return Ok(None);
13341        }
13342        if m > 1 {
13343            let in_f = w0.in_features();
13344            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13345                || !self.batched_supports(QT_NVFP4)
13346                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13347                || (m > 4 && !Self::b8_enabled())
13348                || in_f % 512 != 0
13349                || in_f / 64 > 272
13350            {
13351                return Ok(None);
13352            }
13353        }
13354        let unpack = |w: &crate::model::GpuTensor| match w {
13355            GpuTensor::Quant {
13356                bytes,
13357                qtype,
13358                scale,
13359                rp,
13360                ..
13361            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13362            _ => None,
13363        };
13364        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13365            return Ok(None);
13366        };
13367        let in_f = w0.in_features();
13368        if w1.in_features() != in_f || w2.in_features() != in_f {
13369            return Ok(None);
13370        }
13371        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13372        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13373        const RPW: u32 = 2;
13374        let rows_pb = ROWS_PER_BLOCK * RPW;
13375        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13376        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13377        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13378        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13379        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13380        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13381        // only dereferenced for the launch-arg build inside this call.
13382        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13383        if m > 1 {
13384            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
13385            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13386                return Ok(None);
13387            }
13388            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13389            let cfg = LaunchConfig {
13390                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13391                block_dim: (32, ROWS_PER_BLOCK, 1),
13392                shared_mem_bytes: 0,
13393            };
13394            let __s_b = self.gpu.stream();
13395            let mut b = __s_b.launch_builder(&f);
13396            b.arg(b0)
13397                .arg(b1)
13398                .arg(b2)
13399                .arg(aq)
13400                .arg(ad)
13401                .arg(&mut y0)
13402                .arg(&mut y1)
13403                .arg(&mut y2)
13404                .arg(&inf)
13405                .arg(&oi0)
13406                .arg(&oi1)
13407                .arg(&oi2)
13408                .arg(&mi);
13409            unsafe {
13410                b.launch(cfg)?;
13411            }
13412            return Ok(Some((y0, y1, y2)));
13413        }
13414        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13415        let cfg = LaunchConfig {
13416            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13417            block_dim: (32, ROWS_PER_BLOCK, 1),
13418            shared_mem_bytes: 0,
13419        };
13420        let __s_b = self.gpu.stream();
13421        let mut b = __s_b.launch_builder(&f);
13422        b.arg(b0)
13423            .arg(b1)
13424            .arg(b2)
13425            .arg(aq)
13426            .arg(ad)
13427            .arg(&mut y0)
13428            .arg(&mut y1)
13429            .arg(&mut y2)
13430            .arg(&inf)
13431            .arg(&oi0)
13432            .arg(&oi1)
13433            .arg(&oi2)
13434            .arg(&mi)
13435            .arg(&p0.1)
13436            .arg(&p1.1)
13437            .arg(&p2.1);
13438        unsafe {
13439            b.launch(cfg)?;
13440        }
13441        Ok(Some((y0, y1, y2)))
13442    }
13443
13444    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
13445    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
13446    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
13447    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
13448    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
13449    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
13450    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
13451    /// same-binary interleaved A/B arm.
13452    pub fn matmul_nvfp4_fused2(
13453        &self,
13454        w0: &crate::model::GpuTensor,
13455        w1: &crate::model::GpuTensor,
13456        aq: &CudaSlice<i8>,
13457        ad: &CudaSlice<f32>,
13458        m: usize,
13459    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13460        use crate::model::GpuTensor;
13461        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13462        let off =
13463            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13464        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
13465        // read serves all m rows); the fused segments would re-read the weight per row.
13466        if off
13467            || m != 1
13468            || !self.mmvq_supports(QT_NVFP4)
13469            || !self.uses_q8_1_fast(w0)
13470            || !self.uses_q8_1_fast(w1)
13471        {
13472            return Ok(None);
13473        }
13474        let unpack = |w: &crate::model::GpuTensor| match w {
13475            GpuTensor::Quant {
13476                bytes,
13477                qtype,
13478                scale,
13479                rp,
13480                ..
13481            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13482            _ => None,
13483        };
13484        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13485            return Ok(None);
13486        };
13487        let in_f = w0.in_features();
13488        if w1.in_features() != in_f {
13489            return Ok(None);
13490        }
13491        let (o0, o1) = (w0.out_features(), w1.out_features());
13492        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13493        const RPW: u32 = 2;
13494        let rows_pb = ROWS_PER_BLOCK * RPW;
13495        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13496        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13497        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13498        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13499        let cfg = LaunchConfig {
13500            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13501            block_dim: (32, ROWS_PER_BLOCK, 1),
13502            shared_mem_bytes: 0,
13503        };
13504        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13505        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13506        // only dereferenced for the launch-arg build inside this call.
13507        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13508        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
13509        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
13510        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13511            {
13512                use cudarc::driver::{DevicePtr, DevicePtrMut};
13513                let s = &self.gpu.stream();
13514                let (pw0, _g0) = b0.device_ptr(s);
13515                let (pw1, _g1) = b1.device_ptr(s);
13516                let (paq, _g2) = aq.device_ptr(s);
13517                let (pad, _g3) = ad.device_ptr(s);
13518                let (py0, _g4) = y0.device_ptr_mut(s);
13519                let (py1, _g5) = y1.device_ptr_mut(s);
13520                let (s0, s1) = (p0.1, p1.1);
13521                let mut ps = [
13522                    &pw0 as *const _ as *mut std::ffi::c_void,
13523                    &pw1 as *const _ as *mut _,
13524                    &paq as *const _ as *mut _,
13525                    &pad as *const _ as *mut _,
13526                    &py0 as *const _ as *mut _,
13527                    &py1 as *const _ as *mut _,
13528                    &inf as *const _ as *mut _,
13529                    &oi0 as *const _ as *mut _,
13530                    &oi1 as *const _ as *mut _,
13531                    &mi as *const _ as *mut _,
13532                    &s0 as *const _ as *mut _,
13533                    &s1 as *const _ as *mut _,
13534                ];
13535                unsafe {
13536                    self.launch_pdl(
13537                        "qmatvec_nvfp4_mmvq_fused2_rp",
13538                        cfg.grid_dim,
13539                        cfg.block_dim,
13540                        &mut ps,
13541                    )?;
13542                }
13543            }
13544            return Ok(Some((y0, y1)));
13545        }
13546        let __s_b = self.gpu.stream();
13547        let mut b = __s_b.launch_builder(&f);
13548        b.arg(b0)
13549            .arg(b1)
13550            .arg(aq)
13551            .arg(ad)
13552            .arg(&mut y0)
13553            .arg(&mut y1)
13554            .arg(&inf)
13555            .arg(&oi0)
13556            .arg(&oi1)
13557            .arg(&mi)
13558            .arg(&p0.1)
13559            .arg(&p1.1);
13560        unsafe {
13561            b.launch(cfg)?;
13562        }
13563        Ok(Some((y0, y1)))
13564    }
13565
13566    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
13567    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
13568    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
13569    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
13570    pub fn matmul_nvfp4_fused2_into(
13571        &self,
13572        w0: &crate::model::GpuTensor,
13573        w1: &crate::model::GpuTensor,
13574        aq: &CudaSlice<i8>,
13575        ad: &CudaSlice<f32>,
13576        y0: &mut CudaSlice<f32>,
13577        y1: &mut CudaSlice<f32>,
13578    ) -> Result<bool, Box<dyn std::error::Error>> {
13579        use crate::model::GpuTensor;
13580        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13581        let off =
13582            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13583        if off
13584            || !self.mmvq_supports(QT_NVFP4)
13585            || !self.uses_q8_1_fast(w0)
13586            || !self.uses_q8_1_fast(w1)
13587        {
13588            return Ok(false);
13589        }
13590        let unpack = |w: &crate::model::GpuTensor| match w {
13591            GpuTensor::Quant {
13592                bytes,
13593                qtype,
13594                scale,
13595                rp,
13596                ..
13597            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13598            _ => None,
13599        };
13600        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13601            return Ok(false);
13602        };
13603        let in_f = w0.in_features();
13604        if w1.in_features() != in_f {
13605            return Ok(false);
13606        }
13607        let (o0, o1) = (w0.out_features(), w1.out_features());
13608        if y0.len() < o0 || y1.len() < o1 {
13609            return Ok(false);
13610        }
13611        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13612        const RPW: u32 = 2;
13613        let rows_pb = ROWS_PER_BLOCK * RPW;
13614        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13615        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13616        let cfg = LaunchConfig {
13617            grid_dim: (nb(o0) + nb(o1), 1, 1),
13618            block_dim: (32, ROWS_PER_BLOCK, 1),
13619            shared_mem_bytes: 0,
13620        };
13621        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13622        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13623        // only dereferenced for the launch-arg build inside this call.
13624        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13625        let __s_b = self.gpu.stream();
13626        let mut b = __s_b.launch_builder(&f);
13627        b.arg(b0)
13628            .arg(b1)
13629            .arg(aq)
13630            .arg(ad)
13631            .arg(&mut *y0)
13632            .arg(&mut *y1)
13633            .arg(&inf)
13634            .arg(&oi0)
13635            .arg(&oi1)
13636            .arg(&mi)
13637            .arg(&p0.1)
13638            .arg(&p1.1);
13639        unsafe {
13640            b.launch(cfg)?;
13641        }
13642        Ok(true)
13643    }
13644
13645    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
13646    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
13647    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
13648    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
13649    #[allow(clippy::type_complexity)]
13650    pub fn matmul_nvfp4_fused4(
13651        &self,
13652        w0: &crate::model::GpuTensor,
13653        w1: &crate::model::GpuTensor,
13654        w2: &crate::model::GpuTensor,
13655        w3: &crate::model::GpuTensor,
13656        aq: &CudaSlice<i8>,
13657        ad: &CudaSlice<f32>,
13658        m: usize,
13659    ) -> Result<
13660        Option<(
13661            CudaSlice<f32>,
13662            CudaSlice<f32>,
13663            CudaSlice<f32>,
13664            CudaSlice<f32>,
13665        )>,
13666        Box<dyn std::error::Error>,
13667    > {
13668        use crate::model::GpuTensor;
13669        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
13670        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13671        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
13672        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
13673        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
13674        // Admission mirrors the singles' batched gates below.
13675        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13676            || !self.mmvq_supports(QT_NVFP4)
13677            || !self.uses_q8_1_fast(w0)
13678            || !self.uses_q8_1_fast(w1)
13679            || !self.uses_q8_1_fast(w2)
13680            || !self.uses_q8_1_fast(w3)
13681        {
13682            return Ok(None);
13683        }
13684        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
13685        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
13686        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
13687        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
13688        if (9..=16).contains(&m) {
13689            return Ok(
13690                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13691                    Some(mut ys) => {
13692                        let y3 = ys.pop().unwrap();
13693                        let y2 = ys.pop().unwrap();
13694                        let y1 = ys.pop().unwrap();
13695                        let y0 = ys.pop().unwrap();
13696                        Some((y0, y1, y2, y3))
13697                    }
13698                    None => None,
13699                },
13700            );
13701        }
13702        if !(1..=8).contains(&m) {
13703            return Ok(None);
13704        }
13705        if m > 1 {
13706            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
13707            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
13708            let in_f = w0.in_features();
13709            if !self.batched_supports(QT_NVFP4)
13710                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13711                || (m > 4 && !Self::b8_enabled())
13712                || in_f % 512 != 0
13713                || in_f / 64 > 272
13714            {
13715                return Ok(None);
13716            }
13717        }
13718        let unpack = |w: &crate::model::GpuTensor| match w {
13719            GpuTensor::Quant {
13720                bytes,
13721                qtype,
13722                scale,
13723                rp,
13724                ..
13725            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13726            _ => None,
13727        };
13728        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13729            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13730        else {
13731            return Ok(None);
13732        };
13733        let in_f = w0.in_features();
13734        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13735            return Ok(None);
13736        }
13737        let (o0, o1, o2, o3) = (
13738            w0.out_features(),
13739            w1.out_features(),
13740            w2.out_features(),
13741            w3.out_features(),
13742        );
13743        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13744        const RPW: u32 = 2;
13745        let rows_pb = ROWS_PER_BLOCK * RPW;
13746        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13747        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13748        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13749        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13750        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13751        let (inf, oi0, oi1, oi2, oi3, mi) = (
13752            in_f as i32,
13753            o0 as i32,
13754            o1 as i32,
13755            o2 as i32,
13756            o3 as i32,
13757            m as i32,
13758        );
13759        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13760        // only dereferenced for the launch-arg build inside this call.
13761        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13762        if m > 1 {
13763            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
13764            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
13765            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13766                return Ok(None);
13767            }
13768            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13769            let cfg = LaunchConfig {
13770                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13771                block_dim: (32, ROWS_PER_BLOCK, 1),
13772                shared_mem_bytes: 0,
13773            };
13774            let __s_b = self.gpu.stream();
13775            let mut b = __s_b.launch_builder(&f);
13776            b.arg(b0)
13777                .arg(b1)
13778                .arg(b2)
13779                .arg(b3)
13780                .arg(aq)
13781                .arg(ad)
13782                .arg(&mut y0)
13783                .arg(&mut y1)
13784                .arg(&mut y2)
13785                .arg(&mut y3)
13786                .arg(&inf)
13787                .arg(&oi0)
13788                .arg(&oi1)
13789                .arg(&oi2)
13790                .arg(&oi3)
13791                .arg(&mi);
13792            unsafe {
13793                b.launch(cfg)?;
13794            }
13795            return Ok(Some((y0, y1, y2, y3)));
13796        }
13797        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13798        let cfg = LaunchConfig {
13799            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13800            block_dim: (32, ROWS_PER_BLOCK, 1),
13801            shared_mem_bytes: 0,
13802        };
13803        let __s_b = self.gpu.stream();
13804        let mut b = __s_b.launch_builder(&f);
13805        b.arg(b0)
13806            .arg(b1)
13807            .arg(b2)
13808            .arg(b3)
13809            .arg(aq)
13810            .arg(ad)
13811            .arg(&mut y0)
13812            .arg(&mut y1)
13813            .arg(&mut y2)
13814            .arg(&mut y3)
13815            .arg(&inf)
13816            .arg(&oi0)
13817            .arg(&oi1)
13818            .arg(&oi2)
13819            .arg(&oi3)
13820            .arg(&mi)
13821            .arg(&p0.1)
13822            .arg(&p1.1)
13823            .arg(&p2.1)
13824            .arg(&p3.1);
13825        unsafe {
13826            b.launch(cfg)?;
13827        }
13828        Ok(Some((y0, y1, y2, y3)))
13829    }
13830
13831    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
13832    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
13833    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
13834    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
13835    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
13836    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
13837    /// back to the per-tensor path.
13838    pub fn matmul_q8_fused2(
13839        &self,
13840        w0: &crate::model::GpuTensor,
13841        w1: &crate::model::GpuTensor,
13842        aq: &CudaSlice<i8>,
13843        ad: &CudaSlice<f32>,
13844    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13845        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
13846        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
13847        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
13848        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
13849        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
13850        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13851            return Ok(Some(self.e4m3_fused2_core(
13852                p0.0,
13853                p1.0,
13854                aq,
13855                ad,
13856                w0.in_features(),
13857                p0.1,
13858                p1.1,
13859                p0.2,
13860                p0.3,
13861                p1.3,
13862            )?));
13863        }
13864        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13865            return Ok(None);
13866        };
13867        Ok(Some(self.q8_fused2_core(
13868            p0.0,
13869            p1.0,
13870            aq,
13871            ad,
13872            w0.in_features(),
13873            p0.1,
13874            p1.1,
13875            p0.2,
13876        )?))
13877    }
13878
13879    #[allow(clippy::too_many_arguments)]
13880    fn q8_fused2_core(
13881        &self,
13882        b0: &CudaSlice<u8>,
13883        b1: &CudaSlice<u8>,
13884        aq: &CudaSlice<i8>,
13885        ad: &CudaSlice<f32>,
13886        in_f: usize,
13887        out0: usize,
13888        out1: usize,
13889        row_bytes: usize,
13890    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13891        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13892        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13893        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13894        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13895        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13896        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13897        let cfg = LaunchConfig {
13898            grid_dim: (nb0 + nb1, 1, 1),
13899            block_dim: (32, ROWS_PER_BLOCK, 1),
13900            shared_mem_bytes: 0,
13901        };
13902        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13903        let __s_b = self.gpu.stream();
13904        let mut b = __s_b.launch_builder(&f);
13905        b.arg(b0)
13906            .arg(b1)
13907            .arg(aq)
13908            .arg(ad)
13909            .arg(&mut y0)
13910            .arg(&mut y1)
13911            .arg(&inf)
13912            .arg(&o0)
13913            .arg(&o1)
13914            .arg(&rbl);
13915        unsafe {
13916            b.launch(cfg)?;
13917        }
13918        Ok((y0, y1))
13919    }
13920
13921    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13922    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13923    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13924    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13925    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13926    pub fn matmul_q8_fused2_x(
13927        &self,
13928        w0: &crate::model::GpuTensor,
13929        w1: &crate::model::GpuTensor,
13930        x: &CudaSlice<f32>,
13931    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13932        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13933            return Ok(None);
13934        }
13935        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13936            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13937            return Ok(Some(self.e4m3_fused2_core(
13938                p0.0,
13939                p1.0,
13940                &aq,
13941                &ad,
13942                w0.in_features(),
13943                p0.1,
13944                p1.1,
13945                p0.2,
13946                p0.3,
13947                p1.3,
13948            )?));
13949        }
13950        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13951            return Ok(None);
13952        };
13953        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13954        Ok(Some(self.q8_fused2_core(
13955            p0.0,
13956            p1.0,
13957            &aq,
13958            &ad,
13959            w0.in_features(),
13960            p0.1,
13961            p1.1,
13962            p0.2,
13963        )?))
13964    }
13965
13966    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13967    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13968    #[allow(clippy::too_many_arguments)]
13969    pub fn qmatvec_q8_fused2_raw(
13970        &self,
13971        b0: &CudaSlice<u8>,
13972        b1: &CudaSlice<u8>,
13973        x: &CudaSlice<f32>,
13974        in_f: usize,
13975        out0: usize,
13976        out1: usize,
13977        row_bytes: usize,
13978    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13979        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13980        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13981    }
13982
13983    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13984    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13985    /// (tensor,row) to three separate m=1 MMVQ launches.
13986    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13987    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13988    pub fn matmul_q4_fused3(
13989        &self,
13990        w0: &crate::model::GpuTensor,
13991        w1: &crate::model::GpuTensor,
13992        w2: &crate::model::GpuTensor,
13993        aq: &CudaSlice<i8>,
13994        ad: &CudaSlice<f32>,
13995    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13996    {
13997        use crate::model::GpuTensor;
13998        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13999            match w {
14000                GpuTensor::Quant {
14001                    qtype, row_bytes, ..
14002                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14003                _ => None,
14004            }
14005        };
14006        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14007            return Ok(None);
14008        };
14009        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14010            return Ok(None);
14011        }
14012        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
14013        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
14014        // the separate matvecs (each routes its own rp).
14015        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14016            match w {
14017                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14018                    Some(m) => (m, true),
14019                    None => (bytes, *rp),
14020                },
14021                _ => unreachable!(),
14022            }
14023        }
14024        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14025        if rp0 != rp1 || rp1 != rp2 {
14026            return Ok(None);
14027        }
14028        let rp = rp0;
14029        let rpb: u32 = 4;
14030        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
14031        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
14032        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
14033        let mr1 = rp && Self::q40_mr1_on();
14034        let nb = |o: usize| {
14035            if mr1 {
14036                (o as u32).div_ceil(rpb)
14037            } else {
14038                (o as u32).div_ceil(2).div_ceil(rpb)
14039            }
14040        };
14041        let grid = nb(o0) + nb(o1) + nb(o2);
14042        let mut y0 = self.alloc_uninit::<f32>(o0)?;
14043        let mut y1 = self.alloc_uninit::<f32>(o1)?;
14044        let mut y2 = self.alloc_uninit::<f32>(o2)?;
14045        let f = self.func(if mr1 {
14046            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14047        } else if rp {
14048            "qmatvec_q4_0_mmvq_fused3_rp"
14049        } else {
14050            "qmatvec_q4_0_mmvq_fused3"
14051        });
14052        let cfg = LaunchConfig {
14053            grid_dim: (grid, 1, 1),
14054            block_dim: (32, rpb, 1),
14055            shared_mem_bytes: 0,
14056        };
14057        let inf = w0.in_features() as i32;
14058        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14059        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14060        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
14061        // variant may take the programmatic-serialization launch.
14062        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14063            {
14064                use cudarc::driver::{DevicePtr, DevicePtrMut};
14065                let s = &self.gpu.stream();
14066                let (p0, _g0) = b0.device_ptr(s);
14067                let (p1, _g1) = b1.device_ptr(s);
14068                let (p2, _g2) = b2.device_ptr(s);
14069                let (paq, _g3) = aq.device_ptr(s);
14070                let (pad, _g4) = ad.device_ptr(s);
14071                let (py0, _g5) = y0.device_ptr_mut(s);
14072                let (py1, _g6) = y1.device_ptr_mut(s);
14073                let (py2, _g7) = y2.device_ptr_mut(s);
14074                let mut ps = [
14075                    &p0 as *const _ as *mut std::ffi::c_void,
14076                    &p1 as *const _ as *mut _,
14077                    &p2 as *const _ as *mut _,
14078                    &paq as *const _ as *mut _,
14079                    &pad as *const _ as *mut _,
14080                    &py0 as *const _ as *mut _,
14081                    &py1 as *const _ as *mut _,
14082                    &py2 as *const _ as *mut _,
14083                    &inf as *const _ as *mut _,
14084                    &oo0 as *const _ as *mut _,
14085                    &oo1 as *const _ as *mut _,
14086                    &oo2 as *const _ as *mut _,
14087                    &r0 as *const _ as *mut _,
14088                    &r1 as *const _ as *mut _,
14089                    &r2 as *const _ as *mut _,
14090                ];
14091                unsafe {
14092                    self.launch_pdl(
14093                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14094                        (grid, 1, 1),
14095                        (32, rpb, 1),
14096                        &mut ps,
14097                    )?;
14098                }
14099            }
14100            return Ok(Some((y0, y1, y2)));
14101        }
14102        let __s_b = self.gpu.stream();
14103        let mut b = __s_b.launch_builder(&f);
14104        b.arg(b0)
14105            .arg(b1)
14106            .arg(b2)
14107            .arg(aq)
14108            .arg(ad)
14109            .arg(&mut y0)
14110            .arg(&mut y1)
14111            .arg(&mut y2)
14112            .arg(&inf)
14113            .arg(&oo0)
14114            .arg(&oo1)
14115            .arg(&oo2)
14116            .arg(&r0)
14117            .arg(&r1)
14118            .arg(&r2);
14119        unsafe {
14120            b.launch(cfg)?;
14121        }
14122        Ok(Some((y0, y1, y2)))
14123    }
14124
14125    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
14126    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
14127    #[allow(clippy::too_many_arguments)]
14128    pub fn matmul_q4_fused3_into(
14129        &self,
14130        w0: &crate::model::GpuTensor,
14131        w1: &crate::model::GpuTensor,
14132        w2: &crate::model::GpuTensor,
14133        aq: &CudaSlice<i8>,
14134        ad: &CudaSlice<f32>,
14135        y0: &mut CudaSlice<f32>,
14136        y1: &mut CudaSlice<f32>,
14137        y2: &mut CudaSlice<f32>,
14138    ) -> Result<bool, Box<dyn std::error::Error>> {
14139        use crate::model::GpuTensor;
14140        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14141            match w {
14142                GpuTensor::Quant {
14143                    qtype, row_bytes, ..
14144                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14145                _ => None,
14146            }
14147        };
14148        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14149            return Ok(false);
14150        };
14151        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14152            return Ok(false);
14153        }
14154        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14155            match w {
14156                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14157                    Some(m) => (m, true),
14158                    None => (bytes, *rp),
14159                },
14160                _ => unreachable!(),
14161            }
14162        }
14163        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14164        if rp0 != rp1 || rp1 != rp2 {
14165            return Ok(false);
14166        }
14167        let rp = rp0;
14168        let rpb: u32 = 4;
14169        let mr1 = rp && Self::q40_mr1_on();
14170        let nb = |o: usize| {
14171            if mr1 {
14172                (o as u32).div_ceil(rpb)
14173            } else {
14174                (o as u32).div_ceil(2).div_ceil(rpb)
14175            }
14176        };
14177        let grid = nb(o0) + nb(o1) + nb(o2);
14178        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14179        let f = self.func(if mr1 {
14180            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14181        } else if rp {
14182            "qmatvec_q4_0_mmvq_fused3_rp"
14183        } else {
14184            "qmatvec_q4_0_mmvq_fused3"
14185        });
14186        let cfg = LaunchConfig {
14187            grid_dim: (grid, 1, 1),
14188            block_dim: (32, rpb, 1),
14189            shared_mem_bytes: 0,
14190        };
14191        let inf = w0.in_features() as i32;
14192        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14193        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14194        // PDL wave-A: identical to the owned twin (capture-lane parity).
14195        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14196            use cudarc::driver::{DevicePtr, DevicePtrMut};
14197            let s = &self.gpu.stream();
14198            let (p0, _g0) = b0.device_ptr(s);
14199            let (p1, _g1) = b1.device_ptr(s);
14200            let (p2, _g2) = b2.device_ptr(s);
14201            let (paq, _g3) = aq.device_ptr(s);
14202            let (pad, _g4) = ad.device_ptr(s);
14203            let (py0, _g5) = y0.device_ptr_mut(s);
14204            let (py1, _g6) = y1.device_ptr_mut(s);
14205            let (py2, _g7) = y2.device_ptr_mut(s);
14206            let mut ps = [
14207                &p0 as *const _ as *mut std::ffi::c_void,
14208                &p1 as *const _ as *mut _,
14209                &p2 as *const _ as *mut _,
14210                &paq as *const _ as *mut _,
14211                &pad as *const _ as *mut _,
14212                &py0 as *const _ as *mut _,
14213                &py1 as *const _ as *mut _,
14214                &py2 as *const _ as *mut _,
14215                &inf as *const _ as *mut _,
14216                &oo0 as *const _ as *mut _,
14217                &oo1 as *const _ as *mut _,
14218                &oo2 as *const _ as *mut _,
14219                &r0 as *const _ as *mut _,
14220                &r1 as *const _ as *mut _,
14221                &r2 as *const _ as *mut _,
14222            ];
14223            unsafe {
14224                self.launch_pdl(
14225                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14226                    (grid, 1, 1),
14227                    (32, rpb, 1),
14228                    &mut ps,
14229                )?;
14230            }
14231            return Ok(true);
14232        }
14233        let __s_b = self.gpu.stream();
14234        let mut b = __s_b.launch_builder(&f);
14235        b.arg(b0)
14236            .arg(b1)
14237            .arg(b2)
14238            .arg(aq)
14239            .arg(ad)
14240            .arg(&mut *y0)
14241            .arg(&mut *y1)
14242            .arg(&mut *y2)
14243            .arg(&inf)
14244            .arg(&oo0)
14245            .arg(&oo1)
14246            .arg(&oo2)
14247            .arg(&r0)
14248            .arg(&r1)
14249            .arg(&r2);
14250        unsafe {
14251            b.launch(cfg)?;
14252        }
14253        Ok(true)
14254    }
14255
14256    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
14257    pub fn matmul_q4_fused2(
14258        &self,
14259        w0: &crate::model::GpuTensor,
14260        w1: &crate::model::GpuTensor,
14261        aq: &CudaSlice<i8>,
14262        ad: &CudaSlice<f32>,
14263    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14264        use crate::model::GpuTensor;
14265        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14266            match w {
14267                GpuTensor::Quant {
14268                    qtype, row_bytes, ..
14269                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14270                _ => None,
14271            }
14272        };
14273        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14274            return Ok(None);
14275        };
14276        if w0.in_features() != w1.in_features() {
14277            return Ok(None);
14278        }
14279        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
14280        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14281            match w {
14282                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14283                    Some(m) => (m, 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 rp = rp0;
14294        let rpb: u32 = 4;
14295        // mr1 twin — see matmul_q4_fused3.
14296        let mr1 = rp && Self::q40_mr1_on();
14297        let nb = |o: usize| {
14298            if mr1 {
14299                (o as u32).div_ceil(rpb)
14300            } else {
14301                (o as u32).div_ceil(2).div_ceil(rpb)
14302            }
14303        };
14304        let grid = nb(o0) + nb(o1);
14305        let mut y0 = self.alloc_uninit::<f32>(o0)?;
14306        let mut y1 = self.alloc_uninit::<f32>(o1)?;
14307        let f = self.func(if mr1 {
14308            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14309        } else if rp {
14310            "qmatvec_q4_0_mmvq_fused2_rp"
14311        } else {
14312            "qmatvec_q4_0_mmvq_fused2"
14313        });
14314        let cfg = LaunchConfig {
14315            grid_dim: (grid, 1, 1),
14316            block_dim: (32, rpb, 1),
14317            shared_mem_bytes: 0,
14318        };
14319        let inf = w0.in_features() as i32;
14320        let (oo0, oo1) = (o0 as i32, o1 as i32);
14321        let (r0, r1) = (rb0 as i64, rb1 as i64);
14322        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
14323        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14324            {
14325                use cudarc::driver::{DevicePtr, DevicePtrMut};
14326                let s = &self.gpu.stream();
14327                let (p0, _g0) = b0.device_ptr(s);
14328                let (p1, _g1) = b1.device_ptr(s);
14329                let (paq, _g2) = aq.device_ptr(s);
14330                let (pad, _g3) = ad.device_ptr(s);
14331                let (py0, _g4) = y0.device_ptr_mut(s);
14332                let (py1, _g5) = y1.device_ptr_mut(s);
14333                let mut ps = [
14334                    &p0 as *const _ as *mut std::ffi::c_void,
14335                    &p1 as *const _ as *mut _,
14336                    &paq as *const _ as *mut _,
14337                    &pad as *const _ as *mut _,
14338                    &py0 as *const _ as *mut _,
14339                    &py1 as *const _ as *mut _,
14340                    &inf as *const _ as *mut _,
14341                    &oo0 as *const _ as *mut _,
14342                    &oo1 as *const _ as *mut _,
14343                    &r0 as *const _ as *mut _,
14344                    &r1 as *const _ as *mut _,
14345                ];
14346                unsafe {
14347                    self.launch_pdl(
14348                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14349                        (grid, 1, 1),
14350                        (32, rpb, 1),
14351                        &mut ps,
14352                    )?;
14353                }
14354            }
14355            return Ok(Some((y0, y1)));
14356        }
14357        let __s_b = self.gpu.stream();
14358        let mut b = __s_b.launch_builder(&f);
14359        b.arg(b0)
14360            .arg(b1)
14361            .arg(aq)
14362            .arg(ad)
14363            .arg(&mut y0)
14364            .arg(&mut y1)
14365            .arg(&inf)
14366            .arg(&oo0)
14367            .arg(&oo1)
14368            .arg(&r0)
14369            .arg(&r1);
14370        unsafe {
14371            b.launch(cfg)?;
14372        }
14373        Ok(Some((y0, y1)))
14374    }
14375
14376    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
14377    pub fn matmul_q4_fused2_into(
14378        &self,
14379        w0: &crate::model::GpuTensor,
14380        w1: &crate::model::GpuTensor,
14381        aq: &CudaSlice<i8>,
14382        ad: &CudaSlice<f32>,
14383        y0: &mut CudaSlice<f32>,
14384        y1: &mut CudaSlice<f32>,
14385    ) -> Result<bool, Box<dyn std::error::Error>> {
14386        use crate::model::GpuTensor;
14387        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14388            match w {
14389                GpuTensor::Quant {
14390                    qtype, row_bytes, ..
14391                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14392                _ => None,
14393            }
14394        };
14395        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14396            return Ok(false);
14397        };
14398        if w0.in_features() != w1.in_features() {
14399            return Ok(false);
14400        }
14401        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14402            match w {
14403                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14404                    Some(m) => (m, true),
14405                    None => (bytes, *rp),
14406                },
14407                _ => unreachable!(),
14408            }
14409        }
14410        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14411        if rp0 != rp1 {
14412            return Ok(false);
14413        }
14414        let rp = rp0;
14415        let rpb: u32 = 4;
14416        let mr1 = rp && Self::q40_mr1_on();
14417        let nb = |o: usize| {
14418            if mr1 {
14419                (o as u32).div_ceil(rpb)
14420            } else {
14421                (o as u32).div_ceil(2).div_ceil(rpb)
14422            }
14423        };
14424        let grid = nb(o0) + nb(o1);
14425        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14426        let f = self.func(if mr1 {
14427            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14428        } else if rp {
14429            "qmatvec_q4_0_mmvq_fused2_rp"
14430        } else {
14431            "qmatvec_q4_0_mmvq_fused2"
14432        });
14433        let cfg = LaunchConfig {
14434            grid_dim: (grid, 1, 1),
14435            block_dim: (32, rpb, 1),
14436            shared_mem_bytes: 0,
14437        };
14438        let inf = w0.in_features() as i32;
14439        let (oo0, oo1) = (o0 as i32, o1 as i32);
14440        let (r0, r1) = (rb0 as i64, rb1 as i64);
14441        // PDL wave-A: identical to the owned twin (capture-lane parity).
14442        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14443            use cudarc::driver::{DevicePtr, DevicePtrMut};
14444            let s = &self.gpu.stream();
14445            let (p0, _g0) = b0.device_ptr(s);
14446            let (p1, _g1) = b1.device_ptr(s);
14447            let (paq, _g2) = aq.device_ptr(s);
14448            let (pad, _g3) = ad.device_ptr(s);
14449            let (py0, _g4) = y0.device_ptr_mut(s);
14450            let (py1, _g5) = y1.device_ptr_mut(s);
14451            let mut ps = [
14452                &p0 as *const _ as *mut std::ffi::c_void,
14453                &p1 as *const _ as *mut _,
14454                &paq as *const _ as *mut _,
14455                &pad as *const _ as *mut _,
14456                &py0 as *const _ as *mut _,
14457                &py1 as *const _ as *mut _,
14458                &inf as *const _ as *mut _,
14459                &oo0 as *const _ as *mut _,
14460                &oo1 as *const _ as *mut _,
14461                &r0 as *const _ as *mut _,
14462                &r1 as *const _ as *mut _,
14463            ];
14464            unsafe {
14465                self.launch_pdl(
14466                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14467                    (grid, 1, 1),
14468                    (32, rpb, 1),
14469                    &mut ps,
14470                )?;
14471            }
14472            return Ok(true);
14473        }
14474        let __s_b = self.gpu.stream();
14475        let mut b = __s_b.launch_builder(&f);
14476        b.arg(b0)
14477            .arg(b1)
14478            .arg(aq)
14479            .arg(ad)
14480            .arg(&mut *y0)
14481            .arg(&mut *y1)
14482            .arg(&inf)
14483            .arg(&oo0)
14484            .arg(&oo1)
14485            .arg(&r0)
14486            .arg(&r1);
14487        unsafe {
14488            b.launch(cfg)?;
14489        }
14490        Ok(true)
14491    }
14492
14493    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
14494    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
14495    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
14496    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
14497    pub fn matmul_q4_fused2_batched(
14498        &self,
14499        w0: &crate::model::GpuTensor,
14500        w1: &crate::model::GpuTensor,
14501        aq: &CudaSlice<i8>,
14502        ad: &CudaSlice<f32>,
14503        m: usize,
14504    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14505        use crate::model::GpuTensor;
14506        if m < 2 || m > 8 {
14507            return Ok(None);
14508        }
14509        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14510            match w {
14511                GpuTensor::Quant {
14512                    qtype, row_bytes, ..
14513                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14514                _ => None,
14515            }
14516        };
14517        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14518            return Ok(None);
14519        };
14520        if w0.in_features() != w1.in_features() {
14521            return Ok(None);
14522        }
14523        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14524            match w {
14525                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14526                    Some(mr) => (mr, true),
14527                    None => (bytes, *rp),
14528                },
14529                _ => unreachable!(),
14530            }
14531        }
14532        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14533        if !rp0 || !rp1 {
14534            return Ok(None);
14535        }
14536        let mcols = Self::batched_mcols(m);
14537        let rpb: u32 = 4;
14538        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14539        let grid = nb(o0) + nb(o1);
14540        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14541        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14542        let f = self.func(match mcols {
14543            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14544            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14545            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14546        });
14547        let cfg = LaunchConfig {
14548            grid_dim: (grid, 1, 1),
14549            block_dim: (32, rpb, 1),
14550            shared_mem_bytes: 0,
14551        };
14552        let inf = w0.in_features() as i32;
14553        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14554        let rb = rb0 as i64;
14555        let __s_b = self.gpu.stream();
14556        let mut b = __s_b.launch_builder(&f);
14557        b.arg(b0)
14558            .arg(b1)
14559            .arg(aq)
14560            .arg(ad)
14561            .arg(&mut y0)
14562            .arg(&mut y1)
14563            .arg(&inf)
14564            .arg(&oo0)
14565            .arg(&oo1)
14566            .arg(&mi)
14567            .arg(&rb);
14568        unsafe {
14569            b.launch(cfg)?;
14570        }
14571        Ok(Some((y0, y1)))
14572    }
14573
14574    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
14575    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
14576    #[allow(clippy::too_many_arguments)]
14577    pub fn matmul_q4_fused3_batched(
14578        &self,
14579        w0: &crate::model::GpuTensor,
14580        w1: &crate::model::GpuTensor,
14581        w2: &crate::model::GpuTensor,
14582        aq: &CudaSlice<i8>,
14583        ad: &CudaSlice<f32>,
14584        m: usize,
14585    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14586    {
14587        use crate::model::GpuTensor;
14588        if m < 2 || m > 8 {
14589            return Ok(None);
14590        }
14591        let q4 = |w: &GpuTensor| -> Option<usize> {
14592            match w {
14593                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14594                _ => None,
14595            }
14596        };
14597        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14598            return Ok(None);
14599        };
14600        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14601            return Ok(None);
14602        }
14603        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14604            match w {
14605                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14606                    Some(mr) => (mr, true),
14607                    None => (bytes, *rp),
14608                },
14609                _ => unreachable!(),
14610            }
14611        }
14612        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14613        if !rp0 || !rp1 || !rp2 {
14614            return Ok(None);
14615        }
14616        let mcols = Self::batched_mcols(m);
14617        let rpb: u32 = 4;
14618        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14619        let grid = nb(o0) + nb(o1) + nb(o2);
14620        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14621        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14622        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14623        let f = self.func(match mcols {
14624            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14625            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14626            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14627        });
14628        let cfg = LaunchConfig {
14629            grid_dim: (grid, 1, 1),
14630            block_dim: (32, rpb, 1),
14631            shared_mem_bytes: 0,
14632        };
14633        let inf = w0.in_features() as i32;
14634        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14635        let rb = 0i64;
14636        let __s_b = self.gpu.stream();
14637        let mut b = __s_b.launch_builder(&f);
14638        b.arg(b0)
14639            .arg(b1)
14640            .arg(b2)
14641            .arg(aq)
14642            .arg(ad)
14643            .arg(&mut y0)
14644            .arg(&mut y1)
14645            .arg(&mut y2)
14646            .arg(&inf)
14647            .arg(&oo0)
14648            .arg(&oo1)
14649            .arg(&oo2)
14650            .arg(&mi)
14651            .arg(&rb);
14652        unsafe {
14653            b.launch(cfg)?;
14654        }
14655        Ok(Some((y0, y1, y2)))
14656    }
14657
14658    pub fn matmul_q8_fused3(
14659        &self,
14660        w0: &crate::model::GpuTensor,
14661        w1: &crate::model::GpuTensor,
14662        w2: &crate::model::GpuTensor,
14663        aq: &CudaSlice<i8>,
14664        ad: &CudaSlice<f32>,
14665    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14666    {
14667        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
14668        // are per-tensor FP8, so native residency without this arm meant three separate launches.
14669        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14670            return Ok(Some(self.e4m3_fused3_core(
14671                p0.0,
14672                p1.0,
14673                p2.0,
14674                aq,
14675                ad,
14676                w0.in_features(),
14677                p0.1,
14678                p1.1,
14679                p2.1,
14680                p0.2,
14681                p0.3,
14682                p1.3,
14683                p2.3,
14684            )?));
14685        }
14686        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14687            return Ok(None);
14688        };
14689        Ok(Some(self.q8_fused3_core(
14690            p0.0,
14691            p1.0,
14692            p2.0,
14693            aq,
14694            ad,
14695            w0.in_features(),
14696            p0.1,
14697            p1.1,
14698            p2.1,
14699            p0.2,
14700        )?))
14701    }
14702
14703    #[allow(clippy::too_many_arguments)]
14704    fn q8_fused3_core(
14705        &self,
14706        b0: &CudaSlice<u8>,
14707        b1: &CudaSlice<u8>,
14708        b2: &CudaSlice<u8>,
14709        aq: &CudaSlice<i8>,
14710        ad: &CudaSlice<f32>,
14711        in_f: usize,
14712        out0: usize,
14713        out1: usize,
14714        out2: usize,
14715        row_bytes: usize,
14716    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14717        const ROWS_PER_BLOCK: u32 = 4;
14718        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14719        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14720        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14721        let f = self.func("qmatvec_q8_0_mmvq_fused3");
14722        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14723        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14724        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14725        let cfg = LaunchConfig {
14726            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14727            block_dim: (32, ROWS_PER_BLOCK, 1),
14728            shared_mem_bytes: 0,
14729        };
14730        let (inf, o0, o1, o2, rbl) = (
14731            in_f as i32,
14732            out0 as i32,
14733            out1 as i32,
14734            out2 as i32,
14735            row_bytes as i64,
14736        );
14737        let __s_b = self.gpu.stream();
14738        let mut b = __s_b.launch_builder(&f);
14739        b.arg(b0)
14740            .arg(b1)
14741            .arg(b2)
14742            .arg(aq)
14743            .arg(ad)
14744            .arg(&mut y0)
14745            .arg(&mut y1)
14746            .arg(&mut y2)
14747            .arg(&inf)
14748            .arg(&o0)
14749            .arg(&o1)
14750            .arg(&o2)
14751            .arg(&rbl);
14752        unsafe {
14753            b.launch(cfg)?;
14754        }
14755        Ok((y0, y1, y2))
14756    }
14757
14758    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
14759    #[allow(clippy::too_many_arguments)]
14760    pub fn qmatvec_q8_fused3_raw(
14761        &self,
14762        b0: &CudaSlice<u8>,
14763        b1: &CudaSlice<u8>,
14764        b2: &CudaSlice<u8>,
14765        x: &CudaSlice<f32>,
14766        in_f: usize,
14767        out0: usize,
14768        out1: usize,
14769        out2: usize,
14770        row_bytes: usize,
14771    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14772        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14773        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14774    }
14775
14776    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
14777    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
14778    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
14779    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
14780    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
14781    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
14782    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
14783    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
14784    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
14785    /// twin must not introduce a batched program the reference path would not run).
14786    pub fn matmul_q8_fused2_t(
14787        &self,
14788        w0: &crate::model::GpuTensor,
14789        w1: &crate::model::GpuTensor,
14790        aq: &CudaSlice<i8>,
14791        ad: &CudaSlice<f32>,
14792        m: usize,
14793    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14794        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
14795        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
14796        // fuses too — same template body, still bit-identical to the two _b8 launches.
14797        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14798            return Ok(None);
14799        }
14800        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
14801        // so the fused b8 launch would introduce a batched program the reference path would not run.
14802        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14803            if m > 4 && !Self::b8_enabled() {
14804                return Ok(None);
14805            }
14806            return Ok(Some(self.e4m3_fused2_t_core(
14807                p0.0,
14808                p1.0,
14809                aq,
14810                ad,
14811                m,
14812                w0.in_features(),
14813                p0.1,
14814                p1.1,
14815                p0.2,
14816                p0.3,
14817                p1.3,
14818            )?));
14819        }
14820        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14821            return Ok(None);
14822        };
14823        Ok(Some(self.q8_fused2_t_core(
14824            p0.0,
14825            p1.0,
14826            aq,
14827            ad,
14828            m,
14829            w0.in_features(),
14830            p0.1,
14831            p1.1,
14832            p0.2,
14833        )?))
14834    }
14835
14836    #[allow(clippy::too_many_arguments)]
14837    fn q8_fused2_t_core(
14838        &self,
14839        b0: &CudaSlice<u8>,
14840        b1: &CudaSlice<u8>,
14841        aq: &CudaSlice<i8>,
14842        ad: &CudaSlice<f32>,
14843        m: usize,
14844        in_f: usize,
14845        out0: usize,
14846        out1: usize,
14847        row_bytes: usize,
14848    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14849        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14850        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14851        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14852        let f = self.func(match Self::batched_mcols(m) {
14853            2 => "qmatvec_q8_0_mmvq_fused2_b2",
14854            4 => "qmatvec_q8_0_mmvq_fused2_b4",
14855            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
14856            _ => "qmatvec_q8_0_mmvq_fused2_b8",
14857        });
14858        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14859        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14860        let cfg = LaunchConfig {
14861            grid_dim: (nb0 + nb1, 1, 1),
14862            block_dim: (32, ROWS_PER_BLOCK, 1),
14863            shared_mem_bytes: 0,
14864        };
14865        let (inf, o0, o1, mi, rbl) = (
14866            in_f as i32,
14867            out0 as i32,
14868            out1 as i32,
14869            m as i32,
14870            row_bytes as i64,
14871        );
14872        let __s_b = self.gpu.stream();
14873        let mut b = __s_b.launch_builder(&f);
14874        b.arg(b0)
14875            .arg(b1)
14876            .arg(aq)
14877            .arg(ad)
14878            .arg(&mut y0)
14879            .arg(&mut y1)
14880            .arg(&inf)
14881            .arg(&o0)
14882            .arg(&o1)
14883            .arg(&mi)
14884            .arg(&rbl);
14885        unsafe {
14886            b.launch(cfg)?;
14887        }
14888        Ok((y0, y1))
14889    }
14890
14891    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14892    /// q8_1 quant of the [m, in_f] activation), no env gating.
14893    #[allow(clippy::too_many_arguments)]
14894    pub fn qmatvec_q8_fused2_t_raw(
14895        &self,
14896        b0: &CudaSlice<u8>,
14897        b1: &CudaSlice<u8>,
14898        x: &CudaSlice<f32>,
14899        m: usize,
14900        in_f: usize,
14901        out0: usize,
14902        out1: usize,
14903        row_bytes: usize,
14904    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14905        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14906        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14907    }
14908
14909    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14910    /// `matmul_q8_fused2_t` with three ranges.
14911    #[allow(clippy::too_many_arguments)]
14912    pub fn matmul_q8_fused3_t(
14913        &self,
14914        w0: &crate::model::GpuTensor,
14915        w1: &crate::model::GpuTensor,
14916        w2: &crate::model::GpuTensor,
14917        aq: &CudaSlice<i8>,
14918        ad: &CudaSlice<f32>,
14919        m: usize,
14920    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14921    {
14922        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14923            return Ok(None);
14924        }
14925        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14926            return Ok(Some(self.e4m3_fused3_t_core(
14927                p0.0,
14928                p1.0,
14929                p2.0,
14930                aq,
14931                ad,
14932                m,
14933                w0.in_features(),
14934                p0.1,
14935                p1.1,
14936                p2.1,
14937                p0.2,
14938                p0.3,
14939                p1.3,
14940                p2.3,
14941            )?));
14942        }
14943        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14944            return Ok(None);
14945        };
14946        Ok(Some(self.q8_fused3_t_core(
14947            p0.0,
14948            p1.0,
14949            p2.0,
14950            aq,
14951            ad,
14952            m,
14953            w0.in_features(),
14954            p0.1,
14955            p1.1,
14956            p2.1,
14957            p0.2,
14958        )?))
14959    }
14960
14961    #[allow(clippy::too_many_arguments)]
14962    fn q8_fused3_t_core(
14963        &self,
14964        b0: &CudaSlice<u8>,
14965        b1: &CudaSlice<u8>,
14966        b2: &CudaSlice<u8>,
14967        aq: &CudaSlice<i8>,
14968        ad: &CudaSlice<f32>,
14969        m: usize,
14970        in_f: usize,
14971        out0: usize,
14972        out1: usize,
14973        out2: usize,
14974        row_bytes: usize,
14975    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14976        const ROWS_PER_BLOCK: u32 = 4;
14977        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14978        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14979        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14980        let f = self.func(if Self::batched_mcols(m) == 2 {
14981            "qmatvec_q8_0_mmvq_fused3_b2"
14982        } else {
14983            "qmatvec_q8_0_mmvq_fused3_b4"
14984        });
14985        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14986        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14987        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14988        let cfg = LaunchConfig {
14989            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14990            block_dim: (32, ROWS_PER_BLOCK, 1),
14991            shared_mem_bytes: 0,
14992        };
14993        let (inf, o0, o1, o2, mi, rbl) = (
14994            in_f as i32,
14995            out0 as i32,
14996            out1 as i32,
14997            out2 as i32,
14998            m as i32,
14999            row_bytes as i64,
15000        );
15001        let __s_b = self.gpu.stream();
15002        let mut b = __s_b.launch_builder(&f);
15003        b.arg(b0)
15004            .arg(b1)
15005            .arg(b2)
15006            .arg(aq)
15007            .arg(ad)
15008            .arg(&mut y0)
15009            .arg(&mut y1)
15010            .arg(&mut y2)
15011            .arg(&inf)
15012            .arg(&o0)
15013            .arg(&o1)
15014            .arg(&o2)
15015            .arg(&mi)
15016            .arg(&rbl);
15017        unsafe {
15018            b.launch(cfg)?;
15019        }
15020        Ok((y0, y1, y2))
15021    }
15022
15023    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
15024    #[allow(clippy::too_many_arguments)]
15025    pub fn qmatvec_q8_fused3_t_raw(
15026        &self,
15027        b0: &CudaSlice<u8>,
15028        b1: &CudaSlice<u8>,
15029        b2: &CudaSlice<u8>,
15030        x: &CudaSlice<f32>,
15031        m: usize,
15032        in_f: usize,
15033        out0: usize,
15034        out1: usize,
15035        out2: usize,
15036        row_bytes: usize,
15037    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15038        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15039        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
15040    }
15041
15042    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
15043    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
15044    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
15045    pub fn q8_ffn_fuse2_on(&self) -> bool {
15046        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15047        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15048    }
15049
15050    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
15051    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
15052    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
15053    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
15054    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
15055    #[allow(clippy::type_complexity)]
15056    fn q8_fused_params<'w, const N: usize>(
15057        &self,
15058        ws: &[&'w crate::model::GpuTensor; N],
15059    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15060        use crate::model::GpuTensor;
15061        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15062            return None;
15063        }
15064        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15065            return None;
15066        }
15067        let in_f = ws[0].in_features();
15068        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15069        for (i, w) in ws.iter().enumerate() {
15070            match w {
15071                GpuTensor::Quant {
15072                    bytes,
15073                    qtype,
15074                    row_bytes,
15075                    scale,
15076                    ..
15077                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15078                    out[i] = Some((bytes, w.out_features(), *row_bytes))
15079                }
15080                _ => return None,
15081            }
15082        }
15083        Some(out.map(|o| o.unwrap()))
15084    }
15085
15086    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
15087    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
15088    pub fn e4m3_dual_on(&self) -> bool {
15089        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15090        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15091    }
15092
15093    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
15094    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
15095    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
15096    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
15097    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
15098    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
15099    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
15100    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
15101    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
15102    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
15103    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
15104    #[allow(clippy::type_complexity)]
15105    fn e4m3_fused_params<'w, const N: usize>(
15106        &self,
15107        ws: &[&'w crate::model::GpuTensor; N],
15108    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15109        use crate::model::GpuTensor;
15110        if !self.e4m3_dual_on() {
15111            return None;
15112        }
15113        let in_f = ws[0].in_features();
15114        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15115        for (i, w) in ws.iter().enumerate() {
15116            match w {
15117                GpuTensor::Quant {
15118                    bytes,
15119                    qtype,
15120                    row_bytes,
15121                    scale,
15122                    rp,
15123                    rp4,
15124                    ..
15125                } if *qtype == QT_F8_E4M3
15126                    && w.in_features() == in_f
15127                    && *row_bytes == in_f
15128                    && !*rp
15129                    && rp4.is_none() =>
15130                {
15131                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15132                }
15133                _ => return None,
15134            }
15135        }
15136        Some(out.map(|o| o.unwrap()))
15137    }
15138
15139    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
15140    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
15141    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
15142    #[allow(clippy::too_many_arguments)]
15143    fn e4m3_fused2_core(
15144        &self,
15145        b0: &CudaSlice<u8>,
15146        b1: &CudaSlice<u8>,
15147        aq: &CudaSlice<i8>,
15148        ad: &CudaSlice<f32>,
15149        in_f: usize,
15150        out0: usize,
15151        out1: usize,
15152        row_bytes: usize,
15153        ws0: f32,
15154        ws1: f32,
15155    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15156        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15157        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15158        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15159        let f = self.func("qmatvec_e4m3_mmvq_fused2");
15160        let mut y0 = self.alloc_uninit::<f32>(out0)?;
15161        let mut y1 = self.alloc_uninit::<f32>(out1)?;
15162        let cfg = LaunchConfig {
15163            grid_dim: (nb0 + nb1, 1, 1),
15164            block_dim: (32, ROWS_PER_BLOCK, 1),
15165            shared_mem_bytes: 0,
15166        };
15167        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15168        let __s_b = self.gpu.stream();
15169        let mut b = __s_b.launch_builder(&f);
15170        b.arg(b0)
15171            .arg(b1)
15172            .arg(aq)
15173            .arg(ad)
15174            .arg(&mut y0)
15175            .arg(&mut y1)
15176            .arg(&inf)
15177            .arg(&o0)
15178            .arg(&o1)
15179            .arg(&rbl)
15180            .arg(&ws0)
15181            .arg(&ws1);
15182        unsafe {
15183            b.launch(cfg)?;
15184        }
15185        Ok((y0, y1))
15186    }
15187
15188    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
15189    #[allow(clippy::too_many_arguments)]
15190    fn e4m3_fused3_core(
15191        &self,
15192        b0: &CudaSlice<u8>,
15193        b1: &CudaSlice<u8>,
15194        b2: &CudaSlice<u8>,
15195        aq: &CudaSlice<i8>,
15196        ad: &CudaSlice<f32>,
15197        in_f: usize,
15198        out0: usize,
15199        out1: usize,
15200        out2: usize,
15201        row_bytes: usize,
15202        ws0: f32,
15203        ws1: f32,
15204        ws2: f32,
15205    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15206        const ROWS_PER_BLOCK: u32 = 4;
15207        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15208        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15209        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15210        let f = self.func("qmatvec_e4m3_mmvq_fused3");
15211        let mut y0 = self.alloc_uninit::<f32>(out0)?;
15212        let mut y1 = self.alloc_uninit::<f32>(out1)?;
15213        let mut y2 = self.alloc_uninit::<f32>(out2)?;
15214        let cfg = LaunchConfig {
15215            grid_dim: (nb0 + nb1 + nb2, 1, 1),
15216            block_dim: (32, ROWS_PER_BLOCK, 1),
15217            shared_mem_bytes: 0,
15218        };
15219        let (inf, o0, o1, o2, rbl) = (
15220            in_f as i32,
15221            out0 as i32,
15222            out1 as i32,
15223            out2 as i32,
15224            row_bytes as i64,
15225        );
15226        let __s_b = self.gpu.stream();
15227        let mut b = __s_b.launch_builder(&f);
15228        b.arg(b0)
15229            .arg(b1)
15230            .arg(b2)
15231            .arg(aq)
15232            .arg(ad)
15233            .arg(&mut y0)
15234            .arg(&mut y1)
15235            .arg(&mut y2)
15236            .arg(&inf)
15237            .arg(&o0)
15238            .arg(&o1)
15239            .arg(&o2)
15240            .arg(&rbl)
15241            .arg(&ws0)
15242            .arg(&ws1)
15243            .arg(&ws2);
15244        unsafe {
15245            b.launch(cfg)?;
15246        }
15247        Ok((y0, y1, y2))
15248    }
15249
15250    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
15251    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
15252    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
15253    #[allow(clippy::too_many_arguments)]
15254    fn e4m3_fused2_t_core(
15255        &self,
15256        b0: &CudaSlice<u8>,
15257        b1: &CudaSlice<u8>,
15258        aq: &CudaSlice<i8>,
15259        ad: &CudaSlice<f32>,
15260        m: usize,
15261        in_f: usize,
15262        out0: usize,
15263        out1: usize,
15264        row_bytes: usize,
15265        ws0: f32,
15266        ws1: f32,
15267    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15268        const ROWS_PER_BLOCK: u32 = 4;
15269        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15270        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15271        let f = self.func(match Self::batched_mcols(m) {
15272            2 => "qmatvec_e4m3_mmvq_fused2_b2",
15273            4 => "qmatvec_e4m3_mmvq_fused2_b4",
15274            _ => "qmatvec_e4m3_mmvq_fused2_b8",
15275        });
15276        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15277        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15278        let cfg = LaunchConfig {
15279            grid_dim: (nb0 + nb1, 1, 1),
15280            block_dim: (32, ROWS_PER_BLOCK, 1),
15281            shared_mem_bytes: 0,
15282        };
15283        let (inf, o0, o1, mi, rbl) = (
15284            in_f as i32,
15285            out0 as i32,
15286            out1 as i32,
15287            m as i32,
15288            row_bytes as i64,
15289        );
15290        let __s_b = self.gpu.stream();
15291        let mut b = __s_b.launch_builder(&f);
15292        b.arg(b0)
15293            .arg(b1)
15294            .arg(aq)
15295            .arg(ad)
15296            .arg(&mut y0)
15297            .arg(&mut y1)
15298            .arg(&inf)
15299            .arg(&o0)
15300            .arg(&o1)
15301            .arg(&mi)
15302            .arg(&rbl);
15303        unsafe {
15304            b.launch(cfg)?;
15305        }
15306        if ws0 != 1.0 {
15307            self.scale_inplace(&mut y0, ws0, m * out0)?;
15308        }
15309        if ws1 != 1.0 {
15310            self.scale_inplace(&mut y1, ws1, m * out1)?;
15311        }
15312        Ok((y0, y1))
15313    }
15314
15315    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
15316    #[allow(clippy::too_many_arguments)]
15317    fn e4m3_fused3_t_core(
15318        &self,
15319        b0: &CudaSlice<u8>,
15320        b1: &CudaSlice<u8>,
15321        b2: &CudaSlice<u8>,
15322        aq: &CudaSlice<i8>,
15323        ad: &CudaSlice<f32>,
15324        m: usize,
15325        in_f: usize,
15326        out0: usize,
15327        out1: usize,
15328        out2: usize,
15329        row_bytes: usize,
15330        ws0: f32,
15331        ws1: f32,
15332        ws2: f32,
15333    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15334        const ROWS_PER_BLOCK: u32 = 4;
15335        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15336        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15337        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15338        let f = self.func(if Self::batched_mcols(m) == 2 {
15339            "qmatvec_e4m3_mmvq_fused3_b2"
15340        } else {
15341            "qmatvec_e4m3_mmvq_fused3_b4"
15342        });
15343        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15344        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15345        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15346        let cfg = LaunchConfig {
15347            grid_dim: (nb0 + nb1 + nb2, 1, 1),
15348            block_dim: (32, ROWS_PER_BLOCK, 1),
15349            shared_mem_bytes: 0,
15350        };
15351        let (inf, o0, o1, o2, mi, rbl) = (
15352            in_f as i32,
15353            out0 as i32,
15354            out1 as i32,
15355            out2 as i32,
15356            m as i32,
15357            row_bytes as i64,
15358        );
15359        let __s_b = self.gpu.stream();
15360        let mut b = __s_b.launch_builder(&f);
15361        b.arg(b0)
15362            .arg(b1)
15363            .arg(b2)
15364            .arg(aq)
15365            .arg(ad)
15366            .arg(&mut y0)
15367            .arg(&mut y1)
15368            .arg(&mut y2)
15369            .arg(&inf)
15370            .arg(&o0)
15371            .arg(&o1)
15372            .arg(&o2)
15373            .arg(&mi)
15374            .arg(&rbl);
15375        unsafe {
15376            b.launch(cfg)?;
15377        }
15378        if ws0 != 1.0 {
15379            self.scale_inplace(&mut y0, ws0, m * out0)?;
15380        }
15381        if ws1 != 1.0 {
15382            self.scale_inplace(&mut y1, ws1, m * out1)?;
15383        }
15384        if ws2 != 1.0 {
15385            self.scale_inplace(&mut y2, ws2, m * out2)?;
15386        }
15387        Ok((y0, y1, y2))
15388    }
15389
15390    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
15391    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
15392    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
15393    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
15394    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
15395    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
15396    ///
15397    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
15398    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
15399    pub fn qmatvec_e4m3_blk_mmvq(
15400        &self,
15401        bytes: &CudaSlice<u8>,
15402        aq: &CudaSlice<i8>,
15403        ad: &CudaSlice<f32>,
15404        scales: &CudaSlice<f32>,
15405        m: usize,
15406        in_f: usize,
15407        out_f: usize,
15408        row_bytes: usize,
15409        scale_cols: usize,
15410    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15411        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
15412        self.qmatvec_e4m3_blk_mmvq_into(
15413            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15414        )?;
15415        Ok(y)
15416    }
15417
15418    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
15419    #[allow(clippy::too_many_arguments)]
15420    pub fn qmatvec_e4m3_blk_mmvq_into(
15421        &self,
15422        bytes: &CudaSlice<u8>,
15423        aq: &CudaSlice<i8>,
15424        ad: &CudaSlice<f32>,
15425        scales: &CudaSlice<f32>,
15426        m: usize,
15427        in_f: usize,
15428        out_f: usize,
15429        row_bytes: usize,
15430        scale_cols: usize,
15431        y: &mut CudaSlice<f32>,
15432    ) -> Result<(), Box<dyn std::error::Error>> {
15433        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15434        let f = self.func("qmatvec_e4m3_blk_mmvq");
15435        let cfg = LaunchConfig {
15436            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15437            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
15438            shared_mem_bytes: 0,                // warp-only reduce
15439        };
15440        let (inf, outf, mi, rb, sc) = (
15441            in_f as i32,
15442            out_f as i32,
15443            m as i32,
15444            row_bytes as i64,
15445            scale_cols as i32,
15446        );
15447        let __s_b = self.gpu.stream();
15448        let mut b = __s_b.launch_builder(&f);
15449        b.arg(bytes)
15450            .arg(aq)
15451            .arg(ad)
15452            .arg(scales)
15453            .arg(&mut *y)
15454            .arg(&inf)
15455            .arg(&outf)
15456            .arg(&mi)
15457            .arg(&rb)
15458            .arg(&sc);
15459        unsafe {
15460            b.launch(cfg)?;
15461        }
15462        Ok(())
15463    }
15464
15465    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
15466    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
15467    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
15468    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
15469    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
15470    #[allow(clippy::too_many_arguments)]
15471    pub fn qmatvec_e4m3_blk_mmvq_batched(
15472        &self,
15473        bytes: &CudaSlice<u8>,
15474        aq: &CudaSlice<i8>,
15475        ad: &CudaSlice<f32>,
15476        scales: &CudaSlice<f32>,
15477        m: usize,
15478        in_f: usize,
15479        out_f: usize,
15480        row_bytes: usize,
15481        scale_cols: usize,
15482        mcols: usize,
15483    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15484        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15485        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15486        let name = match mcols {
15487            2 => "qmatvec_e4m3_blk_mmvq_b2",
15488            4 => "qmatvec_e4m3_blk_mmvq_b4",
15489            8 => "qmatvec_e4m3_blk_mmvq_b8",
15490            16 => "qmatvec_e4m3_blk_mmvq_b16",
15491            _ => {
15492                return Err(
15493                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15494                );
15495            }
15496        };
15497        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15498        let f = self.func(name);
15499        let cfg = LaunchConfig {
15500            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15501            block_dim: (32, ROWS_PER_BLOCK, 1),
15502            shared_mem_bytes: 0,
15503        };
15504        let (inf, outf, mi, rb, sc) = (
15505            in_f as i32,
15506            out_f as i32,
15507            m as i32,
15508            row_bytes as i64,
15509            scale_cols as i32,
15510        );
15511        let __s_b = self.gpu.stream();
15512        let mut b = __s_b.launch_builder(&f);
15513        b.arg(bytes)
15514            .arg(aq)
15515            .arg(ad)
15516            .arg(scales)
15517            .arg(&mut y)
15518            .arg(&inf)
15519            .arg(&outf)
15520            .arg(&mi)
15521            .arg(&rb)
15522            .arg(&sc);
15523        unsafe {
15524            b.launch(cfg)?;
15525        }
15526        Ok(y)
15527    }
15528
15529    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
15530    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
15531    #[allow(clippy::too_many_arguments)]
15532    pub fn qmatvec_e4m3_blk_batched_raw(
15533        &self,
15534        bytes: &CudaSlice<u8>,
15535        x: &CudaSlice<f32>,
15536        scales: &CudaSlice<f32>,
15537        m: usize,
15538        in_f: usize,
15539        out_f: usize,
15540        row_bytes: usize,
15541        scale_cols: usize,
15542        mcols: usize,
15543    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15544        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15545        self.qmatvec_e4m3_blk_mmvq_batched(
15546            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15547        )
15548    }
15549
15550    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
15551    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
15552    #[allow(clippy::too_many_arguments)]
15553    pub fn qmatvec_e4m3_blk_mmvq_raw(
15554        &self,
15555        bytes: &CudaSlice<u8>,
15556        x: &CudaSlice<f32>,
15557        scales: &CudaSlice<f32>,
15558        m: usize,
15559        in_f: usize,
15560        out_f: usize,
15561        row_bytes: usize,
15562        scale_cols: usize,
15563    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15564        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15565        self.qmatvec_e4m3_blk_mmvq(
15566            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15567        )
15568    }
15569
15570    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
15571    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
15572    #[allow(clippy::too_many_arguments)]
15573    pub fn qmatvec_e4m3_fused2_raw(
15574        &self,
15575        b0: &CudaSlice<u8>,
15576        b1: &CudaSlice<u8>,
15577        x: &CudaSlice<f32>,
15578        in_f: usize,
15579        out0: usize,
15580        out1: usize,
15581        row_bytes: usize,
15582        ws0: f32,
15583        ws1: f32,
15584    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15585        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15586        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15587    }
15588
15589    #[allow(clippy::too_many_arguments)]
15590    pub fn qmatvec_e4m3_fused3_raw(
15591        &self,
15592        b0: &CudaSlice<u8>,
15593        b1: &CudaSlice<u8>,
15594        b2: &CudaSlice<u8>,
15595        x: &CudaSlice<f32>,
15596        in_f: usize,
15597        out0: usize,
15598        out1: usize,
15599        out2: usize,
15600        row_bytes: usize,
15601        ws0: f32,
15602        ws1: f32,
15603        ws2: f32,
15604    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15605        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15606        self.e4m3_fused3_core(
15607            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15608        )
15609    }
15610
15611    #[allow(clippy::too_many_arguments)]
15612    pub fn qmatvec_e4m3_fused2_t_raw(
15613        &self,
15614        b0: &CudaSlice<u8>,
15615        b1: &CudaSlice<u8>,
15616        x: &CudaSlice<f32>,
15617        m: usize,
15618        in_f: usize,
15619        out0: usize,
15620        out1: usize,
15621        row_bytes: usize,
15622        ws0: f32,
15623        ws1: f32,
15624    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15625        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15626        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15627    }
15628
15629    #[allow(clippy::too_many_arguments)]
15630    pub fn qmatvec_e4m3_fused3_t_raw(
15631        &self,
15632        b0: &CudaSlice<u8>,
15633        b1: &CudaSlice<u8>,
15634        b2: &CudaSlice<u8>,
15635        x: &CudaSlice<f32>,
15636        m: usize,
15637        in_f: usize,
15638        out0: usize,
15639        out1: usize,
15640        out2: usize,
15641        row_bytes: usize,
15642        ws0: f32,
15643        ws1: f32,
15644        ws2: f32,
15645    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15646        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15647        self.e4m3_fused3_t_core(
15648            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15649        )
15650    }
15651
15652    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
15653    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
15654    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
15655    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
15656    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
15657    ///
15658    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
15659    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
15660    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
15661    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
15662    fn try_e4m3_blk_pre(
15663        &self,
15664        w: &crate::model::GpuTensor,
15665        aq: &CudaSlice<i8>,
15666        ad: &CudaSlice<f32>,
15667        m: usize,
15668    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15669        use crate::model::GpuTensor;
15670        if let GpuTensor::Quant {
15671            bytes,
15672            qtype,
15673            row_bytes,
15674            blk: Some(g),
15675            ..
15676        } = w
15677        {
15678            if *qtype == QT_F8_E4M3_BLK {
15679                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
15680                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
15681                // below, so the decode-exactness contract is preserved at every width. Gated by
15682                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
15683                // one rollback door covers every dtype's batched tier.
15684                if (2..=16).contains(&m)
15685                    && std::env::var("MEMRA_NO_BATCHED").is_err()
15686                    && (m <= 4 || Self::b8_enabled())
15687                {
15688                    let mcols = Self::batched_mcols(m);
15689                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15690                        bytes,
15691                        aq,
15692                        ad,
15693                        &g.scales,
15694                        m,
15695                        w.in_features(),
15696                        w.out_features(),
15697                        *row_bytes,
15698                        g.cols,
15699                        mcols,
15700                    )?));
15701                }
15702                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15703                    bytes,
15704                    aq,
15705                    ad,
15706                    &g.scales,
15707                    m,
15708                    w.in_features(),
15709                    w.out_features(),
15710                    *row_bytes,
15711                    g.cols,
15712                )?));
15713            }
15714        }
15715        Ok(None)
15716    }
15717
15718    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
15719    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
15720    ///
15721    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
15722    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
15723    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
15724    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
15725    /// prefill keeps the floor's arithmetic and the floor's kernels.
15726    ///
15727    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
15728    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
15729    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
15730    /// (projection, prefill call) and frees immediately.
15731    ///
15732    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
15733    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
15734    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
15735    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
15736    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
15737    /// single-variable comparison instead of a two-variable one.
15738    ///
15739    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
15740    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
15741    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
15742    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
15743    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
15744    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
15745    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
15746    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
15747    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
15748    ///
15749    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
15750    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
15751    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
15752    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
15753    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
15754    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
15755    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
15756    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
15757    /// because v2's denominator had its slab already resident while this class's floor must build it
15758    /// every call; same tile, opposite sign, because the question changed.
15759    ///
15760    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
15761    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
15762    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
15763    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
15764    fn try_e4m3_blk_prefill(
15765        &self,
15766        w: &crate::model::GpuTensor,
15767        x: &CudaSlice<f32>,
15768        m: usize,
15769    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15770        use crate::model::GpuTensor;
15771        let GpuTensor::Quant {
15772            bytes,
15773            qtype,
15774            blk: Some(g),
15775            ..
15776        } = w
15777        else {
15778            return Ok(None);
15779        };
15780        if *qtype != QT_F8_E4M3_BLK {
15781            return Ok(None);
15782        }
15783        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
15784        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
15785        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
15786        // through to the dequant below when they do, never silently produce nothing.
15787        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15788            return Ok(Some(y));
15789        }
15790        let (in_f, out_f) = (w.in_features(), w.out_features());
15791        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15792        let tmp = GpuTensor::Quant {
15793            bytes: slab,
15794            qtype: QT_Q8_0,
15795            row_bytes: in_f / 32 * 34,
15796            ne: vec![in_f as u64, out_f as u64],
15797            scale: 1.0,
15798            rp: false,
15799            #[cfg(memra_cutlass)]
15800            cutlass: None,
15801            fp8: None,
15802            blk: None,
15803            f16: None,
15804            rp4: None,
15805        };
15806        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
15807        Ok(Some(self.matmul(&tmp, x, m)?))
15808    }
15809
15810    pub fn matmul_pre_noscale(
15811        &self,
15812        w: &crate::model::GpuTensor,
15813        aq: &CudaSlice<i8>,
15814        ad: &CudaSlice<f32>,
15815        m: usize,
15816    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15817        use crate::model::GpuTensor;
15818        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
15819        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
15820        // rather than let the tail below refuse and cost the caller a re-dispatch.
15821        if m == 1 {
15822            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15823                return Ok(Some((y, 1.0)));
15824            }
15825        }
15826        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
15827        if m != 1 || !self.uses_q8_1_fast(w) {
15828            return Ok(None);
15829        }
15830        let in_f = w.in_features();
15831        let out_f = w.out_features();
15832        let (bytes, qtype, row_bytes, scale, rp) = match w {
15833            GpuTensor::Quant {
15834                bytes,
15835                qtype,
15836                row_bytes,
15837                scale,
15838                rp,
15839                ..
15840            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15841            _ => return Ok(None),
15842        };
15843        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
15844        if self.mmvq_supports(qtype) {
15845            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
15846            let (mbytes, mrp) = match w {
15847                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15848                _ => (bytes, rp),
15849            };
15850            let y = self.qmatvec_mmvq(
15851                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
15852            )?;
15853            return Ok(Some((y, scale)));
15854        }
15855        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
15856        let name = match qtype {
15857            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15858            QT_Q4_K => "qmatvec_q4_K_dp4a",
15859            QT_Q6_K => "qmatvec_q6_K_dp4a",
15860            QT_Q5_K => "qmatvec_q5_K_dp4a",
15861            QT_Q3_K => "qmatvec_q3_K_dp4a",
15862            QT_NVFP4 => {
15863                if rp {
15864                    "qmatvec_nvfp4_dp4a_rp"
15865                } else {
15866                    "qmatvec_nvfp4_dp4a"
15867                }
15868            }
15869            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15870            _ => return Ok(None),
15871        };
15872        let f = self.func(name);
15873        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15874        let cfg = LaunchConfig {
15875            grid_dim: (out_f as u32, m as u32, 1),
15876            block_dim: (128, 1, 1),
15877            shared_mem_bytes: 0,
15878        };
15879        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15880        let __s_b = self.gpu.stream();
15881        let mut b = __s_b.launch_builder(&f);
15882        b.arg(bytes)
15883            .arg(aq)
15884            .arg(ad)
15885            .arg(&mut y)
15886            .arg(&inf)
15887            .arg(&outf)
15888            .arg(&mi)
15889            .arg(&rb);
15890        unsafe {
15891            b.launch(cfg)?;
15892        }
15893        Ok(Some((y, scale)))
15894    }
15895
15896    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15897    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15898    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15899        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15900        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15901        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15902        // is a pure function of the dtype — the decode-parity law holds under every env.
15903        if qtype == QT_F8_E4M3 {
15904            return true;
15905        }
15906        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15907            return false;
15908        }
15909        matches!(
15910            qtype,
15911            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15912        )
15913    }
15914
15915    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15916    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15917    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15918    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15919    pub fn qmatvec_mmvq(
15920        &self,
15921        bytes: &CudaSlice<u8>,
15922        aq: &CudaSlice<i8>,
15923        ad: &CudaSlice<f32>,
15924        m: usize,
15925        in_f: usize,
15926        out_f: usize,
15927        qtype: i32,
15928        row_bytes: usize,
15929        scale: f32,
15930        rp: bool,
15931    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15932        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15933        self.qmatvec_mmvq_into(
15934            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15935        )?;
15936        Ok(y)
15937    }
15938
15939    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15940    #[allow(clippy::too_many_arguments)]
15941    pub fn qmatvec_mmvq_into(
15942        &self,
15943        bytes: &CudaSlice<u8>,
15944        aq: &CudaSlice<i8>,
15945        ad: &CudaSlice<f32>,
15946        m: usize,
15947        in_f: usize,
15948        out_f: usize,
15949        qtype: i32,
15950        row_bytes: usize,
15951        scale: f32,
15952        rp: bool,
15953        y: &mut CudaSlice<f32>,
15954    ) -> Result<(), Box<dyn std::error::Error>> {
15955        debug_assert!(y.len() >= m * out_f);
15956        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15957        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15958        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15959        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15960        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15961        if qtype == QT_Q8_0
15962            && rp
15963            && m == 1
15964            && out_f >= 64
15965            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15966            && {
15967                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15968                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15969            }
15970        {
15971            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15972            let cfg = LaunchConfig {
15973                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15974                block_dim: (32, 2, 1),
15975                shared_mem_bytes: 0,
15976            };
15977            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15978            let __s_b = self.gpu.stream();
15979            let mut b = __s_b.launch_builder(&f);
15980            b.arg(bytes)
15981                .arg(aq)
15982                .arg(ad)
15983                .arg(&mut *y)
15984                .arg(&inf)
15985                .arg(&outf)
15986                .arg(&mi)
15987                .arg(&rb);
15988            unsafe {
15989                b.launch(cfg)?;
15990            }
15991            if scale != 1.0 {
15992                self.scale_inplace(y, scale, out_f)?;
15993            }
15994            return Ok(());
15995        }
15996        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15997        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15998        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15999        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
16000        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
16001        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
16002        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
16003        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
16004        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
16005            2
16006        } else {
16007            1
16008        };
16009        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
16010        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
16011        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
16012        // valid-window interleaved, bit-identical per row — same dot program).
16013        if m == 1 && qtype == QT_Q4_0 {
16014            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16015            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
16016            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
16017            mr = *Q40MR.get_or_init(|| {
16018                std::env::var("MEMRA_Q40_MR")
16019                    .ok()
16020                    .and_then(|v| v.parse().ok())
16021                    .unwrap_or(1)
16022            });
16023        }
16024        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
16025        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
16026        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
16027        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
16028        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
16029        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
16030        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
16031        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
16032        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
16033        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
16034        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
16035        let q5_force = q5_mode.as_deref() == Some("2");
16036        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
16037        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
16038        let q5_il = qtype == QT_Q5_K
16039            && m == 1
16040            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16041        if q5_il && !q5_force && out_f > 65536 {
16042            mr = 1;
16043        }
16044        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
16045        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
16046        if qtype == QT_Q4_0 && rp && mr != 1 {
16047            mr = 2;
16048        }
16049        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
16050        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
16051        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
16052        if qtype == QT_Q8_0 && rp {
16053            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16054            mr = *Q80MR.get_or_init(|| {
16055                std::env::var("MEMRA_Q80_MR")
16056                    .ok()
16057                    .and_then(|v| v.parse().ok())
16058                    .unwrap_or(1)
16059            });
16060        }
16061        let name = match (qtype, mr, rp) {
16062            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16063            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16064            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16065            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16066            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16067            (QT_Q5_K, 2, _) => {
16068                if q5_il {
16069                    "qmatvec_q5_K_mmvq_mr2_il"
16070                } else {
16071                    "qmatvec_q5_K_mmvq_mr2"
16072                }
16073            }
16074            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16075            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
16076            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
16077            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
16078            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
16079            (QT_Q8_0, _, true)
16080                if in_f % 1024 == 0 && {
16081                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16082                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16083                } =>
16084            {
16085                "qmatvec_q8_0_mmvq_rpca"
16086            }
16087            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16088            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16089            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
16090            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
16091            // reach a GGUF-layout kernel or vice versa.
16092            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16093            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16094            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16095            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16096            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16097            (QT_Q5_K, _, _) => {
16098                if q5_il {
16099                    "qmatvec_q5_K_mmvq_il"
16100                } else {
16101                    "qmatvec_q5_K_mmvq"
16102                }
16103            }
16104            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16105            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16106            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16107            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16108        };
16109        let f = self.func(name);
16110        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
16111        let rows_per_block = ROWS_PER_BLOCK * mr;
16112        let cfg = LaunchConfig {
16113            grid_dim: (
16114                (out_f as u32 + rows_per_block - 1) / rows_per_block,
16115                m as u32,
16116                1,
16117            ),
16118            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
16119            shared_mem_bytes: 0,                // warp-only reduce at m=1
16120        };
16121        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16122        let __s_b = self.gpu.stream();
16123        let mut b = __s_b.launch_builder(&f);
16124        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
16125        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
16126        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
16127        // weight_scale). Other mmvq kernels keep the 8-arg signature.
16128        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16129            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
16130            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
16131            if Self::pdl_on()
16132                && Self::pdl_mmvq_on()
16133                && Self::pdl_nvfp4q8_on()
16134                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16135            {
16136                use cudarc::driver::{DevicePtr, DevicePtrMut};
16137                let s = &self.gpu.stream();
16138                let (pw, _g0) = bytes.device_ptr(s);
16139                let (paq, _g1) = aq.device_ptr(s);
16140                let (pad, _g2) = ad.device_ptr(s);
16141                let (py, _g3) = y.device_ptr_mut(s);
16142                let mut ps = [
16143                    &pw as *const _ as *mut std::ffi::c_void,
16144                    &paq as *const _ as *mut _,
16145                    &pad as *const _ as *mut _,
16146                    &py as *const _ as *mut _,
16147                    &inf as *const _ as *mut _,
16148                    &outf as *const _ as *mut _,
16149                    &mi as *const _ as *mut _,
16150                    &rb as *const _ as *mut _,
16151                    &scale as *const _ as *mut _,
16152                ];
16153                unsafe {
16154                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16155                }
16156                return Ok(());
16157            }
16158            b.arg(bytes)
16159                .arg(aq)
16160                .arg(ad)
16161                .arg(&mut *y)
16162                .arg(&inf)
16163                .arg(&outf)
16164                .arg(&mi)
16165                .arg(&rb)
16166                .arg(&scale);
16167            unsafe {
16168                b.launch(cfg)?;
16169            }
16170        } else if Self::pdl_on()
16171            && Self::pdl_mmvq_on()
16172            && (matches!(
16173                name,
16174                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16175            ) || (Self::pdl_nvfp4q8_on()
16176                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16177        {
16178            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
16179            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
16180            // names may take this launch (unmarked kernels would read unordered).
16181            {
16182                use cudarc::driver::{DevicePtr, DevicePtrMut};
16183                let s = &self.gpu.stream();
16184                let (pw, _g0) = bytes.device_ptr(s);
16185                let (paq, _g1) = aq.device_ptr(s);
16186                let (pad, _g2) = ad.device_ptr(s);
16187                let (py, _g3) = y.device_ptr_mut(s);
16188                let mut ps = [
16189                    &pw as *const _ as *mut std::ffi::c_void,
16190                    &paq as *const _ as *mut _,
16191                    &pad as *const _ as *mut _,
16192                    &py as *const _ as *mut _,
16193                    &inf as *const _ as *mut _,
16194                    &outf as *const _ as *mut _,
16195                    &mi as *const _ as *mut _,
16196                    &rb as *const _ as *mut _,
16197                ];
16198                unsafe {
16199                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16200                }
16201            }
16202            if scale != 1.0 {
16203                self.scale_inplace(y, scale, m * out_f)?;
16204            }
16205        } else {
16206            b.arg(bytes)
16207                .arg(aq)
16208                .arg(ad)
16209                .arg(&mut *y)
16210                .arg(&inf)
16211                .arg(&outf)
16212                .arg(&mi)
16213                .arg(&rb);
16214            unsafe {
16215                b.launch(cfg)?;
16216            }
16217            if scale != 1.0 {
16218                self.scale_inplace(y, scale, m * out_f)?;
16219            }
16220        }
16221        Ok(())
16222    }
16223
16224    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
16225    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
16226    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
16227    pub fn qmatvec_mmvq_raw(
16228        &self,
16229        bytes: &CudaSlice<u8>,
16230        x: &CudaSlice<f32>,
16231        m: usize,
16232        in_f: usize,
16233        out_f: usize,
16234        qtype: i32,
16235        row_bytes: usize,
16236        rp: bool,
16237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16238        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16239        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16240    }
16241
16242    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
16243    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
16244    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
16245    pub fn batched_supports(&self, qtype: i32) -> bool {
16246        matches!(
16247            qtype,
16248            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16249        )
16250    }
16251
16252    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
16253    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
16254    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
16255    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
16256    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
16257    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
16258    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
16259    pub fn iq_fast_enabled() -> bool {
16260        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16261        *ON.get_or_init(|| {
16262            std::env::var("MEMRA_IQ_FAST")
16263                .map(|v| v != "0")
16264                .unwrap_or(true)
16265        })
16266    }
16267
16268    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
16269    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
16270    pub fn b8_enabled() -> bool {
16271        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16272        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16273    }
16274
16275    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
16276    pub fn batched_mcols(m: usize) -> usize {
16277        if m == 2 {
16278            2
16279        } else if m <= 4 {
16280            4
16281        } else if m <= 8 {
16282            8
16283        } else {
16284            16
16285        }
16286    }
16287
16288    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
16289    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
16290    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
16291    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
16292    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16293        Some(match (qtype, mcols) {
16294            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16295            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16296            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16297            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
16298            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
16299            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
16300            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
16301            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
16302            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16303            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16304            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16305            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16306            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
16307            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
16308            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16309            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16310            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16311            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16312            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
16313            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
16314            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16315            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16316            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16317            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16318            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16319            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16320            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16321            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16322            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
16323            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
16324            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
16325            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16326            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16327            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16328            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16329            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
16330            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
16331            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
16332            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16333            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16334            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16335            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16336            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16337            _ => return None,
16338        })
16339    }
16340
16341    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
16342    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
16343    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
16344    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
16345    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
16346    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
16347    ///
16348    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
16349    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
16350    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
16351    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
16352    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
16353    /// msweep on all six 27B shapes (2026-07-03):
16354    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
16355    ///          it applies for b4 (-3..-14%), never loses;
16356    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
16357    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
16358    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
16359    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
16360    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
16361    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
16362    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
16363    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
16364    /// b2: in_f>=6144 -> r2, else base.
16365    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
16366    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
16367    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
16368    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
16369    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
16370    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
16371    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
16372    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
16373    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
16374    /// Device SM count (cached) — grid-fill policy input.
16375    pub fn sm_count(&self) -> i32 {
16376        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16377        *SMS.get_or_init(|| {
16378            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16379            self.gpu
16380                .ctx
16381                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16382                .unwrap_or(82)
16383        })
16384    }
16385
16386    pub fn batched_variant(
16387        &self,
16388        _m: usize,
16389        in_f: usize,
16390        out_f: usize,
16391        qtype: i32,
16392        row_bytes: usize,
16393        mcols: usize,
16394        rp: bool,
16395    ) -> &'static str {
16396        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
16397        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
16398        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
16399        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
16400        if qtype == QT_Q8_0 {
16401            return if rp { "rp" } else { "base" };
16402        }
16403        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16404        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16405            Ok("base") => "base",
16406            Ok("pf") => "pf",
16407            Ok("r2") => "r2",
16408            Ok("r2w8") => "r2w8",
16409            Ok("pfr2") => "pfr2",
16410            Ok("ca") => "ca",
16411            Ok("car2") => "car2",
16412            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
16413            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
16414            Ok("rp") => "rp",
16415            Ok("rpr2") => "rpr2",
16416            Ok("rpr2w8") => "rpr2w8",
16417            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
16418            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
16419            Ok("rpca") => "rpca",
16420            Ok("rpcar2") => "rpcar2",
16421            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
16422            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
16423            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
16424            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
16425            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
16426            // bit-identical to the decode path — measurement corpus ONLY, never auto).
16427            Ok("rpsc") => "rpsc",
16428            Ok("rpms") => "rpms",
16429            Ok("rpmsc") => "rpmsc",
16430            Ok("rpks") => "rpks",
16431            Ok("rpksc") => "rpksc",
16432            _ => "auto",
16433        });
16434        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
16435        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
16436        // shapes qualify; anything else falls back to the register variants.
16437        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16438        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
16439        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
16440        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
16441        // forced MEMRA_MMVQ_BV values still work).
16442        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16443        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16444        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16445        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16446        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16447        let sms = *SMS.get_or_init(|| {
16448            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16449            self.gpu
16450                .ctx
16451                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16452                .unwrap_or(82)
16453        });
16454        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
16455        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
16456        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
16457        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
16458        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
16459        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
16460        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
16461        // AUTO RULE = the measured winners table (differs from NVFP4's!):
16462        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
16463        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
16464        //     r2 1258us) — kernels kept behind the force seam for the corpus;
16465        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
16466        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
16467        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
16468        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
16469        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
16470        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
16471        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
16472        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
16473        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16474        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
16475        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
16476        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16477        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16478            Ok("base") => "base",
16479            Ok("r2") => "r2",
16480            Ok("r2w8") => "r2w8",
16481            _ => "auto",
16482        });
16483        let variant: &'static str = if qtype == QT_Q4_0 {
16484            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
16485            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
16486            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
16487            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16488            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16489                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
16490                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
16491                // + syncs cost more than the stalls, bank-pad made no difference);
16492                // register load-ahead flat (nvcc already reorders). The b-tier limiter
16493                // is still unidentified — see the jsonl row.
16494                Ok("base") => "base",
16495                Ok("r2") => "r2",
16496                Ok("ms") => "ms",
16497                Ok("sm") => "sm",
16498                Ok("la") => "la",
16499                _ => "auto",
16500            });
16501            let v = if q40 != "auto" {
16502                q40
16503            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16504                "r2"
16505            } else {
16506                "base"
16507            };
16508            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
16509            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
16510            // and the limiter is the per-column activation load chain (long_scoreboard
16511            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
16512            if rp {
16513                match v {
16514                    "ms" => "r2ms_rp",
16515                    "sm" => "r2sm_rp",
16516                    "la" => "r2la_rp",
16517                    "r2" => "r2_rp",
16518                    _ => "rp",
16519                }
16520            } else if matches!(v, "ms" | "sm" | "la") {
16521                "r2"
16522            } else {
16523                v
16524            }
16525        } else if qtype != QT_NVFP4 && !kq_r2 {
16526            "base"
16527        } else if kq_r2 && rp {
16528            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
16529            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
16530            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
16531            "rp"
16532        } else if kq_r2 {
16533            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
16534            // mcols != 4 forced r2w8 falls to unbounded r2.
16535            if kq_bv != "auto" {
16536                if kq_bv == "r2w8" && mcols != 4 {
16537                    "r2"
16538                } else {
16539                    kq_bv
16540                }
16541            } else if bv != "auto" {
16542                match bv {
16543                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16544                    "r2w8" | "rpr2w8" => {
16545                        if mcols != 4 {
16546                            "r2"
16547                        } else {
16548                            "r2w8"
16549                        }
16550                    }
16551                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
16552                }
16553            } else {
16554                let blocks = (out_f + 7) / 8;
16555                let waves = blocks as f64 / (7 * sms as usize) as f64;
16556                let filled = blocks >= 4 * sms as usize;
16557                let use_r2 = if qtype == QT_Q4_K {
16558                    filled
16559                } else {
16560                    waves >= 2.0
16561                };
16562                if use_r2 { "r2" } else { "base" }
16563            }
16564        } else if bv != "auto" {
16565            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
16566            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
16567            // unsupported (shape, mcols) combos fall back to pf/r2.
16568            // On rp buffers, forced legacy names map to their rp twins (layout law).
16569            let v = if bv == "r2w8" && mcols == 2 {
16570                "r2"
16571            } else if bv == "ca" && (!ca_ok || mcols == 8) {
16572                "pf"
16573            } else if bv == "car2" && (!ca_ok || mcols == 8) {
16574                "r2"
16575            } else if bv == "pfr2" && mcols == 8 {
16576                "r2"
16577            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16578                "rpr2"
16579            }
16580            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
16581            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16582                if mcols == 8 { "rpr2w8" } else { "rpr2" }
16583            } else if bv == "rpcar2" && mcols == 2 {
16584                "rpca"
16585            }
16586            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
16587            // (rpms has no smem and no alignment need — always valid on rp buffers).
16588            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16589                "rpr2"
16590            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16591                "rpr2"
16592            } else {
16593                bv
16594            };
16595            if rp {
16596                match v {
16597                    "base" | "pf" | "ca" | "rp" => "rp",
16598                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16599                    "r2w8" | "rpr2w8" => {
16600                        if mcols == 2 {
16601                            "rpr2"
16602                        } else {
16603                            "rpr2w8"
16604                        }
16605                    }
16606                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
16607                }
16608            } else {
16609                v
16610            }
16611        } else if mcols == 8 {
16612            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
16613            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
16614            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
16615            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
16616            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
16617            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
16618            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
16619            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
16620            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
16621            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
16622            if rp {
16623                if sc_ok { "rpsc" } else { "rpr2w8" }
16624            } else {
16625                "r2w8"
16626            }
16627        } else if mcols >= 4 {
16628            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
16629            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
16630            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
16631            let blocks = (out_f + 7) / 8;
16632            let r7 = 7 * sms as usize;
16633            let r8 = 8 * sms as usize;
16634            let waves = blocks as f64 / r7 as f64;
16635            let filled = blocks >= 4 * sms as usize;
16636            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
16637            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
16638            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
16639            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16640                // the extra residency drops the INTEGER wave count -> the straggler wave a
16641                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
16642                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
16643                if rp { "rpr2w8" } else { "r2w8" }
16644            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16645                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
16646                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
16647                if rp { "rpr2" } else { "r2" }
16648            } else {
16649                // fractional straggler-wave window with no crossing, or grid too small to fill
16650                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
16651                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
16652                if rp { "rp" } else { "pf" }
16653            }
16654        } else if in_f >= 6144 {
16655            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
16656            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
16657            // stays.
16658            if rp { "rpr2" } else { "r2" }
16659        } else if rp {
16660            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
16661            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
16662            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
16663            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
16664            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16665            if sc_ok && waves >= 0.9 && waves <= 1.1 {
16666                "rpsc"
16667            } else {
16668                "rp"
16669            }
16670        } else {
16671            "base"
16672        };
16673        variant
16674    }
16675
16676    pub fn qmatvec_mmvq_batched(
16677        &self,
16678        bytes: &CudaSlice<u8>,
16679        aq: &CudaSlice<i8>,
16680        ad: &CudaSlice<f32>,
16681        m: usize,
16682        in_f: usize,
16683        out_f: usize,
16684        qtype: i32,
16685        row_bytes: usize,
16686        mcols: usize,
16687        scale: f32,
16688        rp: bool,
16689    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16690        const ROWS_PER_BLOCK: u32 = 4;
16691        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
16692        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
16693        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
16694        // weight keeps its rp-layout kernel family regardless of the override.
16695        let forced: Option<&'static str> = {
16696            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16697            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16698                .as_deref()
16699                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16700        };
16701        let variant = match forced {
16702            Some(v) if !rp || v.contains("rp") => v,
16703            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16704        };
16705        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16706            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16707        })?;
16708        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
16709        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
16710        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
16711        let variant = if mcols == 16 {
16712            if rp { "rp" } else { "base" }
16713        } else {
16714            variant
16715        };
16716        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
16717        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
16718        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
16719        // per-(token,row) chain (columns c >= m never execute in either form) ->
16720        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
16721        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
16722        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16723        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16724        if b567
16725            && qtype == QT_NVFP4
16726            && rp
16727            && mcols == 8
16728            && (5..=7).contains(&m)
16729            && matches!(variant, "rpsc" | "rpr2w8")
16730        {
16731            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16732            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
16733            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16734            let cfg = LaunchConfig {
16735                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16736                block_dim: (32, ROWS_PER_BLOCK, 1),
16737                shared_mem_bytes: 0,
16738            };
16739            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16740            let __s_b = self.gpu.stream();
16741            let mut b = __s_b.launch_builder(&f);
16742            b.arg(bytes)
16743                .arg(aq)
16744                .arg(ad)
16745                .arg(&mut y)
16746                .arg(&inf)
16747                .arg(&outf)
16748                .arg(&mi)
16749                .arg(&rb);
16750            unsafe {
16751                b.launch(cfg)?;
16752            }
16753            if scale != 1.0 {
16754                self.scale_inplace(&mut y, scale, m * out_f)?;
16755            }
16756            return Ok(y);
16757        }
16758        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16759            "base" => (base_name.into(), ROWS_PER_BLOCK),
16760            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16761            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16762            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16763            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
16764            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
16765            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
16766            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16767            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16768            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16769            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16770            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16771            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16772            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16773            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
16774        };
16775        debug_assert!(
16776            !rp || name.contains("_rp"),
16777            "rp weight dispatched to a GGUF-layout kernel"
16778        );
16779        let f = self.func(&name);
16780        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16781        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
16782        let smem = if name.contains("_r2sm_rp") {
16783            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16784        } else {
16785            0
16786        };
16787        let cfg = LaunchConfig {
16788            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16789            block_dim: (32, ROWS_PER_BLOCK, 1),
16790            shared_mem_bytes: smem,
16791        };
16792        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16793        let __s_b = self.gpu.stream();
16794        let mut b = __s_b.launch_builder(&f);
16795        b.arg(bytes)
16796            .arg(aq)
16797            .arg(ad)
16798            .arg(&mut y)
16799            .arg(&inf)
16800            .arg(&outf)
16801            .arg(&mi)
16802            .arg(&rb);
16803        unsafe {
16804            b.launch(cfg)?;
16805        }
16806        if scale != 1.0 {
16807            self.scale_inplace(&mut y, scale, m * out_f)?;
16808        }
16809        Ok(y)
16810    }
16811
16812    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
16813    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
16814    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
16815    pub fn qmatvec_batched_raw(
16816        &self,
16817        bytes: &CudaSlice<u8>,
16818        x: &CudaSlice<f32>,
16819        m: usize,
16820        in_f: usize,
16821        out_f: usize,
16822        qtype: i32,
16823        row_bytes: usize,
16824        mcols: usize,
16825        rp: bool,
16826    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16827        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16828        self.qmatvec_mmvq_batched(
16829            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16830        )
16831    }
16832
16833    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
16834    pub fn qmatvec_nvfp4_batched_raw(
16835        &self,
16836        bytes: &CudaSlice<u8>,
16837        x: &CudaSlice<f32>,
16838        m: usize,
16839        in_f: usize,
16840        out_f: usize,
16841        row_bytes: usize,
16842        mcols: usize,
16843        rp: bool,
16844    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16845        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16846    }
16847
16848    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
16849    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
16850    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
16851    fn try_fp4_gemm(
16852        &self,
16853        w: &crate::model::GpuTensor,
16854        x: &CudaSlice<f32>,
16855        m: usize,
16856        in_f: usize,
16857        out_f: usize,
16858    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16859        use crate::model::GpuTensor;
16860        if cfg!(memra_portable_cuda) {
16861            return Ok(None);
16862        }
16863        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which cu/qmatvec_gemm.cu:1234 omits on a
16864        // portable build (the mxf4 block-scale MMA is sm_120a-only). Refuse at the door.
16865        if std::env::var("MEMRA_FP4").is_ok() {
16866            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16867        }
16868        if std::env::var("MEMRA_FP4").is_err() {
16869            return Ok(None);
16870        }
16871        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
16872        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
16873        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
16874        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
16875        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
16876        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
16877        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
16878        // for the common no-macro-scale case.
16879        #[cfg(memra_cutlass)]
16880        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16881            if let GpuTensor::Quant {
16882                bytes,
16883                qtype,
16884                scale,
16885                row_bytes,
16886                cutlass,
16887                ..
16888            } = w
16889            {
16890                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16891                    if let Some(cw) = cutlass {
16892                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16893                        let y = self.cutlass_fp4_gemm(
16894                            &cw.b_packed,
16895                            &cw.sfb_swizzled,
16896                            x,
16897                            *scale,
16898                            m,
16899                            out_f,
16900                            in_f,
16901                        )?;
16902                        return Ok(Some(y));
16903                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16904                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16905                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16906                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16907                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16908                        let (b_packed, sfb_sw) =
16909                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16910                        let y =
16911                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16912                        return Ok(Some(y));
16913                    }
16914                }
16915            }
16916        }
16917        if let GpuTensor::Quant {
16918            bytes,
16919            qtype,
16920            row_bytes,
16921            scale,
16922            rp,
16923            ..
16924        } = w
16925        {
16926            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16927            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16928            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16929                let y =
16930                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16931                return Ok(Some(y));
16932            }
16933        }
16934        Ok(None)
16935    }
16936
16937    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16938    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16939    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16940    pub fn rms_norm_f16out(
16941        &self,
16942        x: &CudaSlice<f32>,
16943        w: &CudaSlice<f32>,
16944        dst: &mut CudaSlice<f32>,
16945        dst16: &mut CudaSlice<u8>,
16946        ncols: usize,
16947        nrows: usize,
16948        eps: f32,
16949    ) -> Result<(), Box<dyn std::error::Error>> {
16950        let f = self.func("rms_norm_f16out_f32");
16951        let cfg = LaunchConfig {
16952            grid_dim: (nrows as u32, 1, 1),
16953            block_dim: (rms_block(), 1, 1),
16954            shared_mem_bytes: 0,
16955        };
16956        let (nc, e) = (ncols as i32, eps);
16957        let __s_b = self.gpu.stream();
16958        let mut b = __s_b.launch_builder(&f);
16959        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16960        unsafe {
16961            b.launch(cfg)?;
16962        }
16963        Ok(())
16964    }
16965
16966    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16967    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16968    #[allow(clippy::too_many_arguments)]
16969    pub fn add_rms_norm_f16out(
16970        &self,
16971        a: &CudaSlice<f32>,
16972        b: &CudaSlice<f32>,
16973        w: &CudaSlice<f32>,
16974        res: &mut CudaSlice<f32>,
16975        dst: &mut CudaSlice<f32>,
16976        dst16: &mut CudaSlice<u8>,
16977        ncols: usize,
16978        nrows: usize,
16979        eps: f32,
16980    ) -> Result<(), Box<dyn std::error::Error>> {
16981        let f = self.func("add_rms_norm_f16out_f32");
16982        let cfg = LaunchConfig {
16983            grid_dim: (nrows as u32, 1, 1),
16984            block_dim: (rms_block(), 1, 1),
16985            shared_mem_bytes: 0,
16986        };
16987        let (nc, e) = (ncols as i32, eps);
16988        let __s_lb = self.gpu.stream();
16989        let mut lb = __s_lb.launch_builder(&f);
16990        lb.arg(a)
16991            .arg(b)
16992            .arg(w)
16993            .arg(res)
16994            .arg(dst)
16995            .arg(dst16)
16996            .arg(&nc)
16997            .arg(&e);
16998        unsafe {
16999            lb.launch(cfg)?;
17000        }
17001        Ok(())
17002    }
17003
17004    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
17005    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
17006    pub fn matmul_group_xh(
17007        &self,
17008        ws: &[&crate::model::GpuTensor],
17009        x: &CudaSlice<f32>,
17010        xh: &CudaSlice<u8>,
17011        m: usize,
17012    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17013        let mut out = Vec::with_capacity(ws.len());
17014        let in_f = ws[0].in_features();
17015        for w in ws {
17016            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
17017                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
17018                    out.push(y);
17019                    continue;
17020                }
17021            }
17022            out.push(self.matmul(w, x, m)?);
17023        }
17024        Ok(out)
17025    }
17026
17027    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
17028    /// GDN steps). Layouts [T, H].
17029    pub fn gdn_pad_mask(
17030        &self,
17031        beta: &mut CudaSlice<f32>,
17032        g_log: &mut CudaSlice<f32>,
17033        len_d: &CudaSlice<i32>,
17034        h: usize,
17035        t: usize,
17036    ) -> Result<(), Box<dyn std::error::Error>> {
17037        let f = self.func("gdn_pad_mask_f32");
17038        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
17039        let (hi, ti) = (h as i32, t as i32);
17040        let __s_b = self.gpu.stream();
17041        let mut b = __s_b.launch_builder(&f);
17042        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17043        unsafe {
17044            b.launch(cfg)?;
17045        }
17046        Ok(())
17047    }
17048
17049    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
17050    /// gather for the padded prime graph's h_seed/hlast.
17051    pub fn row_gather_dev(
17052        &self,
17053        src: &CudaSlice<f32>,
17054        dst: &mut CudaSlice<f32>,
17055        len_d: &CudaSlice<i32>,
17056        ncols: usize,
17057    ) -> Result<(), Box<dyn std::error::Error>> {
17058        let f = self.func("row_gather_dev_f32");
17059        let cfg = LaunchConfig::for_num_elems(ncols as u32);
17060        let nc = ncols as i32;
17061        let __s_b = self.gpu.stream();
17062        let mut b = __s_b.launch_builder(&f);
17063        b.arg(src).arg(dst).arg(len_d).arg(&nc);
17064        unsafe {
17065            b.launch(cfg)?;
17066        }
17067        Ok(())
17068    }
17069
17070    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
17071    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
17072    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
17073    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
17074    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
17075    /// different in_f) falls back to its own `matmul` — behavior unchanged.
17076    pub fn matmul_group(
17077        &self,
17078        ws: &[&crate::model::GpuTensor],
17079        x: &CudaSlice<f32>,
17080        m: usize,
17081    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17082        use crate::model::GpuTensor;
17083        let mut out = Vec::with_capacity(ws.len());
17084        let any_mirror = ws
17085            .iter()
17086            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17087        if m >= 16 && any_mirror && !self.verify_exact_on() {
17088            let in_f = ws[0].in_features();
17089            let xh = self.f16_act(x, m * in_f, in_f)?;
17090            for w in ws {
17091                if w.in_features() == in_f {
17092                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17093                        out.push(y);
17094                        continue;
17095                    }
17096                }
17097                out.push(self.matmul(w, x, m)?);
17098            }
17099            return Ok(out);
17100        }
17101        for w in ws {
17102            out.push(self.matmul(w, x, m)?);
17103        }
17104        Ok(out)
17105    }
17106
17107    /// Cross-request grouped matmul (task #13): run ONE projection group over the
17108    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
17109    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
17110    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
17111    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
17112    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
17113    pub fn matmul_group_multi(
17114        &self,
17115        ws: &[&crate::model::GpuTensor],
17116        xs: &[&CudaSlice<f32>],
17117        ms: &[usize],
17118    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17119        assert_eq!(xs.len(), ms.len());
17120        let in_f = ws[0].in_features();
17121        let total: usize = ms.iter().sum();
17122        let mut xcat = self.uninit(total * in_f)?;
17123        let mut off = 0usize;
17124        for (x, &m) in xs.iter().zip(ms) {
17125            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17126            off += m;
17127        }
17128        let ys = self.matmul_group(ws, &xcat, total)?;
17129        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17130        for (w, y) in ws.iter().zip(ys) {
17131            let out_f = w.out_features();
17132            let mut off = 0usize;
17133            for (s, &m) in ms.iter().enumerate() {
17134                let mut ys_s = self.uninit(m * out_f)?;
17135                let src = y.slice(off * out_f..(off + m) * out_f);
17136                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17137                out[s].push(ys_s);
17138                off += m;
17139            }
17140        }
17141        Ok(out)
17142    }
17143
17144    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
17145    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
17146    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
17147    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
17148    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
17149    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
17150    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
17151    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
17152    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
17153    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17154        use crate::model::GpuTensor;
17155        if !legacy_quant_gemm_allowed(
17156            cfg!(memra_portable_cuda),
17157            cfg!(memra_hopper_mma),
17158            std::env::var_os("MEMRA_NO_GEMM").is_some(),
17159        ) {
17160            return false;
17161        }
17162        match w {
17163            GpuTensor::Quant { qtype, .. } => {
17164                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17165                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17166            }
17167            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17168        }
17169    }
17170
17171    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
17172    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
17173    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
17174    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
17175    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
17176    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
17177    pub fn qmatvec_gemm(
17178        &self,
17179        w: &crate::model::GpuTensor,
17180        aq: &CudaSlice<i8>,
17181        ad: &CudaSlice<f32>,
17182        m: usize,
17183    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17184        use crate::model::GpuTensor;
17185        let in_f = w.in_features();
17186        let out_f = w.out_features();
17187        let (bytes, qtype, row_bytes, scale, rp) = match w {
17188            GpuTensor::Quant {
17189                bytes,
17190                qtype,
17191                row_bytes,
17192                scale,
17193                rp,
17194                ..
17195            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17196            _ => unreachable!("gemm_supports guaranteed Quant"),
17197        };
17198        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
17199        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
17200        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
17201        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
17202        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
17203        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17204            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17205                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17206                if scale != 1.0 {
17207                    self.scale_inplace(&mut y, scale, m * out_f)?;
17208                }
17209                return Ok(y);
17210            }
17211        }
17212        let name = match qtype {
17213            QT_Q8_0 => "qmatvec_gemm_q8_0",
17214            QT_Q4_K => "qmatvec_gemm_q4_K",
17215            QT_Q4_0 => {
17216                if rp {
17217                    "qmatvec_gemm_q4_0_rp"
17218                } else {
17219                    "qmatvec_gemm_q4_0"
17220                }
17221            }
17222            QT_Q5_K => "qmatvec_gemm_q5_K",
17223            QT_Q6_K => "qmatvec_gemm_q6_K",
17224            QT_NVFP4 => {
17225                if rp {
17226                    "qmatvec_gemm_nvfp4_rp"
17227                } else {
17228                    "qmatvec_gemm_nvfp4"
17229                }
17230            }
17231            _ => unreachable!(),
17232        };
17233        let f = self.func(name);
17234        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17235        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
17236        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
17237        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
17238        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
17239        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
17240        let k1_tile = if is_k1 {
17241            k1_launch_override().unwrap_or((128, 128, 8))
17242        } else {
17243            (128, 128, 8)
17244        };
17245        let (bm, bn): (u32, u32) = if is_k1 {
17246            (k1_tile.0, k1_tile.1)
17247        } else {
17248            (64, 256)
17249        };
17250        let warps: u32 = if is_k1 {
17251            k1_tile.2
17252        } else {
17253            match qtype {
17254                QT_NVFP4 => 8,
17255                _ => 4,
17256            }
17257        };
17258        let cfg = LaunchConfig {
17259            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17260            block_dim: (32, warps, 1),
17261            shared_mem_bytes: 0,
17262        };
17263        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17264        let __s_b = self.gpu.stream();
17265        let mut b = __s_b.launch_builder(&f);
17266        b.arg(bytes)
17267            .arg(aq)
17268            .arg(ad)
17269            .arg(&mut y)
17270            .arg(&inf)
17271            .arg(&outf)
17272            .arg(&mi)
17273            .arg(&rb);
17274        unsafe {
17275            b.launch(cfg)?;
17276        }
17277        if scale != 1.0 {
17278            self.scale_inplace(&mut y, scale, m * out_f)?;
17279        }
17280        Ok(y)
17281    }
17282
17283    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
17284    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
17285    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
17286    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
17287    pub fn qmatvec_gemm_raw(
17288        &self,
17289        bytes: &CudaSlice<u8>,
17290        x: &CudaSlice<f32>,
17291        m: usize,
17292        in_f: usize,
17293        out_f: usize,
17294        qtype: i32,
17295        row_bytes: usize,
17296    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17297        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17298        let name = match qtype {
17299            QT_Q8_0 => "qmatvec_gemm_q8_0",
17300            QT_Q4_K => "qmatvec_gemm_q4_K",
17301            QT_Q4_0 => "qmatvec_gemm_q4_0",
17302            QT_Q5_K => "qmatvec_gemm_q5_K",
17303            QT_Q6_K => "qmatvec_gemm_q6_K",
17304            QT_NVFP4 => "qmatvec_gemm_nvfp4",
17305            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17306            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17307        };
17308        let f = self.func(name);
17309        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17310        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
17311        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
17312        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
17313        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
17314        let k1_tile = if is_k1 {
17315            k1_launch_override().unwrap_or((128, 128, 8))
17316        } else {
17317            (128, 128, 8)
17318        };
17319        let (bm, bn): (u32, u32) = if is_k1 {
17320            (k1_tile.0, k1_tile.1)
17321        } else {
17322            (64, 256)
17323        };
17324        let warps: u32 = if is_k1 {
17325            k1_tile.2
17326        } else {
17327            match qtype {
17328                QT_NVFP4 | QT_NVFP4_RP => 8,
17329                _ => 4,
17330            }
17331        };
17332        let cfg = LaunchConfig {
17333            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17334            block_dim: (32, warps, 1),
17335            shared_mem_bytes: 0,
17336        };
17337        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17338        let __s_b = self.gpu.stream();
17339        let mut b = __s_b.launch_builder(&f);
17340        b.arg(bytes)
17341            .arg(&aq)
17342            .arg(&ad)
17343            .arg(&mut y)
17344            .arg(&inf)
17345            .arg(&outf)
17346            .arg(&mi)
17347            .arg(&rb);
17348        unsafe {
17349            b.launch(cfg)?;
17350        }
17351        Ok(y)
17352    }
17353
17354    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
17355    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
17356    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
17357    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
17358    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
17359    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
17360    pub fn qmatvec_gemm_q8_0_wgmma_raw(
17361        &self,
17362        rp4: &CudaSlice<u8>,
17363        aq: &CudaSlice<i8>,
17364        ad: &CudaSlice<f32>,
17365        m: usize,
17366        in_f: usize,
17367        out_f: usize,
17368    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17369        assert!(
17370            out_f % 64 == 0 && in_f % 32 == 0,
17371            "wgmma GEMM needs out_f%64==0, in_f%32==0"
17372        );
17373        let f = self.func("qmatvec_gemm_q8_0_wgmma");
17374        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
17375        let cfg = LaunchConfig {
17376            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17377            block_dim: (128, 1, 1),
17378            shared_mem_bytes: 0,
17379        };
17380        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17381        let __s_b = self.gpu.stream();
17382        let mut b = __s_b.launch_builder(&f);
17383        b.arg(rp4)
17384            .arg(aq)
17385            .arg(ad)
17386            .arg(&mut y)
17387            .arg(&inf)
17388            .arg(&outf)
17389            .arg(&mi);
17390        unsafe {
17391            b.launch(cfg)?;
17392        }
17393        Ok(y)
17394    }
17395
17396    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
17397    pub fn scale_inplace(
17398        &self,
17399        y: &mut CudaSlice<f32>,
17400        s: f32,
17401        n: usize,
17402    ) -> Result<(), Box<dyn std::error::Error>> {
17403        let f = self.func("scale_f32");
17404        let cfg = LaunchConfig::for_num_elems(n as u32);
17405        let (sf, ni) = (s, n as i32);
17406        let __s_b = self.gpu.stream();
17407        let mut b = __s_b.launch_builder(&f);
17408        b.arg(y).arg(&sf).arg(&ni);
17409        unsafe {
17410            b.launch(cfg)?;
17411        }
17412        Ok(())
17413    }
17414
17415    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
17416    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
17417    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
17418    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
17419    pub fn bf16_to_f32(
17420        &self,
17421        data: &cudarc::driver::CudaView<'_, u8>,
17422        n: usize,
17423    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17424        let mut out = self.alloc_uninit::<f32>(n)?;
17425        let f = self.func("bf16_to_f32");
17426        let cfg = LaunchConfig::for_num_elems(n as u32);
17427        let ni = n as i32;
17428        let __s_b = self.gpu.stream();
17429        let mut b = __s_b.launch_builder(&f);
17430        b.arg(data).arg(&mut out).arg(&ni);
17431        unsafe {
17432            b.launch(cfg)?;
17433        }
17434        Ok(out)
17435    }
17436
17437    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
17438    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
17439    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
17440    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
17441    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
17442    /// calls, the spec-verify contract) vs plain linear.
17443    fn linear_bf16_chunked(
17444        &self,
17445        x: &CudaSlice<f32>,
17446        data: &CudaSlice<u8>,
17447        m: usize,
17448        in_f: usize,
17449        out_f: usize,
17450        exact: bool,
17451        canonical_chunk_rows: Option<usize>,
17452    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17453        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
17454        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
17455        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
17456        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17457        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17458        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17459        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17460        let started = timing.then(std::time::Instant::now);
17461        let result =
17462            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17463        if let Some(started) = started {
17464            use std::sync::atomic::Ordering;
17465            self.stream().synchronize()?;
17466            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17467                + started.elapsed().as_nanos() as u64;
17468            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17469                + (in_f * out_f * 2) as u64;
17470            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17471            if calls % 1024 == 0 {
17472                eprintln!(
17473                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17474                     weight_gb={:.2}",
17475                    ns as f64 / 1.0e6,
17476                    ns as f64 / calls as f64 / 1.0e3,
17477                    wb as f64 / 1.0e9,
17478                );
17479            }
17480        }
17481        result
17482    }
17483
17484    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
17485    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
17486    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
17487    /// numeric-class doors (DEV_ROUTES precedent).
17488    pub(crate) fn bf16_mmv_on() -> bool {
17489        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17490        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17491    }
17492
17493    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
17494    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
17495    fn matvec_bf16(
17496        &self,
17497        data: &CudaSlice<u8>,
17498        x: &CudaSlice<f32>,
17499        in_f: usize,
17500        out_f: usize,
17501    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17502        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17503            return Err(format!(
17504                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17505                data.len(),
17506                x.len()
17507            )
17508            .into());
17509        }
17510        let mut y = self.alloc_uninit::<f32>(out_f)?;
17511        let f = self.func("matvec_bf16_f32acc");
17512        let cfg = LaunchConfig {
17513            grid_dim: (out_f as u32, 1, 1),
17514            block_dim: (mmv_block(), 1, 1),
17515            shared_mem_bytes: 0,
17516        };
17517        let ini = in_f as i32;
17518        let __s_bld = self.gpu.stream();
17519        let mut bld = __s_bld.launch_builder(&f);
17520        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17521        unsafe {
17522            bld.launch(cfg)?;
17523        }
17524        Ok(y)
17525    }
17526
17527    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
17528    /// launches, a position upload, and the rope launch; the position is read directly from
17529    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
17530    #[allow(clippy::too_many_arguments)]
17531    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
17532    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
17533    /// Bit-identical to the split kernels; requires head_dim == 128 and
17534    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
17535    #[allow(clippy::too_many_arguments)]
17536    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
17537    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
17538    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
17539    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
17540    #[allow(clippy::too_many_arguments)]
17541    pub fn qk_norm_rope_append_inc_dcw_rows(
17542        &self,
17543        q_raw_t: &CudaSlice<f32>,
17544        k_raw_t: &CudaSlice<f32>,
17545        v_raw_t: &CudaSlice<f32>,
17546        qw: &CudaSlice<f32>,
17547        kw: &CudaSlice<f32>,
17548        q_out_t: &mut CudaSlice<f32>,
17549        k_out_t: &mut CudaSlice<f32>,
17550        tab: &CudaSlice<u64>,
17551        pos_t: &CudaSlice<i32>,
17552        same_session: bool,
17553        t: usize,
17554        kv_dim_k: usize,
17555        kv_dim_v: usize,
17556        k_tok_bytes: usize,
17557        v_tok_bytes: usize,
17558        head_dim: usize,
17559        n_dims: usize,
17560        nh_q: usize,
17561        nh_k: usize,
17562        eps: f32,
17563        freq_base: f32,
17564        freq_scale: f32,
17565        ff: Option<&CudaSlice<f32>>,
17566    ) -> Result<(), Box<dyn std::error::Error>> {
17567        if head_dim != 128
17568            || kv_dim_v != kv_dim_k
17569            || kv_dim_k != nh_k * head_dim
17570            || t == 0
17571            || t > 32
17572            || tab.len() < t * 6
17573            || pos_t.len() < t
17574            || q_raw_t.len() < t * nh_q * head_dim
17575            || k_raw_t.len() < t * nh_k * head_dim
17576            || v_raw_t.len() < t * kv_dim_v
17577            || q_out_t.len() < t * nh_q * head_dim
17578            || k_out_t.len() < t * nh_k * head_dim
17579        {
17580            return Err(format!(
17581                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17582                 nh_q={nh_q} nh_k={nh_k}"
17583            )
17584            .into());
17585        }
17586        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17587        let same_t: i32 = if same_session { t as i32 } else { 0 };
17588        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17589        let cfg = LaunchConfig {
17590            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17591            block_dim: (128, 1, 1),
17592            shared_mem_bytes: 0,
17593        };
17594        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17595        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17596        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17597        let null: u64 = 0;
17598        let __s_b = self.gpu.stream();
17599        let mut b = __s_b.launch_builder(&f);
17600        b.arg(q_raw_t)
17601            .arg(k_raw_t)
17602            .arg(v_raw_t)
17603            .arg(qw)
17604            .arg(kw)
17605            .arg(q_out_t)
17606            .arg(k_out_t)
17607            .arg(tab)
17608            .arg(pos_t)
17609            .arg(&same_t)
17610            .arg(&kvk)
17611            .arg(&kvv)
17612            .arg(&ktb)
17613            .arg(&vtb)
17614            .arg(&hd)
17615            .arg(&nd)
17616            .arg(&nq)
17617            .arg(&nk)
17618            .arg(&eps)
17619            .arg(&theta_scale)
17620            .arg(&freq_scale);
17621        match ff {
17622            Some(freqs) => {
17623                b.arg(freqs);
17624            }
17625            None => {
17626                b.arg(&null);
17627            }
17628        }
17629        unsafe {
17630            b.launch(cfg)?;
17631        }
17632        Ok(())
17633    }
17634
17635    pub fn qk_norm_rope_append_inc_dcw(
17636        &self,
17637        q_raw: &CudaSlice<f32>,
17638        k_raw: &CudaSlice<f32>,
17639        v_raw: &CudaSlice<f32>,
17640        qw: &CudaSlice<f32>,
17641        kw: &CudaSlice<f32>,
17642        q_out: &mut CudaSlice<f32>,
17643        k_out: &mut CudaSlice<f32>,
17644        pos: &CudaSlice<i32>,
17645        k_plane: &mut CudaSlice<u8>,
17646        v_plane: &mut CudaSlice<u8>,
17647        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
17648        // (single) writer, exactly like the split append+inc pair it replaces.
17649        len_dev: &CudaSlice<i32>,
17650        base_dev: Option<&CudaSlice<i32>>,
17651        done_ctr: &mut CudaSlice<u32>,
17652        kv_dim_k: usize,
17653        kv_dim_v: usize,
17654        k_tok_bytes: usize,
17655        v_tok_bytes: usize,
17656        head_dim: usize,
17657        n_dims: usize,
17658        nh_q: usize,
17659        nh_k: usize,
17660        eps: f32,
17661        freq_base: f32,
17662        freq_scale: f32,
17663        ff: Option<&CudaSlice<f32>>,
17664    ) -> Result<(), Box<dyn std::error::Error>> {
17665        if head_dim != 128
17666            || kv_dim_v != kv_dim_k
17667            || kv_dim_k != nh_k * head_dim
17668            || q_raw.len() < nh_q * head_dim
17669            || k_raw.len() < nh_k * head_dim
17670            || v_raw.len() < kv_dim_v
17671            || q_out.len() < nh_q * head_dim
17672            || k_out.len() < nh_k * head_dim
17673            || pos.is_empty()
17674            || done_ctr.is_empty()
17675        {
17676            return Err(format!(
17677                "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}"
17678            )
17679            .into());
17680        }
17681        let f = self.func("qk_norm_rope_append_inc_dcw");
17682        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17683        let cfg = LaunchConfig {
17684            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17685            block_dim: (128, 1, 1),
17686            shared_mem_bytes: 0,
17687        };
17688        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17689        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17690        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17691        let null: u64 = 0;
17692        let __s_b = self.gpu.stream();
17693        let mut b = __s_b.launch_builder(&f);
17694        b.arg(q_raw)
17695            .arg(k_raw)
17696            .arg(v_raw)
17697            .arg(qw)
17698            .arg(kw)
17699            .arg(q_out)
17700            .arg(k_out)
17701            .arg(pos)
17702            .arg(&mut *k_plane)
17703            .arg(&mut *v_plane)
17704            .arg(len_dev);
17705        match base_dev {
17706            Some(base) => {
17707                b.arg(base);
17708            }
17709            None => {
17710                b.arg(&null);
17711            }
17712        }
17713        b.arg(&mut *done_ctr)
17714            .arg(&kvk)
17715            .arg(&kvv)
17716            .arg(&ktb)
17717            .arg(&vtb)
17718            .arg(&hd)
17719            .arg(&nd)
17720            .arg(&nq)
17721            .arg(&eps)
17722            .arg(&theta_scale)
17723            .arg(&freq_scale);
17724        match ff {
17725            Some(freqs) => {
17726                b.arg(freqs);
17727            }
17728            None => {
17729                b.arg(&null);
17730            }
17731        }
17732        unsafe {
17733            b.launch(cfg)?;
17734        }
17735        Ok(())
17736    }
17737
17738    pub fn qk_norm_rope_into(
17739        &self,
17740        q_raw: &CudaSlice<f32>,
17741        k_raw: &CudaSlice<f32>,
17742        qw: &CudaSlice<f32>,
17743        kw: &CudaSlice<f32>,
17744        q_out: &mut CudaSlice<f32>,
17745        k_out: &mut CudaSlice<f32>,
17746        pos: &CudaSlice<i32>,
17747        head_dim: usize,
17748        n_dims: usize,
17749        nh_q: usize,
17750        nh_k: usize,
17751        eps: f32,
17752        freq_base: f32,
17753        freq_scale: f32,
17754        ff: Option<&CudaSlice<f32>>,
17755    ) -> Result<(), Box<dyn std::error::Error>> {
17756        if head_dim > 512
17757            || q_raw.len() < nh_q * head_dim
17758            || k_raw.len() < nh_k * head_dim
17759            || q_out.len() < nh_q * head_dim
17760            || k_out.len() < nh_k * head_dim
17761            || qw.len() < head_dim
17762            || kw.len() < head_dim
17763            || pos.is_empty()
17764        {
17765            return Err(format!(
17766                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17767            )
17768            .into());
17769        }
17770        let f = self.func("qk_norm_rope_f32");
17771        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17772        let cfg = LaunchConfig {
17773            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17774            block_dim: (128, 1, 1),
17775            shared_mem_bytes: 0,
17776        };
17777        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17778        let __s_b = self.gpu.stream();
17779        let mut b = __s_b.launch_builder(&f);
17780        b.arg(q_raw)
17781            .arg(k_raw)
17782            .arg(qw)
17783            .arg(kw)
17784            .arg(q_out)
17785            .arg(k_out)
17786            .arg(pos)
17787            .arg(&hd)
17788            .arg(&nd)
17789            .arg(&nq)
17790            .arg(&eps)
17791            .arg(&theta_scale)
17792            .arg(&freq_scale);
17793        match ff {
17794            Some(ffv) => {
17795                b.arg(ffv);
17796                unsafe {
17797                    b.launch(cfg)?;
17798                }
17799            }
17800            None => {
17801                let null: u64 = 0;
17802                b.arg(&null);
17803                unsafe {
17804                    b.launch(cfg)?;
17805                }
17806            }
17807        }
17808        Ok(())
17809    }
17810
17811    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
17812    /// launch computes a rank's whole O partial from its four canonical column blocks.
17813    #[allow(clippy::too_many_arguments)]
17814    pub fn matvec_f32_b4_into(
17815        &self,
17816        w: [&CudaSlice<f32>; 4],
17817        x: &CudaSlice<f32>,
17818        y: &mut CudaSlice<f32>,
17819        block_cols: usize,
17820        out_f: usize,
17821    ) -> Result<(), Box<dyn std::error::Error>> {
17822        if block_cols % 4 != 0
17823            || x.len() < 4 * block_cols
17824            || y.len() < out_f
17825            || w.iter().any(|w| w.len() != out_f * block_cols)
17826        {
17827            return Err(format!(
17828                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17829                x.len()
17830            )
17831            .into());
17832        }
17833        let f = self.func("matvec_f32_b4");
17834        let cfg = LaunchConfig {
17835            grid_dim: (out_f as u32, 1, 1),
17836            block_dim: (128, 1, 1),
17837            shared_mem_bytes: 0,
17838        };
17839        let (bc, of) = (block_cols as i32, out_f as i32);
17840        let __s_b = self.gpu.stream();
17841        let mut b = __s_b.launch_builder(&f);
17842        b.arg(w[0])
17843            .arg(w[1])
17844            .arg(w[2])
17845            .arg(w[3])
17846            .arg(x)
17847            .arg(y)
17848            .arg(&bc)
17849            .arg(&of);
17850        unsafe {
17851            b.launch(cfg)?;
17852        }
17853        Ok(())
17854    }
17855
17856    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
17857    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
17858    pub fn axpy_rows_seq_into(
17859        &self,
17860        x: &CudaSlice<f32>,
17861        w: &CudaSlice<f32>,
17862        y: &mut CudaSlice<f32>,
17863        width: usize,
17864        n_rows: usize,
17865    ) -> Result<(), Box<dyn std::error::Error>> {
17866        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17867            return Err(format!(
17868                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17869                x.len(),
17870                w.len(),
17871                y.len()
17872            )
17873            .into());
17874        }
17875        let f = self.func("axpy_rows_seq_f32");
17876        let cfg = LaunchConfig::for_num_elems(width as u32);
17877        let (wi, nr) = (width as i32, n_rows as i32);
17878        let __s_b = self.gpu.stream();
17879        let mut b = __s_b.launch_builder(&f);
17880        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17881        unsafe {
17882            b.launch(cfg)?;
17883        }
17884        Ok(())
17885    }
17886
17887    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
17888    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
17889    /// exact sequential FP chain of the base kernel over that window.
17890    #[allow(clippy::too_many_arguments)]
17891    pub fn axpy_rows_seq_md_off_into(
17892        &self,
17893        x: &CudaSlice<f32>,
17894        w_route: &CudaSlice<f32>,
17895        md: &CudaSlice<f32>,
17896        sel: &CudaSlice<i32>,
17897        y: &mut CudaSlice<f32>,
17898        width: usize,
17899        n_rows: usize,
17900        row0: usize,
17901    ) -> Result<(), Box<dyn std::error::Error>> {
17902        if x.len() < (row0 + n_rows) * width
17903            || w_route.len() < row0 + n_rows
17904            || sel.len() < row0 + n_rows
17905            || y.len() < width
17906        {
17907            return Err(format!(
17908                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17909                 rows={n_rows} row0={row0}",
17910                x.len(),
17911                w_route.len(),
17912                sel.len(),
17913                y.len()
17914            )
17915            .into());
17916        }
17917        let f = self.func("axpy_rows_seq_md_off_f32");
17918        let cfg = LaunchConfig::for_num_elems(width as u32);
17919        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17920        let __s_b = self.gpu.stream();
17921        let mut b = __s_b.launch_builder(&f);
17922        b.arg(x)
17923            .arg(w_route)
17924            .arg(md)
17925            .arg(sel)
17926            .arg(y)
17927            .arg(&wi)
17928            .arg(&nr)
17929            .arg(&r0);
17930        unsafe {
17931            b.launch(cfg)?;
17932        }
17933        Ok(())
17934    }
17935
17936    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
17937    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
17938    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
17939    /// outputs are bit-equal to its own t=1 launch.
17940    #[allow(clippy::too_many_arguments)]
17941    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17942        &self,
17943        gate_bank: &CudaSlice<u8>,
17944        up_bank: &CudaSlice<u8>,
17945        sel: &CudaSlice<i32>,
17946        aq: &CudaSlice<i8>,
17947        ad: &CudaSlice<f32>,
17948        yg: &mut CudaSlice<f32>,
17949        yu: &mut CudaSlice<f32>,
17950        n_sel: usize,
17951        n_sel_col: usize,
17952        in_f: usize,
17953        out_f: usize,
17954        row_bytes: usize,
17955        expert_stride: usize,
17956        act_row_stride: usize,
17957        ad_row_stride: usize,
17958    ) -> Result<(), Box<dyn std::error::Error>> {
17959        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17960        if yg.len() < n_sel * out_f
17961            || yu.len() < n_sel * out_f
17962            || sel.len() < n_sel
17963            || n_sel_col == 0
17964            || n_sel % n_sel_col != 0
17965        {
17966            return Err("NVFP4 gu tcol geometry".into());
17967        }
17968        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17969        let cfg = LaunchConfig {
17970            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17971            block_dim: (128, 1, 1),
17972            shared_mem_bytes: 0,
17973        };
17974        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17975        let (rb, es) = (row_bytes as i64, expert_stride as i64);
17976        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17977        let __s_b = self.gpu.stream();
17978        let mut b = __s_b.launch_builder(&f);
17979        b.arg(gate_bank)
17980            .arg(up_bank)
17981            .arg(sel)
17982            .arg(aq)
17983            .arg(ad)
17984            .arg(yg)
17985            .arg(yu)
17986            .arg(&inf)
17987            .arg(&outf)
17988            .arg(&ns)
17989            .arg(&rb)
17990            .arg(&es)
17991            .arg(&ars)
17992            .arg(&adrs)
17993            .arg(&nsc);
17994        unsafe {
17995            b.launch(cfg)?;
17996        }
17997        Ok(())
17998    }
17999
18000    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
18001    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
18002    #[allow(clippy::too_many_arguments)]
18003    pub fn axpy_rows_seq_md_into(
18004        &self,
18005        x: &CudaSlice<f32>,
18006        w_route: &CudaSlice<f32>,
18007        md: &CudaSlice<f32>,
18008        sel: &CudaSlice<i32>,
18009        y: &mut CudaSlice<f32>,
18010        width: usize,
18011        n_rows: usize,
18012    ) -> Result<(), Box<dyn std::error::Error>> {
18013        if x.len() < n_rows * width
18014            || w_route.len() < n_rows
18015            || sel.len() < n_rows
18016            || y.len() < width
18017        {
18018            return Err(format!(
18019                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
18020                x.len(),
18021                w_route.len(),
18022                sel.len(),
18023                y.len()
18024            )
18025            .into());
18026        }
18027        let f = self.func("axpy_rows_seq_md_f32");
18028        let cfg = LaunchConfig::for_num_elems(width as u32);
18029        let (wi, nr) = (width as i32, n_rows as i32);
18030        let __s_b = self.gpu.stream();
18031        let mut b = __s_b.launch_builder(&f);
18032        b.arg(x)
18033            .arg(w_route)
18034            .arg(md)
18035            .arg(sel)
18036            .arg(y)
18037            .arg(&wi)
18038            .arg(&nr);
18039        unsafe {
18040            b.launch(cfg)?;
18041        }
18042        Ok(())
18043    }
18044
18045    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
18046    #[allow(clippy::too_many_arguments)]
18047    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
18048    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
18049    /// land column-major-of-rows: yq[c*out_q + row] etc.
18050    #[allow(clippy::too_many_arguments)]
18051    pub fn matvec_bf16_qkvg_tcol_into(
18052        &self,
18053        wq: &CudaSlice<u8>,
18054        wk: &CudaSlice<u8>,
18055        wv: &CudaSlice<u8>,
18056        wg: &CudaSlice<u8>,
18057        x_t: &CudaSlice<f32>,
18058        yq: &mut CudaSlice<f32>,
18059        yk: &mut CudaSlice<f32>,
18060        yv: &mut CudaSlice<f32>,
18061        yg: &mut CudaSlice<f32>,
18062        in_f: usize,
18063        out_q: usize,
18064        out_kv: usize,
18065        out_g: usize,
18066        t: usize,
18067    ) -> Result<(), Box<dyn std::error::Error>> {
18068        if t == 0
18069            || t > 8
18070            || in_f % 8 != 0
18071            || x_t.len() < t * in_f
18072            || yq.len() < t * out_q
18073            || yk.len() < t * out_kv
18074            || yv.len() < t * out_kv
18075            || (out_g > 0 && yg.len() < t * out_g)
18076        {
18077            return Err("matvec_bf16_qkvg_tcol geometry".into());
18078        }
18079        let grid = out_q + 2 * out_kv + out_g;
18080        let cfg = LaunchConfig {
18081            grid_dim: (grid as u32, 1, 1),
18082            block_dim: (mmv_block(), 1, 1),
18083            shared_mem_bytes: 0,
18084        };
18085        let (ini, oq, okv, og, ti) = (
18086            in_f as i32,
18087            out_q as i32,
18088            out_kv as i32,
18089            out_g as i32,
18090            t as i32,
18091        );
18092        let __s_b = self.gpu.stream();
18093        // One runtime-T program for every live width. The compile-time 2/4/8 twins are
18094        // retained in the fatbin as research controls, but dispatching them by the current
18095        // batch width changes kernels inside a request when peers arrive or retire. That is
18096        // a load-history numeric-program switch, and their pre-twin TOKFP receipts did not
18097        // qualify it (Hermes `64fa2b55baf0d887`).
18098        let f = self.func("matvec_bf16_qkvg_tcol");
18099        let mut b = __s_b.launch_builder(&f);
18100        b.arg(wq)
18101            .arg(wk)
18102            .arg(wv)
18103            .arg(wg)
18104            .arg(x_t)
18105            .arg(yq)
18106            .arg(yk)
18107            .arg(yv)
18108            .arg(yg)
18109            .arg(&ini)
18110            .arg(&oq)
18111            .arg(&okv)
18112            .arg(&og)
18113            .arg(&ti);
18114        unsafe {
18115            b.launch(cfg)?;
18116        }
18117        Ok(())
18118    }
18119
18120    pub fn matvec_bf16_qkvg_into(
18121        &self,
18122        wq: &CudaSlice<u8>,
18123        wk: &CudaSlice<u8>,
18124        wv: &CudaSlice<u8>,
18125        wg: &CudaSlice<u8>,
18126        x: &CudaSlice<f32>,
18127        yq: &mut CudaSlice<f32>,
18128        yk: &mut CudaSlice<f32>,
18129        yv: &mut CudaSlice<f32>,
18130        yg: &mut CudaSlice<f32>,
18131        in_f: usize,
18132        out_q: usize,
18133        out_kv: usize,
18134        out_g: usize,
18135    ) -> Result<(), Box<dyn std::error::Error>> {
18136        if in_f % 8 != 0
18137            || wq.len() != out_q * in_f * 2
18138            || wk.len() != out_kv * in_f * 2
18139            || wv.len() != out_kv * in_f * 2
18140            || wg.len() < out_g * in_f * 2
18141            || x.len() < in_f
18142            || yq.len() < out_q
18143            || yk.len() < out_kv
18144            || yv.len() < out_kv
18145            || (out_g > 0 && yg.len() < out_g)
18146        {
18147            return Err(format!(
18148                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18149            )
18150            .into());
18151        }
18152        let f = self.func("matvec_bf16_qkvg");
18153        let cfg = LaunchConfig {
18154            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18155            block_dim: (mmv_block(), 1, 1),
18156            shared_mem_bytes: 0,
18157        };
18158        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18159        let __s_b = self.gpu.stream();
18160        let mut b = __s_b.launch_builder(&f);
18161        b.arg(wq)
18162            .arg(wk)
18163            .arg(wv)
18164            .arg(wg)
18165            .arg(x)
18166            .arg(yq)
18167            .arg(yk)
18168            .arg(yv)
18169            .arg(yg)
18170            .arg(&inf)
18171            .arg(&oq)
18172            .arg(&okv)
18173            .arg(&og);
18174        unsafe {
18175            b.launch(cfg)?;
18176        }
18177        Ok(())
18178    }
18179
18180    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
18181    pub fn matvec_bf16_b4_into(
18182        &self,
18183        w: [&CudaSlice<u8>; 4],
18184        x: &CudaSlice<f32>,
18185        y: &mut CudaSlice<f32>,
18186        block_cols: usize,
18187        out_f: usize,
18188    ) -> Result<(), Box<dyn std::error::Error>> {
18189        if block_cols % 8 != 0
18190            || x.len() < 4 * block_cols
18191            || y.len() < out_f
18192            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18193        {
18194            return Err(format!(
18195                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18196                x.len()
18197            )
18198            .into());
18199        }
18200        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
18201        // bit-identical per row (the second row's stream hides the first's reduce tail).
18202        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18203        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18204        let f = self.func(if x2 {
18205            "matvec_bf16_b4_x2"
18206        } else {
18207            "matvec_bf16_b4"
18208        });
18209        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18210        let cfg = LaunchConfig {
18211            grid_dim: (grid as u32, 1, 1),
18212            block_dim: (mmv_block(), 1, 1),
18213            shared_mem_bytes: 0,
18214        };
18215        let (bc, of) = (block_cols as i32, out_f as i32);
18216        let __s_b = self.gpu.stream();
18217        let mut b = __s_b.launch_builder(&f);
18218        b.arg(w[0])
18219            .arg(w[1])
18220            .arg(w[2])
18221            .arg(w[3])
18222            .arg(x)
18223            .arg(y)
18224            .arg(&bc)
18225            .arg(&of);
18226        unsafe {
18227            b.launch(cfg)?;
18228        }
18229        Ok(())
18230    }
18231
18232    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
18233    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
18234    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
18235    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
18236    /// t=1 program).
18237    pub fn matvec_bf16_b4_tcol_into(
18238        &self,
18239        w: [&CudaSlice<u8>; 4],
18240        x_t: &CudaSlice<f32>,
18241        y_t: &mut CudaSlice<f32>,
18242        block_cols: usize,
18243        out_f: usize,
18244        t: usize,
18245    ) -> Result<(), Box<dyn std::error::Error>> {
18246        if block_cols % 8 != 0
18247            || t == 0
18248            || t > 8
18249            || x_t.len() < t * 4 * block_cols
18250            || y_t.len() < t * out_f
18251            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18252        {
18253            return Err(format!(
18254                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18255                x_t.len()
18256            )
18257            .into());
18258        }
18259        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18260            return Err(
18261                "b4 tcol verify is qualified against the plain b4 kernel only \
18262                        (MEMRA_B4_X2=1 is a different t=1 program)"
18263                    .into(),
18264            );
18265        }
18266        // Keep one runtime-T program at every live width. Compile-time twins remain research
18267        // controls only; selecting them from the changing batch width switches programs
18268        // mid-request.
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 (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18275        let __s_b = self.gpu.stream();
18276        let f = self.func("matvec_bf16_b4_tcol");
18277        let mut b = __s_b.launch_builder(&f);
18278        b.arg(w[0])
18279            .arg(w[1])
18280            .arg(w[2])
18281            .arg(w[3])
18282            .arg(x_t)
18283            .arg(y_t)
18284            .arg(&bc)
18285            .arg(&of)
18286            .arg(&ti);
18287        unsafe {
18288            b.launch(cfg)?;
18289        }
18290        Ok(())
18291    }
18292
18293    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
18294    /// q8_0 row bytes for an `in_f`-wide weight row: one 34-byte block per 32 weights.
18295    pub fn q8_0_row_bytes(in_f: usize) -> usize {
18296        in_f / 32 * 34
18297    }
18298
18299    /// Encode a resident bf16 weight slab into its q8_0 mirror (MEMRA_STEP_TP_W8). Runs once
18300    /// per matrix at load; the block program is the one `quant_K_block` writes for the KV
18301    /// cache, so the two formats cannot drift apart.
18302    pub fn encode_q8_0_from_bf16(
18303        &self,
18304        w_bf16: &CudaSlice<u8>,
18305        out: &mut CudaSlice<u8>,
18306        in_f: usize,
18307        out_f: usize,
18308    ) -> Result<(), Box<dyn std::error::Error>> {
18309        if in_f % 32 != 0
18310            || w_bf16.len() < in_f * out_f * 2
18311            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18312        {
18313            return Err(format!(
18314                "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18315                w_bf16.len(),
18316                out.len()
18317            )
18318            .into());
18319        }
18320        let f = self.func("encode_q8_0_rows_from_bf16");
18321        // Flat 1D grid of (row, 32-block) pairs, 4 pairs per block: rows on grid.y would cap
18322        // at 65535 and the LM head has 128896 rows.
18323        const PAIRS_PER_BLOCK: u32 = 4;
18324        let pairs = (out_f * (in_f / 32)) as u64;
18325        let cfg = LaunchConfig {
18326            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18327            block_dim: (32, PAIRS_PER_BLOCK, 1),
18328            shared_mem_bytes: 0,
18329        };
18330        let (ini, outi) = (in_f as i32, out_f as i32);
18331        let __s_b = self.gpu.stream();
18332        let mut b = __s_b.launch_builder(&f);
18333        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18334        unsafe {
18335            b.launch(cfg)?;
18336        }
18337        Ok(())
18338    }
18339
18340    /// Fused q8_0 QKV against a q8_1 activation (MEMRA_STEP_TP_W8): one launch over the
18341    /// stacked q/k/v rows, each row running the exact `qmatvec_q8_0_mmvq_rp` per-row program.
18342    /// Bit-identical to three per-matrix mmvq calls; it exists because those three launches
18343    /// plus the activation quantize measured SLOWER than the bf16 fused kernel.
18344    #[allow(clippy::too_many_arguments)]
18345    pub fn qmatvec_q8_0_qkv_rp_into(
18346        &self,
18347        wq: &CudaSlice<u8>,
18348        wk: &CudaSlice<u8>,
18349        wv: &CudaSlice<u8>,
18350        aq: &CudaSlice<i8>,
18351        ad: &CudaSlice<f32>,
18352        yq: &mut CudaSlice<f32>,
18353        yk: &mut CudaSlice<f32>,
18354        yv: &mut CudaSlice<f32>,
18355        in_f: usize,
18356        out_q: usize,
18357        out_kv: usize,
18358    ) -> Result<(), Box<dyn std::error::Error>> {
18359        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18360        let rows = out_q + 2 * out_kv;
18361        let nblk = in_f / 32;
18362        if in_f % 32 != 0
18363            || aq.len() < in_f
18364            || ad.len() < nblk
18365            || yq.len() < out_q
18366            || yk.len() < out_kv
18367            || yv.len() < out_kv
18368            || wq.len() < out_q * nblk * 34
18369            || wk.len() < out_kv * nblk * 34
18370            || wv.len() < out_kv * nblk * 34
18371        {
18372            return Err(
18373                format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18374            );
18375        }
18376        let f = self.func("qmatvec_q8_0_qkv_rp");
18377        let cfg = LaunchConfig {
18378            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18379            block_dim: (32, ROWS_PER_BLOCK, 1),
18380            shared_mem_bytes: 0,
18381        };
18382        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18383        let __s_b = self.gpu.stream();
18384        let mut b = __s_b.launch_builder(&f);
18385        b.arg(wq)
18386            .arg(wk)
18387            .arg(wv)
18388            .arg(aq)
18389            .arg(ad)
18390            .arg(yq)
18391            .arg(yk)
18392            .arg(yv)
18393            .arg(&ini)
18394            .arg(&oq)
18395            .arg(&okv);
18396        unsafe {
18397            b.launch(cfg)?;
18398        }
18399        Ok(())
18400    }
18401
18402    /// Fused q8_0 O projection over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8): one
18403    /// launch, one warp per output row, per-block reduce then add — the same shape
18404    /// `matvec_bf16_b4` uses, against a q8_1 activation.
18405    #[allow(clippy::too_many_arguments)]
18406    pub fn qmatvec_q8_0_b4_rp_into(
18407        &self,
18408        w: [&CudaSlice<u8>; 4],
18409        aq: &CudaSlice<i8>,
18410        ad: &CudaSlice<f32>,
18411        y: &mut CudaSlice<f32>,
18412        block_cols: usize,
18413        out_f: usize,
18414    ) -> Result<(), Box<dyn std::error::Error>> {
18415        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18416        let nblk = block_cols / 32;
18417        if block_cols % 32 != 0
18418            || aq.len() < 4 * block_cols
18419            || ad.len() < 4 * nblk
18420            || y.len() < out_f
18421            || w.iter().any(|p| p.len() < out_f * nblk * 34)
18422        {
18423            return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18424        }
18425        let f = self.func("qmatvec_q8_0_b4_rp");
18426        let cfg = LaunchConfig {
18427            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18428            block_dim: (32, ROWS_PER_BLOCK, 1),
18429            shared_mem_bytes: 0,
18430        };
18431        let (bc, of) = (block_cols as i32, out_f as i32);
18432        let __s_b = self.gpu.stream();
18433        let mut b = __s_b.launch_builder(&f);
18434        b.arg(w[0])
18435            .arg(w[1])
18436            .arg(w[2])
18437            .arg(w[3])
18438            .arg(aq)
18439            .arg(ad)
18440            .arg(y)
18441            .arg(&bc)
18442            .arg(&of);
18443        unsafe {
18444            b.launch(cfg)?;
18445        }
18446        Ok(())
18447    }
18448
18449    /// Get-or-build this bf16 weight's q8_0 mirror and run the GEMV through it. Returns
18450    /// `None` when the shape has no mirror form, so the caller falls back to bf16.
18451    fn matvec_bf16_via_q8_mirror(
18452        &self,
18453        data: &CudaSlice<u8>,
18454        x: &CudaSlice<f32>,
18455        y: &mut CudaSlice<f32>,
18456        in_f: usize,
18457        out_f: usize,
18458    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18459        use cudarc::driver::DevicePtr;
18460        let key = {
18461            let s = self.gpu.stream();
18462            let (p, _g) = data.device_ptr(&s);
18463            p as u64
18464        };
18465        {
18466            let mut mirrors = self
18467                .w8_mirrors
18468                .lock()
18469                .map_err(|_| "w8 mirror map is poisoned")?;
18470            if !mirrors.contains_key(&key) {
18471                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18472                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18473                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18474                mirrors.insert(key, planar);
18475                // Which weights this half actually covers is not obvious from the call graph:
18476                // the head and the shared expert may reach the GPU through the rows fast path
18477                // or the fused dual-silu launcher instead of here. One line per mirror answers
18478                // that without a profiler (the hybrid half measured +0.1% and this is how we
18479                // find out whether it even fired).
18480                if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18481                    eprintln!(
18482                        "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18483                        mirrors.len()
18484                    );
18485                }
18486            }
18487        }
18488        let nblk = in_f / 32;
18489        {
18490            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18491            if !act.contains_key(&in_f) {
18492                let aq = self.alloc_uninit::<i8>(in_f)?;
18493                let ad = self.alloc_uninit::<f32>(nblk)?;
18494                act.insert(in_f, (aq, ad));
18495            }
18496            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18497            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18498        }
18499        let mirrors = self
18500            .w8_mirrors
18501            .lock()
18502            .map_err(|_| "w8 mirror map is poisoned")?;
18503        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18504        let mirror = mirrors.get(&key).expect("built above");
18505        let (aq, ad) = act.get(&in_f).expect("built above");
18506        self.qmatvec_mmvq_into(
18507            mirror,
18508            aq,
18509            ad,
18510            1,
18511            in_f,
18512            out_f,
18513            QT_Q8_0,
18514            Self::q8_0_row_bytes(in_f),
18515            1.0,
18516            true,
18517            y,
18518        )?;
18519        Ok(Some(()))
18520    }
18521
18522    pub fn matvec_bf16_into(
18523        &self,
18524        data: &CudaSlice<u8>,
18525        x: &CudaSlice<f32>,
18526        y: &mut CudaSlice<f32>,
18527        in_f: usize,
18528        out_f: usize,
18529    ) -> Result<(), Box<dyn std::error::Error>> {
18530        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18531            return Err(format!(
18532                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18533                data.len(),
18534                x.len(),
18535                y.len()
18536            )
18537            .into());
18538        }
18539        // MEMRA_STEP_TP_W8, hybrid half: route this GEMV through a q8_0 mirror of the same
18540        // weight. Covers exactly the bf16 GEMVs that are NOT in a TP resident bank — the LM
18541        // head (324.4 -> 163.7 us measured), the shared-expert down rows (13.0 -> 5.6 us) and
18542        // the dense-FFN layers. Same numeric class as the QKV/o_proj arms (int8 dp4a with
18543        // per-32 scales), so it rides the same argmax acceptance; the bf16 slab stays resident
18544        // for prefill. The mirror builds on first use and is keyed by the slab's pointer.
18545        if step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18546            if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
18547                return Ok(());
18548            }
18549        }
18550        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
18551        // block, exact f32acc per-row program — cures the 1-iteration latency
18552        // starvation (shexp down measured 420GB/s at in_f=1280).
18553        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18554        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18555            && in_f <= 2048;
18556        if x4 {
18557            let f = self.func("matvec_bf16_f32acc_x4");
18558            let cfg = LaunchConfig {
18559                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18560                block_dim: (mmv_block(), 1, 1),
18561                shared_mem_bytes: 0,
18562            };
18563            let (ini, outi) = (in_f as i32, out_f as i32);
18564            let __s_b = self.gpu.stream();
18565            let mut b = __s_b.launch_builder(&f);
18566            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18567            unsafe {
18568                b.launch(cfg)?;
18569            }
18570            return Ok(());
18571        }
18572        let f = self.func("matvec_bf16_f32acc");
18573        let cfg = LaunchConfig {
18574            grid_dim: (out_f as u32, 1, 1),
18575            block_dim: (mmv_block(), 1, 1),
18576            shared_mem_bytes: 0,
18577        };
18578        let ini = in_f as i32;
18579        let __s_b = self.gpu.stream();
18580        let mut b = __s_b.launch_builder(&f);
18581        b.arg(data).arg(x).arg(y).arg(&ini);
18582        unsafe {
18583            b.launch(cfg)?;
18584        }
18585        Ok(())
18586    }
18587
18588    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
18589    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
18590    pub fn matvec_bf16_view_into(
18591        &self,
18592        data: &cudarc::driver::CudaView<'_, u8>,
18593        x: &CudaSlice<f32>,
18594        y: &mut CudaSlice<f32>,
18595        in_f: usize,
18596        out_f: usize,
18597    ) -> Result<(), Box<dyn std::error::Error>> {
18598        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18599            return Err(format!(
18600                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18601                data.len(),
18602                x.len(),
18603                y.len()
18604            )
18605            .into());
18606        }
18607        let f = self.func("matvec_bf16_f32acc");
18608        let cfg = LaunchConfig {
18609            grid_dim: (out_f as u32, 1, 1),
18610            block_dim: (mmv_block(), 1, 1),
18611            shared_mem_bytes: 0,
18612        };
18613        let ini = in_f as i32;
18614        let __s_b = self.gpu.stream();
18615        let mut b = __s_b.launch_builder(&f);
18616        b.arg(data).arg(x).arg(y).arg(&ini);
18617        unsafe {
18618            b.launch(cfg)?;
18619        }
18620        Ok(())
18621    }
18622
18623    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
18624    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
18625    pub fn matvec_bf16_raw_out(
18626        &self,
18627        w: &CudaSlice<u8>,
18628        x: &CudaSlice<f32>,
18629        y_raw: u64,
18630        in_f: usize,
18631        out_f: usize,
18632    ) -> Result<(), Box<dyn std::error::Error>> {
18633        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18634            return Err("matvec_bf16_raw_out geometry".into());
18635        }
18636        let f = self.func("matvec_bf16_f32acc");
18637        let cfg = LaunchConfig {
18638            grid_dim: (out_f as u32, 1, 1),
18639            block_dim: (mmv_block(), 1, 1),
18640            shared_mem_bytes: 0,
18641        };
18642        let ini = in_f as i32;
18643        let __s_b = self.gpu.stream();
18644        let mut b = __s_b.launch_builder(&f);
18645        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18646        unsafe {
18647            b.launch(cfg)?;
18648        }
18649        Ok(())
18650    }
18651
18652    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
18653    /// UVA pointers so the caller passes persistent-static rows without holding locks).
18654    /// Exact per-element sequence of the split add + add_scaled_rows pair.
18655    pub fn add3_raw(
18656        &self,
18657        a: &CudaSlice<f32>,
18658        b: &CudaSlice<f32>,
18659        sh_raw: u64,
18660        scale_raw: u64,
18661        dst: &mut CudaSlice<f32>,
18662        n: usize,
18663    ) -> Result<(), Box<dyn std::error::Error>> {
18664        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18665            return Err("add3_raw geometry".into());
18666        }
18667        let f = self.func("add3_f32");
18668        let cfg = LaunchConfig {
18669            grid_dim: ((n as u32).div_ceil(256), 1, 1),
18670            block_dim: (256, 1, 1),
18671            shared_mem_bytes: 0,
18672        };
18673        let ni = n as i32;
18674        let __s_b = self.gpu.stream();
18675        let mut bld = __s_b.launch_builder(&f);
18676        bld.arg(a)
18677            .arg(b)
18678            .arg(&sh_raw)
18679            .arg(&scale_raw)
18680            .arg(dst)
18681            .arg(&ni);
18682        unsafe {
18683            bld.launch(cfg)?;
18684        }
18685        Ok(())
18686    }
18687
18688    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
18689    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
18690    pub fn matvec_bf16_down_addscale_into(
18691        &self,
18692        w: &CudaSlice<u8>,
18693        x: &CudaSlice<f32>,
18694        scale: &CudaSlice<f32>,
18695        dst: &mut CudaSlice<f32>,
18696        in_f: usize,
18697        out_f: usize,
18698    ) -> Result<(), Box<dyn std::error::Error>> {
18699        if w.len() != in_f * out_f * 2
18700            || x.len() < in_f
18701            || in_f % 8 != 0
18702            || dst.len() < out_f
18703            || scale.is_empty()
18704        {
18705            return Err("matvec_bf16_down_addscale geometry".into());
18706        }
18707        let f = self.func("matvec_bf16_down_addscale");
18708        let cfg = LaunchConfig {
18709            grid_dim: (out_f as u32, 1, 1),
18710            block_dim: (mmv_block(), 1, 1),
18711            shared_mem_bytes: 0,
18712        };
18713        let ini = in_f as i32;
18714        let __s_b = self.gpu.stream();
18715        let mut b = __s_b.launch_builder(&f);
18716        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18717        unsafe {
18718            b.launch(cfg)?;
18719        }
18720        Ok(())
18721    }
18722
18723    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
18724    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
18725    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
18726    #[allow(clippy::too_many_arguments)]
18727    pub fn matvec_bf16_dual_silu_rows_into(
18728        &self,
18729        wg: &CudaSlice<u8>,
18730        wu: &CudaSlice<u8>,
18731        x: &CudaSlice<f32>,
18732        act: &mut CudaSlice<f32>,
18733        in_f: usize,
18734        out_f: usize,
18735        limit: Option<f32>,
18736        t: usize,
18737    ) -> Result<(), Box<dyn std::error::Error>> {
18738        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18739            return Err("matvec_bf16_dual_silu_rows geometry".into());
18740        }
18741        let f = self.func("matvec_bf16_dual_silu_rows");
18742        let cfg = LaunchConfig {
18743            grid_dim: (out_f as u32, t as u32, 1),
18744            block_dim: (mmv_block(), 1, 1),
18745            shared_mem_bytes: 0,
18746        };
18747        let (ini, outi) = (in_f as i32, out_f as i32);
18748        let lim = limit.unwrap_or(0.0);
18749        let __s_b = self.gpu.stream();
18750        let mut b = __s_b.launch_builder(&f);
18751        b.arg(wg)
18752            .arg(wu)
18753            .arg(x)
18754            .arg(&mut *act)
18755            .arg(&ini)
18756            .arg(&outi)
18757            .arg(&lim);
18758        unsafe {
18759            b.launch(cfg)?;
18760        }
18761        Ok(())
18762    }
18763
18764    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
18765    pub fn matvec_bf16_rows_into(
18766        &self,
18767        w: &CudaSlice<u8>,
18768        x: &CudaSlice<f32>,
18769        y: &mut CudaSlice<f32>,
18770        in_f: usize,
18771        out_f: usize,
18772        t: usize,
18773    ) -> Result<(), Box<dyn std::error::Error>> {
18774        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18775            return Err("matvec_bf16_rows geometry".into());
18776        }
18777        // MEMRA_STEP_TP_W8: the LM head reaches the device HERE, not through
18778        // matvec_bf16_into — the W8 trace showed the hybrid half building mirrors only for
18779        // in_f=1280 out_f=4096 (the shared-expert down rows, which SHEXP_OVERLAP already
18780        // hides, hence its +0.1%). Route the t=1 decode row through the q8 mirror; wider t
18781        // (the verify walk) keeps bf16 so the prefill class is untouched.
18782        if t == 1 && step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18783            if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
18784                return Ok(());
18785            }
18786        }
18787        let f = self.func("matvec_bf16_f32acc_x4_rows");
18788        let cfg = LaunchConfig {
18789            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18790            block_dim: (mmv_block(), 1, 1),
18791            shared_mem_bytes: 0,
18792        };
18793        let (ini, outi) = (in_f as i32, out_f as i32);
18794        let __s_b = self.gpu.stream();
18795        let mut b = __s_b.launch_builder(&f);
18796        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18797        unsafe {
18798            b.launch(cfg)?;
18799        }
18800        Ok(())
18801    }
18802
18803    pub fn matvec_bf16_dual_silu_into(
18804        &self,
18805        wg: &CudaSlice<u8>,
18806        wu: &CudaSlice<u8>,
18807        x: &CudaSlice<f32>,
18808        act: &mut CudaSlice<f32>,
18809        in_f: usize,
18810        out_f: usize,
18811        limit: Option<f32>,
18812    ) -> Result<(), Box<dyn std::error::Error>> {
18813        if wg.len() != in_f * out_f * 2
18814            || wu.len() != in_f * out_f * 2
18815            || x.len() < in_f
18816            || in_f % 8 != 0
18817            || act.len() < out_f
18818        {
18819            return Err("matvec_bf16_dual_silu geometry".into());
18820        }
18821        let f = self.func("matvec_bf16_dual_silu");
18822        let cfg = LaunchConfig {
18823            grid_dim: (out_f as u32, 1, 1),
18824            block_dim: (mmv_block(), 1, 1),
18825            shared_mem_bytes: 0,
18826        };
18827        let (ini, outi) = (in_f as i32, out_f as i32);
18828        let lim = limit.unwrap_or(0.0);
18829        let __s_b = self.gpu.stream();
18830        let mut b = __s_b.launch_builder(&f);
18831        b.arg(wg)
18832            .arg(wu)
18833            .arg(x)
18834            .arg(act)
18835            .arg(&ini)
18836            .arg(&outi)
18837            .arg(&lim);
18838        unsafe {
18839            b.launch(cfg)?;
18840        }
18841        Ok(())
18842    }
18843
18844    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
18845    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
18846    #[allow(clippy::too_many_arguments)]
18847    pub fn matvec_bf16_dual_view_into(
18848        &self,
18849        wg: &cudarc::driver::CudaView<'_, u8>,
18850        wu: &cudarc::driver::CudaView<'_, u8>,
18851        x: &CudaSlice<f32>,
18852        yg: &mut CudaSlice<f32>,
18853        yu: &mut CudaSlice<f32>,
18854        in_f: usize,
18855        out_f: usize,
18856    ) -> Result<(), Box<dyn std::error::Error>> {
18857        if wg.len() != in_f * out_f * 2
18858            || wu.len() != in_f * out_f * 2
18859            || x.len() < in_f
18860            || in_f % 8 != 0
18861            || yg.len() < out_f
18862            || yu.len() < out_f
18863        {
18864            return Err(format!(
18865                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18866                wg.len(),
18867                wu.len(),
18868                x.len()
18869            )
18870            .into());
18871        }
18872        let f = self.func("matvec_bf16_dual");
18873        let cfg = LaunchConfig {
18874            grid_dim: ((2 * out_f) as u32, 1, 1),
18875            block_dim: (mmv_block(), 1, 1),
18876            shared_mem_bytes: 0,
18877        };
18878        let (ini, outi) = (in_f as i32, out_f as i32);
18879        let __s_b = self.gpu.stream();
18880        let mut b = __s_b.launch_builder(&f);
18881        b.arg(wg)
18882            .arg(wu)
18883            .arg(x)
18884            .arg(yg)
18885            .arg(yu)
18886            .arg(&ini)
18887            .arg(&outi);
18888        unsafe {
18889            b.launch(cfg)?;
18890        }
18891        Ok(())
18892    }
18893
18894    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
18895    #[allow(clippy::too_many_arguments)]
18896    pub fn matvec_bf16_dual_into(
18897        &self,
18898        wg: &CudaSlice<u8>,
18899        wu: &CudaSlice<u8>,
18900        x: &CudaSlice<f32>,
18901        yg: &mut CudaSlice<f32>,
18902        yu: &mut CudaSlice<f32>,
18903        in_f: usize,
18904        out_f: usize,
18905    ) -> Result<(), Box<dyn std::error::Error>> {
18906        if wg.len() != in_f * out_f * 2
18907            || wu.len() != in_f * out_f * 2
18908            || x.len() < in_f
18909            || in_f % 8 != 0
18910            || yg.len() < out_f
18911            || yu.len() < out_f
18912        {
18913            return Err(format!(
18914                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18915                wg.len(),
18916                wu.len(),
18917                x.len()
18918            )
18919            .into());
18920        }
18921        let f = self.func("matvec_bf16_dual");
18922        let cfg = LaunchConfig {
18923            grid_dim: ((2 * out_f) as u32, 1, 1),
18924            block_dim: (mmv_block(), 1, 1),
18925            shared_mem_bytes: 0,
18926        };
18927        let (ini, outi) = (in_f as i32, out_f as i32);
18928        let __s_b = self.gpu.stream();
18929        let mut b = __s_b.launch_builder(&f);
18930        b.arg(wg)
18931            .arg(wu)
18932            .arg(x)
18933            .arg(yg)
18934            .arg(yu)
18935            .arg(&ini)
18936            .arg(&outi);
18937        unsafe {
18938            b.launch(cfg)?;
18939        }
18940        Ok(())
18941    }
18942
18943    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
18944    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
18945    pub(crate) fn matvec_bf16_dual(
18946        &self,
18947        wg: &CudaSlice<u8>,
18948        wu: &CudaSlice<u8>,
18949        x: &CudaSlice<f32>,
18950        in_f: usize,
18951        out_f: usize,
18952    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18953        if wg.len() != in_f * out_f * 2
18954            || wu.len() != in_f * out_f * 2
18955            || x.len() < in_f
18956            || in_f % 8 != 0
18957        {
18958            return Err(format!(
18959                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18960                wg.len(),
18961                wu.len(),
18962                x.len()
18963            )
18964            .into());
18965        }
18966        let mut yg = self.alloc_uninit::<f32>(out_f)?;
18967        let mut yu = self.alloc_uninit::<f32>(out_f)?;
18968        let f = self.func("matvec_bf16_dual");
18969        let cfg = LaunchConfig {
18970            grid_dim: ((2 * out_f) as u32, 1, 1),
18971            block_dim: (mmv_block(), 1, 1),
18972            shared_mem_bytes: 0,
18973        };
18974        let (ini, outi) = (in_f as i32, out_f as i32);
18975        let __s_b = self.gpu.stream();
18976        let mut b = __s_b.launch_builder(&f);
18977        b.arg(wg)
18978            .arg(wu)
18979            .arg(x)
18980            .arg(&mut yg)
18981            .arg(&mut yu)
18982            .arg(&ini)
18983            .arg(&outi);
18984        unsafe {
18985            b.launch(cfg)?;
18986        }
18987        Ok((yg, yu))
18988    }
18989
18990    #[allow(clippy::too_many_arguments)]
18991    fn linear_bf16_chunked_inner(
18992        &self,
18993        x: &CudaSlice<f32>,
18994        data: &CudaSlice<u8>,
18995        m: usize,
18996        in_f: usize,
18997        out_f: usize,
18998        exact: bool,
18999        canonical_chunk_rows: Option<usize>,
19000    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19001        const CHUNK_BYTES: usize = 256 << 20;
19002        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
19003        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
19004        if m == 1
19005            && !exact
19006            && canonical_chunk_rows.is_none()
19007            && in_f % 8 == 0
19008            && Self::bf16_mmv_on()
19009        {
19010            return self.matvec_bf16(data, x, in_f, out_f);
19011        }
19012        let row_bytes = in_f
19013            .checked_mul(std::mem::size_of::<f32>())
19014            .ok_or("BF16 chunk row byte count overflow")?;
19015        if row_bytes == 0 || out_f == 0 {
19016            return Err("BF16 chunk dimensions must be nonzero".into());
19017        }
19018        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
19019        let chunk_rows = match canonical_chunk_rows {
19020            Some(rows) if rows == 0 => {
19021                return Err("canonical BF16 chunk rows must be nonzero".into());
19022            }
19023            Some(rows) if rows > max_chunk_rows => {
19024                return Err(format!(
19025                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
19026                )
19027                .into());
19028            }
19029            Some(rows) if out_f % rows != 0 => {
19030                return Err(format!(
19031                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
19032                )
19033                .into());
19034            }
19035            Some(rows) => rows,
19036            None => max_chunk_rows,
19037        };
19038        if chunk_rows >= out_f {
19039            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
19040            return if exact {
19041                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19042            } else {
19043                self.linear(x, &wf32, m, in_f, out_f)
19044            };
19045        }
19046        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19047        let mut r0 = 0usize;
19048        while r0 < out_f {
19049            let rows = chunk_rows.min(out_f - r0);
19050            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19051            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19052            let yc = if exact {
19053                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19054            } else {
19055                self.linear(x, &wf32, m, in_f, rows)?
19056            };
19057            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
19058            for mi in 0..m {
19059                let src = yc.slice(mi * rows..(mi + 1) * rows);
19060                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19061                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19062            }
19063            r0 += rows;
19064        }
19065        Ok(y)
19066    }
19067
19068    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
19069    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
19070    /// chunked BF16 numerical program instead of re-encoding the weight.
19071    pub fn linear_bf16_resident(
19072        &self,
19073        x: &CudaSlice<f32>,
19074        data: &CudaSlice<u8>,
19075        m: usize,
19076        in_f: usize,
19077        out_f: usize,
19078    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19079        if data.len() != in_f * out_f * 2 {
19080            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19081        }
19082        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19083    }
19084
19085    /// Execute a resident BF16 projection as fixed-width output-row chunks.
19086    ///
19087    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
19088    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
19089    /// model topology rather than the active rank count.
19090    pub fn linear_bf16_resident_canonical_rows(
19091        &self,
19092        x: &CudaSlice<f32>,
19093        data: &CudaSlice<u8>,
19094        m: usize,
19095        in_f: usize,
19096        out_f: usize,
19097        canonical_chunk_rows: usize,
19098    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19099        if data.len() != in_f * out_f * 2 {
19100            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19101        }
19102        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19103    }
19104
19105    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
19106    ///
19107    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
19108    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
19109    pub fn linear_f32_resident_canonical_rows(
19110        &self,
19111        x: &CudaSlice<f32>,
19112        data: &CudaSlice<f32>,
19113        m: usize,
19114        in_f: usize,
19115        out_f: usize,
19116        canonical_chunk_rows: usize,
19117    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19118        self.linear_f32_resident_canonical_rows_inner(
19119            x,
19120            data,
19121            m,
19122            in_f,
19123            out_f,
19124            canonical_chunk_rows,
19125            false,
19126        )
19127    }
19128
19129    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
19130    ///
19131    /// The projection shapes and values are identical to
19132    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
19133    /// changes, replacing one device copy per token with one placement kernel per output chunk.
19134    pub fn linear_f32_resident_canonical_rows_strided(
19135        &self,
19136        x: &CudaSlice<f32>,
19137        data: &CudaSlice<f32>,
19138        m: usize,
19139        in_f: usize,
19140        out_f: usize,
19141        canonical_chunk_rows: usize,
19142    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19143        self.linear_f32_resident_canonical_rows_inner(
19144            x,
19145            data,
19146            m,
19147            in_f,
19148            out_f,
19149            canonical_chunk_rows,
19150            true,
19151        )
19152    }
19153
19154    fn linear_f32_resident_canonical_rows_inner(
19155        &self,
19156        x: &CudaSlice<f32>,
19157        data: &CudaSlice<f32>,
19158        m: usize,
19159        in_f: usize,
19160        out_f: usize,
19161        canonical_chunk_rows: usize,
19162        strided_output: bool,
19163    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19164        if data.len() != in_f * out_f {
19165            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19166        }
19167        if canonical_chunk_rows == 0
19168            || canonical_chunk_rows > out_f
19169            || out_f % canonical_chunk_rows != 0
19170        {
19171            return Err(format!(
19172                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19173            )
19174            .into());
19175        }
19176        if canonical_chunk_rows == out_f {
19177            return self.linear(x, data, m, in_f, out_f);
19178        }
19179
19180        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19181        let input = x.slice(0..x.len());
19182        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19183            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19184            if m == 1 {
19185                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19186                self.linear_device_into(
19187                    &input,
19188                    &weights,
19189                    &mut destination,
19190                    1,
19191                    in_f,
19192                    canonical_chunk_rows,
19193                )?;
19194                continue;
19195            }
19196            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19197            if strided_output {
19198                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19199            } else {
19200                for token in 0..m {
19201                    let source = chunk
19202                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19203                    let mut destination =
19204                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19205                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19206                }
19207            }
19208        }
19209        Ok(y)
19210    }
19211
19212    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
19213    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
19214    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
19215    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
19216    pub fn linear_f32_resident_canonical_rows_t1_into(
19217        &self,
19218        x: &CudaSlice<f32>,
19219        data: &CudaSlice<f32>,
19220        y: &mut CudaSlice<f32>,
19221        in_f: usize,
19222        out_f: usize,
19223        canonical_chunk_rows: usize,
19224    ) -> Result<(), Box<dyn std::error::Error>> {
19225        if data.len() != in_f * out_f {
19226            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19227        }
19228        if y.len() != out_f || x.len() != in_f {
19229            return Err(format!(
19230                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19231                x.len(),
19232                y.len()
19233            )
19234            .into());
19235        }
19236        if canonical_chunk_rows == 0
19237            || canonical_chunk_rows > out_f
19238            || out_f % canonical_chunk_rows != 0
19239        {
19240            return Err(format!(
19241                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19242            )
19243            .into());
19244        }
19245        let input = x.slice(0..x.len());
19246        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19247            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19248            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19249            self.linear_device_into(
19250                &input,
19251                &weights,
19252                &mut destination,
19253                1,
19254                in_f,
19255                canonical_chunk_rows,
19256            )?;
19257        }
19258        Ok(())
19259    }
19260
19261    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
19262    /// without the allocation, for workspace-resident operands.
19263    pub fn linear_t1_into(
19264        &self,
19265        x: &cudarc::driver::CudaView<'_, f32>,
19266        w: &cudarc::driver::CudaView<'_, f32>,
19267        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19268        in_f: usize,
19269        out_f: usize,
19270    ) -> Result<(), Box<dyn std::error::Error>> {
19271        self.linear_device_into(x, w, y, 1, in_f, out_f)
19272    }
19273
19274    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
19275    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
19276    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
19277    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
19278    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
19279    /// router/shexp sites and matmul_decode_exact's Float arm.
19280    pub fn linear_decode_exact(
19281        &self,
19282        x: &CudaSlice<f32>,
19283        w: &CudaSlice<f32>,
19284        m_tokens: usize,
19285        in_f: usize,
19286        out_f: usize,
19287    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19288        if m_tokens == 1 {
19289            return self.linear(x, w, 1, in_f, out_f);
19290        }
19291        let xv = self.view(x, m_tokens * in_f);
19292        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19293        for t in 0..m_tokens {
19294            let row = xv.slice(t * in_f..(t + 1) * in_f);
19295            let mut xr = self.alloc_uninit::<f32>(in_f)?;
19296            self.copy_view_into(&mut xr, 0, &row, in_f)?;
19297            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19298            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19299        }
19300        Ok(y)
19301    }
19302
19303    pub fn linear(
19304        &self,
19305        x: &CudaSlice<f32>,
19306        w: &CudaSlice<f32>,
19307        m_tokens: usize,
19308        in_f: usize,
19309        out_f: usize,
19310    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19311        self.linear_device(x, w, m_tokens, in_f, out_f)
19312    }
19313
19314    fn linear_device<I>(
19315        &self,
19316        x: &I,
19317        w: &I,
19318        m_tokens: usize,
19319        in_f: usize,
19320        out_f: usize,
19321    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19322    where
19323        I: cudarc::driver::DevicePtr<f32>,
19324    {
19325        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
19326        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
19327        Ok(c)
19328    }
19329
19330    fn linear_device_into<I, O>(
19331        &self,
19332        x: &I,
19333        w: &I,
19334        c: &mut O,
19335        m_tokens: usize,
19336        in_f: usize,
19337        out_f: usize,
19338    ) -> Result<(), Box<dyn std::error::Error>>
19339    where
19340        I: cudarc::driver::DevicePtr<f32>,
19341        O: cudarc::driver::DevicePtrMut<f32>,
19342    {
19343        use cudarc::cublaslt::{Matmul, MatmulConfig};
19344        let cfg = MatmulConfig {
19345            transa: true,
19346            transb: false,
19347            transc: false,
19348            m: out_f as u64,
19349            n: m_tokens as u64,
19350            k: in_f as u64,
19351            alpha: 1.0,
19352            lda: in_f as i64,
19353            ldb: in_f as i64,
19354            beta: 0.0,
19355            ldc: out_f as i64,
19356            stride_a: None,
19357            stride_b: None,
19358            stride_c: None,
19359            stride_bias: None,
19360            batch_size: None,
19361        };
19362        let blas = self.gpu.blas();
19363        unsafe {
19364            blas.matmul(cfg, w, x, c, None, None)?;
19365        }
19366        Ok(())
19367    }
19368
19369    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
19370    ///
19371    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
19372    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
19373    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
19374    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
19375    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
19376    /// launch error mid-request.
19377    pub fn sdpa_naive(
19378        &self,
19379        q: &CudaSlice<f32>,
19380        k: &CudaSlice<f32>,
19381        v: &CudaSlice<f32>,
19382        o: &mut CudaSlice<f32>,
19383        head_dim: usize,
19384        n_head: usize,
19385        n_head_kv: usize,
19386        t: usize,
19387        t_kv: usize,
19388        scale: f32,
19389        causal: bool,
19390    ) -> Result<(), Box<dyn std::error::Error>> {
19391        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19392            return self.sdpa_naive_gmem(
19393                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19394            );
19395        }
19396        let f = self.func("sdpa_naive_f32");
19397        let cfg = LaunchConfig {
19398            grid_dim: (n_head as u32, t as u32, 1),
19399            block_dim: (128, 1, 1),
19400            shared_mem_bytes: (t_kv * 4) as u32,
19401        };
19402        let (hd, nh, nhkv, ti, tkvi, cz) = (
19403            head_dim as i32,
19404            n_head as i32,
19405            n_head_kv as i32,
19406            t as i32,
19407            t_kv as i32,
19408            causal as i32,
19409        );
19410        let __s_b = self.gpu.stream();
19411        let mut b = __s_b.launch_builder(&f);
19412        b.arg(q)
19413            .arg(k)
19414            .arg(v)
19415            .arg(o)
19416            .arg(&hd)
19417            .arg(&nh)
19418            .arg(&nhkv)
19419            .arg(&ti)
19420            .arg(&tkvi)
19421            .arg(&scale)
19422            .arg(&cz);
19423        unsafe {
19424            b.launch(cfg)?;
19425        }
19426        Ok(())
19427    }
19428
19429    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
19430    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
19431    /// of dynamic shared memory: identical loop structure and reduction order, so the output
19432    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
19433    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
19434    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
19435    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
19436    /// T==T_kv caller cannot silently allocate tens of GB.
19437    #[allow(clippy::too_many_arguments)]
19438    pub fn sdpa_naive_gmem(
19439        &self,
19440        q: &CudaSlice<f32>,
19441        k: &CudaSlice<f32>,
19442        v: &CudaSlice<f32>,
19443        o: &mut CudaSlice<f32>,
19444        head_dim: usize,
19445        n_head: usize,
19446        n_head_kv: usize,
19447        t: usize,
19448        t_kv: usize,
19449        scale: f32,
19450        causal: bool,
19451    ) -> Result<(), Box<dyn std::error::Error>> {
19452        let ws_len = n_head
19453            .checked_mul(t)
19454            .and_then(|x| x.checked_mul(t_kv))
19455            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19456        let ws_bytes = ws_len
19457            .checked_mul(std::mem::size_of::<f32>())
19458            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19459        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19460            return Err(format!(
19461                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19462                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19463                 needs a tiled/flash kernel, not the naive oracle"
19464            )
19465            .into());
19466        }
19467        let mut scores = self.uninit(ws_len)?;
19468        let f = self.func("sdpa_naive_gmem_f32");
19469        let cfg = LaunchConfig {
19470            grid_dim: (n_head as u32, t as u32, 1),
19471            block_dim: (128, 1, 1),
19472            shared_mem_bytes: 0,
19473        };
19474        let (hd, nh, nhkv, ti, tkvi, cz) = (
19475            head_dim as i32,
19476            n_head as i32,
19477            n_head_kv as i32,
19478            t as i32,
19479            t_kv as i32,
19480            causal as i32,
19481        );
19482        let __s_b = self.gpu.stream();
19483        let mut b = __s_b.launch_builder(&f);
19484        b.arg(q)
19485            .arg(k)
19486            .arg(v)
19487            .arg(o)
19488            .arg(&mut scores)
19489            .arg(&hd)
19490            .arg(&nh)
19491            .arg(&nhkv)
19492            .arg(&ti)
19493            .arg(&tkvi)
19494            .arg(&scale)
19495            .arg(&cz);
19496        unsafe {
19497            b.launch(cfg)?;
19498        }
19499        Ok(())
19500    }
19501
19502    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
19503    /// bidirectional image islands. `span_id` labels each absolute kv position
19504    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
19505    /// reproducing the reference's non-causal image batch. window 0 = no window.
19506    #[allow(clippy::too_many_arguments)]
19507    pub fn sdpa_naive_island(
19508        &self,
19509        q: &CudaSlice<f32>,
19510        k: &CudaSlice<f32>,
19511        v: &CudaSlice<f32>,
19512        o: &mut CudaSlice<f32>,
19513        span_id: &CudaSlice<i32>,
19514        head_dim: usize,
19515        n_head: usize,
19516        n_head_kv: usize,
19517        t: usize,
19518        t_kv: usize,
19519        scale: f32,
19520        window: usize,
19521    ) -> Result<(), Box<dyn std::error::Error>> {
19522        let f = self.func("sdpa_naive_island_f32");
19523        let cfg = LaunchConfig {
19524            grid_dim: (n_head as u32, t as u32, 1),
19525            block_dim: (128, 1, 1),
19526            shared_mem_bytes: (t_kv * 4) as u32,
19527        };
19528        let (hd, nh, nhkv, ti, tkvi, wi) = (
19529            head_dim as i32,
19530            n_head as i32,
19531            n_head_kv as i32,
19532            t as i32,
19533            t_kv as i32,
19534            window as i32,
19535        );
19536        let __s_b = self.gpu.stream();
19537        let mut b = __s_b.launch_builder(&f);
19538        b.arg(q)
19539            .arg(k)
19540            .arg(v)
19541            .arg(o)
19542            .arg(span_id)
19543            .arg(&hd)
19544            .arg(&nh)
19545            .arg(&nhkv)
19546            .arg(&ti)
19547            .arg(&tkvi)
19548            .arg(&scale)
19549            .arg(&wi);
19550        unsafe {
19551            b.launch(cfg)?;
19552        }
19553        Ok(())
19554    }
19555
19556    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
19557    #[allow(clippy::too_many_arguments)]
19558    pub fn sdpa_naive_w(
19559        &self,
19560        q: &CudaSlice<f32>,
19561        k: &CudaSlice<f32>,
19562        v: &CudaSlice<f32>,
19563        o: &mut CudaSlice<f32>,
19564        head_dim: usize,
19565        n_head: usize,
19566        n_head_kv: usize,
19567        t: usize,
19568        t_kv: usize,
19569        scale: f32,
19570        causal: bool,
19571        window: usize,
19572    ) -> Result<(), Box<dyn std::error::Error>> {
19573        let f = self.func("sdpa_naive_w_f32");
19574        let cfg = LaunchConfig {
19575            grid_dim: (n_head as u32, t as u32, 1),
19576            block_dim: (128, 1, 1),
19577            shared_mem_bytes: (t_kv * 4) as u32,
19578        };
19579        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19580            head_dim as i32,
19581            n_head as i32,
19582            n_head_kv as i32,
19583            t as i32,
19584            t_kv as i32,
19585            causal as i32,
19586            window as i32,
19587        );
19588        let __s_b = self.gpu.stream();
19589        let mut b = __s_b.launch_builder(&f);
19590        b.arg(q)
19591            .arg(k)
19592            .arg(v)
19593            .arg(o)
19594            .arg(&hd)
19595            .arg(&nh)
19596            .arg(&nhkv)
19597            .arg(&ti)
19598            .arg(&tkvi)
19599            .arg(&scale)
19600            .arg(&cz)
19601            .arg(&wi);
19602        unsafe {
19603            b.launch(cfg)?;
19604        }
19605        Ok(())
19606    }
19607
19608    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
19609    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
19610    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
19611    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
19612    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
19613    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
19614    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
19615    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
19616    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
19617    #[allow(clippy::too_many_arguments)]
19618    pub fn sdpa_naive_w_lo(
19619        &self,
19620        q: &CudaSlice<f32>,
19621        k: &CudaSlice<f32>,
19622        v: &CudaSlice<f32>,
19623        o: &mut CudaSlice<f32>,
19624        head_dim: usize,
19625        n_head: usize,
19626        n_head_kv: usize,
19627        t: usize,
19628        t_kv: usize,
19629        scale: f32,
19630        causal: bool,
19631        window: usize,
19632    ) -> Result<(), Box<dyn std::error::Error>> {
19633        let kv_lo = if window > 0 {
19634            (t_kv - t + 1).saturating_sub(window)
19635        } else {
19636            0
19637        };
19638        let smem = (t_kv - kv_lo) * 4;
19639        if smem > 48 * 1024 {
19640            return Err(format!(
19641                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19642                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19643                 a window this wide needs the multi-pass long-ctx kernel"
19644            )
19645            .into());
19646        }
19647        let f = self.func("sdpa_naive_w_lo_f32");
19648        let cfg = LaunchConfig {
19649            grid_dim: (n_head as u32, t as u32, 1),
19650            block_dim: (128, 1, 1),
19651            shared_mem_bytes: smem as u32,
19652        };
19653        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19654            head_dim as i32,
19655            n_head as i32,
19656            n_head_kv as i32,
19657            t as i32,
19658            t_kv as i32,
19659            causal as i32,
19660            window as i32,
19661            kv_lo as i32,
19662        );
19663        let __s_b = self.gpu.stream();
19664        let mut b = __s_b.launch_builder(&f);
19665        b.arg(q)
19666            .arg(k)
19667            .arg(v)
19668            .arg(o)
19669            .arg(&hd)
19670            .arg(&nh)
19671            .arg(&nhkv)
19672            .arg(&ti)
19673            .arg(&tkvi)
19674            .arg(&scale)
19675            .arg(&cz)
19676            .arg(&wi)
19677            .arg(&lo);
19678        unsafe {
19679            b.launch(cfg)?;
19680        }
19681        Ok(())
19682    }
19683
19684    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
19685    pub fn sdpa_naive_view(
19686        &self,
19687        q: &CudaSlice<f32>,
19688        k: &cudarc::driver::CudaView<f32>,
19689        v: &cudarc::driver::CudaView<f32>,
19690        o: &mut CudaSlice<f32>,
19691        head_dim: usize,
19692        n_head: usize,
19693        n_head_kv: usize,
19694        t: usize,
19695        t_kv: usize,
19696        scale: f32,
19697        causal: bool,
19698    ) -> Result<(), Box<dyn std::error::Error>> {
19699        let f = self.func("sdpa_naive_f32");
19700        let cfg = LaunchConfig {
19701            grid_dim: (n_head as u32, t as u32, 1),
19702            block_dim: (128, 1, 1),
19703            shared_mem_bytes: (t_kv * 4) as u32,
19704        };
19705        let (hd, nh, nhkv, ti, tkvi, cz) = (
19706            head_dim as i32,
19707            n_head as i32,
19708            n_head_kv as i32,
19709            t as i32,
19710            t_kv as i32,
19711            causal as i32,
19712        );
19713        let __s_b = self.gpu.stream();
19714        let mut b = __s_b.launch_builder(&f);
19715        b.arg(q)
19716            .arg(k)
19717            .arg(v)
19718            .arg(o)
19719            .arg(&hd)
19720            .arg(&nh)
19721            .arg(&nhkv)
19722            .arg(&ti)
19723            .arg(&tkvi)
19724            .arg(&scale)
19725            .arg(&cz);
19726        unsafe {
19727            b.launch(cfg)?;
19728        }
19729        Ok(())
19730    }
19731
19732    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
19733    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
19734    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
19735    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
19736    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
19737    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
19738    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
19739    #[allow(clippy::too_many_arguments)]
19740    pub fn fa_dequant_kv_view_f32(
19741        &self,
19742        k: &cudarc::driver::CudaView<u8>,
19743        v: &cudarc::driver::CudaView<u8>,
19744        kf: &mut CudaSlice<f32>,
19745        vf: &mut CudaSlice<f32>,
19746        kv_dim_k: usize,
19747        kv_dim_v: usize,
19748        t_kv: usize,
19749        k_tok_bytes: usize,
19750        v_tok_bytes: usize,
19751        g: bool,
19752    ) -> Result<(), Box<dyn std::error::Error>> {
19753        let f = if g {
19754            self.func_g("fa_dequant_kv_ws_f32")
19755        } else {
19756            self.func("fa_dequant_kv_ws_f32")
19757        };
19758        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19759        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19760        let cfg = LaunchConfig {
19761            grid_dim: (nblk.max(1), 1, 1),
19762            block_dim: (256, 1, 1),
19763            shared_mem_bytes: 0,
19764        };
19765        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19766        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19767        let __s_b = self.gpu.stream();
19768        let mut b = __s_b.launch_builder(&f);
19769        b.arg(k)
19770            .arg(v)
19771            .arg(&mut *kf)
19772            .arg(&mut *vf)
19773            .arg(&kdk)
19774            .arg(&kdv)
19775            .arg(&tkvi)
19776            .arg(&ktb)
19777            .arg(&vtb);
19778        unsafe {
19779            b.launch(cfg)?;
19780        }
19781        Ok(())
19782    }
19783
19784    #[allow(clippy::too_many_arguments)]
19785    pub fn sdpa_naive_quantized_view(
19786        &self,
19787        q: &CudaSlice<f32>,
19788        k: &cudarc::driver::CudaView<u8>,
19789        v: &cudarc::driver::CudaView<u8>,
19790        o: &mut CudaSlice<f32>,
19791        head_dim: usize,
19792        n_head: usize,
19793        n_head_kv: usize,
19794        t: usize,
19795        t_kv: usize,
19796        scale: f32,
19797        causal: bool,
19798        k_tok_bytes: usize,
19799        v_tok_bytes: usize,
19800    ) -> Result<(), Box<dyn std::error::Error>> {
19801        let kv_dim = n_head_kv * head_dim;
19802        let mut kf = self.uninit(t_kv * kv_dim)?;
19803        let mut vf = self.uninit(t_kv * kv_dim)?;
19804        let f = self.func("fa_dequant_kv_ws_f32");
19805        let total = (2 * t_kv * kv_dim) as u64;
19806        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19807        let cfg = LaunchConfig {
19808            grid_dim: (nblk.max(1), 1, 1),
19809            block_dim: (256, 1, 1),
19810            shared_mem_bytes: 0,
19811        };
19812        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19813        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19814        let __s_b = self.gpu.stream();
19815        let mut b = __s_b.launch_builder(&f);
19816        b.arg(k)
19817            .arg(v)
19818            .arg(&mut kf)
19819            .arg(&mut vf)
19820            .arg(&kv_dim_i)
19821            .arg(&kv_dim_i)
19822            .arg(&t_kv_i)
19823            .arg(&k_tok_bytes_i)
19824            .arg(&v_tok_bytes_i);
19825        unsafe { b.launch(cfg)? };
19826        self.sdpa_naive(
19827            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19828        )
19829    }
19830
19831    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
19832    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
19833    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
19834    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
19835    /// unwindowed function above and produces bit-identical output at window == 0.
19836    ///
19837    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
19838    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
19839    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
19840    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
19841    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
19842    #[allow(clippy::too_many_arguments)]
19843    pub fn sdpa_naive_w_quantized_view(
19844        &self,
19845        q: &CudaSlice<f32>,
19846        k: &cudarc::driver::CudaView<u8>,
19847        v: &cudarc::driver::CudaView<u8>,
19848        o: &mut CudaSlice<f32>,
19849        head_dim: usize,
19850        n_head: usize,
19851        n_head_kv: usize,
19852        t: usize,
19853        t_kv: usize,
19854        scale: f32,
19855        causal: bool,
19856        window: usize,
19857        k_tok_bytes: usize,
19858        v_tok_bytes: usize,
19859    ) -> Result<(), Box<dyn std::error::Error>> {
19860        let kv_dim = n_head_kv * head_dim;
19861        let mut kf = self.uninit(t_kv * kv_dim)?;
19862        let mut vf = self.uninit(t_kv * kv_dim)?;
19863        let f = self.func("fa_dequant_kv_ws_f32");
19864        let total = (2 * t_kv * kv_dim) as u64;
19865        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19866        let cfg = LaunchConfig {
19867            grid_dim: (nblk.max(1), 1, 1),
19868            block_dim: (256, 1, 1),
19869            shared_mem_bytes: 0,
19870        };
19871        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19872        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19873        let __s_b = self.gpu.stream();
19874        let mut b = __s_b.launch_builder(&f);
19875        b.arg(k)
19876            .arg(v)
19877            .arg(&mut kf)
19878            .arg(&mut vf)
19879            .arg(&kv_dim_i)
19880            .arg(&kv_dim_i)
19881            .arg(&t_kv_i)
19882            .arg(&k_tok_bytes_i)
19883            .arg(&v_tok_bytes_i);
19884        unsafe { b.launch(cfg)? };
19885        self.sdpa_naive_w(
19886            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19887        )
19888    }
19889
19890    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
19891    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
19892    /// Q/K/V/O [head_dim, n_head(_kv), T].
19893    pub fn fa_prefill(
19894        &self,
19895        q: &CudaSlice<f32>,
19896        k: &CudaSlice<f32>,
19897        v: &CudaSlice<f32>,
19898        o: &mut CudaSlice<f32>,
19899        head_dim: usize,
19900        n_head: usize,
19901        n_head_kv: usize,
19902        t: usize,
19903        t_kv: usize,
19904        scale: f32,
19905        causal: bool,
19906    ) -> Result<(), Box<dyn std::error::Error>> {
19907        if portable_mma_gated() {
19908            return self.sdpa_naive(
19909                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19910            );
19911        }
19912        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
19913        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
19914        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
19915        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
19916        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
19917        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
19918        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
19919        let fa3_on = head_dim == 256
19920            && causal
19921            && t == t_kv
19922            && match std::env::var("MEMRA_FA3").as_deref() {
19923                Ok("0") => false,
19924                // The force arm consults the arch now: the bf16 stage below calls
19925                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
19926                // a portable build. Refuse at the switch, not at the lookup.
19927                Ok("1") => {
19928                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19929                    true
19930                }
19931                _ => cfg!(memra_hopper_mma),
19932            };
19933        if fa3_on {
19934            let n = t * n_head * head_dim;
19935            let nkv = t * n_head_kv * head_dim;
19936            let mut q16 = self.alloc_u8_uninit(n * 2)?;
19937            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19938            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19939            self.f32_to_bf16_into(q, &mut q16, n)?;
19940            self.f32_to_bf16_into(k, &mut k16, nkv)?;
19941            self.f32_to_bf16_into(v, &mut v16, nkv)?;
19942            let rc = {
19943                use cudarc::driver::{DevicePtr, DevicePtrMut};
19944                let stream = self.gpu.stream();
19945                let (qp, _g1) = q16.device_ptr(&stream);
19946                let (kp, _g2) = k16.device_ptr(&stream);
19947                let (vp, _g3) = v16.device_ptr(&stream);
19948                let (op, _g4) = o.device_ptr_mut(&stream);
19949                unsafe {
19950                    memra_fa3_prefill(
19951                        qp as *const core::ffi::c_void,
19952                        kp as *const core::ffi::c_void,
19953                        vp as *const core::ffi::c_void,
19954                        op as *mut f32,
19955                        t as i32,
19956                        n_head as i32,
19957                        n_head_kv as i32,
19958                        head_dim as i32,
19959                        scale,
19960                        stream.cu_stream() as *mut core::ffi::c_void,
19961                    )
19962                }
19963            };
19964            if rc != 0 {
19965                return Err(format!("memra_fa3_prefill rc={rc}").into());
19966            }
19967            return Ok(());
19968        }
19969        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
19970        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
19971        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
19972        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
19973        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19974        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19975        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19976            const BLOCK_Q: usize = 64;
19977            const BKX: usize = 32;
19978            let f = self.func("fa_prefill_bf16_p1");
19979            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19980                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19981            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19982            f.set_attribute(
19983                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19984                shmem as i32,
19985            )?;
19986            let cfg = LaunchConfig {
19987                grid_dim: (
19988                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19989                    n_head as u32,
19990                    1,
19991                ),
19992                block_dim: (32, 4, 1),
19993                shared_mem_bytes: shmem,
19994            };
19995            let (hd, nh, nhkv, ti, tkvi, cz) = (
19996                head_dim as i32,
19997                n_head as i32,
19998                n_head_kv as i32,
19999                t as i32,
20000                t_kv as i32,
20001                causal as i32,
20002            );
20003            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20004            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20005            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20006            let __s_b = self.gpu.stream();
20007            let mut b = __s_b.launch_builder(&f);
20008            b.arg(&qb)
20009                .arg(&kb)
20010                .arg(&vb)
20011                .arg(o)
20012                .arg(&hd)
20013                .arg(&nh)
20014                .arg(&nhkv)
20015                .arg(&ti)
20016                .arg(&tkvi)
20017                .arg(&scale)
20018                .arg(&cz);
20019            unsafe {
20020                b.launch(cfg)?;
20021            }
20022            return Ok(());
20023        }
20024        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
20025        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
20026        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
20027        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
20028        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
20029        const BK: usize = 32;
20030        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
20031        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
20032        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
20033        let (block_q, warps, w2_sfx): (usize, u32, &str) =
20034            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
20035        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
20036        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
20037        // other head_dims to sdpa_naive before reaching here.
20038        let hd_sfx = fa_hd_suffix(head_dim)?;
20039        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20040        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
20041        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
20042        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
20043        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
20044        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20045        let (kb16, vb16) = if bf16kv {
20046            let n = t_kv * n_head_kv * head_dim;
20047            let mut kb = self.alloc_u8_uninit(n * 2)?;
20048            let mut vb = self.alloc_u8_uninit(n * 2)?;
20049            let fcv = self.func("f32_to_bf16_bulk");
20050            let ni = n as i64;
20051            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20052            let __s_b = self.gpu.stream();
20053            let mut b = __s_b.launch_builder(&fcv);
20054            b.arg(k).arg(&mut kb).arg(&ni);
20055            unsafe {
20056                b.launch(cfgc)?;
20057            }
20058            let __s_b = self.gpu.stream();
20059            let mut b = __s_b.launch_builder(&fcv);
20060            b.arg(v).arg(&mut vb).arg(&ni);
20061            unsafe {
20062                b.launch(cfgc)?;
20063            }
20064            (Some(kb), Some(vb))
20065        } else {
20066            (None, None)
20067        };
20068        let f = self.func(&if bf16kv {
20069            format!("fa_prefill_bf16kv_pp{hd_sfx}")
20070        } else {
20071            format!(
20072                "fa_prefill_f32{}{}{hd_sfx}",
20073                if floor { "" } else { "_pp" },
20074                if floor { "" } else { w2_sfx }
20075            )
20076        });
20077        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
20078        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
20079        let kv_stages = if bf16kv { 2 } else { 1 };
20080        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20081            + 4 * (block_q * BK + 2 * block_q)) as u32;
20082        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20083        f.set_attribute(
20084            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20085            shmem as i32,
20086        )?;
20087        let cfg = LaunchConfig {
20088            grid_dim: (
20089                (t as u32 + block_q as u32 - 1) / block_q as u32,
20090                n_head as u32,
20091                1,
20092            ),
20093            block_dim: (32, warps, 1),
20094            shared_mem_bytes: shmem,
20095        };
20096        let (hd, nh, nhkv, ti, tkvi, cz) = (
20097            head_dim as i32,
20098            n_head as i32,
20099            n_head_kv as i32,
20100            t as i32,
20101            t_kv as i32,
20102            causal as i32,
20103        );
20104        let __s_b = self.gpu.stream();
20105        let mut b = __s_b.launch_builder(&f);
20106        b.arg(q);
20107        match (&kb16, &vb16) {
20108            (Some(kb), Some(vb)) => {
20109                b.arg(kb).arg(vb);
20110            }
20111            _ => {
20112                b.arg(k).arg(v);
20113            }
20114        }
20115        b.arg(o)
20116            .arg(&hd)
20117            .arg(&nh)
20118            .arg(&nhkv)
20119            .arg(&ti)
20120            .arg(&tkvi)
20121            .arg(&scale)
20122            .arg(&cz);
20123        unsafe {
20124            b.launch(cfg)?;
20125        }
20126        Ok(())
20127    }
20128
20129    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
20130    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
20131    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
20132    #[allow(clippy::too_many_arguments)]
20133    pub fn fa_prefill_w(
20134        &self,
20135        q: &CudaSlice<f32>,
20136        k: &CudaSlice<f32>,
20137        v: &CudaSlice<f32>,
20138        o: &mut CudaSlice<f32>,
20139        head_dim: usize,
20140        n_head: usize,
20141        n_head_kv: usize,
20142        t: usize,
20143        t_kv: usize,
20144        scale: f32,
20145        causal: bool,
20146        window: usize,
20147    ) -> Result<(), Box<dyn std::error::Error>> {
20148        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
20149        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
20150        if portable_mma_gated() {
20151            return self.sdpa_naive_w(
20152                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20153            );
20154        }
20155        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
20156        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
20157        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
20158        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20159        let faw_f32 =
20160            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20161        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20162        self.fa_prefill_w_arm(
20163            q,
20164            k,
20165            v,
20166            o,
20167            head_dim,
20168            n_head,
20169            n_head_kv,
20170            t,
20171            t_kv,
20172            scale,
20173            causal,
20174            window,
20175            floor || faw_f32,
20176            floor,
20177        )
20178    }
20179
20180    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
20181    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
20182    #[allow(clippy::too_many_arguments)]
20183    pub fn fa_prefill_w_pre(
20184        &self,
20185        qb: &CudaSlice<u8>,
20186        kb: &CudaSlice<u8>,
20187        vb: &CudaSlice<u8>,
20188        o: &mut CudaSlice<f32>,
20189        head_dim: usize,
20190        n_head: usize,
20191        n_head_kv: usize,
20192        t: usize,
20193        t_kv: usize,
20194        scale: f32,
20195        causal: bool,
20196        window: usize,
20197        v_f16: bool,
20198    ) -> Result<(), Box<dyn std::error::Error>> {
20199        const BLOCK_Q: usize = 64;
20200        const BK: usize = 32;
20201        debug_assert_eq!(head_dim, 256);
20202        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20203        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20204        if hp {
20205            const BLOCK_QH: usize = 32;
20206            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
20207            // else re-encode through the pooled scratch (stream-ordered reuse).
20208            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20209            let vh: &CudaSlice<u8> = if v_f16 {
20210                vb
20211            } else {
20212                let n = t_kv * n_head_kv * head_dim;
20213                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20214                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20215                }
20216                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20217                vguard.as_ref().unwrap()
20218            };
20219            let f = self.func("fa_prefill_w_bf16_p1h2");
20220            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20221            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20222            f.set_attribute(
20223                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20224                shmem as i32,
20225            )?;
20226            let cfg = LaunchConfig {
20227                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20228                block_dim: (32, 4, 1),
20229                shared_mem_bytes: shmem,
20230            };
20231            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20232                head_dim as i32,
20233                n_head as i32,
20234                n_head_kv as i32,
20235                t as i32,
20236                t_kv as i32,
20237                causal as i32,
20238                window as i32,
20239            );
20240            let __s_b = self.gpu.stream();
20241            let mut b = __s_b.launch_builder(&f);
20242            b.arg(qb)
20243                .arg(kb)
20244                .arg(vh)
20245                .arg(o)
20246                .arg(&hd)
20247                .arg(&nh)
20248                .arg(&nhkv)
20249                .arg(&ti)
20250                .arg(&tkvi)
20251                .arg(&scale)
20252                .arg(&cz)
20253                .arg(&wi);
20254            unsafe {
20255                b.launch(cfg)?;
20256            }
20257            return Ok(());
20258        }
20259        let f = self.func("fa_prefill_w_bf16_p1");
20260        let shmem =
20261            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20262        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20263        f.set_attribute(
20264            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20265            shmem as i32,
20266        )?;
20267        let cfg = LaunchConfig {
20268            grid_dim: (
20269                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20270                n_head as u32,
20271                1,
20272            ),
20273            block_dim: (32, 4, 1),
20274            shared_mem_bytes: shmem,
20275        };
20276        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20277            head_dim as i32,
20278            n_head as i32,
20279            n_head_kv as i32,
20280            t as i32,
20281            t_kv as i32,
20282            causal as i32,
20283            window as i32,
20284        );
20285        let __s_b = self.gpu.stream();
20286        let mut b = __s_b.launch_builder(&f);
20287        b.arg(qb)
20288            .arg(kb)
20289            .arg(vb)
20290            .arg(o)
20291            .arg(&hd)
20292            .arg(&nh)
20293            .arg(&nhkv)
20294            .arg(&ti)
20295            .arg(&tkvi)
20296            .arg(&scale)
20297            .arg(&cz)
20298            .arg(&wi);
20299        unsafe {
20300            b.launch(cfg)?;
20301        }
20302        Ok(())
20303    }
20304
20305    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
20306    #[allow(clippy::too_many_arguments)]
20307    pub fn fa_prefill_w_arm(
20308        &self,
20309        q: &CudaSlice<f32>,
20310        k: &CudaSlice<f32>,
20311        v: &CudaSlice<f32>,
20312        o: &mut CudaSlice<f32>,
20313        head_dim: usize,
20314        n_head: usize,
20315        n_head_kv: usize,
20316        t: usize,
20317        t_kv: usize,
20318        scale: f32,
20319        causal: bool,
20320        window: usize,
20321        f32_stage: bool,
20322        floor: bool,
20323    ) -> Result<(), Box<dyn std::error::Error>> {
20324        const BLOCK_Q: usize = 64;
20325        const BK: usize = 32;
20326        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20327        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
20328        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
20329        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
20330        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20331        let p1 = !floor
20332            && !f32_stage
20333            && *P1_ON.get_or_init(|| {
20334                std::env::var("MEMRA_FAW_P1")
20335                    .map(|v| v != "0")
20336                    .unwrap_or(true)
20337            });
20338        let hp =
20339            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20340        if hp {
20341            const BLOCK_QH: usize = 32;
20342            let f = self.func("fa_prefill_w_bf16_p1h2");
20343            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20344            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20345            f.set_attribute(
20346                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20347                shmem as i32,
20348            )?;
20349            let cfg = LaunchConfig {
20350                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20351                block_dim: (32, 4, 1),
20352                shared_mem_bytes: shmem,
20353            };
20354            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20355                head_dim as i32,
20356                n_head as i32,
20357                n_head_kv as i32,
20358                t as i32,
20359                t_kv as i32,
20360                causal as i32,
20361                window as i32,
20362            );
20363            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20364            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20365            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20366            let __s_b = self.gpu.stream();
20367            let mut b = __s_b.launch_builder(&f);
20368            b.arg(&qb)
20369                .arg(&kb)
20370                .arg(&vh)
20371                .arg(o)
20372                .arg(&hd)
20373                .arg(&nh)
20374                .arg(&nhkv)
20375                .arg(&ti)
20376                .arg(&tkvi)
20377                .arg(&scale)
20378                .arg(&cz)
20379                .arg(&wi);
20380            unsafe {
20381                b.launch(cfg)?;
20382            }
20383            return Ok(());
20384        }
20385        if p1 {
20386            let f = self.func("fa_prefill_w_bf16_p1");
20387            let shmem =
20388                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20389            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20390            f.set_attribute(
20391                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20392                shmem as i32,
20393            )?;
20394            let cfg = LaunchConfig {
20395                grid_dim: (
20396                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20397                    n_head as u32,
20398                    1,
20399                ),
20400                block_dim: (32, 4, 1),
20401                shared_mem_bytes: shmem,
20402            };
20403            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20404                head_dim as i32,
20405                n_head as i32,
20406                n_head_kv as i32,
20407                t as i32,
20408                t_kv as i32,
20409                causal as i32,
20410                window as i32,
20411            );
20412            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20413            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20414            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20415            let __s_b = self.gpu.stream();
20416            let mut b = __s_b.launch_builder(&f);
20417            b.arg(&qb)
20418                .arg(&kb)
20419                .arg(&vb)
20420                .arg(o)
20421                .arg(&hd)
20422                .arg(&nh)
20423                .arg(&nhkv)
20424                .arg(&ti)
20425                .arg(&tkvi)
20426                .arg(&scale)
20427                .arg(&cz)
20428                .arg(&wi);
20429            unsafe {
20430                b.launch(cfg)?;
20431            }
20432            return Ok(());
20433        }
20434        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
20435        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
20436        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20437        let g4 = !floor
20438            && !f32_stage
20439            && n_head_kv == 1
20440            && n_head % 4 == 0
20441            && *G4_ON.get_or_init(|| {
20442                std::env::var("MEMRA_FAW_G4")
20443                    .map(|v| v != "0")
20444                    .unwrap_or(true)
20445            });
20446        if g4 {
20447            const SP_M: usize = 16;
20448            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
20449            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
20450            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20451            let o2 = *O2_ON.get_or_init(|| {
20452                std::env::var("MEMRA_FAW_O2")
20453                    .map(|v| v != "0")
20454                    .unwrap_or(true)
20455            });
20456            let f = self.func(if o2 {
20457                "fa_prefill_w_bf16_g4o2"
20458            } else {
20459                "fa_prefill_w_bf16_g4"
20460            });
20461            let shmem = if o2 {
20462                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20463            } else {
20464                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20465                    as u32
20466            };
20467            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20468            f.set_attribute(
20469                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20470                shmem as i32,
20471            )?;
20472            let cfg = LaunchConfig {
20473                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20474                block_dim: (32, 4, 1),
20475                shared_mem_bytes: shmem,
20476            };
20477            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20478                head_dim as i32,
20479                n_head as i32,
20480                n_head_kv as i32,
20481                t as i32,
20482                t_kv as i32,
20483                causal as i32,
20484                window as i32,
20485            );
20486            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20487            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20488            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20489            let __s_b = self.gpu.stream();
20490            let mut b = __s_b.launch_builder(&f);
20491            b.arg(&qb)
20492                .arg(&kb)
20493                .arg(&vb)
20494                .arg(o)
20495                .arg(&hd)
20496                .arg(&nh)
20497                .arg(&nhkv)
20498                .arg(&ti)
20499                .arg(&tkvi)
20500                .arg(&scale)
20501                .arg(&cz)
20502                .arg(&wi);
20503            unsafe {
20504                b.launch(cfg)?;
20505            }
20506            return Ok(());
20507        }
20508        let f = self.func(if floor {
20509            "fa_prefill_w_f32"
20510        } else if f32_stage {
20511            "fa_prefill_w_f32_pp"
20512        } else {
20513            "fa_prefill_w_bf16_pp"
20514        });
20515        let shmem =
20516            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20517        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20518        f.set_attribute(
20519            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20520            shmem as i32,
20521        )?;
20522        let cfg = LaunchConfig {
20523            grid_dim: (
20524                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20525                n_head as u32,
20526                1,
20527            ),
20528            block_dim: (32, 4, 1),
20529            shared_mem_bytes: shmem,
20530        };
20531        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20532            head_dim as i32,
20533            n_head as i32,
20534            n_head_kv as i32,
20535            t as i32,
20536            t_kv as i32,
20537            causal as i32,
20538            window as i32,
20539        );
20540        if f32_stage {
20541            let __s_b = self.gpu.stream();
20542            let mut b = __s_b.launch_builder(&f);
20543            b.arg(q)
20544                .arg(k)
20545                .arg(v)
20546                .arg(o)
20547                .arg(&hd)
20548                .arg(&nh)
20549                .arg(&nhkv)
20550                .arg(&ti)
20551                .arg(&tkvi)
20552                .arg(&scale)
20553                .arg(&cz)
20554                .arg(&wi);
20555            unsafe {
20556                b.launch(cfg)?;
20557            }
20558        } else {
20559            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20560            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20561            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20562            let __s_b = self.gpu.stream();
20563            let mut b = __s_b.launch_builder(&f);
20564            b.arg(&qb)
20565                .arg(&kb)
20566                .arg(&vb)
20567                .arg(o)
20568                .arg(&hd)
20569                .arg(&nh)
20570                .arg(&nhkv)
20571                .arg(&ti)
20572                .arg(&tkvi)
20573                .arg(&scale)
20574                .arg(&cz)
20575                .arg(&wi);
20576            unsafe {
20577                b.launch(cfg)?;
20578            }
20579        }
20580        Ok(())
20581    }
20582
20583    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
20584    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
20585    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
20586    #[allow(clippy::too_many_arguments)]
20587    pub fn fa_prefill_hd512(
20588        &self,
20589        q: &CudaSlice<f32>,
20590        k: &CudaSlice<f32>,
20591        v: &CudaSlice<f32>,
20592        o: &mut CudaSlice<f32>,
20593        head_dim: usize,
20594        n_head: usize,
20595        n_head_kv: usize,
20596        t: usize,
20597        t_kv: usize,
20598        scale: f32,
20599        causal: bool,
20600    ) -> Result<(), Box<dyn std::error::Error>> {
20601        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
20602        if portable_mma_gated() {
20603            return self.sdpa_naive(
20604                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20605            );
20606        }
20607        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
20608        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
20609        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
20610        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
20611        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
20612        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20613        let f32_stage =
20614            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20615        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
20616        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
20617        // Own numeric config (partial-sum order) — battery-gated.
20618        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20619        let sp = !f32_stage
20620            && *SP_ON.get_or_init(|| {
20621                std::env::var("MEMRA_FA512_SP")
20622                    .map(|v| v != "0")
20623                    .unwrap_or(true)
20624            });
20625        self.fa_prefill_hd512_arm(
20626            q,
20627            k,
20628            v,
20629            o,
20630            head_dim,
20631            n_head,
20632            n_head_kv,
20633            t,
20634            t_kv,
20635            scale,
20636            causal,
20637            f32_stage,
20638            sp,
20639            sp && fa_f16pv_on(),
20640        )
20641    }
20642
20643    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
20644    #[allow(clippy::too_many_arguments)]
20645    pub fn fa_prefill_hd512_pre(
20646        &self,
20647        qb: &CudaSlice<u8>,
20648        kb: &CudaSlice<u8>,
20649        vb: &CudaSlice<u8>,
20650        o: &mut CudaSlice<f32>,
20651        head_dim: usize,
20652        n_head: usize,
20653        n_head_kv: usize,
20654        t: usize,
20655        t_kv: usize,
20656        scale: f32,
20657        causal: bool,
20658        v_f16: bool,
20659    ) -> Result<(), Box<dyn std::error::Error>> {
20660        debug_assert_eq!(head_dim, 512);
20661        const SP_M: usize = 16;
20662        const BKS: usize = 32;
20663        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
20664        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
20665        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
20666        let f16pv = fa_f16pv_on();
20667        let nw = if f16pv { fa512_wide_warps() } else { 2 };
20668        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20669        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20670        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20671        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20672            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
20673            let n = t_kv * n_head_kv * head_dim;
20674            let need = n * 2;
20675            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20676                *vguard = Some(self.alloc_uninit::<u8>(need)?);
20677            }
20678            let dst = vguard.as_mut().unwrap();
20679            self.bf16_to_f16_into(vb, n, dst)?;
20680            vguard.as_ref().unwrap()
20681        } else {
20682            vb
20683        };
20684        let f = self.func(if hp {
20685            "fa_prefill_bf16_hd512_sp16h2"
20686        } else {
20687            match (f16pv, nw) {
20688                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20689                (true, _) => "fa_prefill_bf16_hd512_sp16",
20690                _ => "fa_prefill_bf16_hd512_sp",
20691            }
20692        });
20693        let (nwarp, npart) = if hp {
20694            (4usize, 4usize)
20695        } else if nw > 2 {
20696            (nw, nw)
20697        } else {
20698            (2, 1)
20699        };
20700        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
20701        let shmem = if hp {
20702            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20703                as u32
20704        } else {
20705            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20706                + 4 * (npart * SP_M * BKS + SP_M)) as u32
20707        };
20708        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20709        f.set_attribute(
20710            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20711            shmem as i32,
20712        )?;
20713        let grid_y = if hp {
20714            (n_head / 2) as u32
20715        } else {
20716            n_head as u32
20717        };
20718        let cfg = LaunchConfig {
20719            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20720            block_dim: (32, nwarp as u32, 1),
20721            shared_mem_bytes: shmem,
20722        };
20723        let (hd, nh, nhkv, ti, tkvi, cz) = (
20724            head_dim as i32,
20725            n_head as i32,
20726            n_head_kv as i32,
20727            t as i32,
20728            t_kv as i32,
20729            causal as i32,
20730        );
20731        let __s_b = self.gpu.stream();
20732        let mut b = __s_b.launch_builder(&f);
20733        b.arg(qb)
20734            .arg(kb)
20735            .arg(vref)
20736            .arg(o)
20737            .arg(&hd)
20738            .arg(&nh)
20739            .arg(&nhkv)
20740            .arg(&ti)
20741            .arg(&tkvi)
20742            .arg(&scale)
20743            .arg(&cz);
20744        unsafe {
20745            b.launch(cfg)?;
20746        }
20747        Ok(())
20748    }
20749
20750    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
20751    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
20752    #[allow(clippy::too_many_arguments)]
20753    pub fn fa_prefill_hd512_arm(
20754        &self,
20755        q: &CudaSlice<f32>,
20756        k: &CudaSlice<f32>,
20757        v: &CudaSlice<f32>,
20758        o: &mut CudaSlice<f32>,
20759        head_dim: usize,
20760        n_head: usize,
20761        n_head_kv: usize,
20762        t: usize,
20763        t_kv: usize,
20764        scale: f32,
20765        causal: bool,
20766        f32_stage: bool,
20767        sp: bool,
20768        f16pv: bool,
20769    ) -> Result<(), Box<dyn std::error::Error>> {
20770        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20771        if sp && !f32_stage {
20772            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
20773            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
20774            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
20775            const SP_M: usize = 16;
20776            const BKS: usize = 32;
20777            let nw = if f16pv { fa512_wide_warps() } else { 2 };
20778            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20779            let f = self.func(if hp {
20780                "fa_prefill_bf16_hd512_sp16h2"
20781            } else {
20782                match (f16pv, nw) {
20783                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20784                    (true, _) => "fa_prefill_bf16_hd512_sp16",
20785                    _ => "fa_prefill_bf16_hd512_sp",
20786                }
20787            });
20788            let (nwarp, npart) = if hp {
20789                (4usize, 4usize)
20790            } else if nw > 2 {
20791                (nw, nw)
20792            } else {
20793                (2, 1)
20794            };
20795            let shmem = if hp {
20796                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20797                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20798            } else {
20799                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20800                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
20801            };
20802            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20803            f.set_attribute(
20804                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20805                shmem as i32,
20806            )?;
20807            let grid_y = if hp {
20808                (n_head / 2) as u32
20809            } else {
20810                n_head as u32
20811            };
20812            let cfg = LaunchConfig {
20813                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20814                block_dim: (32, nwarp as u32, 1),
20815                shared_mem_bytes: shmem,
20816            };
20817            let (hd, nh, nhkv, ti, tkvi, cz) = (
20818                head_dim as i32,
20819                n_head as i32,
20820                n_head_kv as i32,
20821                t as i32,
20822                t_kv as i32,
20823                causal as i32,
20824            );
20825            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20826            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20827            let vb = if f16pv {
20828                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20829            } else {
20830                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20831            };
20832            let __s_b = self.gpu.stream();
20833            let mut b = __s_b.launch_builder(&f);
20834            b.arg(&qb)
20835                .arg(&kb)
20836                .arg(&vb)
20837                .arg(o)
20838                .arg(&hd)
20839                .arg(&nh)
20840                .arg(&nhkv)
20841                .arg(&ti)
20842                .arg(&tkvi)
20843                .arg(&scale)
20844                .arg(&cz);
20845            unsafe {
20846                b.launch(cfg)?;
20847            }
20848            return Ok(());
20849        }
20850        const BLOCK_Q: usize = 32;
20851        const BK: usize = 32;
20852        const HALF: usize = 256;
20853        let f = self.func(if f32_stage {
20854            "fa_prefill_f32_hd512"
20855        } else {
20856            "fa_prefill_bf16_hd512"
20857        });
20858        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
20859        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20860            + 4 * BLOCK_Q) as u32;
20861        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20862        f.set_attribute(
20863            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20864            shmem as i32,
20865        )?;
20866        let cfg = LaunchConfig {
20867            grid_dim: (
20868                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20869                n_head as u32,
20870                2,
20871            ),
20872            block_dim: (32, 2, 1),
20873            shared_mem_bytes: shmem,
20874        };
20875        let (hd, nh, nhkv, ti, tkvi, cz) = (
20876            head_dim as i32,
20877            n_head as i32,
20878            n_head_kv as i32,
20879            t as i32,
20880            t_kv as i32,
20881            causal as i32,
20882        );
20883        if f32_stage {
20884            let __s_b = self.gpu.stream();
20885            let mut b = __s_b.launch_builder(&f);
20886            b.arg(q)
20887                .arg(k)
20888                .arg(v)
20889                .arg(o)
20890                .arg(&hd)
20891                .arg(&nh)
20892                .arg(&nhkv)
20893                .arg(&ti)
20894                .arg(&tkvi)
20895                .arg(&scale)
20896                .arg(&cz);
20897            unsafe {
20898                b.launch(cfg)?;
20899            }
20900        } else {
20901            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20902            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20903            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20904            let __s_b = self.gpu.stream();
20905            let mut b = __s_b.launch_builder(&f);
20906            b.arg(&qb)
20907                .arg(&kb)
20908                .arg(&vb)
20909                .arg(o)
20910                .arg(&hd)
20911                .arg(&nh)
20912                .arg(&nhkv)
20913                .arg(&ti)
20914                .arg(&tkvi)
20915                .arg(&scale)
20916                .arg(&cz);
20917            unsafe {
20918                b.launch(cfg)?;
20919            }
20920        }
20921        Ok(())
20922    }
20923
20924    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
20925    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
20926    /// separate f32_to_bf16 the FA entries would run).
20927    #[allow(clippy::too_many_arguments)]
20928    pub fn rope_neox2_bf16e(
20929        &self,
20930        q: &mut CudaSlice<f32>,
20931        k: &mut CudaSlice<f32>,
20932        qb: &mut CudaSlice<u8>,
20933        kb: &mut CudaSlice<u8>,
20934        pos: &CudaSlice<i32>,
20935        head_dim: usize,
20936        n_dims: usize,
20937        nh_q: usize,
20938        nh_k: usize,
20939        n_tokens: usize,
20940        base: f32,
20941        freq_scale: f32,
20942        ff: Option<&CudaSlice<f32>>,
20943    ) -> Result<(), Box<dyn std::error::Error>> {
20944        let f = self.func("rope_neox2_bf16e_f32");
20945        let rows = ((nh_q + nh_k) * n_tokens) as u32;
20946        let cfg = LaunchConfig {
20947            grid_dim: (rows, 1, 1),
20948            block_dim: ((head_dim / 2) as u32, 1, 1),
20949            shared_mem_bytes: 0,
20950        };
20951        let theta_scale = base.powf(-2.0 / n_dims as f32);
20952        let (hd, nd, nhq, nhk, nt) = (
20953            head_dim as i32,
20954            n_dims as i32,
20955            nh_q as i32,
20956            nh_k as i32,
20957            n_tokens as i32,
20958        );
20959        let __s_b = self.gpu.stream();
20960        let mut b = __s_b.launch_builder(&f);
20961        match ff {
20962            Some(t) => {
20963                b.arg(&mut *q)
20964                    .arg(&mut *k)
20965                    .arg(&mut *qb)
20966                    .arg(&mut *kb)
20967                    .arg(pos)
20968                    .arg(&hd)
20969                    .arg(&nd)
20970                    .arg(&nhq)
20971                    .arg(&nhk)
20972                    .arg(&nt)
20973                    .arg(&theta_scale)
20974                    .arg(&freq_scale)
20975                    .arg(t);
20976                unsafe {
20977                    b.launch(cfg)?;
20978                }
20979            }
20980            None => {
20981                let null: u64 = 0;
20982                b.arg(&mut *q)
20983                    .arg(&mut *k)
20984                    .arg(&mut *qb)
20985                    .arg(&mut *kb)
20986                    .arg(pos)
20987                    .arg(&hd)
20988                    .arg(&nd)
20989                    .arg(&nhq)
20990                    .arg(&nhk)
20991                    .arg(&nt)
20992                    .arg(&theta_scale)
20993                    .arg(&freq_scale)
20994                    .arg(&null);
20995                unsafe {
20996                    b.launch(cfg)?;
20997                }
20998            }
20999        }
21000        Ok(())
21001    }
21002
21003    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
21004    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
21005    pub fn f32_to_bf16(
21006        &self,
21007        x: &CudaSlice<f32>,
21008        n: usize,
21009    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21010        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
21011        let mut y = self.alloc_uninit::<u8>(n * 2)?;
21012        let f = self.func("f32_to_bf16_flat");
21013        let n_i = n as i64;
21014        let cfg = LaunchConfig {
21015            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21016            block_dim: (256, 1, 1),
21017            shared_mem_bytes: 0,
21018        };
21019        let __s_b = self.gpu.stream();
21020        let mut b = __s_b.launch_builder(&f);
21021        b.arg(x).arg(&mut y).arg(&n_i);
21022        unsafe {
21023            b.launch(cfg)?;
21024        }
21025        Ok(y)
21026    }
21027
21028    pub fn f32_to_f16(
21029        &self,
21030        x: &CudaSlice<f32>,
21031        n: usize,
21032    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21033        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
21034        let mut y = self.alloc_uninit::<u8>(n * 2)?;
21035        let f = self.func("f32_to_f16_flat");
21036        let n_i = n as i64;
21037        let cfg = LaunchConfig {
21038            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21039            block_dim: (256, 1, 1),
21040            shared_mem_bytes: 0,
21041        };
21042        let __s_b = self.gpu.stream();
21043        let mut b = __s_b.launch_builder(&f);
21044        b.arg(x).arg(&mut y).arg(&n_i);
21045        unsafe {
21046            b.launch(cfg)?;
21047        }
21048        Ok(y)
21049    }
21050
21051    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
21052    pub fn bf16_to_f16(
21053        &self,
21054        xb: &CudaSlice<u8>,
21055        n: usize,
21056    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21057        let mut y = self.alloc_uninit::<u8>(n * 2)?;
21058        self.bf16_to_f16_into(xb, n, &mut y)?;
21059        Ok(y)
21060    }
21061
21062    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
21063    pub fn bf16_to_f16_into(
21064        &self,
21065        xb: &CudaSlice<u8>,
21066        n: usize,
21067        y: &mut CudaSlice<u8>,
21068    ) -> Result<(), Box<dyn std::error::Error>> {
21069        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21070        assert!(y.len() >= n * 2);
21071        let f = self.func("bf16_to_f16_flat");
21072        let n2 = (n / 2) as i64;
21073        let cfg = LaunchConfig {
21074            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21075            block_dim: (256, 1, 1),
21076            shared_mem_bytes: 0,
21077        };
21078        let __s_b = self.gpu.stream();
21079        let mut b = __s_b.launch_builder(&f);
21080        b.arg(xb).arg(y).arg(&n2);
21081        unsafe {
21082            b.launch(cfg)?;
21083        }
21084        Ok(())
21085    }
21086
21087    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
21088    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
21089    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
21090    /// head_dim in {256, 128}, bf16kv lane on.
21091    #[allow(clippy::too_many_arguments)]
21092    pub fn fa_prefill_vl8(
21093        &self,
21094        seqs: &[FaSeqVl],
21095        head_dim: usize,
21096        n_head: usize,
21097        n_head_kv: usize,
21098        scale: f32,
21099    ) -> Result<(), Box<dyn std::error::Error>> {
21100        const BK: usize = 32;
21101        let b = seqs.len();
21102        assert!(b >= 1 && b <= 8);
21103        let mut packed = [FaSeqVl::default(); 8];
21104        packed[..b].copy_from_slice(seqs);
21105        let v = FaVl8(packed);
21106        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21107        let ept = (n_head_kv * head_dim) as i32;
21108        {
21109            let f = self.func("fa_mirror_vl");
21110            let max_n = (max_t as i64) * ept as i64;
21111            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21112            for which in 0..2i32 {
21113                let cfg = LaunchConfig {
21114                    grid_dim: (blocks, 1, b as u32),
21115                    block_dim: (256, 1, 1),
21116                    shared_mem_bytes: 0,
21117                };
21118                let __s_lb = self.gpu.stream();
21119                let mut lb = __s_lb.launch_builder(&f);
21120                lb.arg(&v).arg(&ept).arg(&which);
21121                unsafe {
21122                    lb.launch(cfg)?;
21123                }
21124            }
21125        }
21126        let hd_sfx = fa_hd_suffix(head_dim)?;
21127        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21128        let block_q = 64usize;
21129        let kv_stages = 2usize;
21130        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21131            + 4 * (block_q * BK + 2 * block_q)) as u32;
21132        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21133        f.set_attribute(
21134            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21135            shmem as i32,
21136        )?;
21137        let cfg = LaunchConfig {
21138            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21139            block_dim: (32, 4, 1),
21140            shared_mem_bytes: shmem,
21141        };
21142        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21143        let __s_lb = self.gpu.stream();
21144        let mut lb = __s_lb.launch_builder(&f);
21145        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21146        unsafe {
21147            lb.launch(cfg)?;
21148        }
21149        Ok(())
21150    }
21151
21152    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
21153    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
21154    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
21155    #[allow(clippy::too_many_arguments)]
21156    pub fn attn_pre_vl8(
21157        &self,
21158        seqs: &[AttnPreVl],
21159        wq: &CudaSlice<f32>,
21160        wk: &CudaSlice<f32>,
21161        head_dim: usize,
21162        rope_dims: usize,
21163        n_head: usize,
21164        n_head_kv: usize,
21165        eps: f32,
21166        freq_base: f32,
21167        freq_scale: f32,
21168        kv_dim_k: usize,
21169        kv_dim_v: usize,
21170        k_tok_bytes: usize,
21171        v_tok_bytes: usize,
21172    ) -> Result<(), Box<dyn std::error::Error>> {
21173        let b = seqs.len();
21174        assert!(b >= 1 && b <= 8);
21175        let mut packed = [AttnPreVl::default(); 8];
21176        packed[..b].copy_from_slice(seqs);
21177        let v = AttnPreVl8(packed);
21178        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21179        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21180        {
21181            let f = self.func("q_gate_split_vl");
21182            let n = max_t * (n_head * head_dim) as u32;
21183            let cfg = LaunchConfig {
21184                grid_dim: (n.div_ceil(256), 1, b as u32),
21185                block_dim: (256, 1, 1),
21186                shared_mem_bytes: 0,
21187            };
21188            let __s_lb = self.gpu.stream();
21189            let mut lb = __s_lb.launch_builder(&f);
21190            lb.arg(&v).arg(&hd).arg(&nh);
21191            unsafe {
21192                lb.launch(cfg)?;
21193            }
21194        }
21195        {
21196            let f = self.func("attn_rms_vl");
21197            let cfg = LaunchConfig {
21198                grid_dim: (max_t * n_head as u32, 2, b as u32),
21199                block_dim: (rms_block(), 1, 1),
21200                shared_mem_bytes: 0,
21201            };
21202            let __s_lb = self.gpu.stream();
21203            let mut lb = __s_lb.launch_builder(&f);
21204            lb.arg(&v)
21205                .arg(wq)
21206                .arg(wk)
21207                .arg(&hd)
21208                .arg(&nh)
21209                .arg(&nhkv)
21210                .arg(&eps);
21211            unsafe {
21212                lb.launch(cfg)?;
21213            }
21214        }
21215        {
21216            let f = self.func("attn_rope_vl");
21217            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21218            let nd = rope_dims as i32;
21219            let cfg = LaunchConfig {
21220                grid_dim: (max_t * n_head as u32, 2, b as u32),
21221                block_dim: ((head_dim / 2) as u32, 1, 1),
21222                shared_mem_bytes: 0,
21223            };
21224            let __s_lb = self.gpu.stream();
21225            let mut lb = __s_lb.launch_builder(&f);
21226            lb.arg(&v)
21227                .arg(&hd)
21228                .arg(&nd)
21229                .arg(&nh)
21230                .arg(&nhkv)
21231                .arg(&theta_scale)
21232                .arg(&freq_scale);
21233            unsafe {
21234                lb.launch(cfg)?;
21235            }
21236        }
21237        {
21238            let f = self.func("append_kv_vl");
21239            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21240            let cfg = LaunchConfig {
21241                grid_dim: (nblk, max_t, b as u32),
21242                block_dim: (32, 1, 1),
21243                shared_mem_bytes: 0,
21244            };
21245            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21246            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21247            let __s_lb = self.gpu.stream();
21248            let mut lb = __s_lb.launch_builder(&f);
21249            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21250            unsafe {
21251                lb.launch(cfg)?;
21252            }
21253        }
21254        Ok(())
21255    }
21256
21257    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
21258    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
21259    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
21260    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
21261    pub fn fa_prefill_view(
21262        &self,
21263        q: &CudaSlice<f32>,
21264        k: &cudarc::driver::CudaView<u8>,
21265        v: &cudarc::driver::CudaView<u8>,
21266        o: &mut CudaSlice<f32>,
21267        head_dim: usize,
21268        n_head: usize,
21269        n_head_kv: usize,
21270        t: usize,
21271        t_kv: usize,
21272        scale: f32,
21273        causal: bool,
21274        k_tok_bytes: usize,
21275        v_tok_bytes: usize,
21276        g: bool,
21277    ) -> Result<(), Box<dyn std::error::Error>> {
21278        if portable_mma_gated() {
21279            return self.sdpa_naive_quantized_view(
21280                q,
21281                k,
21282                v,
21283                o,
21284                head_dim,
21285                n_head,
21286                n_head_kv,
21287                t,
21288                t_kv,
21289                scale,
21290                causal,
21291                k_tok_bytes,
21292                v_tok_bytes,
21293            );
21294        }
21295        const BLOCK_Q: usize = 64;
21296        const BK: usize = 32;
21297        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
21298        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
21299        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21300        let f = if g {
21301            self.func_g(&name)
21302        } else {
21303            self.func(&name)
21304        };
21305        let shmem =
21306            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21307        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21308        f.set_attribute(
21309            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21310            shmem as i32,
21311        )?;
21312        let cfg = LaunchConfig {
21313            grid_dim: (
21314                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21315                n_head as u32,
21316                1,
21317            ),
21318            block_dim: (32, 4, 1),
21319            shared_mem_bytes: shmem,
21320        };
21321        let (hd, nh, nhkv, ti, tkvi, cz) = (
21322            head_dim as i32,
21323            n_head as i32,
21324            n_head_kv as i32,
21325            t as i32,
21326            t_kv as i32,
21327            causal as i32,
21328        );
21329        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21330        let __s_b = self.gpu.stream();
21331        let mut b = __s_b.launch_builder(&f);
21332        b.arg(q)
21333            .arg(k)
21334            .arg(v)
21335            .arg(o)
21336            .arg(&hd)
21337            .arg(&nh)
21338            .arg(&nhkv)
21339            .arg(&ti)
21340            .arg(&tkvi)
21341            .arg(&scale)
21342            .arg(&cz)
21343            .arg(&ktb)
21344            .arg(&vtb);
21345        unsafe {
21346            b.launch(cfg)?;
21347        }
21348        Ok(())
21349    }
21350
21351    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
21352    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
21353    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
21354    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
21355    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
21356    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
21357    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
21358    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
21359    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
21360    #[allow(clippy::too_many_arguments)]
21361    pub fn fa_prefill_view_ws(
21362        &self,
21363        q: &CudaSlice<f32>,
21364        k: &cudarc::driver::CudaView<u8>,
21365        v: &cudarc::driver::CudaView<u8>,
21366        o: &mut CudaSlice<f32>,
21367        head_dim: usize,
21368        n_head: usize,
21369        n_head_kv: usize,
21370        t: usize,
21371        t_kv: usize,
21372        scale: f32,
21373        causal: bool,
21374        k_tok_bytes: usize,
21375        v_tok_bytes: usize,
21376        g: bool,
21377    ) -> Result<(), Box<dyn std::error::Error>> {
21378        if portable_mma_gated() {
21379            return self.sdpa_naive_quantized_view(
21380                q,
21381                k,
21382                v,
21383                o,
21384                head_dim,
21385                n_head,
21386                n_head_kv,
21387                t,
21388                t_kv,
21389                scale,
21390                causal,
21391                k_tok_bytes,
21392                v_tok_bytes,
21393            );
21394        }
21395        const BLOCK_Q: usize = 64;
21396        const BK: usize = 32;
21397        let kv_dim_k = n_head_kv * head_dim;
21398        let kv_dim_v = n_head_kv * head_dim;
21399        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
21400        let v_ws_bytes = t_kv * kv_dim_v * 2;
21401        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
21402        let mut guard = self.prime_deqw_ws.lock().unwrap();
21403        let need_grow = match guard.as_ref() {
21404            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21405            None => true,
21406        };
21407        if need_grow {
21408            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21409            let (ck, cv) = guard
21410                .as_ref()
21411                .map(|(a, b)| (a.len(), b.len()))
21412                .unwrap_or((0, 0));
21413            *guard = Some((
21414                self.alloc_u8(grow(ck, k_ws_bytes))?,
21415                self.alloc_u8(grow(cv, v_ws_bytes))?,
21416            ));
21417        }
21418        let (kw, vw) = guard.as_mut().unwrap();
21419        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
21420        {
21421            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
21422            let f = if g {
21423                self.func_g("fa_dequant_kv_ws_bf16")
21424            } else {
21425                self.func("fa_dequant_kv_ws_bf16")
21426            };
21427            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21428            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21429            let cfg = LaunchConfig {
21430                grid_dim: (nblk.max(1), 1, 1),
21431                block_dim: (256, 1, 1),
21432                shared_mem_bytes: 0,
21433            };
21434            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21435            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21436            let __s_b = self.gpu.stream();
21437            let mut b = __s_b.launch_builder(&f);
21438            b.arg(k)
21439                .arg(v)
21440                .arg(&mut *kw)
21441                .arg(&mut *vw)
21442                .arg(&kdk)
21443                .arg(&kdv)
21444                .arg(&tkvi)
21445                .arg(&ktb)
21446                .arg(&vtb);
21447            unsafe {
21448                b.launch(cfg)?;
21449            }
21450        }
21451        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
21452        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
21453        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
21454        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
21455        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
21456        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
21457        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
21458        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21459            .map(|v| v != "0")
21460            .unwrap_or(true);
21461        {
21462            let hd_sfx = fa_hd_suffix(head_dim)?;
21463            let f = self.func(&format!(
21464                "fa_prefill_qw{}{hd_sfx}",
21465                if db { "_db" } else { "" }
21466            ));
21467            let shmem = if db {
21468                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
21469                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21470            } else {
21471                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21472            };
21473            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21474            f.set_attribute(
21475                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21476                shmem as i32,
21477            )?;
21478            let cfg = LaunchConfig {
21479                grid_dim: (
21480                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21481                    n_head as u32,
21482                    1,
21483                ),
21484                block_dim: (32, 4, 1),
21485                shared_mem_bytes: shmem,
21486            };
21487            let (hd, nh, nhkv, ti, tkvi, cz) = (
21488                head_dim as i32,
21489                n_head as i32,
21490                n_head_kv as i32,
21491                t as i32,
21492                t_kv as i32,
21493                causal as i32,
21494            );
21495            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21496            let __s_b = self.gpu.stream();
21497            let mut b = __s_b.launch_builder(&f);
21498            b.arg(q)
21499                .arg(&*kw)
21500                .arg(&*vw)
21501                .arg(o)
21502                .arg(&hd)
21503                .arg(&nh)
21504                .arg(&nhkv)
21505                .arg(&ti)
21506                .arg(&tkvi)
21507                .arg(&scale)
21508                .arg(&cz)
21509                .arg(&kdk)
21510                .arg(&kdv);
21511            unsafe {
21512                b.launch(cfg)?;
21513            }
21514        }
21515        Ok(())
21516    }
21517
21518    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
21519    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
21520    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
21521    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
21522    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
21523    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
21524    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
21525    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
21526    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
21527    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
21528    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
21529    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
21530    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
21531    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
21532    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
21533    #[allow(clippy::too_many_arguments)]
21534    pub fn fa_prefill_view_ws_w_hd128(
21535        &self,
21536        q: &CudaSlice<f32>,
21537        k: &cudarc::driver::CudaView<u8>,
21538        v: &cudarc::driver::CudaView<u8>,
21539        o: &mut CudaSlice<f32>,
21540        head_dim: usize,
21541        n_head: usize,
21542        n_head_kv: usize,
21543        t: usize,
21544        t_kv: usize,
21545        scale: f32,
21546        causal: bool,
21547        window: usize,
21548        k_tok_bytes: usize,
21549        v_tok_bytes: usize,
21550    ) -> Result<(), Box<dyn std::error::Error>> {
21551        assert_eq!(
21552            head_dim, 128,
21553            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21554        );
21555        if portable_mma_gated() {
21556            return self.sdpa_naive_w_quantized_view(
21557                q,
21558                k,
21559                v,
21560                o,
21561                head_dim,
21562                n_head,
21563                n_head_kv,
21564                t,
21565                t_kv,
21566                scale,
21567                causal,
21568                window,
21569                k_tok_bytes,
21570                v_tok_bytes,
21571            );
21572        }
21573        const BLOCK_Q: usize = 64;
21574        const BK: usize = 32;
21575        let kv_dim_k = n_head_kv * head_dim;
21576        let kv_dim_v = n_head_kv * head_dim;
21577        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
21578        let v_ws_bytes = t_kv * kv_dim_v * 2;
21579        let mut guard = self.prime_deqw_ws.lock().unwrap();
21580        let need_grow = match guard.as_ref() {
21581            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21582            None => true,
21583        };
21584        if need_grow {
21585            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21586            let (ck, cv) = guard
21587                .as_ref()
21588                .map(|(a, b)| (a.len(), b.len()))
21589                .unwrap_or((0, 0));
21590            *guard = Some((
21591                self.alloc_u8(grow(ck, k_ws_bytes))?,
21592                self.alloc_u8(grow(cv, v_ws_bytes))?,
21593            ));
21594        }
21595        let (kw, vw) = guard.as_mut().unwrap();
21596        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
21597        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
21598        {
21599            let f = self.func("fa_dequant_kv_ws_bf16");
21600            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21601            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21602            let cfg = LaunchConfig {
21603                grid_dim: (nblk.max(1), 1, 1),
21604                block_dim: (256, 1, 1),
21605                shared_mem_bytes: 0,
21606            };
21607            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21608            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21609            let __s_b = self.gpu.stream();
21610            let mut b = __s_b.launch_builder(&f);
21611            b.arg(k)
21612                .arg(v)
21613                .arg(&mut *kw)
21614                .arg(&mut *vw)
21615                .arg(&kdk)
21616                .arg(&kdv)
21617                .arg(&tkvi)
21618                .arg(&ktb)
21619                .arg(&vtb);
21620            unsafe {
21621                b.launch(cfg)?;
21622            }
21623        }
21624        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
21625        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21626            .map(|v| v != "0")
21627            .unwrap_or(true);
21628        {
21629            let f = self.func(if db {
21630                "fa_prefill_qw_db_w_hd128"
21631            } else {
21632                "fa_prefill_qw_w_hd128"
21633            });
21634            let shmem = if db {
21635                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21636            } else {
21637                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21638            };
21639            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21640            f.set_attribute(
21641                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21642                shmem as i32,
21643            )?;
21644            let cfg = LaunchConfig {
21645                grid_dim: (
21646                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21647                    n_head as u32,
21648                    1,
21649                ),
21650                block_dim: (32, 4, 1),
21651                shared_mem_bytes: shmem,
21652            };
21653            let (hd, nh, nhkv, ti, tkvi, cz) = (
21654                head_dim as i32,
21655                n_head as i32,
21656                n_head_kv as i32,
21657                t as i32,
21658                t_kv as i32,
21659                causal as i32,
21660            );
21661            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21662            let __s_b = self.gpu.stream();
21663            let mut b = __s_b.launch_builder(&f);
21664            b.arg(q)
21665                .arg(&*kw)
21666                .arg(&*vw)
21667                .arg(o)
21668                .arg(&hd)
21669                .arg(&nh)
21670                .arg(&nhkv)
21671                .arg(&ti)
21672                .arg(&tkvi)
21673                .arg(&scale)
21674                .arg(&cz)
21675                .arg(&kdk)
21676                .arg(&kdv)
21677                .arg(&wnd);
21678            unsafe {
21679                b.launch(cfg)?;
21680            }
21681        }
21682        Ok(())
21683    }
21684
21685    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
21686    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
21687    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
21688    pub fn fa_decode(
21689        &self,
21690        q: &CudaSlice<f32>,
21691        k: &cudarc::driver::CudaView<u8>,
21692        v: &cudarc::driver::CudaView<u8>,
21693        o: &mut CudaSlice<f32>,
21694        head_dim: usize,
21695        n_head: usize,
21696        n_head_kv: usize,
21697        t_kv: usize,
21698        scale: f32,
21699        k_tok_bytes: usize,
21700        v_tok_bytes: usize,
21701    ) -> Result<(), Box<dyn std::error::Error>> {
21702        self.fa_decode_kvmod(
21703            q,
21704            k,
21705            v,
21706            o,
21707            head_dim,
21708            n_head,
21709            n_head_kv,
21710            t_kv,
21711            scale,
21712            k_tok_bytes,
21713            v_tok_bytes,
21714            false,
21715        )
21716    }
21717
21718    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
21719    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
21720    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
21721    #[allow(clippy::too_many_arguments)]
21722    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
21723    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
21724    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
21725    #[allow(clippy::too_many_arguments)]
21726    #[allow(clippy::too_many_arguments)]
21727    fn fa_decode_scalar_unified(
21728        &self,
21729        q: &cudarc::driver::CudaView<f32>,
21730        k: &cudarc::driver::CudaView<u8>,
21731        v: &cudarc::driver::CudaView<u8>,
21732        o: &mut cudarc::driver::CudaViewMut<f32>,
21733        head_dim: usize,
21734        n_head: usize,
21735        n_head_kv: usize,
21736        t_kv_host: usize,
21737        t_kv_dev: Option<&CudaSlice<i32>>,
21738        scale: f32,
21739        n_splits: usize,
21740        split_keys: usize,
21741        k_tok_bytes: usize,
21742        v_tok_bytes: usize,
21743        g: bool,
21744        part_o: &mut CudaSlice<f32>,
21745        part_m: &mut CudaSlice<f32>,
21746        part_l: &mut CudaSlice<f32>,
21747        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21748    ) -> Result<(), Box<dyn std::error::Error>> {
21749        let f = if g {
21750            self.func_g("fa_decode_f32")
21751        } else {
21752            self.fa_func("fa_decode_f32", head_dim)
21753        };
21754        let cfg = LaunchConfig {
21755            grid_dim: (n_head as u32, n_splits as u32, 1),
21756            block_dim: (head_dim as u32, 1, 1),
21757            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21758        };
21759        let (hd, nh, nhkv, nsp) = (
21760            head_dim as i32,
21761            n_head as i32,
21762            n_head_kv as i32,
21763            n_splits as i32,
21764        );
21765        let (ktb, vtb, tkvi, ski) = (
21766            k_tok_bytes as i64,
21767            v_tok_bytes as i64,
21768            t_kv_host as i32,
21769            split_keys as i32,
21770        );
21771        let __s_b = self.gpu.stream();
21772        let mut b = __s_b.launch_builder(&f);
21773        match t_kv_dev {
21774            Some(d) => {
21775                b.arg(q)
21776                    .arg(k)
21777                    .arg(v)
21778                    .arg(&mut *part_o)
21779                    .arg(&mut *part_m)
21780                    .arg(&mut *part_l)
21781                    .arg(&hd)
21782                    .arg(&nh)
21783                    .arg(&nhkv)
21784                    .arg(&tkvi)
21785                    .arg(d)
21786                    .arg(&scale)
21787                    .arg(&nsp)
21788                    .arg(&ski)
21789                    .arg(&ktb)
21790                    .arg(&vtb);
21791                unsafe {
21792                    b.launch(cfg)?;
21793                }
21794            }
21795            None => {
21796                let null: u64 = 0;
21797                b.arg(q)
21798                    .arg(k)
21799                    .arg(v)
21800                    .arg(&mut *part_o)
21801                    .arg(&mut *part_m)
21802                    .arg(&mut *part_l)
21803                    .arg(&hd)
21804                    .arg(&nh)
21805                    .arg(&nhkv)
21806                    .arg(&tkvi)
21807                    .arg(&null)
21808                    .arg(&scale)
21809                    .arg(&nsp)
21810                    .arg(&ski)
21811                    .arg(&ktb)
21812                    .arg(&vtb);
21813                unsafe {
21814                    b.launch(cfg)?;
21815                }
21816            }
21817        }
21818        let cfg2 = LaunchConfig {
21819            grid_dim: (n_head as u32, 1, 1),
21820            block_dim: (head_dim as u32, 1, 1),
21821            shared_mem_bytes: 0,
21822        };
21823        if let Some((oq, od)) = q8_out {
21824            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
21825            let fc = if g {
21826                self.func_g("fa_decode_combine_q8_1")
21827            } else {
21828                self.fa_func("fa_decode_combine_q8_1", head_dim)
21829            };
21830            let __s_b2 = self.gpu.stream();
21831            let mut b2 = __s_b2.launch_builder(&fc);
21832            b2.arg(&*part_o)
21833                .arg(&*part_m)
21834                .arg(&*part_l)
21835                .arg(oq)
21836                .arg(od)
21837                .arg(&hd)
21838                .arg(&nh)
21839                .arg(&nsp);
21840            unsafe {
21841                b2.launch(cfg2)?;
21842            }
21843            return Ok(());
21844        }
21845        let fc = if g {
21846            self.func_g("fa_decode_combine_f32")
21847        } else {
21848            self.fa_func("fa_decode_combine_f32", head_dim)
21849        };
21850        let __s_b2 = self.gpu.stream();
21851        let mut b2 = __s_b2.launch_builder(&fc);
21852        b2.arg(&*part_o)
21853            .arg(&*part_m)
21854            .arg(&*part_l)
21855            .arg(o)
21856            .arg(&hd)
21857            .arg(&nh)
21858            .arg(&nsp);
21859        unsafe {
21860            b2.launch(cfg2)?;
21861        }
21862        Ok(())
21863    }
21864
21865    pub fn fa_decode_kvmod(
21866        &self,
21867        q: &CudaSlice<f32>,
21868        k: &cudarc::driver::CudaView<u8>,
21869        v: &cudarc::driver::CudaView<u8>,
21870        o: &mut CudaSlice<f32>,
21871        head_dim: usize,
21872        n_head: usize,
21873        n_head_kv: usize,
21874        t_kv: usize,
21875        scale: f32,
21876        k_tok_bytes: usize,
21877        v_tok_bytes: usize,
21878        g: bool,
21879    ) -> Result<(), Box<dyn std::error::Error>> {
21880        let q_view = q.as_view();
21881        let mut o_view = o.as_view_mut();
21882        self.fa_decode_kvmod_view(
21883            &q_view,
21884            k,
21885            v,
21886            &mut o_view,
21887            head_dim,
21888            n_head,
21889            n_head_kv,
21890            t_kv,
21891            scale,
21892            k_tok_bytes,
21893            v_tok_bytes,
21894            g,
21895        )
21896    }
21897
21898    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
21899    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
21900    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
21901    /// per-session KV view and FA launch.
21902    #[allow(clippy::too_many_arguments)]
21903    pub fn fa_decode_kvmod_view(
21904        &self,
21905        q: &cudarc::driver::CudaView<f32>,
21906        k: &cudarc::driver::CudaView<u8>,
21907        v: &cudarc::driver::CudaView<u8>,
21908        o: &mut cudarc::driver::CudaViewMut<f32>,
21909        head_dim: usize,
21910        n_head: usize,
21911        n_head_kv: usize,
21912        t_kv: usize,
21913        scale: f32,
21914        k_tok_bytes: usize,
21915        v_tok_bytes: usize,
21916        g: bool,
21917    ) -> Result<(), Box<dyn std::error::Error>> {
21918        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
21919        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
21920        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
21921        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
21922        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
21923        //
21924        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
21925        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
21926        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
21927        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
21928        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
21929        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
21930        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
21931        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
21932        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
21933        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
21934        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
21935        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
21936        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
21937        // fall to the exact scalar there instead of the broken register arm.
21938        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21939        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
21940        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
21941        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
21942        if g && head_dim == 256 && !fa_v4_at(t_kv) {
21943            fa_vec = false;
21944        }
21945        let sp = fa_split_keys(t_kv, n_head_kv);
21946        let n_splits = if fa_vec {
21947            ((t_kv + sp - 1) / sp).max(1)
21948        } else {
21949            ((t_kv + 255) / 256).max(1)
21950        };
21951        let o_len = n_head * n_splits * head_dim;
21952        let ml_len = n_head * n_splits;
21953        let mut part_guard = self.fa_part_pool.lock().unwrap();
21954        if part_guard
21955            .as_ref()
21956            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21957            .unwrap_or(true)
21958        {
21959            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21960            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21961            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21962            // later live allocations land at those addresses, and the next graph REPLAY writes
21963            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21964            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21965            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21966            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21967            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21968            // (total retired < final size).
21969            let old = part_guard.take();
21970            let (co, cm) = old
21971                .as_ref()
21972                .map(|pp| (pp.0.len(), pp.1.len()))
21973                .unwrap_or((0, 0));
21974            if let Some(old) = old {
21975                self.fa_part_retired.lock().unwrap().push(old);
21976            }
21977            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21978                eprintln!(
21979                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21980                    co, o_len, cm, ml_len
21981                );
21982            }
21983            *part_guard = Some((
21984                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21985                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21986                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21987            ));
21988        }
21989        let pg = part_guard.as_mut().unwrap();
21990        self.gpu
21991            .stream()
21992            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21993        self.gpu
21994            .stream()
21995            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21996        self.gpu
21997            .stream()
21998            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21999        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22000        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22001        let (hd, nh, nhkv, tkvi, nsp) = (
22002            head_dim as i32,
22003            n_head as i32,
22004            n_head_kv as i32,
22005            t_kv as i32,
22006            n_splits as i32,
22007        );
22008        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22009        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
22010        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
22011        // silently truncating the accumulator.
22012        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22013        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
22014        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
22015        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
22016        // 178.4 -> 173.7 when 512 rode vec unconditionally).
22017        let fa512_min = fa512_min_tkv();
22018        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
22019        // g-module keeps the v4 pick (its class is not the depth-decay class).
22020        let deep = fa_vec
22021            && head_dim == 256
22022            && fa_v4_at(t_kv)
22023            && !g
22024            && fa_deep_at(t_kv)
22025            && !matches!(fa_v4_mode(), "noB3" | "stage");
22026        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
22027            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
22028            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
22029            let gqa = (n_head / n_head_kv).max(1) as u32;
22030            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
22031            (
22032                fv,
22033                LaunchConfig {
22034                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22035                    block_dim: (32, gqa, 1),
22036                    shared_mem_bytes: 0,
22037                },
22038            )
22039        } else if fa_vec && head_dim <= 256 {
22040            let gqa = (n_head / n_head_kv).max(1) as u32;
22041            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
22042            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
22043            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
22044            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
22045            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
22046            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
22047            // dequant each tile ONCE per block.
22048            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
22049            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
22050            // there by 12x — latency, not bandwidth, rules small KV).
22051            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22052            let smem_tkv = *SMEM_TKV.get_or_init(|| {
22053                std::env::var("MEMRA_FA_SMEM_TKV")
22054                    .ok()
22055                    .and_then(|v| v.parse().ok())
22056                    .unwrap_or_else(|| {
22057                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22058                    })
22059            });
22060            if fa_v4_at(t_kv) && head_dim == 256 {
22061                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
22062                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
22063                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
22064                let v4name = match fa_v4_mode() {
22065                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
22066                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
22067                    _ if deep => "fa_decode_vec_q_v4_deep",
22068                    _ => "fa_decode_vec_q_v4",
22069                };
22070                let fv = if g {
22071                    self.func_g(v4name)
22072                } else {
22073                    self.func(v4name)
22074                };
22075                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
22076                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
22077                let shmem = (if deep { 12160 } else { 11520 }
22078                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22079                use cudarc::driver::sys::CUfunction_attribute_enum as A;
22080                fv.set_attribute(
22081                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22082                    shmem as i32,
22083                )?;
22084                (
22085                    fv,
22086                    LaunchConfig {
22087                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22088                        block_dim: (32, gqa, 1),
22089                        shared_mem_bytes: shmem,
22090                    },
22091                )
22092            } else if fa_v3_active(head_dim) {
22093                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
22094                // smem = sV only (half of v2's).
22095                let fv = if g {
22096                    self.func_g("fa_decode_vec_q_v3")
22097                } else {
22098                    self.func("fa_decode_vec_q_v3")
22099                };
22100                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
22101                (
22102                    fv,
22103                    LaunchConfig {
22104                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22105                        block_dim: (32, gqa, 1),
22106                        shared_mem_bytes: shmem,
22107                    },
22108                )
22109            } else if fa_v2_on() {
22110                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
22111                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
22112                // partials; same 32KB sK+sV tile as the smem twin.
22113                let fv = if g {
22114                    self.func_g("fa_decode_vec_q_v2")
22115                } else {
22116                    self.func("fa_decode_vec_q_v2")
22117                };
22118                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22119                (
22120                    fv,
22121                    LaunchConfig {
22122                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22123                        block_dim: (32, gqa, 1),
22124                        shared_mem_bytes: shmem,
22125                    },
22126                )
22127            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22128            {
22129                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
22130                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
22131                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
22132                let fv = if g {
22133                    self.func_g("fa_decode_vec_q_smem")
22134                } else {
22135                    self.func("fa_decode_vec_q_smem")
22136                };
22137                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22138                use cudarc::driver::sys::CUfunction_attribute_enum as A;
22139                fv.set_attribute(
22140                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22141                    shmem as i32,
22142                )?;
22143                (
22144                    fv,
22145                    LaunchConfig {
22146                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22147                        block_dim: (32, gqa, 1),
22148                        shared_mem_bytes: shmem,
22149                    },
22150                )
22151            } else {
22152                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
22153                // dequant, zero dynamic shared memory.
22154                let fv = if g {
22155                    self.func_g("fa_decode_vec_q")
22156                } else {
22157                    self.func("fa_decode_vec_q")
22158                };
22159                (
22160                    fv,
22161                    LaunchConfig {
22162                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22163                        block_dim: (32, gqa, 1),
22164                        shared_mem_bytes: 0,
22165                    },
22166                )
22167            }
22168        } else {
22169            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
22170            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
22171            return self.fa_decode_scalar_unified(
22172                q,
22173                k,
22174                v,
22175                o,
22176                head_dim,
22177                n_head,
22178                n_head_kv,
22179                t_kv,
22180                None,
22181                scale,
22182                n_splits,
22183                if fa_vec { sp } else { 256 },
22184                k_tok_bytes,
22185                v_tok_bytes,
22186                g,
22187                part_o,
22188                part_m,
22189                part_l,
22190                None,
22191            );
22192        };
22193        let __s_b = self.gpu.stream();
22194        let mut b = __s_b.launch_builder(&f);
22195        b.arg(q)
22196            .arg(k)
22197            .arg(v)
22198            .arg(&mut *part_o)
22199            .arg(&mut *part_m)
22200            .arg(&mut *part_l)
22201            .arg(&hd)
22202            .arg(&nh)
22203            .arg(&nhkv)
22204            .arg(&tkvi)
22205            .arg(&scale)
22206            .arg(&nsp)
22207            .arg(&ktb)
22208            .arg(&vtb);
22209        unsafe {
22210            b.launch(cfg)?;
22211        }
22212        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
22213        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
22214        let (fc, cfg2) = (
22215            if g {
22216                self.func_g("fa_decode_combine_f32")
22217            } else {
22218                self.fa_func("fa_decode_combine_f32", head_dim)
22219            },
22220            LaunchConfig {
22221                grid_dim: (n_head as u32, 1, 1),
22222                block_dim: (head_dim as u32, 1, 1),
22223                shared_mem_bytes: 0,
22224            },
22225        );
22226        let __s_b2 = self.gpu.stream();
22227        let mut b2 = __s_b2.launch_builder(&fc);
22228        b2.arg(&*part_o)
22229            .arg(&*part_m)
22230            .arg(&*part_l)
22231            .arg(o)
22232            .arg(&hd)
22233            .arg(&nh)
22234            .arg(&nsp);
22235        unsafe {
22236            b2.launch(cfg2)?;
22237        }
22238        Ok(())
22239    }
22240
22241    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
22242    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
22243    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
22244    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
22245    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
22246    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
22247    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
22248    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
22249    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
22250    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
22251    #[allow(clippy::too_many_arguments)]
22252    pub fn fa_decode_batch_seqs_v4(
22253        &self,
22254        q: &CudaSlice<f32>,
22255        kv_ptrs: &cudarc::driver::CudaView<u64>,
22256        pos_seq: &CudaSlice<i32>,
22257        o: &mut CudaSlice<f32>,
22258        head_dim: usize,
22259        n_head: usize,
22260        n_head_kv: usize,
22261        b_n: usize,
22262        t_kv_max: usize,
22263        scale: f32,
22264        split_keys: usize,
22265        k_tok_bytes: usize,
22266        v_tok_bytes: usize,
22267    ) -> Result<(), Box<dyn std::error::Error>> {
22268        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22269        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22270        let o_len = b_n * n_head * n_splits_max * head_dim;
22271        let ml_len = b_n * n_head * n_splits_max;
22272        let mut part_guard = self.fa_part_pool.lock().unwrap();
22273        if part_guard
22274            .as_ref()
22275            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22276            .unwrap_or(true)
22277        {
22278            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22279            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22280            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22281            // later live allocations land at those addresses, and the next graph REPLAY writes
22282            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22283            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22284            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22285            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22286            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22287            // (total retired < final size).
22288            let old = part_guard.take();
22289            let (co, cm) = old
22290                .as_ref()
22291                .map(|pp| (pp.0.len(), pp.1.len()))
22292                .unwrap_or((0, 0));
22293            if let Some(old) = old {
22294                self.fa_part_retired.lock().unwrap().push(old);
22295            }
22296            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22297                eprintln!(
22298                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22299                    co, o_len, cm, ml_len
22300                );
22301            }
22302            *part_guard = Some((
22303                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22304                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22305                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22306            ));
22307        }
22308        let pg = part_guard.as_mut().unwrap();
22309        self.gpu
22310            .stream()
22311            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22312        self.gpu
22313            .stream()
22314            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22315        self.gpu
22316            .stream()
22317            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22318        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22319        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22320        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22321        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22322        let gqa = (n_head / n_head_kv).max(1) as u32;
22323        let f = self.func("fa_decode_vec_q_seqs_v4");
22324        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
22325        let shmem = (11520 + 32 * head_dim * 2) as u32;
22326        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22327        f.set_attribute(
22328            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22329            shmem as i32,
22330        )?;
22331        let cfg = LaunchConfig {
22332            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22333            block_dim: (32, gqa, 1),
22334            shared_mem_bytes: shmem,
22335        };
22336        {
22337            let __s_b = self.gpu.stream();
22338            let mut b = __s_b.launch_builder(&f);
22339            b.arg(q)
22340                .arg(kv_ptrs)
22341                .arg(pos_seq)
22342                .arg(&mut *part_o)
22343                .arg(&mut *part_m)
22344                .arg(&mut *part_l)
22345                .arg(&hd)
22346                .arg(&nh)
22347                .arg(&nhkv)
22348                .arg(&scale)
22349                .arg(&nspm)
22350                .arg(&spk)
22351                .arg(&ktb)
22352                .arg(&vtb);
22353            unsafe {
22354                b.launch(cfg)?;
22355            }
22356        }
22357        let fc = self.func("fa_decode_combine_seqs");
22358        let cfg2 = LaunchConfig {
22359            grid_dim: (n_head as u32, b_n as u32, 1),
22360            block_dim: (head_dim as u32, 1, 1),
22361            shared_mem_bytes: 0,
22362        };
22363        let __s_b2 = self.gpu.stream();
22364        let mut b2 = __s_b2.launch_builder(&fc);
22365        b2.arg(&*part_o)
22366            .arg(&*part_m)
22367            .arg(&*part_l)
22368            .arg(o)
22369            .arg(&hd)
22370            .arg(&nh)
22371            .arg(pos_seq)
22372            .arg(&nspm)
22373            .arg(&spk);
22374        unsafe {
22375            b2.launch(cfg2)?;
22376        }
22377        Ok(())
22378    }
22379
22380    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
22381    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
22382    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
22383    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
22384    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
22385    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
22386    #[allow(clippy::too_many_arguments)]
22387    pub fn append_kv_quantized_seqs(
22388        &self,
22389        k_rows: &CudaSlice<f32>,
22390        v_rows: &CudaSlice<f32>,
22391        kv_ptrs: &cudarc::driver::CudaView<u64>,
22392        pos_seq: &CudaSlice<i32>,
22393        b_n: usize,
22394        kv_dim_k: usize,
22395        kv_dim_v: usize,
22396        k_tok_bytes: usize,
22397        v_tok_bytes: usize,
22398    ) -> Result<(), Box<dyn std::error::Error>> {
22399        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22400        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22401        let cfg = LaunchConfig {
22402            grid_dim: (nblk, b_n as u32, 1),
22403            block_dim: (32, 1, 1),
22404            shared_mem_bytes: 0,
22405        };
22406        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22407        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22408        let __s_b = self.gpu.stream();
22409        let mut b = __s_b.launch_builder(&f);
22410        b.arg(k_rows)
22411            .arg(v_rows)
22412            .arg(kv_ptrs)
22413            .arg(pos_seq)
22414            .arg(&kdk)
22415            .arg(&kdv)
22416            .arg(&ktb)
22417            .arg(&vtb);
22418        unsafe {
22419            b.launch(cfg)?;
22420        }
22421        Ok(())
22422    }
22423
22424    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
22425    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
22426    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
22427    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
22428    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
22429    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22430        std::env::var("MEMRA_NO_FA_VEC").is_err()
22431            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22432            && base_len + 1 >= fa_vec_min_tkv()
22433            && head_dim <= 256
22434            && head_dim % 32 == 0
22435    }
22436
22437    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
22438    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
22439    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
22440    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
22441    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
22442    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
22443    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
22444    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
22445    #[allow(clippy::too_many_arguments)]
22446    pub fn fa_decode_rows(
22447        &self,
22448        q: &CudaSlice<f32>,
22449        k: &cudarc::driver::CudaView<u8>,
22450        v: &cudarc::driver::CudaView<u8>,
22451        o: &mut CudaSlice<f32>,
22452        head_dim: usize,
22453        n_head: usize,
22454        n_head_kv: usize,
22455        base_len: usize,
22456        t: usize,
22457        scale: f32,
22458        k_tok_bytes: usize,
22459        v_tok_bytes: usize,
22460        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
22461        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
22462        // keep the host arg. None is a bug for hd512 (asserted below).
22463        base_dev: Option<(&CudaSlice<i32>, i32)>,
22464        // K and V planes hold the same values (gemma globals, wv:=wk): pick
22465        // the _kv twin — V plane never read, value rides the q8_0 key dq.
22466        kv_shared: bool,
22467        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
22468        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
22469        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
22470        g: bool,
22471        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
22472        // (hd512 path) — the standalone quantize launch folds away.
22473        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22474    ) -> Result<(), Box<dyn std::error::Error>> {
22475        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22476        let t_kv_max = base_len + t; // LAST row's key bound
22477        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
22478        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
22479        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
22480        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
22481        // (parity law), so the partition is freely tunable — verify and decode move together.
22482        if head_dim == 512 {
22483            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22484            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
22485            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
22486            let v = *SP512.get_or_init(|| {
22487                std::env::var("MEMRA_FA_SP512")
22488                    .ok()
22489                    .and_then(|x| x.parse().ok())
22490                    .unwrap_or(0)
22491            });
22492            sp = if v >= 8 {
22493                v
22494            } else {
22495                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22496            };
22497        }
22498        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22499        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22500        let gqa = (n_head / n_head_kv).max(1) as u32;
22501        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
22502        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
22503        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
22504        // the different partition changes the combine's FP order (greedy tie flips at depth;
22505        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
22506        // consecutive rows by their OWN ladder value and launch once per group — each row then
22507        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
22508        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
22509        // sp override is t_kv-independent by construction).
22510        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
22511        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22512            groups.push((0, t, sp));
22513        } else {
22514            let mut r0 = 0usize;
22515            while r0 < t {
22516                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22517                let mut r1 = r0 + 1;
22518                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22519                    r1 += 1;
22520                }
22521                groups.push((r0, r1 - r0, sp_g));
22522                r0 = r1;
22523            }
22524        }
22525        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
22526        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
22527        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
22528        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22529        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22530            std::env::var("MEMRA_FA_SMEM_TKV")
22531                .ok()
22532                .and_then(|v| v.parse().ok())
22533                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22534        });
22535        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22536        let v3 = fa_v3_active(head_dim);
22537        let smem_rows =
22538            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22539        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
22540        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
22541        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
22542        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
22543        let _ = kv_shared;
22544        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
22545        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
22546        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22547        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
22548        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
22549        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
22550        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
22551        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
22552        // (kv_head, split) stages its tile once and loops the rows over it — kills the
22553        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
22554        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
22555        // shared by every hd512 caller through this wrapper (decode+verify flip together;
22556        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
22557        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
22558        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
22559        // not unpack-bound; jsonl 2026-07-14.
22560        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22561        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
22562        let tb512 = head_dim == 512
22563            && sp <= 32
22564            && n_head / n_head_kv.max(1) <= 16
22565            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22566        let fname = if tb512 {
22567            "fa_decode_vec_q_rows_v4_512_tb"
22568        } else if i2 {
22569            "fa_decode_vec_q_rows_dpl16_i2"
22570        } else if head_dim == 512 {
22571            "fa_decode_vec_q_rows_dpl16"
22572        }
22573        // gemma globals (parity law)
22574        else if v4 {
22575            "fa_decode_vec_q_rows_v4"
22576        } else if v3 {
22577            "fa_decode_vec_q_rows_v3"
22578        } else if fa_v2_on() {
22579            "fa_decode_vec_q_rows_v2"
22580        } else if smem_rows {
22581            "fa_decode_vec_q_rows_smem"
22582        } else {
22583            "fa_decode_vec_q_rows"
22584        };
22585        let f = if head_dim == 512 {
22586            self.fa_func(fname, head_dim)
22587        } else if g {
22588            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
22589            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
22590            // g-module rows against decode's g-module v4 — different programs, short-VG
22591            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
22592            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
22593            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
22594            // dq macros are format-aware.
22595            self.func_g(if smem_rows {
22596                "fa_decode_vec_q_rows"
22597            } else {
22598                fname
22599            })
22600        } else {
22601            self.func(fname)
22602        };
22603        let shmem = if tb512 {
22604            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
22605            let gk = Self::gkv_on();
22606            let sh =
22607                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22608            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22609            f.set_attribute(
22610                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22611                sh as i32,
22612            )?;
22613            sh
22614        } else if v4 || v3 || smem_rows || fa_v2_on() {
22615            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
22616            let sh = (if v4 {
22617                11520 + 32 * head_dim * if g { 1 } else { 2 }
22618            } else if v3 {
22619                32 * head_dim * 2
22620            } else {
22621                2 * 32 * head_dim * 2
22622            }) as u32;
22623            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22624            f.set_attribute(
22625                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22626                sh as i32,
22627            )?;
22628            sh
22629        } else {
22630            0
22631        };
22632        // Per-GROUP launches (single group in the common case — identical to the pre-fix
22633        // single launch there): each group gets its own partials (the rows kernel indexes
22634        // partials by its LOCAL grid.z row) and q/o row-offset views.
22635        for &(r0, t_g, sp_g) in &groups {
22636            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22637            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22638            let base_i = (base_len + r0) as i32;
22639            let o_len = t_g * n_head * n_splits_g * head_dim;
22640            let ml_len = t_g * n_head * n_splits_g;
22641            let mut part_guard = self.fa_part_pool.lock().unwrap();
22642            if part_guard
22643                .as_ref()
22644                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22645                .unwrap_or(true)
22646            {
22647                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22648                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22649                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22650                // later live allocations land at those addresses, and the next graph REPLAY writes
22651                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22652                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22653                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22654                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22655                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22656                // (total retired < final size).
22657                let old = part_guard.take();
22658                let (co, cm) = old
22659                    .as_ref()
22660                    .map(|pp| (pp.0.len(), pp.1.len()))
22661                    .unwrap_or((0, 0));
22662                if let Some(old) = old {
22663                    self.fa_part_retired.lock().unwrap().push(old);
22664                }
22665                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22666                    eprintln!(
22667                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22668                        co, o_len, cm, ml_len
22669                    );
22670                }
22671                *part_guard = Some((
22672                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22673                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22674                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22675                ));
22676            }
22677            let pg = part_guard.as_mut().unwrap();
22678            self.gpu
22679                .stream()
22680                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22681            self.gpu
22682                .stream()
22683                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22684            self.gpu
22685                .stream()
22686                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22687            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22688            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22689            let qv = self.view(q, t * n_head * head_dim);
22690            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22691            let cfg = LaunchConfig {
22692                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22693                block_dim: (32, gqa, 1),
22694                shared_mem_bytes: shmem,
22695            };
22696            {
22697                let __s_b = self.gpu.stream();
22698                let mut b = __s_b.launch_builder(&f);
22699                if tb512 {
22700                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
22701                    let (bd, plus) =
22702                        base_dev.expect("hd512 rows twin requires a device base counter");
22703                    let plus_g = plus + r0 as i32;
22704                    let nr = t_g as i32;
22705                    if Self::pdl_on() && Self::pdl_wb_on() {
22706                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
22707                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22708                        let s = &self.gpu.stream();
22709                        let (pq, _b0) = q_g.device_ptr(s);
22710                        let (pk, _b1) = k.device_ptr(s);
22711                        let (pv, _b2) = v.device_ptr(s);
22712                        let (po, _b3) = part_o.device_ptr_mut(s);
22713                        let (pm, _b4) = part_m.device_ptr_mut(s);
22714                        let (pl, _b5) = part_l.device_ptr_mut(s);
22715                        let (pb, _b6) = bd.device_ptr(s);
22716                        let mut ps = [
22717                            &pq as *const _ as *mut std::ffi::c_void,
22718                            &pk as *const _ as *mut _,
22719                            &pv as *const _ as *mut _,
22720                            &po as *const _ as *mut _,
22721                            &pm as *const _ as *mut _,
22722                            &pl as *const _ as *mut _,
22723                            &hd as *const _ as *mut _,
22724                            &nh as *const _ as *mut _,
22725                            &nhkv as *const _ as *mut _,
22726                            &pb as *const _ as *mut _,
22727                            &plus_g as *const _ as *mut _,
22728                            &scale as *const _ as *mut _,
22729                            &nspm as *const _ as *mut _,
22730                            &spk as *const _ as *mut _,
22731                            &ktb as *const _ as *mut _,
22732                            &vtb as *const _ as *mut _,
22733                            &nr as *const _ as *mut _,
22734                        ];
22735                        unsafe {
22736                            self.launch_pdl_flash(
22737                                Self::gkv_on(),
22738                                "fa_decode_vec_q_rows_v4_512_tb",
22739                                (n_head_kv as u32, n_splits_g as u32, 1),
22740                                (32, gqa, 1),
22741                                shmem,
22742                                &mut ps,
22743                            )?;
22744                        }
22745                    } else {
22746                        let cfg_tb = LaunchConfig {
22747                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22748                            block_dim: (32, gqa, 1),
22749                            shared_mem_bytes: shmem,
22750                        };
22751                        b.arg(&q_g)
22752                            .arg(k)
22753                            .arg(v)
22754                            .arg(&mut *part_o)
22755                            .arg(&mut *part_m)
22756                            .arg(&mut *part_l)
22757                            .arg(&hd)
22758                            .arg(&nh)
22759                            .arg(&nhkv)
22760                            .arg(bd)
22761                            .arg(&plus_g)
22762                            .arg(&scale)
22763                            .arg(&nspm)
22764                            .arg(&spk)
22765                            .arg(&ktb)
22766                            .arg(&vtb)
22767                            .arg(&nr);
22768                        unsafe {
22769                            b.launch(cfg_tb)?;
22770                        }
22771                    }
22772                } else if head_dim == 512 {
22773                    let (bd, plus) =
22774                        base_dev.expect("hd512 rows twin requires a device base counter");
22775                    let plus_g = plus + r0 as i32;
22776                    b.arg(&q_g)
22777                        .arg(k)
22778                        .arg(v)
22779                        .arg(&mut *part_o)
22780                        .arg(&mut *part_m)
22781                        .arg(&mut *part_l)
22782                        .arg(&hd)
22783                        .arg(&nh)
22784                        .arg(&nhkv)
22785                        .arg(bd)
22786                        .arg(&plus_g)
22787                        .arg(&scale)
22788                        .arg(&nspm)
22789                        .arg(&spk)
22790                        .arg(&ktb)
22791                        .arg(&vtb);
22792                    unsafe {
22793                        b.launch(cfg)?;
22794                    }
22795                } else {
22796                    b.arg(&q_g)
22797                        .arg(k)
22798                        .arg(v)
22799                        .arg(&mut *part_o)
22800                        .arg(&mut *part_m)
22801                        .arg(&mut *part_l)
22802                        .arg(&hd)
22803                        .arg(&nh)
22804                        .arg(&nhkv)
22805                        .arg(&base_i)
22806                        .arg(&scale)
22807                        .arg(&nspm)
22808                        .arg(&spk)
22809                        .arg(&ktb)
22810                        .arg(&vtb);
22811                    unsafe {
22812                        b.launch(cfg)?;
22813                    }
22814                }
22815            }
22816            let cfg2 = LaunchConfig {
22817                grid_dim: (n_head as u32, t_g as u32, 1),
22818                block_dim: (head_dim as u32, 1, 1),
22819                shared_mem_bytes: 0,
22820            };
22821            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22822            if head_dim == 512 {
22823                // device-len combine (shared by verify/eager/graph — parity by symbol): the
22824                // per-row n_splits derives from the SAME counter the rows kernel read.
22825                let (bd, plus) = base_dev.unwrap();
22826                let plus_g = plus + r0 as i32;
22827                if let Some((oq, od)) = q8_out.as_mut() {
22828                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
22829                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22830                    if Self::pdl_on() && Self::pdl_wb_on() {
22831                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
22832                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22833                        let s = &self.gpu.stream();
22834                        let (po, _g0) = part_o.device_ptr(s);
22835                        let (pm, _g1) = part_m.device_ptr(s);
22836                        let (pl, _g2) = part_l.device_ptr(s);
22837                        let (pq, _g3) = oq.device_ptr_mut(s);
22838                        let (pd, _g4) = od.device_ptr_mut(s);
22839                        let (pb, _g5) = bd.device_ptr(s);
22840                        let mut ps = [
22841                            &po as *const _ as *mut std::ffi::c_void,
22842                            &pm as *const _ as *mut _,
22843                            &pl as *const _ as *mut _,
22844                            &pq as *const _ as *mut _,
22845                            &pd as *const _ as *mut _,
22846                            &hd as *const _ as *mut _,
22847                            &nh as *const _ as *mut _,
22848                            &pb as *const _ as *mut _,
22849                            &plus_g as *const _ as *mut _,
22850                            &nspm as *const _ as *mut _,
22851                            &spk as *const _ as *mut _,
22852                        ];
22853                        unsafe {
22854                            self.launch_pdl_flash(
22855                                Self::gkv_on(),
22856                                "fa_decode_combine_rows_dc_q8_1",
22857                                cfg2.grid_dim,
22858                                cfg2.block_dim,
22859                                0,
22860                                &mut ps,
22861                            )?;
22862                        }
22863                        continue;
22864                    }
22865                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22866                    let __s_b2 = self.gpu.stream();
22867                    let mut b2 = __s_b2.launch_builder(&fc);
22868                    b2.arg(&*part_o)
22869                        .arg(&*part_m)
22870                        .arg(&*part_l)
22871                        .arg(&mut **oq)
22872                        .arg(&mut **od)
22873                        .arg(&hd)
22874                        .arg(&nh)
22875                        .arg(bd)
22876                        .arg(&plus_g)
22877                        .arg(&nspm)
22878                        .arg(&spk);
22879                    unsafe {
22880                        b2.launch(cfg2)?;
22881                    }
22882                    continue;
22883                }
22884                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22885                let __s_b2 = self.gpu.stream();
22886                let mut b2 = __s_b2.launch_builder(&fc);
22887                b2.arg(&*part_o)
22888                    .arg(&*part_m)
22889                    .arg(&*part_l)
22890                    .arg(&mut o_g)
22891                    .arg(&hd)
22892                    .arg(&nh)
22893                    .arg(bd)
22894                    .arg(&plus_g)
22895                    .arg(&nspm)
22896                    .arg(&spk);
22897                unsafe {
22898                    b2.launch(cfg2)?;
22899                }
22900            } else {
22901                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
22902                // leave the caller's pair unwritten (consumer would read garbage).
22903                assert!(
22904                    q8_out.is_none(),
22905                    "rows q8 emit requires the hd512 dc combine"
22906                );
22907                let fc = self.func("fa_decode_combine_rows");
22908                let __s_b2 = self.gpu.stream();
22909                let mut b2 = __s_b2.launch_builder(&fc);
22910                b2.arg(&*part_o)
22911                    .arg(&*part_m)
22912                    .arg(&*part_l)
22913                    .arg(&mut o_g)
22914                    .arg(&hd)
22915                    .arg(&nh)
22916                    .arg(&base_i)
22917                    .arg(&nspm)
22918                    .arg(&spk);
22919                unsafe {
22920                    b2.launch(cfg2)?;
22921                }
22922            }
22923        }
22924        Ok(())
22925    }
22926
22927    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
22928    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
22929    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
22930    #[allow(clippy::too_many_arguments)]
22931    pub fn fa_decode_rows_w(
22932        &self,
22933        q: &CudaSlice<f32>,
22934        k: &cudarc::driver::CudaView<u8>,
22935        v: &cudarc::driver::CudaView<u8>,
22936        o: &mut CudaSlice<f32>,
22937        head_dim: usize,
22938        n_head: usize,
22939        n_head_kv: usize,
22940        base_dev: &CudaSlice<i32>,
22941        base_plus: i32,
22942        t: usize,
22943        scale: f32,
22944        window: usize,
22945        k_tok_bytes: usize,
22946        v_tok_bytes: usize,
22947        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22948    ) -> Result<(), Box<dyn std::error::Error>> {
22949        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
22950        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
22951        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
22952        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
22953        debug_assert!(head_dim == 256);
22954        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
22955        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
22956        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
22957        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
22958        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
22959        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
22960        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
22961        let sp = {
22962            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22963            let v = *SPW.get_or_init(|| {
22964                std::env::var("MEMRA_FA_SPW")
22965                    .ok()
22966                    .and_then(|x| x.parse().ok())
22967                    .unwrap_or(0)
22968            });
22969            if v >= 8 {
22970                v
22971            } else {
22972                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22973            }
22974        };
22975        let n_splits_max = (window + sp - 1) / sp;
22976        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22977        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22978        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22979        let gqa = (n_head / n_head_kv).max(1) as u32;
22980        let o_len = t * n_head * n_splits_max * head_dim;
22981        let ml_len = t * n_head * n_splits_max;
22982        let mut part_guard = self.fa_part_pool.lock().unwrap();
22983        if part_guard
22984            .as_ref()
22985            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22986            .unwrap_or(true)
22987        {
22988            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22989            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22990            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22991            // later live allocations land at those addresses, and the next graph REPLAY writes
22992            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22993            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22994            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22995            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22996            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22997            // (total retired < final size).
22998            let old = part_guard.take();
22999            let (co, cm) = old
23000                .as_ref()
23001                .map(|pp| (pp.0.len(), pp.1.len()))
23002                .unwrap_or((0, 0));
23003            if let Some(old) = old {
23004                self.fa_part_retired.lock().unwrap().push(old);
23005            }
23006            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23007                eprintln!(
23008                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23009                    co, o_len, cm, ml_len
23010                );
23011            }
23012            *part_guard = Some((
23013                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23014                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23015                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23016            ));
23017        }
23018        let pg = part_guard.as_mut().unwrap();
23019        self.gpu
23020            .stream()
23021            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23022        self.gpu
23023            .stream()
23024            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23025        self.gpu
23026            .stream()
23027            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23028        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23029        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
23030        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
23031        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
23032        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
23033        // floor (deep-ctx broadcast win); register twin between.
23034        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23035        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
23036            std::env::var("MEMRA_FA_SMEM_TKV")
23037                .ok()
23038                .and_then(|v| v.parse().ok())
23039                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23040        });
23041        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
23042        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
23043        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
23044        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
23045        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
23046        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23047        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
23048        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
23049        // per (lane, format-module) keeps parity structural; the old register-i2 detour
23050        // (-33%) is retired.
23051        let wg = Self::wkv_on();
23052        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
23053        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
23054        let sp2 =
23055            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23056        if sp2 {
23057            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23058            if Self::pdl_on() && Self::pdl_wb_on() {
23059                // wave-B2b: flavor mirrors wg.
23060                use cudarc::driver::{DevicePtr, DevicePtrMut};
23061                let s = &self.gpu.stream();
23062                let (pq, _b0) = q.device_ptr(s);
23063                let (pk, _b1) = k.device_ptr(s);
23064                let (pv, _b2) = v.device_ptr(s);
23065                let (po, _b3) = part_o.device_ptr_mut(s);
23066                let (pm, _b4) = part_m.device_ptr_mut(s);
23067                let (pl, _b5) = part_l.device_ptr_mut(s);
23068                let (pb, _b6) = base_dev.device_ptr(s);
23069                let mut ps = [
23070                    &pq as *const _ as *mut std::ffi::c_void,
23071                    &pk as *const _ as *mut _,
23072                    &pv as *const _ as *mut _,
23073                    &po as *const _ as *mut _,
23074                    &pm as *const _ as *mut _,
23075                    &pl as *const _ as *mut _,
23076                    &hd as *const _ as *mut _,
23077                    &nh as *const _ as *mut _,
23078                    &nhkv as *const _ as *mut _,
23079                    &pb as *const _ as *mut _,
23080                    &base_plus as *const _ as *mut _,
23081                    &scale as *const _ as *mut _,
23082                    &nspm as *const _ as *mut _,
23083                    &spk as *const _ as *mut _,
23084                    &ktb as *const _ as *mut _,
23085                    &vtb as *const _ as *mut _,
23086                    &wini as *const _ as *mut _,
23087                ];
23088                unsafe {
23089                    self.launch_pdl_flash(
23090                        wg,
23091                        "fa_decode_vec_q_rows_v4_w_sp",
23092                        (n_head_kv as u32, n_splits_max as u32, t as u32),
23093                        (32, gqa + 1, 1),
23094                        sh,
23095                        &mut ps,
23096                    )?;
23097                }
23098            } else {
23099                let f = if wg {
23100                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23101                } else {
23102                    self.func("fa_decode_vec_q_rows_v4_w_sp")
23103                };
23104                f.set_attribute(
23105                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23106                    sh as i32,
23107                )?;
23108                let cfg = LaunchConfig {
23109                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23110                    block_dim: (32, gqa + 1, 1),
23111                    shared_mem_bytes: sh,
23112                };
23113                let __s_b = self.gpu.stream();
23114                let mut b = __s_b.launch_builder(&f);
23115                b.arg(q)
23116                    .arg(k)
23117                    .arg(v)
23118                    .arg(&mut *part_o)
23119                    .arg(&mut *part_m)
23120                    .arg(&mut *part_l)
23121                    .arg(&hd)
23122                    .arg(&nh)
23123                    .arg(&nhkv)
23124                    .arg(base_dev)
23125                    .arg(&base_plus)
23126                    .arg(&scale)
23127                    .arg(&nspm)
23128                    .arg(&spk)
23129                    .arg(&ktb)
23130                    .arg(&vtb)
23131                    .arg(&wini);
23132                unsafe {
23133                    b.launch(cfg)?;
23134                }
23135            }
23136        } else {
23137            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23138                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
23139                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23140                use cudarc::driver::{DevicePtr, DevicePtrMut};
23141                let s = &self.gpu.stream();
23142                let (pq, _b0) = q.device_ptr(s);
23143                let (pk, _b1) = k.device_ptr(s);
23144                let (pv, _b2) = v.device_ptr(s);
23145                let (po, _b3) = part_o.device_ptr_mut(s);
23146                let (pm, _b4) = part_m.device_ptr_mut(s);
23147                let (pl, _b5) = part_l.device_ptr_mut(s);
23148                let (pb, _b6) = base_dev.device_ptr(s);
23149                let mut ps = [
23150                    &pq as *const _ as *mut std::ffi::c_void,
23151                    &pk as *const _ as *mut _,
23152                    &pv as *const _ as *mut _,
23153                    &po as *const _ as *mut _,
23154                    &pm as *const _ as *mut _,
23155                    &pl as *const _ as *mut _,
23156                    &hd as *const _ as *mut _,
23157                    &nh as *const _ as *mut _,
23158                    &nhkv as *const _ as *mut _,
23159                    &pb as *const _ as *mut _,
23160                    &base_plus as *const _ as *mut _,
23161                    &scale as *const _ as *mut _,
23162                    &nspm as *const _ as *mut _,
23163                    &spk as *const _ as *mut _,
23164                    &ktb as *const _ as *mut _,
23165                    &vtb as *const _ as *mut _,
23166                    &wini as *const _ as *mut _,
23167                ];
23168                unsafe {
23169                    self.launch_pdl_flash(
23170                        wg,
23171                        "fa_decode_vec_q_rows_v4_w",
23172                        (n_head_kv as u32, n_splits_max as u32, t as u32),
23173                        (32, gqa, 1),
23174                        sh,
23175                        &mut ps,
23176                    )?;
23177                }
23178            } else {
23179                let pick = |name: &str| {
23180                    if wg {
23181                        self.func_g(name)
23182                    } else {
23183                        self.func(name)
23184                    }
23185                };
23186                let (f, sh) = if fa_v4_at(window) {
23187                    let f = pick("fa_decode_vec_q_rows_v4_w");
23188                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23189                } else if smem_tkv > 0 && window >= smem_tkv {
23190                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
23191                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
23192                    (
23193                        pick("fa_decode_vec_q_rows_smem_w"),
23194                        (2 * 32 * head_dim * 2) as u32,
23195                    )
23196                } else {
23197                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23198                };
23199                f.set_attribute(
23200                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23201                    sh as i32,
23202                )?;
23203                let cfg = LaunchConfig {
23204                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23205                    block_dim: (32, gqa, 1),
23206                    shared_mem_bytes: sh,
23207                };
23208                let __s_b = self.gpu.stream();
23209                let mut b = __s_b.launch_builder(&f);
23210                b.arg(q)
23211                    .arg(k)
23212                    .arg(v)
23213                    .arg(&mut *part_o)
23214                    .arg(&mut *part_m)
23215                    .arg(&mut *part_l)
23216                    .arg(&hd)
23217                    .arg(&nh)
23218                    .arg(&nhkv)
23219                    .arg(base_dev)
23220                    .arg(&base_plus)
23221                    .arg(&scale)
23222                    .arg(&nspm)
23223                    .arg(&spk)
23224                    .arg(&ktb)
23225                    .arg(&vtb)
23226                    .arg(&wini);
23227                unsafe {
23228                    b.launch(cfg)?;
23229                }
23230            }
23231        }
23232        let cfg2 = LaunchConfig {
23233            grid_dim: (n_head as u32, t as u32, 1),
23234            block_dim: (head_dim as u32, 1, 1),
23235            shared_mem_bytes: 0,
23236        };
23237        if let Some((oq, od)) = q8_out {
23238            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
23239            // consumes the pair directly; the standalone quantize launch folds away.
23240            if Self::pdl_on() && Self::pdl_wb_on() {
23241                // wave-B2: flavor mirrors the builder's wg choice.
23242                use cudarc::driver::{DevicePtr, DevicePtrMut};
23243                let s = &self.gpu.stream();
23244                let (po, _g0) = part_o.device_ptr(s);
23245                let (pm, _g1) = part_m.device_ptr(s);
23246                let (pl, _g2) = part_l.device_ptr(s);
23247                let (pq, _g3) = oq.device_ptr_mut(s);
23248                let (pd, _g4) = od.device_ptr_mut(s);
23249                let mut ps = [
23250                    &po as *const _ as *mut std::ffi::c_void,
23251                    &pm as *const _ as *mut _,
23252                    &pl as *const _ as *mut _,
23253                    &pq as *const _ as *mut _,
23254                    &pd as *const _ as *mut _,
23255                    &hd as *const _ as *mut _,
23256                    &nh as *const _ as *mut _,
23257                    &nspm as *const _ as *mut _,
23258                    &spk as *const _ as *mut _,
23259                    &wini as *const _ as *mut _,
23260                ];
23261                unsafe {
23262                    self.launch_pdl_flash(
23263                        wg,
23264                        "fa_decode_combine_rows_w_q8_1",
23265                        cfg2.grid_dim,
23266                        cfg2.block_dim,
23267                        0,
23268                        &mut ps,
23269                    )?;
23270                }
23271                return Ok(());
23272            }
23273            let fc = if wg {
23274                self.func_g("fa_decode_combine_rows_w_q8_1")
23275            } else {
23276                self.func("fa_decode_combine_rows_w_q8_1")
23277            };
23278            let __s_b2 = self.gpu.stream();
23279            let mut b2 = __s_b2.launch_builder(&fc);
23280            b2.arg(&*part_o)
23281                .arg(&*part_m)
23282                .arg(&*part_l)
23283                .arg(oq)
23284                .arg(od)
23285                .arg(&hd)
23286                .arg(&nh)
23287                .arg(&nspm)
23288                .arg(&spk)
23289                .arg(&wini);
23290            unsafe {
23291                b2.launch(cfg2)?;
23292            }
23293            return Ok(());
23294        }
23295        let fc = if wg {
23296            self.func_g("fa_decode_combine_rows_w")
23297        } else {
23298            self.func("fa_decode_combine_rows_w")
23299        };
23300        let __s_b2 = self.gpu.stream();
23301        let mut b2 = __s_b2.launch_builder(&fc);
23302        b2.arg(&*part_o)
23303            .arg(&*part_m)
23304            .arg(&*part_l)
23305            .arg(o)
23306            .arg(&hd)
23307            .arg(&nh)
23308            .arg(&nspm)
23309            .arg(&spk)
23310            .arg(&wini);
23311        unsafe {
23312            b2.launch(cfg2)?;
23313        }
23314        Ok(())
23315    }
23316
23317    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
23318    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
23319    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
23320    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
23321    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
23322    #[allow(clippy::too_many_arguments)]
23323    pub fn fa_decode_rows_dc(
23324        &self,
23325        q: &CudaSlice<f32>,
23326        k: &cudarc::driver::CudaView<u8>,
23327        v: &cudarc::driver::CudaView<u8>,
23328        o: &mut CudaSlice<f32>,
23329        head_dim: usize,
23330        n_head: usize,
23331        n_head_kv: usize,
23332        base_dev: &CudaSlice<i32>,
23333        t_kv_upper: usize,
23334        t: usize,
23335        scale: f32,
23336        k_tok_bytes: usize,
23337        v_tok_bytes: usize,
23338        base_plus: i32,
23339        g: bool,
23340    ) -> Result<(), Box<dyn std::error::Error>> {
23341        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23342        assert!(
23343            v4 || fa_v3_active(head_dim),
23344            "stream fa rows requires the v3 or v4 lane"
23345        );
23346        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23347        if v4 {
23348            let sp = fa_split_keys(t_kv_upper, n_head_kv);
23349            let n_splits_max = (t_kv_upper + sp - 1) / sp;
23350            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23351            let (nspm, spk) = (n_splits_max as i32, sp as i32);
23352            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23353            let gqa = (n_head / n_head_kv).max(1) as u32;
23354            let o_len = t * n_head * n_splits_max * head_dim;
23355            let ml_len = t * n_head * n_splits_max;
23356            let mut part_guard = self.fa_part_pool.lock().unwrap();
23357            if part_guard
23358                .as_ref()
23359                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23360                .unwrap_or(true)
23361            {
23362                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23363                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23364                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23365                // later live allocations land at those addresses, and the next graph REPLAY writes
23366                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23367                // output corruption began the burst after the trunk's t_kv growth first realloc'd
23368                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23369                // the baked addresses alive (single-stream: eager writes the new buffers, replays
23370                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23371                // (total retired < final size).
23372                let old = part_guard.take();
23373                let (co, cm) = old
23374                    .as_ref()
23375                    .map(|pp| (pp.0.len(), pp.1.len()))
23376                    .unwrap_or((0, 0));
23377                if let Some(old) = old {
23378                    self.fa_part_retired.lock().unwrap().push(old);
23379                }
23380                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23381                    eprintln!(
23382                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23383                        co, o_len, cm, ml_len
23384                    );
23385                }
23386                *part_guard = Some((
23387                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23388                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23389                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23390                ));
23391            }
23392            let pg = part_guard.as_mut().unwrap();
23393            self.gpu
23394                .stream()
23395                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23396            self.gpu
23397                .stream()
23398                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23399            self.gpu
23400                .stream()
23401                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23402            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23403            let f = if g {
23404                self.func_g("fa_decode_vec_q_rows_v4_dc")
23405            } else {
23406                self.func("fa_decode_vec_q_rows_v4_dc")
23407            };
23408            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23409            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23410            f.set_attribute(
23411                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23412                sh as i32,
23413            )?;
23414            let cfg = LaunchConfig {
23415                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23416                block_dim: (32, gqa, 1),
23417                shared_mem_bytes: sh,
23418            };
23419            let __s_b = self.gpu.stream();
23420            let mut b = __s_b.launch_builder(&f);
23421            b.arg(q)
23422                .arg(k)
23423                .arg(v)
23424                .arg(&mut *part_o)
23425                .arg(&mut *part_m)
23426                .arg(&mut *part_l)
23427                .arg(&hd)
23428                .arg(&nh)
23429                .arg(&nhkv)
23430                .arg(base_dev)
23431                .arg(&base_plus)
23432                .arg(&scale)
23433                .arg(&nspm)
23434                .arg(&spk)
23435                .arg(&ktb)
23436                .arg(&vtb);
23437            unsafe {
23438                b.launch(cfg)?;
23439            }
23440            let fc = self.func("fa_decode_combine_rows_dc");
23441            let cfg2 = LaunchConfig {
23442                grid_dim: (n_head as u32, t as u32, 1),
23443                block_dim: (head_dim as u32, 1, 1),
23444                shared_mem_bytes: 0,
23445            };
23446            let __s_b2 = self.gpu.stream();
23447            let mut b2 = __s_b2.launch_builder(&fc);
23448            b2.arg(&*part_o)
23449                .arg(&*part_m)
23450                .arg(&*part_l)
23451                .arg(o)
23452                .arg(&hd)
23453                .arg(&nh)
23454                .arg(base_dev)
23455                .arg(&base_plus)
23456                .arg(&nspm)
23457                .arg(&spk);
23458            unsafe {
23459                b2.launch(cfg2)?;
23460            }
23461            return Ok(());
23462        }
23463        let sp = fa_split_keys(t_kv_upper, n_head_kv);
23464        let n_splits_max = (t_kv_upper + sp - 1) / sp;
23465        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23466        let (nspm, spk) = (n_splits_max as i32, sp as i32);
23467        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23468        let gqa = (n_head / n_head_kv).max(1) as u32;
23469        let o_len = t * n_head * n_splits_max * head_dim;
23470        let ml_len = t * n_head * n_splits_max;
23471        let mut part_guard = self.fa_part_pool.lock().unwrap();
23472        if part_guard
23473            .as_ref()
23474            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23475            .unwrap_or(true)
23476        {
23477            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23478            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23479            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23480            // later live allocations land at those addresses, and the next graph REPLAY writes
23481            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23482            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23483            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23484            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23485            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23486            // (total retired < final size).
23487            let old = part_guard.take();
23488            let (co, cm) = old
23489                .as_ref()
23490                .map(|pp| (pp.0.len(), pp.1.len()))
23491                .unwrap_or((0, 0));
23492            if let Some(old) = old {
23493                self.fa_part_retired.lock().unwrap().push(old);
23494            }
23495            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23496                eprintln!(
23497                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23498                    co, o_len, cm, ml_len
23499                );
23500            }
23501            *part_guard = Some((
23502                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23503                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23504                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23505            ));
23506        }
23507        let pg = part_guard.as_mut().unwrap();
23508        self.gpu
23509            .stream()
23510            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23511        self.gpu
23512            .stream()
23513            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23514        self.gpu
23515            .stream()
23516            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23517        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23518        let f = self.func("fa_decode_vec_q_rows_v3_dc");
23519        let sh = (32 * head_dim * 2) as u32;
23520        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23521        f.set_attribute(
23522            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23523            sh as i32,
23524        )?;
23525        let cfg = LaunchConfig {
23526            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23527            block_dim: (32, gqa, 1),
23528            shared_mem_bytes: sh,
23529        };
23530        let __s_b = self.gpu.stream();
23531        let mut b = __s_b.launch_builder(&f);
23532        b.arg(q)
23533            .arg(k)
23534            .arg(v)
23535            .arg(&mut *part_o)
23536            .arg(&mut *part_m)
23537            .arg(&mut *part_l)
23538            .arg(&hd)
23539            .arg(&nh)
23540            .arg(&nhkv)
23541            .arg(base_dev)
23542            .arg(&scale)
23543            .arg(&nspm)
23544            .arg(&spk)
23545            .arg(&ktb)
23546            .arg(&vtb);
23547        unsafe {
23548            b.launch(cfg)?;
23549        }
23550        let fc = self.func("fa_decode_combine_rows_dc");
23551        let cfg2 = LaunchConfig {
23552            grid_dim: (n_head as u32, t as u32, 1),
23553            block_dim: (head_dim as u32, 1, 1),
23554            shared_mem_bytes: 0,
23555        };
23556        let plus0 = 0i32;
23557        let __s_b2 = self.gpu.stream();
23558        let mut b2 = __s_b2.launch_builder(&fc);
23559        b2.arg(&*part_o)
23560            .arg(&*part_m)
23561            .arg(&*part_l)
23562            .arg(o)
23563            .arg(&hd)
23564            .arg(&nh)
23565            .arg(base_dev)
23566            .arg(&plus0)
23567            .arg(&nspm)
23568            .arg(&spk);
23569        unsafe {
23570            b2.launch(cfg2)?;
23571        }
23572        Ok(())
23573    }
23574
23575    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
23576    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
23577    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
23578    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
23579    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
23580    ///
23581    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
23582    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
23583    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
23584    /// grouping (different but mathematically-equal log-sum-exp merge).
23585    pub fn fa_decode_dc(
23586        &self,
23587        q: &CudaSlice<f32>,
23588        k: &cudarc::driver::CudaView<u8>,
23589        v: &cudarc::driver::CudaView<u8>,
23590        o: &mut CudaSlice<f32>,
23591        head_dim: usize,
23592        n_head: usize,
23593        n_head_kv: usize,
23594        t_kv_dev: &CudaSlice<i32>,
23595        bucket_max: usize,
23596        scale: f32,
23597        k_tok_bytes: usize,
23598        v_tok_bytes: usize,
23599        g: bool,
23600    ) -> Result<(), Box<dyn std::error::Error>> {
23601        self.fa_decode_dc_q8(
23602            q,
23603            k,
23604            v,
23605            o,
23606            head_dim,
23607            n_head,
23608            n_head_kv,
23609            t_kv_dev,
23610            bucket_max,
23611            scale,
23612            k_tok_bytes,
23613            v_tok_bytes,
23614            g,
23615            None,
23616        )
23617    }
23618
23619    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
23620    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
23621    #[allow(clippy::too_many_arguments)]
23622    pub fn fa_decode_dc_q8(
23623        &self,
23624        q: &CudaSlice<f32>,
23625        k: &cudarc::driver::CudaView<u8>,
23626        v: &cudarc::driver::CudaView<u8>,
23627        o: &mut CudaSlice<f32>,
23628        head_dim: usize,
23629        n_head: usize,
23630        n_head_kv: usize,
23631        t_kv_dev: &CudaSlice<i32>,
23632        bucket_max: usize,
23633        scale: f32,
23634        k_tok_bytes: usize,
23635        v_tok_bytes: usize,
23636        g: bool,
23637        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23638    ) -> Result<(), Box<dyn std::error::Error>> {
23639        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
23640        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
23641        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
23642        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
23643        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
23644        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
23645        // 2026-07-12).
23646        let mut fa_vec =
23647            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23648        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23649            fa_vec = false;
23650        } // mirror kvmod/geom
23651        let sp = fa_split_keys(bucket_max, n_head_kv);
23652        let n_splits = if fa_vec {
23653            ((bucket_max + sp - 1) / sp).max(1)
23654        } else {
23655            ((bucket_max + 255) / 256).max(1)
23656        };
23657        let o_len = n_head * n_splits * head_dim;
23658        let ml_len = n_head * n_splits;
23659        let mut part_guard = self.fa_part_pool.lock().unwrap();
23660        if part_guard
23661            .as_ref()
23662            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23663            .unwrap_or(true)
23664        {
23665            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23666            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23667            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23668            // later live allocations land at those addresses, and the next graph REPLAY writes
23669            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23670            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23671            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23672            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23673            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23674            // (total retired < final size).
23675            let old = part_guard.take();
23676            let (co, cm) = old
23677                .as_ref()
23678                .map(|pp| (pp.0.len(), pp.1.len()))
23679                .unwrap_or((0, 0));
23680            if let Some(old) = old {
23681                self.fa_part_retired.lock().unwrap().push(old);
23682            }
23683            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23684                eprintln!(
23685                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23686                    co, o_len, cm, ml_len
23687                );
23688            }
23689            *part_guard = Some((
23690                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23691                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23692                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23693            ));
23694        }
23695        let pg = part_guard.as_mut().unwrap();
23696        self.gpu
23697            .stream()
23698            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23699        self.gpu
23700            .stream()
23701            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23702        self.gpu
23703            .stream()
23704            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23705        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23706        let (hd, nh, nhkv, nsp) = (
23707            head_dim as i32,
23708            n_head as i32,
23709            n_head_kv as i32,
23710            n_splits as i32,
23711        );
23712        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23713        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23714        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
23715        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
23716        let deep = fa_vec
23717            && head_dim == 256
23718            && fa_v4_at(bucket_max)
23719            && !g
23720            && fa_deep_at(bucket_max)
23721            && !matches!(fa_v4_mode(), "noB3" | "stage");
23722        let (f, cfg) = if fa_vec
23723            && head_dim == 512
23724            && bucket_max >= {
23725                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23726                *FA512_MIN_DC.get_or_init(|| {
23727                    std::env::var("MEMRA_FA512_MIN")
23728                        .ok()
23729                        .and_then(|v| v.parse().ok())
23730                        .unwrap_or(512)
23731                })
23732            } {
23733            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
23734            let gqa = (n_head / n_head_kv).max(1) as u32;
23735            (
23736                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23737                LaunchConfig {
23738                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23739                    block_dim: (32, gqa, 1),
23740                    shared_mem_bytes: 0,
23741                },
23742            )
23743        } else if fa_vec && head_dim == 512 {
23744            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
23745            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
23746            let q_view = q.as_view();
23747            let mut o_view = o.as_view_mut();
23748            return self.fa_decode_scalar_unified(
23749                &q_view,
23750                k,
23751                v,
23752                &mut o_view,
23753                head_dim,
23754                n_head,
23755                n_head_kv,
23756                0,
23757                Some(t_kv_dev),
23758                scale,
23759                n_splits,
23760                sp,
23761                k_tok_bytes,
23762                v_tok_bytes,
23763                g,
23764                &mut *part_o,
23765                &mut *part_m,
23766                &mut *part_l,
23767                q8_out,
23768            );
23769        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23770            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
23771            // incl the g-module route + raw-e4m3 sV sizing.
23772            let gqa = (n_head / n_head_kv).max(1) as u32;
23773            let fv = if g {
23774                self.func_g("fa_decode_vec_q_v4_dc")
23775            } else if deep {
23776                self.func("fa_decode_vec_q_v4_deep_dc")
23777            } else {
23778                self.func("fa_decode_vec_q_v4_dc")
23779            };
23780            let shmem =
23781                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23782            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23783            fv.set_attribute(
23784                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23785                shmem as i32,
23786            )?;
23787            (
23788                fv,
23789                LaunchConfig {
23790                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23791                    block_dim: (32, gqa, 1),
23792                    shared_mem_bytes: shmem,
23793                },
23794            )
23795        } else if fa_vec && fa_v3_active(head_dim) {
23796            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
23797            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
23798            let gqa = (n_head / n_head_kv).max(1) as u32;
23799            let fv = if g {
23800                self.func_g("fa_decode_vec_q_v3_dc")
23801            } else {
23802                self.func("fa_decode_vec_q_v3_dc")
23803            };
23804            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
23805            (
23806                fv,
23807                LaunchConfig {
23808                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23809                    block_dim: (32, gqa, 1),
23810                    shared_mem_bytes: shmem,
23811                },
23812            )
23813        } else if fa_vec && fa_v2_on() {
23814            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
23815            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
23816            // a numeric config; eager, rows-verify and graph all switch together).
23817            let gqa = (n_head / n_head_kv).max(1) as u32;
23818            let fv = if g {
23819                self.func_g("fa_decode_vec_q_v2_dc")
23820            } else {
23821                self.func("fa_decode_vec_q_v2_dc")
23822            };
23823            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
23824            (
23825                fv,
23826                LaunchConfig {
23827                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23828                    block_dim: (32, gqa, 1),
23829                    shared_mem_bytes: shmem,
23830                },
23831            )
23832        } else if fa_vec {
23833            let gqa = (n_head / n_head_kv).max(1) as u32;
23834            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
23835            let fv = if g {
23836                self.func_g("fa_decode_vec_q_dc")
23837            } else {
23838                self.func("fa_decode_vec_q_dc")
23839            };
23840            (
23841                fv,
23842                LaunchConfig {
23843                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23844                    block_dim: (32, gqa, 1),
23845                    shared_mem_bytes: 0,
23846                },
23847            )
23848        } else {
23849            let q_view = q.as_view();
23850            let mut o_view = o.as_view_mut();
23851            return self.fa_decode_scalar_unified(
23852                &q_view,
23853                k,
23854                v,
23855                &mut o_view,
23856                head_dim,
23857                n_head,
23858                n_head_kv,
23859                0,
23860                Some(t_kv_dev),
23861                scale,
23862                n_splits,
23863                if fa_vec { sp } else { 256 },
23864                k_tok_bytes,
23865                v_tok_bytes,
23866                g,
23867                &mut *part_o,
23868                &mut *part_m,
23869                &mut *part_l,
23870                q8_out,
23871            );
23872        };
23873        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
23874        let __s_b = self.gpu.stream();
23875        let mut b = __s_b.launch_builder(&f);
23876        b.arg(q)
23877            .arg(k)
23878            .arg(v)
23879            .arg(&mut *part_o)
23880            .arg(&mut *part_m)
23881            .arg(&mut *part_l)
23882            .arg(&hd)
23883            .arg(&nh)
23884            .arg(&nhkv)
23885            .arg(t_kv_dev)
23886            .arg(&scale)
23887            .arg(&nsp)
23888            .arg(&ski)
23889            .arg(&ktb)
23890            .arg(&vtb);
23891        unsafe {
23892            b.launch(cfg)?;
23893        }
23894        let cfg2 = LaunchConfig {
23895            grid_dim: (n_head as u32, 1, 1),
23896            block_dim: (head_dim as u32, 1, 1),
23897            shared_mem_bytes: 0,
23898        };
23899        if let Some((oq, od)) = q8_out {
23900            let fc = if g {
23901                self.func_g("fa_decode_combine_q8_1")
23902            } else {
23903                self.fa_func("fa_decode_combine_q8_1", head_dim)
23904            };
23905            let __s_b2 = self.gpu.stream();
23906            let mut b2 = __s_b2.launch_builder(&fc);
23907            b2.arg(&*part_o)
23908                .arg(&*part_m)
23909                .arg(&*part_l)
23910                .arg(oq)
23911                .arg(od)
23912                .arg(&hd)
23913                .arg(&nh)
23914                .arg(&nsp);
23915            unsafe {
23916                b2.launch(cfg2)?;
23917            }
23918            return Ok(());
23919        }
23920        let fc = if g {
23921            self.func_g("fa_decode_combine_f32")
23922        } else {
23923            self.fa_func("fa_decode_combine_f32", head_dim)
23924        };
23925        let __s_b2 = self.gpu.stream();
23926        let mut b2 = __s_b2.launch_builder(&fc);
23927        b2.arg(&*part_o)
23928            .arg(&*part_m)
23929            .arg(&*part_l)
23930            .arg(o)
23931            .arg(&hd)
23932            .arg(&nh)
23933            .arg(&nsp);
23934        unsafe {
23935            b2.launch(cfg2)?;
23936        }
23937        Ok(())
23938    }
23939
23940    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
23941    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
23942    /// at equal rows.
23943    #[allow(clippy::too_many_arguments)]
23944    pub fn append_kv_quantized_dcw(
23945        &self,
23946        k_row: &CudaSlice<f32>,
23947        v_row: &CudaSlice<f32>,
23948        kc: &mut CudaSlice<u8>,
23949        vc: &mut CudaSlice<u8>,
23950        len_dev: &CudaSlice<i32>,
23951        base_dev: Option<&CudaSlice<i32>>,
23952        kv_dim_k: usize,
23953        kv_dim_v: usize,
23954        k_tok_bytes: usize,
23955        v_tok_bytes: usize,
23956    ) -> Result<(), Box<dyn std::error::Error>> {
23957        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23958        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23959        let cfg = LaunchConfig {
23960            grid_dim: (nblk, 1, 1),
23961            block_dim: (32, 1, 1),
23962            shared_mem_bytes: 0,
23963        };
23964        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23965        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23966        let null: u64 = 0;
23967        let __s_b = self.gpu.stream();
23968        let mut b = __s_b.launch_builder(&f);
23969        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23970        match base_dev {
23971            Some(base) => {
23972                b.arg(base);
23973            }
23974            None => {
23975                b.arg(&null);
23976            }
23977        }
23978        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23979        unsafe {
23980            b.launch(cfg)?;
23981        }
23982        Ok(())
23983    }
23984
23985    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
23986    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23987        let f = self.func("inc_i32");
23988        let cfg = LaunchConfig {
23989            grid_dim: (1, 1, 1),
23990            block_dim: (1, 1, 1),
23991            shared_mem_bytes: 0,
23992        };
23993        let __s_b = self.gpu.stream();
23994        let mut b = __s_b.launch_builder(&f);
23995        b.arg(counter);
23996        unsafe {
23997            b.launch(cfg)?;
23998        }
23999        Ok(())
24000    }
24001
24002    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
24003    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
24004    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
24005    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
24006    /// kernel class on this lane); callers keep eager below the vec floor and for any other
24007    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
24008    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
24009    /// alive across bucket growth.
24010    #[allow(clippy::too_many_arguments)]
24011    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
24012    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
24013    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
24014    fn fa_part_pool_grow(
24015        &self,
24016        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
24017        o_len: usize,
24018        ml_len: usize,
24019    ) -> Result<(), Box<dyn std::error::Error>> {
24020        if part_guard
24021            .as_ref()
24022            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24023            .unwrap_or(true)
24024        {
24025            let old = part_guard.take();
24026            let (co, cm) = old
24027                .as_ref()
24028                .map(|pp| (pp.0.len(), pp.1.len()))
24029                .unwrap_or((0, 0));
24030            if let Some(old) = old {
24031                self.fa_part_retired.lock().unwrap().push(old);
24032            }
24033            *part_guard = Some((
24034                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
24035                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
24036                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
24037            ));
24038        }
24039        Ok(())
24040    }
24041
24042    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
24043    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
24044    pub fn fa_dcw_pool_ensure(
24045        &self,
24046        head_dim: usize,
24047        n_head: usize,
24048        n_head_kv: usize,
24049        bucket_max: usize,
24050    ) -> Result<(), Box<dyn std::error::Error>> {
24051        let sp = fa_split_keys(bucket_max, n_head_kv);
24052        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24053        let o_len = n_head * n_splits * head_dim;
24054        let ml_len = n_head * n_splits;
24055        let mut part_guard = self.fa_part_pool.lock().unwrap();
24056        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24057    }
24058
24059    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
24060    /// appended; one launch walks the KV stream once with two query rows (per-row causal
24061    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
24062    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
24063    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
24064    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
24065    /// outputs (the head gate fuses into the combine as in the t=1 path).
24066    #[allow(clippy::too_many_arguments)]
24067    pub fn fa_decode_dcw2(
24068        &self,
24069        q2: &CudaSlice<f32>,
24070        k_ring: &cudarc::driver::CudaView<u8>,
24071        v_ring: &cudarc::driver::CudaView<u8>,
24072        o2: &mut CudaSlice<f32>,
24073        head_dim: usize,
24074        n_head: usize,
24075        n_head_kv: usize,
24076        len_dev: &CudaSlice<i32>,
24077        base_dev: Option<&CudaSlice<i32>>,
24078        window: usize,
24079        bucket_max: usize,
24080        scale: f32,
24081        k_tok_bytes: usize,
24082        v_tok_bytes: usize,
24083        gate2: &CudaSlice<f32>,
24084    ) -> Result<(), Box<dyn std::error::Error>> {
24085        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24086        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24087            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24088        }
24089        let sp = fa_split_keys(bucket_max, n_head_kv);
24090        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24091        // Partials for BOTH rows: row-major halves.
24092        let o_len = 2 * n_head * n_splits * head_dim;
24093        let ml_len = 2 * n_head * n_splits;
24094        let mut part_guard = self.fa_part_pool.lock().unwrap();
24095        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24096        let pg = part_guard.as_mut().unwrap();
24097        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24098        let (hd, nh, nhkv, nsp) = (
24099            head_dim as i32,
24100            n_head as i32,
24101            n_head_kv as i32,
24102            n_splits as i32,
24103        );
24104        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24105        let (ski, win) = (sp as i32, window as i32);
24106        let gqa = (n_head / n_head_kv).max(1) as u32;
24107        let smem = (32 * head_dim * 2) as u32;
24108        let f = self.func("fa_decode_vec_q_v3_dcw2");
24109        let cfg = LaunchConfig {
24110            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24111            block_dim: (32, gqa, 1),
24112            shared_mem_bytes: smem,
24113        };
24114        let null: u64 = 0;
24115        {
24116            let __s_b = self.gpu.stream();
24117            let mut b = __s_b.launch_builder(&f);
24118            b.arg(q2)
24119                .arg(k_ring)
24120                .arg(v_ring)
24121                .arg(&mut *part_o)
24122                .arg(&mut *part_m)
24123                .arg(&mut *part_l)
24124                .arg(&hd)
24125                .arg(&nh)
24126                .arg(&nhkv)
24127                .arg(len_dev);
24128            match base_dev {
24129                Some(base) => {
24130                    b.arg(base);
24131                }
24132                None => {
24133                    b.arg(&null);
24134                }
24135            }
24136            b.arg(&win)
24137                .arg(&scale)
24138                .arg(&nsp)
24139                .arg(&ski)
24140                .arg(&ktb)
24141                .arg(&vtb);
24142            unsafe {
24143                b.launch(cfg)?;
24144            }
24145        }
24146        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
24147        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
24148        // one launch covers both rows with the exact t=1 program per (row, head).
24149        let fc = {
24150            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24151            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24152                self.func("fa_decode_combine_gate_f32_s")
24153            } else {
24154                self.func("fa_decode_combine_gate_f32")
24155            }
24156        };
24157        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24158        let nh2 = (2 * n_head) as i32;
24159        let cfg2 = LaunchConfig {
24160            grid_dim: ((2 * n_head) as u32, 1, 1),
24161            block_dim: (head_dim as u32, 1, 1),
24162            shared_mem_bytes: if combine_shared {
24163                (2 * n_splits * 4) as u32
24164            } else {
24165                0
24166            },
24167        };
24168        let __s_b2 = self.gpu.stream();
24169        let mut b2 = __s_b2.launch_builder(&fc);
24170        b2.arg(&*part_o)
24171            .arg(&*part_m)
24172            .arg(&*part_l)
24173            .arg(gate2)
24174            .arg(o2)
24175            .arg(&hd)
24176            .arg(&nh2)
24177            .arg(&nsp);
24178        unsafe {
24179            b2.launch(cfg2)?;
24180        }
24181        Ok(())
24182    }
24183
24184    /// T-ROW dcw decode attention over a per-row session table (the per-session
24185    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
24186    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
24187    /// program verbatim with that row's ring/len/base and its own split geometry, so each
24188    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
24189    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
24190    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
24191    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
24192    #[allow(clippy::too_many_arguments)]
24193    pub fn fa_decode_dcw_rows(
24194        &self,
24195        q_rows: &CudaSlice<f32>,
24196        tab: &CudaSlice<u64>,
24197        o_rows: &mut CudaSlice<f32>,
24198        t: usize,
24199        head_dim: usize,
24200        n_head: usize,
24201        n_head_kv: usize,
24202        window: usize,
24203        max_ns: usize,
24204        scale: f32,
24205        k_tok_bytes: usize,
24206        v_tok_bytes: usize,
24207        gate_rows: &CudaSlice<f32>,
24208    ) -> Result<(), Box<dyn std::error::Error>> {
24209        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24210            || head_dim > 256
24211            || head_dim % 32 != 0
24212            || !fa_v3_on()
24213        {
24214            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24215        }
24216        if fa_sm_count() < 128
24217            || std::env::var("MEMRA_FA_SPLIT").is_ok()
24218            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24219            || std::env::var("MEMRA_FA_SP16").is_ok()
24220        {
24221            return Err(
24222                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24223                 (or a <128-SM rig) keep the per-row path"
24224                    .into(),
24225            );
24226        }
24227        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24228            return Err("fa_decode_dcw_rows geometry".into());
24229        }
24230        let o_len = t * n_head * max_ns * head_dim;
24231        let ml_len = t * n_head * max_ns;
24232        let mut part_guard = self.fa_part_pool.lock().unwrap();
24233        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24234        let pg = part_guard.as_mut().unwrap();
24235        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24236        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24237        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24238        let (win, mns) = (window as i32, max_ns as i32);
24239        let gqa = (n_head / n_head_kv).max(1) as u32;
24240        let smem = (32 * head_dim * 2) as u32;
24241        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24242        let cfg = LaunchConfig {
24243            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24244            block_dim: (32, gqa, 1),
24245            shared_mem_bytes: smem,
24246        };
24247        {
24248            let __s_b = self.gpu.stream();
24249            let mut b = __s_b.launch_builder(&f);
24250            b.arg(q_rows)
24251                .arg(tab)
24252                .arg(&mut *part_o)
24253                .arg(&mut *part_m)
24254                .arg(&mut *part_l)
24255                .arg(&hd)
24256                .arg(&nh)
24257                .arg(&nhkv)
24258                .arg(&win)
24259                .arg(&scale)
24260                .arg(&mns)
24261                .arg(&ktb)
24262                .arg(&vtb);
24263            unsafe {
24264                b.launch(cfg)?;
24265            }
24266        }
24267        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
24268        // row r head h reads its own partial bank; splits past a row's ns_eff carry
24269        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
24270        let fc = {
24271            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24272            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24273                self.func("fa_decode_combine_gate_f32_s")
24274            } else {
24275                self.func("fa_decode_combine_gate_f32")
24276            }
24277        };
24278        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24279        let nht = (t * n_head) as i32;
24280        let cfg2 = LaunchConfig {
24281            grid_dim: ((t * n_head) as u32, 1, 1),
24282            block_dim: (head_dim as u32, 1, 1),
24283            shared_mem_bytes: if combine_shared {
24284                (2 * max_ns * 4) as u32
24285            } else {
24286                0
24287            },
24288        };
24289        let __s_b2 = self.gpu.stream();
24290        let mut b2 = __s_b2.launch_builder(&fc);
24291        b2.arg(&*part_o)
24292            .arg(&*part_m)
24293            .arg(&*part_l)
24294            .arg(gate_rows)
24295            .arg(o_rows)
24296            .arg(&hd)
24297            .arg(&nht)
24298            .arg(&mns);
24299        unsafe {
24300            b2.launch(cfg2)?;
24301        }
24302        Ok(())
24303    }
24304
24305    pub fn fa_decode_dcw(
24306        &self,
24307        q: &CudaSlice<f32>,
24308        k_ring: &cudarc::driver::CudaView<u8>,
24309        v_ring: &cudarc::driver::CudaView<u8>,
24310        o: &mut CudaSlice<f32>,
24311        head_dim: usize,
24312        n_head: usize,
24313        n_head_kv: usize,
24314        len_dev: &CudaSlice<i32>,
24315        base_dev: Option<&CudaSlice<i32>>,
24316        window: usize,
24317        bucket_max: usize,
24318        scale: f32,
24319        k_tok_bytes: usize,
24320        v_tok_bytes: usize,
24321        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
24322        // one launch saved); `o` then receives the GATED output and the caller skips its
24323        // attn_head_gate call.
24324        fused_gate: Option<&CudaSlice<f32>>,
24325    ) -> Result<(), Box<dyn std::error::Error>> {
24326        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24327        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24328            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"
24329                .into());
24330        }
24331        let sp = fa_split_keys(bucket_max, n_head_kv);
24332        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24333        let o_len = n_head * n_splits * head_dim;
24334        let ml_len = n_head * n_splits;
24335        let mut part_guard = self.fa_part_pool.lock().unwrap();
24336        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24337        let pg = part_guard.as_mut().unwrap();
24338        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
24339        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
24340        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
24341        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
24342        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24343        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
24344        // finds the attention children BY their three-memset signature and updates the
24345        // memset widths per bucket — capturing without them silently kills retargeting
24346        // (battery-v8 token drift, 2026-08-21).
24347        let memset_on = *MEMSET_ON
24348            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24349            || crate::tp::token_graph_building();
24350        if memset_on {
24351            self.gpu
24352                .stream()
24353                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24354            self.gpu
24355                .stream()
24356                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24357            self.gpu
24358                .stream()
24359                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24360        }
24361        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24362        let (hd, nh, nhkv, nsp) = (
24363            head_dim as i32,
24364            n_head as i32,
24365            n_head_kv as i32,
24366            n_splits as i32,
24367        );
24368        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24369        let (ski, win) = (sp as i32, window as i32);
24370        let gqa = (n_head / n_head_kv).max(1) as u32;
24371        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
24372        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
24373        // see fa_dec_v3_walk_u). Same launch geometry.
24374        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24375        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24376        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24377            Ok("2") => 2,
24378            Ok("1") => 1,
24379            _ => 0,
24380        });
24381        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
24382        // permission-blocked in this container and the module params are not exposed, so this
24383        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
24384        // prints cumulative cycle shares every 430 launches.
24385        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24386        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24387        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24388            std::sync::Mutex::new(None);
24389        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
24390        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
24391        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
24392        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24393        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24394            && (n_head / n_head_kv) % 2 == 0
24395            && (n_head / n_head_kv) >= 2;
24396        let f = if fprof {
24397            self.func("fa_decode_vec_q_v3_dcw_prof")
24398        } else if hs2 {
24399            self.func("fa_decode_vec_q_v3_dcw_hs2")
24400        } else if hoist == 2 {
24401            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
24402            self.func("fa_decode_vec_q_v3_dcw_hc")
24403        } else if hoist == 1 {
24404            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
24405            self.func("fa_decode_vec_q_v3_dcw_h")
24406        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24407            self.func("fa_decode_vec_q_v3_dcw_u8")
24408        } else {
24409            self.func("fa_decode_vec_q_v3_dcw")
24410        };
24411        let cfg = LaunchConfig {
24412            grid_dim: if hs2 {
24413                ((2 * n_head_kv) as u32, n_splits as u32, 1)
24414            } else {
24415                (n_head_kv as u32, n_splits as u32, 1)
24416            },
24417            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24418            shared_mem_bytes: smem,
24419        };
24420        let null: u64 = 0;
24421        let __s_b = self.gpu.stream();
24422        let mut b = __s_b.launch_builder(&f);
24423        b.arg(q)
24424            .arg(k_ring)
24425            .arg(v_ring)
24426            .arg(&mut *part_o)
24427            .arg(&mut *part_m)
24428            .arg(&mut *part_l)
24429            .arg(&hd)
24430            .arg(&nh)
24431            .arg(&nhkv)
24432            .arg(len_dev);
24433        match base_dev {
24434            Some(base) => {
24435                b.arg(base);
24436            }
24437            None => {
24438                b.arg(&null);
24439            }
24440        }
24441        b.arg(&win)
24442            .arg(&scale)
24443            .arg(&nsp)
24444            .arg(&ski)
24445            .arg(&ktb)
24446            .arg(&vtb);
24447        if fprof {
24448            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24449            if guard
24450                .as_ref()
24451                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24452            {
24453                *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24454            }
24455            let (_, buf) = guard.as_mut().expect("armed above");
24456            b.arg(&*buf);
24457            unsafe {
24458                b.launch(cfg)?;
24459            }
24460            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24461            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24462            if n % 430 == 0 {
24463                self.stream().synchronize()?;
24464                let h = self.dtoh_u64(buf)?;
24465                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24466                let tot: u64 = h[..6].iter().sum();
24467                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24468                for (i, name) in phases.iter().enumerate() {
24469                    let pct = if tot > 0 {
24470                        h[i] as f64 / tot as f64 * 100.0
24471                    } else {
24472                        0.0
24473                    };
24474                    line.push_str(&format!(" {name}={pct:.1}%"));
24475                }
24476                if h[6] > 0 {
24477                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24478                }
24479                eprintln!("{line}");
24480            }
24481        } else {
24482            unsafe {
24483                b.launch(cfg)?;
24484            }
24485        }
24486        let mut combine_shared = false;
24487        let fc = if fused_gate.is_some() {
24488            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
24489            // n_splits-deep dependent global load chain every thread used to walk twice).
24490            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24491            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24492                combine_shared = true;
24493                self.func("fa_decode_combine_gate_f32_s")
24494            } else {
24495                self.func("fa_decode_combine_gate_f32")
24496            }
24497        } else {
24498            self.fa_func("fa_decode_combine_f32", head_dim)
24499        };
24500        let cfg2 = LaunchConfig {
24501            grid_dim: (n_head as u32, 1, 1),
24502            block_dim: (head_dim as u32, 1, 1),
24503            shared_mem_bytes: if combine_shared {
24504                (2 * n_splits * 4) as u32
24505            } else {
24506                0
24507            },
24508        };
24509        let __s_b2 = self.gpu.stream();
24510        let mut b2 = __s_b2.launch_builder(&fc);
24511        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24512        if let Some(gate_row) = fused_gate {
24513            b2.arg(gate_row);
24514        }
24515        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24516        unsafe {
24517            b2.launch(cfg2)?;
24518        }
24519        Ok(())
24520    }
24521
24522    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
24523    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
24524    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
24525    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
24526    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
24527    pub fn fa_geom_eager(
24528        &self,
24529        t_kv: usize,
24530        head_dim: usize,
24531        n_head_kv: usize,
24532        g: bool,
24533    ) -> (bool, usize) {
24534        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
24535        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
24536        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
24537        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24538        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
24539        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
24540        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
24541        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
24542        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
24543        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24544        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24545        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
24546        // family; everything else falls to the g-module scalar.
24547        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
24548        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
24549        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
24550        if g && head_dim == 256 && !fa_v4_at(t_kv) {
24551            fa_vec = false;
24552        }
24553        let sp = fa_split_keys(t_kv, n_head_kv);
24554        let n_splits = if fa_vec {
24555            ((t_kv + sp - 1) / sp).max(1)
24556        } else {
24557            ((t_kv + 255) / 256).max(1)
24558        };
24559        (fa_vec, n_splits)
24560    }
24561
24562    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
24563    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
24564    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
24565    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
24566    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
24567    pub fn fa_bucket_key(
24568        &self,
24569        t_kv: usize,
24570        head_dim: usize,
24571        n_head_kv: usize,
24572        g: bool,
24573    ) -> (bool, usize) {
24574        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24575    }
24576
24577    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
24578    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
24579    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
24580    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
24581    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
24582    /// device data) — every per-step varying scalar must come from a device counter. Returns the
24583    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
24584    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
24585    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
24586    /// replays (transients returning to the pool get reused by unrelated work and corrupt
24587    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
24588    pub fn capture_graph_retained<F>(
24589        &self,
24590        step: F,
24591    ) -> Result<
24592        (
24593            cudarc::driver::CudaGraph,
24594            Vec<Box<dyn std::any::Any + Send>>,
24595        ),
24596        Box<dyn std::error::Error>,
24597    >
24598    where
24599        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24600    {
24601        use cudarc::driver::sys::CUgraphInstantiate_flags;
24602        self.capture_graph_retained_flags(
24603            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24604            step,
24605        )
24606    }
24607
24608    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
24609    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
24610    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
24611    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
24612    pub fn capture_graph_retained_flags<F>(
24613        &self,
24614        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24615        mut step: F,
24616    ) -> Result<
24617        (
24618            cudarc::driver::CudaGraph,
24619            Vec<Box<dyn std::any::Any + Send>>,
24620        ),
24621        Box<dyn std::error::Error>,
24622    >
24623    where
24624        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24625    {
24626        use cudarc::driver::sys::CUstreamCaptureMode;
24627        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
24628        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
24629        // while the capture region is open become dead copy NODES replayed every launch
24630        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
24631        // warmup runs allocate the same transient sequence at the same pool addresses, so
24632        // retaining the warmup clones preserves the draft-graph fix without polluting the
24633        // captured graph.
24634        self.capture_keep.lock().unwrap().clear();
24635        let was_tracking = self.gpu.ctx.is_event_tracking();
24636        if was_tracking {
24637            unsafe {
24638                self.gpu.ctx.disable_event_tracking();
24639            }
24640        }
24641        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24642            self.capture_keep_on
24643                .store(true, std::sync::atomic::Ordering::Relaxed);
24644            let w = (|| {
24645                step(self)?;
24646                step(self)
24647            })();
24648            self.capture_keep_on
24649                .store(false, std::sync::atomic::Ordering::Relaxed);
24650            w?;
24651            self.gpu.stream().synchronize()?;
24652            self.gpu
24653                .stream()
24654                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24655            let r = step(self);
24656            let g = self.gpu.stream().end_capture(flags);
24657            r?;
24658            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24659            graph.upload()?;
24660            Ok(graph)
24661        };
24662        let result = run();
24663        self.capture_keep_on
24664            .store(false, std::sync::atomic::Ordering::Relaxed);
24665        if was_tracking {
24666            unsafe {
24667                self.gpu.ctx.enable_event_tracking();
24668            }
24669        }
24670        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24671        Ok((result?, keeper))
24672    }
24673
24674    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
24675    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
24676    /// alloc-free with persistent operands, and their bodies carry device side effects
24677    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
24678    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
24679    pub fn capture_graph_retained_nowarm<F>(
24680        &self,
24681        mut step: F,
24682    ) -> Result<
24683        (
24684            cudarc::driver::CudaGraph,
24685            Vec<Box<dyn std::any::Any + Send>>,
24686        ),
24687        Box<dyn std::error::Error>,
24688    >
24689    where
24690        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24691    {
24692        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24693        let was_tracking = self.gpu.ctx.is_event_tracking();
24694        if was_tracking {
24695            unsafe {
24696                self.gpu.ctx.disable_event_tracking();
24697            }
24698        }
24699        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24700            self.gpu.stream().synchronize()?;
24701            self.gpu
24702                .stream()
24703                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24704            let r = step(self);
24705            let g = self.gpu.stream().end_capture(
24706                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24707            );
24708            r?;
24709            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24710            graph.upload()?;
24711            Ok(graph)
24712        };
24713        let result = run();
24714        if was_tracking {
24715            unsafe {
24716                self.gpu.ctx.enable_event_tracking();
24717            }
24718        }
24719        Ok((result?, Vec::new()))
24720    }
24721
24722    pub fn capture_graph<F>(
24723        &self,
24724        mut step: F,
24725    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24726    where
24727        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24728    {
24729        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24730        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
24731        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
24732        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
24733        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
24734        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
24735        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
24736        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
24737        let was_tracking = self.gpu.ctx.is_event_tracking();
24738        if was_tracking {
24739            unsafe {
24740                self.gpu.ctx.disable_event_tracking();
24741            }
24742        }
24743        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
24744        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
24745        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
24746        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
24747        // measure that scan's real cost on the generic path. Diagnostic door only; the
24748        // default stays AUTO_FREE until a measured A/B justifies moving it.
24749        let iflag = {
24750            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24751            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24752                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
24753                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
24754                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24755                Ok("priority") => {
24756                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24757                }
24758                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24759            })
24760        };
24761        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
24762        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
24763        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
24764        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
24765        // eager step executions and are node-count-invariant. Printing the split bounds the
24766        // refactor's ceiling instead of assuming it.
24767        let ct = {
24768            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24769            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24770        };
24771        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
24772        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
24773        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
24774        // chased, and node-count-invariant, so no capture-body refactor could touch it.
24775        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
24776        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
24777        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
24778        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
24779        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
24780        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
24781        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
24782        // grow and never frees, resident counters/scratch, cache set in place), and the
24783        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
24784        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
24785        // settling and pool mapping. Arbitrated adversarially, not by taste:
24786        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
24787        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
24788        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
24789        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
24790        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
24791        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
24792        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
24793        let warmups = {
24794            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24795            *W.get_or_init(|| {
24796                std::env::var("MEMRA_GRAPH_WARMUPS")
24797                    .ok()
24798                    .and_then(|v| v.parse().ok())
24799                    .filter(|n| *n >= 1)
24800                    .unwrap_or(1)
24801            })
24802        };
24803        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24804            let t_w = std::time::Instant::now();
24805            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
24806            for _ in 0..warmups {
24807                step(self)?;
24808            }
24809            self.gpu.stream().synchronize()?;
24810            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24811            // capture the third run.
24812            let t_c = std::time::Instant::now();
24813            self.gpu
24814                .stream()
24815                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24816            // If the body errors mid-capture, end the capture before propagating so the stream isn't
24817            // left in a capturing state.
24818            let r = step(self);
24819            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24820            let t_i = std::time::Instant::now();
24821            let g = self.gpu.stream().end_capture(iflag);
24822            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24823            r?;
24824            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24825            let t_u = std::time::Instant::now();
24826            graph.upload()?;
24827            if ct {
24828                println!(
24829                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
24830                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
24831                    t_u.elapsed().as_secs_f64() * 1e3
24832                );
24833            }
24834            Ok(graph)
24835        };
24836        let result = run();
24837        if was_tracking {
24838            unsafe {
24839                self.gpu.ctx.enable_event_tracking();
24840            }
24841        }
24842        result
24843    }
24844
24845    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
24846    pub fn gdn_scan_s128_view(
24847        &self,
24848        q: &CudaSlice<f32>,
24849        k: &CudaSlice<f32>,
24850        v: &CudaSlice<f32>,
24851        g: &CudaSlice<f32>,
24852        beta: &CudaSlice<f32>,
24853        state_in: &cudarc::driver::CudaView<f32>,
24854        state_out: &mut cudarc::driver::CudaViewMut<f32>,
24855        o: &mut CudaSlice<f32>,
24856        n_head: usize,
24857        t: usize,
24858        scale: f32,
24859    ) -> Result<(), Box<dyn std::error::Error>> {
24860        let f = self.func("gdn_scan_s128");
24861        const S_V: u32 = 128;
24862        const WARP: u32 = 32;
24863        const COLS: u32 = 4;
24864        let cfg = LaunchConfig {
24865            grid_dim: (n_head as u32, 1, S_V / COLS),
24866            block_dim: (WARP, COLS, 1),
24867            shared_mem_bytes: 0,
24868        };
24869        let (h, ti) = (n_head as i32, t as i32);
24870        let __s_b = self.gpu.stream();
24871        let mut b = __s_b.launch_builder(&f);
24872        b.arg(q)
24873            .arg(k)
24874            .arg(v)
24875            .arg(g)
24876            .arg(beta)
24877            .arg(state_in)
24878            .arg(state_out)
24879            .arg(o)
24880            .arg(&h)
24881            .arg(&ti)
24882            .arg(&scale);
24883        unsafe {
24884            b.launch(cfg)?;
24885        }
24886        Ok(())
24887    }
24888
24889    /// conv1d where the input is a CudaView (resident conv state assembled in place).
24890    pub fn ssm_conv1d_view(
24891        &self,
24892        x: &cudarc::driver::CudaView<f32>,
24893        w: &CudaSlice<f32>,
24894        y: &mut CudaSlice<f32>,
24895        conv_dim: usize,
24896        t: usize,
24897        d_conv: usize,
24898        silu: bool,
24899    ) -> Result<(), Box<dyn std::error::Error>> {
24900        let f = self.func("ssm_conv1d_silu_f32");
24901        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
24902        let cfg = LaunchConfig {
24903            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24904            block_dim: (256, 1, 1),
24905            shared_mem_bytes: 0,
24906        };
24907        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24908        let __s_b = self.gpu.stream();
24909        let mut b = __s_b.launch_builder(&f);
24910        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24911        unsafe {
24912            b.launch(cfg)?;
24913        }
24914        Ok(())
24915    }
24916
24917    /// Depthwise causal conv1d + optional SiLU.
24918    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
24919    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
24920    /// FUSED prefill conv (token-major input, zero left-state): replaces
24921    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
24922    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
24923    pub fn ssm_conv1d_tm(
24924        &self,
24925        qkv_tm: &CudaSlice<f32>,
24926        w: &CudaSlice<f32>,
24927        y: &mut CudaSlice<f32>,
24928        conv_dim: usize,
24929        t: usize,
24930        d_conv: usize,
24931    ) -> Result<(), Box<dyn std::error::Error>> {
24932        let f = self.func("ssm_conv1d_tm_f32");
24933        let cfg = LaunchConfig {
24934            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24935            block_dim: (256, 1, 1),
24936            shared_mem_bytes: 0,
24937        };
24938        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24939        let __s_b = self.gpu.stream();
24940        let mut b = __s_b.launch_builder(&f);
24941        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24942        unsafe {
24943            b.launch(cfg)?;
24944        }
24945        Ok(())
24946    }
24947
24948    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
24949    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
24950    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
24951    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
24952    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
24953    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
24954    /// columns; the final ring == what T sequential decode ring rolls leave).
24955    pub fn ssm_conv1d_tm_state(
24956        &self,
24957        qkv_tm: &CudaSlice<f32>,
24958        conv_state: &mut CudaSlice<f32>,
24959        w: &CudaSlice<f32>,
24960        y: &mut CudaSlice<f32>,
24961        conv_dim: usize,
24962        t: usize,
24963        d_conv: usize,
24964    ) -> Result<(), Box<dyn std::error::Error>> {
24965        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24966    }
24967
24968    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
24969    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
24970    #[allow(clippy::too_many_arguments)]
24971    pub fn ssm_conv1d_tm_state_pad(
24972        &self,
24973        qkv_tm: &CudaSlice<f32>,
24974        conv_state: &mut CudaSlice<f32>,
24975        w: &CudaSlice<f32>,
24976        y: &mut CudaSlice<f32>,
24977        conv_dim: usize,
24978        t: usize,
24979        d_conv: usize,
24980        pad_len: Option<&CudaSlice<i32>>,
24981    ) -> Result<(), Box<dyn std::error::Error>> {
24982        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24983        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
24984        // the window kernel both read the pre-roll ring; the roll launches after both) — but
24985        // cloning first keeps the ordering trivially correct under any future stream split.
24986        let ring_old = if t < d_conv - 1 {
24987            Some(self.clone_dtod(conv_state)?)
24988        } else {
24989            None
24990        };
24991        {
24992            let f = self.func("ssm_conv1d_tm_state_f32");
24993            let cfg = LaunchConfig {
24994                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24995                block_dim: (256, 1, 1),
24996                shared_mem_bytes: 0,
24997            };
24998            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24999            let __s_b = self.gpu.stream();
25000            let mut b = __s_b.launch_builder(&f);
25001            b.arg(qkv_tm)
25002                .arg(&*conv_state)
25003                .arg(w)
25004                .arg(y)
25005                .arg(&cd)
25006                .arg(&ti)
25007                .arg(&dc);
25008            unsafe {
25009                b.launch(cfg)?;
25010            }
25011        }
25012        match (ring_old, pad_len) {
25013            (None, Some(len_d)) => {
25014                let f = self.func("ssm_conv_ring_update_dev_f32");
25015                let n = conv_dim * (d_conv - 1);
25016                let cfg = LaunchConfig::for_num_elems(n as u32);
25017                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25018                let __s_b = self.gpu.stream();
25019                let mut b = __s_b.launch_builder(&f);
25020                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25021                unsafe {
25022                    b.launch(cfg)?;
25023                }
25024            }
25025            (None, None) => {
25026                let f = self.func("ssm_conv_ring_update_f32");
25027                let n = conv_dim * (d_conv - 1);
25028                let cfg = LaunchConfig::for_num_elems(n as u32);
25029                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25030                let __s_b = self.gpu.stream();
25031                let mut b = __s_b.launch_builder(&f);
25032                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25033                unsafe {
25034                    b.launch(cfg)?;
25035                }
25036            }
25037            (Some(old), _) => {
25038                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
25039            }
25040        }
25041        Ok(())
25042    }
25043
25044    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
25045    pub fn ssm_conv1d_tm_state_pad_v(
25046        &self,
25047        qkv_tm: &cudarc::driver::CudaView<f32>,
25048        conv_state: &mut CudaSlice<f32>,
25049        w: &CudaSlice<f32>,
25050        y: &mut CudaSlice<f32>,
25051        conv_dim: usize,
25052        t: usize,
25053        d_conv: usize,
25054        pad_len: Option<&CudaSlice<i32>>,
25055    ) -> Result<(), Box<dyn std::error::Error>> {
25056        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25057        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
25058        // the window kernel both read the pre-roll ring; the roll launches after both) — but
25059        // cloning first keeps the ordering trivially correct under any future stream split.
25060        let ring_old = if t < d_conv - 1 {
25061            Some(self.clone_dtod(conv_state)?)
25062        } else {
25063            None
25064        };
25065        {
25066            let f = self.func("ssm_conv1d_tm_state_f32");
25067            let cfg = LaunchConfig {
25068                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25069                block_dim: (256, 1, 1),
25070                shared_mem_bytes: 0,
25071            };
25072            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25073            let __s_b = self.gpu.stream();
25074            let mut b = __s_b.launch_builder(&f);
25075            b.arg(qkv_tm)
25076                .arg(&*conv_state)
25077                .arg(w)
25078                .arg(y)
25079                .arg(&cd)
25080                .arg(&ti)
25081                .arg(&dc);
25082            unsafe {
25083                b.launch(cfg)?;
25084            }
25085        }
25086        match (ring_old, pad_len) {
25087            (None, Some(len_d)) => {
25088                let f = self.func("ssm_conv_ring_update_dev_f32");
25089                let n = conv_dim * (d_conv - 1);
25090                let cfg = LaunchConfig::for_num_elems(n as u32);
25091                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25092                let __s_b = self.gpu.stream();
25093                let mut b = __s_b.launch_builder(&f);
25094                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25095                unsafe {
25096                    b.launch(cfg)?;
25097                }
25098            }
25099            (None, None) => {
25100                let f = self.func("ssm_conv_ring_update_f32");
25101                let n = conv_dim * (d_conv - 1);
25102                let cfg = LaunchConfig::for_num_elems(n as u32);
25103                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25104                let __s_b = self.gpu.stream();
25105                let mut b = __s_b.launch_builder(&f);
25106                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25107                unsafe {
25108                    b.launch(cfg)?;
25109                }
25110            }
25111            (Some(_), _) => unreachable!(
25112                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25113            ),
25114        }
25115        Ok(())
25116    }
25117
25118    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
25119    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
25120    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
25121    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
25122    pub fn ssm_conv_ring_rebuild(
25123        &self,
25124        qkv_tm: &CudaSlice<f32>,
25125        ring_old: &CudaSlice<f32>,
25126        conv_state: &mut CudaSlice<f32>,
25127        conv_dim: usize,
25128        tc: usize,
25129        d_conv: usize,
25130    ) -> Result<(), Box<dyn std::error::Error>> {
25131        let f = self.func("ssm_conv_ring_rebuild_f32");
25132        let n = conv_dim * (d_conv - 1);
25133        let cfg = LaunchConfig::for_num_elems(n as u32);
25134        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25135        let __s_b = self.gpu.stream();
25136        let mut b = __s_b.launch_builder(&f);
25137        b.arg(qkv_tm)
25138            .arg(ring_old)
25139            .arg(conv_state)
25140            .arg(&cd)
25141            .arg(&ti)
25142            .arg(&dc);
25143        unsafe {
25144            b.launch(cfg)?;
25145        }
25146        Ok(())
25147    }
25148
25149    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
25150    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
25151    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
25152    /// the argmax + run-spec gates are the authority.
25153    #[allow(clippy::too_many_arguments)]
25154    pub fn gdn_prep_decode(
25155        &self,
25156        conv_out: &CudaSlice<f32>,
25157        beta_raw: &CudaSlice<f32>,
25158        alpha: &CudaSlice<f32>,
25159        dt_bias: &CudaSlice<f32>,
25160        a: &CudaSlice<f32>,
25161        q_l2: &mut CudaSlice<f32>,
25162        k_l2: &mut CudaSlice<f32>,
25163        v_g: &mut CudaSlice<f32>,
25164        beta: &mut CudaSlice<f32>,
25165        g_log: &mut CudaSlice<f32>,
25166        d_state: usize,
25167        num_v: usize,
25168        num_k: usize,
25169        key_dim: usize,
25170        eps: f32,
25171    ) -> Result<(), Box<dyn std::error::Error>> {
25172        let f = self.func("gdn_prep_decode_f32");
25173        let cfg = LaunchConfig {
25174            grid_dim: (num_v as u32, 1, 1),
25175            block_dim: (32, 4, 1),
25176            shared_mem_bytes: 0,
25177        };
25178        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25179        let __s_b = self.gpu.stream();
25180        let mut b = __s_b.launch_builder(&f);
25181        b.arg(conv_out)
25182            .arg(beta_raw)
25183            .arg(alpha)
25184            .arg(dt_bias)
25185            .arg(a)
25186            .arg(q_l2)
25187            .arg(k_l2)
25188            .arg(v_g)
25189            .arg(beta)
25190            .arg(g_log)
25191            .arg(&ds)
25192            .arg(&nv)
25193            .arg(&nk)
25194            .arg(&kd)
25195            .arg(&eps);
25196        unsafe {
25197            b.launch(cfg)?;
25198        }
25199        Ok(())
25200    }
25201
25202    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
25203    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
25204    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
25205    #[allow(clippy::too_many_arguments)]
25206    pub fn ssm_conv1d_gdn(
25207        &self,
25208        qkv_tm: &CudaSlice<f32>,
25209        w: &CudaSlice<f32>,
25210        q_g: &mut CudaSlice<f32>,
25211        k_g: &mut CudaSlice<f32>,
25212        v_g: &mut CudaSlice<f32>,
25213        conv_dim: usize,
25214        t: usize,
25215        d_conv: usize,
25216        d_state: usize,
25217        num_v: usize,
25218        num_k: usize,
25219        key_dim: usize,
25220    ) -> Result<(), Box<dyn std::error::Error>> {
25221        let f = self.func("ssm_conv1d_gdn_f32");
25222        let cfg = LaunchConfig {
25223            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25224            block_dim: (256, 1, 1),
25225            shared_mem_bytes: 0,
25226        };
25227        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25228        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25229        let __s_b = self.gpu.stream();
25230        let mut b = __s_b.launch_builder(&f);
25231        b.arg(qkv_tm)
25232            .arg(w)
25233            .arg(q_g)
25234            .arg(k_g)
25235            .arg(v_g)
25236            .arg(&cd)
25237            .arg(&ti)
25238            .arg(&dc)
25239            .arg(&ds)
25240            .arg(&nv)
25241            .arg(&nk)
25242            .arg(&kd);
25243        unsafe {
25244            b.launch(cfg)?;
25245        }
25246        Ok(())
25247    }
25248
25249    pub fn ssm_conv1d(
25250        &self,
25251        x: &CudaSlice<f32>,
25252        w: &CudaSlice<f32>,
25253        y: &mut CudaSlice<f32>,
25254        conv_dim: usize,
25255        t: usize,
25256        d_conv: usize,
25257        silu: bool,
25258    ) -> Result<(), Box<dyn std::error::Error>> {
25259        let f = self.func("ssm_conv1d_silu_f32");
25260        let cfg = LaunchConfig {
25261            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25262            block_dim: (256, 1, 1),
25263            shared_mem_bytes: 0,
25264        };
25265        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25266        let __s_b = self.gpu.stream();
25267        let mut b = __s_b.launch_builder(&f);
25268        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25269        unsafe {
25270            b.launch(cfg)?;
25271        }
25272        Ok(())
25273    }
25274
25275    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
25276    /// o:[128,H,T]. Single sequence.
25277    pub fn gdn_scan_s128(
25278        &self,
25279        q: &CudaSlice<f32>,
25280        k: &CudaSlice<f32>,
25281        v: &CudaSlice<f32>,
25282        g: &CudaSlice<f32>,
25283        beta: &CudaSlice<f32>,
25284        state_in: &CudaSlice<f32>,
25285        state_out: &mut CudaSlice<f32>,
25286        o: &mut CudaSlice<f32>,
25287        n_head: usize,
25288        t: usize,
25289        scale: f32,
25290    ) -> Result<(), Box<dyn std::error::Error>> {
25291        let f = self.func("gdn_scan_s128");
25292        const S_V: u32 = 128;
25293        const WARP: u32 = 32;
25294        const COLS_PER_BLOCK: u32 = 4;
25295        let cfg = LaunchConfig {
25296            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25297            block_dim: (WARP, COLS_PER_BLOCK, 1),
25298            shared_mem_bytes: 0,
25299        };
25300        let (h, ti) = (n_head as i32, t as i32);
25301        let __s_b = self.gpu.stream();
25302        let mut b = __s_b.launch_builder(&f);
25303        b.arg(q)
25304            .arg(k)
25305            .arg(v)
25306            .arg(g)
25307            .arg(beta)
25308            .arg(state_in)
25309            .arg(state_out)
25310            .arg(o)
25311            .arg(&h)
25312            .arg(&ti)
25313            .arg(&scale);
25314        unsafe {
25315            b.launch(cfg)?;
25316        }
25317        Ok(())
25318    }
25319
25320    // ==== B2' batched decode state ops (decode_batch.rs) ====
25321    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
25322    // Bodies are the single-seq kernels per sequence — bit-identical per row.
25323
25324    #[allow(clippy::too_many_arguments)]
25325    pub fn ssm_conv1d_fused_decode_b(
25326        &self,
25327        qkv_cols: &CudaSlice<f32>,
25328        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25329        w: &CudaSlice<f32>,
25330        conv_outs: &mut CudaSlice<f32>,
25331        conv_dim: usize,
25332        d_conv: usize,
25333        b_n: usize,
25334    ) -> Result<(), Box<dyn std::error::Error>> {
25335        let f = self.func("ssm_conv1d_fused_decode_b_f32");
25336        let cfg = LaunchConfig {
25337            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25338            block_dim: (256, 1, 1),
25339            shared_mem_bytes: 0,
25340        };
25341        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25342        let __s_b = self.gpu.stream();
25343        let mut b = __s_b.launch_builder(&f);
25344        b.arg(qkv_cols)
25345            .arg(conv_state_ptrs)
25346            .arg(w)
25347            .arg(conv_outs)
25348            .arg(&cd)
25349            .arg(&dc);
25350        unsafe {
25351            b.launch(cfg)?;
25352        }
25353        Ok(())
25354    }
25355
25356    #[allow(clippy::too_many_arguments)]
25357    pub fn gdn_prep_decode_b(
25358        &self,
25359        conv_outs: &CudaSlice<f32>,
25360        beta_raws: &CudaSlice<f32>,
25361        alphas: &CudaSlice<f32>,
25362        dt_bias: &CudaSlice<f32>,
25363        a: &CudaSlice<f32>,
25364        q_l2: &mut CudaSlice<f32>,
25365        k_l2: &mut CudaSlice<f32>,
25366        v_g: &mut CudaSlice<f32>,
25367        beta: &mut CudaSlice<f32>,
25368        g_log: &mut CudaSlice<f32>,
25369        d_state: usize,
25370        num_v: usize,
25371        num_k: usize,
25372        key_dim: usize,
25373        eps: f32,
25374        conv_dim: usize,
25375        b_n: usize,
25376    ) -> Result<(), Box<dyn std::error::Error>> {
25377        let f = self.func("gdn_prep_decode_b_f32");
25378        let cfg = LaunchConfig {
25379            grid_dim: (num_v as u32, 1, b_n as u32),
25380            block_dim: (32, 4, 1),
25381            shared_mem_bytes: 0,
25382        };
25383        let (ds, nv, nk, kd, cd) = (
25384            d_state as i32,
25385            num_v as i32,
25386            num_k as i32,
25387            key_dim as i32,
25388            conv_dim as i32,
25389        );
25390        let __s_b = self.gpu.stream();
25391        let mut b = __s_b.launch_builder(&f);
25392        b.arg(conv_outs)
25393            .arg(beta_raws)
25394            .arg(alphas)
25395            .arg(dt_bias)
25396            .arg(a)
25397            .arg(q_l2)
25398            .arg(k_l2)
25399            .arg(v_g)
25400            .arg(beta)
25401            .arg(g_log)
25402            .arg(&ds)
25403            .arg(&nv)
25404            .arg(&nk)
25405            .arg(&kd)
25406            .arg(&eps)
25407            .arg(&cd);
25408        unsafe {
25409            b.launch(cfg)?;
25410        }
25411        Ok(())
25412    }
25413
25414    #[allow(clippy::too_many_arguments)]
25415    pub fn gdn_scan_s128_batched(
25416        &self,
25417        q: &CudaSlice<f32>,
25418        k: &CudaSlice<f32>,
25419        v: &CudaSlice<f32>,
25420        g: &CudaSlice<f32>,
25421        beta: &CudaSlice<f32>,
25422        state_in_ptrs: &cudarc::driver::CudaView<u64>,
25423        state_out_ptrs: &cudarc::driver::CudaView<u64>,
25424        o: &mut CudaSlice<f32>,
25425        n_head: usize,
25426        b_n: usize,
25427        scale: f32,
25428    ) -> Result<(), Box<dyn std::error::Error>> {
25429        let f = self.func("gdn_scan_s128_b");
25430        const S_V: u32 = 128;
25431        const WARP: u32 = 32;
25432        const COLS_PER_BLOCK: u32 = 4;
25433        let cfg = LaunchConfig {
25434            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25435            block_dim: (WARP, COLS_PER_BLOCK, 1),
25436            shared_mem_bytes: 0,
25437        };
25438        let h = n_head as i32;
25439        let __s_b = self.gpu.stream();
25440        let mut b = __s_b.launch_builder(&f);
25441        b.arg(q)
25442            .arg(k)
25443            .arg(v)
25444            .arg(g)
25445            .arg(beta)
25446            .arg(state_in_ptrs)
25447            .arg(state_out_ptrs)
25448            .arg(o)
25449            .arg(&h)
25450            .arg(&scale);
25451        unsafe {
25452            b.launch(cfg)?;
25453        }
25454        Ok(())
25455    }
25456
25457    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
25458    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
25459    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
25460    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
25461    /// numeric class; only the pointer arithmetic moved host-side.
25462    #[allow(clippy::too_many_arguments)]
25463    pub fn ssm_conv1d_fused_decode_b_view(
25464        &self,
25465        qkv_cols: &cudarc::driver::CudaView<f32>,
25466        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25467        w: &CudaSlice<f32>,
25468        conv_outs: &mut CudaSlice<f32>,
25469        conv_dim: usize,
25470        d_conv: usize,
25471        b_n: usize,
25472    ) -> Result<(), Box<dyn std::error::Error>> {
25473        let f = self.func("ssm_conv1d_fused_decode_b_f32");
25474        let cfg = LaunchConfig {
25475            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25476            block_dim: (256, 1, 1),
25477            shared_mem_bytes: 0,
25478        };
25479        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25480        let __s_b = self.gpu.stream();
25481        let mut b = __s_b.launch_builder(&f);
25482        b.arg(qkv_cols)
25483            .arg(conv_state_ptrs)
25484            .arg(w)
25485            .arg(conv_outs)
25486            .arg(&cd)
25487            .arg(&dc);
25488        unsafe {
25489            b.launch(cfg)?;
25490        }
25491        Ok(())
25492    }
25493
25494    #[allow(clippy::too_many_arguments)]
25495    pub fn gdn_prep_decode_b_view(
25496        &self,
25497        conv_outs: &CudaSlice<f32>,
25498        beta_raws: &cudarc::driver::CudaView<f32>,
25499        alphas: &cudarc::driver::CudaView<f32>,
25500        dt_bias: &CudaSlice<f32>,
25501        a: &CudaSlice<f32>,
25502        q_l2: &mut CudaSlice<f32>,
25503        k_l2: &mut CudaSlice<f32>,
25504        v_g: &mut CudaSlice<f32>,
25505        beta: &mut CudaSlice<f32>,
25506        g_log: &mut CudaSlice<f32>,
25507        d_state: usize,
25508        num_v: usize,
25509        num_k: usize,
25510        key_dim: usize,
25511        eps: f32,
25512        conv_dim: usize,
25513        b_n: usize,
25514    ) -> Result<(), Box<dyn std::error::Error>> {
25515        let f = self.func("gdn_prep_decode_b_f32");
25516        let cfg = LaunchConfig {
25517            grid_dim: (num_v as u32, 1, b_n as u32),
25518            block_dim: (32, 4, 1),
25519            shared_mem_bytes: 0,
25520        };
25521        let (ds, nv, nk, kd, cd) = (
25522            d_state as i32,
25523            num_v as i32,
25524            num_k as i32,
25525            key_dim as i32,
25526            conv_dim as i32,
25527        );
25528        let __s_b = self.gpu.stream();
25529        let mut b = __s_b.launch_builder(&f);
25530        b.arg(conv_outs)
25531            .arg(beta_raws)
25532            .arg(alphas)
25533            .arg(dt_bias)
25534            .arg(a)
25535            .arg(q_l2)
25536            .arg(k_l2)
25537            .arg(v_g)
25538            .arg(beta)
25539            .arg(g_log)
25540            .arg(&ds)
25541            .arg(&nv)
25542            .arg(&nk)
25543            .arg(&kd)
25544            .arg(&eps)
25545            .arg(&cd);
25546        unsafe {
25547            b.launch(cfg)?;
25548        }
25549        Ok(())
25550    }
25551
25552    #[allow(clippy::too_many_arguments)]
25553    pub fn gdn_scan_s128_batched_view(
25554        &self,
25555        q: &CudaSlice<f32>,
25556        k: &CudaSlice<f32>,
25557        v: &CudaSlice<f32>,
25558        g: &CudaSlice<f32>,
25559        beta: &CudaSlice<f32>,
25560        state_in_ptrs: &cudarc::driver::CudaView<u64>,
25561        state_out_ptrs: &cudarc::driver::CudaView<u64>,
25562        o: &mut cudarc::driver::CudaViewMut<f32>,
25563        n_head: usize,
25564        b_n: usize,
25565        scale: f32,
25566    ) -> Result<(), Box<dyn std::error::Error>> {
25567        let f = self.func("gdn_scan_s128_b");
25568        const S_V: u32 = 128;
25569        const WARP: u32 = 32;
25570        const COLS_PER_BLOCK: u32 = 4;
25571        let cfg = LaunchConfig {
25572            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25573            block_dim: (WARP, COLS_PER_BLOCK, 1),
25574            shared_mem_bytes: 0,
25575        };
25576        let h = n_head as i32;
25577        let __s_b = self.gpu.stream();
25578        let mut b = __s_b.launch_builder(&f);
25579        b.arg(q)
25580            .arg(k)
25581            .arg(v)
25582            .arg(g)
25583            .arg(beta)
25584            .arg(state_in_ptrs)
25585            .arg(state_out_ptrs)
25586            .arg(o)
25587            .arg(&h)
25588            .arg(&scale);
25589        unsafe {
25590            b.launch(cfg)?;
25591        }
25592        Ok(())
25593    }
25594
25595    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
25596    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
25597    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
25598    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
25599    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
25600    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
25601    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
25602    /// identity law); prime_cache/forward/forward_last are the only callers.
25603    pub fn gdn_chunked_enabled() -> bool {
25604        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25605        *E.get_or_init(|| {
25606            std::env::var("MEMRA_GDN_CHUNKED")
25607                .map(|v| v != "0")
25608                .unwrap_or(true)
25609        })
25610    }
25611
25612    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
25613    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
25614    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
25615    /// of 32 in [32, 128] (kernel row mappings require it).
25616    pub fn gdn_chunk_size() -> usize {
25617        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25618        *C.get_or_init(|| {
25619            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25620                .ok()
25621                .and_then(|v| v.parse().ok())
25622                .unwrap_or(32);
25623            c.clamp(32, 128) / 32 * 32
25624        })
25625    }
25626
25627    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
25628    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
25629    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
25630    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
25631    #[allow(clippy::too_many_arguments)]
25632    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
25633    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
25634    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25635    #[allow(clippy::too_many_arguments)]
25636    pub fn gdn_chunk_k123(
25637        &self,
25638        q: &CudaSlice<f32>,
25639        k: &CudaSlice<f32>,
25640        v: &CudaSlice<f32>,
25641        g: &CudaSlice<f32>,
25642        beta: &CudaSlice<f32>,
25643        wb16: Option<&mut CudaSlice<u8>>,
25644        n_head: usize,
25645        t: usize,
25646        c: usize,
25647        hk: usize,
25648        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25649    ) -> Result<
25650        (
25651            CudaSlice<f32>,
25652            CudaSlice<f32>,
25653            CudaSlice<f32>,
25654            CudaSlice<f32>,
25655        ),
25656        Box<dyn std::error::Error>,
25657    > {
25658        const D: usize = 128;
25659        let h = n_head;
25660        let nc = (t + c - 1) / c;
25661        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25662        let mut gcum = self.uninit(t * h)?;
25663        let mut a = self.uninit(nc * h * c * c)?;
25664        let mut p = self.uninit(nc * h * c * c)?;
25665        let mut u = self.uninit(nc * h * c * D)?;
25666        let mut w = self.uninit(nc * h * c * D)?;
25667        {
25668            // K1
25669            let f = self.func("gdn_chunk_cumgate_f32");
25670            let cfg = LaunchConfig {
25671                grid_dim: (nc as u32, h as u32, 1),
25672                block_dim: (32, 1, 1),
25673                shared_mem_bytes: 0,
25674            };
25675            let __s_b = self.gpu.stream();
25676            let mut b = __s_b.launch_builder(&f);
25677            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25678            unsafe {
25679                b.launch(cfg)?;
25680            }
25681        }
25682        if let Some((qb, kb, pb)) = k2w {
25683            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
25684            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
25685            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25686            let f = self.func("gdn_k2_wgmma");
25687            let cfg = LaunchConfig {
25688                grid_dim: (nc as u32, h as u32, 1),
25689                block_dim: (128, 1, 1),
25690                shared_mem_bytes: 0,
25691            };
25692            let hki = hk as i32;
25693            let __s_b = self.gpu.stream();
25694            let mut b = __s_b.launch_builder(&f);
25695            b.arg(qb)
25696                .arg(kb)
25697                .arg(&gcum)
25698                .arg(beta)
25699                .arg(&mut a)
25700                .arg(&mut *pb)
25701                .arg(&hi)
25702                .arg(&ti)
25703                .arg(&ci)
25704                .arg(&hki);
25705            unsafe {
25706                b.launch(cfg)?;
25707            }
25708        } else if c <= 64 && !portable_mma_gated() {
25709            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
25710            let f = self.func("gdn_chunk_attn_f32");
25711            f.set_attribute(
25712                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25713                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25714            )?;
25715            let jt = ((c + 31) / 32) as u32;
25716            let cfg = LaunchConfig {
25717                grid_dim: (nc as u32, h as u32, jt),
25718                block_dim: (256, 1, 1),
25719                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25720            };
25721            let hki = hk as i32;
25722            let __s_b = self.gpu.stream();
25723            let mut b = __s_b.launch_builder(&f);
25724            b.arg(q)
25725                .arg(k)
25726                .arg(&gcum)
25727                .arg(beta)
25728                .arg(&mut a)
25729                .arg(&mut p)
25730                .arg(&hi)
25731                .arg(&ti)
25732                .arg(&ci)
25733                .arg(&hki);
25734            unsafe {
25735                b.launch(cfg)?;
25736            }
25737        } else {
25738            // K2 generic (C = 128, or the portable target's low-smem fallback)
25739            assert!(
25740                hk == h,
25741                "generic K2 is broadcast-only (de-broadcast rides C==32)"
25742            );
25743            let f = self.func("gdn_chunk_attn_g_f32");
25744            let cfg = LaunchConfig {
25745                grid_dim: (nc as u32, h as u32, 1),
25746                block_dim: (32, 8, 1),
25747                shared_mem_bytes: 0,
25748            };
25749            let __s_b = self.gpu.stream();
25750            let mut b = __s_b.launch_builder(&f);
25751            b.arg(q)
25752                .arg(k)
25753                .arg(&gcum)
25754                .arg(beta)
25755                .arg(&mut a)
25756                .arg(&mut p)
25757                .arg(&hi)
25758                .arg(&ti)
25759                .arg(&ci);
25760            unsafe {
25761                b.launch(cfg)?;
25762            }
25763        }
25764        {
25765            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
25766            let cfg = LaunchConfig {
25767                grid_dim: (nc as u32, h as u32, 1),
25768                block_dim: (256, 1, 1),
25769                shared_mem_bytes: 0,
25770            };
25771            match c {
25772                32 | 64 => {
25773                    let f = self.func(if c == 32 {
25774                        "gdn_chunk_solve32_f32"
25775                    } else {
25776                        "gdn_chunk_solve64_f32"
25777                    });
25778                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
25779                    let wb: u64 = match wb16 {
25780                        Some(d) => self.addr_u8(d),
25781                        None => 0,
25782                    };
25783                    let hki = hk as i32;
25784                    let __s_b = self.gpu.stream();
25785                    let mut b = __s_b.launch_builder(&f);
25786                    b.arg(v)
25787                        .arg(k)
25788                        .arg(&a)
25789                        .arg(&gcum)
25790                        .arg(&mut u)
25791                        .arg(&mut w)
25792                        .arg(&wb)
25793                        .arg(&hi)
25794                        .arg(&ti)
25795                        .arg(&hki);
25796                    unsafe {
25797                        b.launch(cfg)?;
25798                    }
25799                }
25800                _ => {
25801                    assert!(hk == h, "generic K3 is broadcast-only");
25802                    let f = self.func("gdn_chunk_solve_f32");
25803                    let __s_b = self.gpu.stream();
25804                    let mut b = __s_b.launch_builder(&f);
25805                    b.arg(v)
25806                        .arg(k)
25807                        .arg(&a)
25808                        .arg(&gcum)
25809                        .arg(&mut u)
25810                        .arg(&mut w)
25811                        .arg(&hi)
25812                        .arg(&ti)
25813                        .arg(&ci);
25814                    unsafe {
25815                        b.launch(cfg)?;
25816                    }
25817                }
25818            }
25819        }
25820        Ok((gcum, p, u, w))
25821    }
25822
25823    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
25824    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
25825    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
25826    pub fn gdn_db_on() -> bool {
25827        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25828    }
25829
25830    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
25831    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
25832    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
25833    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
25834    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
25835    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
25836    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
25837    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
25838    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25839        !portable_mma_gated()
25840            && c == 32
25841            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25842                Ok("1") => true,
25843                Ok("0") => false,
25844                _ => gdn_mma_default_on(),
25845            }
25846    }
25847
25848    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
25849    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
25850    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
25851    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
25852    /// force would silently produce garbage. Required since the sm_120a mma default
25853    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
25854    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25855        cfg!(memra_hopper_mma)
25856            && self.gdn_mma_enabled(c)
25857            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25858    }
25859
25860    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
25861    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
25862    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
25863    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
25864    #[allow(clippy::too_many_arguments)]
25865    pub fn ssm_conv1d_gdn_state_pad(
25866        &self,
25867        qkv_tm: &cudarc::driver::CudaView<f32>,
25868        conv_state: &mut CudaSlice<f32>,
25869        w: &CudaSlice<f32>,
25870        q_g: &mut CudaSlice<f32>,
25871        k_g: &mut CudaSlice<f32>,
25872        v_g: &mut CudaSlice<f32>,
25873        conv_dim: usize,
25874        t: usize,
25875        d_conv: usize,
25876        d_state: usize,
25877        num_v: usize,
25878        num_k: usize,
25879        key_dim: usize,
25880        hk: usize,
25881        pad_len: Option<&CudaSlice<i32>>,
25882    ) -> Result<(), Box<dyn std::error::Error>> {
25883        assert!(
25884            t >= d_conv - 1,
25885            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25886        );
25887        {
25888            let f = self.func("ssm_conv1d_gdn_state_f32");
25889            let cfg = LaunchConfig {
25890                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25891                block_dim: (256, 1, 1),
25892                shared_mem_bytes: 0,
25893            };
25894            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25895            let (ds, nv, nk, kd, hki) = (
25896                d_state as i32,
25897                num_v as i32,
25898                num_k as i32,
25899                key_dim as i32,
25900                hk as i32,
25901            );
25902            let __s_b = self.gpu.stream();
25903            let mut b = __s_b.launch_builder(&f);
25904            b.arg(qkv_tm)
25905                .arg(&*conv_state)
25906                .arg(w)
25907                .arg(q_g)
25908                .arg(k_g)
25909                .arg(v_g)
25910                .arg(&cd)
25911                .arg(&ti)
25912                .arg(&dc)
25913                .arg(&ds)
25914                .arg(&nv)
25915                .arg(&nk)
25916                .arg(&kd)
25917                .arg(&hki);
25918            unsafe {
25919                b.launch(cfg)?;
25920            }
25921        }
25922        match pad_len {
25923            Some(len_d) => {
25924                let f = self.func("ssm_conv_ring_update_dev_f32");
25925                let n = conv_dim * (d_conv - 1);
25926                let cfg = LaunchConfig::for_num_elems(n as u32);
25927                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25928                let __s_b = self.gpu.stream();
25929                let mut b = __s_b.launch_builder(&f);
25930                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25931                unsafe {
25932                    b.launch(cfg)?;
25933                }
25934            }
25935            None => {
25936                let f = self.func("ssm_conv_ring_update_f32");
25937                let n = conv_dim * (d_conv - 1);
25938                let cfg = LaunchConfig::for_num_elems(n as u32);
25939                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25940                let __s_b = self.gpu.stream();
25941                let mut b = __s_b.launch_builder(&f);
25942                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25943                unsafe {
25944                    b.launch(cfg)?;
25945                }
25946            }
25947        }
25948        Ok(())
25949    }
25950
25951    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
25952    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
25953    /// K2/K3 can write them.
25954    pub fn gdn_chunk_alloc(
25955        &self,
25956        n_head: usize,
25957        t: usize,
25958        c: usize,
25959        hk: usize,
25960    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25961        const D: usize = 128;
25962        assert!(
25963            c == 32,
25964            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25965        );
25966        let h = n_head;
25967        let nc = (t + c - 1) / c;
25968        Ok(GdnChunkBufs {
25969            gcum: self.uninit(t * h)?,
25970            a: self.uninit(nc * h * c * c)?,
25971            p: self.uninit(nc * h * c * c)?,
25972            u: self.uninit(nc * h * c * D)?,
25973            w: self.uninit(nc * h * c * D)?,
25974            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25975            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25976            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25977            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25978            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25979            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25980            o: self.uninit(D * h * t)?,
25981            t,
25982            nc,
25983        })
25984    }
25985
25986    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
25987    pub fn f32_to_bf16_v(
25988        &self,
25989        x: &cudarc::driver::CudaView<f32>,
25990        dst: &mut CudaSlice<u8>,
25991        n: usize,
25992    ) -> Result<(), Box<dyn std::error::Error>> {
25993        let f = self.func("f32_to_bf16_bulk");
25994        let ni = n as i64;
25995        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25996        let __s_b = self.gpu.stream();
25997        let mut b = __s_b.launch_builder(&f);
25998        b.arg(x).arg(dst).arg(&ni);
25999        unsafe {
26000            b.launch(cfg)?;
26001        }
26002        Ok(())
26003    }
26004
26005    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
26006    pub fn f32_to_bf16_into(
26007        &self,
26008        x: &CudaSlice<f32>,
26009        dst: &mut CudaSlice<u8>,
26010        n: usize,
26011    ) -> Result<(), Box<dyn std::error::Error>> {
26012        let f = self.func("f32_to_bf16_bulk");
26013        let ni = n as i64;
26014        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26015        let __s_b = self.gpu.stream();
26016        let mut b = __s_b.launch_builder(&f);
26017        b.arg(x).arg(dst).arg(&ni);
26018        unsafe {
26019            b.launch(cfg)?;
26020        }
26021        Ok(())
26022    }
26023
26024    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
26025    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
26026    pub fn gdn_chunk_k123_vl8(
26027        &self,
26028        seqs: &[GdnSeqVl],
26029        n_head: usize,
26030        hk: usize,
26031        wq: Option<&GdnWVl8>,
26032    ) -> Result<(), Box<dyn std::error::Error>> {
26033        let b = seqs.len();
26034        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
26035        let mut packed = [GdnSeqVl::default(); 8];
26036        packed[..b].copy_from_slice(seqs);
26037        let v = GdnVl8(packed);
26038        let (hi, ci) = (n_head as i32, 32i32);
26039        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26040        {
26041            let f = self.func("gdn_chunk_cumgate_vl");
26042            let cfg = LaunchConfig {
26043                grid_dim: (max_nc, n_head as u32, b as u32),
26044                block_dim: (32, 1, 1),
26045                shared_mem_bytes: 0,
26046            };
26047            let __s_lb = self.gpu.stream();
26048            let mut lb = __s_lb.launch_builder(&f);
26049            lb.arg(&v).arg(&hi).arg(&ci);
26050            unsafe {
26051                lb.launch(cfg)?;
26052            }
26053        }
26054        let hki = hk as i32;
26055        if let Some(w) = wq {
26056            // K2-wgmma vl twin (writes A + pre-masked Pb16)
26057            let f = self.func("gdn_k2_wgmma_vl");
26058            let cfg = LaunchConfig {
26059                grid_dim: (max_nc, n_head as u32, b as u32),
26060                block_dim: (128, 1, 1),
26061                shared_mem_bytes: 0,
26062            };
26063            let __s_lb = self.gpu.stream();
26064            let mut lb = __s_lb.launch_builder(&f);
26065            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26066            unsafe {
26067                lb.launch(cfg)?;
26068            }
26069        } else {
26070            let f = self.func("gdn_chunk_attn_vl");
26071            f.set_attribute(
26072                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26073                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26074            )?;
26075            let cfg = LaunchConfig {
26076                grid_dim: (max_nc, n_head as u32, b as u32),
26077                block_dim: (256, 1, 1),
26078                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26079            };
26080            let __s_lb = self.gpu.stream();
26081            let mut lb = __s_lb.launch_builder(&f);
26082            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26083            unsafe {
26084                lb.launch(cfg)?;
26085            }
26086        }
26087        {
26088            let f = self.func("gdn_chunk_solve32_vl");
26089            let cfg = LaunchConfig {
26090                grid_dim: (max_nc, n_head as u32, b as u32),
26091                block_dim: (256, 1, 1),
26092                shared_mem_bytes: 0,
26093            };
26094            let __s_lb = self.gpu.stream();
26095            let mut lb = __s_lb.launch_builder(&f);
26096            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26097            unsafe {
26098                lb.launch(cfg)?;
26099            }
26100        }
26101        Ok(())
26102    }
26103
26104    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
26105    /// fused gate-prep, 5 launches for every sequence (per-element math identical
26106    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
26107    #[allow(clippy::too_many_arguments)]
26108    pub fn gdn_prep_vl8(
26109        &self,
26110        seqs: &[GdnPrepVl],
26111        conv_w: &CudaSlice<f32>,
26112        dt_bias: &CudaSlice<f32>,
26113        a: &CudaSlice<f32>,
26114        conv_dim: usize,
26115        d_conv: usize,
26116        d_state: usize,
26117        num_v: usize,
26118        num_k: usize,
26119        key_dim: usize,
26120        hk: usize,
26121        eps: f32,
26122    ) -> Result<(), Box<dyn std::error::Error>> {
26123        let b = seqs.len();
26124        assert!(b >= 1 && b <= 8);
26125        let mut packed = [GdnPrepVl::default(); 8];
26126        packed[..b].copy_from_slice(seqs);
26127        let v = GdnPrepVl8(packed);
26128        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26129        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26130        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26131        assert!(
26132            conv_fuse || hk == num_v,
26133            "de-broadcast requires the fused conv"
26134        );
26135        if conv_fuse {
26136            let f = self.func("ssm_conv1d_gdn_state_vl");
26137            let cfg = LaunchConfig {
26138                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26139                block_dim: (256, 1, 1),
26140                shared_mem_bytes: 0,
26141            };
26142            let (dsi, nvi, nki, kdi, hki) = (
26143                d_state as i32,
26144                num_v as i32,
26145                num_k as i32,
26146                key_dim as i32,
26147                hk as i32,
26148            );
26149            let __s_lb = self.gpu.stream();
26150            let mut lb = __s_lb.launch_builder(&f);
26151            lb.arg(&v)
26152                .arg(conv_w)
26153                .arg(&cdi)
26154                .arg(&dci)
26155                .arg(&dsi)
26156                .arg(&nvi)
26157                .arg(&nki)
26158                .arg(&kdi)
26159                .arg(&hki);
26160            unsafe {
26161                lb.launch(cfg)?;
26162            }
26163        } else {
26164            let f = self.func("ssm_conv1d_tm_state_vl");
26165            let cfg = LaunchConfig {
26166                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26167                block_dim: (256, 1, 1),
26168                shared_mem_bytes: 0,
26169            };
26170            let __s_lb = self.gpu.stream();
26171            let mut lb = __s_lb.launch_builder(&f);
26172            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26173            unsafe {
26174                lb.launch(cfg)?;
26175            }
26176        }
26177        {
26178            let f = self.func("ssm_conv_ring_update_vl");
26179            let n = (conv_dim * (d_conv - 1)) as u32;
26180            let cfg = LaunchConfig {
26181                grid_dim: (n.div_ceil(256), 1, b as u32),
26182                block_dim: (256, 1, 1),
26183                shared_mem_bytes: 0,
26184            };
26185            let __s_lb = self.gpu.stream();
26186            let mut lb = __s_lb.launch_builder(&f);
26187            lb.arg(&v).arg(&cdi).arg(&dci);
26188            unsafe {
26189                lb.launch(cfg)?;
26190            }
26191        }
26192        if !conv_fuse {
26193            let f = self.func("qkv_to_gdn_repack_vl");
26194            let n = max_t * (num_v * d_state) as u32;
26195            let cfg = LaunchConfig {
26196                grid_dim: (n.div_ceil(256), 1, b as u32),
26197                block_dim: (256, 1, 1),
26198                shared_mem_bytes: 0,
26199            };
26200            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26201            let __s_lb = self.gpu.stream();
26202            let mut lb = __s_lb.launch_builder(&f);
26203            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26204            unsafe {
26205                lb.launch(cfg)?;
26206            }
26207        }
26208        if Self::l2_v2_on(d_state) {
26209            let f = self.func("gdn_l2_v2_vl");
26210            let cfg = LaunchConfig {
26211                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26212                block_dim: (256, 1, 1),
26213                shared_mem_bytes: 0,
26214            };
26215            let (dsi, nvi) = (d_state as i32, hk as i32);
26216            let __s_lb = self.gpu.stream();
26217            let mut lb = __s_lb.launch_builder(&f);
26218            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26219            unsafe {
26220                lb.launch(cfg)?;
26221            }
26222        } else {
26223            let f = self.func("gdn_l2_vl");
26224            let cfg = LaunchConfig {
26225                grid_dim: (max_t * hk as u32, 2, b as u32),
26226                block_dim: (256, 1, 1),
26227                shared_mem_bytes: 0,
26228            };
26229            let (dsi, nvi) = (d_state as i32, hk as i32);
26230            let __s_lb = self.gpu.stream();
26231            let mut lb = __s_lb.launch_builder(&f);
26232            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26233            unsafe {
26234                lb.launch(cfg)?;
26235            }
26236        }
26237        {
26238            let f = self.func("gdn_gate_prep_vl");
26239            let n = max_t * num_v as u32;
26240            let cfg = LaunchConfig {
26241                grid_dim: (n.div_ceil(256), 1, b as u32),
26242                block_dim: (256, 1, 1),
26243                shared_mem_bytes: 0,
26244            };
26245            let nvi = num_v as i32;
26246            let __s_lb = self.gpu.stream();
26247            let mut lb = __s_lb.launch_builder(&f);
26248            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26249            unsafe {
26250                lb.launch(cfg)?;
26251            }
26252        }
26253        Ok(())
26254    }
26255
26256    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
26257    pub fn gdn_mirror_vl8(
26258        &self,
26259        seqs: &[GdnSeqVl],
26260        n_head: usize,
26261        which: i32,
26262        hk: usize,
26263    ) -> Result<(), Box<dyn std::error::Error>> {
26264        let b = seqs.len();
26265        assert!(b >= 1 && b <= 8);
26266        let mut packed = [GdnSeqVl::default(); 8];
26267        packed[..b].copy_from_slice(seqs);
26268        let v = GdnVl8(packed);
26269        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26270        let max_n = seqs
26271            .iter()
26272            .map(|s| {
26273                if which == 0 {
26274                    s.t as i64 * ept as i64
26275                } else {
26276                    s.nc as i64 * ept as i64 * 32
26277                }
26278            })
26279            .max()
26280            .unwrap();
26281        let f = self.func("gdn_mirror_vl");
26282        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26283        let cfg = LaunchConfig {
26284            grid_dim: (blocks, 1, b as u32),
26285            block_dim: (256, 1, 1),
26286            shared_mem_bytes: 0,
26287        };
26288        let __s_lb = self.gpu.stream();
26289        let mut lb = __s_lb.launch_builder(&f);
26290        lb.arg(&v).arg(&ept).arg(&which);
26291        unsafe {
26292            lb.launch(cfg)?;
26293        }
26294        Ok(())
26295    }
26296
26297    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
26298    pub fn gdn_tail_vl8(
26299        &self,
26300        seqs: &[GdnPrepVl],
26301        norm_w: &CudaSlice<f32>,
26302        d_state: usize,
26303        num_v: usize,
26304        eps: f32,
26305    ) -> Result<(), Box<dyn std::error::Error>> {
26306        let b = seqs.len();
26307        assert!(b >= 1 && b <= 8);
26308        let mut packed = [GdnPrepVl::default(); 8];
26309        packed[..b].copy_from_slice(seqs);
26310        let v = GdnPrepVl8(packed);
26311        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26312        let f = self.func("gated_rmsnorm_f16out_vl");
26313        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
26314        let cfg = LaunchConfig {
26315            grid_dim: (max_t * num_v as u32, 1, b as u32),
26316            block_dim: (128, 1, 1),
26317            shared_mem_bytes: 0,
26318        };
26319        let (dsi, nvi) = (d_state as i32, num_v as i32);
26320        let __s_lb = self.gpu.stream();
26321        let mut lb = __s_lb.launch_builder(&f);
26322        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26323        unsafe {
26324            lb.launch(cfg)?;
26325        }
26326        Ok(())
26327    }
26328
26329    /// Raw device address helpers for the varlen by-value arg struct (single-stream
26330    /// launches; every buffer outlives the call — the f16 FFI discipline).
26331    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26332        use cudarc::driver::DevicePtr;
26333        let s = self.gpu.stream();
26334        let (p, _g) = x.device_ptr(&s);
26335        p as u64
26336    }
26337    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26338        use cudarc::driver::DevicePtrMut;
26339        let s = self.gpu.stream();
26340        let (p, _g) = x.device_ptr_mut(&s);
26341        p as u64
26342    }
26343    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26344        use cudarc::driver::DevicePtr;
26345        let s = self.gpu.stream();
26346        let (p, _g) = x.device_ptr(&s);
26347        p as u64
26348    }
26349    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26350        use cudarc::driver::DevicePtr;
26351        let s = self.gpu.stream();
26352        let (p, _g) = x.device_ptr(&s);
26353        p as u64
26354    }
26355
26356    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
26357    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
26358    /// launches, so this is strictly bit-gateable against them).
26359    pub fn gdn_chunk_vl8(
26360        &self,
26361        seqs: &[GdnSeqVl],
26362        n_head: usize,
26363        scale: f32,
26364        hk: usize,
26365        wq: Option<&GdnWVl8>,
26366    ) -> Result<(), Box<dyn std::error::Error>> {
26367        const NSPLIT: u32 = 4;
26368        let b = seqs.len();
26369        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26370        let mut packed = [GdnSeqVl::default(); 8];
26371        packed[..b].copy_from_slice(seqs);
26372        let v = GdnVl8(packed);
26373        let (hi, ci) = (n_head as i32, 32i32);
26374        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26375        let hki = hk as i32;
26376        if let Some(w) = wq {
26377            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
26378            let f = self.func("gdn_k45_wgmma_vl");
26379            let cfg = LaunchConfig {
26380                grid_dim: (n_head as u32, NSPLIT, b as u32),
26381                block_dim: (256, 1, 1),
26382                shared_mem_bytes: 0,
26383            };
26384            let __s_lb = self.gpu.stream();
26385            let mut lb = __s_lb.launch_builder(&f);
26386            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26387            unsafe {
26388                lb.launch(cfg)?;
26389            }
26390            let _ = max_nc;
26391            return Ok(());
26392        }
26393        {
26394            let f = self.func("gdn_chunk_state_mma_vl");
26395            let cfg = LaunchConfig {
26396                grid_dim: (n_head as u32, NSPLIT, b as u32),
26397                block_dim: (256, 1, 1),
26398                shared_mem_bytes: 0,
26399            };
26400            let __s_lb = self.gpu.stream();
26401            let mut lb = __s_lb.launch_builder(&f);
26402            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26403            unsafe {
26404                lb.launch(cfg)?;
26405            }
26406        }
26407        {
26408            let f = self.func("gdn_chunk_output_mma_vl");
26409            let cfg = LaunchConfig {
26410                grid_dim: (max_nc, n_head as u32, b as u32),
26411                block_dim: (256, 1, 1),
26412                shared_mem_bytes: 0,
26413            };
26414            let __s_lb = self.gpu.stream();
26415            let mut lb = __s_lb.launch_builder(&f);
26416            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26417            unsafe {
26418                lb.launch(cfg)?;
26419            }
26420        }
26421        Ok(())
26422    }
26423    pub fn gdn_scan_chunked(
26424        &self,
26425        q: &CudaSlice<f32>,
26426        k: &CudaSlice<f32>,
26427        v: &CudaSlice<f32>,
26428        g: &CudaSlice<f32>,
26429        beta: &CudaSlice<f32>,
26430        kb16_pre: Option<&CudaSlice<u8>>,
26431        qb16_pre: Option<&CudaSlice<u8>>,
26432        state_in: &CudaSlice<f32>,
26433        state_out: &mut CudaSlice<f32>,
26434        o: &mut CudaSlice<f32>,
26435        n_head: usize,
26436        t: usize,
26437        scale: f32,
26438        c: usize,
26439        hk: usize,
26440    ) -> Result<(), Box<dyn std::error::Error>> {
26441        const D: usize = 128;
26442        const NSPLIT: u32 = 4;
26443        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26444        let h = n_head;
26445        let nc = (t + c - 1) / c;
26446        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26447        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
26448        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
26449        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
26450        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
26451        let gdn_mma_pre = !portable_mma_gated()
26452            && c == 32
26453            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26454                Ok("1") => true,
26455                Ok("0") => false,
26456                _ => gdn_mma_default_on(),
26457            };
26458        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26459            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26460        } else {
26461            None
26462        };
26463        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
26464        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
26465        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
26466        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
26467        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26468            && gdn_mma_pre
26469            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26470        let nk = t * hk * D;
26471        let mut kb16_local: Option<CudaSlice<u8>> = None;
26472        if gdn_mma_pre && kb16_pre.is_none() {
26473            let mut kb = self.alloc_u8_uninit(nk * 2)?;
26474            let f = self.func("f32_to_bf16_bulk");
26475            let n2 = nk as i64;
26476            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26477            let __s_b = self.gpu.stream();
26478            let mut b = __s_b.launch_builder(&f);
26479            b.arg(k).arg(&mut kb).arg(&n2);
26480            unsafe {
26481                b.launch(cfg2)?;
26482            }
26483            kb16_local = Some(kb);
26484        }
26485        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26486        if let Some(kb) = kb16_pre {
26487            assert!(kb.len() >= nk * 2, "kb16_pre too small");
26488        }
26489        let mut qb16: Option<CudaSlice<u8>> = None;
26490        let mut pb16: Option<CudaSlice<u8>> = None;
26491        if gdn_wgmma_pre {
26492            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
26493            // the standalone bulk cvt only serves callers without the prep mirror.
26494            if qb16_pre.is_none() {
26495                let mut qb = self.alloc_u8_uninit(nk * 2)?;
26496                let f = self.func("f32_to_bf16_bulk");
26497                let n2 = nk as i64;
26498                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26499                let __s_b = self.gpu.stream();
26500                let mut b = __s_b.launch_builder(&f);
26501                b.arg(q).arg(&mut qb).arg(&n2);
26502                unsafe {
26503                    b.launch(cfg2)?;
26504                }
26505                qb16 = Some(qb);
26506            } else if let Some(qb) = qb16_pre {
26507                assert!(qb.len() >= nk * 2, "qb16_pre too small");
26508            }
26509            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26510        }
26511        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26512        let k2w = if gdn_wgmma_pre {
26513            Some((
26514                *qb16_ref0.as_ref().unwrap(),
26515                *kb16_ref0.as_ref().unwrap(),
26516                pb16.as_mut().unwrap(),
26517            ))
26518        } else {
26519            None
26520        };
26521        let (gcum, p, u, w) =
26522            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26523        let _ = &w;
26524        let mut y = self.uninit(nc * h * c * D)?;
26525        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
26526        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
26527        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
26528        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
26529        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
26530        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
26531        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
26532        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
26533        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
26534        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
26535        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
26536        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
26537        // sites must agree or the pre-work arms while the scan takes the scalar route.
26538        let gdn_mma = !portable_mma_gated()
26539            && c == 32
26540            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26541                Ok("1") => true,
26542                Ok("0") => false,
26543                _ => gdn_mma_default_on(),
26544            };
26545        if gdn_mma {
26546            let wb16 = wb16_pre
26547                .take()
26548                .expect("mma path pre-allocates wb16 (K3 store fold)");
26549            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26550            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
26551            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
26552            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
26553            // pass runs inside the persistent-M kernel; Y and Ssnap are never
26554            // materialized. New numeric class (gk folds into k^T instead of ys) —
26555            // explicit opt-in until the state-carry battery promotes it. Env read per
26556            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
26557            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
26558            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
26559            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
26560            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
26561            if gdn_wgmma_pre {
26562                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
26563                let qb16 = qb16_ref0.unwrap();
26564                let pb16 = pb16.as_ref().unwrap();
26565                {
26566                    let f = self.func("gdn_k45_wgmma");
26567                    let cfg = LaunchConfig {
26568                        grid_dim: (h as u32, 4, 1),
26569                        block_dim: (256, 1, 1),
26570                        shared_mem_bytes: 0,
26571                    };
26572                    let hki = hk as i32;
26573                    let __s_b = self.gpu.stream();
26574                    let mut b = __s_b.launch_builder(&f);
26575                    b.arg(kb16_ref)
26576                        .arg(&gcum)
26577                        .arg(beta)
26578                        .arg(&u)
26579                        .arg(&wb16)
26580                        .arg(qb16)
26581                        .arg(pb16)
26582                        .arg(o)
26583                        .arg(&scale)
26584                        .arg(state_in)
26585                        .arg(&mut *state_out)
26586                        .arg(&hi)
26587                        .arg(&ti)
26588                        .arg(&ci)
26589                        .arg(&hki);
26590                    unsafe {
26591                        b.launch(cfg)?;
26592                    }
26593                }
26594                return Ok(());
26595            }
26596            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
26597            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
26598            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
26599            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26600            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26601            {
26602                let f = self.func("gdn_chunk_state_mma");
26603                let cfg = LaunchConfig {
26604                    grid_dim: (h as u32, NSPLIT, 1),
26605                    block_dim: (256, 1, 1),
26606                    shared_mem_bytes: 0,
26607                };
26608                let hki = hk as i32;
26609                let __s_b = self.gpu.stream();
26610                let mut b = __s_b.launch_builder(&f);
26611                b.arg(kb16_ref)
26612                    .arg(&gcum)
26613                    .arg(beta)
26614                    .arg(&u)
26615                    .arg(&wb16)
26616                    .arg(&mut y16)
26617                    .arg(&mut ssnap16)
26618                    .arg(state_in)
26619                    .arg(&mut *state_out)
26620                    .arg(&hi)
26621                    .arg(&ti)
26622                    .arg(&ci)
26623                    .arg(&hki);
26624                unsafe {
26625                    b.launch(cfg)?;
26626                }
26627            }
26628            {
26629                // K5-mma (bf16 St/Y consumers)
26630                let f = self.func("gdn_chunk_output_mma");
26631                let jt = ((c + 31) / 32) as u32;
26632                let cfg = LaunchConfig {
26633                    grid_dim: (nc as u32, h as u32, jt),
26634                    block_dim: (256, 1, 1),
26635                    shared_mem_bytes: 0,
26636                };
26637                let hki = hk as i32;
26638                let __s_b = self.gpu.stream();
26639                let mut b = __s_b.launch_builder(&f);
26640                b.arg(q)
26641                    .arg(&gcum)
26642                    .arg(&p)
26643                    .arg(&y16)
26644                    .arg(&ssnap16)
26645                    .arg(o)
26646                    .arg(&hi)
26647                    .arg(&ti)
26648                    .arg(&ci)
26649                    .arg(&scale)
26650                    .arg(&hki);
26651                unsafe {
26652                    b.launch(cfg)?;
26653                }
26654            }
26655            return Ok(());
26656        }
26657        {
26658            // K4 (sequential over chunks inside; blocks col-partition the state)
26659            let f = self.func("gdn_chunk_state_f32");
26660            let cfg = LaunchConfig {
26661                grid_dim: (h as u32, NSPLIT, 1),
26662                block_dim: (256, 1, 1),
26663                shared_mem_bytes: 0,
26664            };
26665            let __s_b = self.gpu.stream();
26666            let mut b = __s_b.launch_builder(&f);
26667            b.arg(k)
26668                .arg(&gcum)
26669                .arg(beta)
26670                .arg(&u)
26671                .arg(&w)
26672                .arg(&mut y)
26673                .arg(&mut ssnap)
26674                .arg(state_in)
26675                .arg(&mut *state_out)
26676                .arg(&hi)
26677                .arg(&ti)
26678                .arg(&ci);
26679            unsafe {
26680                b.launch(cfg)?;
26681            }
26682        }
26683        {
26684            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
26685            let f = self.func("gdn_chunk_output_f32");
26686            let jt = ((c + 31) / 32) as u32;
26687            let cfg = LaunchConfig {
26688                grid_dim: (nc as u32, h as u32, jt),
26689                block_dim: (256, 1, 1),
26690                shared_mem_bytes: 0,
26691            };
26692            let __s_b = self.gpu.stream();
26693            let mut b = __s_b.launch_builder(&f);
26694            b.arg(q)
26695                .arg(&gcum)
26696                .arg(&p)
26697                .arg(&y)
26698                .arg(&ssnap)
26699                .arg(o)
26700                .arg(&hi)
26701                .arg(&ti)
26702                .arg(&ci)
26703                .arg(&scale);
26704            unsafe {
26705                b.launch(cfg)?;
26706            }
26707        }
26708        Ok(())
26709    }
26710
26711    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
26712    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
26713    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
26714    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
26715    ///
26716    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
26717    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
26718    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
26719    #[allow(clippy::too_many_arguments)]
26720    #[allow(clippy::too_many_arguments)]
26721    pub fn gdn_scan_prefill(
26722        &self,
26723        q: &CudaSlice<f32>,
26724        k: &CudaSlice<f32>,
26725        v: &CudaSlice<f32>,
26726        g: &CudaSlice<f32>,
26727        beta: &CudaSlice<f32>,
26728        kb16_pre: Option<&CudaSlice<u8>>,
26729        qb16_pre: Option<&CudaSlice<u8>>,
26730        state_in: &CudaSlice<f32>,
26731        state_out: &mut CudaSlice<f32>,
26732        o: &mut CudaSlice<f32>,
26733        n_head: usize,
26734        t: usize,
26735        scale: f32,
26736        hk: usize,
26737    ) -> Result<(), Box<dyn std::error::Error>> {
26738        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26739            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26740            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26741        }
26742        if Self::gdn_chunked_enabled() && t >= 16 {
26743            self.gdn_scan_chunked(
26744                q,
26745                k,
26746                v,
26747                g,
26748                beta,
26749                kb16_pre,
26750                qb16_pre,
26751                state_in,
26752                state_out,
26753                o,
26754                n_head,
26755                t,
26756                scale,
26757                Self::gdn_chunk_size(),
26758                hk,
26759            )
26760        } else {
26761            assert!(
26762                hk == n_head,
26763                "s128 scan is broadcast-only (prep guarantees by predicate)"
26764            );
26765            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26766        }
26767    }
26768
26769    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
26770    #[allow(clippy::too_many_arguments)]
26771    fn gdn_scan_diff(
26772        &self,
26773        q: &CudaSlice<f32>,
26774        k: &CudaSlice<f32>,
26775        v: &CudaSlice<f32>,
26776        g: &CudaSlice<f32>,
26777        beta: &CudaSlice<f32>,
26778        state_in: &CudaSlice<f32>,
26779        state_out: &mut CudaSlice<f32>,
26780        o: &mut CudaSlice<f32>,
26781        n_head: usize,
26782        t: usize,
26783        scale: f32,
26784    ) -> Result<(), Box<dyn std::error::Error>> {
26785        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26786        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26787        let mut o_c = self.uninit(o.len())?;
26788        let mut st_c = self.uninit(state_out.len())?;
26789        self.gdn_scan_chunked(
26790            q,
26791            k,
26792            v,
26793            g,
26794            beta,
26795            None,
26796            None,
26797            state_in,
26798            &mut st_c,
26799            &mut o_c,
26800            n_head,
26801            t,
26802            scale,
26803            Self::gdn_chunk_size(),
26804            n_head,
26805        )?;
26806        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26807        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26808        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26809        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26810            let mut max_abs = 0f32;
26811            let mut max_rel = 0f32;
26812            let mut sum_rel = 0f64;
26813            for (x, y) in a.iter().zip(b) {
26814                let ad = (x - y).abs();
26815                let rel = ad / x.abs().max(y.abs()).max(1e-3);
26816                if ad > max_abs {
26817                    max_abs = ad;
26818                }
26819                if rel > max_rel {
26820                    max_rel = rel;
26821                }
26822                sum_rel += rel as f64;
26823            }
26824            (max_abs, max_rel, sum_rel / a.len() as f64)
26825        };
26826        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26827        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26828        println!(
26829            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26830                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26831            Self::gdn_chunk_size()
26832        );
26833        Ok(())
26834    }
26835
26836    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
26837    pub fn gdn_glog(
26838        &self,
26839        alpha: &CudaSlice<f32>,
26840        dt_bias: &CudaSlice<f32>,
26841        a: &CudaSlice<f32>,
26842        g_log: &mut CudaSlice<f32>,
26843        n_head: usize,
26844        t: usize,
26845    ) -> Result<(), Box<dyn std::error::Error>> {
26846        let f = self.func("gdn_glog_f32");
26847        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26848        let (h, ti) = (n_head as i32, t as i32);
26849        let __s_b = self.gpu.stream();
26850        let mut b = __s_b.launch_builder(&f);
26851        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26852        unsafe {
26853            b.launch(cfg)?;
26854        }
26855        Ok(())
26856    }
26857
26858    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
26859    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
26860    pub fn sigmoid_v(
26861        &self,
26862        x: &cudarc::driver::CudaView<f32>,
26863        y: &mut CudaSlice<f32>,
26864        n: usize,
26865    ) -> Result<(), Box<dyn std::error::Error>> {
26866        let f = self.func("sigmoid_f32");
26867        let cfg = LaunchConfig::for_num_elems(n as u32);
26868        let ni = n as i32;
26869        let __s_b = self.gpu.stream();
26870        let mut b = __s_b.launch_builder(&f);
26871        b.arg(x).arg(y).arg(&ni);
26872        unsafe {
26873            b.launch(cfg)?;
26874        }
26875        Ok(())
26876    }
26877
26878    pub fn gdn_glog_v(
26879        &self,
26880        alpha: &cudarc::driver::CudaView<f32>,
26881        dt_bias: &CudaSlice<f32>,
26882        a: &CudaSlice<f32>,
26883        g_log: &mut CudaSlice<f32>,
26884        n_head: usize,
26885        t: usize,
26886    ) -> Result<(), Box<dyn std::error::Error>> {
26887        let f = self.func("gdn_glog_f32");
26888        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26889        let (h, ti) = (n_head as i32, t as i32);
26890        let __s_b = self.gpu.stream();
26891        let mut b = __s_b.launch_builder(&f);
26892        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26893        unsafe {
26894            b.launch(cfg)?;
26895        }
26896        Ok(())
26897    }
26898
26899    pub fn sigmoid(
26900        &self,
26901        x: &CudaSlice<f32>,
26902        y: &mut CudaSlice<f32>,
26903        n: usize,
26904    ) -> Result<(), Box<dyn std::error::Error>> {
26905        let f = self.func("sigmoid_f32");
26906        let cfg = LaunchConfig::for_num_elems(n as u32);
26907        let ni = n as i32;
26908        let __s_b = self.gpu.stream();
26909        let mut b = __s_b.launch_builder(&f);
26910        b.arg(x).arg(y).arg(&ni);
26911        unsafe {
26912            b.launch(cfg)?;
26913        }
26914        Ok(())
26915    }
26916
26917    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
26918    /// (replaces sigmoid + mul + convert). Bit-identical class.
26919    pub fn sig_mul_f16out(
26920        &self,
26921        a: &CudaSlice<f32>,
26922        g: &CudaSlice<f32>,
26923        dst: &mut CudaSlice<f32>,
26924        dst16: &mut CudaSlice<u8>,
26925        n: usize,
26926    ) -> Result<(), Box<dyn std::error::Error>> {
26927        let f = self.func("sig_mul_f16out_f32");
26928        let cfg = LaunchConfig::for_num_elems(n as u32);
26929        let ni = n as i32;
26930        let __s_b = self.gpu.stream();
26931        let mut b = __s_b.launch_builder(&f);
26932        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26933        unsafe {
26934            b.launch(cfg)?;
26935        }
26936        Ok(())
26937    }
26938
26939    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
26940    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
26941    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
26942    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
26943    ///
26944    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
26945    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
26946    /// applies the wrong number of distinct gate values.
26947    #[allow(clippy::too_many_arguments)]
26948    pub fn attn_head_gate(
26949        &self,
26950        a: &CudaSlice<f32>,
26951        g: &CudaSlice<f32>,
26952        dst: &mut CudaSlice<f32>,
26953        dst16: Option<&mut CudaSlice<u8>>,
26954        head_dim: usize,
26955        n_head: usize,
26956        t: usize,
26957    ) -> Result<(), Box<dyn std::error::Error>> {
26958        let f = self.func("attn_head_gate_f32");
26959        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26960        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26961        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
26962        let d16: u64 = match dst16 {
26963            Some(d) => self.addr_u8(d),
26964            None => 0,
26965        };
26966        let __s_b = self.gpu.stream();
26967        let mut b = __s_b.launch_builder(&f);
26968        b.arg(a)
26969            .arg(g)
26970            .arg(dst)
26971            .arg(&d16)
26972            .arg(&hd)
26973            .arg(&nh)
26974            .arg(&ti);
26975        unsafe {
26976            b.launch(cfg)?;
26977        }
26978        Ok(())
26979    }
26980
26981    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
26982    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
26983    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
26984    ///
26985    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
26986    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
26987    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
26988    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
26989    #[allow(clippy::too_many_arguments)]
26990    pub fn swiglu_clamped_mul_scaled(
26991        &self,
26992        gate: &CudaSlice<f32>,
26993        up: &CudaSlice<f32>,
26994        gs: f32,
26995        us: f32,
26996        limit: f32,
26997        dst: &mut CudaSlice<f32>,
26998        n: usize,
26999    ) -> Result<(), Box<dyn std::error::Error>> {
27000        debug_assert!(
27001            limit > 1e-6,
27002            "swiglu_clamped needs a live limit; use silu_mul_scaled"
27003        );
27004        let f = self.func("swiglu_clamped_mul_scaled_f32");
27005        let cfg = LaunchConfig::for_num_elems(n as u32);
27006        let ni = n as i32;
27007        let __s_b = self.gpu.stream();
27008        let mut b = __s_b.launch_builder(&f);
27009        b.arg(gate)
27010            .arg(up)
27011            .arg(&gs)
27012            .arg(&us)
27013            .arg(&limit)
27014            .arg(dst)
27015            .arg(&ni);
27016        unsafe {
27017            b.launch(cfg)?;
27018        }
27019        Ok(())
27020    }
27021
27022    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
27023    pub fn gated_rmsnorm(
27024        &self,
27025        o: &CudaSlice<f32>,
27026        w: &CudaSlice<f32>,
27027        z: &CudaSlice<f32>,
27028        dst: &mut CudaSlice<f32>,
27029        ncols: usize,
27030        nrows: usize,
27031        eps: f32,
27032    ) -> Result<(), Box<dyn std::error::Error>> {
27033        let f = self.func("gated_rmsnorm_f32");
27034        let cfg = LaunchConfig {
27035            grid_dim: (nrows as u32, 1, 1),
27036            block_dim: (128, 1, 1),
27037            shared_mem_bytes: 0,
27038        };
27039        let (nc, e) = (ncols as i32, eps);
27040        let __s_b = self.gpu.stream();
27041        let mut b = __s_b.launch_builder(&f);
27042        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27043        unsafe {
27044            b.launch(cfg)?;
27045        }
27046        Ok(())
27047    }
27048
27049    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
27050    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
27051    pub fn gated_rmsnorm_f16out(
27052        &self,
27053        o: &CudaSlice<f32>,
27054        w: &CudaSlice<f32>,
27055        z: &CudaSlice<f32>,
27056        dst: &mut CudaSlice<f32>,
27057        dst16: &mut CudaSlice<u8>,
27058        ncols: usize,
27059        nrows: usize,
27060        eps: f32,
27061    ) -> Result<(), Box<dyn std::error::Error>> {
27062        let f = self.func("gated_rmsnorm_f16out_f32");
27063        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
27064        let cfg = LaunchConfig {
27065            grid_dim: (nrows as u32, 1, 1),
27066            block_dim: (128, 1, 1),
27067            shared_mem_bytes: 0,
27068        };
27069        let (nc, e) = (ncols as i32, eps);
27070        let __s_b = self.gpu.stream();
27071        let mut b = __s_b.launch_builder(&f);
27072        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27073        unsafe {
27074            b.launch(cfg)?;
27075        }
27076        Ok(())
27077    }
27078
27079    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
27080    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
27081    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
27082    #[allow(clippy::too_many_arguments)]
27083    pub fn add_rms_norm_zq8(
27084        &self,
27085        a: &CudaSlice<f32>,
27086        b_in: &CudaSlice<f32>,
27087        w: &CudaSlice<f32>,
27088        res: &mut CudaSlice<f32>,
27089        z: &mut CudaSlice<f32>,
27090        ncols: usize,
27091        nrows: usize,
27092        eps: f32,
27093    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27094        assert!(ncols % 32 == 0);
27095        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27096        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27097        let f = self.func("add_rms_norm_zq8");
27098        let cfg = LaunchConfig {
27099            grid_dim: (nrows as u32, 1, 1),
27100            block_dim: (1024, 1, 1),
27101            shared_mem_bytes: 0,
27102        };
27103        let (nc, ep) = (ncols as i32, eps);
27104        let __s_b = self.gpu.stream();
27105        let mut b = __s_b.launch_builder(&f);
27106        b.arg(a)
27107            .arg(b_in)
27108            .arg(w)
27109            .arg(res)
27110            .arg(z)
27111            .arg(&mut q)
27112            .arg(&mut d)
27113            .arg(&nc)
27114            .arg(&ep);
27115        unsafe {
27116            b.launch(cfg)?;
27117        }
27118        Ok((q, d))
27119    }
27120
27121    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
27122    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
27123    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
27124    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
27125    pub fn gated_rmsnorm_zv(
27126        &self,
27127        o: &CudaSlice<f32>,
27128        w: &CudaSlice<f32>,
27129        z: &cudarc::driver::CudaView<f32>,
27130        dst: &mut CudaSlice<f32>,
27131        ncols: usize,
27132        nrows: usize,
27133        eps: f32,
27134    ) -> Result<(), Box<dyn std::error::Error>> {
27135        let f = self.func("gated_rmsnorm_f32");
27136        let cfg = LaunchConfig {
27137            grid_dim: (nrows as u32, 1, 1),
27138            block_dim: (128, 1, 1),
27139            shared_mem_bytes: 0,
27140        };
27141        let (nc, e) = (ncols as i32, eps);
27142        let __s_b = self.gpu.stream();
27143        let mut b = __s_b.launch_builder(&f);
27144        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27145        unsafe {
27146            b.launch(cfg)?;
27147        }
27148        Ok(())
27149    }
27150
27151    pub fn gated_rmsnorm_f16out_zv(
27152        &self,
27153        o: &CudaSlice<f32>,
27154        w: &CudaSlice<f32>,
27155        z: &cudarc::driver::CudaView<f32>,
27156        dst: &mut CudaSlice<f32>,
27157        dst16: &mut CudaSlice<u8>,
27158        ncols: usize,
27159        nrows: usize,
27160        eps: f32,
27161    ) -> Result<(), Box<dyn std::error::Error>> {
27162        let f = self.func("gated_rmsnorm_f16out_f32");
27163        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
27164        let cfg = LaunchConfig {
27165            grid_dim: (nrows as u32, 1, 1),
27166            block_dim: (128, 1, 1),
27167            shared_mem_bytes: 0,
27168        };
27169        let (nc, e) = (ncols as i32, eps);
27170        let __s_b = self.gpu.stream();
27171        let mut b = __s_b.launch_builder(&f);
27172        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27173        unsafe {
27174            b.launch(cfg)?;
27175        }
27176        Ok(())
27177    }
27178
27179    pub fn gated_rmsnorm_q8_1(
27180        &self,
27181        o: &CudaSlice<f32>,
27182        w: &CudaSlice<f32>,
27183        z: &CudaSlice<f32>,
27184        ncols: usize,
27185        nrows: usize,
27186        eps: f32,
27187    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27188        assert!(ncols % 32 == 0);
27189        let f = self.func("gated_rmsnorm_q8_1");
27190        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27191        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27192        let cfg = LaunchConfig {
27193            grid_dim: (nrows as u32, 1, 1),
27194            block_dim: (128, 1, 1),
27195            shared_mem_bytes: 0,
27196        };
27197        let (nc, ep) = (ncols as i32, eps);
27198        let __s_b = self.gpu.stream();
27199        let mut b = __s_b.launch_builder(&f);
27200        b.arg(o)
27201            .arg(w)
27202            .arg(z)
27203            .arg(&mut out_q)
27204            .arg(&mut out_d)
27205            .arg(&nc)
27206            .arg(&ep);
27207        unsafe {
27208            b.launch(cfg)?;
27209        }
27210        Ok((out_q, out_d))
27211    }
27212
27213    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
27214    pub fn transpose(
27215        &self,
27216        inp: &CudaSlice<f32>,
27217        rows: usize,
27218        cols: usize,
27219    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27220        let f = self.func("transpose_f32");
27221        let mut out = self.zeros(rows * cols)?;
27222        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27223        let (r, c) = (rows as i32, cols as i32);
27224        let __s_b = self.gpu.stream();
27225        let mut b = __s_b.launch_builder(&f);
27226        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27227        unsafe {
27228            b.launch(cfg)?;
27229        }
27230        Ok(out)
27231    }
27232
27233    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
27234    pub fn repeat_heads(
27235        &self,
27236        inp: &CudaSlice<f32>,
27237        out: &mut CudaSlice<f32>,
27238        head_dim: usize,
27239        n_in: usize,
27240        n_out: usize,
27241        t: usize,
27242    ) -> Result<(), Box<dyn std::error::Error>> {
27243        let f = self.func("repeat_heads_f32");
27244        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27245        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27246        let __s_b = self.gpu.stream();
27247        let mut b = __s_b.launch_builder(&f);
27248        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27249        unsafe {
27250            b.launch(cfg)?;
27251        }
27252        Ok(())
27253    }
27254
27255    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
27256    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
27257    ///
27258    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
27259    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
27260    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
27261    pub fn q_gate_split(
27262        &self,
27263        qf: &CudaSlice<f32>,
27264        q_out: &mut CudaSlice<f32>,
27265        gate_out: &mut CudaSlice<f32>,
27266        head_dim: usize,
27267        n_head: usize,
27268        t: usize,
27269    ) -> Result<(), Box<dyn std::error::Error>> {
27270        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27271        let out_need = head_dim * n_head * t;
27272        if q_out.len() < out_need || gate_out.len() < out_need {
27273            return Err(format!(
27274                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27275                q_out.len(),
27276                gate_out.len()
27277            )
27278            .into());
27279        }
27280        let f = self.func("q_gate_split_f32");
27281        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27282        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27283        let __s_b = self.gpu.stream();
27284        let mut b = __s_b.launch_builder(&f);
27285        b.arg(qf)
27286            .arg(q_out)
27287            .arg(gate_out)
27288            .arg(&hd)
27289            .arg(&nh)
27290            .arg(&ti);
27291        unsafe {
27292            b.launch(cfg)?;
27293        }
27294        Ok(())
27295    }
27296
27297    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
27298    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
27299    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
27300    pub fn qkv_to_gdn_repack(
27301        &self,
27302        conv_out: &CudaSlice<f32>,
27303        q_g: &mut CudaSlice<f32>,
27304        k_g: &mut CudaSlice<f32>,
27305        v_g: &mut CudaSlice<f32>,
27306        d_state: usize,
27307        num_v: usize,
27308        num_k: usize,
27309        key_dim: usize,
27310        t: usize,
27311    ) -> Result<(), Box<dyn std::error::Error>> {
27312        let f = self.func("qkv_to_gdn_repack_f32");
27313        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27314        let (ds, nv, nk, kd, ti) = (
27315            d_state as i32,
27316            num_v as i32,
27317            num_k as i32,
27318            key_dim as i32,
27319            t as i32,
27320        );
27321        let __s_b = self.gpu.stream();
27322        let mut b = __s_b.launch_builder(&f);
27323        b.arg(conv_out)
27324            .arg(q_g)
27325            .arg(k_g)
27326            .arg(v_g)
27327            .arg(&ds)
27328            .arg(&nv)
27329            .arg(&nk)
27330            .arg(&kd)
27331            .arg(&ti);
27332        unsafe {
27333            b.launch(cfg)?;
27334        }
27335        Ok(())
27336    }
27337
27338    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
27339    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
27340    pub fn conv_left_pad(
27341        &self,
27342        src: &CudaSlice<f32>,
27343        dst: &mut CudaSlice<f32>,
27344        conv_dim: usize,
27345        t: usize,
27346        pad: usize,
27347    ) -> Result<(), Box<dyn std::error::Error>> {
27348        let f = self.func("conv_left_pad_f32");
27349        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27350        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27351        let __s_b = self.gpu.stream();
27352        let mut b = __s_b.launch_builder(&f);
27353        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27354        unsafe {
27355            b.launch(cfg)?;
27356        }
27357        Ok(())
27358    }
27359
27360    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
27361    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
27362    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
27363    pub fn conv_assemble_and_roll(
27364        &self,
27365        qkv_col: &CudaSlice<f32>,
27366        conv_state: &mut CudaSlice<f32>,
27367        conv_in: &mut CudaSlice<f32>,
27368        conv_dim: usize,
27369        pad: usize,
27370    ) -> Result<(), Box<dyn std::error::Error>> {
27371        let f = self.func("conv_assemble_and_roll_f32");
27372        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27373        let (cd, p) = (conv_dim as i32, pad as i32);
27374        let __s_b = self.gpu.stream();
27375        let mut b = __s_b.launch_builder(&f);
27376        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27377        unsafe {
27378            b.launch(cfg)?;
27379        }
27380        Ok(())
27381    }
27382
27383    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
27384    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
27385    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
27386    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
27387    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
27388    pub fn ssm_conv1d_fused_decode(
27389        &self,
27390        qkv_col: &CudaSlice<f32>,
27391        conv_state: &mut CudaSlice<f32>,
27392        w: &CudaSlice<f32>,
27393        conv_out: &mut CudaSlice<f32>,
27394        conv_dim: usize,
27395        d_conv: usize,
27396    ) -> Result<(), Box<dyn std::error::Error>> {
27397        let f = self.func("ssm_conv1d_fused_decode_f32");
27398        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27399        let (cd, dc) = (conv_dim as i32, d_conv as i32);
27400        let __s_b = self.gpu.stream();
27401        let mut b = __s_b.launch_builder(&f);
27402        b.arg(qkv_col)
27403            .arg(conv_state)
27404            .arg(w)
27405            .arg(conv_out)
27406            .arg(&cd)
27407            .arg(&dc);
27408        unsafe {
27409            b.launch(cfg)?;
27410        }
27411        Ok(())
27412    }
27413
27414    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
27415    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
27416    pub fn slice_range(
27417        &self,
27418        src: &CudaSlice<f32>,
27419        start: usize,
27420        len: usize,
27421    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27422        let host = self.gpu.stream().clone_dtoh(src)?;
27423        self.gpu.stream().synchronize()?;
27424        Ok(self.htod(&host[start..start + len])?)
27425    }
27426}
27427
27428#[cfg(test)]
27429mod target_dispatch_tests {
27430    use super::legacy_quant_gemm_allowed;
27431
27432    #[test]
27433    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27434        // sm_120a native lane
27435        assert!(legacy_quant_gemm_allowed(false, false, false));
27436        assert!(!legacy_quant_gemm_allowed(false, false, true));
27437        // pure portable lane (sm_89): gated
27438        assert!(!legacy_quant_gemm_allowed(true, false, false));
27439        assert!(!legacy_quant_gemm_allowed(true, false, true));
27440        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
27441        assert!(legacy_quant_gemm_allowed(true, true, false));
27442        assert!(!legacy_quant_gemm_allowed(true, true, true));
27443    }
27444
27445    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27446    #[test]
27447    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27448        assert!(!legacy_quant_gemm_allowed(
27449            cfg!(memra_portable_cuda),
27450            cfg!(memra_hopper_mma),
27451            false
27452        ));
27453    }
27454
27455    #[cfg(memra_hopper_mma)]
27456    #[test]
27457    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27458        assert!(legacy_quant_gemm_allowed(
27459            cfg!(memra_portable_cuda),
27460            cfg!(memra_hopper_mma),
27461            false
27462        ));
27463        assert!(super::portable_mma_gated() == false);
27464    }
27465}
27466
27467/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
27468/// inherent methods (inherent methods win name resolution, so no recursion).
27469impl memra_kv::KvDev for Engine {
27470    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27471        Engine::zeros(self, n)
27472    }
27473    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27474        Engine::uninit(self, n)
27475    }
27476    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27477        Engine::alloc_u8(self, n)
27478    }
27479    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27480        Engine::htod_i32(self, v)
27481    }
27482    fn clone_dtod(
27483        &self,
27484        src: &CudaSlice<f32>,
27485    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27486        Engine::clone_dtod(self, src)
27487    }
27488    fn copy_into(
27489        &self,
27490        dst: &mut CudaSlice<f32>,
27491        off: usize,
27492        src: &CudaSlice<f32>,
27493        len: usize,
27494    ) -> Result<(), Box<dyn std::error::Error>> {
27495        Engine::copy_into(self, dst, off, src, len)
27496    }
27497    fn set_i32_one(
27498        &self,
27499        d: &mut CudaSlice<i32>,
27500        v: i32,
27501    ) -> Result<(), Box<dyn std::error::Error>> {
27502        Engine::set_i32_one(self, d, v)
27503    }
27504}
27505
27506#[cfg(test)]
27507mod fused_gate_bounds_tests {
27508    use super::*;
27509
27510    /// The fused `[q|gate]` split's read-site guard, on the device.
27511    ///
27512    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
27513    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
27514    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
27515    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
27516    /// `FusedQGateExtent` before the launch.
27517    ///
27518    /// Catch demonstration for this test (guard temporarily removed, then restored):
27519    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
27520    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
27521    /// the call returns `Err`. Receipt in the lane report.
27522    #[test]
27523    #[ignore = "requires a CUDA GPU"]
27524    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27525        let e = Engine::new(0).unwrap();
27526        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27527        let fused = 2 * head_dim * n_head * t;
27528        let out_n = head_dim * n_head * t;
27529
27530        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
27531        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27532        let mut q = e.uninit(out_n).unwrap();
27533        let mut gate = e.uninit(out_n).unwrap();
27534        let err = e
27535            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27536            .expect_err("half-width wq must be refused, not read past")
27537            .to_string();
27538        assert!(err.contains("NO fused gate"), "{err}");
27539        assert!(err.contains(&format!("{fused}")), "{err}");
27540
27541        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
27542        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
27543        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27544        let wide = e.htod(&host).unwrap();
27545        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27546            .expect("full-width wq splits");
27547        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27548        for tok in 0..t {
27549            for hh in 0..n_head {
27550                for d in 0..head_dim {
27551                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27552                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27553                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27554                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27555                }
27556            }
27557        }
27558
27559        // undersized destinations are refused too (the other half of the extent contract)
27560        let mut small = e.uninit(out_n - 1).unwrap();
27561        assert!(
27562            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27563                .is_err()
27564        );
27565    }
27566}
27567
27568/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
27569/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
27570/// any launch, so the refusal is testable without a device.
27571#[cfg(test)]
27572mod fused_rope_width_tests {
27573    use super::Engine;
27574
27575    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
27576    /// safetensors route derives the same), which is why the fusion is legal there today.
27577    #[test]
27578    fn full_width_is_accepted() {
27579        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27580        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27581        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27582    }
27583
27584    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
27585    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
27586    ///
27587    /// ```text
27588    /// attention.key_length     512   rope.dimension_count     512   (global class)
27589    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
27590    /// ```
27591    ///
27592    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
27593    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
27594    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
27595    /// instead of a silently over-rotated head.
27596    #[test]
27597    fn gemma4_official_artifact_widths_pass() {
27598        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27599        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27600    }
27601
27602    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
27603    /// with no `n_dims`, silently rotating the pass-through band.
27604    #[test]
27605    fn partial_rotary_is_refused_with_the_geometry_named() {
27606        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
27607        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27608            .expect_err("partial rotary must refuse");
27609        let msg = err.to_string();
27610        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27611        assert!(msg.contains("n_rot 64"), "{msg}");
27612        assert!(msg.contains("head_dim 256"), "{msg}");
27613        assert!(
27614            msg.contains("64..256"),
27615            "names the band it would corrupt: {msg}"
27616        );
27617        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
27618        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27619        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
27620        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27621    }
27622}