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    pub fn pool_reserved_used(&self) -> (usize, usize) {
1524        use cudarc::driver::sys;
1525        unsafe {
1526            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1527            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1528                != sys::CUresult::CUDA_SUCCESS
1529            {
1530                return (0, 0);
1531            }
1532            let (mut reserved, mut used) = (0u64, 0u64);
1533            if sys::cuMemPoolGetAttribute(
1534                pool,
1535                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1536                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1537            ) != sys::CUresult::CUDA_SUCCESS
1538            {
1539                return (0, 0);
1540            }
1541            if sys::cuMemPoolGetAttribute(
1542                pool,
1543                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1544                &mut used as *mut u64 as *mut core::ffi::c_void,
1545            ) != sys::CUresult::CUDA_SUCCESS
1546            {
1547                return (0, 0);
1548            }
1549            (reserved as usize, used as usize)
1550        }
1551    }
1552
1553    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1554    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1555    pub fn stream(&self) -> Arc<CudaStream> {
1556        self.gpu.stream()
1557    }
1558    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1559    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1560    pub fn gkv_on() -> bool {
1561        memra_kv::gkv_on()
1562    }
1563
1564    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1565    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1566    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1567    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1568    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1569    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1570    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1571    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1572    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1573    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1574    /// ON for both — no acceptance cost measured.
1575    pub fn wkv_on() -> bool {
1576        memra_kv::wkv_on()
1577    }
1578
1579    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1580    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1581    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1582    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1583    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1584    pub fn kv_fp8_on() -> bool {
1585        memra_kv::kv_fp8_on()
1586    }
1587
1588    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1589    /// when the fp8-globals arm is on; everything else from the default flash module.
1590    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1591        if head_dim == 512 && Self::gkv_on() {
1592            self.func_g(name)
1593        } else {
1594            self.func(name)
1595        }
1596    }
1597
1598    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1599    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1600    /// per-format fatbins; fall back to the base modules for those.
1601    fn func_g(&self, name: &str) -> CudaFunction {
1602        let m = self.flash_g.get_or_init(|| {
1603            self.gpu
1604                .ctx
1605                .load_module(cudarc::nvrtc::Ptx::from_binary(
1606                    FLASH_FATBIN_KF8VF8.to_vec(),
1607                ))
1608                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1609        });
1610        let key = format!("g:{name}");
1611        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1612            return f.clone();
1613        }
1614        let f = match m.load_function(name) {
1615            Ok(f) => f,
1616            Err(_) => self.func(name),
1617        };
1618        self.fn_cache.lock().unwrap().insert(key, f.clone());
1619        f
1620    }
1621
1622    fn func(&self, name: &str) -> CudaFunction {
1623        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1624        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1625        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1626            return f.clone();
1627        }
1628        let f = self
1629            .module
1630            .load_function(name)
1631            .or_else(|_| self.hybrid.load_function(name))
1632            .or_else(|_| self.qmatvec.load_function(name))
1633            .or_else(|_| self.flash.load_function(name))
1634            .or_else(|_| self.gemm.load_function(name))
1635            .or_else(|_| self.router.load_function(name))
1636            .or_else(|_| self.sample.load_function(name))
1637            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1638        self.fn_cache
1639            .lock()
1640            .unwrap()
1641            .insert(name.to_string(), f.clone());
1642        f
1643    }
1644
1645    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1646    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1647    pub fn scatter_trim_logits(
1648        &self,
1649        src: &CudaSlice<f32>,
1650        d2t: &CudaSlice<u32>,
1651        dst: &mut CudaSlice<f32>,
1652        d_vocab: usize,
1653        n_vocab: usize,
1654    ) -> Result<(), Box<dyn std::error::Error>> {
1655        let f1 = self.func("scatter_trim_logits_f32");
1656        let f2 = self.func("scatter_trim_logits_pass2_f32");
1657        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1658        let cfg1 = LaunchConfig {
1659            grid_dim: (256, 1, 1),
1660            block_dim: (256, 1, 1),
1661            shared_mem_bytes: 0,
1662        };
1663        let __s_b1 = self.gpu.stream();
1664        let mut b1 = __s_b1.launch_builder(&f1);
1665        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1666        unsafe {
1667            b1.launch(cfg1)?;
1668        }
1669        let cfg2 = LaunchConfig {
1670            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1671            block_dim: (256, 1, 1),
1672            shared_mem_bytes: 0,
1673        };
1674        let __s_b2 = self.gpu.stream();
1675        let mut b2 = __s_b2.launch_builder(&f2);
1676        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1677        unsafe {
1678            b2.launch(cfg2)?;
1679        }
1680        Ok(())
1681    }
1682
1683    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1684    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1685
1686    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1687    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1688    #[allow(clippy::too_many_arguments)]
1689    pub fn filter_stats(
1690        &self,
1691        x: &CudaSlice<f32>,
1692        row_stride: usize,
1693        rows: &CudaSlice<i32>,
1694        out_th: &mut CudaSlice<f32>,
1695        out_z: &mut CudaSlice<f32>,
1696        out_max: &mut CudaSlice<f32>,
1697        n: usize,
1698        nrow: usize,
1699        temp: f32,
1700        top_k: i32,
1701        top_p: f32,
1702        min_p: f32,
1703    ) -> Result<(), Box<dyn std::error::Error>> {
1704        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1705        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1706        // L2-resident, so the extra passes are near-free while the per-thread selection list
1707        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1708        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1709        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1710        //
1711        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1712        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1713        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1714        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1715        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1716        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1717        //
1718        // DETERMINISTIC KEYING (hermes finding, fixed 2026-08-23): the old admission
1719        // `16*nrow <= sm_count` fell back to the single-block program PER CALL when a tick
1720        // carried too many rows — and the two programs are NOT bit-identical (measured
1721        // ~1e-7 rel on the renorm mass: different f32 partial-sum order), so a request's
1722        // sampling threshold arithmetic depended on how many rows shared its serve tick.
1723        // Coop is now THE program on every coop-capable device: rows are CHUNKED to the
1724        // co-residency cap (sm_count/16 rows per cooperative launch) and each row's
1725        // arithmetic uses only its own 16 slices + its own ws region, so the per-row bits
1726        // are independent of batch width by construction — the kernel-check
1727        // FILTER-COOP-CHUNK cell pins exactly that. The single-block program remains only
1728        // behind the deployment-keyed seams: MEMRA_FILTER_COOP=0, or a device with
1729        // sm_count < 16 (fixed per device class, never per call).
1730        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1731        let coop_on =
1732            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1733        if coop_on && self.sm_count() >= 16 {
1734            let cap = self.sm_count() as usize / 16;
1735            let mut done = 0usize;
1736            while done < nrow {
1737                let chunk = cap.min(nrow - done);
1738                self.filter_stats_coop_chunk(
1739                    x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1740                    top_p, min_p,
1741                )?;
1742                done += chunk;
1743            }
1744            return Ok(());
1745        }
1746        self.filter_stats_plain_program(
1747            x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1748        )
1749    }
1750
1751    /// One cooperative `filter_stats` launch over rows `row0..row0+chunk` (pub so the
1752    /// kernel-check FILTER-COOP-CHUNK cell can pin batch-width independence directly).
1753    /// The kernel indexes `rows`/outputs by blockIdx.y, so the chunk is expressed as
1754    /// sub-views at `row0` — per-row arithmetic is untouched by the offset.
1755    #[allow(clippy::too_many_arguments)]
1756    pub fn filter_stats_coop_chunk(
1757        &self,
1758        x: &CudaSlice<f32>,
1759        row_stride: usize,
1760        rows: &CudaSlice<i32>,
1761        row0: usize,
1762        out_th: &mut CudaSlice<f32>,
1763        out_z: &mut CudaSlice<f32>,
1764        out_max: &mut CudaSlice<f32>,
1765        n: usize,
1766        chunk: usize,
1767        temp: f32,
1768        top_k: i32,
1769        top_p: f32,
1770        min_p: f32,
1771    ) -> Result<(), Box<dyn std::error::Error>> {
1772        let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1773        let f = self.func("filter_stats_coop_f32");
1774        let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1775        let cfg = LaunchConfig {
1776            grid_dim: (16, chunk as u32, 1),
1777            block_dim: (512, 1, 1),
1778            shared_mem_bytes: 0,
1779        };
1780        let rows_v = rows.slice(row0..row0 + chunk);
1781        let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1782        let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1783        let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1784        let __s_b = self.gpu.stream();
1785        let mut b = __s_b.launch_builder(&f);
1786        b.arg(x)
1787            .arg(&rs)
1788            .arg(&rows_v)
1789            .arg(&mut th_v)
1790            .arg(&mut z_v)
1791            .arg(&mut mx_v)
1792            .arg(&mut ws)
1793            .arg(&ni)
1794            .arg(&nr)
1795            .arg(&temp)
1796            .arg(&top_k)
1797            .arg(&top_p)
1798            .arg(&min_p);
1799        unsafe {
1800            b.launch_cooperative(cfg)?;
1801        }
1802        Ok(())
1803    }
1804
1805    /// The single-block-per-row `filter_stats` program (the pre-coop form; the
1806    /// MEMRA_FILTER_COOP=0 rollback and the occupancy fallback). Gate-callable twin of
1807    /// `filter_stats_coop_program`.
1808    #[allow(clippy::too_many_arguments)]
1809    pub fn filter_stats_plain_program(
1810        &self,
1811        x: &CudaSlice<f32>,
1812        row_stride: usize,
1813        rows: &CudaSlice<i32>,
1814        out_th: &mut CudaSlice<f32>,
1815        out_z: &mut CudaSlice<f32>,
1816        out_max: &mut CudaSlice<f32>,
1817        n: usize,
1818        nrow: usize,
1819        temp: f32,
1820        top_k: i32,
1821        top_p: f32,
1822        min_p: f32,
1823    ) -> Result<(), Box<dyn std::error::Error>> {
1824        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1825        let f = self.func("filter_stats_f32");
1826        let cfg = LaunchConfig {
1827            grid_dim: (nrow as u32, 1, 1),
1828            block_dim: (1024, 1, 1),
1829            shared_mem_bytes: 0,
1830        };
1831        let __s_b = self.gpu.stream();
1832        let mut b = __s_b.launch_builder(&f);
1833        b.arg(x)
1834            .arg(&rs)
1835            .arg(rows)
1836            .arg(&mut *out_th)
1837            .arg(&mut *out_z)
1838            .arg(&mut *out_max)
1839            .arg(&ni)
1840            .arg(&nr)
1841            .arg(&temp)
1842            .arg(&top_k)
1843            .arg(&top_p)
1844            .arg(&min_p);
1845        unsafe {
1846            b.launch(cfg)?;
1847        }
1848        Ok(())
1849    }
1850
1851    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1852    #[allow(clippy::too_many_arguments)]
1853    pub fn softmax_gather_filtered(
1854        &self,
1855        x: &CudaSlice<f32>,
1856        row_stride: usize,
1857        ids: &CudaSlice<u32>,
1858        rows: &CudaSlice<i32>,
1859        th: &CudaSlice<f32>,
1860        z: &CudaSlice<f32>,
1861        out: &mut CudaSlice<f32>,
1862        n: usize,
1863        npair: usize,
1864        temp: f32,
1865    ) -> Result<(), Box<dyn std::error::Error>> {
1866        let f = self.func("softmax_gather_filtered_f32");
1867        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1868        let cfg = LaunchConfig {
1869            grid_dim: (npair as u32, 1, 1),
1870            block_dim: (256, 1, 1),
1871            shared_mem_bytes: 0,
1872        };
1873        let __s_b = self.gpu.stream();
1874        let mut b = __s_b.launch_builder(&f);
1875        b.arg(x)
1876            .arg(&rs)
1877            .arg(ids)
1878            .arg(rows)
1879            .arg(th)
1880            .arg(z)
1881            .arg(&mut *out)
1882            .arg(&ni)
1883            .arg(&np)
1884            .arg(&temp);
1885        unsafe {
1886            b.launch(cfg)?;
1887        }
1888        Ok(())
1889    }
1890
1891    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1892    #[allow(clippy::too_many_arguments)]
1893    pub fn residual_sample_filtered(
1894        &self,
1895        p: &CudaSlice<f32>,
1896        q: Option<&CudaSlice<f32>>,
1897        n: usize,
1898        temp: f32,
1899        seed: u64,
1900        stream_pos: u32,
1901        p_stats: (f32, f32, f32),
1902        q_stats: (f32, f32, f32),
1903        out_tok: &mut CudaSlice<u32>,
1904    ) -> Result<(), Box<dyn std::error::Error>> {
1905        let f = self.func("residual_sample_filtered_f32");
1906        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1907        let has_q: i32 = q.is_some() as i32;
1908        let qbuf = q.unwrap_or(p);
1909        let (pm, pth, pz) = p_stats;
1910        let (qm, qth, qz) = q_stats;
1911        let cfg = LaunchConfig {
1912            grid_dim: (1, 1, 1),
1913            block_dim: (1024, 1, 1),
1914            shared_mem_bytes: 0,
1915        };
1916        let __s_b = self.gpu.stream();
1917        let mut b = __s_b.launch_builder(&f);
1918        b.arg(p)
1919            .arg(qbuf)
1920            .arg(&has_q)
1921            .arg(&ni)
1922            .arg(&temp)
1923            .arg(&slo)
1924            .arg(&shi)
1925            .arg(&stream_pos)
1926            .arg(&pm)
1927            .arg(&pth)
1928            .arg(&pz)
1929            .arg(&qm)
1930            .arg(&qth)
1931            .arg(&qz)
1932            .arg(&mut *out_tok);
1933        unsafe {
1934            b.launch(cfg)?;
1935        }
1936        Ok(())
1937    }
1938
1939    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1940    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1941    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1942    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1943    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1944    #[allow(clippy::too_many_arguments)]
1945    pub fn residual_sample_sparse_q(
1946        &self,
1947        p: &CudaSlice<f32>,
1948        cand_ids: &CudaSlice<u32>,
1949        q_probs: &CudaSlice<f32>,
1950        n_cand: usize,
1951        n: usize,
1952        temp: f32,
1953        seed: u64,
1954        stream_pos: u32,
1955        p_stats: (f32, f32, f32),
1956        out_tok: &mut CudaSlice<u32>,
1957    ) -> Result<(), Box<dyn std::error::Error>> {
1958        assert!(
1959            n_cand >= 1 && n_cand <= 32,
1960            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1961        );
1962        let f = self.func("residual_sample_sparse_q_f32");
1963        let (ni, nc) = (n as i32, n_cand as i32);
1964        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1965        let (pm, pth, pz) = p_stats;
1966        let cfg = LaunchConfig {
1967            grid_dim: (1, 1, 1),
1968            block_dim: (1024, 1, 1),
1969            shared_mem_bytes: 0,
1970        };
1971        let __s_b = self.gpu.stream();
1972        let mut b = __s_b.launch_builder(&f);
1973        b.arg(p)
1974            .arg(cand_ids)
1975            .arg(q_probs)
1976            .arg(&nc)
1977            .arg(&ni)
1978            .arg(&temp)
1979            .arg(&slo)
1980            .arg(&shi)
1981            .arg(&stream_pos)
1982            .arg(&pm)
1983            .arg(&pth)
1984            .arg(&pz)
1985            .arg(&mut *out_tok);
1986        unsafe {
1987            b.launch(cfg)?;
1988        }
1989        Ok(())
1990    }
1991
1992    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1993    #[allow(clippy::too_many_arguments)]
1994    pub fn gumbel_perturb_filtered(
1995        &self,
1996        x: &CudaSlice<f32>,
1997        y: &mut CudaSlice<f32>,
1998        n: usize,
1999        seed: u64,
2000        stream_pos: u32,
2001        temp: f32,
2002        row_max: f32,
2003        th: f32,
2004    ) -> Result<(), Box<dyn std::error::Error>> {
2005        let f = self.func("gumbel_perturb_filtered_f32");
2006        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2007        let cfg = LaunchConfig {
2008            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2009            block_dim: (256, 1, 1),
2010            shared_mem_bytes: 0,
2011        };
2012        let __s_b = self.gpu.stream();
2013        let mut b = __s_b.launch_builder(&f);
2014        b.arg(x)
2015            .arg(&mut *y)
2016            .arg(&ni)
2017            .arg(&slo)
2018            .arg(&shi)
2019            .arg(&stream_pos)
2020            .arg(&temp)
2021            .arg(&row_max)
2022            .arg(&th);
2023        unsafe {
2024            b.launch(cfg)?;
2025        }
2026        Ok(())
2027    }
2028
2029    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
2030    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
2031    /// filtered rejection sampling exact for the penalized target.
2032    #[allow(clippy::too_many_arguments)]
2033    pub fn penalize_logits(
2034        &self,
2035        x: &mut CudaSlice<f32>,
2036        hist: &CudaSlice<u32>,
2037        n_hist: usize,
2038        rep: f32,
2039        freq: f32,
2040        present: f32,
2041        n: usize,
2042    ) -> Result<(), Box<dyn std::error::Error>> {
2043        if n_hist == 0 {
2044            return Ok(());
2045        }
2046        let f = self.func("penalize_logits_f32");
2047        let (nh, ni) = (n_hist as i32, n as i32);
2048        let cfg = LaunchConfig {
2049            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2050            block_dim: (128, 1, 1),
2051            shared_mem_bytes: 0,
2052        };
2053        let __s_b = self.gpu.stream();
2054        let mut b = __s_b.launch_builder(&f);
2055        b.arg(&mut *x)
2056            .arg(hist)
2057            .arg(&nh)
2058            .arg(&rep)
2059            .arg(&freq)
2060            .arg(&present)
2061            .arg(&ni);
2062        unsafe {
2063            b.launch(cfg)?;
2064        }
2065        Ok(())
2066    }
2067
2068    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
2069    #[allow(clippy::too_many_arguments)]
2070    pub fn penalize_logits_rows(
2071        &self,
2072        x: &mut CudaSlice<f32>,
2073        hist: &CudaSlice<u32>,
2074        n_hist: usize,
2075        rep: f32,
2076        freq: f32,
2077        present: f32,
2078        n: usize,
2079        nrow: usize,
2080    ) -> Result<(), Box<dyn std::error::Error>> {
2081        if n_hist == 0 || nrow == 0 {
2082            return Ok(());
2083        }
2084        let f = self.func("penalize_logits_rows_f32");
2085        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2086        let cfg = LaunchConfig {
2087            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2088            block_dim: (128, 1, 1),
2089            shared_mem_bytes: 0,
2090        };
2091        let __s_b = self.gpu.stream();
2092        let mut b = __s_b.launch_builder(&f);
2093        b.arg(&mut *x)
2094            .arg(hist)
2095            .arg(&nh)
2096            .arg(&rep)
2097            .arg(&freq)
2098            .arg(&present)
2099            .arg(&ni)
2100            .arg(&nr);
2101        unsafe {
2102            b.launch(cfg)?;
2103        }
2104        Ok(())
2105    }
2106
2107    /// Heterogeneous serving-batch penalties over host-maintained sparse window counts.
2108    /// `offsets[r]..offsets[r+1]` indexes the unique positive-count `(id,count)` entries for logits row
2109    /// `rows[r]`; each row may carry independent repetition/frequency/presence coefficients.
2110    /// One thread owns one distinct logit, so the kernel needs neither atomics nor the
2111    /// history-squared dedup scan used by the speculative raw-history oracle.
2112    #[allow(clippy::too_many_arguments)]
2113    pub fn penalize_logits_sparse_rows(
2114        &self,
2115        x: &mut CudaSlice<f32>,
2116        ids: &[u32],
2117        counts: &[u32],
2118        offsets: &[i32],
2119        rows: &[i32],
2120        reps: &[f32],
2121        freqs: &[f32],
2122        presents: &[f32],
2123        n: usize,
2124    ) -> Result<(), Box<dyn std::error::Error>> {
2125        let nrow = rows.len();
2126        if nrow == 0 {
2127            return Ok(());
2128        }
2129        let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2130        let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2131        let entry_count =
2132            i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2133        if ids.len() != counts.len()
2134            || offsets.len() != nrow + 1
2135            || reps.len() != nrow
2136            || freqs.len() != nrow
2137            || presents.len() != nrow
2138            || offsets.first().copied() != Some(0)
2139            || offsets.last().copied() != Some(entry_count)
2140        {
2141            return Err("sparse penalty row metadata shape mismatch".into());
2142        }
2143        if counts.contains(&0) {
2144            return Err("sparse penalty counts must be positive".into());
2145        }
2146        let mut max_len = 0usize;
2147        for pair in offsets.windows(2) {
2148            if pair[0] < 0 || pair[1] < pair[0] {
2149                return Err("sparse penalty offsets must be monotonic".into());
2150            }
2151            max_len = max_len.max((pair[1] - pair[0]) as usize);
2152        }
2153        if max_len == 0 {
2154            return Ok(());
2155        }
2156
2157        let mut seen = std::collections::HashSet::with_capacity(ids.len());
2158        for (r, &row) in rows.iter().enumerate() {
2159            if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2160                return Err("sparse penalty row index exceeds logits shape".into());
2161            }
2162            let begin = offsets[r] as usize;
2163            let end = offsets[r + 1] as usize;
2164            for &id in &ids[begin..end] {
2165                if id as usize >= n {
2166                    return Err("sparse penalty token id exceeds logits row".into());
2167                }
2168                if !seen.insert((row, id)) {
2169                    return Err("sparse penalty entries must be unique per logits row".into());
2170                }
2171            }
2172        }
2173
2174        // SAFETY: the checks above establish every invariant of the launch-only helper.
2175        unsafe {
2176            self.penalize_logits_sparse_rows_unchecked(
2177                x, ids, counts, offsets, rows, reps, freqs, presents, n,
2178            )
2179        }
2180    }
2181
2182    /// Launch-only form for the serving hot path, whose `HashMap`-backed producer already
2183    /// guarantees unique ids and whose rows are enumerated from the live batch.
2184    ///
2185    /// # Safety
2186    ///
2187    /// Shapes must match the safe wrapper, offsets must be monotonic and in bounds, every row
2188    /// must index `x`, and each `(row,id)` pair must occur at most once. Token ids outside the
2189    /// logits row are safe no-ops because the kernel bounds-checks them before computing `x`.
2190    #[allow(clippy::too_many_arguments)]
2191    pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2192        &self,
2193        x: &mut CudaSlice<f32>,
2194        ids: &[u32],
2195        counts: &[u32],
2196        offsets: &[i32],
2197        rows: &[i32],
2198        reps: &[f32],
2199        freqs: &[f32],
2200        presents: &[f32],
2201        n: usize,
2202    ) -> Result<(), Box<dyn std::error::Error>> {
2203        let nrow = rows.len();
2204        if nrow == 0 {
2205            return Ok(());
2206        }
2207        let max_len = offsets
2208            .windows(2)
2209            .map(|pair| (pair[1] - pair[0]) as usize)
2210            .max()
2211            .unwrap_or(0);
2212        if max_len == 0 {
2213            return Ok(());
2214        }
2215        let ids_d = self.htod_u32_v(ids)?;
2216        let counts_d = self.htod_u32_v(counts)?;
2217        let offsets_d = self.htod_i32(offsets)?;
2218        let rows_d = self.htod_i32(rows)?;
2219        let reps_d = self.htod(reps)?;
2220        let freqs_d = self.htod(freqs)?;
2221        let presents_d = self.htod(presents)?;
2222        let f = self.func("penalize_logits_sparse_rows_f32");
2223        let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2224        let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2225        let cfg = LaunchConfig {
2226            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2227            block_dim: (128, 1, 1),
2228            shared_mem_bytes: 0,
2229        };
2230        let __s_b = self.gpu.stream();
2231        let mut b = __s_b.launch_builder(&f);
2232        b.arg(&mut *x)
2233            .arg(&ids_d)
2234            .arg(&counts_d)
2235            .arg(&offsets_d)
2236            .arg(&rows_d)
2237            .arg(&reps_d)
2238            .arg(&freqs_d)
2239            .arg(&presents_d)
2240            .arg(&ni)
2241            .arg(&nr);
2242        unsafe {
2243            b.launch(cfg)?;
2244        }
2245        Ok(())
2246    }
2247
2248    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
2249    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
2250    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
2251    /// is the within-round evolving penalty state block drafting needs: verify row r's
2252    /// target is penalized by every token committed before it INCLUDING same-round
2253    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
2254    /// approximation this exists to replace on the dspark route.
2255    #[allow(clippy::too_many_arguments)]
2256    pub fn penalize_logits_rows_inc(
2257        &self,
2258        x: &mut CudaSlice<f32>,
2259        hist: &CudaSlice<u32>,
2260        n_hist0: usize,
2261        rep: f32,
2262        freq: f32,
2263        present: f32,
2264        n: usize,
2265        nrow: usize,
2266        win: usize,
2267    ) -> Result<(), Box<dyn std::error::Error>> {
2268        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2269            return Ok(());
2270        }
2271        debug_assert!(
2272            hist.len() >= n_hist0 + nrow - 1,
2273            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2274        );
2275        let f = self.func("penalize_logits_rows_inc_f32");
2276        let max_len = win.min(n_hist0 + nrow - 1).max(1);
2277        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2278        let cfg = LaunchConfig {
2279            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2280            block_dim: (128, 1, 1),
2281            shared_mem_bytes: 0,
2282        };
2283        let __s_b = self.gpu.stream();
2284        let mut b = __s_b.launch_builder(&f);
2285        b.arg(&mut *x)
2286            .arg(hist)
2287            .arg(&nh)
2288            .arg(&rep)
2289            .arg(&freq)
2290            .arg(&present)
2291            .arg(&ni)
2292            .arg(&nr)
2293            .arg(&wi);
2294        unsafe {
2295            b.launch(cfg)?;
2296        }
2297        Ok(())
2298    }
2299
2300    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
2301    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
2302    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
2303    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
2304    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
2305    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
2306    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
2307    pub fn wpf_level() -> u32 {
2308        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2309        *ON.get_or_init(|| {
2310            std::env::var("MEMRA_WPF")
2311                .ok()
2312                .and_then(|v| v.parse().ok())
2313                .unwrap_or(1)
2314        })
2315    }
2316
2317    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
2318    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
2319    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
2320    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
2321    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
2322    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
2323    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
2324    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
2325    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
2326    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
2327    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
2328    /// Prefer `exact_scope` — the RAII form — anywhere a `?` can exit the scope: a manual
2329    /// true/false pair leaves the flag LATCHED engine-wide when an error propagates
2330    /// between the two calls (hermes finding on dspark_spec_session_burst, fixed
2331    /// 2026-08-23), and every later request then runs the exact-GEMM program.
2332    pub fn set_verify_exact(&self, on: bool) {
2333        self.verify_exact
2334            .store(on, std::sync::atomic::Ordering::Relaxed);
2335    }
2336    pub(crate) fn verify_exact_on(&self) -> bool {
2337        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2338    }
2339
2340    /// RAII scope over `verify_exact`: sets the flag to `on` now and restores the
2341    /// PREVIOUS value on drop — unwind, early `return`, and every `?` exit included.
2342    /// This is the required form for any scope an error can leave (see
2343    /// `set_verify_exact`); dropping the guard early (`drop(scope)`) ends the scope
2344    /// exactly where the manual `set_verify_exact(false)` used to sit.
2345    pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2346        ExactScope::set(&self.verify_exact, on)
2347    }
2348
2349    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
2350    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
2351    pub fn qkv_append_on() -> bool {
2352        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2353        *ON.get_or_init(|| {
2354            std::env::var("MEMRA_QKV_APPEND")
2355                .map(|v| v != "0")
2356                .unwrap_or(true)
2357        })
2358    }
2359
2360    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
2361    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
2362    pub fn pdl_wb_on() -> bool {
2363        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2364        *ON.get_or_init(|| {
2365            std::env::var("MEMRA_PDL_WB")
2366                .map(|v| v != "0")
2367                .unwrap_or(true)
2368        })
2369    }
2370
2371    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
2372    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
2373    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
2374    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
2375    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
2376    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
2377    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
2378    pub fn norm_ilp_on() -> bool {
2379        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2380        *ON.get_or_init(|| {
2381            std::env::var("MEMRA_NORM_ILP")
2382                .map(|v| v != "0")
2383                .unwrap_or(true)
2384        })
2385    }
2386
2387    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
2388    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
2389    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
2390    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
2391    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
2392    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
2393    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
2394    pub fn tk_ffn_dual_on() -> bool {
2395        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2396        *ON.get_or_init(|| {
2397            std::env::var("MEMRA_TK_FFN_DUAL")
2398                .map(|v| v != "0")
2399                .unwrap_or(true)
2400        })
2401    }
2402
2403    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
2404    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
2405    /// per-model no-harm bisect knob.
2406    pub fn pdl_mmvq_on() -> bool {
2407        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2408        *ON.get_or_init(|| {
2409            std::env::var("MEMRA_PDL_MMVQ")
2410                .map(|v| v != "0")
2411                .unwrap_or(true)
2412        })
2413    }
2414
2415    pub fn pdl_on() -> bool {
2416        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2417        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2418    }
2419
2420    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
2421    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
2422    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
2423    /// on the producer before any read), bit-identical by construction.
2424    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
2425    pub fn pdl_nvfp4q8_on() -> bool {
2426        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2427        *ON.get_or_init(|| {
2428            std::env::var("MEMRA_PDL_NVFP4")
2429                .map(|v| v != "0")
2430                .unwrap_or(true)
2431        })
2432    }
2433
2434    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
2435    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
2436    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
2437    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
2438    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
2439    fn q40_mr1_on() -> bool {
2440        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2441        match *Q40MR.get_or_init(|| {
2442            std::env::var("MEMRA_Q40_MR")
2443                .ok()
2444                .and_then(|v| v.parse().ok())
2445        }) {
2446            Some(v) => v == 1,
2447            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2448        }
2449    }
2450
2451    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
2452    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
2453    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
2454    /// writes wrong bytes silently.
2455    fn pdl_func_flash(
2456        &self,
2457        g: bool,
2458        name: &'static str,
2459    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2460        use cudarc::driver::sys as cu;
2461        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
2462        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
2463        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
2464        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
2465        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
2466        // this engine's CUcontext; single-context runs behave exactly as before.
2467        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2468            std::sync::Mutex::new(None);
2469        static FNS: std::sync::Mutex<
2470            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2471        > = std::sync::Mutex::new(None);
2472        let ctx_key = self.ctx().cu_ctx() as usize;
2473        if let Some(&f) = FNS
2474            .lock()
2475            .unwrap()
2476            .get_or_insert_with(Default::default)
2477            .get(&(ctx_key, g, name))
2478        {
2479            return Ok(f as cu::CUfunction);
2480        }
2481        let module = {
2482            let mut mods = MODS.lock().unwrap();
2483            let map = mods.get_or_insert_with(Default::default);
2484            match map.get(&(ctx_key, g)) {
2485                Some(&m) => m,
2486                None => {
2487                    let m = self.pdl_load_module_in_ctx(if g {
2488                        FLASH_FATBIN_KF8VF8
2489                    } else {
2490                        FLASH_FATBIN
2491                    })?;
2492                    map.insert((ctx_key, g), m);
2493                    m
2494                }
2495            }
2496        };
2497        let cname = std::ffi::CString::new(name)?;
2498        let mut f: cu::CUfunction = std::ptr::null_mut();
2499        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2500        if r != cu::CUresult::CUDA_SUCCESS {
2501            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2502        }
2503        FNS.lock()
2504            .unwrap()
2505            .get_or_insert_with(Default::default)
2506            .insert((ctx_key, g, name), f as usize);
2507        Ok(f)
2508    }
2509
2510    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2511    /// the module to the thread's CURRENT context — a remote-stage engine must not
2512    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2513    /// current context before returning.
2514    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2515        use cudarc::driver::sys as cu;
2516        let mut prev: cu::CUcontext = std::ptr::null_mut();
2517        unsafe {
2518            cu::cuCtxGetCurrent(&mut prev).result()?;
2519        }
2520        self.ctx().bind_to_thread()?;
2521        let mut m: cu::CUmodule = std::ptr::null_mut();
2522        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2523        let restore = if prev.is_null() {
2524            cu::CUresult::CUDA_SUCCESS
2525        } else {
2526            unsafe { cu::cuCtxSetCurrent(prev) }
2527        };
2528        if r != cu::CUresult::CUDA_SUCCESS {
2529            return Err(format!("pdl module load: {r:?}").into());
2530        }
2531        if restore != cu::CUresult::CUDA_SUCCESS {
2532            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2533        }
2534        Ok(m as usize)
2535    }
2536
2537    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
2538    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
2539    pub fn raw_kernel_function(
2540        &self,
2541        name: &'static str,
2542    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2543        self.pdl_func(name)
2544    }
2545
2546    fn pdl_func(
2547        &self,
2548        name: &'static str,
2549    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2550        use cudarc::driver::sys as cu;
2551        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2552        // are context-scoped; key everything by this engine's CUcontext).
2553        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2554            std::sync::Mutex::new(None);
2555        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2556        // duplicate module, loaded lazily on the first kernels-module miss.
2557        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2558            std::sync::Mutex::new(None);
2559        static FNS: std::sync::Mutex<
2560            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2561        > = std::sync::Mutex::new(None);
2562        let ctx_key = self.ctx().cu_ctx() as usize;
2563        if let Some(&f) = FNS
2564            .lock()
2565            .unwrap()
2566            .get_or_insert_with(Default::default)
2567            .get(&(ctx_key, name))
2568        {
2569            return Ok(f as cu::CUfunction);
2570        }
2571        let module = {
2572            let mut mods = MODULES.lock().unwrap();
2573            let map = mods.get_or_insert_with(Default::default);
2574            match map.get(&ctx_key) {
2575                Some(&m) => m,
2576                None => {
2577                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2578                    map.insert(ctx_key, m);
2579                    m
2580                }
2581            }
2582        };
2583        let cname = std::ffi::CString::new(name)?;
2584        let mut f: cu::CUfunction = std::ptr::null_mut();
2585        let mut r =
2586            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2587        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2588            let qmodule = {
2589                let mut mods = QMODULES.lock().unwrap();
2590                let map = mods.get_or_insert_with(Default::default);
2591                match map.get(&ctx_key) {
2592                    Some(&m) => m,
2593                    None => {
2594                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2595                        map.insert(ctx_key, m);
2596                        m
2597                    }
2598                }
2599            };
2600            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2601        }
2602        if r != cu::CUresult::CUDA_SUCCESS {
2603            return Err(format!("pdl_func {name}: {r:?}").into());
2604        }
2605        FNS.lock()
2606            .unwrap()
2607            .get_or_insert_with(Default::default)
2608            .insert((ctx_key, name), f as usize);
2609        Ok(f)
2610    }
2611
2612    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2613    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2614    ///
2615    /// # Safety
2616    /// `params` must match the kernel's exact parameter list (order, types, count) —
2617    /// a mismatch corrupts the launch silently.
2618    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2619    /// builder path's fa_func/func_g choice exactly).
2620    ///
2621    /// # Safety
2622    /// Same contract as `launch_pdl`.
2623    unsafe fn launch_pdl_flash(
2624        &self,
2625        g: bool,
2626        name: &'static str,
2627        grid: (u32, u32, u32),
2628        block: (u32, u32, u32),
2629        smem: u32,
2630        params: &mut [*mut std::ffi::c_void],
2631    ) -> Result<(), Box<dyn std::error::Error>> {
2632        use cudarc::driver::sys as cu;
2633        let f = self.pdl_func_flash(g, name)?;
2634        if smem > 0 {
2635            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2636            let r =
2637                unsafe {
2638                    cu::cuFuncSetAttribute(f,
2639                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2640                smem as i32)
2641                };
2642            if r != cu::CUresult::CUDA_SUCCESS {
2643                return Err(format!("pdl smem attr {name}: {r:?}").into());
2644            }
2645        }
2646        let mut attr = cu::CUlaunchAttribute {
2647            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2648            pad: [0; 4],
2649            value: cu::CUlaunchAttributeValue {
2650                programmaticStreamSerializationAllowed: 1,
2651            },
2652        };
2653        let cfg = cu::CUlaunchConfig {
2654            gridDimX: grid.0,
2655            gridDimY: grid.1,
2656            gridDimZ: grid.2,
2657            blockDimX: block.0,
2658            blockDimY: block.1,
2659            blockDimZ: block.2,
2660            sharedMemBytes: smem,
2661            hStream: self.gpu.stream().cu_stream(),
2662            attrs: &mut attr,
2663            numAttrs: 1,
2664        };
2665        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2666        if r != cu::CUresult::CUDA_SUCCESS {
2667            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2668        }
2669        Ok(())
2670    }
2671
2672    unsafe fn launch_pdl(
2673        &self,
2674        name: &'static str,
2675        grid: (u32, u32, u32),
2676        block: (u32, u32, u32),
2677        params: &mut [*mut std::ffi::c_void],
2678    ) -> Result<(), Box<dyn std::error::Error>> {
2679        use cudarc::driver::sys as cu;
2680        let f = self.pdl_func(name)?;
2681        let mut attr = cu::CUlaunchAttribute {
2682            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2683            pad: [0; 4],
2684            value: cu::CUlaunchAttributeValue {
2685                programmaticStreamSerializationAllowed: 1,
2686            },
2687        };
2688        let cfg = cu::CUlaunchConfig {
2689            gridDimX: grid.0,
2690            gridDimY: grid.1,
2691            gridDimZ: grid.2,
2692            blockDimX: block.0,
2693            blockDimY: block.1,
2694            blockDimZ: block.2,
2695            sharedMemBytes: 0,
2696            hStream: self.gpu.stream().cu_stream(),
2697            attrs: &mut attr,
2698            numAttrs: 1,
2699        };
2700        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2701        if r != cu::CUresult::CUDA_SUCCESS {
2702            return Err(format!("launch_pdl {name}: {r:?}").into());
2703        }
2704        Ok(())
2705    }
2706
2707    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2708    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2709    pub fn prefetch_weight_l2(
2710        &self,
2711        w: &crate::model::GpuTensor,
2712    ) -> Result<(), Box<dyn std::error::Error>> {
2713        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2714            let p = rp4.as_ref().unwrap_or(bytes);
2715            self.prefetch_l2(p, p.len())?;
2716        }
2717        Ok(())
2718    }
2719
2720    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2721    /// by the DEVICE token id at tok[idx] into f32.
2722    pub fn gather_row_bf16(
2723        &self,
2724        table: &CudaSlice<u8>,
2725        tok: &CudaSlice<u32>,
2726        idx: usize,
2727        dst: &mut CudaSlice<f32>,
2728        ncols: usize,
2729    ) -> Result<(), Box<dyn std::error::Error>> {
2730        let f = self.func("gather_row_bf16_f32");
2731        let cfg = LaunchConfig {
2732            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2733            block_dim: (256, 1, 1),
2734            shared_mem_bytes: 0,
2735        };
2736        let (nc, ix) = (ncols as i32, idx as i32);
2737        let __s_b = self.gpu.stream();
2738        let mut b = __s_b.launch_builder(&f);
2739        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2740        unsafe {
2741            b.launch(cfg)?;
2742        }
2743        Ok(())
2744    }
2745
2746    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2747    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2748    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2749    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2750    /// finish(1).
2751    #[allow(clippy::too_many_arguments)]
2752    pub fn dflash2_dynconv(
2753        &self,
2754        x: &CudaSlice<f32>,
2755        dyn_: &CudaSlice<f32>,
2756        base: &CudaSlice<f32>,
2757        out: &mut CudaSlice<f32>,
2758        rows: usize,
2759        hidden: usize,
2760        group_size: usize,
2761        ksize: usize,
2762        half: usize,
2763    ) -> Result<(), Box<dyn std::error::Error>> {
2764        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2765        let f = self.func("dflash2_dynconv_f32");
2766        let n = rows * hidden;
2767        let cfg = LaunchConfig {
2768            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2769            block_dim: (256, 1, 1),
2770            shared_mem_bytes: 0,
2771        };
2772        let (ri, hi, gi, ki, hf) = (
2773            rows as i32,
2774            hidden as i32,
2775            group_size as i32,
2776            ksize as i32,
2777            half as i32,
2778        );
2779        let __s_b = self.gpu.stream();
2780        let mut b = __s_b.launch_builder(&f);
2781        b.arg(x)
2782            .arg(dyn_)
2783            .arg(base)
2784            .arg(out)
2785            .arg(&ri)
2786            .arg(&hi)
2787            .arg(&gi)
2788            .arg(&ki)
2789            .arg(&hf);
2790        unsafe {
2791            b.launch(cfg)?;
2792        }
2793        Ok(())
2794    }
2795
2796    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2797    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2798    /// value-descending, ties to the lower index.
2799    pub fn topk_rows(
2800        &self,
2801        logits: &CudaSlice<f32>,
2802        n_rows: usize,
2803        n_cols: usize,
2804        k: usize,
2805    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2806        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2807        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2808        let f = self.func("topk_rows_f32");
2809        let nth = 256usize;
2810        let mut vals = self.uninit(n_rows * k)?;
2811        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2812        let cfg = LaunchConfig {
2813            grid_dim: (n_rows as u32, 1, 1),
2814            block_dim: (nth as u32, 1, 1),
2815            shared_mem_bytes: (nth * k * 8) as u32,
2816        };
2817        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2818        let __s_b = self.gpu.stream();
2819        let mut b = __s_b.launch_builder(&f);
2820        b.arg(logits)
2821            .arg(&nr)
2822            .arg(&nc)
2823            .arg(&ki)
2824            .arg(&mut vals)
2825            .arg(&mut idxs);
2826        unsafe {
2827            b.launch(cfg)?;
2828        }
2829        Ok((vals, idxs))
2830    }
2831
2832    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2833    pub fn add_row_inplace(
2834        &self,
2835        logits: &mut CudaSlice<f32>,
2836        bias: &CudaSlice<f32>,
2837        n: usize,
2838        row_off: usize,
2839    ) -> Result<(), Box<dyn std::error::Error>> {
2840        let f = self.func("add_row_inplace_f32");
2841        let cfg = LaunchConfig {
2842            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2843            block_dim: (256, 1, 1),
2844            shared_mem_bytes: 0,
2845        };
2846        let (ni, off) = (n as i32, row_off as i64);
2847        let __s_b = self.gpu.stream();
2848        let mut b = __s_b.launch_builder(&f);
2849        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2850        unsafe {
2851            b.launch(cfg)?;
2852        }
2853        Ok(())
2854    }
2855
2856    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2857    pub fn prefetch_l2(
2858        &self,
2859        p: &CudaSlice<u8>,
2860        n: usize,
2861    ) -> Result<(), Box<dyn std::error::Error>> {
2862        let f = self.func("prefetch_l2_bytes");
2863        let lines = n.div_ceil(128);
2864        let ni = n as i64;
2865        let cfg = LaunchConfig {
2866            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2867            block_dim: (256, 1, 1),
2868            shared_mem_bytes: 0,
2869        };
2870        let __s_b = self.gpu.stream();
2871        let mut b = __s_b.launch_builder(&f);
2872        b.arg(p).arg(&ni);
2873        unsafe {
2874            b.launch(cfg)?;
2875        }
2876        Ok(())
2877    }
2878
2879    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2880    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2881    pub fn router_gemv(
2882        &self,
2883        w: &CudaSlice<f32>,
2884        x: &CudaSlice<f32>,
2885        n_embd: usize,
2886        n_experts: usize,
2887        t: usize,
2888    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2889        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2890        // stream differs) — too small to justify a numeric config change; deleted.
2891        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2892        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2893        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2894        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2895            Ok("0") => false,
2896            Ok(_) => true,
2897            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2898        };
2899        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2900        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2901        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2902        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2903        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2904        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2905        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2906        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2907        // (perf-only, bits equal).
2908        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2909        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2910    }
2911
2912    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2913    /// force both forms; `batch` requires `w8`).
2914    pub fn router_gemv_form(
2915        &self,
2916        w: &CudaSlice<f32>,
2917        x: &CudaSlice<f32>,
2918        n_embd: usize,
2919        n_experts: usize,
2920        t: usize,
2921        w8: bool,
2922        batch: bool,
2923    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2924        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2925        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2926        let f = if batch {
2927            self.func("router_gemv_f32_w8_batch")
2928        } else if w8 {
2929            self.func("router_gemv_f32_w8")
2930        } else {
2931            self.func("router_gemv_f32")
2932        };
2933        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2934        let cfg = if batch {
2935            LaunchConfig {
2936                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2937                block_dim: (32, 8, 1),
2938                shared_mem_bytes: 0,
2939            }
2940        } else {
2941            LaunchConfig {
2942                grid_dim: (n_experts as u32, t as u32, 1),
2943                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2944                shared_mem_bytes: 0,
2945            }
2946        };
2947        let __s_b = self.gpu.stream();
2948        let mut b = __s_b.launch_builder(&f);
2949        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2950        unsafe {
2951            b.launch(cfg)?;
2952        }
2953        Ok(y)
2954    }
2955
2956    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
2957    /// buffer — token-graph alloc-free.
2958    pub fn router_gemv_into(
2959        &self,
2960        w: &CudaSlice<f32>,
2961        x: &CudaSlice<f32>,
2962        y: &mut CudaSlice<f32>,
2963        n_embd: usize,
2964        n_experts: usize,
2965        t: usize,
2966    ) -> Result<(), Box<dyn std::error::Error>> {
2967        if y.len() < t * n_experts {
2968            return Err("router_gemv_into output too small".into());
2969        }
2970        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2971            Ok("0") => false,
2972            Ok(_) => true,
2973            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2974        };
2975        let f = if w8 {
2976            self.func("router_gemv_f32_w8")
2977        } else {
2978            self.func("router_gemv_f32")
2979        };
2980        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2981        let cfg = LaunchConfig {
2982            grid_dim: (n_experts as u32, t as u32, 1),
2983            block_dim: (32, if w8 { 8 } else { 1 }, 1),
2984            shared_mem_bytes: 0,
2985        };
2986        let __s_b = self.gpu.stream();
2987        let mut b = __s_b.launch_builder(&f);
2988        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2989        unsafe {
2990            b.launch(cfg)?;
2991        }
2992        Ok(())
2993    }
2994
2995    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2996    pub fn rows_permute(
2997        &self,
2998        src: &CudaSlice<f32>,
2999        idx: &CudaSlice<i32>,
3000        nrows: usize,
3001        ncols: usize,
3002    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3003        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3004        let f = self.func("rows_permute_f32");
3005        let (nc, nr) = (ncols as i32, nrows as i32);
3006        let cfg = LaunchConfig {
3007            grid_dim: (nrows as u32, 1, 1),
3008            block_dim: (256, 1, 1),
3009            shared_mem_bytes: 0,
3010        };
3011        let __s_b = self.gpu.stream();
3012        let mut b = __s_b.launch_builder(&f);
3013        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3014        unsafe {
3015            b.launch(cfg)?;
3016        }
3017        Ok(dst)
3018    }
3019
3020    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
3021    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
3022    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
3023    /// decode chain and the small-t spec-verify chain match per row by construction.
3024    pub fn sigmoid_dot_rows(
3025        &self,
3026        x: &CudaSlice<f32>,
3027        w: &CudaSlice<f32>,
3028        n_embd: usize,
3029        t: usize,
3030    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3031        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
3032        // config; same class as MEMRA_ROUTER_V2).
3033        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3034        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3035            let gs = self.linear(x, w, t, n_embd, 1)?;
3036            let mut g = self.uninit(t)?;
3037            self.sigmoid(&gs, &mut g, t)?;
3038            return Ok(g);
3039        }
3040        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
3041        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
3042        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
3043        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
3044        // flags doctrine; this per-token form serves every t.
3045        let mut g = self.alloc_uninit::<f32>(t)?;
3046        let f = self.func("sigmoid_dot_rows_f32");
3047        let (ne, ti) = (n_embd as i32, t as i32);
3048        let cfg = LaunchConfig {
3049            grid_dim: (t as u32, 1, 1),
3050            block_dim: (32, 8, 1),
3051            shared_mem_bytes: 0,
3052        };
3053        let __s_b = self.gpu.stream();
3054        let mut b = __s_b.launch_builder(&f);
3055        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3056        unsafe {
3057            b.launch(cfg)?;
3058        }
3059        Ok(g)
3060    }
3061
3062    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
3063    pub fn sigmoid_dot_rows_into(
3064        &self,
3065        x: &CudaSlice<f32>,
3066        w: &CudaSlice<f32>,
3067        g: &mut CudaSlice<f32>,
3068        n_embd: usize,
3069        t: usize,
3070    ) -> Result<(), Box<dyn std::error::Error>> {
3071        if g.len() < t {
3072            return Err("sigmoid_dot_rows_into output too small".into());
3073        }
3074        let f = self.func("sigmoid_dot_rows_f32");
3075        let (ne, ti) = (n_embd as i32, t as i32);
3076        let cfg = LaunchConfig {
3077            grid_dim: (t as u32, 1, 1),
3078            block_dim: (32, 8, 1),
3079            shared_mem_bytes: 0,
3080        };
3081        let __s_b = self.gpu.stream();
3082        let mut b = __s_b.launch_builder(&f);
3083        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3084        unsafe {
3085            b.launch(cfg)?;
3086        }
3087        Ok(())
3088    }
3089
3090    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
3091    pub fn spec_rollback_stream(
3092        &self,
3093        len_ptrs: &CudaSlice<u64>,
3094        pos_start: &CudaSlice<i32>,
3095        acc: &CudaSlice<u32>,
3096        base: usize,
3097        n_rows: usize,
3098    ) -> Result<(), Box<dyn std::error::Error>> {
3099        let f = self.func("spec_rollback_stream");
3100        let (b, nr) = (base as i32, n_rows as i32);
3101        let cfg = LaunchConfig {
3102            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3103            block_dim: (64, 1, 1),
3104            shared_mem_bytes: 0,
3105        };
3106        let __s_bl = self.gpu.stream();
3107        let mut bl = __s_bl.launch_builder(&f);
3108        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3109        unsafe {
3110            bl.launch(cfg)?;
3111        }
3112        Ok(())
3113    }
3114
3115    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
3116    pub fn plain_tok_ring(
3117        &self,
3118        vam: &CudaSlice<u32>,
3119        pos_start: &CudaSlice<i32>,
3120        base: usize,
3121        ring: &mut CudaSlice<u32>,
3122    ) -> Result<(), Box<dyn std::error::Error>> {
3123        let f = self.func("plain_tok_ring");
3124        let (b, cap) = (base as i32, ring.len() as i32);
3125        let cfg = LaunchConfig {
3126            grid_dim: (1, 1, 1),
3127            block_dim: (32, 1, 1),
3128            shared_mem_bytes: 0,
3129        };
3130        let __s_bl = self.gpu.stream();
3131        let mut bl = __s_bl.launch_builder(&f);
3132        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3133        unsafe {
3134            bl.launch(cfg)?;
3135        }
3136        Ok(())
3137    }
3138
3139    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
3140    pub fn spec_ring_commit(
3141        &self,
3142        vtok: &CudaSlice<u32>,
3143        acc: &CudaSlice<u32>,
3144        brk: &CudaSlice<u32>,
3145        ring: &mut CudaSlice<u32>,
3146        pend: &mut CudaSlice<u32>,
3147    ) -> Result<(), Box<dyn std::error::Error>> {
3148        let f = self.func("spec_ring_commit");
3149        let cfg = LaunchConfig {
3150            grid_dim: (1, 1, 1),
3151            block_dim: (32, 1, 1),
3152            shared_mem_bytes: 0,
3153        };
3154        let __s_b = self.gpu.stream();
3155        let mut b = __s_b.launch_builder(&f);
3156        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3157        unsafe {
3158            b.launch(cfg)?;
3159        }
3160        Ok(())
3161    }
3162    pub fn i32_copy_add(
3163        &self,
3164        src: &CudaSlice<i32>,
3165        dst: &mut CudaSlice<i32>,
3166        delta: i32,
3167    ) -> Result<(), Box<dyn std::error::Error>> {
3168        let f = self.func("i32_copy_add");
3169        let cfg = LaunchConfig {
3170            grid_dim: (1, 1, 1),
3171            block_dim: (32, 1, 1),
3172            shared_mem_bytes: 0,
3173        };
3174        let __s_b = self.gpu.stream();
3175        let mut b = __s_b.launch_builder(&f);
3176        b.arg(src).arg(dst).arg(&delta);
3177        unsafe {
3178            b.launch(cfg)?;
3179        }
3180        Ok(())
3181    }
3182    pub fn u32_copy(
3183        &self,
3184        src: &CudaSlice<u32>,
3185        dst: &mut CudaSlice<u32>,
3186    ) -> Result<(), Box<dyn std::error::Error>> {
3187        let f = self.func("u32_copy");
3188        let cfg = LaunchConfig {
3189            grid_dim: (1, 1, 1),
3190            block_dim: (32, 1, 1),
3191            shared_mem_bytes: 0,
3192        };
3193        let __s_b = self.gpu.stream();
3194        let mut b = __s_b.launch_builder(&f);
3195        b.arg(src).arg(dst);
3196        unsafe {
3197            b.launch(cfg)?;
3198        }
3199        Ok(())
3200    }
3201
3202    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
3203    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
3204    /// caps acceptance exactly like drafting fewer tokens).
3205    pub fn spec_adapt_k(
3206        &self,
3207        acc: &CudaSlice<u32>,
3208        brk: &mut CudaSlice<u32>,
3209        floor: usize,
3210        cap: usize,
3211    ) -> Result<(), Box<dyn std::error::Error>> {
3212        let f = self.func("spec_adapt_k");
3213        let (fl, cp) = (floor as i32, cap as i32);
3214        let cfg = LaunchConfig {
3215            grid_dim: (1, 1, 1),
3216            block_dim: (32, 1, 1),
3217            shared_mem_bytes: 0,
3218        };
3219        let __s_b = self.gpu.stream();
3220        let mut b = __s_b.launch_builder(&f);
3221        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3222        unsafe {
3223            b.launch(cfg)?;
3224        }
3225        Ok(())
3226    }
3227
3228    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
3229    pub fn spec_accept_greedy_dc(
3230        &self,
3231        preds: &CudaSlice<u32>,
3232        vtok: &CudaSlice<u32>,
3233        last_pred: &CudaSlice<u32>,
3234        brk: &CudaSlice<u32>,
3235        out: &mut CudaSlice<u32>,
3236    ) -> Result<(), Box<dyn std::error::Error>> {
3237        let f = self.func("spec_accept_greedy_dc");
3238        let cfg = LaunchConfig {
3239            grid_dim: (1, 1, 1),
3240            block_dim: (32, 1, 1),
3241            shared_mem_bytes: 0,
3242        };
3243        let __s_b = self.gpu.stream();
3244        let mut b = __s_b.launch_builder(&f);
3245        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3246        unsafe {
3247            b.launch(cfg)?;
3248        }
3249        Ok(())
3250    }
3251
3252    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
3253    pub fn pos_iota(
3254        &self,
3255        pos0: &CudaSlice<i32>,
3256        out: &mut CudaSlice<i32>,
3257        t: usize,
3258    ) -> Result<(), Box<dyn std::error::Error>> {
3259        let f = self.func("pos_iota_i32");
3260        let ti = t as i32;
3261        let cfg = LaunchConfig {
3262            grid_dim: (1, 1, 1),
3263            block_dim: (t.max(1) as u32, 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(pos0).arg(out).arg(&ti);
3269        unsafe {
3270            b.launch(cfg)?;
3271        }
3272        Ok(())
3273    }
3274    #[allow(clippy::too_many_arguments)]
3275    pub fn append_kv_quantized_rows_dc(
3276        &self,
3277        k_rows: &CudaSlice<f32>,
3278        v_rows: &CudaSlice<f32>,
3279        kc: &mut CudaSlice<u8>,
3280        vc: &mut CudaSlice<u8>,
3281        t0_dev: &CudaSlice<i32>,
3282        t: usize,
3283        kv_dim_k: usize,
3284        kv_dim_v: usize,
3285        k_tok_bytes: usize,
3286        v_tok_bytes: usize,
3287        g: bool,
3288    ) -> Result<(), Box<dyn std::error::Error>> {
3289        let f = if g {
3290            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3291        } else {
3292            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3293        };
3294        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3295        let cfg = LaunchConfig {
3296            grid_dim: (nblk, t as u32, 1),
3297            block_dim: (32, 1, 1),
3298            shared_mem_bytes: 0,
3299        };
3300        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3301        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3302        let __s_b = self.gpu.stream();
3303        let mut b = __s_b.launch_builder(&f);
3304        b.arg(k_rows)
3305            .arg(v_rows)
3306            .arg(kc)
3307            .arg(vc)
3308            .arg(t0_dev)
3309            .arg(&kdk)
3310            .arg(&kdv)
3311            .arg(&ktb)
3312            .arg(&vtb);
3313        unsafe {
3314            b.launch(cfg)?;
3315        }
3316        Ok(())
3317    }
3318
3319    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
3320    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
3321    #[allow(clippy::too_many_arguments)]
3322    pub fn append_kv_quantized_row_dc_inc(
3323        &self,
3324        k_row: &CudaSlice<f32>,
3325        v_row: &CudaSlice<f32>,
3326        kc: &mut CudaSlice<u8>,
3327        vc: &mut CudaSlice<u8>,
3328        t0_dev: &mut CudaSlice<i32>,
3329        kv_dim_k: usize,
3330        kv_dim_v: usize,
3331        k_tok_bytes: usize,
3332        v_tok_bytes: usize,
3333        g: bool,
3334    ) -> Result<(), Box<dyn std::error::Error>> {
3335        let f = if g {
3336            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3337        } else {
3338            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3339        };
3340        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3341        let cfg = LaunchConfig {
3342            grid_dim: (1, 1, 1),
3343            block_dim: (nthreads, 1, 1),
3344            shared_mem_bytes: 0,
3345        };
3346        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3347        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3348        let __s_b = self.gpu.stream();
3349        let mut b = __s_b.launch_builder(&f);
3350        b.arg(k_row)
3351            .arg(v_row)
3352            .arg(kc)
3353            .arg(vc)
3354            .arg(t0_dev)
3355            .arg(&kdk)
3356            .arg(&kdv)
3357            .arg(&ktb)
3358            .arg(&vtb);
3359        unsafe {
3360            b.launch(cfg)?;
3361        }
3362        Ok(())
3363    }
3364
3365    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
3366    pub fn pack_tok_p(
3367        &self,
3368        tok: &CudaSlice<u32>,
3369        p: &CudaSlice<f32>,
3370        out: &mut CudaSlice<u32>,
3371        slot: usize,
3372    ) -> Result<(), Box<dyn std::error::Error>> {
3373        let f = self.func("pack_tok_p");
3374        let sl = slot as i32;
3375        let cfg = LaunchConfig {
3376            grid_dim: (1, 1, 1),
3377            block_dim: (32, 1, 1),
3378            shared_mem_bytes: 0,
3379        };
3380        let __s_b = self.gpu.stream();
3381        let mut b = __s_b.launch_builder(&f);
3382        b.arg(tok).arg(p).arg(out).arg(&sl);
3383        unsafe {
3384            b.launch(cfg)?;
3385        }
3386        Ok(())
3387    }
3388    pub fn tok_map_u32(
3389        &self,
3390        tok: &mut CudaSlice<u32>,
3391        map: &CudaSlice<u32>,
3392    ) -> Result<(), Box<dyn std::error::Error>> {
3393        let f = self.func("tok_map_u32");
3394        let cfg = LaunchConfig {
3395            grid_dim: (1, 1, 1),
3396            block_dim: (32, 1, 1),
3397            shared_mem_bytes: 0,
3398        };
3399        let __s_b = self.gpu.stream();
3400        let mut b = __s_b.launch_builder(&f);
3401        b.arg(tok).arg(map);
3402        unsafe {
3403            b.launch(cfg)?;
3404        }
3405        Ok(())
3406    }
3407
3408    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
3409    #[allow(clippy::too_many_arguments)]
3410    pub fn spec_assemble_verify(
3411        &self,
3412        tokp: &CudaSlice<u32>,
3413        pend: &CudaSlice<u32>,
3414        d2t: Option<&CudaSlice<u32>>,
3415        vtok: &mut CudaSlice<u32>,
3416        brk: &mut CudaSlice<u32>,
3417        p_min: f32,
3418        k: usize,
3419        pmin0: bool,
3420    ) -> Result<(), Box<dyn std::error::Error>> {
3421        let f = self.func("spec_assemble_verify");
3422        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3423        let cfg = LaunchConfig {
3424            grid_dim: (1, 1, 1),
3425            block_dim: (32, 1, 1),
3426            shared_mem_bytes: 0,
3427        };
3428        let __s_b = self.gpu.stream();
3429        let mut b = __s_b.launch_builder(&f);
3430        match d2t {
3431            Some(m) => {
3432                b.arg(tokp)
3433                    .arg(pend)
3434                    .arg(m)
3435                    .arg(vtok)
3436                    .arg(brk)
3437                    .arg(&p_min)
3438                    .arg(&ki)
3439                    .arg(&pm);
3440                unsafe {
3441                    b.launch(cfg)?;
3442                }
3443            }
3444            None => {
3445                let null: u64 = 0;
3446                b.arg(tokp)
3447                    .arg(pend)
3448                    .arg(&null)
3449                    .arg(vtok)
3450                    .arg(brk)
3451                    .arg(&p_min)
3452                    .arg(&ki)
3453                    .arg(&pm);
3454                unsafe {
3455                    b.launch(cfg)?;
3456                }
3457            }
3458        }
3459        Ok(())
3460    }
3461
3462    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
3463    #[allow(clippy::too_many_arguments)]
3464    pub fn ssm_conv_ring_rebuild_dc(
3465        &self,
3466        qkv_tm: &CudaSlice<f32>,
3467        ring_old: &CudaSlice<f32>,
3468        conv_state: &mut CudaSlice<f32>,
3469        conv_dim: usize,
3470        acc: &CudaSlice<u32>,
3471        base: usize,
3472        t_v: usize,
3473        d_conv: usize,
3474    ) -> Result<(), Box<dyn std::error::Error>> {
3475        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3476        let n = conv_dim * (d_conv - 1);
3477        let cfg = LaunchConfig::for_num_elems(n as u32);
3478        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3479        let __s_b = self.gpu.stream();
3480        let mut b = __s_b.launch_builder(&f);
3481        b.arg(qkv_tm)
3482            .arg(ring_old)
3483            .arg(conv_state)
3484            .arg(&cd)
3485            .arg(acc)
3486            .arg(&b0)
3487            .arg(&tv)
3488            .arg(&dc);
3489        unsafe {
3490            b.launch(cfg)?;
3491        }
3492        Ok(())
3493    }
3494    #[allow(clippy::too_many_arguments)]
3495    pub fn gdn_scan_s128_dc(
3496        &self,
3497        q: &CudaSlice<f32>,
3498        k: &CudaSlice<f32>,
3499        v: &CudaSlice<f32>,
3500        g: &CudaSlice<f32>,
3501        beta: &CudaSlice<f32>,
3502        state_in: &CudaSlice<f32>,
3503        state_out: &mut CudaSlice<f32>,
3504        o: &mut CudaSlice<f32>,
3505        n_head: usize,
3506        acc: &CudaSlice<u32>,
3507        base: usize,
3508        t_v: usize,
3509        scale: f32,
3510    ) -> Result<(), Box<dyn std::error::Error>> {
3511        let f = self.func("gdn_scan_s128_dc");
3512        const S_V: u32 = 128;
3513        const WARP: u32 = 32;
3514        const COLS_PER_BLOCK: u32 = 4;
3515        let cfg = LaunchConfig {
3516            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3517            block_dim: (WARP, COLS_PER_BLOCK, 1),
3518            shared_mem_bytes: 0,
3519        };
3520        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3521        let __s_b = self.gpu.stream();
3522        let mut b = __s_b.launch_builder(&f);
3523        b.arg(q)
3524            .arg(k)
3525            .arg(v)
3526            .arg(g)
3527            .arg(beta)
3528            .arg(state_in)
3529            .arg(state_out)
3530            .arg(o)
3531            .arg(&h)
3532            .arg(acc)
3533            .arg(&b0)
3534            .arg(&tv)
3535            .arg(&scale);
3536        unsafe {
3537            b.launch(cfg)?;
3538        }
3539        Ok(())
3540    }
3541
3542    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
3543    pub fn spec_rollback_kv(
3544        &self,
3545        len_ptrs: &CudaSlice<u64>,
3546        saved: &CudaSlice<i32>,
3547        acc: &CudaSlice<u32>,
3548        base: usize,
3549        n_layer: usize,
3550    ) -> Result<(), Box<dyn std::error::Error>> {
3551        let f = self.func("spec_rollback_kv");
3552        let (b, nl) = (base as i32, n_layer as i32);
3553        let cfg = LaunchConfig {
3554            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3555            block_dim: (64, 1, 1),
3556            shared_mem_bytes: 0,
3557        };
3558        let __s_bl = self.gpu.stream();
3559        let mut bl = __s_bl.launch_builder(&f);
3560        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3561        unsafe {
3562            bl.launch(cfg)?;
3563        }
3564        Ok(())
3565    }
3566
3567    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3568    pub fn spec_fork_valid(
3569        &self,
3570        acc: &CudaSlice<u32>,
3571        optimistic_pending: u32,
3572        valid: &mut CudaSlice<u32>,
3573    ) -> Result<(), Box<dyn std::error::Error>> {
3574        let f = self.func("spec_fork_valid");
3575        let cfg = LaunchConfig {
3576            grid_dim: (1, 1, 1),
3577            block_dim: (1, 1, 1),
3578            shared_mem_bytes: 0,
3579        };
3580        let __s_bl = self.gpu.stream();
3581        let mut bl = __s_bl.launch_builder(&f);
3582        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3583        unsafe {
3584            bl.launch(cfg)?;
3585        }
3586        Ok(())
3587    }
3588
3589    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3590    pub fn spec_fork_reconcile_kv(
3591        &self,
3592        len_ptrs: &CudaSlice<u64>,
3593        saved: &CudaSlice<i32>,
3594        acc: &CudaSlice<u32>,
3595        valid: &CudaSlice<u32>,
3596        base: usize,
3597        n_layer: usize,
3598    ) -> Result<(), Box<dyn std::error::Error>> {
3599        let f = self.func("spec_fork_reconcile_kv");
3600        let (b, nl) = (base as i32, n_layer as i32);
3601        let cfg = LaunchConfig {
3602            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3603            block_dim: (64, 1, 1),
3604            shared_mem_bytes: 0,
3605        };
3606        let __s_bl = self.gpu.stream();
3607        let mut bl = __s_bl.launch_builder(&f);
3608        bl.arg(len_ptrs)
3609            .arg(saved)
3610            .arg(acc)
3611            .arg(valid)
3612            .arg(&b)
3613            .arg(&nl);
3614        unsafe {
3615            bl.launch(cfg)?;
3616        }
3617        Ok(())
3618    }
3619
3620    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3621    pub fn spec_fork_restore_f32(
3622        &self,
3623        snapshot: &CudaSlice<f32>,
3624        state: &mut CudaSlice<f32>,
3625        valid: &CudaSlice<u32>,
3626    ) -> Result<(), Box<dyn std::error::Error>> {
3627        assert_eq!(
3628            snapshot.len(),
3629            state.len(),
3630            "fork recurrent snapshot shape mismatch"
3631        );
3632        let f = self.func("spec_fork_restore_f32");
3633        let n = state.len() as i32;
3634        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3635        let cfg = LaunchConfig {
3636            grid_dim: (blocks, 1, 1),
3637            block_dim: (256, 1, 1),
3638            shared_mem_bytes: 0,
3639        };
3640        let __s_bl = self.gpu.stream();
3641        let mut bl = __s_bl.launch_builder(&f);
3642        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3643        unsafe {
3644            bl.launch(cfg)?;
3645        }
3646        Ok(())
3647    }
3648
3649    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3650    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3651    pub fn spec_seed_gather(
3652        &self,
3653        vx: &CudaSlice<f32>,
3654        fill_prev: &CudaSlice<f32>,
3655        acc: &CudaSlice<u32>,
3656        h_seed: &mut CudaSlice<f32>,
3657        base: usize,
3658        n_embd: usize,
3659    ) -> Result<(), Box<dyn std::error::Error>> {
3660        let f = self.func("spec_seed_gather");
3661        let (b, ne) = (base as i32, n_embd as i32);
3662        let cfg = LaunchConfig {
3663            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3664            block_dim: (256, 1, 1),
3665            shared_mem_bytes: 0,
3666        };
3667        let __s_bl = self.gpu.stream();
3668        let mut bl = __s_bl.launch_builder(&f);
3669        bl.arg(vx)
3670            .arg(fill_prev)
3671            .arg(acc)
3672            .arg(h_seed)
3673            .arg(&b)
3674            .arg(&ne);
3675        unsafe {
3676            bl.launch(cfg)?;
3677        }
3678        Ok(())
3679    }
3680
3681    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3682    pub fn spec_accept_greedy(
3683        &self,
3684        preds: &CudaSlice<u32>,
3685        draft: &CudaSlice<u32>,
3686        last_pred: u32,
3687        base: usize,
3688        k_round: usize,
3689        out: &mut CudaSlice<u32>,
3690    ) -> Result<(), Box<dyn std::error::Error>> {
3691        let f = self.func("spec_accept_greedy");
3692        let (b, k) = (base as i32, k_round as i32);
3693        let cfg = LaunchConfig {
3694            grid_dim: (1, 1, 1),
3695            block_dim: (32, 1, 1),
3696            shared_mem_bytes: 0,
3697        };
3698        let __s_bl = self.gpu.stream();
3699        let mut bl = __s_bl.launch_builder(&f);
3700        bl.arg(preds)
3701            .arg(draft)
3702            .arg(&last_pred)
3703            .arg(&b)
3704            .arg(&k)
3705            .arg(out);
3706        unsafe {
3707            bl.launch(cfg)?;
3708        }
3709        Ok(())
3710    }
3711
3712    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3713    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3714    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3715
3716    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3717    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3718    pub fn gumbel_perturb(
3719        &self,
3720        x: &CudaSlice<f32>,
3721        y: &mut CudaSlice<f32>,
3722        n: usize,
3723        seed: u64,
3724        stream_pos: u32,
3725        temp: f32,
3726    ) -> Result<(), Box<dyn std::error::Error>> {
3727        let f = self.func("gumbel_perturb_f32");
3728        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3729        let cfg = LaunchConfig {
3730            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3731            block_dim: (256, 1, 1),
3732            shared_mem_bytes: 0,
3733        };
3734        let __s_b = self.gpu.stream();
3735        let mut b = __s_b.launch_builder(&f);
3736        b.arg(x)
3737            .arg(&mut *y)
3738            .arg(&ni)
3739            .arg(&slo)
3740            .arg(&shi)
3741            .arg(&stream_pos)
3742            .arg(&temp);
3743        unsafe {
3744            b.launch(cfg)?;
3745        }
3746        Ok(())
3747    }
3748
3749    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3750    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3751    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3752    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3753    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3754    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3755    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3756    pub fn mask_logits_col(
3757        &self,
3758        logits: &mut CudaSlice<f32>,
3759        mask: &CudaSlice<u32>,
3760        col: usize,
3761        n: usize,
3762        mask_words: usize,
3763    ) -> Result<(), Box<dyn std::error::Error>> {
3764        let f = self.func("mask_logits_f32");
3765        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3766        let cfg = LaunchConfig {
3767            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3768            block_dim: (256, 1, 1),
3769            shared_mem_bytes: 0,
3770        };
3771        let __s_b = self.gpu.stream();
3772        let mut b = __s_b.launch_builder(&f);
3773        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3774        unsafe {
3775            b.launch(cfg)?;
3776        }
3777        Ok(())
3778    }
3779
3780    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3781    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3782    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3783    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3784    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3785    /// pointer-invariance IS the serving isolation contract for sampled rows.
3786    pub fn gumbel_perturb_col(
3787        &self,
3788        x: &CudaSlice<f32>,
3789        col: usize,
3790        y: &mut CudaSlice<f32>,
3791        n: usize,
3792        seed: u64,
3793        stream_pos: u32,
3794        temp: f32,
3795    ) -> Result<(), Box<dyn std::error::Error>> {
3796        let f = self.func("gumbel_perturb_f32");
3797        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3798        let col_view = x.slice(col * n..(col + 1) * n);
3799        let cfg = LaunchConfig {
3800            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3801            block_dim: (256, 1, 1),
3802            shared_mem_bytes: 0,
3803        };
3804        let __s_b = self.gpu.stream();
3805        let mut b = __s_b.launch_builder(&f);
3806        b.arg(&col_view)
3807            .arg(&mut *y)
3808            .arg(&ni)
3809            .arg(&slo)
3810            .arg(&shi)
3811            .arg(&stream_pos)
3812            .arg(&temp);
3813        unsafe {
3814            b.launch(cfg)?;
3815        }
3816        Ok(())
3817    }
3818
3819    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3820    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3821    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3822    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3823    /// the serving isolation contract for sampled rows).
3824    #[allow(clippy::too_many_arguments)]
3825    pub fn gumbel_perturb_filtered_col(
3826        &self,
3827        x: &CudaSlice<f32>,
3828        col: usize,
3829        y: &mut CudaSlice<f32>,
3830        n: usize,
3831        seed: u64,
3832        stream_pos: u32,
3833        temp: f32,
3834        stat_max: &CudaSlice<f32>,
3835        stat_th: &CudaSlice<f32>,
3836        stat_idx: usize,
3837    ) -> Result<(), Box<dyn std::error::Error>> {
3838        let f = self.func("gumbel_perturb_filtered_col_f32");
3839        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3840        let (ci, si) = (col as i32, stat_idx as i32);
3841        let cfg = LaunchConfig {
3842            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3843            block_dim: (256, 1, 1),
3844            shared_mem_bytes: 0,
3845        };
3846        let __s_b = self.gpu.stream();
3847        let mut b = __s_b.launch_builder(&f);
3848        b.arg(x)
3849            .arg(&ci)
3850            .arg(&mut *y)
3851            .arg(&ni)
3852            .arg(&slo)
3853            .arg(&shi)
3854            .arg(&stream_pos)
3855            .arg(&temp)
3856            .arg(stat_max)
3857            .arg(stat_th)
3858            .arg(&si);
3859        unsafe {
3860            b.launch(cfg)?;
3861        }
3862        Ok(())
3863    }
3864
3865    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3866    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3867    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3868    /// reads it (counter is data, not state — graph-replay-safe).
3869    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3870        let f = self.func("memra_sctr_inc");
3871        let cfg = LaunchConfig {
3872            grid_dim: (1, 1, 1),
3873            block_dim: (1, 1, 1),
3874            shared_mem_bytes: 0,
3875        };
3876        let __s_b = self.gpu.stream();
3877        let mut b = __s_b.launch_builder(&f);
3878        b.arg(&mut *ctr);
3879        unsafe {
3880            b.launch(cfg)?;
3881        }
3882        Ok(())
3883    }
3884
3885    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3886    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3887    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3888    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3889    pub fn gumbel_perturb_ctr(
3890        &self,
3891        x: &CudaSlice<f32>,
3892        y: &mut CudaSlice<f32>,
3893        n: usize,
3894        seed: u64,
3895        ctr: &CudaSlice<u32>,
3896        temp: f32,
3897    ) -> Result<(), Box<dyn std::error::Error>> {
3898        let f = self.func("gumbel_perturb_ctr_f32");
3899        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3900        let cfg = LaunchConfig {
3901            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3902            block_dim: (256, 1, 1),
3903            shared_mem_bytes: 0,
3904        };
3905        let __s_b = self.gpu.stream();
3906        let mut b = __s_b.launch_builder(&f);
3907        b.arg(x)
3908            .arg(&mut *y)
3909            .arg(&ni)
3910            .arg(&slo)
3911            .arg(&shi)
3912            .arg(ctr)
3913            .arg(&temp);
3914        unsafe {
3915            b.launch(cfg)?;
3916        }
3917        Ok(())
3918    }
3919
3920    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3921    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3922    /// (smallest-index tie-break — matches the argmax-gate contract).
3923    pub fn softmax_gather(
3924        &self,
3925        x: &CudaSlice<f32>,
3926        row_stride: usize,
3927        ids: &CudaSlice<u32>,
3928        rows: &CudaSlice<i32>,
3929        out: &mut CudaSlice<f32>,
3930        n: usize,
3931        npair: usize,
3932        temp: f32,
3933    ) -> Result<(), Box<dyn std::error::Error>> {
3934        let f = self.func("softmax_gather_f32");
3935        let (ni, rs) = (n as i32, row_stride as i64);
3936        let np = npair as i32;
3937        let cfg = LaunchConfig {
3938            grid_dim: (npair as u32, 1, 1),
3939            block_dim: (256, 1, 1),
3940            shared_mem_bytes: 0,
3941        };
3942        let __s_b = self.gpu.stream();
3943        let mut b = __s_b.launch_builder(&f);
3944        b.arg(x)
3945            .arg(&rs)
3946            .arg(ids)
3947            .arg(rows)
3948            .arg(&mut *out)
3949            .arg(&ni)
3950            .arg(&np)
3951            .arg(&temp);
3952        unsafe {
3953            b.launch(cfg)?;
3954        }
3955        Ok(())
3956    }
3957
3958    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3959    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3960    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3961    pub fn residual_sample(
3962        &self,
3963        p: &CudaSlice<f32>,
3964        q: Option<&CudaSlice<f32>>,
3965        n: usize,
3966        temp: f32,
3967        seed: u64,
3968        stream_pos: u32,
3969        out_tok: &mut CudaSlice<u32>,
3970    ) -> Result<(), Box<dyn std::error::Error>> {
3971        let f = self.func("residual_sample_f32");
3972        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3973        let nth = 1024u32;
3974        let cfg = LaunchConfig {
3975            grid_dim: (1, 1, 1),
3976            block_dim: (nth, 1, 1),
3977            shared_mem_bytes: 0,
3978        };
3979        let has_q: i32 = q.is_some() as i32;
3980        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3981        let __s_b = self.gpu.stream();
3982        let mut b = __s_b.launch_builder(&f);
3983        b.arg(p)
3984            .arg(qbuf)
3985            .arg(&has_q)
3986            .arg(&ni)
3987            .arg(&temp)
3988            .arg(&slo)
3989            .arg(&shi)
3990            .arg(&stream_pos)
3991            .arg(&mut *out_tok);
3992        unsafe {
3993            b.launch(cfg)?;
3994        }
3995        Ok(())
3996    }
3997
3998    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3999    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
4000    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
4001    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
4002    pub fn with_moe_cache<R>(
4003        &self,
4004        max_block_bytes: usize,
4005        f: impl FnOnce(
4006            &mut crate::moe_cache::MoeSlotCache,
4007            &Engine,
4008        ) -> Result<R, Box<dyn std::error::Error>>,
4009    ) -> Result<R, Box<dyn std::error::Error>> {
4010        let mut guard = self.moe_cache.lock().unwrap();
4011        if guard.is_none() {
4012            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4013        }
4014        let cache = guard.as_mut().unwrap();
4015        f(cache, self)
4016    }
4017
4018    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
4019    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
4020    pub fn freeze_moe_cache(&self) {
4021        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4022            cache.freeze();
4023        }
4024    }
4025
4026    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
4027    /// Never constructs a cache.
4028    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4029        self.moe_cache
4030            .lock()
4031            .unwrap()
4032            .as_ref()
4033            .map(crate::moe_cache::MoeSlotCache::export_residency)
4034    }
4035
4036    pub(crate) fn moe_cache_frozen(&self) -> bool {
4037        self.moe_cache
4038            .lock()
4039            .unwrap()
4040            .as_ref()
4041            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4042    }
4043
4044    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
4045    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
4046    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
4047    /// while leaving the profiling warmup's established batched behavior untouched.
4048    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
4049    /// tokenwise arm anyway.)
4050    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4051        crate::cpu_experts::configured()
4052            && self.moe_cache_frozen()
4053            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4054    }
4055
4056    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
4057    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4058        assert!(
4059            self.moe_cache.lock().unwrap().is_none(),
4060            "MoE cache layout configured after cache construction"
4061        );
4062        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4063    }
4064
4065    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4066        self.moe_cache_layout.lock().unwrap().clone()
4067    }
4068
4069    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
4070    pub fn moe_cache_enabled() -> bool {
4071        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4072    }
4073
4074    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
4075    /// Returns None if the cache was never built (disabled or no MoE forward ran).
4076    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4077        let guard = self.moe_cache.lock().unwrap();
4078        guard
4079            .as_ref()
4080            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4081    }
4082
4083    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
4084    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
4085    /// callers compare a before/after snapshot around a decode window.
4086    pub fn cpu_expert_stats(
4087        &self,
4088    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4089        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4090    }
4091
4092    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
4093    /// the backend tail that resident-GPU expert work did not hide.
4094    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4095        crate::cpu_experts::predictor_stats()
4096    }
4097
4098    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4099        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4100    }
4101
4102    /// CPU-routed expert selections grouped by how many of their three projections were already
4103    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
4104    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4105        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4106    }
4107
4108    /// Positioned-read proof-backend counters:
4109    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
4110    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4111        let guard = self.moe_cache.lock().unwrap();
4112        guard
4113            .as_ref()
4114            .and_then(|cache| cache.pread_stats())
4115            .map(|stats| {
4116                (
4117                    stats.reads,
4118                    stats.bytes,
4119                    stats.read_errors,
4120                    stats.short_reads,
4121                    stats.fallbacks,
4122                    stats.buffer_waits,
4123                    stats.ring_full,
4124                )
4125            })
4126    }
4127
4128    /// Spill configuration values that warned and substituted their documented defaults.
4129    pub fn spill_config_fallbacks(&self) -> u64 {
4130        crate::spill_pread::config_fallbacks()
4131    }
4132
4133    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
4134    pub fn moe_cache_reset_counters(&self) {
4135        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4136            c.reset_counters();
4137        }
4138    }
4139
4140    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4141        Ok(self.gpu.stream().clone_htod(v)?)
4142    }
4143
4144    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
4145    /// past the final q4_0 block through their aligned window — the bytes never reach a
4146    /// result (funnelshift discards them) but must be mapped memory.
4147    pub fn htod_bytes_padded(
4148        &self,
4149        v: &[u8],
4150        pad: usize,
4151    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4152        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4153        {
4154            let mut view = d.slice_mut(0..v.len());
4155            self.gpu.stream().memcpy_htod(v, &mut view)?;
4156        }
4157        Ok(d)
4158    }
4159
4160    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
4161    pub fn copy_into(
4162        &self,
4163        dst: &mut CudaSlice<f32>,
4164        off: usize,
4165        src: &CudaSlice<f32>,
4166        len: usize,
4167    ) -> Result<(), Box<dyn std::error::Error>> {
4168        let mut view = dst.slice_mut(off..off + len);
4169        self.gpu
4170            .stream()
4171            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4172        Ok(())
4173    }
4174
4175    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
4176    /// u8 twin of copy_into (D2D byte-range copy at an offset).
4177    pub fn copy_u8_into(
4178        &self,
4179        dst: &mut CudaSlice<u8>,
4180        off: usize,
4181        src: &CudaSlice<u8>,
4182        len: usize,
4183    ) -> Result<(), Box<dyn std::error::Error>> {
4184        let mut view = dst.slice_mut(off..off + len);
4185        self.gpu
4186            .stream()
4187            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4188        Ok(())
4189    }
4190
4191    /// D2D byte-range copy with explicit source and destination offsets.
4192    pub fn copy_u8_range_into(
4193        &self,
4194        dst: &mut CudaSlice<u8>,
4195        dst_off: usize,
4196        src: &CudaSlice<u8>,
4197        src_off: usize,
4198        len: usize,
4199    ) -> Result<(), Box<dyn std::error::Error>> {
4200        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
4201        self.gpu
4202            .stream()
4203            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4204        Ok(())
4205    }
4206
4207    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
4208    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
4209    /// keeping the audited attention range contiguous without changing its absolute start.
4210    pub fn prepare_kv_append(
4211        &self,
4212        kv: &mut crate::cache::KvLayer,
4213        retain_from: usize,
4214        append_rows: usize,
4215    ) -> Result<usize, Box<dyn std::error::Error>> {
4216        let Some(plan) = kv
4217            .ring
4218            .as_ref()
4219            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4220            .transpose()?
4221        else {
4222            return Ok(kv.len);
4223        };
4224        match plan {
4225            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4226            crate::cache::KvRingAppend::Rebase {
4227                src_row,
4228                keep_rows,
4229                new_base,
4230                write_row,
4231            } => {
4232                if keep_rows > 0 {
4233                    let k_len = keep_rows * kv.k_tok_bytes;
4234                    let v_len = keep_rows * kv.v_tok_bytes;
4235                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4236                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4237                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4238                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4239                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4240                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4241                }
4242                kv.ring.as_mut().unwrap().apply_rebase(new_base);
4243                Ok(write_row)
4244            }
4245        }
4246    }
4247
4248    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
4249    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
4250    pub fn htod_u8_into(
4251        &self,
4252        dst: &mut CudaSlice<u8>,
4253        off: usize,
4254        src: &[u8],
4255    ) -> Result<(), Box<dyn std::error::Error>> {
4256        let mut view = dst.slice_mut(off..off + src.len());
4257        self.gpu.stream().memcpy_htod(src, &mut view)?;
4258        Ok(())
4259    }
4260
4261    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4262        b.slice(0..len)
4263    }
4264
4265    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
4266    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
4267    pub fn view_u8_range<'a>(
4268        &self,
4269        b: &'a CudaSlice<u8>,
4270        start: usize,
4271        end: usize,
4272    ) -> cudarc::driver::CudaView<'a, u8> {
4273        b.slice(start..end)
4274    }
4275    pub fn view_u8<'a>(
4276        &self,
4277        b: &'a CudaSlice<u8>,
4278        len: usize,
4279    ) -> cudarc::driver::CudaView<'a, u8> {
4280        b.slice(0..len)
4281    }
4282
4283    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
4284    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
4285    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
4286    pub fn append_kv_quantized(
4287        &self,
4288        k_row: &CudaSlice<f32>,
4289        v_row: &CudaSlice<f32>,
4290        kc: &mut CudaSlice<u8>,
4291        vc: &mut CudaSlice<u8>,
4292        t: usize,
4293        kv_dim_k: usize,
4294        kv_dim_v: usize,
4295        k_tok_bytes: usize,
4296        v_tok_bytes: usize,
4297        g: bool,
4298    ) -> Result<(), Box<dyn std::error::Error>> {
4299        let f = if g {
4300            self.func_g("append_quantize_kv_q8_0_q5_1")
4301        } else {
4302            self.func("append_quantize_kv_q8_0_q5_1")
4303        };
4304        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4305        let cfg = LaunchConfig {
4306            grid_dim: (nblk, 1, 1),
4307            block_dim: (32, 1, 1),
4308            shared_mem_bytes: 0,
4309        };
4310        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4311        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4312        let __s_b = self.gpu.stream();
4313        let mut b = __s_b.launch_builder(&f);
4314        b.arg(k_row)
4315            .arg(v_row)
4316            .arg(kc)
4317            .arg(vc)
4318            .arg(&ti)
4319            .arg(&kdk)
4320            .arg(&kdv)
4321            .arg(&ktb)
4322            .arg(&vtb);
4323        unsafe {
4324            b.launch(cfg)?;
4325        }
4326        Ok(())
4327    }
4328
4329    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
4330    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
4331    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
4332    pub fn append_kv_quantized_dc(
4333        &self,
4334        k_row: &CudaSlice<f32>,
4335        v_row: &CudaSlice<f32>,
4336        kc: &mut CudaSlice<u8>,
4337        vc: &mut CudaSlice<u8>,
4338        t_dev: &CudaSlice<i32>,
4339        kv_dim_k: usize,
4340        kv_dim_v: usize,
4341        k_tok_bytes: usize,
4342        v_tok_bytes: usize,
4343        g: bool,
4344    ) -> Result<(), Box<dyn std::error::Error>> {
4345        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4346        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4347        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4348        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
4349        if Self::pdl_on() && Self::pdl_wb_on() {
4350            use cudarc::driver::{DevicePtr, DevicePtrMut};
4351            let s = &self.gpu.stream();
4352            let (pk, _g0) = k_row.device_ptr(s);
4353            let (pv, _g1) = v_row.device_ptr(s);
4354            let (pkc, _g2) = kc.device_ptr_mut(s);
4355            let (pvc, _g3) = vc.device_ptr_mut(s);
4356            let (pt, _g4) = t_dev.device_ptr(s);
4357            let mut ps = [
4358                &pk as *const _ as *mut std::ffi::c_void,
4359                &pv as *const _ as *mut _,
4360                &pkc as *const _ as *mut _,
4361                &pvc as *const _ as *mut _,
4362                &pt as *const _ as *mut _,
4363                &kdk as *const _ as *mut _,
4364                &kdv as *const _ as *mut _,
4365                &ktb as *const _ as *mut _,
4366                &vtb as *const _ as *mut _,
4367            ];
4368            unsafe {
4369                self.launch_pdl_flash(
4370                    g,
4371                    "append_quantize_kv_q8_0_q5_1_dc",
4372                    (nblk, 1, 1),
4373                    (32, 1, 1),
4374                    0,
4375                    &mut ps,
4376                )?;
4377            }
4378            return Ok(());
4379        }
4380        let f = if g {
4381            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4382        } else {
4383            self.func("append_quantize_kv_q8_0_q5_1_dc")
4384        };
4385        let cfg = LaunchConfig {
4386            grid_dim: (nblk, 1, 1),
4387            block_dim: (32, 1, 1),
4388            shared_mem_bytes: 0,
4389        };
4390        let __s_b = self.gpu.stream();
4391        let mut b = __s_b.launch_builder(&f);
4392        b.arg(k_row)
4393            .arg(v_row)
4394            .arg(kc)
4395            .arg(vc)
4396            .arg(t_dev)
4397            .arg(&kdk)
4398            .arg(&kdv)
4399            .arg(&ktb)
4400            .arg(&vtb);
4401        unsafe {
4402            b.launch(cfg)?;
4403        }
4404        Ok(())
4405    }
4406
4407    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
4408    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
4409    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
4410    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
4411    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
4412    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
4413    #[allow(clippy::too_many_arguments)]
4414    pub fn append_kv_quantized_rows(
4415        &self,
4416        k_rows: &CudaSlice<f32>,
4417        v_rows: &CudaSlice<f32>,
4418        kc: &mut CudaSlice<u8>,
4419        vc: &mut CudaSlice<u8>,
4420        t0: usize,
4421        t: usize,
4422        kv_dim_k: usize,
4423        kv_dim_v: usize,
4424        k_tok_bytes: usize,
4425        v_tok_bytes: usize,
4426        g: bool,
4427    ) -> Result<(), Box<dyn std::error::Error>> {
4428        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4429            for i in 0..t {
4430                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4431                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4432                self.append_kv_quantized_view(
4433                    &k_row,
4434                    &v_row,
4435                    kc,
4436                    vc,
4437                    t0 + i,
4438                    kv_dim_k,
4439                    kv_dim_v,
4440                    k_tok_bytes,
4441                    v_tok_bytes,
4442                    g,
4443                )?;
4444            }
4445            return Ok(());
4446        }
4447        let f = if g {
4448            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4449        } else {
4450            self.func("append_quantize_kv_q8_0_q5_1_rows")
4451        };
4452        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4453        let cfg = LaunchConfig {
4454            grid_dim: (nblk, t as u32, 1),
4455            block_dim: (32, 1, 1),
4456            shared_mem_bytes: 0,
4457        };
4458        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4459        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4460        let __s_b = self.gpu.stream();
4461        let mut b = __s_b.launch_builder(&f);
4462        b.arg(k_rows)
4463            .arg(v_rows)
4464            .arg(kc)
4465            .arg(vc)
4466            .arg(&t0i)
4467            .arg(&kdk)
4468            .arg(&kdv)
4469            .arg(&ktb)
4470            .arg(&vtb);
4471        unsafe {
4472            b.launch(cfg)?;
4473        }
4474        Ok(())
4475    }
4476
4477    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
4478    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
4479    /// later, inside a captured graph) without a host round-trip.
4480    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4481        let f = self.func("inc_i32");
4482        let cfg = LaunchConfig {
4483            grid_dim: (1, 1, 1),
4484            block_dim: (1, 1, 1),
4485            shared_mem_bytes: 0,
4486        };
4487        let __s_b = self.gpu.stream();
4488        let mut b = __s_b.launch_builder(&f);
4489        b.arg(p);
4490        unsafe {
4491            b.launch(cfg)?;
4492        }
4493        Ok(())
4494    }
4495
4496    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
4497    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
4498    pub fn append_kv_quantized_view(
4499        &self,
4500        k_row: &cudarc::driver::CudaView<f32>,
4501        v_row: &cudarc::driver::CudaView<f32>,
4502        kc: &mut CudaSlice<u8>,
4503        vc: &mut CudaSlice<u8>,
4504        t: usize,
4505        kv_dim_k: usize,
4506        kv_dim_v: usize,
4507        k_tok_bytes: usize,
4508        v_tok_bytes: usize,
4509        g: bool,
4510    ) -> Result<(), Box<dyn std::error::Error>> {
4511        let stream = self.gpu.stream();
4512        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4513        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4514        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4515        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4516        let f = if g {
4517            self.func_g("append_quantize_kv_q8_0_q5_1")
4518        } else {
4519            self.func("append_quantize_kv_q8_0_q5_1")
4520        };
4521        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4522        let cfg = LaunchConfig {
4523            grid_dim: (nblk, 1, 1),
4524            block_dim: (32, 1, 1),
4525            shared_mem_bytes: 0,
4526        };
4527        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4528        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4529        let mut b = stream.launch_builder(&f);
4530        b.arg(k_row)
4531            .arg(v_row)
4532            .arg(kc)
4533            .arg(vc)
4534            .arg(&ti)
4535            .arg(&kdk)
4536            .arg(&kdv)
4537            .arg(&ktb)
4538            .arg(&vtb);
4539        unsafe {
4540            b.launch(cfg)?;
4541        }
4542        Ok(())
4543    }
4544
4545    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
4546    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
4547    pub fn copy_view_into(
4548        &self,
4549        dst: &mut CudaSlice<f32>,
4550        off: usize,
4551        src: &cudarc::driver::CudaView<f32>,
4552        len: usize,
4553    ) -> Result<(), Box<dyn std::error::Error>> {
4554        let mut view = dst.slice_mut(off..off + len);
4555        self.gpu
4556            .stream()
4557            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4558        Ok(())
4559    }
4560
4561    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
4562    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
4563    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
4564    pub fn clone_dtod(
4565        &self,
4566        src: &CudaSlice<f32>,
4567    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4568        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4569        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4570        Ok(dst)
4571    }
4572
4573    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4574    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4575    pub fn dtod_copy_view(
4576        &self,
4577        src: &cudarc::driver::CudaView<f32>,
4578        dst: &mut CudaSlice<f32>,
4579    ) -> Result<(), Box<dyn std::error::Error>> {
4580        self.gpu.stream().memcpy_dtod(src, dst)?;
4581        Ok(())
4582    }
4583
4584    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4585    pub fn dtod_copy_view_i8(
4586        &self,
4587        src: &cudarc::driver::CudaView<i8>,
4588        dst: &mut CudaSlice<i8>,
4589    ) -> Result<(), Box<dyn std::error::Error>> {
4590        self.gpu.stream().memcpy_dtod(src, dst)?;
4591        Ok(())
4592    }
4593
4594    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4595    pub fn dtod_copy_into(
4596        &self,
4597        src: &CudaSlice<f32>,
4598        dst: &mut CudaSlice<f32>,
4599        offset: usize,
4600    ) -> Result<(), Box<dyn std::error::Error>> {
4601        let n = src.len();
4602        let mut dv = dst.slice_mut(offset..offset + n);
4603        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4604        Ok(())
4605    }
4606
4607    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4608    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4609    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4610    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4611    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4612    pub fn copy_batch_uniform_f32(
4613        &self,
4614        table: &CudaSlice<u64>,
4615        n: usize,
4616        words: usize,
4617    ) -> Result<(), Box<dyn std::error::Error>> {
4618        if n == 0 || words == 0 {
4619            return Ok(());
4620        }
4621        debug_assert!(
4622            table.len() >= 2 * n,
4623            "pointer table must hold n srcs + n dsts"
4624        );
4625        let f = self.func("copy_batch_uniform_f32");
4626        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4627        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4628        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4629        let (ni, wi) = (n as i32, words as i32);
4630        let cfg = LaunchConfig {
4631            grid_dim: (chunks, n as u32, 1),
4632            block_dim: (256, 1, 1),
4633            shared_mem_bytes: 0,
4634        };
4635        let __s = self.gpu.stream();
4636        let mut b = __s.launch_builder(&f);
4637        b.arg(table).arg(&ni).arg(&wi);
4638        unsafe {
4639            b.launch(cfg)?;
4640        }
4641        Ok(())
4642    }
4643
4644    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4645    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4646    pub fn htod_u64_into(
4647        &self,
4648        v: &[u64],
4649        dst: &mut CudaSlice<u64>,
4650    ) -> Result<(), Box<dyn std::error::Error>> {
4651        let mut view = dst.slice_mut(0..v.len());
4652        self.gpu.stream().memcpy_htod(v, &mut view)?;
4653        Ok(())
4654    }
4655
4656    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4657    /// device pointer-table entry at run time, so a captured graph follows the gdn
4658    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4659    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4660    pub fn copy_indirect_src_f32(
4661        &self,
4662        src_entry: &cudarc::driver::CudaView<u64>,
4663        dst: &mut CudaSlice<f32>,
4664        dst_off: usize,
4665        words: usize,
4666    ) -> Result<(), Box<dyn std::error::Error>> {
4667        let f = self.func("copy_indirect_src_f32");
4668        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4669        let wi = words as i32;
4670        let cfg = LaunchConfig {
4671            grid_dim: (chunks, 1, 1),
4672            block_dim: (256, 1, 1),
4673            shared_mem_bytes: 0,
4674        };
4675        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4676        let __s = self.gpu.stream();
4677        let mut b = __s.launch_builder(&f);
4678        b.arg(src_entry).arg(&mut dv).arg(&wi);
4679        unsafe {
4680            b.launch(cfg)?;
4681        }
4682        Ok(())
4683    }
4684
4685    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4686    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4687        self.alloc_uninit::<i8>(n)
4688    }
4689
4690    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4691    pub fn qmatvec(
4692        &self,
4693        w: &CudaSlice<u8>,
4694        x: &CudaSlice<f32>,
4695        m: usize,
4696        in_f: usize,
4697        out_f: usize,
4698        qtype: i32,
4699        row_bytes: usize,
4700    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4701        let f = self.func("qmatvec_f32");
4702        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4703        let cfg = LaunchConfig {
4704            grid_dim: (out_f as u32, m as u32, 1),
4705            block_dim: (256, 1, 1),
4706            shared_mem_bytes: 0,
4707        };
4708        let (inf, outf, mi, qt, rb) =
4709            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4710        let __s_b = self.gpu.stream();
4711        let mut b = __s_b.launch_builder(&f);
4712        b.arg(w)
4713            .arg(x)
4714            .arg(&mut y)
4715            .arg(&inf)
4716            .arg(&outf)
4717            .arg(&mi)
4718            .arg(&qt)
4719            .arg(&rb);
4720        unsafe {
4721            b.launch(cfg)?;
4722        }
4723        Ok(y)
4724    }
4725
4726    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4727    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4728        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4729        self.keep_if_capturing(&s);
4730        Ok(s)
4731    }
4732
4733    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4734    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4735    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4736    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4737        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4738        self.keep_if_capturing(&s);
4739        Ok(s)
4740    }
4741
4742    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4743    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4744    pub fn memset_zeros_view(
4745        &self,
4746        dst: &mut cudarc::driver::CudaViewMut<f32>,
4747    ) -> Result<(), Box<dyn std::error::Error>> {
4748        self.gpu.stream().memset_zeros(dst)?;
4749        Ok(())
4750    }
4751
4752    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4753    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4754    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4755    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4756    /// stream would require an event).
4757    pub fn stage_expert(
4758        &self,
4759        host_bytes: &[u8],
4760        scratch: &mut CudaSlice<u8>,
4761        off: usize,
4762    ) -> Result<(), Box<dyn std::error::Error>> {
4763        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4764        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4765        Ok(())
4766    }
4767
4768    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4769    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4770    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4771    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4772    /// One CTA per token row, 256 threads (one per expert).
4773    pub fn moe_router_topk(
4774        &self,
4775        logits: &CudaSlice<f32>,
4776        t: usize,
4777        n_expert: usize,
4778        n_used: usize,
4779    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4780        let f = self.func("moe_router_topk_f32");
4781        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4782        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4783        let cfg = LaunchConfig {
4784            grid_dim: (t as u32, 1, 1),
4785            block_dim: (n_expert as u32, 1, 1),
4786            shared_mem_bytes: 0,
4787        };
4788        let (ne, nu) = (n_expert as i32, n_used as i32);
4789        let __s_b = self.gpu.stream();
4790        let mut b = __s_b.launch_builder(&f);
4791        b.arg(logits)
4792            .arg(&mut sel_idx)
4793            .arg(&mut sel_w)
4794            .arg(&ne)
4795            .arg(&nu);
4796        unsafe {
4797            b.launch(cfg)?;
4798        }
4799        Ok((sel_idx, sel_w))
4800    }
4801
4802    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4803    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4804    pub fn moe_router_topk_scaled(
4805        &self,
4806        logits: &CudaSlice<f32>,
4807        t: usize,
4808        n_expert: usize,
4809        n_used: usize,
4810        ex_scale: &CudaSlice<f32>,
4811    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4812        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4813        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4814        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4815        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4816        let f = self.func("moe_router_topk_scaled_f32");
4817        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4818        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4819        let cfg = LaunchConfig {
4820            grid_dim: (t as u32, 1, 1),
4821            block_dim: (n_expert as u32, 1, 1),
4822            shared_mem_bytes: 0,
4823        };
4824        let (ne, nu) = (n_expert as i32, n_used as i32);
4825        let __s_b = self.gpu.stream();
4826        let mut b = __s_b.launch_builder(&f);
4827        b.arg(logits)
4828            .arg(&mut sel_idx)
4829            .arg(&mut sel_w)
4830            .arg(&ne)
4831            .arg(&nu)
4832            .arg(ex_scale);
4833        unsafe {
4834            b.launch(cfg)?;
4835        }
4836        Ok((sel_idx, sel_w))
4837    }
4838
4839    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4840    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4841    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4842    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4843    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4844    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4845    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4846    pub fn moe_router_topk_host(
4847        &self,
4848        logits: &CudaSlice<f32>,
4849        t: usize,
4850        n_expert: usize,
4851        n_used: usize,
4852    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4853        let f = self.func("moe_router_topk_f32");
4854        let n = t * n_used;
4855        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4856        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4857        let cfg = LaunchConfig {
4858            grid_dim: (t as u32, 1, 1),
4859            block_dim: (n_expert as u32, 1, 1),
4860            shared_mem_bytes: 0,
4861        };
4862        let (ne, nu) = (n_expert as i32, n_used as i32);
4863        let __s_b = self.gpu.stream();
4864        let mut b = __s_b.launch_builder(&f);
4865        b.arg(logits)
4866            .arg(&mut sel_idx)
4867            .arg(&mut sel_w)
4868            .arg(&ne)
4869            .arg(&nu);
4870        unsafe {
4871            b.launch(cfg)?;
4872        }
4873        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4874        let bytes = n * 8;
4875        let mut guard = self.router_stage.lock().unwrap();
4876        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4877            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4878        }
4879        let stage = guard.as_mut().unwrap();
4880        let (si, sw) = unsafe {
4881            (
4882                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4883                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4884            )
4885        };
4886        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4887        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4888        self.gpu.stream().synchronize()?; // ONE sync for both
4889        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4890    }
4891
4892    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4893    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4894    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4895    #[allow(clippy::too_many_arguments)]
4896    pub fn moe_router_sigmoid_topk(
4897        &self,
4898        logits: &CudaSlice<f32>,
4899        t: usize,
4900        n_expert: usize,
4901        n_used: usize,
4902        active_count: usize,
4903        correction_bias: &CudaSlice<f32>,
4904        active: &CudaSlice<u8>,
4905        scaling_factor: f32,
4906        route_norm: bool,
4907    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4908        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4909        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4910            return Err(format!(
4911                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4912            )
4913            .into());
4914        }
4915        if logits.len() < t * n_expert
4916            || correction_bias.len() != n_expert
4917            || active.len() != n_expert
4918        {
4919            return Err(format!(
4920                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4921                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4922            ).into());
4923        }
4924        let f = self.func(crate::sigmoid_topk_kernel(
4925            crate::sig_expf_dev_on(),
4926            crate::topk_fast_on(),
4927            n_used,
4928        ));
4929        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4930        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4931        let threads = n_expert.div_ceil(32) * 32;
4932        let cfg = LaunchConfig {
4933            grid_dim: (t as u32, 1, 1),
4934            block_dim: (threads as u32, 1, 1),
4935            shared_mem_bytes: 0,
4936        };
4937        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4938        let __s_b = self.gpu.stream();
4939        let mut b = __s_b.launch_builder(&f);
4940        b.arg(logits)
4941            .arg(correction_bias)
4942            .arg(active)
4943            .arg(&mut sel_idx)
4944            .arg(&mut sel_w)
4945            .arg(&ne)
4946            .arg(&nu)
4947            .arg(&scaling_factor)
4948            .arg(&rn);
4949        unsafe {
4950            b.launch(cfg)?;
4951        }
4952        Ok((sel_idx, sel_w))
4953    }
4954
4955    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
4956    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
4957    #[allow(clippy::too_many_arguments)]
4958    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
4959    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
4960    /// the model engine can wait on it with a same-device stream memop.
4961    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4962        if ptr == 0 {
4963            return Err("ring_flag_raw: unarmed flag".into());
4964        }
4965        let f = self.func("memra_ring_flag");
4966        let cfg = LaunchConfig {
4967            grid_dim: (1, 1, 1),
4968            block_dim: (32, 1, 1),
4969            shared_mem_bytes: 0,
4970        };
4971        let __s_b = self.gpu.stream();
4972        let mut b = __s_b.launch_builder(&f);
4973        b.arg(&ptr).arg(&value);
4974        unsafe {
4975            b.launch(cfg)?;
4976        }
4977        Ok(())
4978    }
4979
4980    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
4981    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
4982    pub fn moe_sel_w_mirror(
4983        &self,
4984        sel_src: &CudaSlice<i32>,
4985        w_src: &CudaSlice<f32>,
4986        sel_dst: &mut CudaSlice<i32>,
4987        w_dst: &mut CudaSlice<f32>,
4988        n: usize,
4989    ) -> Result<(), Box<dyn std::error::Error>> {
4990        if n == 0
4991            || n > 32
4992            || sel_src.len() < n
4993            || w_src.len() < n
4994            || sel_dst.len() < n
4995            || w_dst.len() < n
4996        {
4997            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4998        }
4999        let f = self.func("moe_sel_w_mirror");
5000        let cfg = LaunchConfig {
5001            grid_dim: (1, 1, 1),
5002            block_dim: (32, 1, 1),
5003            shared_mem_bytes: 0,
5004        };
5005        let ni = n as i32;
5006        let __s_b = self.gpu.stream();
5007        let mut b = __s_b.launch_builder(&f);
5008        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5009        unsafe {
5010            b.launch(cfg)?;
5011        }
5012        Ok(())
5013    }
5014
5015    pub fn moe_router_sigmoid_topk_into(
5016        &self,
5017        logits: &CudaSlice<f32>,
5018        t: usize,
5019        n_expert: usize,
5020        n_used: usize,
5021        active_count: usize,
5022        correction_bias: &CudaSlice<f32>,
5023        active: &CudaSlice<u8>,
5024        scaling_factor: f32,
5025        route_norm: bool,
5026        sel_idx: &mut CudaSlice<i32>,
5027        sel_w: &mut CudaSlice<f32>,
5028    ) -> Result<(), Box<dyn std::error::Error>> {
5029        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5030        if n_expert == 0
5031            || n_expert > 1024
5032            || n_used == 0
5033            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
5034            || n_used > n_expert
5035            || logits.len() < t * n_expert
5036            || correction_bias.len() != n_expert
5037            || active.len() != n_expert
5038            || sel_idx.len() < t * n_used
5039            || sel_w.len() < t * n_used
5040        {
5041            return Err("sigmoid router _into geometry mismatch".into());
5042        }
5043        let f = self.func(crate::sigmoid_topk_kernel(
5044            crate::sig_expf_dev_on(),
5045            crate::topk_fast_on(),
5046            n_used,
5047        ));
5048        let threads = n_expert.div_ceil(32) * 32;
5049        let cfg = LaunchConfig {
5050            grid_dim: (t as u32, 1, 1),
5051            block_dim: (threads as u32, 1, 1),
5052            shared_mem_bytes: 0,
5053        };
5054        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5055        let __s_b = self.gpu.stream();
5056        let mut b = __s_b.launch_builder(&f);
5057        b.arg(logits)
5058            .arg(correction_bias)
5059            .arg(active)
5060            .arg(&mut *sel_idx)
5061            .arg(&mut *sel_w)
5062            .arg(&ne)
5063            .arg(&nu)
5064            .arg(&scaling_factor)
5065            .arg(&rn);
5066        unsafe {
5067            b.launch(cfg)?;
5068        }
5069        Ok(())
5070    }
5071
5072    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
5073    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
5074    #[allow(clippy::too_many_arguments)]
5075    pub fn moe_router_sigmoid_topk_host(
5076        &self,
5077        logits: &CudaSlice<f32>,
5078        t: usize,
5079        n_expert: usize,
5080        n_used: usize,
5081        active_count: usize,
5082        correction_bias: &CudaSlice<f32>,
5083        active: &CudaSlice<u8>,
5084        scaling_factor: f32,
5085        route_norm: bool,
5086    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5087        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5088            logits,
5089            t,
5090            n_expert,
5091            n_used,
5092            active_count,
5093            correction_bias,
5094            active,
5095            scaling_factor,
5096            route_norm,
5097        )?;
5098        let n = t * n_used;
5099        let bytes = n * 8;
5100        let mut guard = self.router_stage.lock().unwrap();
5101        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5102            *guard = Some(PinnedStage::new(bytes.max(4096))?);
5103        }
5104        let stage = guard.as_mut().unwrap();
5105        let (si, sw) = unsafe {
5106            (
5107                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5108                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5109            )
5110        };
5111        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5112        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5113        self.gpu.stream().synchronize()?;
5114        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5115    }
5116
5117    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
5118    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
5119    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
5120    pub fn stage_expert_async(
5121        &self,
5122        host_bytes: &[u8],
5123        scratch: &mut CudaSlice<u8>,
5124        off: usize,
5125    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5126        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5127        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5128        Ok(self.copy_stream.record_event(None)?)
5129    }
5130
5131    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
5132    pub fn compute_wait(
5133        &self,
5134        ev: &cudarc::driver::CudaEvent,
5135    ) -> Result<(), Box<dyn std::error::Error>> {
5136        self.gpu.stream().wait(ev)?;
5137        Ok(())
5138    }
5139
5140    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
5141    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
5142    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
5143    /// CudaView base+offset pointer is honored by the launch arg.
5144    pub fn qmatvec_view(
5145        &self,
5146        w: &CudaSlice<u8>,
5147        range: std::ops::Range<usize>,
5148        x: &cudarc::driver::CudaView<f32>,
5149        m: usize,
5150        in_f: usize,
5151        out_f: usize,
5152        qtype: i32,
5153        row_bytes: usize,
5154    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5155        let f = self.func("qmatvec_f32");
5156        let wv = w.slice(range); // CudaView<u8>, offset honored
5157        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
5158        let cfg = LaunchConfig {
5159            grid_dim: (out_f as u32, m as u32, 1),
5160            block_dim: (256, 1, 1),
5161            shared_mem_bytes: 0,
5162        };
5163        let (inf, outf, mi, qt, rb) =
5164            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5165        let __s_b = self.gpu.stream();
5166        let mut b = __s_b.launch_builder(&f);
5167        b.arg(&wv)
5168            .arg(x)
5169            .arg(&mut y)
5170            .arg(&inf)
5171            .arg(&outf)
5172            .arg(&mi)
5173            .arg(&qt)
5174            .arg(&rb);
5175        unsafe {
5176            b.launch(cfg)?;
5177        }
5178        Ok(y)
5179    }
5180
5181    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
5182    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
5183    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
5184    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
5185    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
5186    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
5187    #[allow(clippy::too_many_arguments)]
5188    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
5189    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
5190    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
5191    pub fn moe_gate_up_silu8_q8(
5192        &self,
5193        gp: WPtr8,
5194        up: WPtr8,
5195        aq: &CudaSlice<i8>,
5196        ad: &CudaSlice<f32>,
5197        in_f: usize,
5198        n_ff: usize,
5199        n_used: usize,
5200        qt_g: i32,
5201        qt_u: i32,
5202        rb_g: usize,
5203        rb_u: usize,
5204    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5205        let f = self.func("moe_gate_up_silu8_q8");
5206        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5207        let cfg = LaunchConfig {
5208            grid_dim: (n_ff as u32, n_used as u32, 1),
5209            block_dim: (32, 1, 1),
5210            shared_mem_bytes: 0,
5211        };
5212        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5213        let __s_b = self.gpu.stream();
5214        let mut b = __s_b.launch_builder(&f);
5215        b.arg(&gp)
5216            .arg(&up)
5217            .arg(aq)
5218            .arg(ad)
5219            .arg(&mut act)
5220            .arg(&inf)
5221            .arg(&nff)
5222            .arg(&qt_g)
5223            .arg(&qt_u)
5224            .arg(&rbg)
5225            .arg(&rbu);
5226        unsafe {
5227            b.launch(cfg)?;
5228        }
5229        Ok(act)
5230    }
5231
5232    #[allow(clippy::too_many_arguments)]
5233    pub fn moe_down8_fma_q8(
5234        &self,
5235        dp: WPtr8,
5236        w: F32x8,
5237        aq2: &CudaSlice<i8>,
5238        ad2: &CudaSlice<f32>,
5239        dst: &mut cudarc::driver::CudaViewMut<f32>,
5240        in_f: usize,
5241        out_f: usize,
5242        n_used: usize,
5243        qt: i32,
5244        rb: usize,
5245    ) -> Result<(), Box<dyn std::error::Error>> {
5246        let f = self.func("moe_down8_fma_q8");
5247        let cfg = LaunchConfig {
5248            grid_dim: (out_f as u32, 1, 1),
5249            block_dim: (32, 1, 1),
5250            shared_mem_bytes: 0,
5251        };
5252        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5253        let __s_b = self.gpu.stream();
5254        let mut b = __s_b.launch_builder(&f);
5255        b.arg(&dp)
5256            .arg(&w)
5257            .arg(aq2)
5258            .arg(ad2)
5259            .arg(dst)
5260            .arg(&inf)
5261            .arg(&outf)
5262            .arg(&nu)
5263            .arg(&qt)
5264            .arg(&rbi);
5265        unsafe {
5266            b.launch(cfg)?;
5267        }
5268        Ok(())
5269    }
5270
5271    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
5272    pub fn qmatvec_expert_q8(
5273        &self,
5274        w: &CudaSlice<u8>,
5275        range: std::ops::Range<usize>,
5276        aq: &CudaSlice<i8>,
5277        ad: &CudaSlice<f32>,
5278        m: usize,
5279        in_f: usize,
5280        out_f: usize,
5281        qtype: i32,
5282        row_bytes: usize,
5283    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5284        let f = self.func("qmatvec_expert_q8");
5285        let wv = w.slice(range);
5286        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5287        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
5288        let cfg = LaunchConfig {
5289            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5290            block_dim: (32, ROWS, 1),
5291            shared_mem_bytes: 0,
5292        };
5293        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5294        let __s_b = self.gpu.stream();
5295        let mut b = __s_b.launch_builder(&f);
5296        b.arg(&wv)
5297            .arg(aq)
5298            .arg(ad)
5299            .arg(&mut y)
5300            .arg(&inf)
5301            .arg(&outf)
5302            .arg(&mi)
5303            .arg(&qtype)
5304            .arg(&rbi);
5305        unsafe {
5306            b.launch(cfg)?;
5307        }
5308        Ok(y)
5309    }
5310
5311    pub fn moe_gate_up_silu8(
5312        &self,
5313        gp: WPtr8,
5314        up: WPtr8,
5315        x: &cudarc::driver::CudaView<f32>,
5316        in_f: usize,
5317        n_ff: usize,
5318        n_used: usize,
5319        qt_g: i32,
5320        qt_u: i32,
5321        rb_g: usize,
5322        rb_u: usize,
5323    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5324        let f = self.func("moe_gate_up_silu8_f32");
5325        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
5326        let cfg = LaunchConfig {
5327            grid_dim: (n_ff as u32, n_used as u32, 1),
5328            block_dim: (256, 1, 1),
5329            shared_mem_bytes: 0,
5330        };
5331        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5332        let __s_b = self.gpu.stream();
5333        let mut b = __s_b.launch_builder(&f);
5334        b.arg(&gp)
5335            .arg(&up)
5336            .arg(x)
5337            .arg(&mut act)
5338            .arg(&inf)
5339            .arg(&nff)
5340            .arg(&qt_g)
5341            .arg(&qt_u)
5342            .arg(&rbg)
5343            .arg(&rbu);
5344        unsafe {
5345            b.launch(cfg)?;
5346        }
5347        Ok(act)
5348    }
5349
5350    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
5351    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
5352    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
5353    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
5354    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
5355    #[allow(clippy::too_many_arguments)]
5356    pub fn moe_down8_fma_into(
5357        &self,
5358        dp: WPtr8,
5359        w: F32x8,
5360        act: &CudaSlice<f32>,
5361        dst: &mut cudarc::driver::CudaViewMut<f32>,
5362        in_f: usize,
5363        out_f: usize,
5364        n_used: usize,
5365        qt: i32,
5366        rb: usize,
5367    ) -> Result<(), Box<dyn std::error::Error>> {
5368        let f = self.func("moe_down8_fma_f32");
5369        let cfg = LaunchConfig {
5370            grid_dim: (out_f as u32, 1, 1),
5371            block_dim: (256, 1, 1),
5372            shared_mem_bytes: 0,
5373        };
5374        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5375        let __s_b = self.gpu.stream();
5376        let mut b = __s_b.launch_builder(&f);
5377        b.arg(&dp)
5378            .arg(&w)
5379            .arg(act)
5380            .arg(dst)
5381            .arg(&inf)
5382            .arg(&outf)
5383            .arg(&nu)
5384            .arg(&qt)
5385            .arg(&rbv);
5386        unsafe {
5387            b.launch(cfg)?;
5388        }
5389        Ok(())
5390    }
5391
5392    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
5393    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
5394    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
5395    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
5396    #[allow(clippy::too_many_arguments)]
5397    /// dp4a q8 twin of the _dev pair (resident-experts arc).
5398    ///
5399    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
5400    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
5401    /// down's FMA chain stays slot-ordered serial). Seams:
5402    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
5403    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
5404    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
5405    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
5406    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
5407    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
5408    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
5409    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
5410    ///                       only) | w8h2 (h2 x slot-parallel)
5411    #[allow(clippy::too_many_arguments)]
5412    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
5413    #[allow(clippy::too_many_arguments)]
5414    pub fn moe_pairs_matvec_q8(
5415        &self,
5416        table: &CudaSlice<u64>,
5417        proj: i32,
5418        pair_tok: &CudaSlice<i32>,
5419        pair_ex: &CudaSlice<i32>,
5420        aq: &CudaSlice<i8>,
5421        ad: &CudaSlice<f32>,
5422        in_f: usize,
5423        out_f: usize,
5424        n_expert: usize,
5425        n_pairs: usize,
5426        qtype: i32,
5427        row_bytes: usize,
5428    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5429        let f = self.func("moe_pairs_matvec_q8");
5430        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5431        const ROWS: u32 = 4;
5432        let cfg = LaunchConfig {
5433            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5434            block_dim: (32, ROWS, 1),
5435            shared_mem_bytes: 0,
5436        };
5437        let (inf, outf, ne, np, rbi) = (
5438            in_f as i32,
5439            out_f as i32,
5440            n_expert as i32,
5441            n_pairs as i32,
5442            row_bytes as i64,
5443        );
5444        let __s_b = self.gpu.stream();
5445        let mut b = __s_b.launch_builder(&f);
5446        b.arg(table)
5447            .arg(&proj)
5448            .arg(pair_tok)
5449            .arg(pair_ex)
5450            .arg(aq)
5451            .arg(ad)
5452            .arg(&mut y)
5453            .arg(&inf)
5454            .arg(&outf)
5455            .arg(&ne)
5456            .arg(&np)
5457            .arg(&qtype)
5458            .arg(&rbi);
5459        unsafe {
5460            b.launch(cfg)?;
5461        }
5462        Ok(y)
5463    }
5464
5465    /// Expert-major pair matvec (weight-reuse across each expert's token group).
5466    #[allow(clippy::too_many_arguments)]
5467    pub fn moe_pairs_matvec_q8_em(
5468        &self,
5469        table: &CudaSlice<u64>,
5470        proj: i32,
5471        ex_ids: &CudaSlice<i32>,
5472        ex_off: &CudaSlice<i32>,
5473        ex_pairs: &CudaSlice<i32>,
5474        pair_tok: &CudaSlice<i32>,
5475        aq: &CudaSlice<i8>,
5476        ad: &CudaSlice<f32>,
5477        in_f: usize,
5478        out_f: usize,
5479        n_expert: usize,
5480        n_active: usize,
5481        n_pairs: usize,
5482        qtype: i32,
5483        row_bytes: usize,
5484    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5485        let f = self.func("moe_pairs_matvec_q8_em");
5486        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5487        const ROWS: u32 = 4;
5488        let cfg = LaunchConfig {
5489            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5490            block_dim: (32, ROWS, 1),
5491            shared_mem_bytes: 0,
5492        };
5493        let (inf, outf, ne, na, rbi) = (
5494            in_f as i32,
5495            out_f as i32,
5496            n_expert as i32,
5497            n_active as i32,
5498            row_bytes as i64,
5499        );
5500        let __s_b = self.gpu.stream();
5501        let mut b = __s_b.launch_builder(&f);
5502        b.arg(table)
5503            .arg(&proj)
5504            .arg(ex_ids)
5505            .arg(ex_off)
5506            .arg(ex_pairs)
5507            .arg(pair_tok)
5508            .arg(aq)
5509            .arg(ad)
5510            .arg(&mut y)
5511            .arg(&inf)
5512            .arg(&outf)
5513            .arg(&ne)
5514            .arg(&na)
5515            .arg(&qtype)
5516            .arg(&rbi);
5517        unsafe {
5518            b.launch(cfg)?;
5519        }
5520        Ok(y)
5521    }
5522
5523    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
5524    // weight group once per (row,group) then dp4a's across the expert's token group.
5525    #[allow(clippy::too_many_arguments)]
5526    pub fn moe_pairs_matvec_q8_dec(
5527        &self,
5528        table: &CudaSlice<u64>,
5529        proj: i32,
5530        ex_ids: &CudaSlice<i32>,
5531        ex_off: &CudaSlice<i32>,
5532        ex_pairs: &CudaSlice<i32>,
5533        pair_tok: &CudaSlice<i32>,
5534        aq: &CudaSlice<i8>,
5535        ad: &CudaSlice<f32>,
5536        in_f: usize,
5537        out_f: usize,
5538        n_expert: usize,
5539        n_active: usize,
5540        n_pairs: usize,
5541        qtype: i32,
5542        row_bytes: usize,
5543    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5544        let f = self.func("moe_pairs_matvec_q8_dec");
5545        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5546        const ROWS: u32 = 4;
5547        let cfg = LaunchConfig {
5548            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5549            block_dim: (32, ROWS, 1),
5550            shared_mem_bytes: 0,
5551        };
5552        let (inf, outf, ne, na, rbi) = (
5553            in_f as i32,
5554            out_f as i32,
5555            n_expert as i32,
5556            n_active as i32,
5557            row_bytes as i64,
5558        );
5559        let __s_b = self.gpu.stream();
5560        let mut b = __s_b.launch_builder(&f);
5561        b.arg(table)
5562            .arg(&proj)
5563            .arg(ex_ids)
5564            .arg(ex_off)
5565            .arg(ex_pairs)
5566            .arg(pair_tok)
5567            .arg(aq)
5568            .arg(ad)
5569            .arg(&mut y)
5570            .arg(&inf)
5571            .arg(&outf)
5572            .arg(&ne)
5573            .arg(&na)
5574            .arg(&qtype)
5575            .arg(&rbi);
5576        unsafe {
5577            b.launch(cfg)?;
5578        }
5579        Ok(y)
5580    }
5581
5582    pub fn moe_pairs_gelu_mul(
5583        &self,
5584        gate: &CudaSlice<f32>,
5585        up: &CudaSlice<f32>,
5586        n: usize,
5587    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5588        let f = self.func("moe_pairs_gelu_mul");
5589        let mut act = self.alloc_uninit::<f32>(n)?;
5590        let cfg = LaunchConfig::for_num_elems(n as u32);
5591        let nl = n as i64;
5592        let __s_b = self.gpu.stream();
5593        let mut b = __s_b.launch_builder(&f);
5594        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5595        unsafe {
5596            b.launch(cfg)?;
5597        }
5598        Ok(act)
5599    }
5600
5601    pub fn moe_pairs_silu_mul(
5602        &self,
5603        gate: &CudaSlice<f32>,
5604        up: &CudaSlice<f32>,
5605        n: usize,
5606    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5607        let f = self.func("moe_pairs_silu_mul");
5608        let mut act = self.alloc_uninit::<f32>(n)?;
5609        let cfg = LaunchConfig::for_num_elems(n as u32);
5610        let nl = n as i64;
5611        let __s_b = self.gpu.stream();
5612        let mut b = __s_b.launch_builder(&f);
5613        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5614        unsafe {
5615            b.launch(cfg)?;
5616        }
5617        Ok(act)
5618    }
5619
5620    #[allow(clippy::too_many_arguments)]
5621    pub fn moe_pairs_scatter(
5622        &self,
5623        y_down: &CudaSlice<f32>,
5624        pair_w: &CudaSlice<f32>,
5625        tok_pair_off: &CudaSlice<i32>,
5626        tok_pair_ids: &CudaSlice<i32>,
5627        moe_out: &mut CudaSlice<f32>,
5628        t: usize,
5629        n_embd: usize,
5630    ) -> Result<(), Box<dyn std::error::Error>> {
5631        let f = self.func("moe_pairs_scatter");
5632        let cfg = LaunchConfig {
5633            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5634            block_dim: (256, 1, 1),
5635            shared_mem_bytes: 0,
5636        };
5637        let ne = n_embd as i32;
5638        let __s_b = self.gpu.stream();
5639        let mut b = __s_b.launch_builder(&f);
5640        b.arg(y_down)
5641            .arg(pair_w)
5642            .arg(tok_pair_off)
5643            .arg(tok_pair_ids)
5644            .arg(moe_out)
5645            .arg(&ne);
5646        unsafe {
5647            b.launch(cfg)?;
5648        }
5649        Ok(())
5650    }
5651
5652    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
5653    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
5654    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
5655    #[allow(clippy::too_many_arguments)]
5656    pub fn moe_gate_up_gelu8_dev_q8(
5657        &self,
5658        table: &CudaSlice<u64>,
5659        sel: &cudarc::driver::CudaView<i32>,
5660        aq: &CudaSlice<i8>,
5661        ad: &CudaSlice<f32>,
5662        in_f: usize,
5663        n_ff: usize,
5664        n_used: usize,
5665        n_expert: usize,
5666        qt_g: i32,
5667        qt_u: i32,
5668        rb_g: usize,
5669        rb_u: usize,
5670    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5671        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5672        let (inf, nff, ne, rbg, rbu) = (
5673            in_f as i32,
5674            n_ff as i32,
5675            n_expert as i32,
5676            rb_g as i64,
5677            rb_u as i64,
5678        );
5679        let f = self.func("moe_gate_up_gelu8_dev_q8");
5680        let cfg = LaunchConfig {
5681            grid_dim: (n_ff as u32, n_used as u32, 1),
5682            block_dim: (32, 1, 1),
5683            shared_mem_bytes: 0,
5684        };
5685        let __s_b = self.gpu.stream();
5686        let mut b = __s_b.launch_builder(&f);
5687        b.arg(table)
5688            .arg(sel)
5689            .arg(aq)
5690            .arg(ad)
5691            .arg(&mut act)
5692            .arg(&inf)
5693            .arg(&nff)
5694            .arg(&ne)
5695            .arg(&qt_g)
5696            .arg(&qt_u)
5697            .arg(&rbg)
5698            .arg(&rbu);
5699        unsafe {
5700            b.launch(cfg)?;
5701        }
5702        Ok(act)
5703    }
5704
5705    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5706    #[allow(clippy::too_many_arguments)]
5707    pub fn moe_gate_up_gelu8_dev_q8_rows(
5708        &self,
5709        table: &CudaSlice<u64>,
5710        sel: &CudaSlice<i32>,
5711        aq: &CudaSlice<i8>,
5712        ad: &CudaSlice<f32>,
5713        t: usize,
5714        in_f: usize,
5715        n_ff: usize,
5716        n_used: usize,
5717        n_expert: usize,
5718        qt_g: i32,
5719        qt_u: i32,
5720        rb_g: usize,
5721        rb_u: usize,
5722    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5723        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5724        let (inf, nff, ne, rbg, rbu, nu) = (
5725            in_f as i32,
5726            n_ff as i32,
5727            n_expert as i32,
5728            rb_g as i64,
5729            rb_u as i64,
5730            n_used as i32,
5731        );
5732        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5733        let cfg = LaunchConfig {
5734            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5735            block_dim: (32, 1, 1),
5736            shared_mem_bytes: 0,
5737        };
5738        let __s_b = self.gpu.stream();
5739        let mut b = __s_b.launch_builder(&f);
5740        b.arg(table)
5741            .arg(sel)
5742            .arg(aq)
5743            .arg(ad)
5744            .arg(&mut act)
5745            .arg(&inf)
5746            .arg(&nff)
5747            .arg(&ne)
5748            .arg(&qt_g)
5749            .arg(&qt_u)
5750            .arg(&rbg)
5751            .arg(&rbu)
5752            .arg(&nu);
5753        unsafe {
5754            b.launch(cfg)?;
5755        }
5756        Ok(act)
5757    }
5758
5759    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5760    #[allow(clippy::too_many_arguments)]
5761    pub fn moe_gate_up_gelu8_dev_q8_csr(
5762        &self,
5763        table: &CudaSlice<u64>,
5764        sel: &CudaSlice<i32>,
5765        aq: &CudaSlice<i8>,
5766        ad: &CudaSlice<f32>,
5767        n_pairs: usize,
5768        in_f: usize,
5769        n_ff: usize,
5770        n_used: usize,
5771        n_expert: usize,
5772        qt_g: i32,
5773        qt_u: i32,
5774        rb_g: usize,
5775        rb_u: usize,
5776    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5777        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5778        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5779            in_f as i32,
5780            n_ff as i32,
5781            n_expert as i32,
5782            rb_g as i64,
5783            rb_u as i64,
5784            n_used as i32,
5785            n_pairs as i32,
5786        );
5787        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5788        let cfg = LaunchConfig {
5789            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5790            block_dim: (32, 1, 1),
5791            shared_mem_bytes: 0,
5792        };
5793        let __s_b = self.gpu.stream();
5794        let mut b = __s_b.launch_builder(&f);
5795        b.arg(table)
5796            .arg(sel)
5797            .arg(aq)
5798            .arg(ad)
5799            .arg(&mut act)
5800            .arg(&inf)
5801            .arg(&nff)
5802            .arg(&ne)
5803            .arg(&qt_g)
5804            .arg(&qt_u)
5805            .arg(&rbg)
5806            .arg(&rbu)
5807            .arg(&nu)
5808            .arg(&npi);
5809        unsafe {
5810            b.launch(cfg)?;
5811        }
5812        Ok(act)
5813    }
5814
5815    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5816    #[allow(clippy::too_many_arguments)]
5817    pub fn moe_down8_fma_dev_q8_rows_g(
5818        &self,
5819        table: &CudaSlice<u64>,
5820        sel: &CudaSlice<i32>,
5821        w: &CudaSlice<f32>,
5822        aq2: &CudaSlice<i8>,
5823        ad2: &CudaSlice<f32>,
5824        dst: &mut CudaSlice<f32>,
5825        t: usize,
5826        in_f: usize,
5827        out_f: usize,
5828        n_used: usize,
5829        n_expert: usize,
5830        qt: i32,
5831        rb: usize,
5832    ) -> Result<(), Box<dyn std::error::Error>> {
5833        let (inf, outf, nu, ne, rbi) = (
5834            in_f as i32,
5835            out_f as i32,
5836            n_used as i32,
5837            n_expert as i32,
5838            rb as i64,
5839        );
5840        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5841        // eight warps, then replay the original slot-ordered FMA chain. Every
5842        // other shape retains the generic one-warp rows kernel.
5843        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5844        let f = self.func(if step_b1_w8 {
5845            "moe_down8_fma_dev_q8_rows_w8"
5846        } else {
5847            "moe_down8_fma_dev_q8_rows_g"
5848        });
5849        let cfg = LaunchConfig {
5850            grid_dim: (out_f as u32, 1, t as u32),
5851            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5852            shared_mem_bytes: 0,
5853        };
5854        let __s_b = self.gpu.stream();
5855        let mut b = __s_b.launch_builder(&f);
5856        b.arg(table)
5857            .arg(sel)
5858            .arg(w)
5859            .arg(aq2)
5860            .arg(ad2)
5861            .arg(dst)
5862            .arg(&inf)
5863            .arg(&outf)
5864            .arg(&nu)
5865            .arg(&ne)
5866            .arg(&qt)
5867            .arg(&rbi);
5868        unsafe {
5869            b.launch(cfg)?;
5870        }
5871        Ok(())
5872    }
5873
5874    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5875    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5876    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5877    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5878        let (out_f, in_f) = (2048usize, 2816usize);
5879        let nblk = in_f / 32;
5880        let mut seed = 0x9E3779B97F4A7C15u64;
5881        let mut rng = move || {
5882            seed = seed
5883                .wrapping_mul(6364136223846793005)
5884                .wrapping_add(1442695040888963407);
5885            (seed >> 33) as u8
5886        };
5887        let mut w = vec![0u8; out_f * nblk * 18];
5888        for b in w.iter_mut() {
5889            *b = rng();
5890        }
5891        for r in 0..out_f {
5892            for g in 0..nblk {
5893                let off = (r * nblk + g) * 18;
5894                w[off] = 0x00;
5895                w[off + 1] = 0x2C; // sane half d
5896            }
5897        }
5898        let qplane = out_f * nblk * 16;
5899        let mut wrp = vec![0u8; w.len()];
5900        for r in 0..out_f {
5901            for g in 0..nblk {
5902                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5903                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5904                    .copy_from_slice(&src[0..2]);
5905                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5906            }
5907        }
5908        let w_d = self.htod_bytes(&w)?;
5909        let wrp_d = self.htod_bytes(&wrp)?;
5910        let mut aq = vec![0i8; m * in_f];
5911        for v in aq.iter_mut() {
5912            *v = rng() as i8;
5913        }
5914        let aq_d = self.htod_i8(&aq)?;
5915        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5916        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5917        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5918        const RPB: u32 = 4;
5919        let cfg = LaunchConfig {
5920            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5921            block_dim: (32, RPB, 1),
5922            shared_mem_bytes: 0,
5923        };
5924        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5925        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5926        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5927        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5928        {
5929            let __s_b = self.gpu.stream();
5930            let mut b = __s_b.launch_builder(&fb);
5931            b.arg(&w_d)
5932                .arg(&aq_d)
5933                .arg(&ad_d)
5934                .arg(&mut y0)
5935                .arg(&inf)
5936                .arg(&outf)
5937                .arg(&mi)
5938                .arg(&rb);
5939            unsafe {
5940                b.launch(cfg)?;
5941            }
5942            let __s_b = self.gpu.stream();
5943            let mut b = __s_b.launch_builder(&fr);
5944            b.arg(&wrp_d)
5945                .arg(&aq_d)
5946                .arg(&ad_d)
5947                .arg(&mut y1)
5948                .arg(&inf)
5949                .arg(&outf)
5950                .arg(&mi)
5951                .arg(&qp);
5952            unsafe {
5953                b.launch(cfg)?;
5954            }
5955        }
5956        self.gpu.stream().synchronize()?;
5957        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5958        let nd = h0
5959            .iter()
5960            .zip(&h1)
5961            .filter(|(a, b)| a.to_bits() != b.to_bits())
5962            .count();
5963        if nd != 0 {
5964            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5965        }
5966        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5967            self.gpu.stream().synchronize()?;
5968            let t0 = std::time::Instant::now();
5969            for _ in 0..500 {
5970                if rp {
5971                    let __s_b = self.gpu.stream();
5972                    let mut b = __s_b.launch_builder(&fr);
5973                    b.arg(&wrp_d)
5974                        .arg(&aq_d)
5975                        .arg(&ad_d)
5976                        .arg(&mut y1)
5977                        .arg(&inf)
5978                        .arg(&outf)
5979                        .arg(&mi)
5980                        .arg(&qp);
5981                    unsafe {
5982                        b.launch(cfg)?;
5983                    }
5984                } else {
5985                    let __s_b = self.gpu.stream();
5986                    let mut b = __s_b.launch_builder(&fb);
5987                    b.arg(&w_d)
5988                        .arg(&aq_d)
5989                        .arg(&ad_d)
5990                        .arg(&mut y0)
5991                        .arg(&inf)
5992                        .arg(&outf)
5993                        .arg(&mi)
5994                        .arg(&rb);
5995                    unsafe {
5996                        b.launch(cfg)?;
5997                    }
5998                }
5999            }
6000            self.gpu.stream().synchronize()?;
6001            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6002        };
6003        let _ = time(false)?;
6004        let _ = time(true)?; // warm
6005        Ok((time(false)?, time(true)?))
6006    }
6007
6008    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
6009    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
6010    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
6011    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
6012    pub fn build_q4_rp4(
6013        &self,
6014        t: &mut crate::model::GpuTensor,
6015    ) -> Result<(), Box<dyn std::error::Error>> {
6016        use crate::model::GpuTensor;
6017        let GpuTensor::Quant {
6018            bytes,
6019            qtype,
6020            row_bytes,
6021            ne,
6022            rp4,
6023            ..
6024        } = t
6025        else {
6026            return Ok(());
6027        };
6028        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6029            return Ok(());
6030        }
6031        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6032        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6033            return Ok(());
6034        }
6035        let nblk = in_f / 32;
6036        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6037        let f = self.func("q4_0_split_rp_build");
6038        let n = (out_f * nblk) as i32;
6039        let cfg = LaunchConfig {
6040            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6041            block_dim: (256, 1, 1),
6042            shared_mem_bytes: 0,
6043        };
6044        let (of, nb) = (out_f as i32, nblk as i32);
6045        let _ = n;
6046        let __s_b = self.gpu.stream();
6047        let mut b = __s_b.launch_builder(&f);
6048        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6049        unsafe {
6050            b.launch(cfg)?;
6051        }
6052        *rp4 = Some(dst);
6053        Ok(())
6054    }
6055
6056    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
6057    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
6058    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
6059    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
6060    pub fn build_q8_rp4(
6061        &self,
6062        t: &mut crate::model::GpuTensor,
6063    ) -> Result<(), Box<dyn std::error::Error>> {
6064        use crate::model::GpuTensor;
6065        let GpuTensor::Quant {
6066            bytes,
6067            qtype,
6068            row_bytes,
6069            ne,
6070            rp4,
6071            ..
6072        } = t
6073        else {
6074            return Ok(());
6075        };
6076        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6077            return Ok(());
6078        }
6079        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6080        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6081            return Ok(());
6082        }
6083        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6084        Ok(())
6085    }
6086
6087    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
6088    /// mirror without a GpuTensor (same kernel the loader path above uses).
6089    pub fn build_q8_rp4_raw(
6090        &self,
6091        bytes: &CudaSlice<u8>,
6092        in_f: usize,
6093        out_f: usize,
6094    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6095        assert!(in_f % 32 == 0);
6096        let nblk = in_f / 32;
6097        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6098        let f = self.func("q8_0_split_rp_build");
6099        let cfg = LaunchConfig {
6100            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6101            block_dim: (256, 1, 1),
6102            shared_mem_bytes: 0,
6103        };
6104        let (of, nb) = (out_f as i32, nblk as i32);
6105        let __s_b = self.gpu.stream();
6106        let mut b = __s_b.launch_builder(&f);
6107        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6108        unsafe {
6109            b.launch(cfg)?;
6110        }
6111        Ok(dst)
6112    }
6113
6114    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
6115    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
6116    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
6117    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
6118    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
6119    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
6120    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
6121    pub fn build_q4k_rp4(
6122        &self,
6123        t: &mut crate::model::GpuTensor,
6124    ) -> Result<(), Box<dyn std::error::Error>> {
6125        use crate::model::GpuTensor;
6126        let GpuTensor::Quant {
6127            bytes,
6128            qtype,
6129            row_bytes,
6130            ne,
6131            rp4,
6132            ..
6133        } = t
6134        else {
6135            return Ok(());
6136        };
6137        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6138            return Ok(());
6139        }
6140        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6141        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6142            return Ok(());
6143        }
6144        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6145        Ok(())
6146    }
6147
6148    pub fn build_q6k_rp4(
6149        &self,
6150        t: &mut crate::model::GpuTensor,
6151    ) -> Result<(), Box<dyn std::error::Error>> {
6152        use crate::model::GpuTensor;
6153        let GpuTensor::Quant {
6154            bytes,
6155            qtype,
6156            row_bytes,
6157            ne,
6158            rp4,
6159            ..
6160        } = t
6161        else {
6162            return Ok(());
6163        };
6164        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6165            return Ok(());
6166        }
6167        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6168        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6169            return Ok(());
6170        }
6171        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6172        Ok(())
6173    }
6174
6175    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
6176    pub fn build_kq_rp4_raw(
6177        &self,
6178        bytes: &CudaSlice<u8>,
6179        in_f: usize,
6180        out_f: usize,
6181        qtype: i32,
6182    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6183        assert!(in_f % 256 == 0);
6184        let nsbk = in_f / 256;
6185        let (sb_bytes, kname) = match qtype {
6186            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6187            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6188            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6189        };
6190        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6191        let f = self.func(kname);
6192        let cfg = LaunchConfig {
6193            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6194            block_dim: (256, 1, 1),
6195            shared_mem_bytes: 0,
6196        };
6197        let (of, nb) = (out_f as i32, nsbk as i32);
6198        let __s_b = self.gpu.stream();
6199        let mut b = __s_b.launch_builder(&f);
6200        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6201        unsafe {
6202            b.launch(cfg)?;
6203        }
6204        Ok(dst)
6205    }
6206
6207    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
6208    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
6209    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
6210    pub fn kqrp_enabled() -> bool {
6211        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6212        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6213            Ok("0") => false,
6214            Ok(_) => true,
6215            Err(_) => cfg!(memra_hopper_mma),
6216        })
6217    }
6218
6219    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
6220    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
6221    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
6222    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
6223    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
6224    pub fn build_q4_rp_swap(
6225        &self,
6226        t: &mut crate::model::GpuTensor,
6227    ) -> Result<bool, Box<dyn std::error::Error>> {
6228        use crate::model::GpuTensor;
6229        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
6230        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
6231        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
6232        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
6233        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
6234        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
6235        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
6236        // this fn's OWN builder serves may ever be swapped; everything else refuses
6237        // here, regardless of walk ordering.
6238        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6239            return Ok(false);
6240        }
6241        self.build_q4_rp4(t)?;
6242        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
6243        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6244            return Ok(false);
6245        };
6246        match rp4.take() {
6247            Some(split) => {
6248                *bytes = split; // the GGUF-layout buffer drops here
6249                *rp = true;
6250                Ok(true)
6251            }
6252            None => Ok(false),
6253        }
6254    }
6255
6256    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
6257    pub fn q4rp_enabled() -> bool {
6258        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6259        *ON.get_or_init(|| {
6260            std::env::var("MEMRA_Q4RP")
6261                .map(|v| v != "0")
6262                .unwrap_or(true)
6263        })
6264    }
6265
6266    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
6267    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
6268    pub fn copy_rows_strided(
6269        &self,
6270        src: &CudaSlice<f32>,
6271        dst: &mut CudaSlice<f32>,
6272        row_elems: usize,
6273        n_rows: usize,
6274        src_stride: usize,
6275        src_off: usize,
6276    ) -> Result<(), Box<dyn std::error::Error>> {
6277        let f = self.func("copy_rows_strided_f32");
6278        let cfg = LaunchConfig {
6279            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6280            block_dim: (256, 1, 1),
6281            shared_mem_bytes: 0,
6282        };
6283        let (re, nr) = (row_elems as i32, n_rows as i32);
6284        let (st, off) = (src_stride as i64, src_off as i64);
6285        let __s_b = self.gpu.stream();
6286        let mut b = __s_b.launch_builder(&f);
6287        b.arg(src)
6288            .arg(&mut *dst)
6289            .arg(&re)
6290            .arg(&nr)
6291            .arg(&st)
6292            .arg(&off);
6293        unsafe {
6294            b.launch(cfg)?;
6295        }
6296        Ok(())
6297    }
6298
6299    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
6300    ///
6301    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
6302    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
6303    /// one peer copy per token.
6304    pub fn place_rows_strided(
6305        &self,
6306        src: &CudaSlice<f32>,
6307        dst: &mut CudaSlice<f32>,
6308        row_elems: usize,
6309        n_rows: usize,
6310        dst_stride: usize,
6311        dst_off: usize,
6312    ) -> Result<(), Box<dyn std::error::Error>> {
6313        if row_elems == 0 || n_rows == 0 {
6314            return Err("strided row placement requires nonzero rows and row width".into());
6315        }
6316        let src_len = n_rows
6317            .checked_mul(row_elems)
6318            .ok_or("strided row placement source size overflow")?;
6319        let dst_len = n_rows
6320            .checked_sub(1)
6321            .and_then(|rows| rows.checked_mul(dst_stride))
6322            .and_then(|base| base.checked_add(dst_off))
6323            .and_then(|base| base.checked_add(row_elems))
6324            .ok_or("strided row placement destination size overflow")?;
6325        let row_end = dst_off
6326            .checked_add(row_elems)
6327            .ok_or("strided row placement row size overflow")?;
6328        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6329            return Err(format!(
6330                "strided row placement geometry mismatch: src={} need_src={src_len} \
6331                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6332                 dst_stride={dst_stride} dst_off={dst_off}",
6333                src.len(),
6334                dst.len(),
6335            )
6336            .into());
6337        }
6338        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6339            return Err("strided row placement exceeds CUDA kernel geometry".into());
6340        }
6341        let f = self.func("place_rows_strided_f32");
6342        let cfg = LaunchConfig {
6343            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6344            block_dim: (256, 1, 1),
6345            shared_mem_bytes: 0,
6346        };
6347        let (re, nr) = (row_elems as i32, n_rows as i32);
6348        let (st, off) = (dst_stride as i64, dst_off as i64);
6349        let __s_b = self.gpu.stream();
6350        let mut b = __s_b.launch_builder(&f);
6351        b.arg(src)
6352            .arg(&mut *dst)
6353            .arg(&re)
6354            .arg(&nr)
6355            .arg(&st)
6356            .arg(&off);
6357        unsafe {
6358            b.launch(cfg)?;
6359        }
6360        Ok(())
6361    }
6362
6363    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
6364    pub fn u32_set_k(
6365        &self,
6366        dst: &mut CudaSlice<u32>,
6367        v: u32,
6368        idx: usize,
6369    ) -> Result<(), Box<dyn std::error::Error>> {
6370        let f = self.func("u32_set_k");
6371        let cfg = LaunchConfig {
6372            grid_dim: (1, 1, 1),
6373            block_dim: (1, 1, 1),
6374            shared_mem_bytes: 0,
6375        };
6376        let ii = idx as i32;
6377        let __s_b = self.gpu.stream();
6378        let mut b = __s_b.launch_builder(&f);
6379        b.arg(dst).arg(&v).arg(&ii);
6380        unsafe {
6381            b.launch(cfg)?;
6382        }
6383        Ok(())
6384    }
6385
6386    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
6387    pub fn i32_add_k(
6388        &self,
6389        d: &mut CudaSlice<i32>,
6390        v: i32,
6391    ) -> Result<(), Box<dyn std::error::Error>> {
6392        let f = self.func("i32_add_k");
6393        let cfg = LaunchConfig {
6394            grid_dim: (1, 1, 1),
6395            block_dim: (32, 1, 1),
6396            shared_mem_bytes: 0,
6397        };
6398        let __s_b = self.gpu.stream();
6399        let mut b = __s_b.launch_builder(&f);
6400        b.arg(d).arg(&v);
6401        unsafe {
6402            b.launch(cfg)?;
6403        }
6404        Ok(())
6405    }
6406
6407    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
6408    pub fn i32_iota_from(
6409        &self,
6410        ctr: &CudaSlice<i32>,
6411        dst: &mut CudaSlice<i32>,
6412        n: usize,
6413    ) -> Result<(), Box<dyn std::error::Error>> {
6414        let f = self.func("i32_iota_from");
6415        let cfg = LaunchConfig::for_num_elems(n as u32);
6416        let ni = n as i32;
6417        let __s_b = self.gpu.stream();
6418        let mut b = __s_b.launch_builder(&f);
6419        b.arg(ctr).arg(dst).arg(&ni);
6420        unsafe {
6421            b.launch(cfg)?;
6422        }
6423        Ok(())
6424    }
6425
6426    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
6427    pub fn u32_map_k(
6428        &self,
6429        buf: &mut CudaSlice<u32>,
6430        map: &CudaSlice<u32>,
6431        idx: usize,
6432    ) -> Result<(), Box<dyn std::error::Error>> {
6433        let f = self.func("u32_map_k");
6434        let cfg = LaunchConfig {
6435            grid_dim: (1, 1, 1),
6436            block_dim: (1, 1, 1),
6437            shared_mem_bytes: 0,
6438        };
6439        let ii = idx as i32;
6440        let __s_b = self.gpu.stream();
6441        let mut b = __s_b.launch_builder(&f);
6442        b.arg(buf).arg(map).arg(&ii);
6443        unsafe {
6444            b.launch(cfg)?;
6445        }
6446        Ok(())
6447    }
6448
6449    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
6450    #[allow(clippy::too_many_arguments)]
6451    pub fn u32_pack2(
6452        &self,
6453        a: &CudaSlice<u32>,
6454        off_a: usize,
6455        n1: usize,
6456        b_in: &CudaSlice<u32>,
6457        n2: usize,
6458        out: &mut CudaSlice<u32>,
6459    ) -> Result<(), Box<dyn std::error::Error>> {
6460        let f = self.func("u32_pack2");
6461        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6462        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6463        let __s_b = self.gpu.stream();
6464        let mut b = __s_b.launch_builder(&f);
6465        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6466        unsafe {
6467            b.launch(cfg)?;
6468        }
6469        Ok(())
6470    }
6471
6472    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
6473    pub fn moe_w_exscale(
6474        &self,
6475        w: &mut CudaSlice<f32>,
6476        sel: &CudaSlice<i32>,
6477        s: &CudaSlice<f32>,
6478        n: usize,
6479    ) -> Result<(), Box<dyn std::error::Error>> {
6480        let f = self.func("moe_w_exscale");
6481        let cfg = LaunchConfig::for_num_elems(n as u32);
6482        let ni = n as i32;
6483        let __s_b = self.gpu.stream();
6484        let mut b = __s_b.launch_builder(&f);
6485        b.arg(w).arg(sel).arg(s).arg(&ni);
6486        unsafe {
6487            b.launch(cfg)?;
6488        }
6489        Ok(())
6490    }
6491
6492    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
6493    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
6494    pub fn moe_w_scale_by_expert(
6495        &self,
6496        w: &mut CudaSlice<f32>,
6497        sel: &CudaSlice<i32>,
6498        macros: &CudaSlice<f32>,
6499        n_expert: usize,
6500        n: usize,
6501    ) -> Result<(), Box<dyn std::error::Error>> {
6502        let f = self.func("moe_w_scale_by_expert");
6503        let cfg = LaunchConfig {
6504            grid_dim: (n.div_ceil(64) as u32, 1, 1),
6505            block_dim: (64, 1, 1),
6506            shared_mem_bytes: 0,
6507        };
6508        let (ne, nn) = (n_expert as i32, n as i32);
6509        let __s_b = self.gpu.stream();
6510        let mut b = __s_b.launch_builder(&f);
6511        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6512        unsafe {
6513            b.launch(cfg)?;
6514        }
6515        Ok(())
6516    }
6517
6518    pub fn moe_gate_up_silu8_dev_q8(
6519        &self,
6520        table: &CudaSlice<u64>,
6521        sel: &cudarc::driver::CudaView<i32>,
6522        aq: &CudaSlice<i8>,
6523        ad: &CudaSlice<f32>,
6524        in_f: usize,
6525        n_ff: usize,
6526        n_used: usize,
6527        n_expert: usize,
6528        qt_g: i32,
6529        qt_u: i32,
6530        rb_g: usize,
6531        rb_u: usize,
6532        macros: &CudaSlice<f32>,
6533    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6534        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6535        let (mode, wpb) = GU.get_or_init(|| {
6536            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6537            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6538                .ok()
6539                .and_then(|v| v.parse().ok())
6540                .unwrap_or(4u32)
6541                .clamp(1, 16);
6542            (mode, wpb)
6543        });
6544        let (mode, wpb) = (mode.as_str(), *wpb);
6545        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6546        let (inf, nff, ne, rbg, rbu) = (
6547            in_f as i32,
6548            n_ff as i32,
6549            n_expert as i32,
6550            rb_g as i64,
6551            rb_u as i64,
6552        );
6553        let (f, cfg) = match mode {
6554            "1" | "2" | "4" => {
6555                let rpw: u32 = mode.parse().unwrap();
6556                let f = self.func(match rpw {
6557                    1 => "moe_gate_up_silu8_dev_q8_r1",
6558                    2 => "moe_gate_up_silu8_dev_q8_r2",
6559                    _ => "moe_gate_up_silu8_dev_q8_r4",
6560                });
6561                let rows_per_block = (rpw * wpb) as usize;
6562                let gx = n_ff.div_ceil(rows_per_block) as u32;
6563                (
6564                    f,
6565                    LaunchConfig {
6566                        grid_dim: (gx, n_used as u32, 1),
6567                        block_dim: (32, wpb, 1),
6568                        shared_mem_bytes: 0,
6569                    },
6570                )
6571            }
6572            "j8" if n_used <= 32 => (
6573                self.func("moe_gate_up_silu8_dev_q8_j8"),
6574                LaunchConfig {
6575                    grid_dim: (n_ff as u32, 1, 1),
6576                    block_dim: (32, n_used as u32, 1),
6577                    shared_mem_bytes: 0,
6578                },
6579            ),
6580            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
6581            "vsm2" => {
6582                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6583                let sh = (rb_g + rb_u) as u32;
6584                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6585                f.set_attribute(
6586                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6587                    sh as i32,
6588                )?;
6589                (
6590                    f,
6591                    LaunchConfig {
6592                        grid_dim: (n_ff as u32, n_used as u32, 1),
6593                        block_dim: (32, 1, 1),
6594                        shared_mem_bytes: sh,
6595                    },
6596                )
6597            }
6598            "vsm" => {
6599                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6600                let sh = (rb_g + rb_u) as u32;
6601                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6602                f.set_attribute(
6603                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6604                    sh as i32,
6605                )?;
6606                (
6607                    f,
6608                    LaunchConfig {
6609                        grid_dim: (n_ff as u32, n_used as u32, 1),
6610                        block_dim: (32, 1, 1),
6611                        shared_mem_bytes: sh,
6612                    },
6613                )
6614            }
6615            "sg" => (
6616                self.func("moe_gate_up_silu8_dev_q8_sg"),
6617                LaunchConfig {
6618                    grid_dim: (n_ff as u32, n_used as u32, 1),
6619                    block_dim: (32, 1, 1),
6620                    shared_mem_bytes: 0,
6621                },
6622            ),
6623            "j8sg" if n_used <= 32 => (
6624                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6625                LaunchConfig {
6626                    grid_dim: (n_ff as u32, 1, 1),
6627                    block_dim: (32, n_used as u32, 1),
6628                    shared_mem_bytes: 0,
6629                },
6630            ),
6631            "u64" if in_f == 2048 => (
6632                self.func("moe_gate_up_silu8_dev_q8_u64"),
6633                LaunchConfig {
6634                    grid_dim: (n_ff as u32, n_used as u32, 1),
6635                    block_dim: (32, 1, 1),
6636                    shared_mem_bytes: 0,
6637                },
6638            ),
6639            "gs4" if in_f == 2048 => (
6640                self.func("moe_gate_up_silu8_dev_q8_gs4"),
6641                LaunchConfig {
6642                    grid_dim: (n_ff as u32, n_used as u32, 1),
6643                    block_dim: (32, 4, 1),
6644                    shared_mem_bytes: 0,
6645                },
6646            ),
6647            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
6648            "v" | "" => (
6649                self.func("moe_gate_up_silu8_dev_q8_v"),
6650                LaunchConfig {
6651                    grid_dim: (n_ff as u32, n_used as u32, 1),
6652                    block_dim: (32, 1, 1),
6653                    shared_mem_bytes: 0,
6654                },
6655            ),
6656            "s2" => (
6657                self.func("moe_gate_up_silu8_dev_q8_s2"),
6658                LaunchConfig {
6659                    grid_dim: (n_ff as u32, n_used as u32, 1),
6660                    block_dim: (32, 2, 1),
6661                    shared_mem_bytes: 0,
6662                },
6663            ),
6664            "s2z" => {
6665                let rz = wpb.min(16); // s2z smem tile is [16][2]
6666                (
6667                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
6668                    LaunchConfig {
6669                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6670                        block_dim: (32, 2, rz),
6671                        shared_mem_bytes: 0,
6672                    },
6673                )
6674            }
6675            _ => (
6676                self.func("moe_gate_up_silu8_dev_q8"),
6677                LaunchConfig {
6678                    grid_dim: (n_ff as u32, n_used as u32, 1),
6679                    block_dim: (32, 1, 1),
6680                    shared_mem_bytes: 0,
6681                },
6682            ),
6683        };
6684        let __s_b = self.gpu.stream();
6685        let mut b = __s_b.launch_builder(&f);
6686        b.arg(table)
6687            .arg(sel)
6688            .arg(aq)
6689            .arg(ad)
6690            .arg(&mut act)
6691            .arg(&inf)
6692            .arg(&nff)
6693            .arg(&ne)
6694            .arg(&qt_g)
6695            .arg(&qt_u)
6696            .arg(&rbg)
6697            .arg(&rbu)
6698            .arg(macros);
6699        unsafe {
6700            b.launch(cfg)?;
6701        }
6702        Ok(act)
6703    }
6704
6705    #[allow(clippy::too_many_arguments)]
6706    pub fn moe_down8_fma_dev_q8(
6707        &self,
6708        table: &CudaSlice<u64>,
6709        sel: &cudarc::driver::CudaView<i32>,
6710        w: &cudarc::driver::CudaView<f32>,
6711        aq2: &CudaSlice<i8>,
6712        ad2: &CudaSlice<f32>,
6713        dst: &mut cudarc::driver::CudaViewMut<f32>,
6714        in_f: usize,
6715        out_f: usize,
6716        n_used: usize,
6717        n_expert: usize,
6718        qt: i32,
6719        rb: usize,
6720    ) -> Result<(), Box<dyn std::error::Error>> {
6721        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6722        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6723        let (inf, outf, nu, ne, rbi) = (
6724            in_f as i32,
6725            out_f as i32,
6726            n_used as i32,
6727            n_expert as i32,
6728            rb as i64,
6729        );
6730        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
6731        // the h2 twins are nsb==16 (in_f==512) shape-gated.
6732        let (f, cfg) = match mode.as_str() {
6733            m @ ("1" | "2" | "4") if n_used <= 8 => {
6734                let rpw: usize = m.parse().unwrap();
6735                let f = self.func(match rpw {
6736                    1 => "moe_down8_fma_dev_q8_w8r1",
6737                    2 => "moe_down8_fma_dev_q8_w8r2",
6738                    _ => "moe_down8_fma_dev_q8_w8r4",
6739                });
6740                (
6741                    f,
6742                    LaunchConfig {
6743                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6744                        block_dim: (32, n_used as u32, 1),
6745                        shared_mem_bytes: 0,
6746                    },
6747                )
6748            }
6749            "h2" if in_f == 512 => (
6750                self.func("moe_down8_fma_dev_q8_h2"),
6751                LaunchConfig {
6752                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6753                    block_dim: (32, 1, 1),
6754                    shared_mem_bytes: 0,
6755                },
6756            ),
6757            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6758            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6759            "" if in_f == 704 && n_used <= 8 => (
6760                self.func("moe_down8_fma_dev_q8_w8r2"),
6761                LaunchConfig {
6762                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6763                    block_dim: (32, n_used as u32, 1),
6764                    shared_mem_bytes: 0,
6765                },
6766            ),
6767            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6768            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6769            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6770            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6771                self.func("moe_down8_fma_dev_q8_w8h2v"),
6772                LaunchConfig {
6773                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6774                    block_dim: (32, n_used as u32, 1),
6775                    shared_mem_bytes: 0,
6776                },
6777            ),
6778            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6779                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6780                LaunchConfig {
6781                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6782                    block_dim: (32, n_used as u32, 1),
6783                    shared_mem_bytes: 0,
6784                },
6785            ),
6786            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6787                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6788                LaunchConfig {
6789                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6790                    block_dim: (32, n_used as u32, 1),
6791                    shared_mem_bytes: 0,
6792                },
6793            ),
6794            "w8h2" if in_f == 512 && n_used <= 8 => (
6795                self.func("moe_down8_fma_dev_q8_w8h2"),
6796                LaunchConfig {
6797                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6798                    block_dim: (32, n_used as u32, 1),
6799                    shared_mem_bytes: 0,
6800                },
6801            ),
6802            _ => (
6803                self.func("moe_down8_fma_dev_q8"),
6804                LaunchConfig {
6805                    grid_dim: (out_f as u32, 1, 1),
6806                    block_dim: (32, 1, 1),
6807                    shared_mem_bytes: 0,
6808                },
6809            ),
6810        };
6811        let __s_b = self.gpu.stream();
6812        let mut b = __s_b.launch_builder(&f);
6813        b.arg(table)
6814            .arg(sel)
6815            .arg(w)
6816            .arg(aq2)
6817            .arg(ad2)
6818            .arg(dst)
6819            .arg(&inf)
6820            .arg(&outf)
6821            .arg(&nu)
6822            .arg(&ne)
6823            .arg(&qt)
6824            .arg(&rbi);
6825        unsafe {
6826            b.launch(cfg)?;
6827        }
6828        Ok(())
6829    }
6830
6831    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6832    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6833    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6834    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6835    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6836    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6837    #[allow(clippy::too_many_arguments)]
6838    pub fn moe_gate_up_silu8_dev_q8_rows(
6839        &self,
6840        table: &CudaSlice<u64>,
6841        sel: &CudaSlice<i32>,
6842        aq: &CudaSlice<i8>,
6843        ad: &CudaSlice<f32>,
6844        t: usize,
6845        in_f: usize,
6846        n_ff: usize,
6847        n_used: usize,
6848        n_expert: usize,
6849        qt_g: i32,
6850        qt_u: i32,
6851        rb_g: usize,
6852        rb_u: usize,
6853        macros: &CudaSlice<f32>,
6854    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6855        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6856        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6857        let cfg = LaunchConfig {
6858            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6859            block_dim: (32, 1, 1),
6860            shared_mem_bytes: 0,
6861        };
6862        let (inf, nff, ne, nu, rbg, rbu) = (
6863            in_f as i32,
6864            n_ff as i32,
6865            n_expert as i32,
6866            n_used as i32,
6867            rb_g as i64,
6868            rb_u as i64,
6869        );
6870        let __s_b = self.gpu.stream();
6871        let mut b = __s_b.launch_builder(&f);
6872        b.arg(table)
6873            .arg(sel)
6874            .arg(aq)
6875            .arg(ad)
6876            .arg(&mut act)
6877            .arg(&inf)
6878            .arg(&nff)
6879            .arg(&ne)
6880            .arg(&qt_g)
6881            .arg(&qt_u)
6882            .arg(&rbg)
6883            .arg(&rbu)
6884            .arg(&nu)
6885            .arg(macros);
6886        unsafe {
6887            b.launch(cfg)?;
6888        }
6889        Ok(act)
6890    }
6891
6892    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6893    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6894    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6895    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6896    #[allow(clippy::too_many_arguments)]
6897    pub fn moe_down8_fma_dev_q8_rows(
6898        &self,
6899        table: &CudaSlice<u64>,
6900        sel: &CudaSlice<i32>,
6901        w: &CudaSlice<f32>,
6902        aq2: &CudaSlice<i8>,
6903        ad2: &CudaSlice<f32>,
6904        dst: &mut CudaSlice<f32>,
6905        t: usize,
6906        in_f: usize,
6907        out_f: usize,
6908        n_used: usize,
6909        n_expert: usize,
6910        qt: i32,
6911        rb: usize,
6912    ) -> Result<(), Box<dyn std::error::Error>> {
6913        assert!(
6914            in_f == 512 && n_used <= 8,
6915            "down rows twin is w8h2v shape-gated"
6916        );
6917        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6918        let cfg = LaunchConfig {
6919            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6920            block_dim: (32, n_used as u32, 1),
6921            shared_mem_bytes: 0,
6922        };
6923        let (inf, outf, nu, ne, rbi) = (
6924            in_f as i32,
6925            out_f as i32,
6926            n_used as i32,
6927            n_expert as i32,
6928            rb as i64,
6929        );
6930        let __s_b = self.gpu.stream();
6931        let mut b = __s_b.launch_builder(&f);
6932        b.arg(table)
6933            .arg(sel)
6934            .arg(w)
6935            .arg(aq2)
6936            .arg(ad2)
6937            .arg(dst)
6938            .arg(&inf)
6939            .arg(&outf)
6940            .arg(&nu)
6941            .arg(&ne)
6942            .arg(&qt)
6943            .arg(&rbi);
6944        unsafe {
6945            b.launch(cfg)?;
6946        }
6947        Ok(())
6948    }
6949
6950    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6951    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6952    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6953    #[allow(clippy::too_many_arguments)]
6954    pub fn moe_gate_up_silu8_dev_q8_csr(
6955        &self,
6956        table: &CudaSlice<u64>,
6957        sel: &CudaSlice<i32>,
6958        aq: &CudaSlice<i8>,
6959        ad: &CudaSlice<f32>,
6960        n_pairs: usize,
6961        in_f: usize,
6962        n_ff: usize,
6963        n_used: usize,
6964        n_expert: usize,
6965        qt_g: i32,
6966        qt_u: i32,
6967        rb_g: usize,
6968        rb_u: usize,
6969    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6970        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6971        // host gate guarantees qt_g == qt_u within a supported class.
6972        let f = if qt_g == crate::QT_NVFP4 {
6973            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6974        } else {
6975            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6976        };
6977        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6978        let cfg = LaunchConfig {
6979            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6980            block_dim: (32, 1, 1),
6981            shared_mem_bytes: 0,
6982        };
6983        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6984            in_f as i32,
6985            n_ff as i32,
6986            n_expert as i32,
6987            n_used as i32,
6988            n_pairs as i32,
6989            rb_g as i64,
6990            rb_u as i64,
6991        );
6992        let __s_b = self.gpu.stream();
6993        let mut b = __s_b.launch_builder(&f);
6994        b.arg(table)
6995            .arg(sel)
6996            .arg(aq)
6997            .arg(ad)
6998            .arg(&mut act)
6999            .arg(&inf)
7000            .arg(&nff)
7001            .arg(&ne)
7002            .arg(&qt_g)
7003            .arg(&qt_u)
7004            .arg(&rbg)
7005            .arg(&rbu)
7006            .arg(&nu)
7007            .arg(&npi);
7008        unsafe {
7009            b.launch(cfg)?;
7010        }
7011        Ok(act)
7012    }
7013
7014    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
7015    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
7016    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
7017    #[allow(clippy::too_many_arguments)]
7018    pub fn moe_down8_fma_dev_q8_variant(
7019        &self,
7020        variant: &str,
7021        table: &CudaSlice<u64>,
7022        sel: &cudarc::driver::CudaView<i32>,
7023        w: &cudarc::driver::CudaView<f32>,
7024        aq2: &CudaSlice<i8>,
7025        ad2: &CudaSlice<f32>,
7026        dst: &mut cudarc::driver::CudaViewMut<f32>,
7027        in_f: usize,
7028        out_f: usize,
7029        n_used: usize,
7030        n_expert: usize,
7031        qt: i32,
7032        rb: usize,
7033    ) -> Result<(), Box<dyn std::error::Error>> {
7034        let (inf, outf, nu, ne, rbi) = (
7035            in_f as i32,
7036            out_f as i32,
7037            n_used as i32,
7038            n_expert as i32,
7039            rb as i64,
7040        );
7041        let (f, cfg) = match variant {
7042            "w8h2" | "w8h2v" => (
7043                self.func(if variant == "w8h2" {
7044                    "moe_down8_fma_dev_q8_w8h2"
7045                } else {
7046                    "moe_down8_fma_dev_q8_w8h2v"
7047                }),
7048                LaunchConfig {
7049                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7050                    block_dim: (32, n_used as u32, 1),
7051                    shared_mem_bytes: 0,
7052                },
7053            ),
7054            "w8h2r2" | "w8h2r2v" => (
7055                self.func(if variant == "w8h2r2" {
7056                    "moe_down8_fma_dev_q8_w8h2r2"
7057                } else {
7058                    "moe_down8_fma_dev_q8_w8h2r2v"
7059                }),
7060                LaunchConfig {
7061                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7062                    block_dim: (32, n_used as u32, 1),
7063                    shared_mem_bytes: 0,
7064                },
7065            ),
7066            _ => (
7067                self.func("moe_down8_fma_dev_q8"),
7068                LaunchConfig {
7069                    grid_dim: (out_f as u32, 1, 1),
7070                    block_dim: (32, 1, 1),
7071                    shared_mem_bytes: 0,
7072                },
7073            ),
7074        };
7075        let __s_b = self.gpu.stream();
7076        let mut b = __s_b.launch_builder(&f);
7077        b.arg(table)
7078            .arg(sel)
7079            .arg(w)
7080            .arg(aq2)
7081            .arg(ad2)
7082            .arg(dst)
7083            .arg(&inf)
7084            .arg(&outf)
7085            .arg(&nu)
7086            .arg(&ne)
7087            .arg(&qt)
7088            .arg(&rbi);
7089        unsafe {
7090            b.launch(cfg)?;
7091        }
7092        Ok(())
7093    }
7094
7095    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
7096    #[allow(clippy::too_many_arguments)]
7097    pub fn moe_gate_up_silu8_dev_q8_variant(
7098        &self,
7099        variant: &str,
7100        table: &CudaSlice<u64>,
7101        sel: &cudarc::driver::CudaView<i32>,
7102        aq: &CudaSlice<i8>,
7103        ad: &CudaSlice<f32>,
7104        in_f: usize,
7105        n_ff: usize,
7106        n_used: usize,
7107        n_expert: usize,
7108        qt_g: i32,
7109        qt_u: i32,
7110        rb_g: usize,
7111        rb_u: usize,
7112    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7113        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7114        let (inf, nff, ne, rbg, rbu) = (
7115            in_f as i32,
7116            n_ff as i32,
7117            n_expert as i32,
7118            rb_g as i64,
7119            rb_u as i64,
7120        );
7121        let f = self.func(if variant == "v" {
7122            "moe_gate_up_silu8_dev_q8_v"
7123        } else {
7124            "moe_gate_up_silu8_dev_q8"
7125        });
7126        let cfg = LaunchConfig {
7127            grid_dim: (n_ff as u32, n_used as u32, 1),
7128            block_dim: (32, 1, 1),
7129            shared_mem_bytes: 0,
7130        };
7131        let __s_b = self.gpu.stream();
7132        let mut b = __s_b.launch_builder(&f);
7133        b.arg(table)
7134            .arg(sel)
7135            .arg(aq)
7136            .arg(ad)
7137            .arg(&mut act)
7138            .arg(&inf)
7139            .arg(&nff)
7140            .arg(&ne)
7141            .arg(&qt_g)
7142            .arg(&qt_u)
7143            .arg(&rbg)
7144            .arg(&rbu);
7145        unsafe {
7146            b.launch(cfg)?;
7147        }
7148        Ok(act)
7149    }
7150
7151    pub fn moe_gate_up_silu8_dev(
7152        &self,
7153        table: &CudaSlice<u64>,
7154        sel: &cudarc::driver::CudaView<i32>,
7155        x: &cudarc::driver::CudaView<f32>,
7156        in_f: usize,
7157        n_ff: usize,
7158        n_used: usize,
7159        n_expert: usize,
7160        qt_g: i32,
7161        qt_u: i32,
7162        rb_g: usize,
7163        rb_u: usize,
7164        macros: &CudaSlice<f32>,
7165    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7166        let f = self.func("moe_gate_up_silu8_dev");
7167        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
7168        let cfg = LaunchConfig {
7169            grid_dim: (n_ff as u32, n_used as u32, 1),
7170            block_dim: (256, 1, 1),
7171            shared_mem_bytes: 0,
7172        };
7173        let (inf, nff, ne, rbg, rbu) = (
7174            in_f as i32,
7175            n_ff as i32,
7176            n_expert as i32,
7177            rb_g as i64,
7178            rb_u as i64,
7179        );
7180        let __s_b = self.gpu.stream();
7181        let mut b = __s_b.launch_builder(&f);
7182        b.arg(table)
7183            .arg(sel)
7184            .arg(x)
7185            .arg(&mut act)
7186            .arg(&inf)
7187            .arg(&nff)
7188            .arg(&ne)
7189            .arg(&qt_g)
7190            .arg(&qt_u)
7191            .arg(&rbg)
7192            .arg(&rbu)
7193            .arg(macros);
7194        unsafe {
7195            b.launch(cfg)?;
7196        }
7197        Ok(act)
7198    }
7199
7200    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
7201    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
7202    #[allow(clippy::too_many_arguments)]
7203    pub fn moe_down8_fma_dev(
7204        &self,
7205        table: &CudaSlice<u64>,
7206        sel: &cudarc::driver::CudaView<i32>,
7207        w: &cudarc::driver::CudaView<f32>,
7208        act: &CudaSlice<f32>,
7209        dst: &mut cudarc::driver::CudaViewMut<f32>,
7210        in_f: usize,
7211        out_f: usize,
7212        n_used: usize,
7213        n_expert: usize,
7214        qt: i32,
7215        rb: usize,
7216    ) -> Result<(), Box<dyn std::error::Error>> {
7217        let f = self.func("moe_down8_fma_dev");
7218        let cfg = LaunchConfig {
7219            grid_dim: (out_f as u32, 1, 1),
7220            block_dim: (256, 1, 1),
7221            shared_mem_bytes: 0,
7222        };
7223        let (inf, outf, nu, ne, rbv) = (
7224            in_f as i32,
7225            out_f as i32,
7226            n_used as i32,
7227            n_expert as i32,
7228            rb as i64,
7229        );
7230        let __s_b = self.gpu.stream();
7231        let mut b = __s_b.launch_builder(&f);
7232        b.arg(table)
7233            .arg(sel)
7234            .arg(w)
7235            .arg(act)
7236            .arg(dst)
7237            .arg(&inf)
7238            .arg(&outf)
7239            .arg(&nu)
7240            .arg(&ne)
7241            .arg(&qt)
7242            .arg(&rbv);
7243        unsafe {
7244            b.launch(cfg)?;
7245        }
7246        Ok(())
7247    }
7248
7249    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
7250    pub fn axpy_into(
7251        &self,
7252        src: &CudaSlice<f32>,
7253        alpha: f32,
7254        dst: &mut cudarc::driver::CudaViewMut<f32>,
7255        n: usize,
7256    ) -> Result<(), Box<dyn std::error::Error>> {
7257        let f = self.func("axpy_f32");
7258        let cfg = LaunchConfig::for_num_elems(n as u32);
7259        let (a, ni) = (alpha, n as i32);
7260        let __s_b = self.gpu.stream();
7261        let mut b = __s_b.launch_builder(&f);
7262        b.arg(src).arg(dst).arg(&a).arg(&ni);
7263        unsafe {
7264            b.launch(cfg)?;
7265        }
7266        Ok(())
7267    }
7268
7269    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
7270    pub fn axpy_host_into(
7271        &self,
7272        src: &cudarc::driver::CudaView<'_, f32>,
7273        alpha: f32,
7274        dst: &mut cudarc::driver::CudaViewMut<f32>,
7275        n: usize,
7276    ) -> Result<(), Box<dyn std::error::Error>> {
7277        let f = self.func("axpy_host_f32");
7278        let cfg = LaunchConfig::for_num_elems(n as u32);
7279        let (a, ni) = (alpha, n as i32);
7280        let __s_b = self.gpu.stream();
7281        let mut b = __s_b.launch_builder(&f);
7282        b.arg(src).arg(dst).arg(&a).arg(&ni);
7283        unsafe {
7284            b.launch(cfg)?;
7285        }
7286        Ok(())
7287    }
7288
7289    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
7290    pub fn add_scaled_rows(
7291        &self,
7292        src: &CudaSlice<f32>,
7293        scale: &CudaSlice<f32>,
7294        dst: &mut CudaSlice<f32>,
7295        ncols: usize,
7296        nrows: usize,
7297    ) -> Result<(), Box<dyn std::error::Error>> {
7298        let f = self.func("add_scaled_rows_f32");
7299        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7300        let (nc, nr) = (ncols as i32, nrows as i32);
7301        let __s_b = self.gpu.stream();
7302        let mut b = __s_b.launch_builder(&f);
7303        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7304        unsafe {
7305            b.launch(cfg)?;
7306        }
7307        Ok(())
7308    }
7309
7310    // ======== A2 GROUPED MoE PREFILL KERNELS ========
7311
7312    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
7313    pub fn gather_rows(
7314        &self,
7315        src: &CudaSlice<f32>,
7316        idx: &CudaSlice<i32>,
7317        dst: &mut CudaSlice<f32>,
7318        ncols: usize,
7319        m_e: usize,
7320    ) -> Result<(), Box<dyn std::error::Error>> {
7321        let f = self.func("gather_rows_f32");
7322        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7323        let (nc, me) = (ncols as i32, m_e as i32);
7324        let __s_b = self.gpu.stream();
7325        let mut b = __s_b.launch_builder(&f);
7326        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7327        unsafe {
7328            b.launch(cfg)?;
7329        }
7330        Ok(())
7331    }
7332
7333    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
7334    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
7335    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
7336    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
7337    pub fn scatter_slot(
7338        &self,
7339        src: &CudaSlice<f32>,
7340        tok_idx: &CudaSlice<i32>,
7341        slot_idx: &CudaSlice<i32>,
7342        weight: &CudaSlice<f32>,
7343        dst: &mut CudaSlice<f32>,
7344        wbuf: &mut CudaSlice<f32>,
7345        ncols: usize,
7346        n_used: usize,
7347        m_e: usize,
7348    ) -> Result<(), Box<dyn std::error::Error>> {
7349        let f = self.func("scatter_add_slot_f32");
7350        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7351        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7352        let __s_b = self.gpu.stream();
7353        let mut b = __s_b.launch_builder(&f);
7354        b.arg(src)
7355            .arg(tok_idx)
7356            .arg(slot_idx)
7357            .arg(weight)
7358            .arg(dst)
7359            .arg(wbuf)
7360            .arg(&nc)
7361            .arg(&nu)
7362            .arg(&me);
7363        unsafe {
7364            b.launch(cfg)?;
7365        }
7366        Ok(())
7367    }
7368
7369    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
7370    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
7371    /// Uses FMA for bit-identity with the sequential axpy path.
7372    pub fn reduce_slots(
7373        &self,
7374        slots: &CudaSlice<f32>,
7375        wbuf: &CudaSlice<f32>,
7376        dst: &mut CudaSlice<f32>,
7377        ncols: usize,
7378        n_used: usize,
7379        t: usize,
7380    ) -> Result<(), Box<dyn std::error::Error>> {
7381        let f = self.func("reduce_slots_f32");
7382        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7383        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7384        let __s_b = self.gpu.stream();
7385        let mut b = __s_b.launch_builder(&f);
7386        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7387        unsafe {
7388            b.launch(cfg)?;
7389        }
7390        Ok(())
7391    }
7392
7393    /// Canonical slot-order reduction with separately rounded multiply and add.
7394    ///
7395    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
7396    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
7397    pub fn reduce_slots_host(
7398        &self,
7399        slots: &CudaSlice<f32>,
7400        wbuf: &CudaSlice<f32>,
7401        dst: &mut CudaSlice<f32>,
7402        ncols: usize,
7403        n_used: usize,
7404        t: usize,
7405    ) -> Result<(), Box<dyn std::error::Error>> {
7406        let f = self.func("reduce_slots_host_f32");
7407        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7408        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7409        let __s_b = self.gpu.stream();
7410        let mut b = __s_b.launch_builder(&f);
7411        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7412        unsafe {
7413            b.launch(cfg)?;
7414        }
7415        Ok(())
7416    }
7417
7418    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
7419    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
7420    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
7421    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
7422    /// GPU time, ~half of it redundant re-quantization of the same row.
7423    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
7424    pub fn quantize_q8_1_view(
7425        &self,
7426        x: &cudarc::driver::CudaView<f32>,
7427        m: usize,
7428        in_f: usize,
7429    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7430        let f = self.func("quantize_q8_1");
7431        let nblk = in_f / 32;
7432        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7433        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7434        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7435        let (inf, mi) = (in_f as i32, m as i32);
7436        let __s_b = self.gpu.stream();
7437        let mut b = __s_b.launch_builder(&f);
7438        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7439        unsafe {
7440            b.launch(cfg)?;
7441        }
7442        Ok((q, d))
7443    }
7444
7445    pub fn quantize_q8_1(
7446        &self,
7447        x: &CudaSlice<f32>,
7448        m: usize,
7449        in_f: usize,
7450    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7451        let nblk = in_f / 32;
7452        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
7453        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
7454        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
7455        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7456        let (inf, mi) = (in_f as i32, m as i32);
7457        if Self::pdl_on() && Self::pdl_wb_on() {
7458            {
7459                use cudarc::driver::{DevicePtr, DevicePtrMut};
7460                let s = &self.gpu.stream();
7461                let (px, _g0) = x.device_ptr(s);
7462                let (pq, _g1) = q.device_ptr_mut(s);
7463                let (pd, _g2) = d.device_ptr_mut(s);
7464                let mut ps = [
7465                    &px as *const _ as *mut std::ffi::c_void,
7466                    &pq as *const _ as *mut _,
7467                    &pd as *const _ as *mut _,
7468                    &inf as *const _ as *mut _,
7469                    &mi as *const _ as *mut _,
7470                ];
7471                unsafe {
7472                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7473                }
7474            }
7475            return Ok((q, d));
7476        }
7477        let f = self.func("quantize_q8_1");
7478        let __s_b = self.gpu.stream();
7479        let mut b = __s_b.launch_builder(&f);
7480        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7481        unsafe {
7482            b.launch(cfg)?;
7483        }
7484        Ok((q, d))
7485    }
7486
7487    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
7488    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
7489    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
7490    pub fn quantize_fp4_act(
7491        &self,
7492        x: &CudaSlice<f32>,
7493        m: usize,
7494        in_f: usize,
7495    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7496        let f = self.func("quantize_fp4_act");
7497        let nb16 = in_f / 16;
7498        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
7499        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
7500        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
7501        let (inf, mi) = (in_f as i32, m as i32);
7502        let __s_b = self.gpu.stream();
7503        let mut b = __s_b.launch_builder(&f);
7504        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7505        unsafe {
7506            b.launch(cfg)?;
7507        }
7508        Ok((aq4, ad4))
7509    }
7510
7511    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
7512    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
7513    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
7514    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
7515    pub fn qmatvec_gemm_nvfp4_fp4(
7516        &self,
7517        bytes: &CudaSlice<u8>,
7518        x: &CudaSlice<f32>,
7519        m: usize,
7520        in_f: usize,
7521        out_f: usize,
7522        row_bytes: usize,
7523        scale: f32,
7524    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7525        assert!(
7526            in_f % 64 == 0,
7527            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7528        );
7529        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7530        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7531        if scale != 1.0 {
7532            self.scale_inplace(&mut y, scale, m * out_f)?;
7533        }
7534        Ok(y)
7535    }
7536
7537    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
7538    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
7539    fn fp4_gemm_launch(
7540        &self,
7541        bytes: &CudaSlice<u8>,
7542        aq4: &CudaSlice<u32>,
7543        ad4: &CudaSlice<u8>,
7544        m: usize,
7545        in_f: usize,
7546        out_f: usize,
7547        row_bytes: usize,
7548    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7549        let f = self.func("qmatvec_gemm_nvfp4_fp4");
7550        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7551        const BM: u32 = 64;
7552        const BN: u32 = 256;
7553        let cfg = LaunchConfig {
7554            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7555            block_dim: (32, 4, 1),
7556            shared_mem_bytes: 0,
7557        };
7558        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7559        let __s_b = self.gpu.stream();
7560        let mut b = __s_b.launch_builder(&f);
7561        b.arg(bytes)
7562            .arg(aq4)
7563            .arg(ad4)
7564            .arg(&mut y)
7565            .arg(&inf)
7566            .arg(&outf)
7567            .arg(&mi)
7568            .arg(&rb);
7569        unsafe {
7570            b.launch(cfg)?;
7571        }
7572        Ok(y)
7573    }
7574
7575    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
7576    pub fn qmatvec_gemm_nvfp4_fp4_raw(
7577        &self,
7578        bytes: &CudaSlice<u8>,
7579        x: &CudaSlice<f32>,
7580        m: usize,
7581        in_f: usize,
7582        out_f: usize,
7583        row_bytes: usize,
7584    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7585        assert!(
7586            in_f % 64 == 0,
7587            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7588        );
7589        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7590        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7591    }
7592
7593    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
7594    pub fn qmatvec_q8_0_fast(
7595        &self,
7596        w: &CudaSlice<u8>,
7597        x: &CudaSlice<f32>,
7598        m: usize,
7599        in_f: usize,
7600        out_f: usize,
7601        row_bytes: usize,
7602    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7603        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7604        let f = self.func("qmatvec_q8_0_dp4a");
7605        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7606        let cfg = LaunchConfig {
7607            grid_dim: (out_f as u32, m as u32, 1),
7608            block_dim: (128, 1, 1),
7609            shared_mem_bytes: 0,
7610        };
7611        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7612        let __s_b = self.gpu.stream();
7613        let mut b = __s_b.launch_builder(&f);
7614        b.arg(w)
7615            .arg(&aq)
7616            .arg(&ad)
7617            .arg(&mut y)
7618            .arg(&inf)
7619            .arg(&outf)
7620            .arg(&mi)
7621            .arg(&rb);
7622        unsafe {
7623            b.launch(cfg)?;
7624        }
7625        Ok(y)
7626    }
7627
7628    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7629    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7630    pub fn qmatvec_q4_K_fast(
7631        &self,
7632        w: &CudaSlice<u8>,
7633        x: &CudaSlice<f32>,
7634        m: usize,
7635        in_f: usize,
7636        out_f: usize,
7637        row_bytes: usize,
7638    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7639        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7640        let f = self.func("qmatvec_q4_K_dp4a");
7641        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7642        let cfg = LaunchConfig {
7643            grid_dim: (out_f as u32, m as u32, 1),
7644            block_dim: (128, 1, 1),
7645            shared_mem_bytes: 0,
7646        };
7647        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7648        let __s_b = self.gpu.stream();
7649        let mut b = __s_b.launch_builder(&f);
7650        b.arg(w)
7651            .arg(&aq)
7652            .arg(&ad)
7653            .arg(&mut y)
7654            .arg(&inf)
7655            .arg(&outf)
7656            .arg(&mi)
7657            .arg(&rb);
7658        unsafe {
7659            b.launch(cfg)?;
7660        }
7661        Ok(y)
7662    }
7663
7664    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7665    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7666    pub fn qmatvec_q6_K_fast(
7667        &self,
7668        w: &CudaSlice<u8>,
7669        x: &CudaSlice<f32>,
7670        m: usize,
7671        in_f: usize,
7672        out_f: usize,
7673        row_bytes: usize,
7674    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7675        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7676        let f = self.func("qmatvec_q6_K_dp4a");
7677        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7678        let cfg = LaunchConfig {
7679            grid_dim: (out_f as u32, m as u32, 1),
7680            block_dim: (128, 1, 1),
7681            shared_mem_bytes: 0,
7682        };
7683        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7684        let __s_b = self.gpu.stream();
7685        let mut b = __s_b.launch_builder(&f);
7686        b.arg(w)
7687            .arg(&aq)
7688            .arg(&ad)
7689            .arg(&mut y)
7690            .arg(&inf)
7691            .arg(&outf)
7692            .arg(&mi)
7693            .arg(&rb);
7694        unsafe {
7695            b.launch(cfg)?;
7696        }
7697        Ok(y)
7698    }
7699
7700    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7701    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7702    pub fn qmatvec_q5_K_fast(
7703        &self,
7704        w: &CudaSlice<u8>,
7705        x: &CudaSlice<f32>,
7706        m: usize,
7707        in_f: usize,
7708        out_f: usize,
7709        row_bytes: usize,
7710    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7711        self.qmatvec_dp4a_named(
7712            "qmatvec_q5_K_dp4a",
7713            &w.slice(0..w.len()),
7714            x,
7715            m,
7716            in_f,
7717            out_f,
7718            row_bytes,
7719        )
7720    }
7721    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7722    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7723    pub fn qmatvec_q3_K_fast(
7724        &self,
7725        w: &CudaSlice<u8>,
7726        x: &CudaSlice<f32>,
7727        m: usize,
7728        in_f: usize,
7729        out_f: usize,
7730        row_bytes: usize,
7731    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7732        self.qmatvec_dp4a_named(
7733            "qmatvec_q3_K_dp4a",
7734            &w.slice(0..w.len()),
7735            x,
7736            m,
7737            in_f,
7738            out_f,
7739            row_bytes,
7740        )
7741    }
7742    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
7743    pub fn qmatvec_nvfp4_fast_rp(
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        assert!(
7753            in_f % 64 == 0,
7754            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7755        );
7756        self.qmatvec_dp4a_named(
7757            "qmatvec_nvfp4_dp4a_rp",
7758            &w.slice(0..w.len()),
7759            x,
7760            m,
7761            in_f,
7762            out_f,
7763            row_bytes,
7764        )
7765    }
7766    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
7767    pub fn qmatvec_nvfp4_fast(
7768        &self,
7769        w: &cudarc::driver::CudaView<'_, u8>,
7770        x: &CudaSlice<f32>,
7771        m: usize,
7772        in_f: usize,
7773        out_f: usize,
7774        row_bytes: usize,
7775    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7776        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
7777        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
7778        assert!(
7779            in_f % 64 == 0,
7780            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7781        );
7782        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7783    }
7784    /// v2-layout twin of `qmatvec_nvfp4_fast` for the slot-major expert banks
7785    /// (MEMRA_NVFP4_BANK_V2) — bit-identical per row, coalesced reads.
7786    pub fn qmatvec_nvfp4_fast_v2(
7787        &self,
7788        w: &cudarc::driver::CudaView<'_, u8>,
7789        x: &CudaSlice<f32>,
7790        m: usize,
7791        in_f: usize,
7792        out_f: usize,
7793        row_bytes: usize,
7794    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7795        assert!(
7796            in_f % 64 == 0,
7797            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7798        );
7799        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7800    }
7801    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
7802    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7803    pub fn qmatvec_iq4_XS_fast(
7804        &self,
7805        w: &CudaSlice<u8>,
7806        x: &CudaSlice<f32>,
7807        m: usize,
7808        in_f: usize,
7809        out_f: usize,
7810        row_bytes: usize,
7811    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7812        self.qmatvec_dp4a_named(
7813            "qmatvec_iq4_XS_dp4a",
7814            &w.slice(0..w.len()),
7815            x,
7816            m,
7817            in_f,
7818            out_f,
7819            row_bytes,
7820        )
7821    }
7822
7823    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
7824    fn qmatvec_dp4a_named(
7825        &self,
7826        name: &str,
7827        w: &cudarc::driver::CudaView<'_, u8>,
7828        x: &CudaSlice<f32>,
7829        m: usize,
7830        in_f: usize,
7831        out_f: usize,
7832        row_bytes: usize,
7833    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7834        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7835        let f = self.func(name);
7836        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7837        let cfg = LaunchConfig {
7838            grid_dim: (out_f as u32, m as u32, 1),
7839            block_dim: (128, 1, 1),
7840            shared_mem_bytes: 0,
7841        };
7842        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7843        let __s_b = self.gpu.stream();
7844        let mut b = __s_b.launch_builder(&f);
7845        b.arg(w)
7846            .arg(&aq)
7847            .arg(&ad)
7848            .arg(&mut y)
7849            .arg(&inf)
7850            .arg(&outf)
7851            .arg(&mi)
7852            .arg(&rb);
7853        unsafe {
7854            b.launch(cfg)?;
7855        }
7856        Ok(y)
7857    }
7858
7859    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
7860    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
7861    /// its output); this entry exists so a routed-expert program can quantize one activation
7862    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
7863    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
7864    #[allow(clippy::too_many_arguments)]
7865    pub fn qmatvec_nvfp4_fast_prequant_into(
7866        &self,
7867        w: &CudaSlice<u8>,
7868        aq: &CudaSlice<i8>,
7869        ad: &CudaSlice<f32>,
7870        y: &mut CudaSlice<f32>,
7871        m: usize,
7872        in_f: usize,
7873        out_f: usize,
7874        row_bytes: usize,
7875    ) -> Result<(), Box<dyn std::error::Error>> {
7876        assert!(
7877            in_f % 64 == 0,
7878            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7879        );
7880        if y.len() < m * out_f {
7881            return Err(format!(
7882                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7883                y.len()
7884            )
7885            .into());
7886        }
7887        let f = self.func("qmatvec_nvfp4_dp4a");
7888        let cfg = LaunchConfig {
7889            grid_dim: (out_f as u32, m as u32, 1),
7890            block_dim: (128, 1, 1),
7891            shared_mem_bytes: 0,
7892        };
7893        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7894        let __s_b = self.gpu.stream();
7895        let mut b = __s_b.launch_builder(&f);
7896        b.arg(w)
7897            .arg(aq)
7898            .arg(ad)
7899            .arg(y)
7900            .arg(&inf)
7901            .arg(&outf)
7902            .arg(&mi)
7903            .arg(&rb);
7904        unsafe {
7905            b.launch(cfg)?;
7906        }
7907        Ok(())
7908    }
7909
7910    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
7911    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
7912    #[allow(clippy::too_many_arguments)]
7913    pub fn matvec_f32_qkv_into(
7914        &self,
7915        wq: &CudaSlice<f32>,
7916        wk: &CudaSlice<f32>,
7917        wv: &CudaSlice<f32>,
7918        wg: &CudaSlice<f32>,
7919        x: &CudaSlice<f32>,
7920        yq: &mut CudaSlice<f32>,
7921        yk: &mut CudaSlice<f32>,
7922        yv: &mut CudaSlice<f32>,
7923        yg: &mut CudaSlice<f32>,
7924        in_f: usize,
7925        out_q: usize,
7926        out_kv: usize,
7927        out_g: usize,
7928    ) -> Result<(), Box<dyn std::error::Error>> {
7929        if in_f % 4 != 0
7930            || wq.len() != out_q * in_f
7931            || wk.len() != out_kv * in_f
7932            || wv.len() != out_kv * in_f
7933            || wg.len() < out_g * in_f
7934            || x.len() < in_f
7935            || yq.len() < out_q
7936            || yk.len() < out_kv
7937            || yv.len() < out_kv
7938            || (out_g > 0 && yg.len() < out_g)
7939        {
7940            return Err(format!(
7941                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7942                 wq={} wk={} wv={} wg={}",
7943                wq.len(),
7944                wk.len(),
7945                wv.len(),
7946                wg.len()
7947            )
7948            .into());
7949        }
7950        let f = self.func("matvec_f32_qkv");
7951        let cfg = LaunchConfig {
7952            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7953            block_dim: (128, 1, 1),
7954            shared_mem_bytes: 0,
7955        };
7956        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7957        let __s_b = self.gpu.stream();
7958        let mut b = __s_b.launch_builder(&f);
7959        b.arg(wq)
7960            .arg(wk)
7961            .arg(wv)
7962            .arg(wg)
7963            .arg(x)
7964            .arg(yq)
7965            .arg(yk)
7966            .arg(yv)
7967            .arg(yg)
7968            .arg(&inf)
7969            .arg(&oq)
7970            .arg(&okv)
7971            .arg(&og);
7972        unsafe {
7973            b.launch(cfg)?;
7974        }
7975        Ok(())
7976    }
7977
7978    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
7979    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
7980    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
7981    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
7982    /// kernel — the batching only removes host launch latency.
7983    #[allow(clippy::too_many_arguments)]
7984    /// FUSION #2a: gate+up sweeps in one launch (v2 banks only; identical geometry both
7985    /// banks, caller-guarded). Per-row bit-identical to two qmatvec_nvfp4_sel_into calls.
7986    #[allow(clippy::too_many_arguments)]
7987    pub fn qmatvec_nvfp4_sel_gu_into(
7988        &self,
7989        gate_bank: &CudaSlice<u8>,
7990        up_bank: &CudaSlice<u8>,
7991        sel: &CudaSlice<i32>,
7992        aq: &CudaSlice<i8>,
7993        ad: &CudaSlice<f32>,
7994        yg: &mut CudaSlice<f32>,
7995        yu: &mut CudaSlice<f32>,
7996        n_sel: usize,
7997        in_f: usize,
7998        out_f: usize,
7999        row_bytes: usize,
8000        expert_stride: usize,
8001    ) -> Result<(), Box<dyn std::error::Error>> {
8002        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8003        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8004            return Err("NVFP4 gu sel geometry".into());
8005        }
8006        // MEMRA_SEL_GU_RPW=2|4: multirow twin (activation group read once, reused across
8007        // RPW rows' gate+up dots) — bit-identical per row, one block per RPW rows.
8008        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
8009        let rpw = *RPW.get_or_init(|| {
8010            std::env::var("MEMRA_SEL_GU_RPW")
8011                .ok()
8012                .and_then(|v| v.parse().ok())
8013                .filter(|r| *r == 2 || *r == 4)
8014                .unwrap_or(1)
8015        });
8016        let rpw = if out_f % rpw == 0 { rpw } else { 1 };
8017        // MEMRA_SEL_GU_WPR=1: warp-per-row (NUMERIC-CLASS — per-row reduction order changes;
8018        // acceptance is the argmax gate + battery, the QKV_FUSED/BF16_MMV class).
8019        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8020        let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
8021        let f = self.func(match (wpr, rpw) {
8022            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
8023            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
8024            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
8025            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
8026        });
8027        let cfg = LaunchConfig {
8028            grid_dim: if wpr {
8029                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
8030            } else if rpw == 1 {
8031                ((2 * out_f) as u32, n_sel as u32, 1)
8032            } else {
8033                ((out_f / rpw) as u32, n_sel as u32, 1)
8034            },
8035            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
8036            shared_mem_bytes: 0,
8037        };
8038        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8039        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8040        let (ars, adrs) = (0i64, 0i64);
8041        let __s_b = self.gpu.stream();
8042        let mut b = __s_b.launch_builder(&f);
8043        b.arg(gate_bank)
8044            .arg(up_bank)
8045            .arg(sel)
8046            .arg(aq)
8047            .arg(ad)
8048            .arg(yg)
8049            .arg(yu)
8050            .arg(&inf)
8051            .arg(&outf)
8052            .arg(&ns)
8053            .arg(&rb)
8054            .arg(&es)
8055            .arg(&ars)
8056            .arg(&adrs);
8057        unsafe {
8058            b.launch(cfg)?;
8059        }
8060        Ok(())
8061    }
8062
8063    /// MEMRA_SEL_DOWN8=1: the DOWN sweep and the route-weight combine in ONE launch
8064    /// (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm ported to the
8065    /// NVFP4 banks). Block = (32, n_sel): one warp per slot instead of one warp per
8066    /// (row, slot), and the n_sel x out_f partial buffer disappears. Bit-identical to
8067    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce
8068    /// tree, same slot-ordered chain. Requires the v2 banks and nsb <= 32 (the fit-block
8069    /// class the reduce identity is argued at).
8070    #[allow(clippy::too_many_arguments)]
8071    pub fn qmatvec_nvfp4_sel_down8_into(
8072        &self,
8073        bank: &CudaSlice<u8>,
8074        sel: &CudaSlice<i32>,
8075        aq: &CudaSlice<i8>,
8076        ad: &CudaSlice<f32>,
8077        route_w: &CudaSlice<f32>,
8078        md: &CudaSlice<f32>,
8079        dst: &mut CudaSlice<f32>,
8080        n_sel: usize,
8081        in_f: usize,
8082        out_f: usize,
8083        row_bytes: usize,
8084        expert_stride: usize,
8085        act_row_stride: usize,
8086        ad_row_stride: usize,
8087    ) -> Result<(), Box<dyn std::error::Error>> {
8088        if in_f % 64 != 0
8089            || n_sel == 0
8090            || n_sel > 8
8091            || (in_f >> 5) > 32
8092            || dst.len() < out_f
8093            || sel.len() < n_sel
8094            || route_w.len() < n_sel
8095        {
8096            return Err(format!(
8097                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
8098                dst.len()
8099            )
8100            .into());
8101        }
8102        if !crate::tp::nvfp4_bank_v2_on() {
8103            return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
8104        }
8105        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
8106        let cfg = LaunchConfig {
8107            grid_dim: (out_f as u32, 1, 1),
8108            block_dim: (32, n_sel as u32, 1),
8109            shared_mem_bytes: 0,
8110        };
8111        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8112        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8113        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8114        let __s_b = self.gpu.stream();
8115        let mut b = __s_b.launch_builder(&f);
8116        b.arg(bank)
8117            .arg(sel)
8118            .arg(aq)
8119            .arg(ad)
8120            .arg(route_w)
8121            .arg(md)
8122            .arg(dst)
8123            .arg(&inf)
8124            .arg(&outf)
8125            .arg(&ns)
8126            .arg(&rb)
8127            .arg(&es)
8128            .arg(&ars)
8129            .arg(&adrs);
8130        unsafe {
8131            b.launch(cfg)?;
8132        }
8133        Ok(())
8134    }
8135
8136    /// T-ROW twin of the down8 fusion (spec verify + batched serving MoE): one block per
8137    /// (output row, token row) = the exact t=1 down8 program per token — bit-identical
8138    /// per row to its own down8/axpy pair at any t.
8139    #[allow(clippy::too_many_arguments)]
8140    pub fn qmatvec_nvfp4_sel_down8_rows_into(
8141        &self,
8142        bank: &CudaSlice<u8>,
8143        sel: &CudaSlice<i32>,
8144        aq: &CudaSlice<i8>,
8145        ad: &CudaSlice<f32>,
8146        route_w: &CudaSlice<f32>,
8147        md: &CudaSlice<f32>,
8148        dst: &mut CudaSlice<f32>,
8149        t: usize,
8150        n_sel_col: usize,
8151        in_f: usize,
8152        out_f: usize,
8153        row_bytes: usize,
8154        expert_stride: usize,
8155        act_row_stride: usize,
8156        ad_row_stride: usize,
8157    ) -> Result<(), Box<dyn std::error::Error>> {
8158        let n_sel = t * n_sel_col;
8159        if in_f % 64 != 0
8160            || n_sel_col == 0
8161            || n_sel_col > 8
8162            || t == 0
8163            || t > 64
8164            || (in_f >> 5) > 32
8165            || dst.len() < t * out_f
8166            || sel.len() < n_sel
8167            || route_w.len() < n_sel
8168        {
8169            return Err(format!(
8170                "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
8171                dst.len()
8172            )
8173            .into());
8174        }
8175        if !crate::tp::nvfp4_bank_v2_on() {
8176            return Err(
8177                "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8178            );
8179        }
8180        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8181        let cfg = LaunchConfig {
8182            grid_dim: (out_f as u32, t as u32, 1),
8183            block_dim: (32, n_sel_col as u32, 1),
8184            shared_mem_bytes: 0,
8185        };
8186        let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8187        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8188        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8189        let __s_b = self.gpu.stream();
8190        let mut b = __s_b.launch_builder(&f);
8191        b.arg(bank)
8192            .arg(sel)
8193            .arg(aq)
8194            .arg(ad)
8195            .arg(route_w)
8196            .arg(md)
8197            .arg(dst)
8198            .arg(&inf)
8199            .arg(&outf)
8200            .arg(&nsc)
8201            .arg(&rb)
8202            .arg(&es)
8203            .arg(&ars)
8204            .arg(&adrs);
8205        unsafe {
8206            b.launch(cfg)?;
8207        }
8208        Ok(())
8209    }
8210
8211    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
8212    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
8213    #[allow(clippy::too_many_arguments)]
8214    pub fn qmatvec_nvfp4_sel_gu_ep_into(
8215        &self,
8216        gate_bank: &CudaSlice<u8>,
8217        up_bank: &CudaSlice<u8>,
8218        sel: &CudaSlice<i32>,
8219        aq: &CudaSlice<i8>,
8220        ad: &CudaSlice<f32>,
8221        yg: &mut CudaSlice<f32>,
8222        yu: &mut CudaSlice<f32>,
8223        n_sel: usize,
8224        in_f: usize,
8225        out_f: usize,
8226        row_bytes: usize,
8227        expert_stride: usize,
8228        owner: usize,
8229    ) -> Result<(), Box<dyn std::error::Error>> {
8230        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8231        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8232            return Err("NVFP4 gu ep geometry".into());
8233        }
8234        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8235        let cfg = LaunchConfig {
8236            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8237            block_dim: (128, 1, 1),
8238            shared_mem_bytes: 0,
8239        };
8240        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8241        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8242        let (ars, adrs) = (0i64, 0i64);
8243        let __s_b = self.gpu.stream();
8244        let mut b = __s_b.launch_builder(&f);
8245        b.arg(gate_bank)
8246            .arg(up_bank)
8247            .arg(sel)
8248            .arg(aq)
8249            .arg(ad)
8250            .arg(yg)
8251            .arg(yu)
8252            .arg(&inf)
8253            .arg(&outf)
8254            .arg(&ns)
8255            .arg(&rb)
8256            .arg(&es)
8257            .arg(&ars)
8258            .arg(&adrs)
8259            .arg(&own);
8260        unsafe {
8261            b.launch(cfg)?;
8262        }
8263        Ok(())
8264    }
8265
8266    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
8267    #[allow(clippy::too_many_arguments)]
8268    pub fn silu_mul_scaled_q8_1_sel_ep_into(
8269        &self,
8270        gate: &CudaSlice<f32>,
8271        up: &CudaSlice<f32>,
8272        gmac: &CudaSlice<f32>,
8273        umac: &CudaSlice<f32>,
8274        sel: &CudaSlice<i32>,
8275        limit: Option<f32>,
8276        out_q: &mut CudaSlice<i8>,
8277        out_d: &mut CudaSlice<f32>,
8278        n_per: usize,
8279        n_sel: usize,
8280        owner: usize,
8281    ) -> Result<(), Box<dyn std::error::Error>> {
8282        if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8283            return Err("NVFP4 silu ep geometry".into());
8284        }
8285        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8286        let warps = n_sel * n_per / 32;
8287        let cfg = LaunchConfig {
8288            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8289            block_dim: (128, 1, 1),
8290            shared_mem_bytes: 0,
8291        };
8292        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8293        let (lim, has) = match limit {
8294            Some(l) => (l, 1i32),
8295            None => (0.0f32, 0i32),
8296        };
8297        let __s_b = self.gpu.stream();
8298        let mut b = __s_b.launch_builder(&f);
8299        b.arg(gate)
8300            .arg(up)
8301            .arg(gmac)
8302            .arg(umac)
8303            .arg(sel)
8304            .arg(&lim)
8305            .arg(&has)
8306            .arg(out_q)
8307            .arg(out_d)
8308            .arg(&np)
8309            .arg(&ns)
8310            .arg(&own);
8311        unsafe {
8312            b.launch(cfg)?;
8313        }
8314        Ok(())
8315    }
8316
8317    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
8318    #[allow(clippy::too_many_arguments)]
8319    pub fn qmatvec_nvfp4_sel_down8_ep_into(
8320        &self,
8321        bank: &CudaSlice<u8>,
8322        sel: &CudaSlice<i32>,
8323        aq: &CudaSlice<i8>,
8324        ad: &CudaSlice<f32>,
8325        route_w: &CudaSlice<f32>,
8326        md: &CudaSlice<f32>,
8327        dst: &mut CudaSlice<f32>,
8328        n_sel: usize,
8329        in_f: usize,
8330        out_f: usize,
8331        row_bytes: usize,
8332        expert_stride: usize,
8333        act_row_stride: usize,
8334        ad_row_stride: usize,
8335        owner: usize,
8336    ) -> Result<(), Box<dyn std::error::Error>> {
8337        if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8338            return Err("NVFP4 down8 ep geometry".into());
8339        }
8340        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8341        let cfg = LaunchConfig {
8342            grid_dim: (out_f as u32, 1, 1),
8343            block_dim: (32, n_sel as u32, 1),
8344            shared_mem_bytes: 0,
8345        };
8346        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8347        let (rb, es) = (row_bytes as i64, expert_stride as i64);
8348        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8349        let __s_b = self.gpu.stream();
8350        let mut b = __s_b.launch_builder(&f);
8351        b.arg(bank)
8352            .arg(sel)
8353            .arg(aq)
8354            .arg(ad)
8355            .arg(route_w)
8356            .arg(md)
8357            .arg(dst)
8358            .arg(&inf)
8359            .arg(&outf)
8360            .arg(&ns)
8361            .arg(&rb)
8362            .arg(&es)
8363            .arg(&ars)
8364            .arg(&adrs)
8365            .arg(&own);
8366        unsafe {
8367            b.launch(cfg)?;
8368        }
8369        Ok(())
8370    }
8371
8372    pub fn qmatvec_nvfp4_sel_into(
8373        &self,
8374        bank: &CudaSlice<u8>,
8375        sel: &CudaSlice<i32>,
8376        aq: &CudaSlice<i8>,
8377        ad: &CudaSlice<f32>,
8378        y: &mut CudaSlice<f32>,
8379        n_sel: usize,
8380        in_f: usize,
8381        out_f: usize,
8382        row_bytes: usize,
8383        expert_stride: usize,
8384        act_row_stride: usize,
8385        ad_row_stride: usize,
8386    ) -> Result<(), Box<dyn std::error::Error>> {
8387        assert!(
8388            in_f % 64 == 0,
8389            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8390        );
8391        if y.len() < n_sel * out_f || sel.len() < n_sel {
8392            return Err(format!(
8393                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8394                y.len(),
8395                sel.len()
8396            )
8397            .into());
8398        }
8399        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
8400        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
8401        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
8402        // sequential-rows variant was flat). Default stays the single-row form.
8403        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
8404        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
8405        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8406        let mode = *MR.get_or_init(|| {
8407            if crate::tp::nvfp4_bank_v2_on() {
8408                3
8409            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8410                2
8411            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8412                1
8413            } else {
8414                0
8415            }
8416        });
8417        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8418        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
8419        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
8420        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
8421        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8422        let v2s = mode == 3
8423            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8424            && row_bytes % 16 == 0
8425            && in_f <= 4096;
8426        let f = match (mode, v2s) {
8427            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8428            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8429            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8430            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8431            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8432        };
8433        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
8434        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
8435        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
8436        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
8437        let nsb = in_f >> 5;
8438        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8439            32
8440        } else if mode == 1 {
8441            512
8442        } else {
8443            128
8444        };
8445        let cfg = LaunchConfig {
8446            grid_dim: (
8447                if v2s {
8448                    (out_f as u32).div_ceil(8)
8449                } else {
8450                    match mode {
8451                        2 => (out_f as u32).div_ceil(16),
8452                        1 => (out_f as u32).div_ceil(4),
8453                        _ => out_f as u32,
8454                    }
8455                },
8456                n_sel as u32,
8457                1,
8458            ),
8459            block_dim: (fit_block, 1, 1),
8460            shared_mem_bytes: 0,
8461        };
8462        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8463        let (rb, es, ars, adrs) = (
8464            row_bytes as i64,
8465            expert_stride as i64,
8466            act_row_stride as i64,
8467            ad_row_stride as i64,
8468        );
8469        let __s_b = self.gpu.stream();
8470        let mut b = __s_b.launch_builder(&f);
8471        b.arg(bank)
8472            .arg(sel)
8473            .arg(aq)
8474            .arg(ad)
8475            .arg(y)
8476            .arg(&inf)
8477            .arg(&outf)
8478            .arg(&ns)
8479            .arg(&rb)
8480            .arg(&es)
8481            .arg(&ars)
8482            .arg(&adrs);
8483        unsafe {
8484            b.launch(cfg)?;
8485        }
8486        Ok(())
8487    }
8488
8489    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
8490    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
8491    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
8492    /// takes the plain SiLU kernel.
8493    #[allow(clippy::too_many_arguments)]
8494    pub fn silu_mul_scaled_q8_1_sel_into(
8495        &self,
8496        gate: &CudaSlice<f32>,
8497        up: &CudaSlice<f32>,
8498        gmac: &CudaSlice<f32>,
8499        umac: &CudaSlice<f32>,
8500        sel: &CudaSlice<i32>,
8501        limit: Option<f32>,
8502        out_q: &mut CudaSlice<i8>,
8503        out_d: &mut CudaSlice<f32>,
8504        n_per: usize,
8505        n_sel: usize,
8506    ) -> Result<(), Box<dyn std::error::Error>> {
8507        let n = n_per * n_sel;
8508        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8509            return Err(format!(
8510                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8511                out_q.len(),
8512                out_d.len()
8513            )
8514            .into());
8515        }
8516        if let Some(limit) = limit {
8517            if limit <= 1e-6 {
8518                return Err(format!(
8519                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8520                )
8521                .into());
8522            }
8523            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8524            let cfg = LaunchConfig::for_num_elems(n as u32);
8525            let (np, ns) = (n_per as i32, n_sel as i32);
8526            let __s_b = self.gpu.stream();
8527            let mut b = __s_b.launch_builder(&f);
8528            b.arg(gate)
8529                .arg(up)
8530                .arg(gmac)
8531                .arg(umac)
8532                .arg(sel)
8533                .arg(&limit)
8534                .arg(out_q)
8535                .arg(out_d)
8536                .arg(&np)
8537                .arg(&ns);
8538            unsafe {
8539                b.launch(cfg)?;
8540            }
8541            return Ok(());
8542        }
8543        let f = self.func("silu_mul_scaled_q8_1_sel");
8544        let cfg = LaunchConfig::for_num_elems(n as u32);
8545        let (np, ns) = (n_per as i32, n_sel as i32);
8546        let __s_b = self.gpu.stream();
8547        let mut b = __s_b.launch_builder(&f);
8548        b.arg(gate)
8549            .arg(up)
8550            .arg(gmac)
8551            .arg(umac)
8552            .arg(sel)
8553            .arg(out_q)
8554            .arg(out_d)
8555            .arg(&np)
8556            .arg(&ns);
8557        unsafe {
8558            b.launch(cfg)?;
8559        }
8560        Ok(())
8561    }
8562
8563    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8564        Ok(self.gpu.stream().clone_htod(v)?)
8565    }
8566    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8567        Ok(self.gpu.stream().clone_htod(v)?)
8568    }
8569    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
8570    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8571        Ok(self.gpu.stream().clone_htod(v)?)
8572    }
8573    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8574        Ok(self.gpu.stream().clone_htod(v)?)
8575    }
8576    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
8577    pub fn dtoh_view(
8578        &self,
8579        d: &cudarc::driver::CudaView<f32>,
8580    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8581        let v = self.gpu.stream().clone_dtoh(d)?;
8582        self.gpu.stream().synchronize()?;
8583        Ok(v)
8584    }
8585    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8586        let v = self.gpu.stream().clone_dtoh(d)?;
8587        self.gpu.stream().synchronize()?;
8588        Ok(v)
8589    }
8590    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
8591    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
8592    /// issuing them together avoids a second stream synchronization in every trunk layer.
8593    pub fn dtoh_pair(
8594        &self,
8595        a: &CudaSlice<f32>,
8596        b: &CudaSlice<f32>,
8597    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8598        let av = self.gpu.stream().clone_dtoh(a)?;
8599        let bv = self.gpu.stream().clone_dtoh(b)?;
8600        self.gpu.stream().synchronize()?;
8601        Ok((av, bv))
8602    }
8603    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
8604    /// cross a shape-sensitive host boundary.
8605    pub fn dtoh_pair_views(
8606        &self,
8607        a: &cudarc::driver::CudaView<f32>,
8608        b: &cudarc::driver::CudaView<f32>,
8609    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8610        let av = self.gpu.stream().clone_dtoh(a)?;
8611        let bv = self.gpu.stream().clone_dtoh(b)?;
8612        self.gpu.stream().synchronize()?;
8613        Ok((av, bv))
8614    }
8615    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
8616    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8617        let v = self.gpu.stream().clone_dtoh(d)?;
8618        self.gpu.stream().synchronize()?;
8619        Ok(v)
8620    }
8621    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
8622    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8623        let v = self.gpu.stream().clone_dtoh(d)?;
8624        self.gpu.stream().synchronize()?;
8625        Ok(v)
8626    }
8627    pub fn dtoh_u8_view(
8628        &self,
8629        d: &cudarc::driver::CudaView<u8>,
8630    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8631        let v = self.gpu.stream().clone_dtoh(d)?;
8632        self.gpu.stream().synchronize()?;
8633        Ok(v)
8634    }
8635    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8636        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8637        self.keep_if_capturing(&s);
8638        Ok(s)
8639    }
8640
8641    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
8642    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
8643    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
8644    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
8645    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
8646    /// back (or kept resident for graph replay). Returns the device token buffer.
8647    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
8648    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
8649    pub fn prob_of_token_device(
8650        &self,
8651        logits: &CudaSlice<f32>,
8652        tok: &CudaSlice<u32>,
8653        n_vocab: usize,
8654    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8655        let nb = ARGMAX_NB;
8656        let mut part = self.alloc_uninit::<f32>(nb)?;
8657        let mut p = self.alloc_uninit::<f32>(1)?;
8658        let f1 = self.func("prob_of_token_partial_f32");
8659        let cfg1 = LaunchConfig {
8660            grid_dim: (nb as u32, 1, 1),
8661            block_dim: (256, 1, 1),
8662            shared_mem_bytes: 0,
8663        };
8664        let nv = n_vocab as i32;
8665        let __s_b1 = self.gpu.stream();
8666        let mut b1 = __s_b1.launch_builder(&f1);
8667        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8668        unsafe {
8669            b1.launch(cfg1)?;
8670        }
8671        let f2 = self.func("prob_of_token_final_f32");
8672        let cfg2 = LaunchConfig {
8673            grid_dim: (1, 1, 1),
8674            block_dim: (256, 1, 1),
8675            shared_mem_bytes: 0,
8676        };
8677        let nbi = nb as i32;
8678        let __s_b2 = self.gpu.stream();
8679        let mut b2 = __s_b2.launch_builder(&f2);
8680        b2.arg(&part).arg(&mut p).arg(&nbi);
8681        unsafe {
8682            b2.launch(cfg2)?;
8683        }
8684        Ok(p)
8685    }
8686
8687    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
8688    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
8689    /// where the host reads the p-min confidence between replays. Same kernels, same math.
8690    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
8691    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
8692    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
8693    pub fn prob_of_token_device_col(
8694        &self,
8695        logits: &CudaSlice<f32>,
8696        tok_all: &CudaSlice<u32>,
8697        tok_idx: usize,
8698        p_out: &mut CudaSlice<f32>,
8699        p_idx: usize,
8700        n_vocab: usize,
8701    ) -> Result<(), Box<dyn std::error::Error>> {
8702        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8703        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8704        let nb = ARGMAX_NB;
8705        let mut part = self.alloc_uninit::<f32>(nb)?;
8706        let f1 = self.func("prob_of_token_partial_f32");
8707        let cfg1 = LaunchConfig {
8708            grid_dim: (nb as u32, 1, 1),
8709            block_dim: (256, 1, 1),
8710            shared_mem_bytes: 0,
8711        };
8712        let nv = n_vocab as i32;
8713        let __s_b1 = self.gpu.stream();
8714        let mut b1 = __s_b1.launch_builder(&f1);
8715        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8716        unsafe {
8717            b1.launch(cfg1)?;
8718        }
8719        let f2 = self.func("prob_of_token_final_f32");
8720        let cfg2 = LaunchConfig {
8721            grid_dim: (1, 1, 1),
8722            block_dim: (256, 1, 1),
8723            shared_mem_bytes: 0,
8724        };
8725        let nbi = nb as i32;
8726        let __s_b2 = self.gpu.stream();
8727        let mut b2 = __s_b2.launch_builder(&f2);
8728        b2.arg(&part).arg(&mut p_v).arg(&nbi);
8729        unsafe {
8730            b2.launch(cfg2)?;
8731        }
8732        Ok(())
8733    }
8734
8735    pub fn prob_of_token_device_into(
8736        &self,
8737        logits: &CudaSlice<f32>,
8738        tok: &CudaSlice<u32>,
8739        p_out: &mut CudaSlice<f32>,
8740        n_vocab: usize,
8741    ) -> Result<(), Box<dyn std::error::Error>> {
8742        let nb = ARGMAX_NB;
8743        let mut part = self.alloc_uninit::<f32>(nb)?;
8744        let f1 = self.func("prob_of_token_partial_f32");
8745        let cfg1 = LaunchConfig {
8746            grid_dim: (nb as u32, 1, 1),
8747            block_dim: (256, 1, 1),
8748            shared_mem_bytes: 0,
8749        };
8750        let nv = n_vocab as i32;
8751        let __s_b1 = self.gpu.stream();
8752        let mut b1 = __s_b1.launch_builder(&f1);
8753        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8754        unsafe {
8755            b1.launch(cfg1)?;
8756        }
8757        let f2 = self.func("prob_of_token_final_f32");
8758        let cfg2 = LaunchConfig {
8759            grid_dim: (1, 1, 1),
8760            block_dim: (256, 1, 1),
8761            shared_mem_bytes: 0,
8762        };
8763        let nbi = nb as i32;
8764        let __s_b2 = self.gpu.stream();
8765        let mut b2 = __s_b2.launch_builder(&f2);
8766        b2.arg(&part).arg(p_out).arg(&nbi);
8767        unsafe {
8768            b2.launch(cfg2)?;
8769        }
8770        Ok(())
8771    }
8772
8773    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
8774    /// (graph-constant params, device-varying index). Capture-safe.
8775    pub fn u32_hist_append(
8776        &self,
8777        tok: &CudaSlice<u32>,
8778        hist: &mut CudaSlice<u32>,
8779        idx: &mut CudaSlice<i32>,
8780    ) -> Result<(), Box<dyn std::error::Error>> {
8781        let f = self.func("u32_hist_append");
8782        let cfg = LaunchConfig {
8783            grid_dim: (1, 1, 1),
8784            block_dim: (32, 1, 1),
8785            shared_mem_bytes: 0,
8786        };
8787        let __s_b = self.gpu.stream();
8788        let mut b = __s_b.launch_builder(&f);
8789        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8790        unsafe {
8791            b.launch(cfg)?;
8792        }
8793        Ok(())
8794    }
8795
8796    pub fn argmax_token_device(
8797        &self,
8798        logits: &CudaSlice<f32>,
8799        n_vocab: usize,
8800    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8801        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8802        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8803        Ok(tok)
8804    }
8805    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
8806    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
8807    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
8808    /// pointer is baked once and the token id never round-trips to host inside steady state. The
8809    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
8810    /// captured passes bake fixed addresses.
8811    pub fn argmax_token_device_into(
8812        &self,
8813        logits: &CudaSlice<f32>,
8814        tok: &mut CudaSlice<u32>,
8815        n_vocab: usize,
8816    ) -> Result<(), Box<dyn std::error::Error>> {
8817        let nb = ARGMAX_NB;
8818        let f1 = self.func("argmax_partial_f32");
8819        let f2 = self.func("argmax_final_f32");
8820        let mut guard = self.argmax_partials.lock().unwrap();
8821        if guard.is_none() {
8822            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8823            // buffers carry no cudarc events (illegal inside capture).
8824            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8825            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8826            *guard = Some((pv, pi));
8827        }
8828        let (part_v, part_i) = guard.as_mut().unwrap();
8829        let nv = n_vocab as i32;
8830        let nbi = nb as i32;
8831        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8832        let cfg1 = LaunchConfig {
8833            grid_dim: (nb as u32, 1, 1),
8834            block_dim: (256, 1, 1),
8835            shared_mem_bytes: 0,
8836        };
8837        let __s_b1 = self.gpu.stream();
8838        let mut b1 = __s_b1.launch_builder(&f1);
8839        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8840        unsafe {
8841            b1.launch(cfg1)?;
8842        }
8843        // pass 2: one block reduces NB partials -> token_out[0].
8844        let cfg2 = LaunchConfig {
8845            grid_dim: (1, 1, 1),
8846            block_dim: (256, 1, 1),
8847            shared_mem_bytes: 0,
8848        };
8849        let __s_b2 = self.gpu.stream();
8850        let mut b2 = __s_b2.launch_builder(&f2);
8851        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8852        unsafe {
8853            b2.launch(cfg2)?;
8854        }
8855        Ok(())
8856    }
8857    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8858    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8859    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8860    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8861    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8862    pub fn argmax_token_device_col(
8863        &self,
8864        logits: &CudaSlice<f32>,
8865        col: usize,
8866        n_vocab: usize,
8867        toks: &mut CudaSlice<u32>,
8868        out_idx: usize,
8869    ) -> Result<(), Box<dyn std::error::Error>> {
8870        let nb = ARGMAX_NB;
8871        let f1 = self.func("argmax_partial_f32");
8872        let f2 = self.func("argmax_final_f32");
8873        let mut guard = self.argmax_partials.lock().unwrap();
8874        if guard.is_none() {
8875            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8876            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8877            *guard = Some((pv, pi));
8878        }
8879        let (part_v, part_i) = guard.as_mut().unwrap();
8880        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8881        let nv = n_vocab as i32;
8882        let nbi = nb as i32;
8883        let cfg1 = LaunchConfig {
8884            grid_dim: (nb as u32, 1, 1),
8885            block_dim: (256, 1, 1),
8886            shared_mem_bytes: 0,
8887        };
8888        let __s_b1 = self.gpu.stream();
8889        let mut b1 = __s_b1.launch_builder(&f1);
8890        b1.arg(&col_view)
8891            .arg(&mut *part_v)
8892            .arg(&mut *part_i)
8893            .arg(&nv);
8894        unsafe {
8895            b1.launch(cfg1)?;
8896        }
8897        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8898        let cfg2 = LaunchConfig {
8899            grid_dim: (1, 1, 1),
8900            block_dim: (256, 1, 1),
8901            shared_mem_bytes: 0,
8902        };
8903        let __s_b2 = self.gpu.stream();
8904        let mut b2 = __s_b2.launch_builder(&f2);
8905        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8906        unsafe {
8907            b2.launch(cfg2)?;
8908        }
8909        Ok(())
8910    }
8911    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8912    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8913        Ok(self.gpu.stream().clone_htod(v)?)
8914    }
8915    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8916        let v = self.gpu.stream().clone_dtoh(d)?;
8917        self.gpu.stream().synchronize()?;
8918        Ok(v)
8919    }
8920
8921    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8922        let v = self.gpu.stream().clone_dtoh(d)?;
8923        self.gpu.stream().synchronize()?;
8924        Ok(v)
8925    }
8926    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8927    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8928    /// contents change every step, the address must not, so a captured graph can read it).
8929    pub fn htod_u32_into(
8930        &self,
8931        dst: &mut CudaSlice<u32>,
8932        src: &[u32],
8933    ) -> Result<(), Box<dyn std::error::Error>> {
8934        let mut view = dst.slice_mut(0..src.len());
8935        self.gpu.stream().memcpy_htod(src, &mut view)?;
8936        Ok(())
8937    }
8938
8939    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8940    /// table without changing the device address its reconcile kernel consumes.
8941    pub fn htod_i32_into(
8942        &self,
8943        dst: &mut CudaSlice<i32>,
8944        src: &[i32],
8945    ) -> Result<(), Box<dyn std::error::Error>> {
8946        let mut view = dst.slice_mut(0..src.len());
8947        self.gpu.stream().memcpy_htod(src, &mut view)?;
8948        Ok(())
8949    }
8950
8951    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8952        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8953        self.keep_if_capturing(&s);
8954        Ok(s)
8955    }
8956    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8957    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8958    pub fn embed_gather_device_into(
8959        &self,
8960        embd: &CudaSlice<u8>,
8961        token_d: &CudaSlice<u32>,
8962        x_out: &mut CudaSlice<f32>,
8963        n_embd: usize,
8964        qtype: i32,
8965        row_bytes: usize,
8966    ) -> Result<(), Box<dyn std::error::Error>> {
8967        let f = self.func("embed_gather_u32");
8968        let cfg = LaunchConfig {
8969            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8970            block_dim: (256, 1, 1),
8971            shared_mem_bytes: 0,
8972        };
8973        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8974        let __s_b = self.gpu.stream();
8975        let mut b = __s_b.launch_builder(&f);
8976        b.arg(embd)
8977            .arg(token_d)
8978            .arg(x_out)
8979            .arg(&ne)
8980            .arg(&qt)
8981            .arg(&rb);
8982        unsafe {
8983            b.launch(cfg)?;
8984        }
8985        Ok(())
8986    }
8987    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
8988    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8989        let v = self.gpu.stream().clone_dtoh(d)?;
8990        self.gpu.stream().synchronize()?;
8991        Ok(v[0])
8992    }
8993    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
8994    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
8995    /// the counter value after the throwaway capture warmups corrupt it.
8996    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
8997    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
8998    /// copy (fine at stream-idle boundaries, poison mid-round).
8999    pub fn i32_set_k(
9000        &self,
9001        dst: &mut CudaSlice<i32>,
9002        v: i32,
9003    ) -> Result<(), Box<dyn std::error::Error>> {
9004        let f = self.func("i32_set_k");
9005        let cfg = LaunchConfig {
9006            grid_dim: (1, 1, 1),
9007            block_dim: (1, 1, 1),
9008            shared_mem_bytes: 0,
9009        };
9010        let idx = 0i32;
9011        let __s_b = self.gpu.stream();
9012        let mut b = __s_b.launch_builder(&f);
9013        b.arg(dst).arg(&v).arg(&idx);
9014        unsafe {
9015            b.launch(cfg)?;
9016        }
9017        Ok(())
9018    }
9019
9020    pub fn set_i32_one(
9021        &self,
9022        d: &mut CudaSlice<i32>,
9023        v: i32,
9024    ) -> Result<(), Box<dyn std::error::Error>> {
9025        self.gpu.stream().memcpy_htod(&[v], d)?;
9026        Ok(())
9027    }
9028    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
9029    /// during priming / capture-state restore.
9030    pub fn set_u32_one(
9031        &self,
9032        d: &mut CudaSlice<u32>,
9033        v: u32,
9034    ) -> Result<(), Box<dyn std::error::Error>> {
9035        self.gpu.stream().memcpy_htod(&[v], d)?;
9036        Ok(())
9037    }
9038    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
9039    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9040        let v = self.gpu.stream().clone_dtoh(d)?;
9041        self.gpu.stream().synchronize()?;
9042        Ok(v[0])
9043    }
9044    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
9045    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9046        Ok(self.gpu.stream().clone_htod(bytes)?)
9047    }
9048    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
9049    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
9050    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
9051    pub fn embed_gather_device(
9052        &self,
9053        embd: &CudaSlice<u8>,
9054        token_d: &CudaSlice<u32>,
9055        n_embd: usize,
9056        qtype: i32,
9057        row_bytes: usize,
9058    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9059        let f = self.func("embed_gather_u32");
9060        let mut x = self.alloc_uninit::<f32>(n_embd)?;
9061        let cfg = LaunchConfig {
9062            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9063            block_dim: (256, 1, 1),
9064            shared_mem_bytes: 0,
9065        };
9066        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9067        let __s_b = self.gpu.stream();
9068        let mut b = __s_b.launch_builder(&f);
9069        b.arg(embd)
9070            .arg(token_d)
9071            .arg(&mut x)
9072            .arg(&ne)
9073            .arg(&qt)
9074            .arg(&rb);
9075        unsafe {
9076            b.launch(cfg)?;
9077        }
9078        Ok(x)
9079    }
9080
9081    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
9082    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
9083    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
9084    pub fn embed_gather_device_t(
9085        &self,
9086        embd: &CudaSlice<u8>,
9087        tokens: &[u32],
9088        n_embd: usize,
9089        qtype: i32,
9090        row_bytes: usize,
9091    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9092        let t = tokens.len();
9093        let tok_d = self.gpu.stream().clone_htod(tokens)?;
9094        let f = self.func("embed_gather_u32_t");
9095        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9096        let cfg = LaunchConfig {
9097            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9098            block_dim: (256, 1, 1),
9099            shared_mem_bytes: 0,
9100        };
9101        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9102        let __s_b = self.gpu.stream();
9103        let mut b = __s_b.launch_builder(&f);
9104        b.arg(embd)
9105            .arg(&tok_d)
9106            .arg(&mut x)
9107            .arg(&ne)
9108            .arg(&qt)
9109            .arg(&rb)
9110            .arg(&ti);
9111        unsafe {
9112            b.launch(cfg)?;
9113        }
9114        Ok(x)
9115    }
9116
9117    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
9118    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
9119    /// as embed_gather_device_t — bit-identical rows.
9120    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
9121    pub fn embed_gather_device_tv(
9122        &self,
9123        embd: &CudaSlice<u8>,
9124        tok_v: &cudarc::driver::CudaView<u32>,
9125        t: usize,
9126        n_embd: usize,
9127        qtype: i32,
9128        row_bytes: usize,
9129    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9130        let f = self.func("embed_gather_u32_t");
9131        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9132        let cfg = LaunchConfig {
9133            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9134            block_dim: (256, 1, 1),
9135            shared_mem_bytes: 0,
9136        };
9137        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9138        let __s_b = self.gpu.stream();
9139        let mut b = __s_b.launch_builder(&f);
9140        b.arg(embd)
9141            .arg(tok_v)
9142            .arg(&mut x)
9143            .arg(&ne)
9144            .arg(&qt)
9145            .arg(&rb)
9146            .arg(&ti);
9147        unsafe {
9148            b.launch(cfg)?;
9149        }
9150        Ok(x)
9151    }
9152
9153    pub fn embed_gather_device_td(
9154        &self,
9155        embd: &CudaSlice<u8>,
9156        tok_d: &CudaSlice<u32>,
9157        t: usize,
9158        n_embd: usize,
9159        qtype: i32,
9160        row_bytes: usize,
9161    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9162        let f = self.func("embed_gather_u32_t");
9163        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9164        let cfg = LaunchConfig {
9165            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9166            block_dim: (256, 1, 1),
9167            shared_mem_bytes: 0,
9168        };
9169        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9170        let __s_b = self.gpu.stream();
9171        let mut b = __s_b.launch_builder(&f);
9172        b.arg(embd)
9173            .arg(tok_d)
9174            .arg(&mut x)
9175            .arg(&ne)
9176            .arg(&qt)
9177            .arg(&rb)
9178            .arg(&ti);
9179        unsafe {
9180            b.launch(cfg)?;
9181        }
9182        Ok(x)
9183    }
9184
9185    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
9186    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
9187    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
9188    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
9189    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
9190    #[inline]
9191    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
9192    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9193        if self
9194            .capture_keep_on
9195            .load(std::sync::atomic::Ordering::Relaxed)
9196        {
9197            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9198        }
9199    }
9200
9201    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9202        &self,
9203        n: usize,
9204    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9205        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9206        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
9207        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
9208        // not cover engine-internal buffers). Debug-only: massive launch overhead.
9209        {
9210            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9211            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9212                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
9213                use cudarc::driver::DevicePtrMut;
9214                let n_bytes = s.len() * std::mem::size_of::<T>();
9215                let stream = self.gpu.stream();
9216                let (p_, _g) = s.device_ptr_mut(&stream);
9217                unsafe {
9218                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9219                        .result()?;
9220                }
9221            }
9222        }
9223        self.keep_if_capturing(&s);
9224        Ok(s)
9225    }
9226
9227    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
9228    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
9229    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
9230    /// consumers alloc through this (m=1 decode arms).
9231    pub fn uninit_q8_pair(
9232        &self,
9233        n: usize,
9234    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9235        Ok((
9236            self.alloc_uninit::<i8>(n)?,
9237            self.alloc_uninit::<f32>(n / 32)?,
9238        ))
9239    }
9240
9241    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9242        self.alloc_uninit::<f32>(n)
9243    }
9244
9245    /// i8 uninitialized scratch (same contract as `uninit`).
9246    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9247        self.alloc_uninit::<i8>(n)
9248    }
9249
9250    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
9251    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
9252    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
9253    #[allow(clippy::too_many_arguments)]
9254    pub fn rms_norm3(
9255        &self,
9256        x: &CudaSlice<f32>,
9257        w0: &CudaSlice<f32>,
9258        w1: &CudaSlice<f32>,
9259        w2: &CudaSlice<f32>,
9260        d0: &mut CudaSlice<f32>,
9261        d1: &mut CudaSlice<f32>,
9262        d2: &mut CudaSlice<f32>,
9263        ncols: usize,
9264        nrows: usize,
9265        eps: f32,
9266    ) -> Result<(), Box<dyn std::error::Error>> {
9267        let f = self.func("rms_norm3_f32");
9268        let cfg = LaunchConfig {
9269            grid_dim: (nrows as u32, 1, 1),
9270            block_dim: (rms_block(), 1, 1),
9271            shared_mem_bytes: 0,
9272        };
9273        let (nc, e) = (ncols as i32, eps);
9274        let __s_b = self.gpu.stream();
9275        let mut b = __s_b.launch_builder(&f);
9276        b.arg(x)
9277            .arg(w0)
9278            .arg(w1)
9279            .arg(w2)
9280            .arg(d0)
9281            .arg(d1)
9282            .arg(d2)
9283            .arg(&nc)
9284            .arg(&e);
9285        unsafe {
9286            b.launch(cfg)?;
9287        }
9288        Ok(())
9289    }
9290
9291    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
9292    #[allow(clippy::too_many_arguments)]
9293    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
9294    /// piggybacks on the same conditions.
9295    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9296        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9297        *WARP_ON.get_or_init(|| {
9298            std::env::var("MEMRA_QKVNORM_W")
9299                .map(|v| v != "0")
9300                .unwrap_or(true)
9301        }) && ncols % 4 == 0
9302            && rows >= 64
9303    }
9304
9305    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
9306    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
9307    #[allow(clippy::too_many_arguments)]
9308    pub fn rms_norm_qkv_w4b(
9309        &self,
9310        q: &CudaSlice<f32>,
9311        k: &CudaSlice<f32>,
9312        v: &CudaSlice<f32>,
9313        wq: &CudaSlice<f32>,
9314        wk: &CudaSlice<f32>,
9315        wv: &CudaSlice<f32>,
9316        dq: &mut CudaSlice<f32>,
9317        dk: &mut CudaSlice<f32>,
9318        dv: &mut CudaSlice<f32>,
9319        dvb: &mut CudaSlice<u8>,
9320        ncols: usize,
9321        rq: usize,
9322        rk: usize,
9323        eps: f32,
9324        vf16: bool,
9325    ) -> Result<(), Box<dyn std::error::Error>> {
9326        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9327        let f = self.func("rms_norm_qkv_w4b_f32");
9328        let rows = (rq + 2 * rk) as u32;
9329        let cfg = LaunchConfig {
9330            grid_dim: (rows.div_ceil(8), 1, 1),
9331            block_dim: (256, 1, 1),
9332            shared_mem_bytes: 0,
9333        };
9334        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9335        let vf = vf16 as i32;
9336        let __s_b = self.gpu.stream();
9337        let mut b = __s_b.launch_builder(&f);
9338        b.arg(q)
9339            .arg(k)
9340            .arg(v)
9341            .arg(wq)
9342            .arg(wk)
9343            .arg(wv)
9344            .arg(dq)
9345            .arg(dk)
9346            .arg(dv)
9347            .arg(&mut *dvb)
9348            .arg(&nc)
9349            .arg(&rqi)
9350            .arg(&rki)
9351            .arg(&rvi)
9352            .arg(&e)
9353            .arg(&vf);
9354        unsafe {
9355            b.launch(cfg)?;
9356        }
9357        Ok(())
9358    }
9359
9360    pub fn rms_norm_qkv(
9361        &self,
9362        q: &CudaSlice<f32>,
9363        k: &CudaSlice<f32>,
9364        v: &CudaSlice<f32>,
9365        wq: &CudaSlice<f32>,
9366        wk: &CudaSlice<f32>,
9367        wv: &CudaSlice<f32>,
9368        dq: &mut CudaSlice<f32>,
9369        dk: &mut CudaSlice<f32>,
9370        dv: &mut CudaSlice<f32>,
9371        ncols: usize,
9372        rq: usize,
9373        rk: usize,
9374        eps: f32,
9375    ) -> Result<(), Box<dyn std::error::Error>> {
9376        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
9377        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
9378        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
9379        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9380        let warp_on = *WARP_ON.get_or_init(|| {
9381            std::env::var("MEMRA_QKVNORM_W")
9382                .map(|v| v != "0")
9383                .unwrap_or(true)
9384        });
9385        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
9386        // replay numerics are untouched on every model; only prefill depth takes the new config.
9387        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9388            let f = self.func("rms_norm_qkv_w4_f32");
9389            let rows = (rq + 2 * rk) as u32;
9390            let cfg = LaunchConfig {
9391                grid_dim: (rows.div_ceil(8), 1, 1),
9392                block_dim: (256, 1, 1),
9393                shared_mem_bytes: 0,
9394            };
9395            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9396            let __s_b = self.gpu.stream();
9397            let mut b = __s_b.launch_builder(&f);
9398            b.arg(q)
9399                .arg(k)
9400                .arg(v)
9401                .arg(wq)
9402                .arg(wk)
9403                .arg(wv)
9404                .arg(dq)
9405                .arg(dk)
9406                .arg(dv)
9407                .arg(&nc)
9408                .arg(&rqi)
9409                .arg(&rki)
9410                .arg(&rvi)
9411                .arg(&e);
9412            unsafe {
9413                b.launch(cfg)?;
9414            }
9415            return Ok(());
9416        }
9417        let f = self.func("rms_norm_qkv_f32");
9418        let grid = (rq + 2 * rk) as u32;
9419        let cfg = LaunchConfig {
9420            grid_dim: (grid, 1, 1),
9421            block_dim: (rms_block(), 1, 1),
9422            shared_mem_bytes: 0,
9423        };
9424        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9425        let __s_b = self.gpu.stream();
9426        let mut b = __s_b.launch_builder(&f);
9427        b.arg(q)
9428            .arg(k)
9429            .arg(v)
9430            .arg(wq)
9431            .arg(wk)
9432            .arg(wv)
9433            .arg(dq)
9434            .arg(dk)
9435            .arg(dv)
9436            .arg(&nc)
9437            .arg(&rqi)
9438            .arg(&rki)
9439            .arg(&e);
9440        unsafe {
9441            b.launch(cfg)?;
9442        }
9443        Ok(())
9444    }
9445
9446    /// gemma4 fused pair of rms_norms over two different inputs (same width).
9447    #[allow(clippy::too_many_arguments)]
9448    pub fn rms_norm2x(
9449        &self,
9450        a: &CudaSlice<f32>,
9451        bb: &CudaSlice<f32>,
9452        wa: &CudaSlice<f32>,
9453        wb: &CudaSlice<f32>,
9454        da: &mut CudaSlice<f32>,
9455        db: &mut CudaSlice<f32>,
9456        ncols: usize,
9457        nrows: usize,
9458        eps: f32,
9459    ) -> Result<(), Box<dyn std::error::Error>> {
9460        let f = self.func("rms_norm2x_f32");
9461        let cfg = LaunchConfig {
9462            grid_dim: (2 * nrows as u32, 1, 1),
9463            block_dim: (rms_block(), 1, 1),
9464            shared_mem_bytes: 0,
9465        };
9466        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9467        let __s_b = self.gpu.stream();
9468        let mut b = __s_b.launch_builder(&f);
9469        b.arg(a)
9470            .arg(bb)
9471            .arg(wa)
9472            .arg(wb)
9473            .arg(da)
9474            .arg(db)
9475            .arg(&nc)
9476            .arg(&nr)
9477            .arg(&e);
9478        unsafe {
9479            b.launch(cfg)?;
9480        }
9481        Ok(())
9482    }
9483
9484    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
9485    pub fn softcap(
9486        &self,
9487        y: &mut CudaSlice<f32>,
9488        cap: f32,
9489        n: usize,
9490    ) -> Result<(), Box<dyn std::error::Error>> {
9491        let f = self.func("softcap_f32");
9492        let cfg = LaunchConfig::for_num_elems(n as u32);
9493        let ni = n as i32;
9494        let __s_b = self.gpu.stream();
9495        let mut b = __s_b.launch_builder(&f);
9496        b.arg(y).arg(&cap).arg(&ni);
9497        unsafe {
9498            b.launch(cfg)?;
9499        }
9500        Ok(())
9501    }
9502
9503    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
9504    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
9505    pub fn mask_ids_rows(
9506        &self,
9507        y: &mut CudaSlice<f32>,
9508        ids: &CudaSlice<i32>,
9509        n_ids: usize,
9510        n_vocab: usize,
9511        t: usize,
9512    ) -> Result<(), Box<dyn std::error::Error>> {
9513        let f = self.func("mask_ids_rows_f32");
9514        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9515        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9516        let __s_b = self.gpu.stream();
9517        let mut b = __s_b.launch_builder(&f);
9518        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9519        unsafe {
9520            b.launch(cfg)?;
9521        }
9522        Ok(())
9523    }
9524
9525    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
9526    #[allow(clippy::too_many_arguments)]
9527    pub fn add_scale_rms_norm(
9528        &self,
9529        a: &CudaSlice<f32>,
9530        b_in: &CudaSlice<f32>,
9531        c: f32,
9532        w: &CudaSlice<f32>,
9533        res: &mut CudaSlice<f32>,
9534        dst: &mut CudaSlice<f32>,
9535        ncols: usize,
9536        nrows: usize,
9537        eps: f32,
9538    ) -> Result<(), Box<dyn std::error::Error>> {
9539        let f = self.func("add_scale_rms_norm_f32");
9540        let cfg = LaunchConfig {
9541            grid_dim: (nrows as u32, 1, 1),
9542            block_dim: (rms_block(), 1, 1),
9543            shared_mem_bytes: 0,
9544        };
9545        let (nc, e2) = (ncols as i32, eps);
9546        let __s_b = self.gpu.stream();
9547        let mut b = __s_b.launch_builder(&f);
9548        b.arg(a)
9549            .arg(b_in)
9550            .arg(&c)
9551            .arg(w)
9552            .arg(res)
9553            .arg(dst)
9554            .arg(&nc)
9555            .arg(&e2);
9556        unsafe {
9557            b.launch(cfg)?;
9558        }
9559        Ok(())
9560    }
9561
9562    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
9563    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
9564    #[allow(clippy::too_many_arguments)]
9565    pub fn add_scale_rms_norm_q8_1(
9566        &self,
9567        a: &CudaSlice<f32>,
9568        b_in: &CudaSlice<f32>,
9569        c: f32,
9570        w: &CudaSlice<f32>,
9571        res: &mut CudaSlice<f32>,
9572        ncols: usize,
9573        nrows: usize,
9574        eps: f32,
9575    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9576        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9577        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9578        let (nc, e2) = (ncols as i32, eps);
9579        if Self::pdl_on() && Self::pdl_wb_on() {
9580            {
9581                use cudarc::driver::{DevicePtr, DevicePtrMut};
9582                let s = &self.gpu.stream();
9583                let (pa, _g0) = a.device_ptr(s);
9584                let (pb, _g1) = b_in.device_ptr(s);
9585                let (pw, _g2) = w.device_ptr(s);
9586                let (pr, _g3) = res.device_ptr_mut(s);
9587                let (pq, _g4) = out_q.device_ptr_mut(s);
9588                let (pd, _g5) = out_d.device_ptr_mut(s);
9589                let mut ps = [
9590                    &pa as *const _ as *mut std::ffi::c_void,
9591                    &pb as *const _ as *mut _,
9592                    &c as *const _ as *mut _,
9593                    &pw as *const _ as *mut _,
9594                    &pr as *const _ as *mut _,
9595                    &pq as *const _ as *mut _,
9596                    &pd as *const _ as *mut _,
9597                    &nc as *const _ as *mut _,
9598                    &e2 as *const _ as *mut _,
9599                ];
9600                unsafe {
9601                    self.launch_pdl(
9602                        "add_scale_rms_norm_q8_1",
9603                        (nrows as u32, 1, 1),
9604                        (rms_block(), 1, 1),
9605                        &mut ps,
9606                    )?;
9607                }
9608            }
9609            return Ok((out_q, out_d));
9610        }
9611        let f = self.func("add_scale_rms_norm_q8_1");
9612        let cfg = LaunchConfig {
9613            grid_dim: (nrows as u32, 1, 1),
9614            block_dim: (rms_block(), 1, 1),
9615            shared_mem_bytes: 0,
9616        };
9617        let __s_b = self.gpu.stream();
9618        let mut b = __s_b.launch_builder(&f);
9619        b.arg(a)
9620            .arg(b_in)
9621            .arg(&c)
9622            .arg(w)
9623            .arg(res)
9624            .arg(&mut out_q)
9625            .arg(&mut out_d)
9626            .arg(&nc)
9627            .arg(&e2);
9628        unsafe {
9629            b.launch(cfg)?;
9630        }
9631        Ok((out_q, out_d))
9632    }
9633
9634    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
9635    #[allow(clippy::too_many_arguments)]
9636    pub fn add_scale_rms_norm_q8_1_into(
9637        &self,
9638        a: &CudaSlice<f32>,
9639        b_in: &CudaSlice<f32>,
9640        c: f32,
9641        w: &CudaSlice<f32>,
9642        res: &mut CudaSlice<f32>,
9643        ncols: usize,
9644        nrows: usize,
9645        eps: f32,
9646        out_q: &mut CudaSlice<i8>,
9647        out_d: &mut CudaSlice<f32>,
9648    ) -> Result<(), Box<dyn std::error::Error>> {
9649        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9650        let (nc, e2) = (ncols as i32, eps);
9651        if Self::pdl_on() && Self::pdl_wb_on() {
9652            use cudarc::driver::{DevicePtr, DevicePtrMut};
9653            let s = &self.gpu.stream();
9654            let (pa, _g0) = a.device_ptr(s);
9655            let (pb, _g1) = b_in.device_ptr(s);
9656            let (pw, _g2) = w.device_ptr(s);
9657            let (pr, _g3) = res.device_ptr_mut(s);
9658            let (pq, _g4) = out_q.device_ptr_mut(s);
9659            let (pd, _g5) = out_d.device_ptr_mut(s);
9660            let mut ps = [
9661                &pa as *const _ as *mut std::ffi::c_void,
9662                &pb as *const _ as *mut _,
9663                &c as *const _ as *mut _,
9664                &pw as *const _ as *mut _,
9665                &pr as *const _ as *mut _,
9666                &pq as *const _ as *mut _,
9667                &pd as *const _ as *mut _,
9668                &nc as *const _ as *mut _,
9669                &e2 as *const _ as *mut _,
9670            ];
9671            unsafe {
9672                self.launch_pdl(
9673                    "add_scale_rms_norm_q8_1",
9674                    (nrows as u32, 1, 1),
9675                    (rms_block(), 1, 1),
9676                    &mut ps,
9677                )?;
9678            }
9679            return Ok(());
9680        }
9681        let f = self.func("add_scale_rms_norm_q8_1");
9682        let cfg = LaunchConfig {
9683            grid_dim: (nrows as u32, 1, 1),
9684            block_dim: (rms_block(), 1, 1),
9685            shared_mem_bytes: 0,
9686        };
9687        let __s_b = self.gpu.stream();
9688        let mut b = __s_b.launch_builder(&f);
9689        b.arg(a)
9690            .arg(b_in)
9691            .arg(&c)
9692            .arg(w)
9693            .arg(res)
9694            .arg(&mut *out_q)
9695            .arg(&mut *out_d)
9696            .arg(&nc)
9697            .arg(&e2);
9698        unsafe {
9699            b.launch(cfg)?;
9700        }
9701        Ok(())
9702    }
9703
9704    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
9705    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
9706    #[allow(clippy::too_many_arguments)]
9707    pub fn rms_pre_add_scale_rms_norm_q8_1(
9708        &self,
9709        a: &CudaSlice<f32>,
9710        wa: &CudaSlice<f32>,
9711        b_in: &CudaSlice<f32>,
9712        c: f32,
9713        w: &CudaSlice<f32>,
9714        res: &mut CudaSlice<f32>,
9715        ncols: usize,
9716        nrows: usize,
9717        eps: f32,
9718    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9719        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9720        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9721        let (nc, e2) = (ncols as i32, eps);
9722        if Self::pdl_on() {
9723            {
9724                use cudarc::driver::{DevicePtr, DevicePtrMut};
9725                let s = &self.gpu.stream();
9726                let (pa, _g0) = a.device_ptr(s);
9727                let (pwa, _g1) = wa.device_ptr(s);
9728                let (pb, _g2) = b_in.device_ptr(s);
9729                let (pw, _g3) = w.device_ptr(s);
9730                let (pr, _g4) = res.device_ptr_mut(s);
9731                let (pq, _g5) = out_q.device_ptr_mut(s);
9732                let (pd, _g6) = out_d.device_ptr_mut(s);
9733                let mut ps = [
9734                    &pa as *const _ as *mut std::ffi::c_void,
9735                    &pwa as *const _ as *mut _,
9736                    &pb as *const _ as *mut _,
9737                    &c as *const _ as *mut _,
9738                    &pw as *const _ as *mut _,
9739                    &pr as *const _ as *mut _,
9740                    &pq as *const _ as *mut _,
9741                    &pd as *const _ as *mut _,
9742                    &nc as *const _ as *mut _,
9743                    &e2 as *const _ as *mut _,
9744                ];
9745                unsafe {
9746                    self.launch_pdl(
9747                        "rms_pre_add_scale_rms_norm_q8_1",
9748                        (nrows as u32, 1, 1),
9749                        (rms_block(), 1, 1),
9750                        &mut ps,
9751                    )?;
9752                }
9753            }
9754            return Ok((out_q, out_d));
9755        }
9756        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9757        let cfg = LaunchConfig {
9758            grid_dim: (nrows as u32, 1, 1),
9759            block_dim: (rms_block(), 1, 1),
9760            shared_mem_bytes: 0,
9761        };
9762        let __s_b = self.gpu.stream();
9763        let mut b = __s_b.launch_builder(&f);
9764        b.arg(a)
9765            .arg(wa)
9766            .arg(b_in)
9767            .arg(&c)
9768            .arg(w)
9769            .arg(res)
9770            .arg(&mut out_q)
9771            .arg(&mut out_d)
9772            .arg(&nc)
9773            .arg(&e2);
9774        unsafe {
9775            b.launch(cfg)?;
9776        }
9777        Ok((out_q, out_d))
9778    }
9779
9780    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
9781    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
9782    pub fn gelu_tanh_mul_q8_1(
9783        &self,
9784        gate: &CudaSlice<f32>,
9785        up: &cudarc::driver::CudaView<f32>,
9786        act: &mut CudaSlice<f32>,
9787        ncols: usize,
9788        nrows: usize,
9789    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9790        debug_assert!(ncols % 128 == 0);
9791        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9792        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9793        let nc = ncols as i32;
9794        if Self::pdl_on() {
9795            {
9796                use cudarc::driver::{DevicePtr, DevicePtrMut};
9797                let s = &self.gpu.stream();
9798                let (pg, _g0) = gate.device_ptr(s);
9799                let (pu, _g1) = up.device_ptr(s);
9800                let (pact, _g2) = act.device_ptr_mut(s);
9801                let (pq, _g3) = out_q.device_ptr_mut(s);
9802                let (pd, _g4) = out_d.device_ptr_mut(s);
9803                let mut ps = [
9804                    &pg as *const _ as *mut std::ffi::c_void,
9805                    &pu as *const _ as *mut _,
9806                    &pact as *const _ as *mut _,
9807                    &pq as *const _ as *mut _,
9808                    &pd as *const _ as *mut _,
9809                    &nc as *const _ as *mut _,
9810                ];
9811                unsafe {
9812                    self.launch_pdl(
9813                        "gelu_tanh_mul_q8_1",
9814                        (nrows as u32, 1, 1),
9815                        (rms_block(), 1, 1),
9816                        &mut ps,
9817                    )?;
9818                }
9819            }
9820            return Ok((out_q, out_d));
9821        }
9822        let f = self.func("gelu_tanh_mul_q8_1");
9823        let cfg = LaunchConfig {
9824            grid_dim: (nrows as u32, 1, 1),
9825            block_dim: (rms_block(), 1, 1),
9826            shared_mem_bytes: 0,
9827        };
9828        let __s_b = self.gpu.stream();
9829        let mut b = __s_b.launch_builder(&f);
9830        b.arg(gate)
9831            .arg(up)
9832            .arg(act)
9833            .arg(&mut out_q)
9834            .arg(&mut out_d)
9835            .arg(&nc);
9836        unsafe {
9837            b.launch(cfg)?;
9838        }
9839        Ok((out_q, out_d))
9840    }
9841
9842    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9843    #[allow(clippy::too_many_arguments)]
9844    pub fn gelu_tanh_mul_q8_1_into(
9845        &self,
9846        gate: &CudaSlice<f32>,
9847        up: &cudarc::driver::CudaView<f32>,
9848        act: &mut CudaSlice<f32>,
9849        ncols: usize,
9850        nrows: usize,
9851        out_q: &mut CudaSlice<i8>,
9852        out_d: &mut CudaSlice<f32>,
9853    ) -> Result<(), Box<dyn std::error::Error>> {
9854        debug_assert!(ncols % 128 == 0);
9855        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9856        let nc = ncols as i32;
9857        if Self::pdl_on() {
9858            use cudarc::driver::{DevicePtr, DevicePtrMut};
9859            let s = &self.gpu.stream();
9860            let (pg, _g0) = gate.device_ptr(s);
9861            let (pu, _g1) = up.device_ptr(s);
9862            let (pact, _g2) = act.device_ptr_mut(s);
9863            let (pq, _g3) = out_q.device_ptr_mut(s);
9864            let (pd, _g4) = out_d.device_ptr_mut(s);
9865            let mut ps = [
9866                &pg as *const _ as *mut std::ffi::c_void,
9867                &pu as *const _ as *mut _,
9868                &pact as *const _ as *mut _,
9869                &pq as *const _ as *mut _,
9870                &pd as *const _ as *mut _,
9871                &nc as *const _ as *mut _,
9872            ];
9873            unsafe {
9874                self.launch_pdl(
9875                    "gelu_tanh_mul_q8_1",
9876                    (nrows as u32, 1, 1),
9877                    (rms_block(), 1, 1),
9878                    &mut ps,
9879                )?;
9880            }
9881            return Ok(());
9882        }
9883        let f = self.func("gelu_tanh_mul_q8_1");
9884        let cfg = LaunchConfig {
9885            grid_dim: (nrows as u32, 1, 1),
9886            block_dim: (rms_block(), 1, 1),
9887            shared_mem_bytes: 0,
9888        };
9889        let __s_b = self.gpu.stream();
9890        let mut b = __s_b.launch_builder(&f);
9891        b.arg(gate)
9892            .arg(up)
9893            .arg(&mut *act)
9894            .arg(&mut *out_q)
9895            .arg(&mut *out_d)
9896            .arg(&nc);
9897        unsafe {
9898            b.launch(cfg)?;
9899        }
9900        Ok(())
9901    }
9902
9903    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9904    #[allow(clippy::too_many_arguments)]
9905    pub fn add_rms_norm3_q8z(
9906        &self,
9907        a: &CudaSlice<f32>,
9908        b_in: &CudaSlice<f32>,
9909        w0: &CudaSlice<f32>,
9910        w1: &CudaSlice<f32>,
9911        w2: &CudaSlice<f32>,
9912        res: &mut CudaSlice<f32>,
9913        out1: &mut CudaSlice<f32>,
9914        ncols: usize,
9915        nrows: usize,
9916        eps: f32,
9917    ) -> Result<
9918        (
9919            (CudaSlice<i8>, CudaSlice<f32>),
9920            (CudaSlice<i8>, CudaSlice<f32>),
9921        ),
9922        Box<dyn std::error::Error>,
9923    > {
9924        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9925        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9926        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9927        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9928        let f = self.func("add_rms_norm3_q8z_f32");
9929        let cfg = LaunchConfig {
9930            grid_dim: (nrows as u32, 1, 1),
9931            block_dim: (rms_block(), 1, 1),
9932            shared_mem_bytes: 0,
9933        };
9934        let (nc, e2) = (ncols as i32, eps);
9935        let __s_b = self.gpu.stream();
9936        let mut b = __s_b.launch_builder(&f);
9937        b.arg(a)
9938            .arg(b_in)
9939            .arg(w0)
9940            .arg(w1)
9941            .arg(w2)
9942            .arg(res)
9943            .arg(&mut q0)
9944            .arg(&mut d0)
9945            .arg(out1)
9946            .arg(&mut q2)
9947            .arg(&mut d2)
9948            .arg(&nc)
9949            .arg(&e2);
9950        unsafe {
9951            b.launch(cfg)?;
9952        }
9953        Ok(((q0, d0), (q2, d2)))
9954    }
9955
9956    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9957    #[allow(clippy::too_many_arguments)]
9958    pub fn add_rms_norm3(
9959        &self,
9960        a: &CudaSlice<f32>,
9961        b_in: &CudaSlice<f32>,
9962        w0: &CudaSlice<f32>,
9963        w1: &CudaSlice<f32>,
9964        w2: &CudaSlice<f32>,
9965        res: &mut CudaSlice<f32>,
9966        d0: &mut CudaSlice<f32>,
9967        d1: &mut CudaSlice<f32>,
9968        d2: &mut CudaSlice<f32>,
9969        ncols: usize,
9970        nrows: usize,
9971        eps: f32,
9972    ) -> Result<(), Box<dyn std::error::Error>> {
9973        let f = self.func("add_rms_norm3_f32");
9974        let cfg = LaunchConfig {
9975            grid_dim: (nrows as u32, 1, 1),
9976            block_dim: (rms_block(), 1, 1),
9977            shared_mem_bytes: 0,
9978        };
9979        let (nc, e2) = (ncols as i32, eps);
9980        let __s_b = self.gpu.stream();
9981        let mut b = __s_b.launch_builder(&f);
9982        b.arg(a)
9983            .arg(b_in)
9984            .arg(w0)
9985            .arg(w1)
9986            .arg(w2)
9987            .arg(res)
9988            .arg(d0)
9989            .arg(d1)
9990            .arg(d2)
9991            .arg(&nc)
9992            .arg(&e2);
9993        unsafe {
9994            b.launch(cfg)?;
9995        }
9996        Ok(())
9997    }
9998
9999    /// dst = (a + b) * c (residual add + layer scale, one launch).
10000    pub fn add_scale(
10001        &self,
10002        a: &CudaSlice<f32>,
10003        b_in: &CudaSlice<f32>,
10004        c: f32,
10005        dst: &mut CudaSlice<f32>,
10006        n: usize,
10007    ) -> Result<(), Box<dyn std::error::Error>> {
10008        let f = self.func("add_scale_f32");
10009        let cfg = LaunchConfig::for_num_elems(n as u32);
10010        let ni = n as i32;
10011        let __s_b = self.gpu.stream();
10012        let mut b = __s_b.launch_builder(&f);
10013        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10014        unsafe {
10015            b.launch(cfg)?;
10016        }
10017        Ok(())
10018    }
10019
10020    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
10021    pub fn layer_norm_bias(
10022        &self,
10023        x: &CudaSlice<f32>,
10024        w: &CudaSlice<f32>,
10025        b: &CudaSlice<f32>,
10026        dst: &mut CudaSlice<f32>,
10027        ncols: usize,
10028        nrows: usize,
10029        eps: f32,
10030    ) -> Result<(), Box<dyn std::error::Error>> {
10031        let f = self.func("layer_norm_bias_f32");
10032        let (nc, e) = (ncols as i32, eps);
10033        let cfg = LaunchConfig {
10034            grid_dim: (nrows as u32, 1, 1),
10035            block_dim: (256, 1, 1),
10036            shared_mem_bytes: 0,
10037        };
10038        let __s_b = self.gpu.stream();
10039        let mut lb = __s_b.launch_builder(&f);
10040        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10041        unsafe {
10042            lb.launch(cfg)?;
10043        }
10044        Ok(())
10045    }
10046
10047    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
10048    pub fn gelu_tanh(
10049        &self,
10050        x: &CudaSlice<f32>,
10051        dst: &mut CudaSlice<f32>,
10052        n: usize,
10053    ) -> Result<(), Box<dyn std::error::Error>> {
10054        let f = self.func("gelu_tanh_f32");
10055        let ni = n as i64;
10056        let cfg = LaunchConfig {
10057            grid_dim: (n.div_ceil(256) as u32, 1, 1),
10058            block_dim: (256, 1, 1),
10059            shared_mem_bytes: 0,
10060        };
10061        let __s_b = self.gpu.stream();
10062        let mut lb = __s_b.launch_builder(&f);
10063        lb.arg(x).arg(&mut *dst).arg(&ni);
10064        unsafe {
10065            lb.launch(cfg)?;
10066        }
10067        Ok(())
10068    }
10069
10070    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
10071    pub fn row_softmax(
10072        &self,
10073        x: &mut CudaSlice<f32>,
10074        ncols: usize,
10075        nrows: usize,
10076    ) -> Result<(), Box<dyn std::error::Error>> {
10077        let f = self.func("row_softmax_f32");
10078        let nc = ncols as i32;
10079        let cfg = LaunchConfig {
10080            grid_dim: (nrows as u32, 1, 1),
10081            block_dim: (256, 1, 1),
10082            shared_mem_bytes: 0,
10083        };
10084        let __s_b = self.gpu.stream();
10085        let mut lb = __s_b.launch_builder(&f);
10086        lb.arg(&mut *x).arg(&nc);
10087        unsafe {
10088            lb.launch(cfg)?;
10089        }
10090        Ok(())
10091    }
10092
10093    pub fn rms_norm(
10094        &self,
10095        x: &CudaSlice<f32>,
10096        w: &CudaSlice<f32>,
10097        dst: &mut CudaSlice<f32>,
10098        ncols: usize,
10099        nrows: usize,
10100        eps: f32,
10101    ) -> Result<(), Box<dyn std::error::Error>> {
10102        let (nc, e) = (ncols as i32, eps);
10103        let kname = if Self::norm_ilp_on() {
10104            "rms_norm_f32_v2"
10105        } else {
10106            "rms_norm_f32"
10107        };
10108        if Self::pdl_on() && Self::pdl_wb_on() {
10109            use cudarc::driver::{DevicePtr, DevicePtrMut};
10110            let s = &self.gpu.stream();
10111            let (px, _g0) = x.device_ptr(s);
10112            let (pw, _g1) = w.device_ptr(s);
10113            let (pd, _g2) = dst.device_ptr_mut(s);
10114            let mut ps = [
10115                &px as *const _ as *mut std::ffi::c_void,
10116                &pw as *const _ as *mut _,
10117                &pd as *const _ as *mut _,
10118                &nc as *const _ as *mut _,
10119                &e as *const _ as *mut _,
10120            ];
10121            unsafe {
10122                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10123            }
10124            return Ok(());
10125        }
10126        let f = self.func(kname);
10127        let cfg = LaunchConfig {
10128            grid_dim: (nrows as u32, 1, 1),
10129            block_dim: (rms_block(), 1, 1),
10130            shared_mem_bytes: 0,
10131        };
10132        let __s_b = self.gpu.stream();
10133        let mut b = __s_b.launch_builder(&f);
10134        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10135        unsafe {
10136            b.launch(cfg)?;
10137        }
10138        Ok(())
10139    }
10140
10141    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
10142    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
10143    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
10144    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
10145    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
10146    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
10147    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
10148    pub fn rms_norm_decode(
10149        &self,
10150        x: &CudaSlice<f32>,
10151        w: &CudaSlice<f32>,
10152        dst: &mut CudaSlice<f32>,
10153        ncols: usize,
10154        nrows: usize,
10155        eps: f32,
10156    ) -> Result<(), Box<dyn std::error::Error>> {
10157        let f = self.func(if Self::norm_ilp_on() {
10158            "rms_norm_f32_v2"
10159        } else {
10160            "rms_norm_f32"
10161        });
10162        let cfg = LaunchConfig {
10163            grid_dim: (nrows as u32, 1, 1),
10164            block_dim: (1024, 1, 1),
10165            shared_mem_bytes: 0,
10166        };
10167        let (nc, e) = (ncols as i32, eps);
10168        let __s_b = self.gpu.stream();
10169        let mut b = __s_b.launch_builder(&f);
10170        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10171        unsafe {
10172            b.launch(cfg)?;
10173        }
10174        Ok(())
10175    }
10176
10177    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
10178    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
10179    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
10180    pub fn rms_norm_q8_1(
10181        &self,
10182        x: &CudaSlice<f32>,
10183        w: &CudaSlice<f32>,
10184        ncols: usize,
10185        nrows: usize,
10186        eps: f32,
10187    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10188        let nblk = ncols / 32;
10189        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10190        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10191        let (nc, e) = (ncols as i32, eps);
10192        if Self::pdl_on() {
10193            {
10194                use cudarc::driver::{DevicePtr, DevicePtrMut};
10195                let s = &self.gpu.stream();
10196                let (px, _g0) = x.device_ptr(s);
10197                let (pw, _g1) = w.device_ptr(s);
10198                let (pq, _g2) = q.device_ptr_mut(s);
10199                let (pd, _g3) = d.device_ptr_mut(s);
10200                let mut ps = [
10201                    &px as *const _ as *mut std::ffi::c_void,
10202                    &pw as *const _ as *mut _,
10203                    &pq as *const _ as *mut _,
10204                    &pd as *const _ as *mut _,
10205                    &nc as *const _ as *mut _,
10206                    &e as *const _ as *mut _,
10207                ];
10208                unsafe {
10209                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10210                }
10211            }
10212            return Ok((q, d));
10213        }
10214        let f = self.func("rms_norm_q8_1");
10215        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
10216        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
10217        let cfg = LaunchConfig {
10218            grid_dim: (nrows as u32, 1, 1),
10219            block_dim: (1024, 1, 1),
10220            shared_mem_bytes: 0,
10221        };
10222        let __s_b = self.gpu.stream();
10223        let mut b = __s_b.launch_builder(&f);
10224        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10225        unsafe {
10226            b.launch(cfg)?;
10227        }
10228        Ok((q, d))
10229    }
10230
10231    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
10232    /// PDL arm), caller-owned outputs.
10233    pub fn rms_norm_q8_1_into(
10234        &self,
10235        x: &CudaSlice<f32>,
10236        w: &CudaSlice<f32>,
10237        ncols: usize,
10238        nrows: usize,
10239        eps: f32,
10240        q: &mut CudaSlice<i8>,
10241        d: &mut CudaSlice<f32>,
10242    ) -> Result<(), Box<dyn std::error::Error>> {
10243        let nblk = ncols / 32;
10244        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10245        let (nc, e) = (ncols as i32, eps);
10246        if Self::pdl_on() {
10247            use cudarc::driver::{DevicePtr, DevicePtrMut};
10248            let s = &self.gpu.stream();
10249            let (px, _g0) = x.device_ptr(s);
10250            let (pw, _g1) = w.device_ptr(s);
10251            let (pq, _g2) = q.device_ptr_mut(s);
10252            let (pd, _g3) = d.device_ptr_mut(s);
10253            let mut ps = [
10254                &px as *const _ as *mut std::ffi::c_void,
10255                &pw as *const _ as *mut _,
10256                &pq as *const _ as *mut _,
10257                &pd as *const _ as *mut _,
10258                &nc as *const _ as *mut _,
10259                &e as *const _ as *mut _,
10260            ];
10261            unsafe {
10262                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10263            }
10264            return Ok(());
10265        }
10266        let f = self.func("rms_norm_q8_1");
10267        let cfg = LaunchConfig {
10268            grid_dim: (nrows as u32, 1, 1),
10269            block_dim: (1024, 1, 1),
10270            shared_mem_bytes: 0,
10271        };
10272        let __s_b = self.gpu.stream();
10273        let mut b = __s_b.launch_builder(&f);
10274        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10275        unsafe {
10276            b.launch(cfg)?;
10277        }
10278        Ok(())
10279    }
10280
10281    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
10282    pub fn quantize_q8_1_into(
10283        &self,
10284        x: &CudaSlice<f32>,
10285        m: usize,
10286        in_f: usize,
10287        q: &mut CudaSlice<i8>,
10288        d: &mut CudaSlice<f32>,
10289    ) -> Result<(), Box<dyn std::error::Error>> {
10290        let nblk = in_f / 32;
10291        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10292        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10293        let (inf, mi) = (in_f as i32, m as i32);
10294        if Self::pdl_on() && Self::pdl_wb_on() {
10295            use cudarc::driver::{DevicePtr, DevicePtrMut};
10296            let s = &self.gpu.stream();
10297            let (px, _g0) = x.device_ptr(s);
10298            let (pq, _g1) = q.device_ptr_mut(s);
10299            let (pd, _g2) = d.device_ptr_mut(s);
10300            let mut ps = [
10301                &px as *const _ as *mut std::ffi::c_void,
10302                &pq as *const _ as *mut _,
10303                &pd as *const _ as *mut _,
10304                &inf as *const _ as *mut _,
10305                &mi as *const _ as *mut _,
10306            ];
10307            unsafe {
10308                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10309            }
10310            return Ok(());
10311        }
10312        let f = self.func("quantize_q8_1");
10313        let __s_b = self.gpu.stream();
10314        let mut b = __s_b.launch_builder(&f);
10315        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10316        unsafe {
10317            b.launch(cfg)?;
10318        }
10319        Ok(())
10320    }
10321
10322    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
10323    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
10324    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
10325    pub fn add_rms_norm_q8_1(
10326        &self,
10327        a: &CudaSlice<f32>,
10328        b_in: &CudaSlice<f32>,
10329        w: &CudaSlice<f32>,
10330        res: &mut CudaSlice<f32>,
10331        ncols: usize,
10332        nrows: usize,
10333        eps: f32,
10334    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10335        let nblk = ncols / 32;
10336        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10337        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10338        let f = self.func("add_rms_norm_q8_1");
10339        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
10340        let cfg = LaunchConfig {
10341            grid_dim: (nrows as u32, 1, 1),
10342            block_dim: (1024, 1, 1),
10343            shared_mem_bytes: 0,
10344        };
10345        let (nc, e) = (ncols as i32, eps);
10346        let __s_bld = self.gpu.stream();
10347        let mut bld = __s_bld.launch_builder(&f);
10348        bld.arg(a)
10349            .arg(b_in)
10350            .arg(w)
10351            .arg(res)
10352            .arg(&mut q)
10353            .arg(&mut d)
10354            .arg(&nc)
10355            .arg(&e);
10356        unsafe {
10357            bld.launch(cfg)?;
10358        }
10359        Ok((q, d))
10360    }
10361
10362    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
10363    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
10364    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
10365    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
10366    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
10367    #[allow(clippy::too_many_arguments)]
10368    pub fn join_add_rms_norm_raw(
10369        &self,
10370        a0_raw: u64,
10371        a1_raw: u64,
10372        x: &CudaSlice<f32>,
10373        w: &CudaSlice<f32>,
10374        res: &mut CudaSlice<f32>,
10375        dst: &mut CudaSlice<f32>,
10376        ncols: usize,
10377        eps: f32,
10378    ) -> Result<(), Box<dyn std::error::Error>> {
10379        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10380            return Err("join_add_rms_norm geometry".into());
10381        }
10382        let f = self.func("join_add_rms_norm_f32");
10383        let cfg = LaunchConfig {
10384            grid_dim: (1, 1, 1),
10385            block_dim: (rms_block(), 1, 1),
10386            shared_mem_bytes: 0,
10387        };
10388        let (nc, e) = (ncols as i32, eps);
10389        let __s_b = self.gpu.stream();
10390        let mut b = __s_b.launch_builder(&f);
10391        b.arg(&a0_raw)
10392            .arg(&a1_raw)
10393            .arg(x)
10394            .arg(w)
10395            .arg(&mut *res)
10396            .arg(&mut *dst)
10397            .arg(&nc)
10398            .arg(&e);
10399        unsafe {
10400            b.launch(cfg)?;
10401        }
10402        Ok(())
10403    }
10404
10405    pub fn add_rms_norm(
10406        &self,
10407        a: &CudaSlice<f32>,
10408        b: &CudaSlice<f32>,
10409        w: &CudaSlice<f32>,
10410        res: &mut CudaSlice<f32>,
10411        dst: &mut CudaSlice<f32>,
10412        ncols: usize,
10413        nrows: usize,
10414        eps: f32,
10415    ) -> Result<(), Box<dyn std::error::Error>> {
10416        let (nc, e) = (ncols as i32, eps);
10417        let kname = if Self::norm_ilp_on() {
10418            "add_rms_norm_f32_v2"
10419        } else {
10420            "add_rms_norm_f32"
10421        };
10422        if Self::pdl_on() && Self::pdl_wb_on() {
10423            use cudarc::driver::{DevicePtr, DevicePtrMut};
10424            let s = &self.gpu.stream();
10425            let (pa, _g0) = a.device_ptr(s);
10426            let (pb, _g1) = b.device_ptr(s);
10427            let (pw, _g2) = w.device_ptr(s);
10428            let (pr, _g3) = res.device_ptr_mut(s);
10429            let (pd, _g4) = dst.device_ptr_mut(s);
10430            let mut ps = [
10431                &pa as *const _ as *mut std::ffi::c_void,
10432                &pb as *const _ as *mut _,
10433                &pw as *const _ as *mut _,
10434                &pr as *const _ as *mut _,
10435                &pd as *const _ as *mut _,
10436                &nc as *const _ as *mut _,
10437                &e as *const _ as *mut _,
10438            ];
10439            unsafe {
10440                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10441            }
10442            return Ok(());
10443        }
10444        let f = self.func(kname);
10445        let cfg = LaunchConfig {
10446            grid_dim: (nrows as u32, 1, 1),
10447            block_dim: (rms_block(), 1, 1),
10448            shared_mem_bytes: 0,
10449        };
10450        let __s_b2 = self.gpu.stream();
10451        let mut b2 = __s_b2.launch_builder(&f);
10452        b2.arg(a)
10453            .arg(b)
10454            .arg(w)
10455            .arg(&mut *res)
10456            .arg(&mut *dst)
10457            .arg(&nc)
10458            .arg(&e);
10459        unsafe {
10460            b2.launch(cfg)?;
10461        }
10462        Ok(())
10463    }
10464
10465    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
10466    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
10467    #[allow(clippy::too_many_arguments)]
10468    pub fn rms_pre_add_rms_norm(
10469        &self,
10470        a: &CudaSlice<f32>,
10471        wa: &CudaSlice<f32>,
10472        b: &CudaSlice<f32>,
10473        w: &CudaSlice<f32>,
10474        res: &mut CudaSlice<f32>,
10475        dst: &mut CudaSlice<f32>,
10476        ncols: usize,
10477        nrows: usize,
10478        eps: f32,
10479    ) -> Result<(), Box<dyn std::error::Error>> {
10480        let f = self.func("rms_pre_add_rms_norm_f32");
10481        let cfg = LaunchConfig {
10482            grid_dim: (nrows as u32, 1, 1),
10483            block_dim: (rms_block(), 1, 1),
10484            shared_mem_bytes: 0,
10485        };
10486        let (nc, e) = (ncols as i32, eps);
10487        let __s_b2 = self.gpu.stream();
10488        let mut b2 = __s_b2.launch_builder(&f);
10489        b2.arg(a)
10490            .arg(wa)
10491            .arg(b)
10492            .arg(w)
10493            .arg(&mut *res)
10494            .arg(&mut *dst)
10495            .arg(&nc)
10496            .arg(&e);
10497        unsafe {
10498            b2.launch(cfg)?;
10499        }
10500        Ok(())
10501    }
10502
10503    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
10504    #[allow(clippy::too_many_arguments)]
10505    pub fn rms_pre_add_rms_norm_q8z(
10506        &self,
10507        a: &CudaSlice<f32>,
10508        wa: &CudaSlice<f32>,
10509        b: &CudaSlice<f32>,
10510        w: &CudaSlice<f32>,
10511        res: &mut CudaSlice<f32>,
10512        dst: &mut CudaSlice<f32>,
10513        ncols: usize,
10514        nrows: usize,
10515        eps: f32,
10516    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10517        debug_assert!(ncols % 128 == 0);
10518        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10519        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10520        let (nc, e) = (ncols as i32, eps);
10521        if Self::pdl_on() {
10522            {
10523                use cudarc::driver::{DevicePtr, DevicePtrMut};
10524                let s = &self.gpu.stream();
10525                let (pa, _g0) = a.device_ptr(s);
10526                let (pwa, _g1) = wa.device_ptr(s);
10527                let (pb, _g2) = b.device_ptr(s);
10528                let (pw, _g3) = w.device_ptr(s);
10529                let (pr, _g4) = res.device_ptr_mut(s);
10530                let (pdst, _g5) = dst.device_ptr_mut(s);
10531                let (pq, _g6) = out_q.device_ptr_mut(s);
10532                let (pd, _g7) = out_d.device_ptr_mut(s);
10533                let mut ps = [
10534                    &pa as *const _ as *mut std::ffi::c_void,
10535                    &pwa as *const _ as *mut _,
10536                    &pb as *const _ as *mut _,
10537                    &pw as *const _ as *mut _,
10538                    &pr as *const _ as *mut _,
10539                    &pdst as *const _ as *mut _,
10540                    &pq as *const _ as *mut _,
10541                    &pd as *const _ as *mut _,
10542                    &nc as *const _ as *mut _,
10543                    &e as *const _ as *mut _,
10544                ];
10545                unsafe {
10546                    self.launch_pdl(
10547                        "rms_pre_add_rms_norm_q8z_f32",
10548                        (nrows as u32, 1, 1),
10549                        (rms_block(), 1, 1),
10550                        &mut ps,
10551                    )?;
10552                }
10553            }
10554            return Ok((out_q, out_d));
10555        }
10556        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10557        let cfg = LaunchConfig {
10558            grid_dim: (nrows as u32, 1, 1),
10559            block_dim: (rms_block(), 1, 1),
10560            shared_mem_bytes: 0,
10561        };
10562        let __s_b2 = self.gpu.stream();
10563        let mut b2 = __s_b2.launch_builder(&f);
10564        b2.arg(a)
10565            .arg(wa)
10566            .arg(b)
10567            .arg(w)
10568            .arg(&mut *res)
10569            .arg(&mut *dst)
10570            .arg(&mut out_q)
10571            .arg(&mut out_d)
10572            .arg(&nc)
10573            .arg(&e);
10574        unsafe {
10575            b2.launch(cfg)?;
10576        }
10577        Ok((out_q, out_d))
10578    }
10579
10580    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
10581    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
10582    /// body must stay attribute-free (the fused2_into precedent).
10583    #[allow(clippy::too_many_arguments)]
10584    pub fn rms_pre_add_rms_norm_q8z_into(
10585        &self,
10586        a: &CudaSlice<f32>,
10587        wa: &CudaSlice<f32>,
10588        b: &CudaSlice<f32>,
10589        w: &CudaSlice<f32>,
10590        res: &mut CudaSlice<f32>,
10591        dst: &mut CudaSlice<f32>,
10592        ncols: usize,
10593        nrows: usize,
10594        eps: f32,
10595        out_q: &mut CudaSlice<i8>,
10596        out_d: &mut CudaSlice<f32>,
10597    ) -> Result<(), Box<dyn std::error::Error>> {
10598        debug_assert!(ncols % 128 == 0);
10599        let (nc, e) = (ncols as i32, eps);
10600        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10601        let cfg = LaunchConfig {
10602            grid_dim: (nrows as u32, 1, 1),
10603            block_dim: (rms_block(), 1, 1),
10604            shared_mem_bytes: 0,
10605        };
10606        let __s_b = self.gpu.stream();
10607        let mut b2 = __s_b.launch_builder(&f);
10608        b2.arg(a)
10609            .arg(wa)
10610            .arg(b)
10611            .arg(w)
10612            .arg(&mut *res)
10613            .arg(&mut *dst)
10614            .arg(&mut *out_q)
10615            .arg(&mut *out_d)
10616            .arg(&nc)
10617            .arg(&e);
10618        unsafe {
10619            b2.launch(cfg)?;
10620        }
10621        Ok(())
10622    }
10623
10624    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
10625    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
10626    #[allow(clippy::too_many_arguments)]
10627    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10628        &self,
10629        a: &CudaSlice<f32>,
10630        wa: &CudaSlice<f32>,
10631        b_in: &CudaSlice<f32>,
10632        c: f32,
10633        w: &CudaSlice<f32>,
10634        res: &mut CudaSlice<f32>,
10635        ncols: usize,
10636        nrows: usize,
10637        eps: f32,
10638        out_q: &mut CudaSlice<i8>,
10639        out_d: &mut CudaSlice<f32>,
10640    ) -> Result<(), Box<dyn std::error::Error>> {
10641        debug_assert!(ncols % 128 == 0);
10642        let (nc, e2) = (ncols as i32, eps);
10643        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10644        let cfg = LaunchConfig {
10645            grid_dim: (nrows as u32, 1, 1),
10646            block_dim: (rms_block(), 1, 1),
10647            shared_mem_bytes: 0,
10648        };
10649        let __s_b = self.gpu.stream();
10650        let mut b2 = __s_b.launch_builder(&f);
10651        b2.arg(a)
10652            .arg(wa)
10653            .arg(b_in)
10654            .arg(&c)
10655            .arg(w)
10656            .arg(&mut *res)
10657            .arg(&mut *out_q)
10658            .arg(&mut *out_d)
10659            .arg(&nc)
10660            .arg(&e2);
10661        unsafe {
10662            b2.launch(cfg)?;
10663        }
10664        Ok(())
10665    }
10666
10667    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
10668    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
10669    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
10670    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
10671    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
10672    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
10673    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
10674    pub fn g4_pnfold_on() -> bool {
10675        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10676        *ON.get_or_init(|| {
10677            std::env::var("MEMRA_G4_PNFOLD")
10678                .map(|v| v != "0")
10679                .unwrap_or(true)
10680        })
10681    }
10682
10683    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
10684    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
10685    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
10686    pub fn build_q4_out_concat3(
10687        &self,
10688        w0: &crate::model::GpuTensor,
10689        w1: &crate::model::GpuTensor,
10690        w2: &crate::model::GpuTensor,
10691    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10692        use crate::model::GpuTensor;
10693        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10694            match w {
10695                GpuTensor::Quant {
10696                    qtype,
10697                    row_bytes,
10698                    rp,
10699                    ..
10700                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10701                _ => None,
10702            }
10703        };
10704        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10705        else {
10706            return Ok(None);
10707        };
10708        if rb0 != rb1
10709            || rb0 != rb2
10710            || w0.in_features() != w1.in_features()
10711            || w0.in_features() != w2.in_features()
10712        {
10713            return Ok(None);
10714        }
10715        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10716            match w {
10717                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10718                _ => unreachable!(),
10719            }
10720        }
10721        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10722        let total = rb0 * (o0 + o1 + o2);
10723        let mut cat = self.alloc_u8(total)?;
10724        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10725        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10726        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10727        Ok(Some(GpuTensor::Quant {
10728            bytes: cat,
10729            qtype: QT_Q4_0,
10730            row_bytes: rb0,
10731            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10732            scale: 1.0,
10733            rp: false,
10734            #[cfg(memra_cutlass)]
10735            cutlass: None,
10736            fp8: None,
10737            blk: None,
10738            rp4: None,
10739            f16: None,
10740        }))
10741    }
10742
10743    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
10744    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
10745    ///
10746    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
10747    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
10748    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
10749    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
10750    ///
10751    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
10752    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
10753    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
10754    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
10755    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
10756    ///
10757    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
10758    /// width. A future partial-rotary caller fails at its first launch with the geometry named
10759    /// instead of serving quietly wrong logits.
10760    fn full_width_rope_only(
10761        kernel: &str,
10762        n_rot: usize,
10763        head_dim: usize,
10764    ) -> Result<(), Box<dyn std::error::Error>> {
10765        if n_rot == head_dim {
10766            return Ok(());
10767        }
10768        Err(format!(
10769            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10770             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10771             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10772             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10773             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10774        )
10775        .into())
10776    }
10777
10778    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
10779    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10780    /// ([`Engine::full_width_rope_only`]).
10781    #[allow(clippy::too_many_arguments)]
10782    pub fn rms_norm_qkv_rope_cat(
10783        &self,
10784        qkv: &CudaSlice<f32>,
10785        wq: &CudaSlice<f32>,
10786        wk: &CudaSlice<f32>,
10787        wv: &CudaSlice<f32>,
10788        q: &mut CudaSlice<f32>,
10789        k: &mut CudaSlice<f32>,
10790        v: &mut CudaSlice<f32>,
10791        head_dim: usize,
10792        n_rot: usize,
10793        rq: usize,
10794        rk: usize,
10795        pos: &CudaSlice<i32>,
10796        nh_q: usize,
10797        nh_k: usize,
10798        base: f32,
10799        freq_scale: f32,
10800        ff: Option<&CudaSlice<f32>>,
10801        eps: f32,
10802    ) -> Result<(), Box<dyn std::error::Error>> {
10803        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10804        let rows = rq + rk + rk;
10805        let theta_scale = base.powf(-2.0 / head_dim as f32);
10806        let (nc, rqi, rki, nhq, nhk) = (
10807            head_dim as i32,
10808            rq as i32,
10809            rk as i32,
10810            nh_q as i32,
10811            nh_k as i32,
10812        );
10813        if Self::pdl_on() {
10814            use cudarc::driver::{DevicePtr, DevicePtrMut};
10815            let s = &self.gpu.stream();
10816            let (pqkv, _g0) = qkv.device_ptr(s);
10817            let (pwq, _g1) = wq.device_ptr(s);
10818            let (pwk, _g2) = wk.device_ptr(s);
10819            let (pwv, _g3) = wv.device_ptr(s);
10820            let (pq, _g4) = q.device_ptr_mut(s);
10821            let (pk, _g5) = k.device_ptr_mut(s);
10822            let (pv, _g6) = v.device_ptr_mut(s);
10823            let (ppos, _g7) = pos.device_ptr(s);
10824            let (pff, _g8) = match ff {
10825                Some(t) => {
10826                    let (p, g) = t.device_ptr(s);
10827                    (p, Some(g))
10828                }
10829                None => (0, None),
10830            };
10831            let mut ps = [
10832                &pqkv as *const _ as *mut std::ffi::c_void,
10833                &pwq as *const _ as *mut _,
10834                &pwk as *const _ as *mut _,
10835                &pwv as *const _ as *mut _,
10836                &pq as *const _ as *mut _,
10837                &pk as *const _ as *mut _,
10838                &pv as *const _ as *mut _,
10839                &nc as *const _ as *mut _,
10840                &rqi as *const _ as *mut _,
10841                &rki as *const _ as *mut _,
10842                &ppos as *const _ as *mut _,
10843                &nhq as *const _ as *mut _,
10844                &nhk as *const _ as *mut _,
10845                &theta_scale as *const _ as *mut _,
10846                &freq_scale as *const _ as *mut _,
10847                &pff as *const _ as *mut _,
10848                &eps as *const _ as *mut _,
10849            ];
10850            unsafe {
10851                self.launch_pdl(
10852                    "rms_norm_qkv_rope_cat_f32",
10853                    (rows as u32, 1, 1),
10854                    (rms_block(), 1, 1),
10855                    &mut ps,
10856                )?;
10857            }
10858            return Ok(());
10859        }
10860        let f = self.func("rms_norm_qkv_rope_cat_f32");
10861        let cfg = LaunchConfig {
10862            grid_dim: (rows as u32, 1, 1),
10863            block_dim: (rms_block(), 1, 1),
10864            shared_mem_bytes: 0,
10865        };
10866        let __s_b = self.gpu.stream();
10867        let mut b = __s_b.launch_builder(&f);
10868        match ff {
10869            Some(t) => {
10870                b.arg(qkv)
10871                    .arg(wq)
10872                    .arg(wk)
10873                    .arg(wv)
10874                    .arg(&mut *q)
10875                    .arg(&mut *k)
10876                    .arg(&mut *v)
10877                    .arg(&nc)
10878                    .arg(&rqi)
10879                    .arg(&rki)
10880                    .arg(pos)
10881                    .arg(&nhq)
10882                    .arg(&nhk)
10883                    .arg(&theta_scale)
10884                    .arg(&freq_scale)
10885                    .arg(t)
10886                    .arg(&eps);
10887                unsafe {
10888                    b.launch(cfg)?;
10889                }
10890            }
10891            None => {
10892                let null: u64 = 0;
10893                b.arg(qkv)
10894                    .arg(wq)
10895                    .arg(wk)
10896                    .arg(wv)
10897                    .arg(&mut *q)
10898                    .arg(&mut *k)
10899                    .arg(&mut *v)
10900                    .arg(&nc)
10901                    .arg(&rqi)
10902                    .arg(&rki)
10903                    .arg(pos)
10904                    .arg(&nhq)
10905                    .arg(&nhk)
10906                    .arg(&theta_scale)
10907                    .arg(&freq_scale)
10908                    .arg(&null)
10909                    .arg(&eps);
10910                unsafe {
10911                    b.launch(cfg)?;
10912                }
10913            }
10914        }
10915        Ok(())
10916    }
10917
10918    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10919    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10920    /// ([`Engine::full_width_rope_only`]).
10921    #[allow(clippy::too_many_arguments)]
10922    pub fn rms_norm_qkv_rope(
10923        &self,
10924        q0: &CudaSlice<f32>,
10925        k0: &CudaSlice<f32>,
10926        v0: &CudaSlice<f32>,
10927        wq: &CudaSlice<f32>,
10928        wk: &CudaSlice<f32>,
10929        wv: &CudaSlice<f32>,
10930        q: &mut CudaSlice<f32>,
10931        k: &mut CudaSlice<f32>,
10932        v: &mut CudaSlice<f32>,
10933        head_dim: usize,
10934        n_rot: usize,
10935        rq: usize,
10936        rk: usize,
10937        pos: &CudaSlice<i32>,
10938        nh_q: usize,
10939        nh_k: usize,
10940        base: f32,
10941        freq_scale: f32,
10942        ff: Option<&CudaSlice<f32>>,
10943        eps: f32,
10944    ) -> Result<(), Box<dyn std::error::Error>> {
10945        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10946        let f = self.func("rms_norm_qkv_rope_f32");
10947        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10948        let cfg = LaunchConfig {
10949            grid_dim: (rows as u32, 1, 1),
10950            block_dim: (rms_block(), 1, 1),
10951            shared_mem_bytes: 0,
10952        };
10953        let theta_scale = base.powf(-2.0 / head_dim as f32);
10954        let (nc, rqi, rki, nhq, nhk) = (
10955            head_dim as i32,
10956            rq as i32,
10957            rk as i32,
10958            nh_q as i32,
10959            nh_k as i32,
10960        );
10961        let __s_b = self.gpu.stream();
10962        let mut b = __s_b.launch_builder(&f);
10963        match ff {
10964            Some(t) => {
10965                b.arg(q0)
10966                    .arg(k0)
10967                    .arg(v0)
10968                    .arg(wq)
10969                    .arg(wk)
10970                    .arg(wv)
10971                    .arg(&mut *q)
10972                    .arg(&mut *k)
10973                    .arg(&mut *v)
10974                    .arg(&nc)
10975                    .arg(&rqi)
10976                    .arg(&rki)
10977                    .arg(pos)
10978                    .arg(&nhq)
10979                    .arg(&nhk)
10980                    .arg(&theta_scale)
10981                    .arg(&freq_scale)
10982                    .arg(t)
10983                    .arg(&eps);
10984                unsafe {
10985                    b.launch(cfg)?;
10986                }
10987            }
10988            None => {
10989                let null: u64 = 0;
10990                b.arg(q0)
10991                    .arg(k0)
10992                    .arg(v0)
10993                    .arg(wq)
10994                    .arg(wk)
10995                    .arg(wv)
10996                    .arg(&mut *q)
10997                    .arg(&mut *k)
10998                    .arg(&mut *v)
10999                    .arg(&nc)
11000                    .arg(&rqi)
11001                    .arg(&rki)
11002                    .arg(pos)
11003                    .arg(&nhq)
11004                    .arg(&nhk)
11005                    .arg(&theta_scale)
11006                    .arg(&freq_scale)
11007                    .arg(&null)
11008                    .arg(&eps);
11009                unsafe {
11010                    b.launch(cfg)?;
11011                }
11012            }
11013        }
11014        Ok(())
11015    }
11016
11017    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
11018    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
11019    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
11020    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
11021    /// ([`Engine::full_width_rope_only`]).
11022    #[allow(clippy::too_many_arguments)]
11023    pub fn rms_norm_qkv_rope_append_dc(
11024        &self,
11025        q0: &CudaSlice<f32>,
11026        k0: &CudaSlice<f32>,
11027        v0: &CudaSlice<f32>,
11028        wq: &CudaSlice<f32>,
11029        wk: &CudaSlice<f32>,
11030        wv: &CudaSlice<f32>,
11031        q: &mut CudaSlice<f32>,
11032        k: &mut CudaSlice<f32>,
11033        v: &mut CudaSlice<f32>,
11034        head_dim: usize,
11035        n_rot: usize,
11036        rq: usize,
11037        rk: usize,
11038        pos: &CudaSlice<i32>,
11039        nh_q: usize,
11040        nh_k: usize,
11041        base: f32,
11042        freq_scale: f32,
11043        ff: Option<&CudaSlice<f32>>,
11044        eps: f32,
11045        kc: &mut CudaSlice<u8>,
11046        vc: &mut CudaSlice<u8>,
11047        t_dev: &CudaSlice<i32>,
11048        k_tok_bytes: usize,
11049        v_tok_bytes: usize,
11050        g: bool,
11051    ) -> Result<(), Box<dyn std::error::Error>> {
11052        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11053        let rows = rq + rk + rk;
11054        let theta_scale = base.powf(-2.0 / head_dim as f32);
11055        let (nc, rqi, rki, nhq, nhk) = (
11056            head_dim as i32,
11057            rq as i32,
11058            rk as i32,
11059            nh_q as i32,
11060            nh_k as i32,
11061        );
11062        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11063        if Self::pdl_on() && Self::pdl_wb_on() {
11064            use cudarc::driver::{DevicePtr, DevicePtrMut};
11065            let s = &self.gpu.stream();
11066            let (p0, _a0) = q0.device_ptr(s);
11067            let (p1, _a1) = k0.device_ptr(s);
11068            let (p2, _a2) = v0.device_ptr(s);
11069            let (pwq, _a3) = wq.device_ptr(s);
11070            let (pwk, _a4) = wk.device_ptr(s);
11071            let (pwv, _a5) = wv.device_ptr(s);
11072            let (pq, _a6) = q.device_ptr_mut(s);
11073            let (pk, _a7) = k.device_ptr_mut(s);
11074            let (pv, _a8) = v.device_ptr_mut(s);
11075            let (pp, _a9) = pos.device_ptr(s);
11076            let pff: u64 = match ff {
11077                Some(t) => {
11078                    let (p, _gg) = t.device_ptr(s);
11079                    p as u64
11080                }
11081                None => 0,
11082            };
11083            let (pkc, _a10) = kc.device_ptr_mut(s);
11084            let (pvc, _a11) = vc.device_ptr_mut(s);
11085            let (pt, _a12) = t_dev.device_ptr(s);
11086            let mut ps = [
11087                &p0 as *const _ as *mut std::ffi::c_void,
11088                &p1 as *const _ as *mut _,
11089                &p2 as *const _ as *mut _,
11090                &pwq as *const _ as *mut _,
11091                &pwk as *const _ as *mut _,
11092                &pwv as *const _ as *mut _,
11093                &pq as *const _ as *mut _,
11094                &pk as *const _ as *mut _,
11095                &pv as *const _ as *mut _,
11096                &nc as *const _ as *mut _,
11097                &rqi as *const _ as *mut _,
11098                &rki as *const _ as *mut _,
11099                &pp as *const _ as *mut _,
11100                &nhq as *const _ as *mut _,
11101                &nhk as *const _ as *mut _,
11102                &theta_scale as *const _ as *mut _,
11103                &freq_scale as *const _ as *mut _,
11104                &pff as *const _ as *mut _,
11105                &eps as *const _ as *mut _,
11106                &pkc as *const _ as *mut _,
11107                &pvc as *const _ as *mut _,
11108                &pt as *const _ as *mut _,
11109                &ktb as *const _ as *mut _,
11110                &vtb as *const _ as *mut _,
11111            ];
11112            unsafe {
11113                self.launch_pdl_flash(
11114                    g,
11115                    "rms_norm_qkv_rope_append_dc_f32",
11116                    (rows as u32, 1, 1),
11117                    (rms_block(), 1, 1),
11118                    0,
11119                    &mut ps,
11120                )?;
11121            }
11122            return Ok(());
11123        }
11124        let f = if g {
11125            self.func_g("rms_norm_qkv_rope_append_dc_f32")
11126        } else {
11127            self.func("rms_norm_qkv_rope_append_dc_f32")
11128        };
11129        let cfg = LaunchConfig {
11130            grid_dim: (rows as u32, 1, 1),
11131            block_dim: (rms_block(), 1, 1),
11132            shared_mem_bytes: 0,
11133        };
11134        let __s_b = self.gpu.stream();
11135        let mut b = __s_b.launch_builder(&f);
11136        match ff {
11137            Some(t) => {
11138                b.arg(q0)
11139                    .arg(k0)
11140                    .arg(v0)
11141                    .arg(wq)
11142                    .arg(wk)
11143                    .arg(wv)
11144                    .arg(&mut *q)
11145                    .arg(&mut *k)
11146                    .arg(&mut *v)
11147                    .arg(&nc)
11148                    .arg(&rqi)
11149                    .arg(&rki)
11150                    .arg(pos)
11151                    .arg(&nhq)
11152                    .arg(&nhk)
11153                    .arg(&theta_scale)
11154                    .arg(&freq_scale)
11155                    .arg(t)
11156                    .arg(&eps)
11157                    .arg(&mut *kc)
11158                    .arg(&mut *vc)
11159                    .arg(t_dev)
11160                    .arg(&ktb)
11161                    .arg(&vtb);
11162                unsafe {
11163                    b.launch(cfg)?;
11164                }
11165            }
11166            None => {
11167                let null: u64 = 0;
11168                b.arg(q0)
11169                    .arg(k0)
11170                    .arg(v0)
11171                    .arg(wq)
11172                    .arg(wk)
11173                    .arg(wv)
11174                    .arg(&mut *q)
11175                    .arg(&mut *k)
11176                    .arg(&mut *v)
11177                    .arg(&nc)
11178                    .arg(&rqi)
11179                    .arg(&rki)
11180                    .arg(pos)
11181                    .arg(&nhq)
11182                    .arg(&nhk)
11183                    .arg(&theta_scale)
11184                    .arg(&freq_scale)
11185                    .arg(&null)
11186                    .arg(&eps)
11187                    .arg(&mut *kc)
11188                    .arg(&mut *vc)
11189                    .arg(t_dev)
11190                    .arg(&ktb)
11191                    .arg(&vtb);
11192                unsafe {
11193                    b.launch(cfg)?;
11194                }
11195            }
11196        }
11197        Ok(())
11198    }
11199
11200    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
11201    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
11202    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
11203    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
11204    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
11205    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
11206    /// `head_dim` ([`Engine::full_width_rope_only`]).
11207    #[allow(clippy::too_many_arguments)]
11208    pub fn rms_norm_qkv_rope_append(
11209        &self,
11210        q0: &CudaSlice<f32>,
11211        k0: &CudaSlice<f32>,
11212        v0: &CudaSlice<f32>,
11213        wq: &CudaSlice<f32>,
11214        wk: &CudaSlice<f32>,
11215        wv: &CudaSlice<f32>,
11216        q: &mut CudaSlice<f32>,
11217        k: &mut CudaSlice<f32>,
11218        v: &mut CudaSlice<f32>,
11219        head_dim: usize,
11220        n_rot: usize,
11221        rq: usize,
11222        rk: usize,
11223        pos: &CudaSlice<i32>,
11224        nh_q: usize,
11225        nh_k: usize,
11226        base: f32,
11227        freq_scale: f32,
11228        ff: Option<&CudaSlice<f32>>,
11229        eps: f32,
11230        kc: &mut CudaSlice<u8>,
11231        vc: &mut CudaSlice<u8>,
11232        t: usize,
11233        k_tok_bytes: usize,
11234        v_tok_bytes: usize,
11235        g: bool,
11236    ) -> Result<(), Box<dyn std::error::Error>> {
11237        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11238        let rows = rq + rk + rk;
11239        let theta_scale = base.powf(-2.0 / head_dim as f32);
11240        let (nc, rqi, rki, nhq, nhk) = (
11241            head_dim as i32,
11242            rq as i32,
11243            rk as i32,
11244            nh_q as i32,
11245            nh_k as i32,
11246        );
11247        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11248        let ti = t as i32;
11249        if Self::pdl_on() && Self::pdl_wb_on() {
11250            use cudarc::driver::{DevicePtr, DevicePtrMut};
11251            let s = &self.gpu.stream();
11252            let (p0, _a0) = q0.device_ptr(s);
11253            let (p1, _a1) = k0.device_ptr(s);
11254            let (p2, _a2) = v0.device_ptr(s);
11255            let (pwq, _a3) = wq.device_ptr(s);
11256            let (pwk, _a4) = wk.device_ptr(s);
11257            let (pwv, _a5) = wv.device_ptr(s);
11258            let (pq, _a6) = q.device_ptr_mut(s);
11259            let (pk, _a7) = k.device_ptr_mut(s);
11260            let (pv, _a8) = v.device_ptr_mut(s);
11261            let (pp, _a9) = pos.device_ptr(s);
11262            let pff: u64 = match ff {
11263                Some(t) => {
11264                    let (p, _gg) = t.device_ptr(s);
11265                    p as u64
11266                }
11267                None => 0,
11268            };
11269            let (pkc, _a10) = kc.device_ptr_mut(s);
11270            let (pvc, _a11) = vc.device_ptr_mut(s);
11271            let mut ps = [
11272                &p0 as *const _ as *mut std::ffi::c_void,
11273                &p1 as *const _ as *mut _,
11274                &p2 as *const _ as *mut _,
11275                &pwq as *const _ as *mut _,
11276                &pwk as *const _ as *mut _,
11277                &pwv as *const _ as *mut _,
11278                &pq as *const _ as *mut _,
11279                &pk as *const _ as *mut _,
11280                &pv as *const _ as *mut _,
11281                &nc as *const _ as *mut _,
11282                &rqi as *const _ as *mut _,
11283                &rki as *const _ as *mut _,
11284                &pp as *const _ as *mut _,
11285                &nhq as *const _ as *mut _,
11286                &nhk as *const _ as *mut _,
11287                &theta_scale as *const _ as *mut _,
11288                &freq_scale as *const _ as *mut _,
11289                &pff as *const _ as *mut _,
11290                &eps as *const _ as *mut _,
11291                &pkc as *const _ as *mut _,
11292                &pvc as *const _ as *mut _,
11293                &ti as *const _ as *mut _,
11294                &ktb as *const _ as *mut _,
11295                &vtb as *const _ as *mut _,
11296            ];
11297            unsafe {
11298                self.launch_pdl_flash(
11299                    g,
11300                    "rms_norm_qkv_rope_append_f32",
11301                    (rows as u32, 1, 1),
11302                    (rms_block(), 1, 1),
11303                    0,
11304                    &mut ps,
11305                )?;
11306            }
11307            return Ok(());
11308        }
11309        let f = if g {
11310            self.func_g("rms_norm_qkv_rope_append_f32")
11311        } else {
11312            self.func("rms_norm_qkv_rope_append_f32")
11313        };
11314        let cfg = LaunchConfig {
11315            grid_dim: (rows as u32, 1, 1),
11316            block_dim: (rms_block(), 1, 1),
11317            shared_mem_bytes: 0,
11318        };
11319        let __s_b = self.gpu.stream();
11320        let mut b = __s_b.launch_builder(&f);
11321        let null: u64 = 0;
11322        b.arg(q0)
11323            .arg(k0)
11324            .arg(v0)
11325            .arg(wq)
11326            .arg(wk)
11327            .arg(wv)
11328            .arg(&mut *q)
11329            .arg(&mut *k)
11330            .arg(&mut *v)
11331            .arg(&nc)
11332            .arg(&rqi)
11333            .arg(&rki)
11334            .arg(pos)
11335            .arg(&nhq)
11336            .arg(&nhk)
11337            .arg(&theta_scale)
11338            .arg(&freq_scale);
11339        match ff {
11340            Some(t) => {
11341                b.arg(t);
11342            }
11343            None => {
11344                b.arg(&null);
11345            }
11346        }
11347        b.arg(&eps)
11348            .arg(&mut *kc)
11349            .arg(&mut *vc)
11350            .arg(&ti)
11351            .arg(&ktb)
11352            .arg(&vtb);
11353        unsafe {
11354            b.launch(cfg)?;
11355        }
11356        Ok(())
11357    }
11358
11359    pub fn add_q8_1(
11360        &self,
11361        a: &CudaSlice<f32>,
11362        b: &CudaSlice<f32>,
11363        res: &mut CudaSlice<f32>,
11364        ncols: usize,
11365        nrows: usize,
11366    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11367        debug_assert!(ncols % 128 == 0);
11368        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11369        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11370        let f = self.func("add_q8_1_f32");
11371        let cfg = LaunchConfig {
11372            grid_dim: (nrows as u32, 1, 1),
11373            block_dim: (rms_block(), 1, 1),
11374            shared_mem_bytes: 0,
11375        };
11376        let nc = ncols as i32;
11377        let __s_b2 = self.gpu.stream();
11378        let mut b2 = __s_b2.launch_builder(&f);
11379        b2.arg(a)
11380            .arg(b)
11381            .arg(&mut *res)
11382            .arg(&mut out_q)
11383            .arg(&mut out_d)
11384            .arg(&nc);
11385        unsafe {
11386            b2.launch(cfg)?;
11387        }
11388        Ok((out_q, out_d))
11389    }
11390
11391    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
11392    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
11393    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
11394    pub fn rms_pre_add_q8_1(
11395        &self,
11396        a: &CudaSlice<f32>,
11397        wa: &CudaSlice<f32>,
11398        b: &CudaSlice<f32>,
11399        res: &mut CudaSlice<f32>,
11400        ncols: usize,
11401        nrows: usize,
11402        eps: f32,
11403    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11404        debug_assert!(ncols % 128 == 0);
11405        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11406        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11407        let f = self.func("rms_pre_add_q8_1_f32");
11408        let cfg = LaunchConfig {
11409            grid_dim: (nrows as u32, 1, 1),
11410            block_dim: (rms_block(), 1, 1),
11411            shared_mem_bytes: 0,
11412        };
11413        let (nc, ep) = (ncols as i32, eps);
11414        let __s_b2 = self.gpu.stream();
11415        let mut b2 = __s_b2.launch_builder(&f);
11416        b2.arg(a)
11417            .arg(wa)
11418            .arg(b)
11419            .arg(&mut *res)
11420            .arg(&mut out_q)
11421            .arg(&mut out_d)
11422            .arg(&nc)
11423            .arg(&ep);
11424        unsafe {
11425            b2.launch(cfg)?;
11426        }
11427        Ok((out_q, out_d))
11428    }
11429
11430    /// L2 norm per row (head_dim), no weight.
11431    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
11432    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
11433    pub fn l2_v2_on(ncols: usize) -> bool {
11434        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11435    }
11436
11437    pub fn l2_norm_pp(
11438        &self,
11439        x: &CudaSlice<f32>,
11440        dst: &mut CudaSlice<f32>,
11441        dst16: Option<&mut CudaSlice<u8>>,
11442        ncols: usize,
11443        nrows: usize,
11444        eps: f32,
11445    ) -> Result<(), Box<dyn std::error::Error>> {
11446        if Self::l2_v2_on(ncols) {
11447            let f = self.func("l2_norm_pp_v2_f32");
11448            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
11449            let cfg = LaunchConfig {
11450                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11451                block_dim: (256, 1, 1),
11452                shared_mem_bytes: 0,
11453            };
11454            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11455            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
11456            let d16: u64 = match dst16 {
11457                Some(d) => self.addr_u8(d),
11458                None => 0,
11459            };
11460            let __s_b = self.gpu.stream();
11461            let mut b = __s_b.launch_builder(&f);
11462            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11463            unsafe {
11464                b.launch(cfg)?;
11465            }
11466            return Ok(());
11467        }
11468        self.l2_norm(x, dst, ncols, nrows, eps)
11469    }
11470
11471    pub fn l2_norm(
11472        &self,
11473        x: &CudaSlice<f32>,
11474        dst: &mut CudaSlice<f32>,
11475        ncols: usize,
11476        nrows: usize,
11477        eps: f32,
11478    ) -> Result<(), Box<dyn std::error::Error>> {
11479        let f = self.func("l2_norm_f32");
11480        let cfg = LaunchConfig {
11481            grid_dim: (nrows as u32, 1, 1),
11482            block_dim: (256, 1, 1),
11483            shared_mem_bytes: 0,
11484        };
11485        let (nc, e) = (ncols as i32, eps);
11486        let __s_b = self.gpu.stream();
11487        let mut b = __s_b.launch_builder(&f);
11488        b.arg(x).arg(dst).arg(&nc).arg(&e);
11489        unsafe {
11490            b.launch(cfg)?;
11491        }
11492        Ok(())
11493    }
11494
11495    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
11496    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
11497    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
11498    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
11499    /// propagate through gdn_scan and flip argmax on marginal logits.
11500    pub fn l2_norm_decode(
11501        &self,
11502        x: &CudaSlice<f32>,
11503        dst: &mut CudaSlice<f32>,
11504        ncols: usize,
11505        nrows: usize,
11506        eps: f32,
11507    ) -> Result<(), Box<dyn std::error::Error>> {
11508        let f = self.func("l2_norm_f32");
11509        let cfg = LaunchConfig {
11510            grid_dim: (nrows as u32, 1, 1),
11511            block_dim: (32, 1, 1),
11512            shared_mem_bytes: 0,
11513        };
11514        let (nc, e) = (ncols as i32, eps);
11515        let __s_b = self.gpu.stream();
11516        let mut b = __s_b.launch_builder(&f);
11517        b.arg(x).arg(dst).arg(&nc).arg(&e);
11518        unsafe {
11519            b.launch(cfg)?;
11520        }
11521        Ok(())
11522    }
11523
11524    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
11525    pub fn rope_neox(
11526        &self,
11527        x: &mut CudaSlice<f32>,
11528        pos: &CudaSlice<i32>,
11529        head_dim: usize,
11530        n_dims: usize,
11531        n_heads: usize,
11532        n_tokens: usize,
11533        freq_base: f32,
11534        freq_scale: f32,
11535    ) -> Result<(), Box<dyn std::error::Error>> {
11536        let f = self.func("rope_neox_f32");
11537        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11538        let grid = (n_heads * n_tokens) as u32;
11539        let cfg = LaunchConfig {
11540            grid_dim: (grid, 1, 1),
11541            block_dim: ((head_dim / 2) as u32, 1, 1),
11542            shared_mem_bytes: 0,
11543        };
11544        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11545        let __s_b = self.gpu.stream();
11546        let mut b = __s_b.launch_builder(&f);
11547        b.arg(x)
11548            .arg(pos)
11549            .arg(&hd)
11550            .arg(&nd)
11551            .arg(&nh)
11552            .arg(&theta_scale)
11553            .arg(&freq_scale);
11554        unsafe {
11555            b.launch(cfg)?;
11556        }
11557        Ok(())
11558    }
11559
11560    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
11561    pub fn rope_neox_ff(
11562        &self,
11563        x: &mut CudaSlice<f32>,
11564        pos: &CudaSlice<i32>,
11565        head_dim: usize,
11566        n_dims: usize,
11567        n_heads: usize,
11568        n_tokens: usize,
11569        freq_base: f32,
11570        freq_scale: f32,
11571        ff: &CudaSlice<f32>,
11572    ) -> Result<(), Box<dyn std::error::Error>> {
11573        let f = self.func("rope_neox_ff_f32");
11574        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11575        let grid = (n_heads * n_tokens) as u32;
11576        let cfg = LaunchConfig {
11577            grid_dim: (grid, 1, 1),
11578            block_dim: ((head_dim / 2) as u32, 1, 1),
11579            shared_mem_bytes: 0,
11580        };
11581        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11582        let __s_b = self.gpu.stream();
11583        let mut b = __s_b.launch_builder(&f);
11584        b.arg(x)
11585            .arg(pos)
11586            .arg(&hd)
11587            .arg(&nd)
11588            .arg(&nh)
11589            .arg(&theta_scale)
11590            .arg(&freq_scale)
11591            .arg(ff);
11592        unsafe {
11593            b.launch(cfg)?;
11594        }
11595        Ok(())
11596    }
11597
11598    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
11599    #[allow(clippy::too_many_arguments)]
11600    pub fn rope_neox2(
11601        &self,
11602        q: &mut CudaSlice<f32>,
11603        k: &mut CudaSlice<f32>,
11604        pos: &CudaSlice<i32>,
11605        head_dim: usize,
11606        n_dims: usize,
11607        nh_q: usize,
11608        nh_k: usize,
11609        n_tokens: usize,
11610        freq_base: f32,
11611        freq_scale: f32,
11612        ff: Option<&CudaSlice<f32>>,
11613    ) -> Result<(), Box<dyn std::error::Error>> {
11614        let f = self.func("rope_neox2_f32");
11615        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11616        let grid = ((nh_q + nh_k) * 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, nq, nk, nt) = (
11623            head_dim as i32,
11624            n_dims as i32,
11625            nh_q as i32,
11626            nh_k as i32,
11627            n_tokens as i32,
11628        );
11629        let __s_b = self.gpu.stream();
11630        let mut b = __s_b.launch_builder(&f);
11631        b.arg(q)
11632            .arg(k)
11633            .arg(pos)
11634            .arg(&hd)
11635            .arg(&nd)
11636            .arg(&nq)
11637            .arg(&nk)
11638            .arg(&nt)
11639            .arg(&theta_scale)
11640            .arg(&freq_scale);
11641        match ff {
11642            Some(ffv) => {
11643                b.arg(ffv);
11644                unsafe {
11645                    b.launch(cfg)?;
11646                }
11647            }
11648            None => {
11649                let null: u64 = 0;
11650                b.arg(&null);
11651                unsafe {
11652                    b.launch(cfg)?;
11653                }
11654            }
11655        }
11656        Ok(())
11657    }
11658
11659    /// gemma4 R1: dst = GELU_tanh(gate) * up.
11660    pub fn gelu_tanh_mul(
11661        &self,
11662        gate: &CudaSlice<f32>,
11663        up: &CudaSlice<f32>,
11664        dst: &mut CudaSlice<f32>,
11665        n: usize,
11666    ) -> Result<(), Box<dyn std::error::Error>> {
11667        let f = self.func("gelu_tanh_mul_f32");
11668        let cfg = LaunchConfig::for_num_elems(n as u32);
11669        let ni = n as i32;
11670        let __s_b = self.gpu.stream();
11671        let mut b = __s_b.launch_builder(&f);
11672        b.arg(gate).arg(up).arg(dst).arg(&ni);
11673        unsafe {
11674            b.launch(cfg)?;
11675        }
11676        Ok(())
11677    }
11678
11679    pub fn silu_mul(
11680        &self,
11681        gate: &CudaSlice<f32>,
11682        up: &CudaSlice<f32>,
11683        dst: &mut CudaSlice<f32>,
11684        n: usize,
11685    ) -> Result<(), Box<dyn std::error::Error>> {
11686        let f = self.func("silu_mul_f32");
11687        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11688        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11689        let ni = n as i32;
11690        let __s_b = self.gpu.stream();
11691        let mut b = __s_b.launch_builder(&f);
11692        b.arg(gate).arg(up).arg(dst).arg(&ni);
11693        unsafe {
11694            b.launch(cfg)?;
11695        }
11696        Ok(())
11697    }
11698
11699    /// SwiGLU twin using Memra's host-matching expf transcription.
11700    pub fn silu_mul_host_expf(
11701        &self,
11702        gate: &CudaSlice<f32>,
11703        up: &CudaSlice<f32>,
11704        dst: &mut CudaSlice<f32>,
11705        n: usize,
11706    ) -> Result<(), Box<dyn std::error::Error>> {
11707        let f = self.func("silu_mul_host_expf_f32");
11708        let cfg = LaunchConfig::for_num_elems(n as u32);
11709        let ni = n as i32;
11710        let __s_b = self.gpu.stream();
11711        let mut b = __s_b.launch_builder(&f);
11712        b.arg(gate).arg(up).arg(dst).arg(&ni);
11713        unsafe {
11714            b.launch(cfg)?;
11715        }
11716        Ok(())
11717    }
11718
11719    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
11720    pub fn silu_clamped_mul_host_expf(
11721        &self,
11722        gate: &CudaSlice<f32>,
11723        up: &CudaSlice<f32>,
11724        limit: f32,
11725        dst: &mut CudaSlice<f32>,
11726        n: usize,
11727    ) -> Result<(), Box<dyn std::error::Error>> {
11728        if !limit.is_finite() || limit <= 0.0 {
11729            return Err(
11730                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11731            );
11732        }
11733        let f = self.func("silu_clamped_mul_host_expf_f32");
11734        let cfg = LaunchConfig::for_num_elems(n as u32);
11735        let ni = n as i32;
11736        let __s_b = self.gpu.stream();
11737        let mut b = __s_b.launch_builder(&f);
11738        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11739        unsafe {
11740            b.launch(cfg)?;
11741        }
11742        Ok(())
11743    }
11744
11745    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
11746    /// for the down projection — kills the standalone convert pass. Bit-identical class.
11747    pub fn silu_mul_f16out(
11748        &self,
11749        gate: &CudaSlice<f32>,
11750        up: &CudaSlice<f32>,
11751        dst: &mut CudaSlice<f32>,
11752        dst16: &mut CudaSlice<u8>,
11753        n: usize,
11754    ) -> Result<(), Box<dyn std::error::Error>> {
11755        let f = self.func("silu_mul_f16out_f32");
11756        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11757        let ni = n as i32;
11758        let __s_b = self.gpu.stream();
11759        let mut b = __s_b.launch_builder(&f);
11760        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11761        unsafe {
11762            b.launch(cfg)?;
11763        }
11764        Ok(())
11765    }
11766
11767    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
11768    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
11769    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
11770    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
11771    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
11772    /// launches per dense FFN layer (the gate+up post-matmul scales).
11773    pub fn silu_mul_scaled(
11774        &self,
11775        gate: &CudaSlice<f32>,
11776        up: &CudaSlice<f32>,
11777        gs: f32,
11778        us: f32,
11779        dst: &mut CudaSlice<f32>,
11780        n: usize,
11781    ) -> Result<(), Box<dyn std::error::Error>> {
11782        let f = self.func("silu_mul_scaled_f32");
11783        let cfg = LaunchConfig::for_num_elems(n as u32);
11784        let ni = n as i32;
11785        let (gsf, usf) = (gs, us);
11786        let __s_b = self.gpu.stream();
11787        let mut b = __s_b.launch_builder(&f);
11788        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11789        unsafe {
11790            b.launch(cfg)?;
11791        }
11792        Ok(())
11793    }
11794
11795    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
11796    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
11797    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
11798    #[allow(clippy::too_many_arguments)]
11799    pub fn swigluoai_mul_scaled(
11800        &self,
11801        gate: &CudaSlice<f32>,
11802        up: &CudaSlice<f32>,
11803        gs: f32,
11804        us: f32,
11805        alpha: f32,
11806        limit: f32,
11807        dst: &mut CudaSlice<f32>,
11808        n: usize,
11809    ) -> Result<(), Box<dyn std::error::Error>> {
11810        let f = self.func("swigluoai_mul_scaled_f32");
11811        let cfg = LaunchConfig::for_num_elems(n as u32);
11812        let ni = n as i32;
11813        let __s_b = self.gpu.stream();
11814        let mut b = __s_b.launch_builder(&f);
11815        b.arg(gate)
11816            .arg(up)
11817            .arg(&gs)
11818            .arg(&us)
11819            .arg(&alpha)
11820            .arg(&limit)
11821            .arg(dst)
11822            .arg(&ni);
11823        unsafe {
11824            b.launch(cfg)?;
11825        }
11826        Ok(())
11827    }
11828
11829    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11830    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11831    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11832    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11833    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11834    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11835    /// n must be a multiple of 32 (n_ff always is).
11836    pub fn silu_mul_scaled_q8_1(
11837        &self,
11838        gate: &CudaSlice<f32>,
11839        up: &CudaSlice<f32>,
11840        gs: f32,
11841        us: f32,
11842        n: usize,
11843    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11844        let f = self.func("silu_mul_scaled_q8_1");
11845        let nblk = n / 32;
11846        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11847        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11848        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11849        let cfg = LaunchConfig::for_num_elems(n as u32);
11850        let (gsf, usf, ni) = (gs, us, n as i32);
11851        let __s_b = self.gpu.stream();
11852        let mut b = __s_b.launch_builder(&f);
11853        b.arg(gate)
11854            .arg(up)
11855            .arg(&gsf)
11856            .arg(&usf)
11857            .arg(&mut aq)
11858            .arg(&mut ad)
11859            .arg(&ni);
11860        unsafe {
11861            b.launch(cfg)?;
11862        }
11863        Ok((aq, ad))
11864    }
11865
11866    pub fn add(
11867        &self,
11868        a: &CudaSlice<f32>,
11869        b_in: &CudaSlice<f32>,
11870        dst: &mut CudaSlice<f32>,
11871        n: usize,
11872    ) -> Result<(), Box<dyn std::error::Error>> {
11873        let f = self.func("add_f32");
11874        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11875        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11876        let ni = n as i32;
11877        let __s_bld = self.gpu.stream();
11878        let mut bld = __s_bld.launch_builder(&f);
11879        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11880        unsafe {
11881            bld.launch(cfg)?;
11882        }
11883        Ok(())
11884    }
11885
11886    pub fn mul(
11887        &self,
11888        a: &CudaSlice<f32>,
11889        b_in: &CudaSlice<f32>,
11890        dst: &mut CudaSlice<f32>,
11891        n: usize,
11892    ) -> Result<(), Box<dyn std::error::Error>> {
11893        let f = self.func("mul_f32");
11894        let cfg = LaunchConfig::for_num_elems(n as u32);
11895        let ni = n as i32;
11896        let __s_bld = self.gpu.stream();
11897        let mut bld = __s_bld.launch_builder(&f);
11898        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11899        unsafe {
11900            bld.launch(cfg)?;
11901        }
11902        Ok(())
11903    }
11904
11905    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11906    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11907    pub fn matmul(
11908        &self,
11909        w: &crate::model::GpuTensor,
11910        x: &CudaSlice<f32>,
11911        m: usize,
11912    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11913        use crate::model::GpuTensor;
11914        let in_f = w.in_features();
11915        let out_f = w.out_features();
11916        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11917        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11918        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11919        // gives nothing). Quantize the activation once here then call the GEMM.
11920        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11921        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11922        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11923        #[allow(non_snake_case)]
11924        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11925        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11926        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11927            usize::MAX
11928        } else {
11929            16usize
11930        };
11931
11932        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11933        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11934        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11935        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11936        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11937        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11938        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11939        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11940        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11941        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11942        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11943        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11944        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11945        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11946        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11947        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11948        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11949        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11950        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11951        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11952        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11953        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11954        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11955        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11956        if m >= GEMM_M_THRESHOLD {
11957            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11958                return Ok(y);
11959            }
11960            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
11961            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
11962            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
11963            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
11964            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
11965            // tile defaults differently by operand source.
11966            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11967                return Ok(y);
11968            }
11969            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
11970            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
11971            if let Some(y) = self.try_f16_gemm(w, x, m)? {
11972                return Ok(y);
11973            }
11974        }
11975        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
11976        // m threshold the rest of this method uses:
11977        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
11978        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
11979        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
11980        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
11981        //     across every tier by construction with no batched twin needed.
11982        //
11983        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
11984        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
11985        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
11986        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
11987        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
11988        // arms is what makes sure it never gets there.
11989        if let GpuTensor::Quant { qtype, .. } = w {
11990            if *qtype == QT_F8_E4M3_BLK {
11991                if m >= GEMM_M_THRESHOLD {
11992                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11993                        return Ok(y);
11994                    }
11995                }
11996                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11997                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11998                    return Ok(y);
11999                }
12000            }
12001        }
12002        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12003            return self.qmatvec_mmq(w, x, m);
12004        }
12005        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12006            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12007            return self.qmatvec_gemm(w, &aq, &ad, m);
12008        }
12009        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
12010        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
12011        if m >= GEMM_M_THRESHOLD {
12012            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12013                return Ok(y);
12014            }
12015        }
12016        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
12017        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
12018        // to Stage-A f32-dequant (the correctness oracle path).
12019        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12020        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
12021        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
12022        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
12023        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
12024        if m == 1 && fast {
12025            if let GpuTensor::Quant {
12026                bytes,
12027                qtype,
12028                row_bytes,
12029                rp,
12030                rp4,
12031                scale,
12032                ..
12033            } = w
12034            {
12035                if self.mmvq_supports(*qtype) {
12036                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
12037                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
12038                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
12039                    let (bytes, rp) = match rp4 {
12040                        Some(m4) => (m4, true),
12041                        None => (bytes, *rp),
12042                    };
12043                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12044                    return self.qmatvec_mmvq(
12045                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12046                    );
12047                }
12048            }
12049        }
12050        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
12051        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
12052        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
12053        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
12054        // block below. MEMRA_NO_BATCHED -> per-m path.
12055        //
12056        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
12057        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
12058        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
12059        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
12060        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
12061        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
12062        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
12063        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
12064        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
12065        if (2..=16).contains(&m)
12066            && fast
12067            && std::env::var("MEMRA_NO_BATCHED").is_err()
12068            && (m <= 4 || Self::b8_enabled())
12069        {
12070            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
12071            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
12072            // is present (rp4) — the mirror pick below then routes to the _rp family.
12073            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
12074            // because the native e4m3 row layout is already aligned and needs no mirror.
12075            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
12076            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
12077            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
12078            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
12079            let m_ok = m <= 8
12080                || matches!(w, GpuTensor::Quant { qtype, .. }
12081                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12082                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12083            if m_ok {
12084                if let GpuTensor::Quant {
12085                    bytes,
12086                    qtype,
12087                    row_bytes,
12088                    rp,
12089                    rp4,
12090                    ..
12091                } = w
12092                {
12093                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12094                        let (bytes, rp) = match rp4 {
12095                            Some(m4) => (m4, true),
12096                            None => (bytes, *rp),
12097                        };
12098                        let mcols = Self::batched_mcols(m);
12099                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12100                        let mut y = self.qmatvec_mmvq_batched(
12101                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12102                        )?;
12103                        if let GpuTensor::Quant { scale, .. } = w {
12104                            if *scale != 1.0 {
12105                                self.scale_inplace(&mut y, *scale, m * out_f)?;
12106                            }
12107                        }
12108                        return Ok(y);
12109                    }
12110                }
12111            }
12112        }
12113        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
12114        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
12115        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
12116        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
12117        // for this dtype, so the generic match below must never see it under `fast`.
12118        if fast {
12119            if let GpuTensor::Quant {
12120                bytes,
12121                qtype,
12122                row_bytes,
12123                scale,
12124                ..
12125            } = w
12126            {
12127                if *qtype == QT_F8_E4M3 {
12128                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12129                    return self.qmatvec_mmvq(
12130                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12131                    );
12132                }
12133            }
12134        }
12135        let mut y = match w {
12136            GpuTensor::Quant {
12137                bytes,
12138                qtype,
12139                row_bytes,
12140                ..
12141            } if fast && *qtype == QT_Q8_0 => {
12142                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12143            }
12144            GpuTensor::Quant {
12145                bytes,
12146                qtype,
12147                row_bytes,
12148                ..
12149            } if fast && *qtype == QT_Q4_K => {
12150                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12151            }
12152            GpuTensor::Quant {
12153                bytes,
12154                qtype,
12155                row_bytes,
12156                ..
12157            } if fast && *qtype == QT_Q6_K => {
12158                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12159            }
12160            GpuTensor::Quant {
12161                bytes,
12162                qtype,
12163                row_bytes,
12164                ..
12165            } if fast && *qtype == QT_Q5_K => {
12166                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12167            }
12168            GpuTensor::Quant {
12169                bytes,
12170                qtype,
12171                row_bytes,
12172                ..
12173            } if fast && *qtype == QT_Q3_K => {
12174                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12175            }
12176            GpuTensor::Quant {
12177                bytes,
12178                qtype,
12179                row_bytes,
12180                rp,
12181                ..
12182            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12183                if *rp {
12184                    "qmatvec_nvfp4_dp4a_rp"
12185                } else {
12186                    "qmatvec_nvfp4_dp4a"
12187                },
12188                &bytes.slice(0..bytes.len()),
12189                x,
12190                m,
12191                in_f,
12192                out_f,
12193                *row_bytes,
12194            )?,
12195            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
12196            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
12197            // anomaly (research/kat-anomaly-20260802/).
12198            GpuTensor::Quant {
12199                bytes,
12200                qtype,
12201                row_bytes,
12202                ..
12203            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12204                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12205            }
12206            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
12207            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
12208            // without first writing the matching kernel, or func() will panic
12209            // "kernel ... not in any fatbin".
12210            GpuTensor::Quant {
12211                bytes,
12212                qtype,
12213                row_bytes,
12214                rp,
12215                ..
12216            } =>
12217            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
12218            // deq(row,j) form cannot address the planes; same value/product order).
12219            {
12220                self.qmatvec(
12221                    bytes,
12222                    x,
12223                    m,
12224                    in_f,
12225                    out_f,
12226                    if *rp && *qtype == QT_NVFP4 {
12227                        QT_NVFP4_RP
12228                    } else {
12229                        *qtype
12230                    },
12231                    *row_bytes,
12232                )?
12233            }
12234            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12235            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
12236            // cuBLASLt f32 GEMV as the Float arm.
12237            GpuTensor::FloatBf16 { data, .. } => {
12238                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
12239                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
12240                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
12241                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
12242                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
12243                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
12244                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12245                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12246                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12247                    y
12248                } else {
12249                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12250                }
12251            }
12252        };
12253        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
12254        if let GpuTensor::Quant { scale, .. } = w {
12255            if *scale != 1.0 {
12256                self.scale_inplace(&mut y, *scale, m * out_f)?;
12257            }
12258        }
12259        Ok(y)
12260    }
12261
12262    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
12263    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
12264    ///
12265    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
12266    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
12267    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
12268    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
12269    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
12270    /// path must not pay an env lookup for a flag that is off.
12271    pub fn stage_a_raw_needed() -> bool {
12272        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12273        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12274    }
12275
12276    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
12277    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
12278    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12279        use crate::model::GpuTensor;
12280        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12281            return false;
12282        }
12283        match w {
12284            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
12285            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
12286            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
12287            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
12288            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
12289            // block class has no fused twin yet, so each of its projections takes its own launch.
12290            GpuTensor::Quant { qtype, .. } => {
12291                matches!(
12292                    *qtype,
12293                    QT_Q8_0
12294                        | QT_Q4_K
12295                        | QT_Q6_K
12296                        | QT_Q5_K
12297                        | QT_Q3_K
12298                        | QT_NVFP4
12299                        | QT_F8_E4M3
12300                        | QT_F8_E4M3_BLK
12301                        | QT_Q4_0
12302                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12303            }
12304            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12305        }
12306    }
12307
12308    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
12309    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
12310    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
12311    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
12312    pub fn matmul_pre(
12313        &self,
12314        w: &crate::model::GpuTensor,
12315        aq: &CudaSlice<i8>,
12316        ad: &CudaSlice<f32>,
12317        x_fallback: &CudaSlice<f32>,
12318        m: usize,
12319    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12320        use crate::model::GpuTensor;
12321        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
12322        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
12323        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
12324        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
12325        // rc=30013 dig, 2026-07-31).
12326        let x_raw_ok = x_fallback.len() >= m * w.in_features();
12327        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
12328        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
12329        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12330            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12331                return Ok(y);
12332            }
12333            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
12334            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
12335            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12336                return Ok(y);
12337            }
12338            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
12339            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12340                return Ok(y);
12341            }
12342        }
12343        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
12344        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
12345        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
12346        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
12347        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
12348        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12349            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12350                return Ok(y);
12351            }
12352        }
12353        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12354            return Ok(y);
12355        }
12356        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
12357        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
12358        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
12359        // aq/ad.
12360        if m >= 16
12361            && w.out_features() >= 128
12362            && self.mmq_supports(w)
12363            && !self.verify_exact_on()
12364            && x_raw_ok
12365        {
12366            return self.qmatvec_mmq(w, x_fallback, m);
12367        }
12368        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
12369        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
12370        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12371            if let Some(y) =
12372                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12373            {
12374                return Ok(y);
12375            }
12376        }
12377        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
12378        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
12379        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12380            return self.qmatvec_gemm(w, aq, ad, m);
12381        }
12382        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
12383        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
12384        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
12385        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
12386        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
12387        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
12388        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
12389        // which reads `m * in_f` floats out of a 0-byte allocation ->
12390        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
12391        // it poisons the context, so every LATER request in that process fails with an unrelated
12392        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
12393        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
12394        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
12395        // dense artifact and left the arm with no working truth instrument.
12396        //
12397        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
12398        // strictly better than an illegal address surfacing later at an unrelated sync point, and
12399        // an oracle that cannot run must say so rather than corrupt the context it runs in.
12400        if !self.uses_q8_1_fast(w) {
12401            if !x_raw_ok {
12402                return Err(format!(
12403                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12404                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12405                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
12406                     activation (see Engine::rms_norm_decode, which is bit-identical to \
12407                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12408                    x_fallback.len(),
12409                    m,
12410                    w.in_features(),
12411                    m * w.in_features()
12412                )
12413                .into());
12414            }
12415            return self.matmul(w, x_fallback, m);
12416        }
12417        let in_f = w.in_features();
12418        let out_f = w.out_features();
12419        let (bytes, qtype, row_bytes, scale, rp) = match w {
12420            GpuTensor::Quant {
12421                bytes,
12422                qtype,
12423                row_bytes,
12424                scale,
12425                rp,
12426                ..
12427            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12428            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12429        };
12430        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
12431        // the dp4a/oracle tails below keep the raw GGUF bytes.
12432        let (mbytes, mrp) = match w {
12433            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12434            _ => (bytes, rp),
12435        };
12436        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
12437        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
12438        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
12439        if m == 1 && self.mmvq_supports(qtype) {
12440            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12441        }
12442        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
12443        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
12444        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
12445        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
12446        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
12447        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
12448        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
12449        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
12450        // m=5..8 on the old per-m path (b8-tier-only seam).
12451        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
12452        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
12453        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
12454        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12455            && std::env::var("MEMRA_NO_BATCHED").is_err()
12456            && (m <= 4 || Self::b8_enabled())
12457            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
12458            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
12459            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
12460            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12461                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12462        {
12463            let mcols = Self::batched_mcols(m);
12464            return self.qmatvec_mmvq_batched(
12465                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12466            );
12467        }
12468        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
12469        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
12470        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
12471        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
12472        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
12473        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12474            let (b2, r2) = if qtype == QT_Q4_0 {
12475                (mbytes, mrp)
12476            } else {
12477                (bytes, rp)
12478            };
12479            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12480        }
12481        let name = match qtype {
12482            QT_Q8_0 => "qmatvec_q8_0_dp4a",
12483            QT_Q4_K => "qmatvec_q4_K_dp4a",
12484            QT_Q6_K => "qmatvec_q6_K_dp4a",
12485            QT_Q5_K => "qmatvec_q5_K_dp4a",
12486            QT_Q3_K => "qmatvec_q3_K_dp4a",
12487            QT_NVFP4 => {
12488                if rp {
12489                    "qmatvec_nvfp4_dp4a_rp"
12490                } else {
12491                    "qmatvec_nvfp4_dp4a"
12492                }
12493            }
12494            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12495            _ => unreachable!(),
12496        };
12497        let f = self.func(name);
12498        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
12499        let cfg = LaunchConfig {
12500            grid_dim: (out_f as u32, m as u32, 1),
12501            block_dim: (128, 1, 1),
12502            shared_mem_bytes: 0,
12503        };
12504        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12505        let __s_b = self.gpu.stream();
12506        let mut b = __s_b.launch_builder(&f);
12507        b.arg(bytes)
12508            .arg(aq)
12509            .arg(ad)
12510            .arg(&mut y)
12511            .arg(&inf)
12512            .arg(&outf)
12513            .arg(&mi)
12514            .arg(&rb);
12515        unsafe {
12516            b.launch(cfg)?;
12517        }
12518        if scale != 1.0 {
12519            self.scale_inplace(&mut y, scale, m * out_f)?;
12520        }
12521        Ok(y)
12522    }
12523
12524    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
12525    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
12526    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
12527    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
12528    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
12529    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
12530    /// reduce as m=1); this method just forces that path unconditionally.
12531    pub fn matmul_decode_exact(
12532        &self,
12533        w: &crate::model::GpuTensor,
12534        x: &CudaSlice<f32>,
12535        m: usize,
12536    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12537        use crate::model::GpuTensor;
12538        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
12539        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
12540        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
12541        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
12542        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
12543        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
12544        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
12545        if let GpuTensor::Float { data, .. } = w {
12546            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12547        }
12548        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
12549        // float linear (same n-independent reduction contract as the Float arm above).
12550        if let GpuTensor::FloatBf16 { data, .. } = w {
12551            let (in_f, out_f) = (w.in_features(), w.out_features());
12552            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
12553            // contract — the whole-weight f32 dequant disappears too).
12554            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12555                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12556                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12557                return Ok(y);
12558            }
12559            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12560        }
12561        if !self.uses_q8_1_fast(w) {
12562            return self.matmul(w, x, m);
12563        }
12564        let in_f = w.in_features();
12565        let out_f = w.out_features();
12566        let (bytes, qtype, row_bytes, scale, rp) = match w {
12567            GpuTensor::Quant {
12568                bytes,
12569                qtype,
12570                row_bytes,
12571                scale,
12572                rp,
12573                ..
12574            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12575            _ => return self.matmul(w, x, m),
12576        };
12577        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
12578        // which does its own mirror pick).
12579        let (bytes, rp) = match w {
12580            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12581            _ => (bytes, rp),
12582        };
12583        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12584        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
12585        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
12586        // (token,row) by construction, which is exactly what this method exists to guarantee.
12587        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12588            return Ok(y);
12589        }
12590        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
12591        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
12592        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
12593        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
12594        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
12595        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
12596        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
12597        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
12598        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12599            && std::env::var("MEMRA_NO_BATCHED").is_err()
12600            && (m <= 4 || Self::b8_enabled())
12601            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
12602            // no mirror precondition, `rp` selects the layout only.
12603            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12604                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12605        {
12606            let mcols = Self::batched_mcols(m);
12607            return self.qmatvec_mmvq_batched(
12608                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12609            );
12610        }
12611        if self.mmvq_supports(qtype) {
12612            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
12613            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
12614            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12615        }
12616        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
12617        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
12618        self.matmul_pre(w, &aq, &ad, x, m)
12619    }
12620
12621    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
12622    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
12623    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
12624    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
12625    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
12626    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
12627    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
12628    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
12629    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
12630    pub fn matmul_decode_exact_pre(
12631        &self,
12632        w: &crate::model::GpuTensor,
12633        aq: &CudaSlice<i8>,
12634        ad: &CudaSlice<f32>,
12635        m: usize,
12636    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12637        use crate::model::GpuTensor;
12638        debug_assert!(
12639            self.uses_q8_1_fast(w),
12640            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12641        );
12642        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
12643        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12644            return Ok(y);
12645        }
12646        let in_f = w.in_features();
12647        let out_f = w.out_features();
12648        let (bytes, qtype, row_bytes, scale, rp) = match w {
12649            GpuTensor::Quant {
12650                bytes,
12651                qtype,
12652                row_bytes,
12653                scale,
12654                rp,
12655                ..
12656            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12657            _ => {
12658                return Err(
12659                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12660                );
12661            }
12662        };
12663        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
12664        let (bytes, rp) = match w {
12665            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12666            _ => (bytes, rp),
12667        };
12668        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
12669        if (2..=16).contains(&m)
12670            && self.batched_supports(qtype)
12671            && self.mmvq_supports(qtype)
12672            && std::env::var("MEMRA_NO_BATCHED").is_err()
12673            && (m <= 4 || Self::b8_enabled())
12674            && (m <= 8
12675                || qtype == QT_Q4_0
12676                || qtype == QT_Q6_K
12677                || qtype == QT_F8_E4M3
12678                || qtype == QT_NVFP4
12679                || qtype == QT_Q4_K
12680                || qtype == QT_Q5_K
12681                || qtype == QT_Q8_0)
12682        {
12683            let mcols = Self::batched_mcols(m);
12684            return self.qmatvec_mmvq_batched(
12685                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12686            );
12687        }
12688        if self.mmvq_supports(qtype) {
12689            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12690        }
12691        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
12692        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
12693        let x0 = self.zeros(0)?;
12694        self.matmul_pre(w, aq, ad, &x0, m)
12695    }
12696
12697    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
12698    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
12699    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
12700    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
12701    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
12702    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
12703    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
12704    /// per-tensor path.
12705    pub fn matmul_decode_exact_dual_pre(
12706        &self,
12707        w0: &crate::model::GpuTensor,
12708        w1: &crate::model::GpuTensor,
12709        aq: &CudaSlice<i8>,
12710        ad: &CudaSlice<f32>,
12711        m: usize,
12712    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12713    {
12714        use crate::model::GpuTensor;
12715        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12716        let on = *ON.get_or_init(|| {
12717            std::env::var("MEMRA_SPEC_DUAL_T")
12718                .map(|v| v != "0")
12719                .unwrap_or(true)
12720        });
12721        if !on
12722            || !(2..=7).contains(&m)
12723            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12724            || !self.uses_q8_1_fast(w0)
12725            || !self.uses_q8_1_fast(w1)
12726        {
12727            return Ok(None);
12728        }
12729        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
12730        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
12731        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
12732        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
12733        if !self.mmvq_supports(QT_NVFP4) {
12734            return Ok(None);
12735        }
12736        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12737        if w1.in_features() != in_f || w1.out_features() != out_f {
12738            return Ok(None);
12739        }
12740        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12741            (
12742                GpuTensor::Quant {
12743                    bytes: b0,
12744                    qtype: q0,
12745                    row_bytes: rb0,
12746                    scale: s0,
12747                    rp: rp0,
12748                    rp4: None,
12749                    ..
12750                },
12751                GpuTensor::Quant {
12752                    bytes: b1,
12753                    qtype: q1,
12754                    row_bytes: rb1,
12755                    scale: s1,
12756                    rp: rp1,
12757                    rp4: None,
12758                    ..
12759                },
12760            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12761                (b0, b1, *rb0, *s0, *s1, *rp0)
12762            }
12763            _ => return Ok(None),
12764        };
12765        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
12766        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
12767        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12768        {
12769            return Ok(None);
12770        }
12771        let (y0, y1) =
12772            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12773        Ok(Some(((y0, s0), (y1, s1))))
12774    }
12775
12776    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
12777    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
12778    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
12779    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
12780    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
12781    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
12782    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
12783    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
12784    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
12785    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
12786    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
12787    pub fn matmul_decode_exact_group4_pre(
12788        &self,
12789        ws: [&crate::model::GpuTensor; 4],
12790        aq: &CudaSlice<i8>,
12791        ad: &CudaSlice<f32>,
12792        m: usize,
12793    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12794        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12795        let on = *ON.get_or_init(|| {
12796            std::env::var("MEMRA_TK_GDN_GROUP")
12797                .map(|v| v != "0")
12798                .unwrap_or(true)
12799        });
12800        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12801    }
12802
12803    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
12804    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
12805    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
12806    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
12807    pub fn matmul_decode_exact_group3_pre(
12808        &self,
12809        ws: [&crate::model::GpuTensor; 3],
12810        aq: &CudaSlice<i8>,
12811        ad: &CudaSlice<f32>,
12812        m: usize,
12813    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12814        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12815        let on = *ON.get_or_init(|| {
12816            std::env::var("MEMRA_TK_FA_GROUP")
12817                .map(|v| v != "0")
12818                .unwrap_or(true)
12819        });
12820        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12821    }
12822
12823    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
12824    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
12825    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
12826    fn matmul_decode_exact_group_pre(
12827        &self,
12828        ws: &[&crate::model::GpuTensor],
12829        aq: &CudaSlice<i8>,
12830        ad: &CudaSlice<f32>,
12831        m: usize,
12832        on: bool,
12833        tag: &'static str,
12834    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12835        use crate::model::GpuTensor;
12836        if !on
12837            || !(2..=16).contains(&m)
12838            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12839            || (m > 4 && !Self::b8_enabled())
12840            || !self.mmvq_supports(QT_NVFP4)
12841            || !self.batched_supports(QT_NVFP4)
12842        {
12843            return Ok(None);
12844        }
12845        let in_f = ws[0].in_features();
12846        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12847        for w in ws {
12848            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12849                return Ok(None);
12850            }
12851            match w {
12852                GpuTensor::Quant {
12853                    bytes,
12854                    qtype,
12855                    scale,
12856                    rp: true,
12857                    rp4: None,
12858                    ..
12859                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12860                    parts.push((bytes, w.out_features(), *scale));
12861                }
12862                _ => return Ok(None),
12863            }
12864        }
12865        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12866        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12867        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12868        let mcols = if (5..=7).contains(&m) && b567 {
12869            m
12870        } else {
12871            Self::batched_mcols(m)
12872        };
12873        let kname: &'static str = match mcols {
12874            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12875            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12876            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12877            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12878            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12879            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12880            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12881            _ => return Ok(None),
12882        };
12883        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12884        // the second door's print on the slice-D battery — key the once-set by tag.
12885        if std::env::var("MEMRA_DEBUG").is_ok() {
12886            use std::sync::Mutex;
12887            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12888            let mut seen = SEEN.lock().unwrap();
12889            if !seen.contains(&tag) {
12890                seen.push(tag);
12891                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12892            }
12893        }
12894        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12895        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12896        let total: usize = parts.iter().map(|p| p.1).sum();
12897        let three = parts.len() == 3;
12898        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12899        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12900        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12901        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12902        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12903        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12904        let cfg = LaunchConfig {
12905            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12906            block_dim: (32, ROWS_PER_BLOCK, 1),
12907            shared_mem_bytes: 0,
12908        };
12909        let (inf, mi) = (in_f as i32, m as i32);
12910        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12911        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12912        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12913        let s3 = if three { 1.0f32 } else { parts[3].2 };
12914        let w3 = if three { parts[0].0 } else { parts[3].0 };
12915        let f = self.func(kname);
12916        let __s_b = self.gpu.stream();
12917        let mut b = __s_b.launch_builder(&f);
12918        b.arg(parts[0].0)
12919            .arg(parts[1].0)
12920            .arg(parts[2].0)
12921            .arg(w3)
12922            .arg(aq)
12923            .arg(ad)
12924            .arg(&mut y0)
12925            .arg(&mut y1)
12926            .arg(&mut y2)
12927            .arg(&mut y3)
12928            .arg(&inf)
12929            .arg(&n0)
12930            .arg(&n1)
12931            .arg(&n2)
12932            .arg(&n3)
12933            .arg(&mi)
12934            .arg(&s0)
12935            .arg(&s1)
12936            .arg(&s2)
12937            .arg(&s3);
12938        unsafe {
12939            b.launch(cfg)?;
12940        }
12941        Ok(Some(if three {
12942            vec![y0, y1, y2]
12943        } else {
12944            vec![y0, y1, y2, y3]
12945        }))
12946    }
12947
12948    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12949    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12950    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12951    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12952    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12953    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12954    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12955    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12956    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12957    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12958    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
12959    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
12960    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
12961    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
12962    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
12963    pub fn matmul_decode_exact_dual(
12964        &self,
12965        w0: &crate::model::GpuTensor,
12966        w1: &crate::model::GpuTensor,
12967        x: &CudaSlice<f32>,
12968        m: usize,
12969    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12970        use crate::model::GpuTensor;
12971        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12972        let on = *ON.get_or_init(|| {
12973            std::env::var("MEMRA_SPEC_DUAL_T")
12974                .map(|v| v != "0")
12975                .unwrap_or(true)
12976        });
12977        if !on
12978            || !(2..=4).contains(&m)
12979            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12980            || !self.uses_q8_1_fast(w0)
12981            || !self.uses_q8_1_fast(w1)
12982        {
12983            return Ok(None);
12984        }
12985        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
12986        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
12987        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
12988        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
12989        if !self.mmvq_supports(QT_NVFP4) {
12990            return Ok(None);
12991        }
12992        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12993        if w1.in_features() != in_f || w1.out_features() != out_f {
12994            return Ok(None);
12995        }
12996        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12997            (
12998                GpuTensor::Quant {
12999                    bytes: b0,
13000                    qtype: q0,
13001                    row_bytes: rb0,
13002                    scale: s0,
13003                    rp: rp0,
13004                    rp4: None,
13005                    ..
13006                },
13007                GpuTensor::Quant {
13008                    bytes: b1,
13009                    qtype: q1,
13010                    row_bytes: rb1,
13011                    scale: s1,
13012                    rp: rp1,
13013                    rp4: None,
13014                    ..
13015                },
13016            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13017                (b0, b1, *rb0, *s0, *s1, *rp0)
13018            }
13019            _ => return Ok(None),
13020        };
13021        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
13022        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
13023        if std::env::var("MEMRA_DEBUG").is_ok() {
13024            static ONCE: std::sync::Once = std::sync::Once::new();
13025            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13026        }
13027        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13028        let (y0, y1) =
13029            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13030        let mut y0 = y0;
13031        let mut y1 = y1;
13032        if s0 != 1.0 {
13033            self.scale_inplace(&mut y0, s0, m * out_f)?;
13034        }
13035        if s1 != 1.0 {
13036            self.scale_inplace(&mut y1, s1, m * out_f)?;
13037        }
13038        Ok(Some((y0, y1)))
13039    }
13040
13041    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
13042    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
13043    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
13044    /// twins (both buffers must be the repacked layout).
13045    #[allow(clippy::too_many_arguments)]
13046    pub fn qmatvec_batched_dual_raw(
13047        &self,
13048        b0: &CudaSlice<u8>,
13049        b1: &CudaSlice<u8>,
13050        aq: &CudaSlice<i8>,
13051        ad: &CudaSlice<f32>,
13052        m: usize,
13053        in_f: usize,
13054        out_f: usize,
13055        row_bytes: usize,
13056        rp: bool,
13057    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13058        const ROWS_PER_BLOCK: u32 = 4;
13059        let mcols = Self::batched_mcols(m);
13060        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
13061        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
13062        let tiny_rp1 = rp
13063            && mcols == 4
13064            && out_f <= 128
13065            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13066        let (name, rows_per_block) = if tiny_rp1 {
13067            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13068        } else {
13069            match (mcols, rp, m) {
13070                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13071                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13072                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13073                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13074                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13075                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13076                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13077                _ => {
13078                    return Err(
13079                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13080                    );
13081                }
13082            }
13083        };
13084        let f = self.func(name);
13085        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13086        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13087        let cfg = LaunchConfig {
13088            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13089            block_dim: (32, ROWS_PER_BLOCK, 1),
13090            shared_mem_bytes: 0,
13091        };
13092        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13093        let __s_b = self.gpu.stream();
13094        let mut b = __s_b.launch_builder(&f);
13095        b.arg(b0)
13096            .arg(b1)
13097            .arg(aq)
13098            .arg(ad)
13099            .arg(&mut y0)
13100            .arg(&mut y1)
13101            .arg(&inf)
13102            .arg(&outf)
13103            .arg(&mi)
13104            .arg(&rb);
13105        unsafe {
13106            b.launch(cfg)?;
13107        }
13108        Ok((y0, y1))
13109    }
13110
13111    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
13112    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
13113    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
13114    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
13115    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
13116    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
13117    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
13118    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
13119    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
13120    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
13121    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
13122    pub fn matmul_pre_dual_noscale(
13123        &self,
13124        w0: &crate::model::GpuTensor,
13125        w1: &crate::model::GpuTensor,
13126        aq: &CudaSlice<i8>,
13127        ad: &CudaSlice<f32>,
13128        m: usize,
13129    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13130    {
13131        use crate::model::GpuTensor;
13132        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13133            return Ok(None);
13134        }
13135        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
13136        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
13137        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
13138        // would mix dispatch families across the pair — the exact class `q8_fused_params`
13139        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
13140        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
13141        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
13142        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
13143        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
13144        if !self.mmvq_supports(QT_NVFP4) {
13145            return Ok(None);
13146        }
13147        let (in_f, out_f) = (w0.in_features(), w0.out_features());
13148        if w1.in_features() != in_f || w1.out_features() != out_f {
13149            return Ok(None);
13150        }
13151        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
13152        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
13153        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
13154        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
13155        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
13156        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
13157        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
13158        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
13159        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
13160        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
13161        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
13162        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
13163        let no_mirror =
13164            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13165        if self.q8_ffn_fuse2_on()
13166            && no_mirror(w0)
13167            && no_mirror(w1)
13168            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13169        {
13170            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13171            return Ok(Some(((y0, 1.0), (y1, 1.0))));
13172        }
13173        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
13174        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
13175        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
13176        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
13177        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
13178        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
13179        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
13180        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
13181        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
13182        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13183            let (y0, y1) =
13184                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13185            return Ok(Some(((y0, p0.3), (y1, p1.3))));
13186        }
13187        let (b0, q0, rb0, s0, rp0) = match w0 {
13188            GpuTensor::Quant {
13189                bytes,
13190                qtype,
13191                row_bytes,
13192                scale,
13193                rp,
13194                ..
13195            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13196            _ => return Ok(None),
13197        };
13198        let (b1, q1, rb1, s1, rp1) = match w1 {
13199            GpuTensor::Quant {
13200                bytes,
13201                qtype,
13202                row_bytes,
13203                scale,
13204                rp,
13205                ..
13206            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13207            _ => return Ok(None),
13208        };
13209        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13210            return Ok(None);
13211        }
13212        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13213        const RPW: u32 = 2;
13214        let rows_per_block = ROWS_PER_BLOCK * RPW;
13215        let f = self.func(if rp0 {
13216            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13217        } else {
13218            "qmatvec_nvfp4_mmvq_dual_mr2"
13219        });
13220        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13221        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13222        let cfg = LaunchConfig {
13223            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13224            block_dim: (32, ROWS_PER_BLOCK, 1),
13225            shared_mem_bytes: 0,
13226        };
13227        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13228        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
13229        // yscale args stay 1.0 here (they exist for the single-tensor callers).
13230        let one = 1.0f32;
13231        let __s_b = self.gpu.stream();
13232        let mut b = __s_b.launch_builder(&f);
13233        b.arg(b0)
13234            .arg(b1)
13235            .arg(aq)
13236            .arg(ad)
13237            .arg(&mut y0)
13238            .arg(&mut y1)
13239            .arg(&inf)
13240            .arg(&outf)
13241            .arg(&mi)
13242            .arg(&rb)
13243            .arg(&one)
13244            .arg(&one);
13245        unsafe {
13246            b.launch(cfg)?;
13247        }
13248        Ok(Some(((y0, s0), (y1, s1))))
13249    }
13250
13251    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
13252    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
13253    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
13254    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
13255    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
13256    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
13257    /// back to the three singles.
13258    #[allow(clippy::too_many_arguments)]
13259    pub fn matmul_nvfp4_fused3(
13260        &self,
13261        w0: &crate::model::GpuTensor,
13262        w1: &crate::model::GpuTensor,
13263        w2: &crate::model::GpuTensor,
13264        aq: &CudaSlice<i8>,
13265        ad: &CudaSlice<f32>,
13266        m: usize,
13267    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13268    {
13269        use crate::model::GpuTensor;
13270        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13271        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
13272        // verbatim, weight rows read once for all m columns, bit-identical per
13273        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
13274        // segments would re-read the weight per row" note described the grid.y=m lift,
13275        // which this twin deliberately is NOT.
13276        if !self.mmvq_supports(QT_NVFP4)
13277            || !self.uses_q8_1_fast(w0)
13278            || !self.uses_q8_1_fast(w1)
13279            || !self.uses_q8_1_fast(w2)
13280        {
13281            return Ok(None);
13282        }
13283        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
13284        // door — same family and bit-identity law as the fused4 delegate above.
13285        if (9..=16).contains(&m) {
13286            return Ok(
13287                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13288                    Some(mut ys) => {
13289                        let y2 = ys.pop().unwrap();
13290                        let y1 = ys.pop().unwrap();
13291                        let y0 = ys.pop().unwrap();
13292                        Some((y0, y1, y2))
13293                    }
13294                    None => None,
13295                },
13296            );
13297        }
13298        if !(1..=8).contains(&m) {
13299            return Ok(None);
13300        }
13301        if m > 1 {
13302            let in_f = w0.in_features();
13303            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13304                || !self.batched_supports(QT_NVFP4)
13305                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13306                || (m > 4 && !Self::b8_enabled())
13307                || in_f % 512 != 0
13308                || in_f / 64 > 272
13309            {
13310                return Ok(None);
13311            }
13312        }
13313        let unpack = |w: &crate::model::GpuTensor| match w {
13314            GpuTensor::Quant {
13315                bytes,
13316                qtype,
13317                scale,
13318                rp,
13319                ..
13320            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13321            _ => None,
13322        };
13323        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13324            return Ok(None);
13325        };
13326        let in_f = w0.in_features();
13327        if w1.in_features() != in_f || w2.in_features() != in_f {
13328            return Ok(None);
13329        }
13330        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13331        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13332        const RPW: u32 = 2;
13333        let rows_pb = ROWS_PER_BLOCK * RPW;
13334        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13335        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13336        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13337        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13338        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13339        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13340        // only dereferenced for the launch-arg build inside this call.
13341        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13342        if m > 1 {
13343            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
13344            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13345                return Ok(None);
13346            }
13347            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13348            let cfg = LaunchConfig {
13349                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13350                block_dim: (32, ROWS_PER_BLOCK, 1),
13351                shared_mem_bytes: 0,
13352            };
13353            let __s_b = self.gpu.stream();
13354            let mut b = __s_b.launch_builder(&f);
13355            b.arg(b0)
13356                .arg(b1)
13357                .arg(b2)
13358                .arg(aq)
13359                .arg(ad)
13360                .arg(&mut y0)
13361                .arg(&mut y1)
13362                .arg(&mut y2)
13363                .arg(&inf)
13364                .arg(&oi0)
13365                .arg(&oi1)
13366                .arg(&oi2)
13367                .arg(&mi);
13368            unsafe {
13369                b.launch(cfg)?;
13370            }
13371            return Ok(Some((y0, y1, y2)));
13372        }
13373        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13374        let cfg = LaunchConfig {
13375            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13376            block_dim: (32, ROWS_PER_BLOCK, 1),
13377            shared_mem_bytes: 0,
13378        };
13379        let __s_b = self.gpu.stream();
13380        let mut b = __s_b.launch_builder(&f);
13381        b.arg(b0)
13382            .arg(b1)
13383            .arg(b2)
13384            .arg(aq)
13385            .arg(ad)
13386            .arg(&mut y0)
13387            .arg(&mut y1)
13388            .arg(&mut y2)
13389            .arg(&inf)
13390            .arg(&oi0)
13391            .arg(&oi1)
13392            .arg(&oi2)
13393            .arg(&mi)
13394            .arg(&p0.1)
13395            .arg(&p1.1)
13396            .arg(&p2.1);
13397        unsafe {
13398            b.launch(cfg)?;
13399        }
13400        Ok(Some((y0, y1, y2)))
13401    }
13402
13403    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
13404    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
13405    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
13406    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
13407    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
13408    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
13409    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
13410    /// same-binary interleaved A/B arm.
13411    pub fn matmul_nvfp4_fused2(
13412        &self,
13413        w0: &crate::model::GpuTensor,
13414        w1: &crate::model::GpuTensor,
13415        aq: &CudaSlice<i8>,
13416        ad: &CudaSlice<f32>,
13417        m: usize,
13418    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13419        use crate::model::GpuTensor;
13420        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13421        let off =
13422            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13423        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
13424        // read serves all m rows); the fused segments would re-read the weight per row.
13425        if off
13426            || m != 1
13427            || !self.mmvq_supports(QT_NVFP4)
13428            || !self.uses_q8_1_fast(w0)
13429            || !self.uses_q8_1_fast(w1)
13430        {
13431            return Ok(None);
13432        }
13433        let unpack = |w: &crate::model::GpuTensor| match w {
13434            GpuTensor::Quant {
13435                bytes,
13436                qtype,
13437                scale,
13438                rp,
13439                ..
13440            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13441            _ => None,
13442        };
13443        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13444            return Ok(None);
13445        };
13446        let in_f = w0.in_features();
13447        if w1.in_features() != in_f {
13448            return Ok(None);
13449        }
13450        let (o0, o1) = (w0.out_features(), w1.out_features());
13451        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13452        const RPW: u32 = 2;
13453        let rows_pb = ROWS_PER_BLOCK * RPW;
13454        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13455        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13456        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13457        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13458        let cfg = LaunchConfig {
13459            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13460            block_dim: (32, ROWS_PER_BLOCK, 1),
13461            shared_mem_bytes: 0,
13462        };
13463        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13464        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13465        // only dereferenced for the launch-arg build inside this call.
13466        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13467        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
13468        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
13469        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13470            {
13471                use cudarc::driver::{DevicePtr, DevicePtrMut};
13472                let s = &self.gpu.stream();
13473                let (pw0, _g0) = b0.device_ptr(s);
13474                let (pw1, _g1) = b1.device_ptr(s);
13475                let (paq, _g2) = aq.device_ptr(s);
13476                let (pad, _g3) = ad.device_ptr(s);
13477                let (py0, _g4) = y0.device_ptr_mut(s);
13478                let (py1, _g5) = y1.device_ptr_mut(s);
13479                let (s0, s1) = (p0.1, p1.1);
13480                let mut ps = [
13481                    &pw0 as *const _ as *mut std::ffi::c_void,
13482                    &pw1 as *const _ as *mut _,
13483                    &paq as *const _ as *mut _,
13484                    &pad as *const _ as *mut _,
13485                    &py0 as *const _ as *mut _,
13486                    &py1 as *const _ as *mut _,
13487                    &inf as *const _ as *mut _,
13488                    &oi0 as *const _ as *mut _,
13489                    &oi1 as *const _ as *mut _,
13490                    &mi as *const _ as *mut _,
13491                    &s0 as *const _ as *mut _,
13492                    &s1 as *const _ as *mut _,
13493                ];
13494                unsafe {
13495                    self.launch_pdl(
13496                        "qmatvec_nvfp4_mmvq_fused2_rp",
13497                        cfg.grid_dim,
13498                        cfg.block_dim,
13499                        &mut ps,
13500                    )?;
13501                }
13502            }
13503            return Ok(Some((y0, y1)));
13504        }
13505        let __s_b = self.gpu.stream();
13506        let mut b = __s_b.launch_builder(&f);
13507        b.arg(b0)
13508            .arg(b1)
13509            .arg(aq)
13510            .arg(ad)
13511            .arg(&mut y0)
13512            .arg(&mut y1)
13513            .arg(&inf)
13514            .arg(&oi0)
13515            .arg(&oi1)
13516            .arg(&mi)
13517            .arg(&p0.1)
13518            .arg(&p1.1);
13519        unsafe {
13520            b.launch(cfg)?;
13521        }
13522        Ok(Some((y0, y1)))
13523    }
13524
13525    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
13526    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
13527    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
13528    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
13529    pub fn matmul_nvfp4_fused2_into(
13530        &self,
13531        w0: &crate::model::GpuTensor,
13532        w1: &crate::model::GpuTensor,
13533        aq: &CudaSlice<i8>,
13534        ad: &CudaSlice<f32>,
13535        y0: &mut CudaSlice<f32>,
13536        y1: &mut CudaSlice<f32>,
13537    ) -> Result<bool, Box<dyn std::error::Error>> {
13538        use crate::model::GpuTensor;
13539        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13540        let off =
13541            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13542        if off
13543            || !self.mmvq_supports(QT_NVFP4)
13544            || !self.uses_q8_1_fast(w0)
13545            || !self.uses_q8_1_fast(w1)
13546        {
13547            return Ok(false);
13548        }
13549        let unpack = |w: &crate::model::GpuTensor| match w {
13550            GpuTensor::Quant {
13551                bytes,
13552                qtype,
13553                scale,
13554                rp,
13555                ..
13556            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13557            _ => None,
13558        };
13559        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13560            return Ok(false);
13561        };
13562        let in_f = w0.in_features();
13563        if w1.in_features() != in_f {
13564            return Ok(false);
13565        }
13566        let (o0, o1) = (w0.out_features(), w1.out_features());
13567        if y0.len() < o0 || y1.len() < o1 {
13568            return Ok(false);
13569        }
13570        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13571        const RPW: u32 = 2;
13572        let rows_pb = ROWS_PER_BLOCK * RPW;
13573        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13574        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13575        let cfg = LaunchConfig {
13576            grid_dim: (nb(o0) + nb(o1), 1, 1),
13577            block_dim: (32, ROWS_PER_BLOCK, 1),
13578            shared_mem_bytes: 0,
13579        };
13580        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13581        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13582        // only dereferenced for the launch-arg build inside this call.
13583        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13584        let __s_b = self.gpu.stream();
13585        let mut b = __s_b.launch_builder(&f);
13586        b.arg(b0)
13587            .arg(b1)
13588            .arg(aq)
13589            .arg(ad)
13590            .arg(&mut *y0)
13591            .arg(&mut *y1)
13592            .arg(&inf)
13593            .arg(&oi0)
13594            .arg(&oi1)
13595            .arg(&mi)
13596            .arg(&p0.1)
13597            .arg(&p1.1);
13598        unsafe {
13599            b.launch(cfg)?;
13600        }
13601        Ok(true)
13602    }
13603
13604    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
13605    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
13606    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
13607    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
13608    #[allow(clippy::type_complexity)]
13609    pub fn matmul_nvfp4_fused4(
13610        &self,
13611        w0: &crate::model::GpuTensor,
13612        w1: &crate::model::GpuTensor,
13613        w2: &crate::model::GpuTensor,
13614        w3: &crate::model::GpuTensor,
13615        aq: &CudaSlice<i8>,
13616        ad: &CudaSlice<f32>,
13617        m: usize,
13618    ) -> Result<
13619        Option<(
13620            CudaSlice<f32>,
13621            CudaSlice<f32>,
13622            CudaSlice<f32>,
13623            CudaSlice<f32>,
13624        )>,
13625        Box<dyn std::error::Error>,
13626    > {
13627        use crate::model::GpuTensor;
13628        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
13629        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
13630        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
13631        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
13632        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
13633        // Admission mirrors the singles' batched gates below.
13634        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13635            || !self.mmvq_supports(QT_NVFP4)
13636            || !self.uses_q8_1_fast(w0)
13637            || !self.uses_q8_1_fast(w1)
13638            || !self.uses_q8_1_fast(w2)
13639            || !self.uses_q8_1_fast(w3)
13640        {
13641            return Ok(None);
13642        }
13643        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
13644        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
13645        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
13646        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
13647        if (9..=16).contains(&m) {
13648            return Ok(
13649                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13650                    Some(mut ys) => {
13651                        let y3 = ys.pop().unwrap();
13652                        let y2 = ys.pop().unwrap();
13653                        let y1 = ys.pop().unwrap();
13654                        let y0 = ys.pop().unwrap();
13655                        Some((y0, y1, y2, y3))
13656                    }
13657                    None => None,
13658                },
13659            );
13660        }
13661        if !(1..=8).contains(&m) {
13662            return Ok(None);
13663        }
13664        if m > 1 {
13665            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
13666            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
13667            let in_f = w0.in_features();
13668            if !self.batched_supports(QT_NVFP4)
13669                || std::env::var("MEMRA_NO_BATCHED").is_ok()
13670                || (m > 4 && !Self::b8_enabled())
13671                || in_f % 512 != 0
13672                || in_f / 64 > 272
13673            {
13674                return Ok(None);
13675            }
13676        }
13677        let unpack = |w: &crate::model::GpuTensor| match w {
13678            GpuTensor::Quant {
13679                bytes,
13680                qtype,
13681                scale,
13682                rp,
13683                ..
13684            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13685            _ => None,
13686        };
13687        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13688            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13689        else {
13690            return Ok(None);
13691        };
13692        let in_f = w0.in_features();
13693        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13694            return Ok(None);
13695        }
13696        let (o0, o1, o2, o3) = (
13697            w0.out_features(),
13698            w1.out_features(),
13699            w2.out_features(),
13700            w3.out_features(),
13701        );
13702        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
13703        const RPW: u32 = 2;
13704        let rows_pb = ROWS_PER_BLOCK * RPW;
13705        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13706        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13707        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13708        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13709        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13710        let (inf, oi0, oi1, oi2, oi3, mi) = (
13711            in_f as i32,
13712            o0 as i32,
13713            o1 as i32,
13714            o2 as i32,
13715            o3 as i32,
13716            m as i32,
13717        );
13718        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
13719        // only dereferenced for the launch-arg build inside this call.
13720        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13721        if m > 1 {
13722            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
13723            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
13724            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13725                return Ok(None);
13726            }
13727            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13728            let cfg = LaunchConfig {
13729                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13730                block_dim: (32, ROWS_PER_BLOCK, 1),
13731                shared_mem_bytes: 0,
13732            };
13733            let __s_b = self.gpu.stream();
13734            let mut b = __s_b.launch_builder(&f);
13735            b.arg(b0)
13736                .arg(b1)
13737                .arg(b2)
13738                .arg(b3)
13739                .arg(aq)
13740                .arg(ad)
13741                .arg(&mut y0)
13742                .arg(&mut y1)
13743                .arg(&mut y2)
13744                .arg(&mut y3)
13745                .arg(&inf)
13746                .arg(&oi0)
13747                .arg(&oi1)
13748                .arg(&oi2)
13749                .arg(&oi3)
13750                .arg(&mi);
13751            unsafe {
13752                b.launch(cfg)?;
13753            }
13754            return Ok(Some((y0, y1, y2, y3)));
13755        }
13756        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13757        let cfg = LaunchConfig {
13758            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13759            block_dim: (32, ROWS_PER_BLOCK, 1),
13760            shared_mem_bytes: 0,
13761        };
13762        let __s_b = self.gpu.stream();
13763        let mut b = __s_b.launch_builder(&f);
13764        b.arg(b0)
13765            .arg(b1)
13766            .arg(b2)
13767            .arg(b3)
13768            .arg(aq)
13769            .arg(ad)
13770            .arg(&mut y0)
13771            .arg(&mut y1)
13772            .arg(&mut y2)
13773            .arg(&mut y3)
13774            .arg(&inf)
13775            .arg(&oi0)
13776            .arg(&oi1)
13777            .arg(&oi2)
13778            .arg(&oi3)
13779            .arg(&mi)
13780            .arg(&p0.1)
13781            .arg(&p1.1)
13782            .arg(&p2.1)
13783            .arg(&p3.1);
13784        unsafe {
13785            b.launch(cfg)?;
13786        }
13787        Ok(Some((y0, y1, y2, y3)))
13788    }
13789
13790    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
13791    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
13792    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
13793    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
13794    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
13795    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
13796    /// back to the per-tensor path.
13797    pub fn matmul_q8_fused2(
13798        &self,
13799        w0: &crate::model::GpuTensor,
13800        w1: &crate::model::GpuTensor,
13801        aq: &CudaSlice<i8>,
13802        ad: &CudaSlice<f32>,
13803    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13804        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
13805        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
13806        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
13807        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
13808        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
13809        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13810            return Ok(Some(self.e4m3_fused2_core(
13811                p0.0,
13812                p1.0,
13813                aq,
13814                ad,
13815                w0.in_features(),
13816                p0.1,
13817                p1.1,
13818                p0.2,
13819                p0.3,
13820                p1.3,
13821            )?));
13822        }
13823        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13824            return Ok(None);
13825        };
13826        Ok(Some(self.q8_fused2_core(
13827            p0.0,
13828            p1.0,
13829            aq,
13830            ad,
13831            w0.in_features(),
13832            p0.1,
13833            p1.1,
13834            p0.2,
13835        )?))
13836    }
13837
13838    #[allow(clippy::too_many_arguments)]
13839    fn q8_fused2_core(
13840        &self,
13841        b0: &CudaSlice<u8>,
13842        b1: &CudaSlice<u8>,
13843        aq: &CudaSlice<i8>,
13844        ad: &CudaSlice<f32>,
13845        in_f: usize,
13846        out0: usize,
13847        out1: usize,
13848        row_bytes: usize,
13849    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13850        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13851        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13852        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13853        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13854        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13855        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13856        let cfg = LaunchConfig {
13857            grid_dim: (nb0 + nb1, 1, 1),
13858            block_dim: (32, ROWS_PER_BLOCK, 1),
13859            shared_mem_bytes: 0,
13860        };
13861        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13862        let __s_b = self.gpu.stream();
13863        let mut b = __s_b.launch_builder(&f);
13864        b.arg(b0)
13865            .arg(b1)
13866            .arg(aq)
13867            .arg(ad)
13868            .arg(&mut y0)
13869            .arg(&mut y1)
13870            .arg(&inf)
13871            .arg(&o0)
13872            .arg(&o1)
13873            .arg(&rbl);
13874        unsafe {
13875            b.launch(cfg)?;
13876        }
13877        Ok((y0, y1))
13878    }
13879
13880    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13881    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13882    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13883    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13884    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13885    pub fn matmul_q8_fused2_x(
13886        &self,
13887        w0: &crate::model::GpuTensor,
13888        w1: &crate::model::GpuTensor,
13889        x: &CudaSlice<f32>,
13890    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13891        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13892            return Ok(None);
13893        }
13894        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13895            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13896            return Ok(Some(self.e4m3_fused2_core(
13897                p0.0,
13898                p1.0,
13899                &aq,
13900                &ad,
13901                w0.in_features(),
13902                p0.1,
13903                p1.1,
13904                p0.2,
13905                p0.3,
13906                p1.3,
13907            )?));
13908        }
13909        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13910            return Ok(None);
13911        };
13912        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13913        Ok(Some(self.q8_fused2_core(
13914            p0.0,
13915            p1.0,
13916            &aq,
13917            &ad,
13918            w0.in_features(),
13919            p0.1,
13920            p1.1,
13921            p0.2,
13922        )?))
13923    }
13924
13925    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13926    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13927    #[allow(clippy::too_many_arguments)]
13928    pub fn qmatvec_q8_fused2_raw(
13929        &self,
13930        b0: &CudaSlice<u8>,
13931        b1: &CudaSlice<u8>,
13932        x: &CudaSlice<f32>,
13933        in_f: usize,
13934        out0: usize,
13935        out1: usize,
13936        row_bytes: usize,
13937    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13938        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13939        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13940    }
13941
13942    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13943    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13944    /// (tensor,row) to three separate m=1 MMVQ launches.
13945    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13946    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13947    pub fn matmul_q4_fused3(
13948        &self,
13949        w0: &crate::model::GpuTensor,
13950        w1: &crate::model::GpuTensor,
13951        w2: &crate::model::GpuTensor,
13952        aq: &CudaSlice<i8>,
13953        ad: &CudaSlice<f32>,
13954    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13955    {
13956        use crate::model::GpuTensor;
13957        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13958            match w {
13959                GpuTensor::Quant {
13960                    qtype, row_bytes, ..
13961                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13962                _ => None,
13963            }
13964        };
13965        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13966            return Ok(None);
13967        };
13968        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13969            return Ok(None);
13970        }
13971        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
13972        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
13973        // the separate matvecs (each routes its own rp).
13974        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13975            match w {
13976                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13977                    Some(m) => (m, true),
13978                    None => (bytes, *rp),
13979                },
13980                _ => unreachable!(),
13981            }
13982        }
13983        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13984        if rp0 != rp1 || rp1 != rp2 {
13985            return Ok(None);
13986        }
13987        let rp = rp0;
13988        let rpb: u32 = 4;
13989        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
13990        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
13991        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
13992        let mr1 = rp && Self::q40_mr1_on();
13993        let nb = |o: usize| {
13994            if mr1 {
13995                (o as u32).div_ceil(rpb)
13996            } else {
13997                (o as u32).div_ceil(2).div_ceil(rpb)
13998            }
13999        };
14000        let grid = nb(o0) + nb(o1) + nb(o2);
14001        let mut y0 = self.alloc_uninit::<f32>(o0)?;
14002        let mut y1 = self.alloc_uninit::<f32>(o1)?;
14003        let mut y2 = self.alloc_uninit::<f32>(o2)?;
14004        let f = self.func(if mr1 {
14005            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14006        } else if rp {
14007            "qmatvec_q4_0_mmvq_fused3_rp"
14008        } else {
14009            "qmatvec_q4_0_mmvq_fused3"
14010        });
14011        let cfg = LaunchConfig {
14012            grid_dim: (grid, 1, 1),
14013            block_dim: (32, rpb, 1),
14014            shared_mem_bytes: 0,
14015        };
14016        let inf = w0.in_features() as i32;
14017        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14018        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14019        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
14020        // variant may take the programmatic-serialization launch.
14021        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14022            {
14023                use cudarc::driver::{DevicePtr, DevicePtrMut};
14024                let s = &self.gpu.stream();
14025                let (p0, _g0) = b0.device_ptr(s);
14026                let (p1, _g1) = b1.device_ptr(s);
14027                let (p2, _g2) = b2.device_ptr(s);
14028                let (paq, _g3) = aq.device_ptr(s);
14029                let (pad, _g4) = ad.device_ptr(s);
14030                let (py0, _g5) = y0.device_ptr_mut(s);
14031                let (py1, _g6) = y1.device_ptr_mut(s);
14032                let (py2, _g7) = y2.device_ptr_mut(s);
14033                let mut ps = [
14034                    &p0 as *const _ as *mut std::ffi::c_void,
14035                    &p1 as *const _ as *mut _,
14036                    &p2 as *const _ as *mut _,
14037                    &paq as *const _ as *mut _,
14038                    &pad as *const _ as *mut _,
14039                    &py0 as *const _ as *mut _,
14040                    &py1 as *const _ as *mut _,
14041                    &py2 as *const _ as *mut _,
14042                    &inf as *const _ as *mut _,
14043                    &oo0 as *const _ as *mut _,
14044                    &oo1 as *const _ as *mut _,
14045                    &oo2 as *const _ as *mut _,
14046                    &r0 as *const _ as *mut _,
14047                    &r1 as *const _ as *mut _,
14048                    &r2 as *const _ as *mut _,
14049                ];
14050                unsafe {
14051                    self.launch_pdl(
14052                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14053                        (grid, 1, 1),
14054                        (32, rpb, 1),
14055                        &mut ps,
14056                    )?;
14057                }
14058            }
14059            return Ok(Some((y0, y1, y2)));
14060        }
14061        let __s_b = self.gpu.stream();
14062        let mut b = __s_b.launch_builder(&f);
14063        b.arg(b0)
14064            .arg(b1)
14065            .arg(b2)
14066            .arg(aq)
14067            .arg(ad)
14068            .arg(&mut y0)
14069            .arg(&mut y1)
14070            .arg(&mut y2)
14071            .arg(&inf)
14072            .arg(&oo0)
14073            .arg(&oo1)
14074            .arg(&oo2)
14075            .arg(&r0)
14076            .arg(&r1)
14077            .arg(&r2);
14078        unsafe {
14079            b.launch(cfg)?;
14080        }
14081        Ok(Some((y0, y1, y2)))
14082    }
14083
14084    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
14085    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
14086    #[allow(clippy::too_many_arguments)]
14087    pub fn matmul_q4_fused3_into(
14088        &self,
14089        w0: &crate::model::GpuTensor,
14090        w1: &crate::model::GpuTensor,
14091        w2: &crate::model::GpuTensor,
14092        aq: &CudaSlice<i8>,
14093        ad: &CudaSlice<f32>,
14094        y0: &mut CudaSlice<f32>,
14095        y1: &mut CudaSlice<f32>,
14096        y2: &mut CudaSlice<f32>,
14097    ) -> Result<bool, Box<dyn std::error::Error>> {
14098        use crate::model::GpuTensor;
14099        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14100            match w {
14101                GpuTensor::Quant {
14102                    qtype, row_bytes, ..
14103                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14104                _ => None,
14105            }
14106        };
14107        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14108            return Ok(false);
14109        };
14110        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14111            return Ok(false);
14112        }
14113        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14114            match w {
14115                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14116                    Some(m) => (m, true),
14117                    None => (bytes, *rp),
14118                },
14119                _ => unreachable!(),
14120            }
14121        }
14122        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14123        if rp0 != rp1 || rp1 != rp2 {
14124            return Ok(false);
14125        }
14126        let rp = rp0;
14127        let rpb: u32 = 4;
14128        let mr1 = rp && Self::q40_mr1_on();
14129        let nb = |o: usize| {
14130            if mr1 {
14131                (o as u32).div_ceil(rpb)
14132            } else {
14133                (o as u32).div_ceil(2).div_ceil(rpb)
14134            }
14135        };
14136        let grid = nb(o0) + nb(o1) + nb(o2);
14137        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14138        let f = self.func(if mr1 {
14139            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14140        } else if rp {
14141            "qmatvec_q4_0_mmvq_fused3_rp"
14142        } else {
14143            "qmatvec_q4_0_mmvq_fused3"
14144        });
14145        let cfg = LaunchConfig {
14146            grid_dim: (grid, 1, 1),
14147            block_dim: (32, rpb, 1),
14148            shared_mem_bytes: 0,
14149        };
14150        let inf = w0.in_features() as i32;
14151        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14152        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14153        // PDL wave-A: identical to the owned twin (capture-lane parity).
14154        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14155            use cudarc::driver::{DevicePtr, DevicePtrMut};
14156            let s = &self.gpu.stream();
14157            let (p0, _g0) = b0.device_ptr(s);
14158            let (p1, _g1) = b1.device_ptr(s);
14159            let (p2, _g2) = b2.device_ptr(s);
14160            let (paq, _g3) = aq.device_ptr(s);
14161            let (pad, _g4) = ad.device_ptr(s);
14162            let (py0, _g5) = y0.device_ptr_mut(s);
14163            let (py1, _g6) = y1.device_ptr_mut(s);
14164            let (py2, _g7) = y2.device_ptr_mut(s);
14165            let mut ps = [
14166                &p0 as *const _ as *mut std::ffi::c_void,
14167                &p1 as *const _ as *mut _,
14168                &p2 as *const _ as *mut _,
14169                &paq as *const _ as *mut _,
14170                &pad as *const _ as *mut _,
14171                &py0 as *const _ as *mut _,
14172                &py1 as *const _ as *mut _,
14173                &py2 as *const _ as *mut _,
14174                &inf as *const _ as *mut _,
14175                &oo0 as *const _ as *mut _,
14176                &oo1 as *const _ as *mut _,
14177                &oo2 as *const _ as *mut _,
14178                &r0 as *const _ as *mut _,
14179                &r1 as *const _ as *mut _,
14180                &r2 as *const _ as *mut _,
14181            ];
14182            unsafe {
14183                self.launch_pdl(
14184                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14185                    (grid, 1, 1),
14186                    (32, rpb, 1),
14187                    &mut ps,
14188                )?;
14189            }
14190            return Ok(true);
14191        }
14192        let __s_b = self.gpu.stream();
14193        let mut b = __s_b.launch_builder(&f);
14194        b.arg(b0)
14195            .arg(b1)
14196            .arg(b2)
14197            .arg(aq)
14198            .arg(ad)
14199            .arg(&mut *y0)
14200            .arg(&mut *y1)
14201            .arg(&mut *y2)
14202            .arg(&inf)
14203            .arg(&oo0)
14204            .arg(&oo1)
14205            .arg(&oo2)
14206            .arg(&r0)
14207            .arg(&r1)
14208            .arg(&r2);
14209        unsafe {
14210            b.launch(cfg)?;
14211        }
14212        Ok(true)
14213    }
14214
14215    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
14216    pub fn matmul_q4_fused2(
14217        &self,
14218        w0: &crate::model::GpuTensor,
14219        w1: &crate::model::GpuTensor,
14220        aq: &CudaSlice<i8>,
14221        ad: &CudaSlice<f32>,
14222    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14223        use crate::model::GpuTensor;
14224        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14225            match w {
14226                GpuTensor::Quant {
14227                    qtype, row_bytes, ..
14228                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14229                _ => None,
14230            }
14231        };
14232        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14233            return Ok(None);
14234        };
14235        if w0.in_features() != w1.in_features() {
14236            return Ok(None);
14237        }
14238        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
14239        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14240            match w {
14241                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14242                    Some(m) => (m, true),
14243                    None => (bytes, *rp),
14244                },
14245                _ => unreachable!(),
14246            }
14247        }
14248        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14249        if rp0 != rp1 {
14250            return Ok(None);
14251        }
14252        let rp = rp0;
14253        let rpb: u32 = 4;
14254        // mr1 twin — see matmul_q4_fused3.
14255        let mr1 = rp && Self::q40_mr1_on();
14256        let nb = |o: usize| {
14257            if mr1 {
14258                (o as u32).div_ceil(rpb)
14259            } else {
14260                (o as u32).div_ceil(2).div_ceil(rpb)
14261            }
14262        };
14263        let grid = nb(o0) + nb(o1);
14264        let mut y0 = self.alloc_uninit::<f32>(o0)?;
14265        let mut y1 = self.alloc_uninit::<f32>(o1)?;
14266        let f = self.func(if mr1 {
14267            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14268        } else if rp {
14269            "qmatvec_q4_0_mmvq_fused2_rp"
14270        } else {
14271            "qmatvec_q4_0_mmvq_fused2"
14272        });
14273        let cfg = LaunchConfig {
14274            grid_dim: (grid, 1, 1),
14275            block_dim: (32, rpb, 1),
14276            shared_mem_bytes: 0,
14277        };
14278        let inf = w0.in_features() as i32;
14279        let (oo0, oo1) = (o0 as i32, o1 as i32);
14280        let (r0, r1) = (rb0 as i64, rb1 as i64);
14281        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
14282        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14283            {
14284                use cudarc::driver::{DevicePtr, DevicePtrMut};
14285                let s = &self.gpu.stream();
14286                let (p0, _g0) = b0.device_ptr(s);
14287                let (p1, _g1) = b1.device_ptr(s);
14288                let (paq, _g2) = aq.device_ptr(s);
14289                let (pad, _g3) = ad.device_ptr(s);
14290                let (py0, _g4) = y0.device_ptr_mut(s);
14291                let (py1, _g5) = y1.device_ptr_mut(s);
14292                let mut ps = [
14293                    &p0 as *const _ as *mut std::ffi::c_void,
14294                    &p1 as *const _ as *mut _,
14295                    &paq as *const _ as *mut _,
14296                    &pad as *const _ as *mut _,
14297                    &py0 as *const _ as *mut _,
14298                    &py1 as *const _ as *mut _,
14299                    &inf as *const _ as *mut _,
14300                    &oo0 as *const _ as *mut _,
14301                    &oo1 as *const _ as *mut _,
14302                    &r0 as *const _ as *mut _,
14303                    &r1 as *const _ as *mut _,
14304                ];
14305                unsafe {
14306                    self.launch_pdl(
14307                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14308                        (grid, 1, 1),
14309                        (32, rpb, 1),
14310                        &mut ps,
14311                    )?;
14312                }
14313            }
14314            return Ok(Some((y0, y1)));
14315        }
14316        let __s_b = self.gpu.stream();
14317        let mut b = __s_b.launch_builder(&f);
14318        b.arg(b0)
14319            .arg(b1)
14320            .arg(aq)
14321            .arg(ad)
14322            .arg(&mut y0)
14323            .arg(&mut y1)
14324            .arg(&inf)
14325            .arg(&oo0)
14326            .arg(&oo1)
14327            .arg(&r0)
14328            .arg(&r1);
14329        unsafe {
14330            b.launch(cfg)?;
14331        }
14332        Ok(Some((y0, y1)))
14333    }
14334
14335    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
14336    pub fn matmul_q4_fused2_into(
14337        &self,
14338        w0: &crate::model::GpuTensor,
14339        w1: &crate::model::GpuTensor,
14340        aq: &CudaSlice<i8>,
14341        ad: &CudaSlice<f32>,
14342        y0: &mut CudaSlice<f32>,
14343        y1: &mut CudaSlice<f32>,
14344    ) -> Result<bool, Box<dyn std::error::Error>> {
14345        use crate::model::GpuTensor;
14346        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14347            match w {
14348                GpuTensor::Quant {
14349                    qtype, row_bytes, ..
14350                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14351                _ => None,
14352            }
14353        };
14354        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14355            return Ok(false);
14356        };
14357        if w0.in_features() != w1.in_features() {
14358            return Ok(false);
14359        }
14360        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14361            match w {
14362                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14363                    Some(m) => (m, true),
14364                    None => (bytes, *rp),
14365                },
14366                _ => unreachable!(),
14367            }
14368        }
14369        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14370        if rp0 != rp1 {
14371            return Ok(false);
14372        }
14373        let rp = rp0;
14374        let rpb: u32 = 4;
14375        let mr1 = rp && Self::q40_mr1_on();
14376        let nb = |o: usize| {
14377            if mr1 {
14378                (o as u32).div_ceil(rpb)
14379            } else {
14380                (o as u32).div_ceil(2).div_ceil(rpb)
14381            }
14382        };
14383        let grid = nb(o0) + nb(o1);
14384        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14385        let f = self.func(if mr1 {
14386            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14387        } else if rp {
14388            "qmatvec_q4_0_mmvq_fused2_rp"
14389        } else {
14390            "qmatvec_q4_0_mmvq_fused2"
14391        });
14392        let cfg = LaunchConfig {
14393            grid_dim: (grid, 1, 1),
14394            block_dim: (32, rpb, 1),
14395            shared_mem_bytes: 0,
14396        };
14397        let inf = w0.in_features() as i32;
14398        let (oo0, oo1) = (o0 as i32, o1 as i32);
14399        let (r0, r1) = (rb0 as i64, rb1 as i64);
14400        // PDL wave-A: identical to the owned twin (capture-lane parity).
14401        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14402            use cudarc::driver::{DevicePtr, DevicePtrMut};
14403            let s = &self.gpu.stream();
14404            let (p0, _g0) = b0.device_ptr(s);
14405            let (p1, _g1) = b1.device_ptr(s);
14406            let (paq, _g2) = aq.device_ptr(s);
14407            let (pad, _g3) = ad.device_ptr(s);
14408            let (py0, _g4) = y0.device_ptr_mut(s);
14409            let (py1, _g5) = y1.device_ptr_mut(s);
14410            let mut ps = [
14411                &p0 as *const _ as *mut std::ffi::c_void,
14412                &p1 as *const _ as *mut _,
14413                &paq as *const _ as *mut _,
14414                &pad as *const _ as *mut _,
14415                &py0 as *const _ as *mut _,
14416                &py1 as *const _ as *mut _,
14417                &inf as *const _ as *mut _,
14418                &oo0 as *const _ as *mut _,
14419                &oo1 as *const _ as *mut _,
14420                &r0 as *const _ as *mut _,
14421                &r1 as *const _ as *mut _,
14422            ];
14423            unsafe {
14424                self.launch_pdl(
14425                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14426                    (grid, 1, 1),
14427                    (32, rpb, 1),
14428                    &mut ps,
14429                )?;
14430            }
14431            return Ok(true);
14432        }
14433        let __s_b = self.gpu.stream();
14434        let mut b = __s_b.launch_builder(&f);
14435        b.arg(b0)
14436            .arg(b1)
14437            .arg(aq)
14438            .arg(ad)
14439            .arg(&mut *y0)
14440            .arg(&mut *y1)
14441            .arg(&inf)
14442            .arg(&oo0)
14443            .arg(&oo1)
14444            .arg(&r0)
14445            .arg(&r1);
14446        unsafe {
14447            b.launch(cfg)?;
14448        }
14449        Ok(true)
14450    }
14451
14452    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
14453    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
14454    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
14455    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
14456    pub fn matmul_q4_fused2_batched(
14457        &self,
14458        w0: &crate::model::GpuTensor,
14459        w1: &crate::model::GpuTensor,
14460        aq: &CudaSlice<i8>,
14461        ad: &CudaSlice<f32>,
14462        m: usize,
14463    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14464        use crate::model::GpuTensor;
14465        if m < 2 || m > 8 {
14466            return Ok(None);
14467        }
14468        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14469            match w {
14470                GpuTensor::Quant {
14471                    qtype, row_bytes, ..
14472                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14473                _ => None,
14474            }
14475        };
14476        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14477            return Ok(None);
14478        };
14479        if w0.in_features() != w1.in_features() {
14480            return Ok(None);
14481        }
14482        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14483            match w {
14484                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14485                    Some(mr) => (mr, true),
14486                    None => (bytes, *rp),
14487                },
14488                _ => unreachable!(),
14489            }
14490        }
14491        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14492        if !rp0 || !rp1 {
14493            return Ok(None);
14494        }
14495        let mcols = Self::batched_mcols(m);
14496        let rpb: u32 = 4;
14497        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14498        let grid = nb(o0) + nb(o1);
14499        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14500        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14501        let f = self.func(match mcols {
14502            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14503            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14504            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14505        });
14506        let cfg = LaunchConfig {
14507            grid_dim: (grid, 1, 1),
14508            block_dim: (32, rpb, 1),
14509            shared_mem_bytes: 0,
14510        };
14511        let inf = w0.in_features() as i32;
14512        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14513        let rb = rb0 as i64;
14514        let __s_b = self.gpu.stream();
14515        let mut b = __s_b.launch_builder(&f);
14516        b.arg(b0)
14517            .arg(b1)
14518            .arg(aq)
14519            .arg(ad)
14520            .arg(&mut y0)
14521            .arg(&mut y1)
14522            .arg(&inf)
14523            .arg(&oo0)
14524            .arg(&oo1)
14525            .arg(&mi)
14526            .arg(&rb);
14527        unsafe {
14528            b.launch(cfg)?;
14529        }
14530        Ok(Some((y0, y1)))
14531    }
14532
14533    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
14534    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
14535    #[allow(clippy::too_many_arguments)]
14536    pub fn matmul_q4_fused3_batched(
14537        &self,
14538        w0: &crate::model::GpuTensor,
14539        w1: &crate::model::GpuTensor,
14540        w2: &crate::model::GpuTensor,
14541        aq: &CudaSlice<i8>,
14542        ad: &CudaSlice<f32>,
14543        m: usize,
14544    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14545    {
14546        use crate::model::GpuTensor;
14547        if m < 2 || m > 8 {
14548            return Ok(None);
14549        }
14550        let q4 = |w: &GpuTensor| -> Option<usize> {
14551            match w {
14552                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14553                _ => None,
14554            }
14555        };
14556        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14557            return Ok(None);
14558        };
14559        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14560            return Ok(None);
14561        }
14562        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14563            match w {
14564                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14565                    Some(mr) => (mr, true),
14566                    None => (bytes, *rp),
14567                },
14568                _ => unreachable!(),
14569            }
14570        }
14571        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14572        if !rp0 || !rp1 || !rp2 {
14573            return Ok(None);
14574        }
14575        let mcols = Self::batched_mcols(m);
14576        let rpb: u32 = 4;
14577        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14578        let grid = nb(o0) + nb(o1) + nb(o2);
14579        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14580        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14581        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14582        let f = self.func(match mcols {
14583            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14584            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14585            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14586        });
14587        let cfg = LaunchConfig {
14588            grid_dim: (grid, 1, 1),
14589            block_dim: (32, rpb, 1),
14590            shared_mem_bytes: 0,
14591        };
14592        let inf = w0.in_features() as i32;
14593        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14594        let rb = 0i64;
14595        let __s_b = self.gpu.stream();
14596        let mut b = __s_b.launch_builder(&f);
14597        b.arg(b0)
14598            .arg(b1)
14599            .arg(b2)
14600            .arg(aq)
14601            .arg(ad)
14602            .arg(&mut y0)
14603            .arg(&mut y1)
14604            .arg(&mut y2)
14605            .arg(&inf)
14606            .arg(&oo0)
14607            .arg(&oo1)
14608            .arg(&oo2)
14609            .arg(&mi)
14610            .arg(&rb);
14611        unsafe {
14612            b.launch(cfg)?;
14613        }
14614        Ok(Some((y0, y1, y2)))
14615    }
14616
14617    pub fn matmul_q8_fused3(
14618        &self,
14619        w0: &crate::model::GpuTensor,
14620        w1: &crate::model::GpuTensor,
14621        w2: &crate::model::GpuTensor,
14622        aq: &CudaSlice<i8>,
14623        ad: &CudaSlice<f32>,
14624    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14625    {
14626        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
14627        // are per-tensor FP8, so native residency without this arm meant three separate launches.
14628        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14629            return Ok(Some(self.e4m3_fused3_core(
14630                p0.0,
14631                p1.0,
14632                p2.0,
14633                aq,
14634                ad,
14635                w0.in_features(),
14636                p0.1,
14637                p1.1,
14638                p2.1,
14639                p0.2,
14640                p0.3,
14641                p1.3,
14642                p2.3,
14643            )?));
14644        }
14645        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14646            return Ok(None);
14647        };
14648        Ok(Some(self.q8_fused3_core(
14649            p0.0,
14650            p1.0,
14651            p2.0,
14652            aq,
14653            ad,
14654            w0.in_features(),
14655            p0.1,
14656            p1.1,
14657            p2.1,
14658            p0.2,
14659        )?))
14660    }
14661
14662    #[allow(clippy::too_many_arguments)]
14663    fn q8_fused3_core(
14664        &self,
14665        b0: &CudaSlice<u8>,
14666        b1: &CudaSlice<u8>,
14667        b2: &CudaSlice<u8>,
14668        aq: &CudaSlice<i8>,
14669        ad: &CudaSlice<f32>,
14670        in_f: usize,
14671        out0: usize,
14672        out1: usize,
14673        out2: usize,
14674        row_bytes: usize,
14675    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14676        const ROWS_PER_BLOCK: u32 = 4;
14677        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14678        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14679        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14680        let f = self.func("qmatvec_q8_0_mmvq_fused3");
14681        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14682        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14683        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14684        let cfg = LaunchConfig {
14685            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14686            block_dim: (32, ROWS_PER_BLOCK, 1),
14687            shared_mem_bytes: 0,
14688        };
14689        let (inf, o0, o1, o2, rbl) = (
14690            in_f as i32,
14691            out0 as i32,
14692            out1 as i32,
14693            out2 as i32,
14694            row_bytes as i64,
14695        );
14696        let __s_b = self.gpu.stream();
14697        let mut b = __s_b.launch_builder(&f);
14698        b.arg(b0)
14699            .arg(b1)
14700            .arg(b2)
14701            .arg(aq)
14702            .arg(ad)
14703            .arg(&mut y0)
14704            .arg(&mut y1)
14705            .arg(&mut y2)
14706            .arg(&inf)
14707            .arg(&o0)
14708            .arg(&o1)
14709            .arg(&o2)
14710            .arg(&rbl);
14711        unsafe {
14712            b.launch(cfg)?;
14713        }
14714        Ok((y0, y1, y2))
14715    }
14716
14717    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
14718    #[allow(clippy::too_many_arguments)]
14719    pub fn qmatvec_q8_fused3_raw(
14720        &self,
14721        b0: &CudaSlice<u8>,
14722        b1: &CudaSlice<u8>,
14723        b2: &CudaSlice<u8>,
14724        x: &CudaSlice<f32>,
14725        in_f: usize,
14726        out0: usize,
14727        out1: usize,
14728        out2: usize,
14729        row_bytes: usize,
14730    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14731        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14732        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14733    }
14734
14735    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
14736    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
14737    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
14738    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
14739    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
14740    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
14741    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
14742    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
14743    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
14744    /// twin must not introduce a batched program the reference path would not run).
14745    pub fn matmul_q8_fused2_t(
14746        &self,
14747        w0: &crate::model::GpuTensor,
14748        w1: &crate::model::GpuTensor,
14749        aq: &CudaSlice<i8>,
14750        ad: &CudaSlice<f32>,
14751        m: usize,
14752    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14753        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
14754        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
14755        // fuses too — same template body, still bit-identical to the two _b8 launches.
14756        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14757            return Ok(None);
14758        }
14759        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
14760        // so the fused b8 launch would introduce a batched program the reference path would not run.
14761        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14762            if m > 4 && !Self::b8_enabled() {
14763                return Ok(None);
14764            }
14765            return Ok(Some(self.e4m3_fused2_t_core(
14766                p0.0,
14767                p1.0,
14768                aq,
14769                ad,
14770                m,
14771                w0.in_features(),
14772                p0.1,
14773                p1.1,
14774                p0.2,
14775                p0.3,
14776                p1.3,
14777            )?));
14778        }
14779        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14780            return Ok(None);
14781        };
14782        Ok(Some(self.q8_fused2_t_core(
14783            p0.0,
14784            p1.0,
14785            aq,
14786            ad,
14787            m,
14788            w0.in_features(),
14789            p0.1,
14790            p1.1,
14791            p0.2,
14792        )?))
14793    }
14794
14795    #[allow(clippy::too_many_arguments)]
14796    fn q8_fused2_t_core(
14797        &self,
14798        b0: &CudaSlice<u8>,
14799        b1: &CudaSlice<u8>,
14800        aq: &CudaSlice<i8>,
14801        ad: &CudaSlice<f32>,
14802        m: usize,
14803        in_f: usize,
14804        out0: usize,
14805        out1: usize,
14806        row_bytes: usize,
14807    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14808        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14809        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14810        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14811        let f = self.func(match Self::batched_mcols(m) {
14812            2 => "qmatvec_q8_0_mmvq_fused2_b2",
14813            4 => "qmatvec_q8_0_mmvq_fused2_b4",
14814            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
14815            _ => "qmatvec_q8_0_mmvq_fused2_b8",
14816        });
14817        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14818        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14819        let cfg = LaunchConfig {
14820            grid_dim: (nb0 + nb1, 1, 1),
14821            block_dim: (32, ROWS_PER_BLOCK, 1),
14822            shared_mem_bytes: 0,
14823        };
14824        let (inf, o0, o1, mi, rbl) = (
14825            in_f as i32,
14826            out0 as i32,
14827            out1 as i32,
14828            m as i32,
14829            row_bytes as i64,
14830        );
14831        let __s_b = self.gpu.stream();
14832        let mut b = __s_b.launch_builder(&f);
14833        b.arg(b0)
14834            .arg(b1)
14835            .arg(aq)
14836            .arg(ad)
14837            .arg(&mut y0)
14838            .arg(&mut y1)
14839            .arg(&inf)
14840            .arg(&o0)
14841            .arg(&o1)
14842            .arg(&mi)
14843            .arg(&rbl);
14844        unsafe {
14845            b.launch(cfg)?;
14846        }
14847        Ok((y0, y1))
14848    }
14849
14850    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14851    /// q8_1 quant of the [m, in_f] activation), no env gating.
14852    #[allow(clippy::too_many_arguments)]
14853    pub fn qmatvec_q8_fused2_t_raw(
14854        &self,
14855        b0: &CudaSlice<u8>,
14856        b1: &CudaSlice<u8>,
14857        x: &CudaSlice<f32>,
14858        m: usize,
14859        in_f: usize,
14860        out0: usize,
14861        out1: usize,
14862        row_bytes: usize,
14863    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14864        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14865        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14866    }
14867
14868    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14869    /// `matmul_q8_fused2_t` with three ranges.
14870    #[allow(clippy::too_many_arguments)]
14871    pub fn matmul_q8_fused3_t(
14872        &self,
14873        w0: &crate::model::GpuTensor,
14874        w1: &crate::model::GpuTensor,
14875        w2: &crate::model::GpuTensor,
14876        aq: &CudaSlice<i8>,
14877        ad: &CudaSlice<f32>,
14878        m: usize,
14879    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14880    {
14881        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14882            return Ok(None);
14883        }
14884        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14885            return Ok(Some(self.e4m3_fused3_t_core(
14886                p0.0,
14887                p1.0,
14888                p2.0,
14889                aq,
14890                ad,
14891                m,
14892                w0.in_features(),
14893                p0.1,
14894                p1.1,
14895                p2.1,
14896                p0.2,
14897                p0.3,
14898                p1.3,
14899                p2.3,
14900            )?));
14901        }
14902        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14903            return Ok(None);
14904        };
14905        Ok(Some(self.q8_fused3_t_core(
14906            p0.0,
14907            p1.0,
14908            p2.0,
14909            aq,
14910            ad,
14911            m,
14912            w0.in_features(),
14913            p0.1,
14914            p1.1,
14915            p2.1,
14916            p0.2,
14917        )?))
14918    }
14919
14920    #[allow(clippy::too_many_arguments)]
14921    fn q8_fused3_t_core(
14922        &self,
14923        b0: &CudaSlice<u8>,
14924        b1: &CudaSlice<u8>,
14925        b2: &CudaSlice<u8>,
14926        aq: &CudaSlice<i8>,
14927        ad: &CudaSlice<f32>,
14928        m: usize,
14929        in_f: usize,
14930        out0: usize,
14931        out1: usize,
14932        out2: usize,
14933        row_bytes: usize,
14934    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14935        const ROWS_PER_BLOCK: u32 = 4;
14936        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14937        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14938        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14939        let f = self.func(if Self::batched_mcols(m) == 2 {
14940            "qmatvec_q8_0_mmvq_fused3_b2"
14941        } else {
14942            "qmatvec_q8_0_mmvq_fused3_b4"
14943        });
14944        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14945        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14946        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14947        let cfg = LaunchConfig {
14948            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14949            block_dim: (32, ROWS_PER_BLOCK, 1),
14950            shared_mem_bytes: 0,
14951        };
14952        let (inf, o0, o1, o2, mi, rbl) = (
14953            in_f as i32,
14954            out0 as i32,
14955            out1 as i32,
14956            out2 as i32,
14957            m as i32,
14958            row_bytes as i64,
14959        );
14960        let __s_b = self.gpu.stream();
14961        let mut b = __s_b.launch_builder(&f);
14962        b.arg(b0)
14963            .arg(b1)
14964            .arg(b2)
14965            .arg(aq)
14966            .arg(ad)
14967            .arg(&mut y0)
14968            .arg(&mut y1)
14969            .arg(&mut y2)
14970            .arg(&inf)
14971            .arg(&o0)
14972            .arg(&o1)
14973            .arg(&o2)
14974            .arg(&mi)
14975            .arg(&rbl);
14976        unsafe {
14977            b.launch(cfg)?;
14978        }
14979        Ok((y0, y1, y2))
14980    }
14981
14982    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
14983    #[allow(clippy::too_many_arguments)]
14984    pub fn qmatvec_q8_fused3_t_raw(
14985        &self,
14986        b0: &CudaSlice<u8>,
14987        b1: &CudaSlice<u8>,
14988        b2: &CudaSlice<u8>,
14989        x: &CudaSlice<f32>,
14990        m: usize,
14991        in_f: usize,
14992        out0: usize,
14993        out1: usize,
14994        out2: usize,
14995        row_bytes: usize,
14996    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14997        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14998        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14999    }
15000
15001    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
15002    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
15003    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
15004    pub fn q8_ffn_fuse2_on(&self) -> bool {
15005        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15006        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15007    }
15008
15009    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
15010    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
15011    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
15012    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
15013    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
15014    #[allow(clippy::type_complexity)]
15015    fn q8_fused_params<'w, const N: usize>(
15016        &self,
15017        ws: &[&'w crate::model::GpuTensor; N],
15018    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15019        use crate::model::GpuTensor;
15020        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15021            return None;
15022        }
15023        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15024            return None;
15025        }
15026        let in_f = ws[0].in_features();
15027        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15028        for (i, w) in ws.iter().enumerate() {
15029            match w {
15030                GpuTensor::Quant {
15031                    bytes,
15032                    qtype,
15033                    row_bytes,
15034                    scale,
15035                    ..
15036                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15037                    out[i] = Some((bytes, w.out_features(), *row_bytes))
15038                }
15039                _ => return None,
15040            }
15041        }
15042        Some(out.map(|o| o.unwrap()))
15043    }
15044
15045    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
15046    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
15047    pub fn e4m3_dual_on(&self) -> bool {
15048        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15049        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15050    }
15051
15052    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
15053    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
15054    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
15055    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
15056    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
15057    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
15058    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
15059    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
15060    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
15061    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
15062    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
15063    #[allow(clippy::type_complexity)]
15064    fn e4m3_fused_params<'w, const N: usize>(
15065        &self,
15066        ws: &[&'w crate::model::GpuTensor; N],
15067    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15068        use crate::model::GpuTensor;
15069        if !self.e4m3_dual_on() {
15070            return None;
15071        }
15072        let in_f = ws[0].in_features();
15073        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15074        for (i, w) in ws.iter().enumerate() {
15075            match w {
15076                GpuTensor::Quant {
15077                    bytes,
15078                    qtype,
15079                    row_bytes,
15080                    scale,
15081                    rp,
15082                    rp4,
15083                    ..
15084                } if *qtype == QT_F8_E4M3
15085                    && w.in_features() == in_f
15086                    && *row_bytes == in_f
15087                    && !*rp
15088                    && rp4.is_none() =>
15089                {
15090                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15091                }
15092                _ => return None,
15093            }
15094        }
15095        Some(out.map(|o| o.unwrap()))
15096    }
15097
15098    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
15099    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
15100    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
15101    #[allow(clippy::too_many_arguments)]
15102    fn e4m3_fused2_core(
15103        &self,
15104        b0: &CudaSlice<u8>,
15105        b1: &CudaSlice<u8>,
15106        aq: &CudaSlice<i8>,
15107        ad: &CudaSlice<f32>,
15108        in_f: usize,
15109        out0: usize,
15110        out1: usize,
15111        row_bytes: usize,
15112        ws0: f32,
15113        ws1: f32,
15114    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15115        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15116        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15117        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15118        let f = self.func("qmatvec_e4m3_mmvq_fused2");
15119        let mut y0 = self.alloc_uninit::<f32>(out0)?;
15120        let mut y1 = self.alloc_uninit::<f32>(out1)?;
15121        let cfg = LaunchConfig {
15122            grid_dim: (nb0 + nb1, 1, 1),
15123            block_dim: (32, ROWS_PER_BLOCK, 1),
15124            shared_mem_bytes: 0,
15125        };
15126        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15127        let __s_b = self.gpu.stream();
15128        let mut b = __s_b.launch_builder(&f);
15129        b.arg(b0)
15130            .arg(b1)
15131            .arg(aq)
15132            .arg(ad)
15133            .arg(&mut y0)
15134            .arg(&mut y1)
15135            .arg(&inf)
15136            .arg(&o0)
15137            .arg(&o1)
15138            .arg(&rbl)
15139            .arg(&ws0)
15140            .arg(&ws1);
15141        unsafe {
15142            b.launch(cfg)?;
15143        }
15144        Ok((y0, y1))
15145    }
15146
15147    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
15148    #[allow(clippy::too_many_arguments)]
15149    fn e4m3_fused3_core(
15150        &self,
15151        b0: &CudaSlice<u8>,
15152        b1: &CudaSlice<u8>,
15153        b2: &CudaSlice<u8>,
15154        aq: &CudaSlice<i8>,
15155        ad: &CudaSlice<f32>,
15156        in_f: usize,
15157        out0: usize,
15158        out1: usize,
15159        out2: usize,
15160        row_bytes: usize,
15161        ws0: f32,
15162        ws1: f32,
15163        ws2: f32,
15164    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15165        const ROWS_PER_BLOCK: u32 = 4;
15166        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15167        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15168        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15169        let f = self.func("qmatvec_e4m3_mmvq_fused3");
15170        let mut y0 = self.alloc_uninit::<f32>(out0)?;
15171        let mut y1 = self.alloc_uninit::<f32>(out1)?;
15172        let mut y2 = self.alloc_uninit::<f32>(out2)?;
15173        let cfg = LaunchConfig {
15174            grid_dim: (nb0 + nb1 + nb2, 1, 1),
15175            block_dim: (32, ROWS_PER_BLOCK, 1),
15176            shared_mem_bytes: 0,
15177        };
15178        let (inf, o0, o1, o2, rbl) = (
15179            in_f as i32,
15180            out0 as i32,
15181            out1 as i32,
15182            out2 as i32,
15183            row_bytes as i64,
15184        );
15185        let __s_b = self.gpu.stream();
15186        let mut b = __s_b.launch_builder(&f);
15187        b.arg(b0)
15188            .arg(b1)
15189            .arg(b2)
15190            .arg(aq)
15191            .arg(ad)
15192            .arg(&mut y0)
15193            .arg(&mut y1)
15194            .arg(&mut y2)
15195            .arg(&inf)
15196            .arg(&o0)
15197            .arg(&o1)
15198            .arg(&o2)
15199            .arg(&rbl)
15200            .arg(&ws0)
15201            .arg(&ws1)
15202            .arg(&ws2);
15203        unsafe {
15204            b.launch(cfg)?;
15205        }
15206        Ok((y0, y1, y2))
15207    }
15208
15209    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
15210    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
15211    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
15212    #[allow(clippy::too_many_arguments)]
15213    fn e4m3_fused2_t_core(
15214        &self,
15215        b0: &CudaSlice<u8>,
15216        b1: &CudaSlice<u8>,
15217        aq: &CudaSlice<i8>,
15218        ad: &CudaSlice<f32>,
15219        m: usize,
15220        in_f: usize,
15221        out0: usize,
15222        out1: usize,
15223        row_bytes: usize,
15224        ws0: f32,
15225        ws1: f32,
15226    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15227        const ROWS_PER_BLOCK: u32 = 4;
15228        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15229        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15230        let f = self.func(match Self::batched_mcols(m) {
15231            2 => "qmatvec_e4m3_mmvq_fused2_b2",
15232            4 => "qmatvec_e4m3_mmvq_fused2_b4",
15233            _ => "qmatvec_e4m3_mmvq_fused2_b8",
15234        });
15235        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15236        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15237        let cfg = LaunchConfig {
15238            grid_dim: (nb0 + nb1, 1, 1),
15239            block_dim: (32, ROWS_PER_BLOCK, 1),
15240            shared_mem_bytes: 0,
15241        };
15242        let (inf, o0, o1, mi, rbl) = (
15243            in_f as i32,
15244            out0 as i32,
15245            out1 as i32,
15246            m as i32,
15247            row_bytes as i64,
15248        );
15249        let __s_b = self.gpu.stream();
15250        let mut b = __s_b.launch_builder(&f);
15251        b.arg(b0)
15252            .arg(b1)
15253            .arg(aq)
15254            .arg(ad)
15255            .arg(&mut y0)
15256            .arg(&mut y1)
15257            .arg(&inf)
15258            .arg(&o0)
15259            .arg(&o1)
15260            .arg(&mi)
15261            .arg(&rbl);
15262        unsafe {
15263            b.launch(cfg)?;
15264        }
15265        if ws0 != 1.0 {
15266            self.scale_inplace(&mut y0, ws0, m * out0)?;
15267        }
15268        if ws1 != 1.0 {
15269            self.scale_inplace(&mut y1, ws1, m * out1)?;
15270        }
15271        Ok((y0, y1))
15272    }
15273
15274    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
15275    #[allow(clippy::too_many_arguments)]
15276    fn e4m3_fused3_t_core(
15277        &self,
15278        b0: &CudaSlice<u8>,
15279        b1: &CudaSlice<u8>,
15280        b2: &CudaSlice<u8>,
15281        aq: &CudaSlice<i8>,
15282        ad: &CudaSlice<f32>,
15283        m: usize,
15284        in_f: usize,
15285        out0: usize,
15286        out1: usize,
15287        out2: usize,
15288        row_bytes: usize,
15289        ws0: f32,
15290        ws1: f32,
15291        ws2: f32,
15292    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15293        const ROWS_PER_BLOCK: u32 = 4;
15294        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15295        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15296        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15297        let f = self.func(if Self::batched_mcols(m) == 2 {
15298            "qmatvec_e4m3_mmvq_fused3_b2"
15299        } else {
15300            "qmatvec_e4m3_mmvq_fused3_b4"
15301        });
15302        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15303        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15304        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15305        let cfg = LaunchConfig {
15306            grid_dim: (nb0 + nb1 + nb2, 1, 1),
15307            block_dim: (32, ROWS_PER_BLOCK, 1),
15308            shared_mem_bytes: 0,
15309        };
15310        let (inf, o0, o1, o2, mi, rbl) = (
15311            in_f as i32,
15312            out0 as i32,
15313            out1 as i32,
15314            out2 as i32,
15315            m as i32,
15316            row_bytes as i64,
15317        );
15318        let __s_b = self.gpu.stream();
15319        let mut b = __s_b.launch_builder(&f);
15320        b.arg(b0)
15321            .arg(b1)
15322            .arg(b2)
15323            .arg(aq)
15324            .arg(ad)
15325            .arg(&mut y0)
15326            .arg(&mut y1)
15327            .arg(&mut y2)
15328            .arg(&inf)
15329            .arg(&o0)
15330            .arg(&o1)
15331            .arg(&o2)
15332            .arg(&mi)
15333            .arg(&rbl);
15334        unsafe {
15335            b.launch(cfg)?;
15336        }
15337        if ws0 != 1.0 {
15338            self.scale_inplace(&mut y0, ws0, m * out0)?;
15339        }
15340        if ws1 != 1.0 {
15341            self.scale_inplace(&mut y1, ws1, m * out1)?;
15342        }
15343        if ws2 != 1.0 {
15344            self.scale_inplace(&mut y2, ws2, m * out2)?;
15345        }
15346        Ok((y0, y1, y2))
15347    }
15348
15349    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
15350    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
15351    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
15352    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
15353    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
15354    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
15355    ///
15356    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
15357    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
15358    pub fn qmatvec_e4m3_blk_mmvq(
15359        &self,
15360        bytes: &CudaSlice<u8>,
15361        aq: &CudaSlice<i8>,
15362        ad: &CudaSlice<f32>,
15363        scales: &CudaSlice<f32>,
15364        m: usize,
15365        in_f: usize,
15366        out_f: usize,
15367        row_bytes: usize,
15368        scale_cols: usize,
15369    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15370        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
15371        self.qmatvec_e4m3_blk_mmvq_into(
15372            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15373        )?;
15374        Ok(y)
15375    }
15376
15377    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
15378    #[allow(clippy::too_many_arguments)]
15379    pub fn qmatvec_e4m3_blk_mmvq_into(
15380        &self,
15381        bytes: &CudaSlice<u8>,
15382        aq: &CudaSlice<i8>,
15383        ad: &CudaSlice<f32>,
15384        scales: &CudaSlice<f32>,
15385        m: usize,
15386        in_f: usize,
15387        out_f: usize,
15388        row_bytes: usize,
15389        scale_cols: usize,
15390        y: &mut CudaSlice<f32>,
15391    ) -> Result<(), Box<dyn std::error::Error>> {
15392        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15393        let f = self.func("qmatvec_e4m3_blk_mmvq");
15394        let cfg = LaunchConfig {
15395            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15396            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
15397            shared_mem_bytes: 0,                // warp-only reduce
15398        };
15399        let (inf, outf, mi, rb, sc) = (
15400            in_f as i32,
15401            out_f as i32,
15402            m as i32,
15403            row_bytes as i64,
15404            scale_cols as i32,
15405        );
15406        let __s_b = self.gpu.stream();
15407        let mut b = __s_b.launch_builder(&f);
15408        b.arg(bytes)
15409            .arg(aq)
15410            .arg(ad)
15411            .arg(scales)
15412            .arg(&mut *y)
15413            .arg(&inf)
15414            .arg(&outf)
15415            .arg(&mi)
15416            .arg(&rb)
15417            .arg(&sc);
15418        unsafe {
15419            b.launch(cfg)?;
15420        }
15421        Ok(())
15422    }
15423
15424    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
15425    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
15426    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
15427    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
15428    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
15429    #[allow(clippy::too_many_arguments)]
15430    pub fn qmatvec_e4m3_blk_mmvq_batched(
15431        &self,
15432        bytes: &CudaSlice<u8>,
15433        aq: &CudaSlice<i8>,
15434        ad: &CudaSlice<f32>,
15435        scales: &CudaSlice<f32>,
15436        m: usize,
15437        in_f: usize,
15438        out_f: usize,
15439        row_bytes: usize,
15440        scale_cols: usize,
15441        mcols: usize,
15442    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15443        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15444        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15445        let name = match mcols {
15446            2 => "qmatvec_e4m3_blk_mmvq_b2",
15447            4 => "qmatvec_e4m3_blk_mmvq_b4",
15448            8 => "qmatvec_e4m3_blk_mmvq_b8",
15449            16 => "qmatvec_e4m3_blk_mmvq_b16",
15450            _ => {
15451                return Err(
15452                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15453                );
15454            }
15455        };
15456        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15457        let f = self.func(name);
15458        let cfg = LaunchConfig {
15459            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15460            block_dim: (32, ROWS_PER_BLOCK, 1),
15461            shared_mem_bytes: 0,
15462        };
15463        let (inf, outf, mi, rb, sc) = (
15464            in_f as i32,
15465            out_f as i32,
15466            m as i32,
15467            row_bytes as i64,
15468            scale_cols as i32,
15469        );
15470        let __s_b = self.gpu.stream();
15471        let mut b = __s_b.launch_builder(&f);
15472        b.arg(bytes)
15473            .arg(aq)
15474            .arg(ad)
15475            .arg(scales)
15476            .arg(&mut y)
15477            .arg(&inf)
15478            .arg(&outf)
15479            .arg(&mi)
15480            .arg(&rb)
15481            .arg(&sc);
15482        unsafe {
15483            b.launch(cfg)?;
15484        }
15485        Ok(y)
15486    }
15487
15488    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
15489    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
15490    #[allow(clippy::too_many_arguments)]
15491    pub fn qmatvec_e4m3_blk_batched_raw(
15492        &self,
15493        bytes: &CudaSlice<u8>,
15494        x: &CudaSlice<f32>,
15495        scales: &CudaSlice<f32>,
15496        m: usize,
15497        in_f: usize,
15498        out_f: usize,
15499        row_bytes: usize,
15500        scale_cols: usize,
15501        mcols: usize,
15502    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15503        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15504        self.qmatvec_e4m3_blk_mmvq_batched(
15505            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15506        )
15507    }
15508
15509    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
15510    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
15511    #[allow(clippy::too_many_arguments)]
15512    pub fn qmatvec_e4m3_blk_mmvq_raw(
15513        &self,
15514        bytes: &CudaSlice<u8>,
15515        x: &CudaSlice<f32>,
15516        scales: &CudaSlice<f32>,
15517        m: usize,
15518        in_f: usize,
15519        out_f: usize,
15520        row_bytes: usize,
15521        scale_cols: usize,
15522    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15523        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15524        self.qmatvec_e4m3_blk_mmvq(
15525            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15526        )
15527    }
15528
15529    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
15530    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
15531    #[allow(clippy::too_many_arguments)]
15532    pub fn qmatvec_e4m3_fused2_raw(
15533        &self,
15534        b0: &CudaSlice<u8>,
15535        b1: &CudaSlice<u8>,
15536        x: &CudaSlice<f32>,
15537        in_f: usize,
15538        out0: usize,
15539        out1: usize,
15540        row_bytes: usize,
15541        ws0: f32,
15542        ws1: f32,
15543    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15544        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15545        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15546    }
15547
15548    #[allow(clippy::too_many_arguments)]
15549    pub fn qmatvec_e4m3_fused3_raw(
15550        &self,
15551        b0: &CudaSlice<u8>,
15552        b1: &CudaSlice<u8>,
15553        b2: &CudaSlice<u8>,
15554        x: &CudaSlice<f32>,
15555        in_f: usize,
15556        out0: usize,
15557        out1: usize,
15558        out2: usize,
15559        row_bytes: usize,
15560        ws0: f32,
15561        ws1: f32,
15562        ws2: f32,
15563    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15564        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15565        self.e4m3_fused3_core(
15566            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15567        )
15568    }
15569
15570    #[allow(clippy::too_many_arguments)]
15571    pub fn qmatvec_e4m3_fused2_t_raw(
15572        &self,
15573        b0: &CudaSlice<u8>,
15574        b1: &CudaSlice<u8>,
15575        x: &CudaSlice<f32>,
15576        m: usize,
15577        in_f: usize,
15578        out0: usize,
15579        out1: usize,
15580        row_bytes: usize,
15581        ws0: f32,
15582        ws1: f32,
15583    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15584        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15585        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15586    }
15587
15588    #[allow(clippy::too_many_arguments)]
15589    pub fn qmatvec_e4m3_fused3_t_raw(
15590        &self,
15591        b0: &CudaSlice<u8>,
15592        b1: &CudaSlice<u8>,
15593        b2: &CudaSlice<u8>,
15594        x: &CudaSlice<f32>,
15595        m: usize,
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, m, in_f)?;
15606        self.e4m3_fused3_t_core(
15607            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15608        )
15609    }
15610
15611    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
15612    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
15613    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
15614    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
15615    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
15616    ///
15617    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
15618    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
15619    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
15620    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
15621    fn try_e4m3_blk_pre(
15622        &self,
15623        w: &crate::model::GpuTensor,
15624        aq: &CudaSlice<i8>,
15625        ad: &CudaSlice<f32>,
15626        m: usize,
15627    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15628        use crate::model::GpuTensor;
15629        if let GpuTensor::Quant {
15630            bytes,
15631            qtype,
15632            row_bytes,
15633            blk: Some(g),
15634            ..
15635        } = w
15636        {
15637            if *qtype == QT_F8_E4M3_BLK {
15638                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
15639                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
15640                // below, so the decode-exactness contract is preserved at every width. Gated by
15641                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
15642                // one rollback door covers every dtype's batched tier.
15643                if (2..=16).contains(&m)
15644                    && std::env::var("MEMRA_NO_BATCHED").is_err()
15645                    && (m <= 4 || Self::b8_enabled())
15646                {
15647                    let mcols = Self::batched_mcols(m);
15648                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15649                        bytes,
15650                        aq,
15651                        ad,
15652                        &g.scales,
15653                        m,
15654                        w.in_features(),
15655                        w.out_features(),
15656                        *row_bytes,
15657                        g.cols,
15658                        mcols,
15659                    )?));
15660                }
15661                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15662                    bytes,
15663                    aq,
15664                    ad,
15665                    &g.scales,
15666                    m,
15667                    w.in_features(),
15668                    w.out_features(),
15669                    *row_bytes,
15670                    g.cols,
15671                )?));
15672            }
15673        }
15674        Ok(None)
15675    }
15676
15677    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
15678    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
15679    ///
15680    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
15681    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
15682    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
15683    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
15684    /// prefill keeps the floor's arithmetic and the floor's kernels.
15685    ///
15686    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
15687    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
15688    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
15689    /// (projection, prefill call) and frees immediately.
15690    ///
15691    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
15692    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
15693    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
15694    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
15695    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
15696    /// single-variable comparison instead of a two-variable one.
15697    ///
15698    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
15699    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
15700    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
15701    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
15702    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
15703    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
15704    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
15705    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
15706    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
15707    ///
15708    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
15709    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
15710    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
15711    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
15712    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
15713    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
15714    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
15715    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
15716    /// because v2's denominator had its slab already resident while this class's floor must build it
15717    /// every call; same tile, opposite sign, because the question changed.
15718    ///
15719    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
15720    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
15721    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
15722    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
15723    fn try_e4m3_blk_prefill(
15724        &self,
15725        w: &crate::model::GpuTensor,
15726        x: &CudaSlice<f32>,
15727        m: usize,
15728    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15729        use crate::model::GpuTensor;
15730        let GpuTensor::Quant {
15731            bytes,
15732            qtype,
15733            blk: Some(g),
15734            ..
15735        } = w
15736        else {
15737            return Ok(None);
15738        };
15739        if *qtype != QT_F8_E4M3_BLK {
15740            return Ok(None);
15741        }
15742        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
15743        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
15744        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
15745        // through to the dequant below when they do, never silently produce nothing.
15746        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15747            return Ok(Some(y));
15748        }
15749        let (in_f, out_f) = (w.in_features(), w.out_features());
15750        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15751        let tmp = GpuTensor::Quant {
15752            bytes: slab,
15753            qtype: QT_Q8_0,
15754            row_bytes: in_f / 32 * 34,
15755            ne: vec![in_f as u64, out_f as u64],
15756            scale: 1.0,
15757            rp: false,
15758            #[cfg(memra_cutlass)]
15759            cutlass: None,
15760            fp8: None,
15761            blk: None,
15762            f16: None,
15763            rp4: None,
15764        };
15765        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
15766        Ok(Some(self.matmul(&tmp, x, m)?))
15767    }
15768
15769    pub fn matmul_pre_noscale(
15770        &self,
15771        w: &crate::model::GpuTensor,
15772        aq: &CudaSlice<i8>,
15773        ad: &CudaSlice<f32>,
15774        m: usize,
15775    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15776        use crate::model::GpuTensor;
15777        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
15778        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
15779        // rather than let the tail below refuse and cost the caller a re-dispatch.
15780        if m == 1 {
15781            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15782                return Ok(Some((y, 1.0)));
15783            }
15784        }
15785        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
15786        if m != 1 || !self.uses_q8_1_fast(w) {
15787            return Ok(None);
15788        }
15789        let in_f = w.in_features();
15790        let out_f = w.out_features();
15791        let (bytes, qtype, row_bytes, scale, rp) = match w {
15792            GpuTensor::Quant {
15793                bytes,
15794                qtype,
15795                row_bytes,
15796                scale,
15797                rp,
15798                ..
15799            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15800            _ => return Ok(None),
15801        };
15802        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
15803        if self.mmvq_supports(qtype) {
15804            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
15805            let (mbytes, mrp) = match w {
15806                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15807                _ => (bytes, rp),
15808            };
15809            let y = self.qmatvec_mmvq(
15810                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
15811            )?;
15812            return Ok(Some((y, scale)));
15813        }
15814        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
15815        let name = match qtype {
15816            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15817            QT_Q4_K => "qmatvec_q4_K_dp4a",
15818            QT_Q6_K => "qmatvec_q6_K_dp4a",
15819            QT_Q5_K => "qmatvec_q5_K_dp4a",
15820            QT_Q3_K => "qmatvec_q3_K_dp4a",
15821            QT_NVFP4 => {
15822                if rp {
15823                    "qmatvec_nvfp4_dp4a_rp"
15824                } else {
15825                    "qmatvec_nvfp4_dp4a"
15826                }
15827            }
15828            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15829            _ => return Ok(None),
15830        };
15831        let f = self.func(name);
15832        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15833        let cfg = LaunchConfig {
15834            grid_dim: (out_f as u32, m as u32, 1),
15835            block_dim: (128, 1, 1),
15836            shared_mem_bytes: 0,
15837        };
15838        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15839        let __s_b = self.gpu.stream();
15840        let mut b = __s_b.launch_builder(&f);
15841        b.arg(bytes)
15842            .arg(aq)
15843            .arg(ad)
15844            .arg(&mut y)
15845            .arg(&inf)
15846            .arg(&outf)
15847            .arg(&mi)
15848            .arg(&rb);
15849        unsafe {
15850            b.launch(cfg)?;
15851        }
15852        Ok(Some((y, scale)))
15853    }
15854
15855    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15856    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15857    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15858        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15859        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15860        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15861        // is a pure function of the dtype — the decode-parity law holds under every env.
15862        if qtype == QT_F8_E4M3 {
15863            return true;
15864        }
15865        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15866            return false;
15867        }
15868        matches!(
15869            qtype,
15870            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15871        )
15872    }
15873
15874    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15875    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15876    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15877    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15878    pub fn qmatvec_mmvq(
15879        &self,
15880        bytes: &CudaSlice<u8>,
15881        aq: &CudaSlice<i8>,
15882        ad: &CudaSlice<f32>,
15883        m: usize,
15884        in_f: usize,
15885        out_f: usize,
15886        qtype: i32,
15887        row_bytes: usize,
15888        scale: f32,
15889        rp: bool,
15890    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15891        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15892        self.qmatvec_mmvq_into(
15893            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15894        )?;
15895        Ok(y)
15896    }
15897
15898    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15899    #[allow(clippy::too_many_arguments)]
15900    pub fn qmatvec_mmvq_into(
15901        &self,
15902        bytes: &CudaSlice<u8>,
15903        aq: &CudaSlice<i8>,
15904        ad: &CudaSlice<f32>,
15905        m: usize,
15906        in_f: usize,
15907        out_f: usize,
15908        qtype: i32,
15909        row_bytes: usize,
15910        scale: f32,
15911        rp: bool,
15912        y: &mut CudaSlice<f32>,
15913    ) -> Result<(), Box<dyn std::error::Error>> {
15914        debug_assert!(y.len() >= m * out_f);
15915        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15916        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15917        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15918        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15919        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15920        if qtype == QT_Q8_0
15921            && rp
15922            && m == 1
15923            && out_f >= 64
15924            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15925            && {
15926                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15927                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15928            }
15929        {
15930            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15931            let cfg = LaunchConfig {
15932                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15933                block_dim: (32, 2, 1),
15934                shared_mem_bytes: 0,
15935            };
15936            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15937            let __s_b = self.gpu.stream();
15938            let mut b = __s_b.launch_builder(&f);
15939            b.arg(bytes)
15940                .arg(aq)
15941                .arg(ad)
15942                .arg(&mut *y)
15943                .arg(&inf)
15944                .arg(&outf)
15945                .arg(&mi)
15946                .arg(&rb);
15947            unsafe {
15948                b.launch(cfg)?;
15949            }
15950            if scale != 1.0 {
15951                self.scale_inplace(y, scale, out_f)?;
15952            }
15953            return Ok(());
15954        }
15955        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15956        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15957        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15958        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
15959        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
15960        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
15961        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
15962        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
15963        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15964            2
15965        } else {
15966            1
15967        };
15968        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
15969        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
15970        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
15971        // valid-window interleaved, bit-identical per row — same dot program).
15972        if m == 1 && qtype == QT_Q4_0 {
15973            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15974            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
15975            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
15976            mr = *Q40MR.get_or_init(|| {
15977                std::env::var("MEMRA_Q40_MR")
15978                    .ok()
15979                    .and_then(|v| v.parse().ok())
15980                    .unwrap_or(1)
15981            });
15982        }
15983        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
15984        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
15985        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
15986        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
15987        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
15988        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
15989        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
15990        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
15991        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
15992        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
15993        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15994        let q5_force = q5_mode.as_deref() == Some("2");
15995        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
15996        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
15997        let q5_il = qtype == QT_Q5_K
15998            && m == 1
15999            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16000        if q5_il && !q5_force && out_f > 65536 {
16001            mr = 1;
16002        }
16003        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
16004        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
16005        if qtype == QT_Q4_0 && rp && mr != 1 {
16006            mr = 2;
16007        }
16008        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
16009        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
16010        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
16011        if qtype == QT_Q8_0 && rp {
16012            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16013            mr = *Q80MR.get_or_init(|| {
16014                std::env::var("MEMRA_Q80_MR")
16015                    .ok()
16016                    .and_then(|v| v.parse().ok())
16017                    .unwrap_or(1)
16018            });
16019        }
16020        let name = match (qtype, mr, rp) {
16021            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16022            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16023            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16024            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16025            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16026            (QT_Q5_K, 2, _) => {
16027                if q5_il {
16028                    "qmatvec_q5_K_mmvq_mr2_il"
16029                } else {
16030                    "qmatvec_q5_K_mmvq_mr2"
16031                }
16032            }
16033            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16034            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
16035            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
16036            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
16037            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
16038            (QT_Q8_0, _, true)
16039                if in_f % 1024 == 0 && {
16040                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16041                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16042                } =>
16043            {
16044                "qmatvec_q8_0_mmvq_rpca"
16045            }
16046            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16047            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16048            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
16049            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
16050            // reach a GGUF-layout kernel or vice versa.
16051            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16052            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16053            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16054            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16055            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16056            (QT_Q5_K, _, _) => {
16057                if q5_il {
16058                    "qmatvec_q5_K_mmvq_il"
16059                } else {
16060                    "qmatvec_q5_K_mmvq"
16061                }
16062            }
16063            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16064            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16065            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16066            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16067        };
16068        let f = self.func(name);
16069        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
16070        let rows_per_block = ROWS_PER_BLOCK * mr;
16071        let cfg = LaunchConfig {
16072            grid_dim: (
16073                (out_f as u32 + rows_per_block - 1) / rows_per_block,
16074                m as u32,
16075                1,
16076            ),
16077            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
16078            shared_mem_bytes: 0,                // warp-only reduce at m=1
16079        };
16080        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16081        let __s_b = self.gpu.stream();
16082        let mut b = __s_b.launch_builder(&f);
16083        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
16084        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
16085        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
16086        // weight_scale). Other mmvq kernels keep the 8-arg signature.
16087        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16088            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
16089            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
16090            if Self::pdl_on()
16091                && Self::pdl_mmvq_on()
16092                && Self::pdl_nvfp4q8_on()
16093                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16094            {
16095                use cudarc::driver::{DevicePtr, DevicePtrMut};
16096                let s = &self.gpu.stream();
16097                let (pw, _g0) = bytes.device_ptr(s);
16098                let (paq, _g1) = aq.device_ptr(s);
16099                let (pad, _g2) = ad.device_ptr(s);
16100                let (py, _g3) = y.device_ptr_mut(s);
16101                let mut ps = [
16102                    &pw as *const _ as *mut std::ffi::c_void,
16103                    &paq as *const _ as *mut _,
16104                    &pad as *const _ as *mut _,
16105                    &py as *const _ as *mut _,
16106                    &inf as *const _ as *mut _,
16107                    &outf as *const _ as *mut _,
16108                    &mi as *const _ as *mut _,
16109                    &rb as *const _ as *mut _,
16110                    &scale as *const _ as *mut _,
16111                ];
16112                unsafe {
16113                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16114                }
16115                return Ok(());
16116            }
16117            b.arg(bytes)
16118                .arg(aq)
16119                .arg(ad)
16120                .arg(&mut *y)
16121                .arg(&inf)
16122                .arg(&outf)
16123                .arg(&mi)
16124                .arg(&rb)
16125                .arg(&scale);
16126            unsafe {
16127                b.launch(cfg)?;
16128            }
16129        } else if Self::pdl_on()
16130            && Self::pdl_mmvq_on()
16131            && (matches!(
16132                name,
16133                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16134            ) || (Self::pdl_nvfp4q8_on()
16135                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16136        {
16137            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
16138            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
16139            // names may take this launch (unmarked kernels would read unordered).
16140            {
16141                use cudarc::driver::{DevicePtr, DevicePtrMut};
16142                let s = &self.gpu.stream();
16143                let (pw, _g0) = bytes.device_ptr(s);
16144                let (paq, _g1) = aq.device_ptr(s);
16145                let (pad, _g2) = ad.device_ptr(s);
16146                let (py, _g3) = y.device_ptr_mut(s);
16147                let mut ps = [
16148                    &pw as *const _ as *mut std::ffi::c_void,
16149                    &paq as *const _ as *mut _,
16150                    &pad as *const _ as *mut _,
16151                    &py as *const _ as *mut _,
16152                    &inf as *const _ as *mut _,
16153                    &outf as *const _ as *mut _,
16154                    &mi as *const _ as *mut _,
16155                    &rb as *const _ as *mut _,
16156                ];
16157                unsafe {
16158                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16159                }
16160            }
16161            if scale != 1.0 {
16162                self.scale_inplace(y, scale, m * out_f)?;
16163            }
16164        } else {
16165            b.arg(bytes)
16166                .arg(aq)
16167                .arg(ad)
16168                .arg(&mut *y)
16169                .arg(&inf)
16170                .arg(&outf)
16171                .arg(&mi)
16172                .arg(&rb);
16173            unsafe {
16174                b.launch(cfg)?;
16175            }
16176            if scale != 1.0 {
16177                self.scale_inplace(y, scale, m * out_f)?;
16178            }
16179        }
16180        Ok(())
16181    }
16182
16183    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
16184    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
16185    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
16186    pub fn qmatvec_mmvq_raw(
16187        &self,
16188        bytes: &CudaSlice<u8>,
16189        x: &CudaSlice<f32>,
16190        m: usize,
16191        in_f: usize,
16192        out_f: usize,
16193        qtype: i32,
16194        row_bytes: usize,
16195        rp: bool,
16196    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16197        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16198        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16199    }
16200
16201    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
16202    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
16203    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
16204    pub fn batched_supports(&self, qtype: i32) -> bool {
16205        matches!(
16206            qtype,
16207            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16208        )
16209    }
16210
16211    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
16212    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
16213    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
16214    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
16215    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
16216    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
16217    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
16218    pub fn iq_fast_enabled() -> bool {
16219        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16220        *ON.get_or_init(|| {
16221            std::env::var("MEMRA_IQ_FAST")
16222                .map(|v| v != "0")
16223                .unwrap_or(true)
16224        })
16225    }
16226
16227    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
16228    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
16229    pub fn b8_enabled() -> bool {
16230        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16231        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16232    }
16233
16234    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
16235    pub fn batched_mcols(m: usize) -> usize {
16236        if m == 2 {
16237            2
16238        } else if m <= 4 {
16239            4
16240        } else if m <= 8 {
16241            8
16242        } else {
16243            16
16244        }
16245    }
16246
16247    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
16248    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
16249    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
16250    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
16251    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16252        Some(match (qtype, mcols) {
16253            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16254            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16255            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16256            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
16257            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
16258            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
16259            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
16260            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
16261            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16262            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16263            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16264            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16265            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
16266            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
16267            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16268            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16269            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16270            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16271            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
16272            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
16273            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16274            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16275            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16276            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16277            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16278            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16279            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16280            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16281            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
16282            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
16283            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
16284            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16285            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16286            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16287            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16288            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
16289            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
16290            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
16291            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16292            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16293            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16294            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16295            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16296            _ => return None,
16297        })
16298    }
16299
16300    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
16301    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
16302    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
16303    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
16304    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
16305    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
16306    ///
16307    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
16308    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
16309    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
16310    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
16311    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
16312    /// msweep on all six 27B shapes (2026-07-03):
16313    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
16314    ///          it applies for b4 (-3..-14%), never loses;
16315    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
16316    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
16317    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
16318    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
16319    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
16320    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
16321    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
16322    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
16323    /// b2: in_f>=6144 -> r2, else base.
16324    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
16325    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
16326    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
16327    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
16328    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
16329    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
16330    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
16331    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
16332    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
16333    /// Device SM count (cached) — grid-fill policy input.
16334    pub fn sm_count(&self) -> i32 {
16335        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16336        *SMS.get_or_init(|| {
16337            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16338            self.gpu
16339                .ctx
16340                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16341                .unwrap_or(82)
16342        })
16343    }
16344
16345    pub fn batched_variant(
16346        &self,
16347        _m: usize,
16348        in_f: usize,
16349        out_f: usize,
16350        qtype: i32,
16351        row_bytes: usize,
16352        mcols: usize,
16353        rp: bool,
16354    ) -> &'static str {
16355        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
16356        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
16357        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
16358        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
16359        if qtype == QT_Q8_0 {
16360            return if rp { "rp" } else { "base" };
16361        }
16362        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16363        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16364            Ok("base") => "base",
16365            Ok("pf") => "pf",
16366            Ok("r2") => "r2",
16367            Ok("r2w8") => "r2w8",
16368            Ok("pfr2") => "pfr2",
16369            Ok("ca") => "ca",
16370            Ok("car2") => "car2",
16371            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
16372            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
16373            Ok("rp") => "rp",
16374            Ok("rpr2") => "rpr2",
16375            Ok("rpr2w8") => "rpr2w8",
16376            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
16377            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
16378            Ok("rpca") => "rpca",
16379            Ok("rpcar2") => "rpcar2",
16380            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
16381            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
16382            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
16383            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
16384            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
16385            // bit-identical to the decode path — measurement corpus ONLY, never auto).
16386            Ok("rpsc") => "rpsc",
16387            Ok("rpms") => "rpms",
16388            Ok("rpmsc") => "rpmsc",
16389            Ok("rpks") => "rpks",
16390            Ok("rpksc") => "rpksc",
16391            _ => "auto",
16392        });
16393        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
16394        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
16395        // shapes qualify; anything else falls back to the register variants.
16396        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16397        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
16398        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
16399        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
16400        // forced MEMRA_MMVQ_BV values still work).
16401        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16402        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16403        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16404        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16405        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16406        let sms = *SMS.get_or_init(|| {
16407            use cudarc::driver::sys::CUdevice_attribute_enum as A;
16408            self.gpu
16409                .ctx
16410                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16411                .unwrap_or(82)
16412        });
16413        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
16414        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
16415        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
16416        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
16417        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
16418        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
16419        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
16420        // AUTO RULE = the measured winners table (differs from NVFP4's!):
16421        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
16422        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
16423        //     r2 1258us) — kernels kept behind the force seam for the corpus;
16424        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
16425        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
16426        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
16427        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
16428        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
16429        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
16430        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
16431        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
16432        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16433        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
16434        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
16435        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16436        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16437            Ok("base") => "base",
16438            Ok("r2") => "r2",
16439            Ok("r2w8") => "r2w8",
16440            _ => "auto",
16441        });
16442        let variant: &'static str = if qtype == QT_Q4_0 {
16443            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
16444            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
16445            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
16446            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16447            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16448                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
16449                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
16450                // + syncs cost more than the stalls, bank-pad made no difference);
16451                // register load-ahead flat (nvcc already reorders). The b-tier limiter
16452                // is still unidentified — see the jsonl row.
16453                Ok("base") => "base",
16454                Ok("r2") => "r2",
16455                Ok("ms") => "ms",
16456                Ok("sm") => "sm",
16457                Ok("la") => "la",
16458                _ => "auto",
16459            });
16460            let v = if q40 != "auto" {
16461                q40
16462            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16463                "r2"
16464            } else {
16465                "base"
16466            };
16467            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
16468            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
16469            // and the limiter is the per-column activation load chain (long_scoreboard
16470            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
16471            if rp {
16472                match v {
16473                    "ms" => "r2ms_rp",
16474                    "sm" => "r2sm_rp",
16475                    "la" => "r2la_rp",
16476                    "r2" => "r2_rp",
16477                    _ => "rp",
16478                }
16479            } else if matches!(v, "ms" | "sm" | "la") {
16480                "r2"
16481            } else {
16482                v
16483            }
16484        } else if qtype != QT_NVFP4 && !kq_r2 {
16485            "base"
16486        } else if kq_r2 && rp {
16487            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
16488            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
16489            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
16490            "rp"
16491        } else if kq_r2 {
16492            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
16493            // mcols != 4 forced r2w8 falls to unbounded r2.
16494            if kq_bv != "auto" {
16495                if kq_bv == "r2w8" && mcols != 4 {
16496                    "r2"
16497                } else {
16498                    kq_bv
16499                }
16500            } else if bv != "auto" {
16501                match bv {
16502                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16503                    "r2w8" | "rpr2w8" => {
16504                        if mcols != 4 {
16505                            "r2"
16506                        } else {
16507                            "r2w8"
16508                        }
16509                    }
16510                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
16511                }
16512            } else {
16513                let blocks = (out_f + 7) / 8;
16514                let waves = blocks as f64 / (7 * sms as usize) as f64;
16515                let filled = blocks >= 4 * sms as usize;
16516                let use_r2 = if qtype == QT_Q4_K {
16517                    filled
16518                } else {
16519                    waves >= 2.0
16520                };
16521                if use_r2 { "r2" } else { "base" }
16522            }
16523        } else if bv != "auto" {
16524            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
16525            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
16526            // unsupported (shape, mcols) combos fall back to pf/r2.
16527            // On rp buffers, forced legacy names map to their rp twins (layout law).
16528            let v = if bv == "r2w8" && mcols == 2 {
16529                "r2"
16530            } else if bv == "ca" && (!ca_ok || mcols == 8) {
16531                "pf"
16532            } else if bv == "car2" && (!ca_ok || mcols == 8) {
16533                "r2"
16534            } else if bv == "pfr2" && mcols == 8 {
16535                "r2"
16536            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16537                "rpr2"
16538            }
16539            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
16540            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16541                if mcols == 8 { "rpr2w8" } else { "rpr2" }
16542            } else if bv == "rpcar2" && mcols == 2 {
16543                "rpca"
16544            }
16545            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
16546            // (rpms has no smem and no alignment need — always valid on rp buffers).
16547            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16548                "rpr2"
16549            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16550                "rpr2"
16551            } else {
16552                bv
16553            };
16554            if rp {
16555                match v {
16556                    "base" | "pf" | "ca" | "rp" => "rp",
16557                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16558                    "r2w8" | "rpr2w8" => {
16559                        if mcols == 2 {
16560                            "rpr2"
16561                        } else {
16562                            "rpr2w8"
16563                        }
16564                    }
16565                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
16566                }
16567            } else {
16568                v
16569            }
16570        } else if mcols == 8 {
16571            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
16572            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
16573            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
16574            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
16575            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
16576            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
16577            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
16578            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
16579            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
16580            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
16581            if rp {
16582                if sc_ok { "rpsc" } else { "rpr2w8" }
16583            } else {
16584                "r2w8"
16585            }
16586        } else if mcols >= 4 {
16587            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
16588            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
16589            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
16590            let blocks = (out_f + 7) / 8;
16591            let r7 = 7 * sms as usize;
16592            let r8 = 8 * sms as usize;
16593            let waves = blocks as f64 / r7 as f64;
16594            let filled = blocks >= 4 * sms as usize;
16595            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
16596            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
16597            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
16598            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16599                // the extra residency drops the INTEGER wave count -> the straggler wave a
16600                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
16601                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
16602                if rp { "rpr2w8" } else { "r2w8" }
16603            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16604                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
16605                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
16606                if rp { "rpr2" } else { "r2" }
16607            } else {
16608                // fractional straggler-wave window with no crossing, or grid too small to fill
16609                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
16610                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
16611                if rp { "rp" } else { "pf" }
16612            }
16613        } else if in_f >= 6144 {
16614            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
16615            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
16616            // stays.
16617            if rp { "rpr2" } else { "r2" }
16618        } else if rp {
16619            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
16620            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
16621            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
16622            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
16623            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16624            if sc_ok && waves >= 0.9 && waves <= 1.1 {
16625                "rpsc"
16626            } else {
16627                "rp"
16628            }
16629        } else {
16630            "base"
16631        };
16632        variant
16633    }
16634
16635    pub fn qmatvec_mmvq_batched(
16636        &self,
16637        bytes: &CudaSlice<u8>,
16638        aq: &CudaSlice<i8>,
16639        ad: &CudaSlice<f32>,
16640        m: usize,
16641        in_f: usize,
16642        out_f: usize,
16643        qtype: i32,
16644        row_bytes: usize,
16645        mcols: usize,
16646        scale: f32,
16647        rp: bool,
16648    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16649        const ROWS_PER_BLOCK: u32 = 4;
16650        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
16651        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
16652        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
16653        // weight keeps its rp-layout kernel family regardless of the override.
16654        let forced: Option<&'static str> = {
16655            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16656            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16657                .as_deref()
16658                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16659        };
16660        let variant = match forced {
16661            Some(v) if !rp || v.contains("rp") => v,
16662            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16663        };
16664        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16665            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16666        })?;
16667        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
16668        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
16669        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
16670        let variant = if mcols == 16 {
16671            if rp { "rp" } else { "base" }
16672        } else {
16673            variant
16674        };
16675        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
16676        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
16677        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
16678        // per-(token,row) chain (columns c >= m never execute in either form) ->
16679        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
16680        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
16681        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16682        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16683        if b567
16684            && qtype == QT_NVFP4
16685            && rp
16686            && mcols == 8
16687            && (5..=7).contains(&m)
16688            && matches!(variant, "rpsc" | "rpr2w8")
16689        {
16690            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16691            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
16692            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16693            let cfg = LaunchConfig {
16694                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16695                block_dim: (32, ROWS_PER_BLOCK, 1),
16696                shared_mem_bytes: 0,
16697            };
16698            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16699            let __s_b = self.gpu.stream();
16700            let mut b = __s_b.launch_builder(&f);
16701            b.arg(bytes)
16702                .arg(aq)
16703                .arg(ad)
16704                .arg(&mut y)
16705                .arg(&inf)
16706                .arg(&outf)
16707                .arg(&mi)
16708                .arg(&rb);
16709            unsafe {
16710                b.launch(cfg)?;
16711            }
16712            if scale != 1.0 {
16713                self.scale_inplace(&mut y, scale, m * out_f)?;
16714            }
16715            return Ok(y);
16716        }
16717        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16718            "base" => (base_name.into(), ROWS_PER_BLOCK),
16719            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16720            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16721            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16722            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
16723            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
16724            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
16725            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16726            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16727            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16728            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16729            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16730            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16731            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16732            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
16733        };
16734        debug_assert!(
16735            !rp || name.contains("_rp"),
16736            "rp weight dispatched to a GGUF-layout kernel"
16737        );
16738        let f = self.func(&name);
16739        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16740        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
16741        let smem = if name.contains("_r2sm_rp") {
16742            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16743        } else {
16744            0
16745        };
16746        let cfg = LaunchConfig {
16747            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16748            block_dim: (32, ROWS_PER_BLOCK, 1),
16749            shared_mem_bytes: smem,
16750        };
16751        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16752        let __s_b = self.gpu.stream();
16753        let mut b = __s_b.launch_builder(&f);
16754        b.arg(bytes)
16755            .arg(aq)
16756            .arg(ad)
16757            .arg(&mut y)
16758            .arg(&inf)
16759            .arg(&outf)
16760            .arg(&mi)
16761            .arg(&rb);
16762        unsafe {
16763            b.launch(cfg)?;
16764        }
16765        if scale != 1.0 {
16766            self.scale_inplace(&mut y, scale, m * out_f)?;
16767        }
16768        Ok(y)
16769    }
16770
16771    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
16772    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
16773    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
16774    pub fn qmatvec_batched_raw(
16775        &self,
16776        bytes: &CudaSlice<u8>,
16777        x: &CudaSlice<f32>,
16778        m: usize,
16779        in_f: usize,
16780        out_f: usize,
16781        qtype: i32,
16782        row_bytes: usize,
16783        mcols: usize,
16784        rp: bool,
16785    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16786        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16787        self.qmatvec_mmvq_batched(
16788            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16789        )
16790    }
16791
16792    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
16793    pub fn qmatvec_nvfp4_batched_raw(
16794        &self,
16795        bytes: &CudaSlice<u8>,
16796        x: &CudaSlice<f32>,
16797        m: usize,
16798        in_f: usize,
16799        out_f: usize,
16800        row_bytes: usize,
16801        mcols: usize,
16802        rp: bool,
16803    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16804        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16805    }
16806
16807    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
16808    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
16809    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
16810    fn try_fp4_gemm(
16811        &self,
16812        w: &crate::model::GpuTensor,
16813        x: &CudaSlice<f32>,
16814        m: usize,
16815        in_f: usize,
16816        out_f: usize,
16817    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16818        use crate::model::GpuTensor;
16819        if cfg!(memra_portable_cuda) {
16820            return Ok(None);
16821        }
16822        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which cu/qmatvec_gemm.cu:1234 omits on a
16823        // portable build (the mxf4 block-scale MMA is sm_120a-only). Refuse at the door.
16824        if std::env::var("MEMRA_FP4").is_ok() {
16825            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16826        }
16827        if std::env::var("MEMRA_FP4").is_err() {
16828            return Ok(None);
16829        }
16830        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
16831        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
16832        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
16833        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
16834        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
16835        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
16836        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
16837        // for the common no-macro-scale case.
16838        #[cfg(memra_cutlass)]
16839        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16840            if let GpuTensor::Quant {
16841                bytes,
16842                qtype,
16843                scale,
16844                row_bytes,
16845                cutlass,
16846                ..
16847            } = w
16848            {
16849                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16850                    if let Some(cw) = cutlass {
16851                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16852                        let y = self.cutlass_fp4_gemm(
16853                            &cw.b_packed,
16854                            &cw.sfb_swizzled,
16855                            x,
16856                            *scale,
16857                            m,
16858                            out_f,
16859                            in_f,
16860                        )?;
16861                        return Ok(Some(y));
16862                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16863                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16864                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16865                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16866                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16867                        let (b_packed, sfb_sw) =
16868                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16869                        let y =
16870                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16871                        return Ok(Some(y));
16872                    }
16873                }
16874            }
16875        }
16876        if let GpuTensor::Quant {
16877            bytes,
16878            qtype,
16879            row_bytes,
16880            scale,
16881            rp,
16882            ..
16883        } = w
16884        {
16885            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16886            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16887            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16888                let y =
16889                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16890                return Ok(Some(y));
16891            }
16892        }
16893        Ok(None)
16894    }
16895
16896    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16897    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16898    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16899    pub fn rms_norm_f16out(
16900        &self,
16901        x: &CudaSlice<f32>,
16902        w: &CudaSlice<f32>,
16903        dst: &mut CudaSlice<f32>,
16904        dst16: &mut CudaSlice<u8>,
16905        ncols: usize,
16906        nrows: usize,
16907        eps: f32,
16908    ) -> Result<(), Box<dyn std::error::Error>> {
16909        let f = self.func("rms_norm_f16out_f32");
16910        let cfg = LaunchConfig {
16911            grid_dim: (nrows as u32, 1, 1),
16912            block_dim: (rms_block(), 1, 1),
16913            shared_mem_bytes: 0,
16914        };
16915        let (nc, e) = (ncols as i32, eps);
16916        let __s_b = self.gpu.stream();
16917        let mut b = __s_b.launch_builder(&f);
16918        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16919        unsafe {
16920            b.launch(cfg)?;
16921        }
16922        Ok(())
16923    }
16924
16925    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16926    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16927    #[allow(clippy::too_many_arguments)]
16928    pub fn add_rms_norm_f16out(
16929        &self,
16930        a: &CudaSlice<f32>,
16931        b: &CudaSlice<f32>,
16932        w: &CudaSlice<f32>,
16933        res: &mut CudaSlice<f32>,
16934        dst: &mut CudaSlice<f32>,
16935        dst16: &mut CudaSlice<u8>,
16936        ncols: usize,
16937        nrows: usize,
16938        eps: f32,
16939    ) -> Result<(), Box<dyn std::error::Error>> {
16940        let f = self.func("add_rms_norm_f16out_f32");
16941        let cfg = LaunchConfig {
16942            grid_dim: (nrows as u32, 1, 1),
16943            block_dim: (rms_block(), 1, 1),
16944            shared_mem_bytes: 0,
16945        };
16946        let (nc, e) = (ncols as i32, eps);
16947        let __s_lb = self.gpu.stream();
16948        let mut lb = __s_lb.launch_builder(&f);
16949        lb.arg(a)
16950            .arg(b)
16951            .arg(w)
16952            .arg(res)
16953            .arg(dst)
16954            .arg(dst16)
16955            .arg(&nc)
16956            .arg(&e);
16957        unsafe {
16958            lb.launch(cfg)?;
16959        }
16960        Ok(())
16961    }
16962
16963    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
16964    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
16965    pub fn matmul_group_xh(
16966        &self,
16967        ws: &[&crate::model::GpuTensor],
16968        x: &CudaSlice<f32>,
16969        xh: &CudaSlice<u8>,
16970        m: usize,
16971    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16972        let mut out = Vec::with_capacity(ws.len());
16973        let in_f = ws[0].in_features();
16974        for w in ws {
16975            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16976                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16977                    out.push(y);
16978                    continue;
16979                }
16980            }
16981            out.push(self.matmul(w, x, m)?);
16982        }
16983        Ok(out)
16984    }
16985
16986    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
16987    /// GDN steps). Layouts [T, H].
16988    pub fn gdn_pad_mask(
16989        &self,
16990        beta: &mut CudaSlice<f32>,
16991        g_log: &mut CudaSlice<f32>,
16992        len_d: &CudaSlice<i32>,
16993        h: usize,
16994        t: usize,
16995    ) -> Result<(), Box<dyn std::error::Error>> {
16996        let f = self.func("gdn_pad_mask_f32");
16997        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16998        let (hi, ti) = (h as i32, t as i32);
16999        let __s_b = self.gpu.stream();
17000        let mut b = __s_b.launch_builder(&f);
17001        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17002        unsafe {
17003            b.launch(cfg)?;
17004        }
17005        Ok(())
17006    }
17007
17008    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
17009    /// gather for the padded prime graph's h_seed/hlast.
17010    pub fn row_gather_dev(
17011        &self,
17012        src: &CudaSlice<f32>,
17013        dst: &mut CudaSlice<f32>,
17014        len_d: &CudaSlice<i32>,
17015        ncols: usize,
17016    ) -> Result<(), Box<dyn std::error::Error>> {
17017        let f = self.func("row_gather_dev_f32");
17018        let cfg = LaunchConfig::for_num_elems(ncols as u32);
17019        let nc = ncols as i32;
17020        let __s_b = self.gpu.stream();
17021        let mut b = __s_b.launch_builder(&f);
17022        b.arg(src).arg(dst).arg(len_d).arg(&nc);
17023        unsafe {
17024            b.launch(cfg)?;
17025        }
17026        Ok(())
17027    }
17028
17029    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
17030    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
17031    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
17032    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
17033    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
17034    /// different in_f) falls back to its own `matmul` — behavior unchanged.
17035    pub fn matmul_group(
17036        &self,
17037        ws: &[&crate::model::GpuTensor],
17038        x: &CudaSlice<f32>,
17039        m: usize,
17040    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17041        use crate::model::GpuTensor;
17042        let mut out = Vec::with_capacity(ws.len());
17043        let any_mirror = ws
17044            .iter()
17045            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17046        if m >= 16 && any_mirror && !self.verify_exact_on() {
17047            let in_f = ws[0].in_features();
17048            let xh = self.f16_act(x, m * in_f, in_f)?;
17049            for w in ws {
17050                if w.in_features() == in_f {
17051                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17052                        out.push(y);
17053                        continue;
17054                    }
17055                }
17056                out.push(self.matmul(w, x, m)?);
17057            }
17058            return Ok(out);
17059        }
17060        for w in ws {
17061            out.push(self.matmul(w, x, m)?);
17062        }
17063        Ok(out)
17064    }
17065
17066    /// Cross-request grouped matmul (task #13): run ONE projection group over the
17067    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
17068    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
17069    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
17070    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
17071    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
17072    pub fn matmul_group_multi(
17073        &self,
17074        ws: &[&crate::model::GpuTensor],
17075        xs: &[&CudaSlice<f32>],
17076        ms: &[usize],
17077    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17078        assert_eq!(xs.len(), ms.len());
17079        let in_f = ws[0].in_features();
17080        let total: usize = ms.iter().sum();
17081        let mut xcat = self.uninit(total * in_f)?;
17082        let mut off = 0usize;
17083        for (x, &m) in xs.iter().zip(ms) {
17084            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17085            off += m;
17086        }
17087        let ys = self.matmul_group(ws, &xcat, total)?;
17088        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17089        for (w, y) in ws.iter().zip(ys) {
17090            let out_f = w.out_features();
17091            let mut off = 0usize;
17092            for (s, &m) in ms.iter().enumerate() {
17093                let mut ys_s = self.uninit(m * out_f)?;
17094                let src = y.slice(off * out_f..(off + m) * out_f);
17095                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17096                out[s].push(ys_s);
17097                off += m;
17098            }
17099        }
17100        Ok(out)
17101    }
17102
17103    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
17104    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
17105    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
17106    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
17107    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
17108    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
17109    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
17110    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
17111    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
17112    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17113        use crate::model::GpuTensor;
17114        if !legacy_quant_gemm_allowed(
17115            cfg!(memra_portable_cuda),
17116            cfg!(memra_hopper_mma),
17117            std::env::var_os("MEMRA_NO_GEMM").is_some(),
17118        ) {
17119            return false;
17120        }
17121        match w {
17122            GpuTensor::Quant { qtype, .. } => {
17123                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17124                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17125            }
17126            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17127        }
17128    }
17129
17130    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
17131    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
17132    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
17133    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
17134    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
17135    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
17136    pub fn qmatvec_gemm(
17137        &self,
17138        w: &crate::model::GpuTensor,
17139        aq: &CudaSlice<i8>,
17140        ad: &CudaSlice<f32>,
17141        m: usize,
17142    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17143        use crate::model::GpuTensor;
17144        let in_f = w.in_features();
17145        let out_f = w.out_features();
17146        let (bytes, qtype, row_bytes, scale, rp) = match w {
17147            GpuTensor::Quant {
17148                bytes,
17149                qtype,
17150                row_bytes,
17151                scale,
17152                rp,
17153                ..
17154            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17155            _ => unreachable!("gemm_supports guaranteed Quant"),
17156        };
17157        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
17158        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
17159        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
17160        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
17161        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
17162        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17163            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17164                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17165                if scale != 1.0 {
17166                    self.scale_inplace(&mut y, scale, m * out_f)?;
17167                }
17168                return Ok(y);
17169            }
17170        }
17171        let name = match qtype {
17172            QT_Q8_0 => "qmatvec_gemm_q8_0",
17173            QT_Q4_K => "qmatvec_gemm_q4_K",
17174            QT_Q4_0 => {
17175                if rp {
17176                    "qmatvec_gemm_q4_0_rp"
17177                } else {
17178                    "qmatvec_gemm_q4_0"
17179                }
17180            }
17181            QT_Q5_K => "qmatvec_gemm_q5_K",
17182            QT_Q6_K => "qmatvec_gemm_q6_K",
17183            QT_NVFP4 => {
17184                if rp {
17185                    "qmatvec_gemm_nvfp4_rp"
17186                } else {
17187                    "qmatvec_gemm_nvfp4"
17188                }
17189            }
17190            _ => unreachable!(),
17191        };
17192        let f = self.func(name);
17193        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17194        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
17195        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
17196        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
17197        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
17198        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
17199        let k1_tile = if is_k1 {
17200            k1_launch_override().unwrap_or((128, 128, 8))
17201        } else {
17202            (128, 128, 8)
17203        };
17204        let (bm, bn): (u32, u32) = if is_k1 {
17205            (k1_tile.0, k1_tile.1)
17206        } else {
17207            (64, 256)
17208        };
17209        let warps: u32 = if is_k1 {
17210            k1_tile.2
17211        } else {
17212            match qtype {
17213                QT_NVFP4 => 8,
17214                _ => 4,
17215            }
17216        };
17217        let cfg = LaunchConfig {
17218            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17219            block_dim: (32, warps, 1),
17220            shared_mem_bytes: 0,
17221        };
17222        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17223        let __s_b = self.gpu.stream();
17224        let mut b = __s_b.launch_builder(&f);
17225        b.arg(bytes)
17226            .arg(aq)
17227            .arg(ad)
17228            .arg(&mut y)
17229            .arg(&inf)
17230            .arg(&outf)
17231            .arg(&mi)
17232            .arg(&rb);
17233        unsafe {
17234            b.launch(cfg)?;
17235        }
17236        if scale != 1.0 {
17237            self.scale_inplace(&mut y, scale, m * out_f)?;
17238        }
17239        Ok(y)
17240    }
17241
17242    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
17243    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
17244    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
17245    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
17246    pub fn qmatvec_gemm_raw(
17247        &self,
17248        bytes: &CudaSlice<u8>,
17249        x: &CudaSlice<f32>,
17250        m: usize,
17251        in_f: usize,
17252        out_f: usize,
17253        qtype: i32,
17254        row_bytes: usize,
17255    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17256        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17257        let name = match qtype {
17258            QT_Q8_0 => "qmatvec_gemm_q8_0",
17259            QT_Q4_K => "qmatvec_gemm_q4_K",
17260            QT_Q4_0 => "qmatvec_gemm_q4_0",
17261            QT_Q5_K => "qmatvec_gemm_q5_K",
17262            QT_Q6_K => "qmatvec_gemm_q6_K",
17263            QT_NVFP4 => "qmatvec_gemm_nvfp4",
17264            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17265            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17266        };
17267        let f = self.func(name);
17268        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17269        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
17270        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
17271        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
17272        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
17273        let k1_tile = if is_k1 {
17274            k1_launch_override().unwrap_or((128, 128, 8))
17275        } else {
17276            (128, 128, 8)
17277        };
17278        let (bm, bn): (u32, u32) = if is_k1 {
17279            (k1_tile.0, k1_tile.1)
17280        } else {
17281            (64, 256)
17282        };
17283        let warps: u32 = if is_k1 {
17284            k1_tile.2
17285        } else {
17286            match qtype {
17287                QT_NVFP4 | QT_NVFP4_RP => 8,
17288                _ => 4,
17289            }
17290        };
17291        let cfg = LaunchConfig {
17292            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17293            block_dim: (32, warps, 1),
17294            shared_mem_bytes: 0,
17295        };
17296        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17297        let __s_b = self.gpu.stream();
17298        let mut b = __s_b.launch_builder(&f);
17299        b.arg(bytes)
17300            .arg(&aq)
17301            .arg(&ad)
17302            .arg(&mut y)
17303            .arg(&inf)
17304            .arg(&outf)
17305            .arg(&mi)
17306            .arg(&rb);
17307        unsafe {
17308            b.launch(cfg)?;
17309        }
17310        Ok(y)
17311    }
17312
17313    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
17314    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
17315    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
17316    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
17317    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
17318    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
17319    pub fn qmatvec_gemm_q8_0_wgmma_raw(
17320        &self,
17321        rp4: &CudaSlice<u8>,
17322        aq: &CudaSlice<i8>,
17323        ad: &CudaSlice<f32>,
17324        m: usize,
17325        in_f: usize,
17326        out_f: usize,
17327    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17328        assert!(
17329            out_f % 64 == 0 && in_f % 32 == 0,
17330            "wgmma GEMM needs out_f%64==0, in_f%32==0"
17331        );
17332        let f = self.func("qmatvec_gemm_q8_0_wgmma");
17333        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
17334        let cfg = LaunchConfig {
17335            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17336            block_dim: (128, 1, 1),
17337            shared_mem_bytes: 0,
17338        };
17339        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17340        let __s_b = self.gpu.stream();
17341        let mut b = __s_b.launch_builder(&f);
17342        b.arg(rp4)
17343            .arg(aq)
17344            .arg(ad)
17345            .arg(&mut y)
17346            .arg(&inf)
17347            .arg(&outf)
17348            .arg(&mi);
17349        unsafe {
17350            b.launch(cfg)?;
17351        }
17352        Ok(y)
17353    }
17354
17355    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
17356    pub fn scale_inplace(
17357        &self,
17358        y: &mut CudaSlice<f32>,
17359        s: f32,
17360        n: usize,
17361    ) -> Result<(), Box<dyn std::error::Error>> {
17362        let f = self.func("scale_f32");
17363        let cfg = LaunchConfig::for_num_elems(n as u32);
17364        let (sf, ni) = (s, n as i32);
17365        let __s_b = self.gpu.stream();
17366        let mut b = __s_b.launch_builder(&f);
17367        b.arg(y).arg(&sf).arg(&ni);
17368        unsafe {
17369            b.launch(cfg)?;
17370        }
17371        Ok(())
17372    }
17373
17374    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
17375    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
17376    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
17377    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
17378    pub fn bf16_to_f32(
17379        &self,
17380        data: &cudarc::driver::CudaView<'_, u8>,
17381        n: usize,
17382    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17383        let mut out = self.alloc_uninit::<f32>(n)?;
17384        let f = self.func("bf16_to_f32");
17385        let cfg = LaunchConfig::for_num_elems(n as u32);
17386        let ni = n as i32;
17387        let __s_b = self.gpu.stream();
17388        let mut b = __s_b.launch_builder(&f);
17389        b.arg(data).arg(&mut out).arg(&ni);
17390        unsafe {
17391            b.launch(cfg)?;
17392        }
17393        Ok(out)
17394    }
17395
17396    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
17397    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
17398    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
17399    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
17400    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
17401    /// calls, the spec-verify contract) vs plain linear.
17402    fn linear_bf16_chunked(
17403        &self,
17404        x: &CudaSlice<f32>,
17405        data: &CudaSlice<u8>,
17406        m: usize,
17407        in_f: usize,
17408        out_f: usize,
17409        exact: bool,
17410        canonical_chunk_rows: Option<usize>,
17411    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17412        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
17413        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
17414        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
17415        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17416        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17417        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17418        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17419        let started = timing.then(std::time::Instant::now);
17420        let result =
17421            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17422        if let Some(started) = started {
17423            use std::sync::atomic::Ordering;
17424            self.stream().synchronize()?;
17425            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17426                + started.elapsed().as_nanos() as u64;
17427            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17428                + (in_f * out_f * 2) as u64;
17429            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17430            if calls % 1024 == 0 {
17431                eprintln!(
17432                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17433                     weight_gb={:.2}",
17434                    ns as f64 / 1.0e6,
17435                    ns as f64 / calls as f64 / 1.0e3,
17436                    wb as f64 / 1.0e9,
17437                );
17438            }
17439        }
17440        result
17441    }
17442
17443    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
17444    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
17445    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
17446    /// numeric-class doors (DEV_ROUTES precedent).
17447    pub(crate) fn bf16_mmv_on() -> bool {
17448        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17449        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17450    }
17451
17452    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
17453    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
17454    fn matvec_bf16(
17455        &self,
17456        data: &CudaSlice<u8>,
17457        x: &CudaSlice<f32>,
17458        in_f: usize,
17459        out_f: usize,
17460    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17461        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17462            return Err(format!(
17463                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17464                data.len(),
17465                x.len()
17466            )
17467            .into());
17468        }
17469        let mut y = self.alloc_uninit::<f32>(out_f)?;
17470        let f = self.func("matvec_bf16_f32acc");
17471        let cfg = LaunchConfig {
17472            grid_dim: (out_f as u32, 1, 1),
17473            block_dim: (mmv_block(), 1, 1),
17474            shared_mem_bytes: 0,
17475        };
17476        let ini = in_f as i32;
17477        let __s_bld = self.gpu.stream();
17478        let mut bld = __s_bld.launch_builder(&f);
17479        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17480        unsafe {
17481            bld.launch(cfg)?;
17482        }
17483        Ok(y)
17484    }
17485
17486    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
17487    /// launches, a position upload, and the rope launch; the position is read directly from
17488    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
17489    #[allow(clippy::too_many_arguments)]
17490    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
17491    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
17492    /// Bit-identical to the split kernels; requires head_dim == 128 and
17493    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
17494    #[allow(clippy::too_many_arguments)]
17495    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
17496    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
17497    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
17498    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
17499    #[allow(clippy::too_many_arguments)]
17500    pub fn qk_norm_rope_append_inc_dcw_rows(
17501        &self,
17502        q_raw_t: &CudaSlice<f32>,
17503        k_raw_t: &CudaSlice<f32>,
17504        v_raw_t: &CudaSlice<f32>,
17505        qw: &CudaSlice<f32>,
17506        kw: &CudaSlice<f32>,
17507        q_out_t: &mut CudaSlice<f32>,
17508        k_out_t: &mut CudaSlice<f32>,
17509        tab: &CudaSlice<u64>,
17510        pos_t: &CudaSlice<i32>,
17511        same_session: bool,
17512        t: usize,
17513        kv_dim_k: usize,
17514        kv_dim_v: usize,
17515        k_tok_bytes: usize,
17516        v_tok_bytes: usize,
17517        head_dim: usize,
17518        n_dims: usize,
17519        nh_q: usize,
17520        nh_k: usize,
17521        eps: f32,
17522        freq_base: f32,
17523        freq_scale: f32,
17524        ff: Option<&CudaSlice<f32>>,
17525    ) -> Result<(), Box<dyn std::error::Error>> {
17526        if head_dim != 128
17527            || kv_dim_v != kv_dim_k
17528            || kv_dim_k != nh_k * head_dim
17529            || t == 0
17530            || t > 32
17531            || tab.len() < t * 6
17532            || pos_t.len() < t
17533            || q_raw_t.len() < t * nh_q * head_dim
17534            || k_raw_t.len() < t * nh_k * head_dim
17535            || v_raw_t.len() < t * kv_dim_v
17536            || q_out_t.len() < t * nh_q * head_dim
17537            || k_out_t.len() < t * nh_k * head_dim
17538        {
17539            return Err(format!(
17540                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17541                 nh_q={nh_q} nh_k={nh_k}"
17542            )
17543            .into());
17544        }
17545        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17546        let same_t: i32 = if same_session { t as i32 } else { 0 };
17547        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17548        let cfg = LaunchConfig {
17549            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17550            block_dim: (128, 1, 1),
17551            shared_mem_bytes: 0,
17552        };
17553        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17554        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17555        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17556        let null: u64 = 0;
17557        let __s_b = self.gpu.stream();
17558        let mut b = __s_b.launch_builder(&f);
17559        b.arg(q_raw_t)
17560            .arg(k_raw_t)
17561            .arg(v_raw_t)
17562            .arg(qw)
17563            .arg(kw)
17564            .arg(q_out_t)
17565            .arg(k_out_t)
17566            .arg(tab)
17567            .arg(pos_t)
17568            .arg(&same_t)
17569            .arg(&kvk)
17570            .arg(&kvv)
17571            .arg(&ktb)
17572            .arg(&vtb)
17573            .arg(&hd)
17574            .arg(&nd)
17575            .arg(&nq)
17576            .arg(&nk)
17577            .arg(&eps)
17578            .arg(&theta_scale)
17579            .arg(&freq_scale);
17580        match ff {
17581            Some(freqs) => {
17582                b.arg(freqs);
17583            }
17584            None => {
17585                b.arg(&null);
17586            }
17587        }
17588        unsafe {
17589            b.launch(cfg)?;
17590        }
17591        Ok(())
17592    }
17593
17594    pub fn qk_norm_rope_append_inc_dcw(
17595        &self,
17596        q_raw: &CudaSlice<f32>,
17597        k_raw: &CudaSlice<f32>,
17598        v_raw: &CudaSlice<f32>,
17599        qw: &CudaSlice<f32>,
17600        kw: &CudaSlice<f32>,
17601        q_out: &mut CudaSlice<f32>,
17602        k_out: &mut CudaSlice<f32>,
17603        pos: &CudaSlice<i32>,
17604        k_plane: &mut CudaSlice<u8>,
17605        v_plane: &mut CudaSlice<u8>,
17606        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
17607        // (single) writer, exactly like the split append+inc pair it replaces.
17608        len_dev: &CudaSlice<i32>,
17609        base_dev: Option<&CudaSlice<i32>>,
17610        done_ctr: &mut CudaSlice<u32>,
17611        kv_dim_k: usize,
17612        kv_dim_v: usize,
17613        k_tok_bytes: usize,
17614        v_tok_bytes: usize,
17615        head_dim: usize,
17616        n_dims: usize,
17617        nh_q: usize,
17618        nh_k: usize,
17619        eps: f32,
17620        freq_base: f32,
17621        freq_scale: f32,
17622        ff: Option<&CudaSlice<f32>>,
17623    ) -> Result<(), Box<dyn std::error::Error>> {
17624        if head_dim != 128
17625            || kv_dim_v != kv_dim_k
17626            || kv_dim_k != nh_k * head_dim
17627            || q_raw.len() < nh_q * head_dim
17628            || k_raw.len() < nh_k * head_dim
17629            || v_raw.len() < kv_dim_v
17630            || q_out.len() < nh_q * head_dim
17631            || k_out.len() < nh_k * head_dim
17632            || pos.is_empty()
17633            || done_ctr.is_empty()
17634        {
17635            return Err(format!(
17636                "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}"
17637            )
17638            .into());
17639        }
17640        let f = self.func("qk_norm_rope_append_inc_dcw");
17641        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17642        let cfg = LaunchConfig {
17643            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17644            block_dim: (128, 1, 1),
17645            shared_mem_bytes: 0,
17646        };
17647        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17648        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17649        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17650        let null: u64 = 0;
17651        let __s_b = self.gpu.stream();
17652        let mut b = __s_b.launch_builder(&f);
17653        b.arg(q_raw)
17654            .arg(k_raw)
17655            .arg(v_raw)
17656            .arg(qw)
17657            .arg(kw)
17658            .arg(q_out)
17659            .arg(k_out)
17660            .arg(pos)
17661            .arg(&mut *k_plane)
17662            .arg(&mut *v_plane)
17663            .arg(len_dev);
17664        match base_dev {
17665            Some(base) => {
17666                b.arg(base);
17667            }
17668            None => {
17669                b.arg(&null);
17670            }
17671        }
17672        b.arg(&mut *done_ctr)
17673            .arg(&kvk)
17674            .arg(&kvv)
17675            .arg(&ktb)
17676            .arg(&vtb)
17677            .arg(&hd)
17678            .arg(&nd)
17679            .arg(&nq)
17680            .arg(&eps)
17681            .arg(&theta_scale)
17682            .arg(&freq_scale);
17683        match ff {
17684            Some(freqs) => {
17685                b.arg(freqs);
17686            }
17687            None => {
17688                b.arg(&null);
17689            }
17690        }
17691        unsafe {
17692            b.launch(cfg)?;
17693        }
17694        Ok(())
17695    }
17696
17697    pub fn qk_norm_rope_into(
17698        &self,
17699        q_raw: &CudaSlice<f32>,
17700        k_raw: &CudaSlice<f32>,
17701        qw: &CudaSlice<f32>,
17702        kw: &CudaSlice<f32>,
17703        q_out: &mut CudaSlice<f32>,
17704        k_out: &mut CudaSlice<f32>,
17705        pos: &CudaSlice<i32>,
17706        head_dim: usize,
17707        n_dims: usize,
17708        nh_q: usize,
17709        nh_k: usize,
17710        eps: f32,
17711        freq_base: f32,
17712        freq_scale: f32,
17713        ff: Option<&CudaSlice<f32>>,
17714    ) -> Result<(), Box<dyn std::error::Error>> {
17715        if head_dim > 512
17716            || q_raw.len() < nh_q * head_dim
17717            || k_raw.len() < nh_k * head_dim
17718            || q_out.len() < nh_q * head_dim
17719            || k_out.len() < nh_k * head_dim
17720            || qw.len() < head_dim
17721            || kw.len() < head_dim
17722            || pos.is_empty()
17723        {
17724            return Err(format!(
17725                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17726            )
17727            .into());
17728        }
17729        let f = self.func("qk_norm_rope_f32");
17730        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17731        let cfg = LaunchConfig {
17732            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17733            block_dim: (128, 1, 1),
17734            shared_mem_bytes: 0,
17735        };
17736        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17737        let __s_b = self.gpu.stream();
17738        let mut b = __s_b.launch_builder(&f);
17739        b.arg(q_raw)
17740            .arg(k_raw)
17741            .arg(qw)
17742            .arg(kw)
17743            .arg(q_out)
17744            .arg(k_out)
17745            .arg(pos)
17746            .arg(&hd)
17747            .arg(&nd)
17748            .arg(&nq)
17749            .arg(&eps)
17750            .arg(&theta_scale)
17751            .arg(&freq_scale);
17752        match ff {
17753            Some(ffv) => {
17754                b.arg(ffv);
17755                unsafe {
17756                    b.launch(cfg)?;
17757                }
17758            }
17759            None => {
17760                let null: u64 = 0;
17761                b.arg(&null);
17762                unsafe {
17763                    b.launch(cfg)?;
17764                }
17765            }
17766        }
17767        Ok(())
17768    }
17769
17770    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
17771    /// launch computes a rank's whole O partial from its four canonical column blocks.
17772    #[allow(clippy::too_many_arguments)]
17773    pub fn matvec_f32_b4_into(
17774        &self,
17775        w: [&CudaSlice<f32>; 4],
17776        x: &CudaSlice<f32>,
17777        y: &mut CudaSlice<f32>,
17778        block_cols: usize,
17779        out_f: usize,
17780    ) -> Result<(), Box<dyn std::error::Error>> {
17781        if block_cols % 4 != 0
17782            || x.len() < 4 * block_cols
17783            || y.len() < out_f
17784            || w.iter().any(|w| w.len() != out_f * block_cols)
17785        {
17786            return Err(format!(
17787                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17788                x.len()
17789            )
17790            .into());
17791        }
17792        let f = self.func("matvec_f32_b4");
17793        let cfg = LaunchConfig {
17794            grid_dim: (out_f as u32, 1, 1),
17795            block_dim: (128, 1, 1),
17796            shared_mem_bytes: 0,
17797        };
17798        let (bc, of) = (block_cols as i32, out_f as i32);
17799        let __s_b = self.gpu.stream();
17800        let mut b = __s_b.launch_builder(&f);
17801        b.arg(w[0])
17802            .arg(w[1])
17803            .arg(w[2])
17804            .arg(w[3])
17805            .arg(x)
17806            .arg(y)
17807            .arg(&bc)
17808            .arg(&of);
17809        unsafe {
17810            b.launch(cfg)?;
17811        }
17812        Ok(())
17813    }
17814
17815    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
17816    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
17817    pub fn axpy_rows_seq_into(
17818        &self,
17819        x: &CudaSlice<f32>,
17820        w: &CudaSlice<f32>,
17821        y: &mut CudaSlice<f32>,
17822        width: usize,
17823        n_rows: usize,
17824    ) -> Result<(), Box<dyn std::error::Error>> {
17825        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17826            return Err(format!(
17827                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17828                x.len(),
17829                w.len(),
17830                y.len()
17831            )
17832            .into());
17833        }
17834        let f = self.func("axpy_rows_seq_f32");
17835        let cfg = LaunchConfig::for_num_elems(width as u32);
17836        let (wi, nr) = (width as i32, n_rows as i32);
17837        let __s_b = self.gpu.stream();
17838        let mut b = __s_b.launch_builder(&f);
17839        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17840        unsafe {
17841            b.launch(cfg)?;
17842        }
17843        Ok(())
17844    }
17845
17846    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
17847    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
17848    /// exact sequential FP chain of the base kernel over that window.
17849    #[allow(clippy::too_many_arguments)]
17850    pub fn axpy_rows_seq_md_off_into(
17851        &self,
17852        x: &CudaSlice<f32>,
17853        w_route: &CudaSlice<f32>,
17854        md: &CudaSlice<f32>,
17855        sel: &CudaSlice<i32>,
17856        y: &mut CudaSlice<f32>,
17857        width: usize,
17858        n_rows: usize,
17859        row0: usize,
17860    ) -> Result<(), Box<dyn std::error::Error>> {
17861        if x.len() < (row0 + n_rows) * width
17862            || w_route.len() < row0 + n_rows
17863            || sel.len() < row0 + n_rows
17864            || y.len() < width
17865        {
17866            return Err(format!(
17867                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17868                 rows={n_rows} row0={row0}",
17869                x.len(),
17870                w_route.len(),
17871                sel.len(),
17872                y.len()
17873            )
17874            .into());
17875        }
17876        let f = self.func("axpy_rows_seq_md_off_f32");
17877        let cfg = LaunchConfig::for_num_elems(width as u32);
17878        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17879        let __s_b = self.gpu.stream();
17880        let mut b = __s_b.launch_builder(&f);
17881        b.arg(x)
17882            .arg(w_route)
17883            .arg(md)
17884            .arg(sel)
17885            .arg(y)
17886            .arg(&wi)
17887            .arg(&nr)
17888            .arg(&r0);
17889        unsafe {
17890            b.launch(cfg)?;
17891        }
17892        Ok(())
17893    }
17894
17895    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
17896    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
17897    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
17898    /// outputs are bit-equal to its own t=1 launch.
17899    #[allow(clippy::too_many_arguments)]
17900    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17901        &self,
17902        gate_bank: &CudaSlice<u8>,
17903        up_bank: &CudaSlice<u8>,
17904        sel: &CudaSlice<i32>,
17905        aq: &CudaSlice<i8>,
17906        ad: &CudaSlice<f32>,
17907        yg: &mut CudaSlice<f32>,
17908        yu: &mut CudaSlice<f32>,
17909        n_sel: usize,
17910        n_sel_col: usize,
17911        in_f: usize,
17912        out_f: usize,
17913        row_bytes: usize,
17914        expert_stride: usize,
17915        act_row_stride: usize,
17916        ad_row_stride: usize,
17917    ) -> Result<(), Box<dyn std::error::Error>> {
17918        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17919        if yg.len() < n_sel * out_f
17920            || yu.len() < n_sel * out_f
17921            || sel.len() < n_sel
17922            || n_sel_col == 0
17923            || n_sel % n_sel_col != 0
17924        {
17925            return Err("NVFP4 gu tcol geometry".into());
17926        }
17927        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17928        let cfg = LaunchConfig {
17929            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17930            block_dim: (128, 1, 1),
17931            shared_mem_bytes: 0,
17932        };
17933        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17934        let (rb, es) = (row_bytes as i64, expert_stride as i64);
17935        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17936        let __s_b = self.gpu.stream();
17937        let mut b = __s_b.launch_builder(&f);
17938        b.arg(gate_bank)
17939            .arg(up_bank)
17940            .arg(sel)
17941            .arg(aq)
17942            .arg(ad)
17943            .arg(yg)
17944            .arg(yu)
17945            .arg(&inf)
17946            .arg(&outf)
17947            .arg(&ns)
17948            .arg(&rb)
17949            .arg(&es)
17950            .arg(&ars)
17951            .arg(&adrs)
17952            .arg(&nsc);
17953        unsafe {
17954            b.launch(cfg)?;
17955        }
17956        Ok(())
17957    }
17958
17959    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
17960    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
17961    #[allow(clippy::too_many_arguments)]
17962    pub fn axpy_rows_seq_md_into(
17963        &self,
17964        x: &CudaSlice<f32>,
17965        w_route: &CudaSlice<f32>,
17966        md: &CudaSlice<f32>,
17967        sel: &CudaSlice<i32>,
17968        y: &mut CudaSlice<f32>,
17969        width: usize,
17970        n_rows: usize,
17971    ) -> Result<(), Box<dyn std::error::Error>> {
17972        if x.len() < n_rows * width
17973            || w_route.len() < n_rows
17974            || sel.len() < n_rows
17975            || y.len() < width
17976        {
17977            return Err(format!(
17978                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17979                x.len(),
17980                w_route.len(),
17981                sel.len(),
17982                y.len()
17983            )
17984            .into());
17985        }
17986        let f = self.func("axpy_rows_seq_md_f32");
17987        let cfg = LaunchConfig::for_num_elems(width as u32);
17988        let (wi, nr) = (width as i32, n_rows as i32);
17989        let __s_b = self.gpu.stream();
17990        let mut b = __s_b.launch_builder(&f);
17991        b.arg(x)
17992            .arg(w_route)
17993            .arg(md)
17994            .arg(sel)
17995            .arg(y)
17996            .arg(&wi)
17997            .arg(&nr);
17998        unsafe {
17999            b.launch(cfg)?;
18000        }
18001        Ok(())
18002    }
18003
18004    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
18005    #[allow(clippy::too_many_arguments)]
18006    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
18007    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
18008    /// land column-major-of-rows: yq[c*out_q + row] etc.
18009    #[allow(clippy::too_many_arguments)]
18010    pub fn matvec_bf16_qkvg_tcol_into(
18011        &self,
18012        wq: &CudaSlice<u8>,
18013        wk: &CudaSlice<u8>,
18014        wv: &CudaSlice<u8>,
18015        wg: &CudaSlice<u8>,
18016        x_t: &CudaSlice<f32>,
18017        yq: &mut CudaSlice<f32>,
18018        yk: &mut CudaSlice<f32>,
18019        yv: &mut CudaSlice<f32>,
18020        yg: &mut CudaSlice<f32>,
18021        in_f: usize,
18022        out_q: usize,
18023        out_kv: usize,
18024        out_g: usize,
18025        t: usize,
18026    ) -> Result<(), Box<dyn std::error::Error>> {
18027        if t == 0
18028            || t > 8
18029            || in_f % 8 != 0
18030            || x_t.len() < t * in_f
18031            || yq.len() < t * out_q
18032            || yk.len() < t * out_kv
18033            || yv.len() < t * out_kv
18034            || (out_g > 0 && yg.len() < t * out_g)
18035        {
18036            return Err("matvec_bf16_qkvg_tcol geometry".into());
18037        }
18038        let grid = out_q + 2 * out_kv + out_g;
18039        let cfg = LaunchConfig {
18040            grid_dim: (grid as u32, 1, 1),
18041            block_dim: (mmv_block(), 1, 1),
18042            shared_mem_bytes: 0,
18043        };
18044        let (ini, oq, okv, og, ti) = (
18045            in_f as i32,
18046            out_q as i32,
18047            out_kv as i32,
18048            out_g as i32,
18049            t as i32,
18050        );
18051        let __s_b = self.gpu.stream();
18052        // One runtime-T program for every live width. The compile-time 2/4/8 twins are
18053        // retained in the fatbin as research controls, but dispatching them by the current
18054        // batch width changes kernels inside a request when peers arrive or retire. That is
18055        // a load-history numeric-program switch, and their pre-twin TOKFP receipts did not
18056        // qualify it (Hermes `64fa2b55baf0d887`).
18057        let f = self.func("matvec_bf16_qkvg_tcol");
18058        let mut b = __s_b.launch_builder(&f);
18059        b.arg(wq)
18060            .arg(wk)
18061            .arg(wv)
18062            .arg(wg)
18063            .arg(x_t)
18064            .arg(yq)
18065            .arg(yk)
18066            .arg(yv)
18067            .arg(yg)
18068            .arg(&ini)
18069            .arg(&oq)
18070            .arg(&okv)
18071            .arg(&og)
18072            .arg(&ti);
18073        unsafe {
18074            b.launch(cfg)?;
18075        }
18076        Ok(())
18077    }
18078
18079    pub fn matvec_bf16_qkvg_into(
18080        &self,
18081        wq: &CudaSlice<u8>,
18082        wk: &CudaSlice<u8>,
18083        wv: &CudaSlice<u8>,
18084        wg: &CudaSlice<u8>,
18085        x: &CudaSlice<f32>,
18086        yq: &mut CudaSlice<f32>,
18087        yk: &mut CudaSlice<f32>,
18088        yv: &mut CudaSlice<f32>,
18089        yg: &mut CudaSlice<f32>,
18090        in_f: usize,
18091        out_q: usize,
18092        out_kv: usize,
18093        out_g: usize,
18094    ) -> Result<(), Box<dyn std::error::Error>> {
18095        if in_f % 8 != 0
18096            || wq.len() != out_q * in_f * 2
18097            || wk.len() != out_kv * in_f * 2
18098            || wv.len() != out_kv * in_f * 2
18099            || wg.len() < out_g * in_f * 2
18100            || x.len() < in_f
18101            || yq.len() < out_q
18102            || yk.len() < out_kv
18103            || yv.len() < out_kv
18104            || (out_g > 0 && yg.len() < out_g)
18105        {
18106            return Err(format!(
18107                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18108            )
18109            .into());
18110        }
18111        let f = self.func("matvec_bf16_qkvg");
18112        let cfg = LaunchConfig {
18113            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18114            block_dim: (mmv_block(), 1, 1),
18115            shared_mem_bytes: 0,
18116        };
18117        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18118        let __s_b = self.gpu.stream();
18119        let mut b = __s_b.launch_builder(&f);
18120        b.arg(wq)
18121            .arg(wk)
18122            .arg(wv)
18123            .arg(wg)
18124            .arg(x)
18125            .arg(yq)
18126            .arg(yk)
18127            .arg(yv)
18128            .arg(yg)
18129            .arg(&inf)
18130            .arg(&oq)
18131            .arg(&okv)
18132            .arg(&og);
18133        unsafe {
18134            b.launch(cfg)?;
18135        }
18136        Ok(())
18137    }
18138
18139    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
18140    pub fn matvec_bf16_b4_into(
18141        &self,
18142        w: [&CudaSlice<u8>; 4],
18143        x: &CudaSlice<f32>,
18144        y: &mut CudaSlice<f32>,
18145        block_cols: usize,
18146        out_f: usize,
18147    ) -> Result<(), Box<dyn std::error::Error>> {
18148        if block_cols % 8 != 0
18149            || x.len() < 4 * block_cols
18150            || y.len() < out_f
18151            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18152        {
18153            return Err(format!(
18154                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18155                x.len()
18156            )
18157            .into());
18158        }
18159        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
18160        // bit-identical per row (the second row's stream hides the first's reduce tail).
18161        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18162        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18163        let f = self.func(if x2 {
18164            "matvec_bf16_b4_x2"
18165        } else {
18166            "matvec_bf16_b4"
18167        });
18168        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18169        let cfg = LaunchConfig {
18170            grid_dim: (grid as u32, 1, 1),
18171            block_dim: (mmv_block(), 1, 1),
18172            shared_mem_bytes: 0,
18173        };
18174        let (bc, of) = (block_cols as i32, out_f as i32);
18175        let __s_b = self.gpu.stream();
18176        let mut b = __s_b.launch_builder(&f);
18177        b.arg(w[0])
18178            .arg(w[1])
18179            .arg(w[2])
18180            .arg(w[3])
18181            .arg(x)
18182            .arg(y)
18183            .arg(&bc)
18184            .arg(&of);
18185        unsafe {
18186            b.launch(cfg)?;
18187        }
18188        Ok(())
18189    }
18190
18191    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
18192    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
18193    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
18194    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
18195    /// t=1 program).
18196    pub fn matvec_bf16_b4_tcol_into(
18197        &self,
18198        w: [&CudaSlice<u8>; 4],
18199        x_t: &CudaSlice<f32>,
18200        y_t: &mut CudaSlice<f32>,
18201        block_cols: usize,
18202        out_f: usize,
18203        t: usize,
18204    ) -> Result<(), Box<dyn std::error::Error>> {
18205        if block_cols % 8 != 0
18206            || t == 0
18207            || t > 8
18208            || x_t.len() < t * 4 * block_cols
18209            || y_t.len() < t * out_f
18210            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18211        {
18212            return Err(format!(
18213                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18214                x_t.len()
18215            )
18216            .into());
18217        }
18218        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18219            return Err(
18220                "b4 tcol verify is qualified against the plain b4 kernel only \
18221                        (MEMRA_B4_X2=1 is a different t=1 program)"
18222                    .into(),
18223            );
18224        }
18225        // Keep one runtime-T program at every live width. Compile-time twins remain research
18226        // controls only; selecting them from the changing batch width switches programs
18227        // mid-request.
18228        let cfg = LaunchConfig {
18229            grid_dim: (out_f as u32, 1, 1),
18230            block_dim: (mmv_block(), 1, 1),
18231            shared_mem_bytes: 0,
18232        };
18233        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18234        let __s_b = self.gpu.stream();
18235        let f = self.func("matvec_bf16_b4_tcol");
18236        let mut b = __s_b.launch_builder(&f);
18237        b.arg(w[0])
18238            .arg(w[1])
18239            .arg(w[2])
18240            .arg(w[3])
18241            .arg(x_t)
18242            .arg(y_t)
18243            .arg(&bc)
18244            .arg(&of)
18245            .arg(&ti);
18246        unsafe {
18247            b.launch(cfg)?;
18248        }
18249        Ok(())
18250    }
18251
18252    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
18253    /// q8_0 row bytes for an `in_f`-wide weight row: one 34-byte block per 32 weights.
18254    pub fn q8_0_row_bytes(in_f: usize) -> usize {
18255        in_f / 32 * 34
18256    }
18257
18258    /// Encode a resident bf16 weight slab into its q8_0 mirror (MEMRA_STEP_TP_W8). Runs once
18259    /// per matrix at load; the block program is the one `quant_K_block` writes for the KV
18260    /// cache, so the two formats cannot drift apart.
18261    pub fn encode_q8_0_from_bf16(
18262        &self,
18263        w_bf16: &CudaSlice<u8>,
18264        out: &mut CudaSlice<u8>,
18265        in_f: usize,
18266        out_f: usize,
18267    ) -> Result<(), Box<dyn std::error::Error>> {
18268        if in_f % 32 != 0
18269            || w_bf16.len() < in_f * out_f * 2
18270            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18271        {
18272            return Err(format!(
18273                "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18274                w_bf16.len(),
18275                out.len()
18276            )
18277            .into());
18278        }
18279        let f = self.func("encode_q8_0_rows_from_bf16");
18280        // Flat 1D grid of (row, 32-block) pairs, 4 pairs per block: rows on grid.y would cap
18281        // at 65535 and the LM head has 128896 rows.
18282        const PAIRS_PER_BLOCK: u32 = 4;
18283        let pairs = (out_f * (in_f / 32)) as u64;
18284        let cfg = LaunchConfig {
18285            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18286            block_dim: (32, PAIRS_PER_BLOCK, 1),
18287            shared_mem_bytes: 0,
18288        };
18289        let (ini, outi) = (in_f as i32, out_f as i32);
18290        let __s_b = self.gpu.stream();
18291        let mut b = __s_b.launch_builder(&f);
18292        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18293        unsafe {
18294            b.launch(cfg)?;
18295        }
18296        Ok(())
18297    }
18298
18299    /// Fused q8_0 QKV against a q8_1 activation (MEMRA_STEP_TP_W8): one launch over the
18300    /// stacked q/k/v rows, each row running the exact `qmatvec_q8_0_mmvq_rp` per-row program.
18301    /// Bit-identical to three per-matrix mmvq calls; it exists because those three launches
18302    /// plus the activation quantize measured SLOWER than the bf16 fused kernel.
18303    #[allow(clippy::too_many_arguments)]
18304    pub fn qmatvec_q8_0_qkv_rp_into(
18305        &self,
18306        wq: &CudaSlice<u8>,
18307        wk: &CudaSlice<u8>,
18308        wv: &CudaSlice<u8>,
18309        aq: &CudaSlice<i8>,
18310        ad: &CudaSlice<f32>,
18311        yq: &mut CudaSlice<f32>,
18312        yk: &mut CudaSlice<f32>,
18313        yv: &mut CudaSlice<f32>,
18314        in_f: usize,
18315        out_q: usize,
18316        out_kv: usize,
18317    ) -> Result<(), Box<dyn std::error::Error>> {
18318        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18319        let rows = out_q + 2 * out_kv;
18320        let nblk = in_f / 32;
18321        if in_f % 32 != 0
18322            || aq.len() < in_f
18323            || ad.len() < nblk
18324            || yq.len() < out_q
18325            || yk.len() < out_kv
18326            || yv.len() < out_kv
18327            || wq.len() < out_q * nblk * 34
18328            || wk.len() < out_kv * nblk * 34
18329            || wv.len() < out_kv * nblk * 34
18330        {
18331            return Err(
18332                format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18333            );
18334        }
18335        let f = self.func("qmatvec_q8_0_qkv_rp");
18336        let cfg = LaunchConfig {
18337            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18338            block_dim: (32, ROWS_PER_BLOCK, 1),
18339            shared_mem_bytes: 0,
18340        };
18341        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18342        let __s_b = self.gpu.stream();
18343        let mut b = __s_b.launch_builder(&f);
18344        b.arg(wq)
18345            .arg(wk)
18346            .arg(wv)
18347            .arg(aq)
18348            .arg(ad)
18349            .arg(yq)
18350            .arg(yk)
18351            .arg(yv)
18352            .arg(&ini)
18353            .arg(&oq)
18354            .arg(&okv);
18355        unsafe {
18356            b.launch(cfg)?;
18357        }
18358        Ok(())
18359    }
18360
18361    /// Fused q8_0 O projection over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8): one
18362    /// launch, one warp per output row, per-block reduce then add — the same shape
18363    /// `matvec_bf16_b4` uses, against a q8_1 activation.
18364    #[allow(clippy::too_many_arguments)]
18365    pub fn qmatvec_q8_0_b4_rp_into(
18366        &self,
18367        w: [&CudaSlice<u8>; 4],
18368        aq: &CudaSlice<i8>,
18369        ad: &CudaSlice<f32>,
18370        y: &mut CudaSlice<f32>,
18371        block_cols: usize,
18372        out_f: usize,
18373    ) -> Result<(), Box<dyn std::error::Error>> {
18374        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18375        let nblk = block_cols / 32;
18376        if block_cols % 32 != 0
18377            || aq.len() < 4 * block_cols
18378            || ad.len() < 4 * nblk
18379            || y.len() < out_f
18380            || w.iter().any(|p| p.len() < out_f * nblk * 34)
18381        {
18382            return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18383        }
18384        let f = self.func("qmatvec_q8_0_b4_rp");
18385        let cfg = LaunchConfig {
18386            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18387            block_dim: (32, ROWS_PER_BLOCK, 1),
18388            shared_mem_bytes: 0,
18389        };
18390        let (bc, of) = (block_cols as i32, out_f as i32);
18391        let __s_b = self.gpu.stream();
18392        let mut b = __s_b.launch_builder(&f);
18393        b.arg(w[0])
18394            .arg(w[1])
18395            .arg(w[2])
18396            .arg(w[3])
18397            .arg(aq)
18398            .arg(ad)
18399            .arg(y)
18400            .arg(&bc)
18401            .arg(&of);
18402        unsafe {
18403            b.launch(cfg)?;
18404        }
18405        Ok(())
18406    }
18407
18408    /// Get-or-build this bf16 weight's q8_0 mirror and run the GEMV through it. Returns
18409    /// `None` when the shape has no mirror form, so the caller falls back to bf16.
18410    fn matvec_bf16_via_q8_mirror(
18411        &self,
18412        data: &CudaSlice<u8>,
18413        x: &CudaSlice<f32>,
18414        y: &mut CudaSlice<f32>,
18415        in_f: usize,
18416        out_f: usize,
18417    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18418        use cudarc::driver::DevicePtr;
18419        let key = {
18420            let s = self.gpu.stream();
18421            let (p, _g) = data.device_ptr(&s);
18422            p as u64
18423        };
18424        {
18425            let mut mirrors = self
18426                .w8_mirrors
18427                .lock()
18428                .map_err(|_| "w8 mirror map is poisoned")?;
18429            if !mirrors.contains_key(&key) {
18430                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18431                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18432                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18433                mirrors.insert(key, planar);
18434                // Which weights this half actually covers is not obvious from the call graph:
18435                // the head and the shared expert may reach the GPU through the rows fast path
18436                // or the fused dual-silu launcher instead of here. One line per mirror answers
18437                // that without a profiler (the hybrid half measured +0.1% and this is how we
18438                // find out whether it even fired).
18439                if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18440                    eprintln!(
18441                        "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18442                        mirrors.len()
18443                    );
18444                }
18445            }
18446        }
18447        let nblk = in_f / 32;
18448        {
18449            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18450            if !act.contains_key(&in_f) {
18451                let aq = self.alloc_uninit::<i8>(in_f)?;
18452                let ad = self.alloc_uninit::<f32>(nblk)?;
18453                act.insert(in_f, (aq, ad));
18454            }
18455            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18456            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18457        }
18458        let mirrors = self
18459            .w8_mirrors
18460            .lock()
18461            .map_err(|_| "w8 mirror map is poisoned")?;
18462        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18463        let mirror = mirrors.get(&key).expect("built above");
18464        let (aq, ad) = act.get(&in_f).expect("built above");
18465        self.qmatvec_mmvq_into(
18466            mirror,
18467            aq,
18468            ad,
18469            1,
18470            in_f,
18471            out_f,
18472            QT_Q8_0,
18473            Self::q8_0_row_bytes(in_f),
18474            1.0,
18475            true,
18476            y,
18477        )?;
18478        Ok(Some(()))
18479    }
18480
18481    pub fn matvec_bf16_into(
18482        &self,
18483        data: &CudaSlice<u8>,
18484        x: &CudaSlice<f32>,
18485        y: &mut CudaSlice<f32>,
18486        in_f: usize,
18487        out_f: usize,
18488    ) -> Result<(), Box<dyn std::error::Error>> {
18489        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18490            return Err(format!(
18491                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18492                data.len(),
18493                x.len(),
18494                y.len()
18495            )
18496            .into());
18497        }
18498        // MEMRA_STEP_TP_W8, hybrid half: route this GEMV through a q8_0 mirror of the same
18499        // weight. Covers exactly the bf16 GEMVs that are NOT in a TP resident bank — the LM
18500        // head (324.4 -> 163.7 us measured), the shared-expert down rows (13.0 -> 5.6 us) and
18501        // the dense-FFN layers. Same numeric class as the QKV/o_proj arms (int8 dp4a with
18502        // per-32 scales), so it rides the same argmax acceptance; the bf16 slab stays resident
18503        // for prefill. The mirror builds on first use and is keyed by the slab's pointer.
18504        if step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18505            if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
18506                return Ok(());
18507            }
18508        }
18509        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
18510        // block, exact f32acc per-row program — cures the 1-iteration latency
18511        // starvation (shexp down measured 420GB/s at in_f=1280).
18512        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18513        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18514            && in_f <= 2048;
18515        if x4 {
18516            let f = self.func("matvec_bf16_f32acc_x4");
18517            let cfg = LaunchConfig {
18518                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18519                block_dim: (mmv_block(), 1, 1),
18520                shared_mem_bytes: 0,
18521            };
18522            let (ini, outi) = (in_f as i32, out_f as i32);
18523            let __s_b = self.gpu.stream();
18524            let mut b = __s_b.launch_builder(&f);
18525            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18526            unsafe {
18527                b.launch(cfg)?;
18528            }
18529            return Ok(());
18530        }
18531        let f = self.func("matvec_bf16_f32acc");
18532        let cfg = LaunchConfig {
18533            grid_dim: (out_f as u32, 1, 1),
18534            block_dim: (mmv_block(), 1, 1),
18535            shared_mem_bytes: 0,
18536        };
18537        let ini = in_f as i32;
18538        let __s_b = self.gpu.stream();
18539        let mut b = __s_b.launch_builder(&f);
18540        b.arg(data).arg(x).arg(y).arg(&ini);
18541        unsafe {
18542            b.launch(cfg)?;
18543        }
18544        Ok(())
18545    }
18546
18547    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
18548    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
18549    pub fn matvec_bf16_view_into(
18550        &self,
18551        data: &cudarc::driver::CudaView<'_, u8>,
18552        x: &CudaSlice<f32>,
18553        y: &mut CudaSlice<f32>,
18554        in_f: usize,
18555        out_f: usize,
18556    ) -> Result<(), Box<dyn std::error::Error>> {
18557        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18558            return Err(format!(
18559                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18560                data.len(),
18561                x.len(),
18562                y.len()
18563            )
18564            .into());
18565        }
18566        let f = self.func("matvec_bf16_f32acc");
18567        let cfg = LaunchConfig {
18568            grid_dim: (out_f as u32, 1, 1),
18569            block_dim: (mmv_block(), 1, 1),
18570            shared_mem_bytes: 0,
18571        };
18572        let ini = in_f as i32;
18573        let __s_b = self.gpu.stream();
18574        let mut b = __s_b.launch_builder(&f);
18575        b.arg(data).arg(x).arg(y).arg(&ini);
18576        unsafe {
18577            b.launch(cfg)?;
18578        }
18579        Ok(())
18580    }
18581
18582    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
18583    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
18584    pub fn matvec_bf16_raw_out(
18585        &self,
18586        w: &CudaSlice<u8>,
18587        x: &CudaSlice<f32>,
18588        y_raw: u64,
18589        in_f: usize,
18590        out_f: usize,
18591    ) -> Result<(), Box<dyn std::error::Error>> {
18592        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18593            return Err("matvec_bf16_raw_out geometry".into());
18594        }
18595        let f = self.func("matvec_bf16_f32acc");
18596        let cfg = LaunchConfig {
18597            grid_dim: (out_f as u32, 1, 1),
18598            block_dim: (mmv_block(), 1, 1),
18599            shared_mem_bytes: 0,
18600        };
18601        let ini = in_f as i32;
18602        let __s_b = self.gpu.stream();
18603        let mut b = __s_b.launch_builder(&f);
18604        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18605        unsafe {
18606            b.launch(cfg)?;
18607        }
18608        Ok(())
18609    }
18610
18611    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
18612    /// UVA pointers so the caller passes persistent-static rows without holding locks).
18613    /// Exact per-element sequence of the split add + add_scaled_rows pair.
18614    pub fn add3_raw(
18615        &self,
18616        a: &CudaSlice<f32>,
18617        b: &CudaSlice<f32>,
18618        sh_raw: u64,
18619        scale_raw: u64,
18620        dst: &mut CudaSlice<f32>,
18621        n: usize,
18622    ) -> Result<(), Box<dyn std::error::Error>> {
18623        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18624            return Err("add3_raw geometry".into());
18625        }
18626        let f = self.func("add3_f32");
18627        let cfg = LaunchConfig {
18628            grid_dim: ((n as u32).div_ceil(256), 1, 1),
18629            block_dim: (256, 1, 1),
18630            shared_mem_bytes: 0,
18631        };
18632        let ni = n as i32;
18633        let __s_b = self.gpu.stream();
18634        let mut bld = __s_b.launch_builder(&f);
18635        bld.arg(a)
18636            .arg(b)
18637            .arg(&sh_raw)
18638            .arg(&scale_raw)
18639            .arg(dst)
18640            .arg(&ni);
18641        unsafe {
18642            bld.launch(cfg)?;
18643        }
18644        Ok(())
18645    }
18646
18647    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
18648    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
18649    pub fn matvec_bf16_down_addscale_into(
18650        &self,
18651        w: &CudaSlice<u8>,
18652        x: &CudaSlice<f32>,
18653        scale: &CudaSlice<f32>,
18654        dst: &mut CudaSlice<f32>,
18655        in_f: usize,
18656        out_f: usize,
18657    ) -> Result<(), Box<dyn std::error::Error>> {
18658        if w.len() != in_f * out_f * 2
18659            || x.len() < in_f
18660            || in_f % 8 != 0
18661            || dst.len() < out_f
18662            || scale.is_empty()
18663        {
18664            return Err("matvec_bf16_down_addscale geometry".into());
18665        }
18666        let f = self.func("matvec_bf16_down_addscale");
18667        let cfg = LaunchConfig {
18668            grid_dim: (out_f as u32, 1, 1),
18669            block_dim: (mmv_block(), 1, 1),
18670            shared_mem_bytes: 0,
18671        };
18672        let ini = in_f as i32;
18673        let __s_b = self.gpu.stream();
18674        let mut b = __s_b.launch_builder(&f);
18675        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18676        unsafe {
18677            b.launch(cfg)?;
18678        }
18679        Ok(())
18680    }
18681
18682    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
18683    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
18684    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
18685    #[allow(clippy::too_many_arguments)]
18686    pub fn matvec_bf16_dual_silu_rows_into(
18687        &self,
18688        wg: &CudaSlice<u8>,
18689        wu: &CudaSlice<u8>,
18690        x: &CudaSlice<f32>,
18691        act: &mut CudaSlice<f32>,
18692        in_f: usize,
18693        out_f: usize,
18694        limit: Option<f32>,
18695        t: usize,
18696    ) -> Result<(), Box<dyn std::error::Error>> {
18697        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18698            return Err("matvec_bf16_dual_silu_rows geometry".into());
18699        }
18700        let f = self.func("matvec_bf16_dual_silu_rows");
18701        let cfg = LaunchConfig {
18702            grid_dim: (out_f as u32, t as u32, 1),
18703            block_dim: (mmv_block(), 1, 1),
18704            shared_mem_bytes: 0,
18705        };
18706        let (ini, outi) = (in_f as i32, out_f as i32);
18707        let lim = limit.unwrap_or(0.0);
18708        let __s_b = self.gpu.stream();
18709        let mut b = __s_b.launch_builder(&f);
18710        b.arg(wg)
18711            .arg(wu)
18712            .arg(x)
18713            .arg(&mut *act)
18714            .arg(&ini)
18715            .arg(&outi)
18716            .arg(&lim);
18717        unsafe {
18718            b.launch(cfg)?;
18719        }
18720        Ok(())
18721    }
18722
18723    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
18724    pub fn matvec_bf16_rows_into(
18725        &self,
18726        w: &CudaSlice<u8>,
18727        x: &CudaSlice<f32>,
18728        y: &mut CudaSlice<f32>,
18729        in_f: usize,
18730        out_f: usize,
18731        t: usize,
18732    ) -> Result<(), Box<dyn std::error::Error>> {
18733        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18734            return Err("matvec_bf16_rows geometry".into());
18735        }
18736        // MEMRA_STEP_TP_W8: the LM head reaches the device HERE, not through
18737        // matvec_bf16_into — the W8 trace showed the hybrid half building mirrors only for
18738        // in_f=1280 out_f=4096 (the shared-expert down rows, which SHEXP_OVERLAP already
18739        // hides, hence its +0.1%). Route the t=1 decode row through the q8 mirror; wider t
18740        // (the verify walk) keeps bf16 so the prefill class is untouched.
18741        if t == 1 && step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18742            if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
18743                return Ok(());
18744            }
18745        }
18746        let f = self.func("matvec_bf16_f32acc_x4_rows");
18747        let cfg = LaunchConfig {
18748            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18749            block_dim: (mmv_block(), 1, 1),
18750            shared_mem_bytes: 0,
18751        };
18752        let (ini, outi) = (in_f as i32, out_f as i32);
18753        let __s_b = self.gpu.stream();
18754        let mut b = __s_b.launch_builder(&f);
18755        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18756        unsafe {
18757            b.launch(cfg)?;
18758        }
18759        Ok(())
18760    }
18761
18762    pub fn matvec_bf16_dual_silu_into(
18763        &self,
18764        wg: &CudaSlice<u8>,
18765        wu: &CudaSlice<u8>,
18766        x: &CudaSlice<f32>,
18767        act: &mut CudaSlice<f32>,
18768        in_f: usize,
18769        out_f: usize,
18770        limit: Option<f32>,
18771    ) -> Result<(), Box<dyn std::error::Error>> {
18772        if wg.len() != in_f * out_f * 2
18773            || wu.len() != in_f * out_f * 2
18774            || x.len() < in_f
18775            || in_f % 8 != 0
18776            || act.len() < out_f
18777        {
18778            return Err("matvec_bf16_dual_silu geometry".into());
18779        }
18780        let f = self.func("matvec_bf16_dual_silu");
18781        let cfg = LaunchConfig {
18782            grid_dim: (out_f as u32, 1, 1),
18783            block_dim: (mmv_block(), 1, 1),
18784            shared_mem_bytes: 0,
18785        };
18786        let (ini, outi) = (in_f as i32, out_f as i32);
18787        let lim = limit.unwrap_or(0.0);
18788        let __s_b = self.gpu.stream();
18789        let mut b = __s_b.launch_builder(&f);
18790        b.arg(wg)
18791            .arg(wu)
18792            .arg(x)
18793            .arg(act)
18794            .arg(&ini)
18795            .arg(&outi)
18796            .arg(&lim);
18797        unsafe {
18798            b.launch(cfg)?;
18799        }
18800        Ok(())
18801    }
18802
18803    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
18804    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
18805    #[allow(clippy::too_many_arguments)]
18806    pub fn matvec_bf16_dual_view_into(
18807        &self,
18808        wg: &cudarc::driver::CudaView<'_, u8>,
18809        wu: &cudarc::driver::CudaView<'_, u8>,
18810        x: &CudaSlice<f32>,
18811        yg: &mut CudaSlice<f32>,
18812        yu: &mut CudaSlice<f32>,
18813        in_f: usize,
18814        out_f: usize,
18815    ) -> Result<(), Box<dyn std::error::Error>> {
18816        if wg.len() != in_f * out_f * 2
18817            || wu.len() != in_f * out_f * 2
18818            || x.len() < in_f
18819            || in_f % 8 != 0
18820            || yg.len() < out_f
18821            || yu.len() < out_f
18822        {
18823            return Err(format!(
18824                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18825                wg.len(),
18826                wu.len(),
18827                x.len()
18828            )
18829            .into());
18830        }
18831        let f = self.func("matvec_bf16_dual");
18832        let cfg = LaunchConfig {
18833            grid_dim: ((2 * out_f) as u32, 1, 1),
18834            block_dim: (mmv_block(), 1, 1),
18835            shared_mem_bytes: 0,
18836        };
18837        let (ini, outi) = (in_f as i32, out_f as i32);
18838        let __s_b = self.gpu.stream();
18839        let mut b = __s_b.launch_builder(&f);
18840        b.arg(wg)
18841            .arg(wu)
18842            .arg(x)
18843            .arg(yg)
18844            .arg(yu)
18845            .arg(&ini)
18846            .arg(&outi);
18847        unsafe {
18848            b.launch(cfg)?;
18849        }
18850        Ok(())
18851    }
18852
18853    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
18854    #[allow(clippy::too_many_arguments)]
18855    pub fn matvec_bf16_dual_into(
18856        &self,
18857        wg: &CudaSlice<u8>,
18858        wu: &CudaSlice<u8>,
18859        x: &CudaSlice<f32>,
18860        yg: &mut CudaSlice<f32>,
18861        yu: &mut CudaSlice<f32>,
18862        in_f: usize,
18863        out_f: usize,
18864    ) -> Result<(), Box<dyn std::error::Error>> {
18865        if wg.len() != in_f * out_f * 2
18866            || wu.len() != in_f * out_f * 2
18867            || x.len() < in_f
18868            || in_f % 8 != 0
18869            || yg.len() < out_f
18870            || yu.len() < out_f
18871        {
18872            return Err(format!(
18873                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18874                wg.len(),
18875                wu.len(),
18876                x.len()
18877            )
18878            .into());
18879        }
18880        let f = self.func("matvec_bf16_dual");
18881        let cfg = LaunchConfig {
18882            grid_dim: ((2 * out_f) as u32, 1, 1),
18883            block_dim: (mmv_block(), 1, 1),
18884            shared_mem_bytes: 0,
18885        };
18886        let (ini, outi) = (in_f as i32, out_f as i32);
18887        let __s_b = self.gpu.stream();
18888        let mut b = __s_b.launch_builder(&f);
18889        b.arg(wg)
18890            .arg(wu)
18891            .arg(x)
18892            .arg(yg)
18893            .arg(yu)
18894            .arg(&ini)
18895            .arg(&outi);
18896        unsafe {
18897            b.launch(cfg)?;
18898        }
18899        Ok(())
18900    }
18901
18902    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
18903    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
18904    pub(crate) fn matvec_bf16_dual(
18905        &self,
18906        wg: &CudaSlice<u8>,
18907        wu: &CudaSlice<u8>,
18908        x: &CudaSlice<f32>,
18909        in_f: usize,
18910        out_f: usize,
18911    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18912        if wg.len() != in_f * out_f * 2
18913            || wu.len() != in_f * out_f * 2
18914            || x.len() < in_f
18915            || in_f % 8 != 0
18916        {
18917            return Err(format!(
18918                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18919                wg.len(),
18920                wu.len(),
18921                x.len()
18922            )
18923            .into());
18924        }
18925        let mut yg = self.alloc_uninit::<f32>(out_f)?;
18926        let mut yu = self.alloc_uninit::<f32>(out_f)?;
18927        let f = self.func("matvec_bf16_dual");
18928        let cfg = LaunchConfig {
18929            grid_dim: ((2 * out_f) as u32, 1, 1),
18930            block_dim: (mmv_block(), 1, 1),
18931            shared_mem_bytes: 0,
18932        };
18933        let (ini, outi) = (in_f as i32, out_f as i32);
18934        let __s_b = self.gpu.stream();
18935        let mut b = __s_b.launch_builder(&f);
18936        b.arg(wg)
18937            .arg(wu)
18938            .arg(x)
18939            .arg(&mut yg)
18940            .arg(&mut yu)
18941            .arg(&ini)
18942            .arg(&outi);
18943        unsafe {
18944            b.launch(cfg)?;
18945        }
18946        Ok((yg, yu))
18947    }
18948
18949    #[allow(clippy::too_many_arguments)]
18950    fn linear_bf16_chunked_inner(
18951        &self,
18952        x: &CudaSlice<f32>,
18953        data: &CudaSlice<u8>,
18954        m: usize,
18955        in_f: usize,
18956        out_f: usize,
18957        exact: bool,
18958        canonical_chunk_rows: Option<usize>,
18959    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18960        const CHUNK_BYTES: usize = 256 << 20;
18961        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
18962        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
18963        if m == 1
18964            && !exact
18965            && canonical_chunk_rows.is_none()
18966            && in_f % 8 == 0
18967            && Self::bf16_mmv_on()
18968        {
18969            return self.matvec_bf16(data, x, in_f, out_f);
18970        }
18971        let row_bytes = in_f
18972            .checked_mul(std::mem::size_of::<f32>())
18973            .ok_or("BF16 chunk row byte count overflow")?;
18974        if row_bytes == 0 || out_f == 0 {
18975            return Err("BF16 chunk dimensions must be nonzero".into());
18976        }
18977        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18978        let chunk_rows = match canonical_chunk_rows {
18979            Some(rows) if rows == 0 => {
18980                return Err("canonical BF16 chunk rows must be nonzero".into());
18981            }
18982            Some(rows) if rows > max_chunk_rows => {
18983                return Err(format!(
18984                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18985                )
18986                .into());
18987            }
18988            Some(rows) if out_f % rows != 0 => {
18989                return Err(format!(
18990                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18991                )
18992                .into());
18993            }
18994            Some(rows) => rows,
18995            None => max_chunk_rows,
18996        };
18997        if chunk_rows >= out_f {
18998            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18999            return if exact {
19000                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19001            } else {
19002                self.linear(x, &wf32, m, in_f, out_f)
19003            };
19004        }
19005        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19006        let mut r0 = 0usize;
19007        while r0 < out_f {
19008            let rows = chunk_rows.min(out_f - r0);
19009            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19010            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19011            let yc = if exact {
19012                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19013            } else {
19014                self.linear(x, &wf32, m, in_f, rows)?
19015            };
19016            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
19017            for mi in 0..m {
19018                let src = yc.slice(mi * rows..(mi + 1) * rows);
19019                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19020                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19021            }
19022            r0 += rows;
19023        }
19024        Ok(y)
19025    }
19026
19027    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
19028    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
19029    /// chunked BF16 numerical program instead of re-encoding the weight.
19030    pub fn linear_bf16_resident(
19031        &self,
19032        x: &CudaSlice<f32>,
19033        data: &CudaSlice<u8>,
19034        m: usize,
19035        in_f: usize,
19036        out_f: usize,
19037    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19038        if data.len() != in_f * out_f * 2 {
19039            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19040        }
19041        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19042    }
19043
19044    /// Execute a resident BF16 projection as fixed-width output-row chunks.
19045    ///
19046    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
19047    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
19048    /// model topology rather than the active rank count.
19049    pub fn linear_bf16_resident_canonical_rows(
19050        &self,
19051        x: &CudaSlice<f32>,
19052        data: &CudaSlice<u8>,
19053        m: usize,
19054        in_f: usize,
19055        out_f: usize,
19056        canonical_chunk_rows: usize,
19057    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19058        if data.len() != in_f * out_f * 2 {
19059            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19060        }
19061        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19062    }
19063
19064    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
19065    ///
19066    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
19067    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
19068    pub fn linear_f32_resident_canonical_rows(
19069        &self,
19070        x: &CudaSlice<f32>,
19071        data: &CudaSlice<f32>,
19072        m: usize,
19073        in_f: usize,
19074        out_f: usize,
19075        canonical_chunk_rows: usize,
19076    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19077        self.linear_f32_resident_canonical_rows_inner(
19078            x,
19079            data,
19080            m,
19081            in_f,
19082            out_f,
19083            canonical_chunk_rows,
19084            false,
19085        )
19086    }
19087
19088    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
19089    ///
19090    /// The projection shapes and values are identical to
19091    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
19092    /// changes, replacing one device copy per token with one placement kernel per output chunk.
19093    pub fn linear_f32_resident_canonical_rows_strided(
19094        &self,
19095        x: &CudaSlice<f32>,
19096        data: &CudaSlice<f32>,
19097        m: usize,
19098        in_f: usize,
19099        out_f: usize,
19100        canonical_chunk_rows: usize,
19101    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19102        self.linear_f32_resident_canonical_rows_inner(
19103            x,
19104            data,
19105            m,
19106            in_f,
19107            out_f,
19108            canonical_chunk_rows,
19109            true,
19110        )
19111    }
19112
19113    fn linear_f32_resident_canonical_rows_inner(
19114        &self,
19115        x: &CudaSlice<f32>,
19116        data: &CudaSlice<f32>,
19117        m: usize,
19118        in_f: usize,
19119        out_f: usize,
19120        canonical_chunk_rows: usize,
19121        strided_output: bool,
19122    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19123        if data.len() != in_f * out_f {
19124            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19125        }
19126        if canonical_chunk_rows == 0
19127            || canonical_chunk_rows > out_f
19128            || out_f % canonical_chunk_rows != 0
19129        {
19130            return Err(format!(
19131                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19132            )
19133            .into());
19134        }
19135        if canonical_chunk_rows == out_f {
19136            return self.linear(x, data, m, in_f, out_f);
19137        }
19138
19139        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19140        let input = x.slice(0..x.len());
19141        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19142            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19143            if m == 1 {
19144                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19145                self.linear_device_into(
19146                    &input,
19147                    &weights,
19148                    &mut destination,
19149                    1,
19150                    in_f,
19151                    canonical_chunk_rows,
19152                )?;
19153                continue;
19154            }
19155            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19156            if strided_output {
19157                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19158            } else {
19159                for token in 0..m {
19160                    let source = chunk
19161                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19162                    let mut destination =
19163                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19164                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19165                }
19166            }
19167        }
19168        Ok(y)
19169    }
19170
19171    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
19172    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
19173    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
19174    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
19175    pub fn linear_f32_resident_canonical_rows_t1_into(
19176        &self,
19177        x: &CudaSlice<f32>,
19178        data: &CudaSlice<f32>,
19179        y: &mut CudaSlice<f32>,
19180        in_f: usize,
19181        out_f: usize,
19182        canonical_chunk_rows: usize,
19183    ) -> Result<(), Box<dyn std::error::Error>> {
19184        if data.len() != in_f * out_f {
19185            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19186        }
19187        if y.len() != out_f || x.len() != in_f {
19188            return Err(format!(
19189                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19190                x.len(),
19191                y.len()
19192            )
19193            .into());
19194        }
19195        if canonical_chunk_rows == 0
19196            || canonical_chunk_rows > out_f
19197            || out_f % canonical_chunk_rows != 0
19198        {
19199            return Err(format!(
19200                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19201            )
19202            .into());
19203        }
19204        let input = x.slice(0..x.len());
19205        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19206            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19207            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19208            self.linear_device_into(
19209                &input,
19210                &weights,
19211                &mut destination,
19212                1,
19213                in_f,
19214                canonical_chunk_rows,
19215            )?;
19216        }
19217        Ok(())
19218    }
19219
19220    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
19221    /// without the allocation, for workspace-resident operands.
19222    pub fn linear_t1_into(
19223        &self,
19224        x: &cudarc::driver::CudaView<'_, f32>,
19225        w: &cudarc::driver::CudaView<'_, f32>,
19226        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19227        in_f: usize,
19228        out_f: usize,
19229    ) -> Result<(), Box<dyn std::error::Error>> {
19230        self.linear_device_into(x, w, y, 1, in_f, out_f)
19231    }
19232
19233    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
19234    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
19235    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
19236    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
19237    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
19238    /// router/shexp sites and matmul_decode_exact's Float arm.
19239    pub fn linear_decode_exact(
19240        &self,
19241        x: &CudaSlice<f32>,
19242        w: &CudaSlice<f32>,
19243        m_tokens: usize,
19244        in_f: usize,
19245        out_f: usize,
19246    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19247        if m_tokens == 1 {
19248            return self.linear(x, w, 1, in_f, out_f);
19249        }
19250        let xv = self.view(x, m_tokens * in_f);
19251        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19252        for t in 0..m_tokens {
19253            let row = xv.slice(t * in_f..(t + 1) * in_f);
19254            let mut xr = self.alloc_uninit::<f32>(in_f)?;
19255            self.copy_view_into(&mut xr, 0, &row, in_f)?;
19256            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19257            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19258        }
19259        Ok(y)
19260    }
19261
19262    pub fn linear(
19263        &self,
19264        x: &CudaSlice<f32>,
19265        w: &CudaSlice<f32>,
19266        m_tokens: usize,
19267        in_f: usize,
19268        out_f: usize,
19269    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19270        self.linear_device(x, w, m_tokens, in_f, out_f)
19271    }
19272
19273    fn linear_device<I>(
19274        &self,
19275        x: &I,
19276        w: &I,
19277        m_tokens: usize,
19278        in_f: usize,
19279        out_f: usize,
19280    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19281    where
19282        I: cudarc::driver::DevicePtr<f32>,
19283    {
19284        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
19285        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
19286        Ok(c)
19287    }
19288
19289    fn linear_device_into<I, O>(
19290        &self,
19291        x: &I,
19292        w: &I,
19293        c: &mut O,
19294        m_tokens: usize,
19295        in_f: usize,
19296        out_f: usize,
19297    ) -> Result<(), Box<dyn std::error::Error>>
19298    where
19299        I: cudarc::driver::DevicePtr<f32>,
19300        O: cudarc::driver::DevicePtrMut<f32>,
19301    {
19302        use cudarc::cublaslt::{Matmul, MatmulConfig};
19303        let cfg = MatmulConfig {
19304            transa: true,
19305            transb: false,
19306            transc: false,
19307            m: out_f as u64,
19308            n: m_tokens as u64,
19309            k: in_f as u64,
19310            alpha: 1.0,
19311            lda: in_f as i64,
19312            ldb: in_f as i64,
19313            beta: 0.0,
19314            ldc: out_f as i64,
19315            stride_a: None,
19316            stride_b: None,
19317            stride_c: None,
19318            stride_bias: None,
19319            batch_size: None,
19320        };
19321        let blas = self.gpu.blas();
19322        unsafe {
19323            blas.matmul(cfg, w, x, c, None, None)?;
19324        }
19325        Ok(())
19326    }
19327
19328    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
19329    ///
19330    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
19331    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
19332    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
19333    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
19334    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
19335    /// launch error mid-request.
19336    pub fn sdpa_naive(
19337        &self,
19338        q: &CudaSlice<f32>,
19339        k: &CudaSlice<f32>,
19340        v: &CudaSlice<f32>,
19341        o: &mut CudaSlice<f32>,
19342        head_dim: usize,
19343        n_head: usize,
19344        n_head_kv: usize,
19345        t: usize,
19346        t_kv: usize,
19347        scale: f32,
19348        causal: bool,
19349    ) -> Result<(), Box<dyn std::error::Error>> {
19350        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19351            return self.sdpa_naive_gmem(
19352                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19353            );
19354        }
19355        let f = self.func("sdpa_naive_f32");
19356        let cfg = LaunchConfig {
19357            grid_dim: (n_head as u32, t as u32, 1),
19358            block_dim: (128, 1, 1),
19359            shared_mem_bytes: (t_kv * 4) as u32,
19360        };
19361        let (hd, nh, nhkv, ti, tkvi, cz) = (
19362            head_dim as i32,
19363            n_head as i32,
19364            n_head_kv as i32,
19365            t as i32,
19366            t_kv as i32,
19367            causal as i32,
19368        );
19369        let __s_b = self.gpu.stream();
19370        let mut b = __s_b.launch_builder(&f);
19371        b.arg(q)
19372            .arg(k)
19373            .arg(v)
19374            .arg(o)
19375            .arg(&hd)
19376            .arg(&nh)
19377            .arg(&nhkv)
19378            .arg(&ti)
19379            .arg(&tkvi)
19380            .arg(&scale)
19381            .arg(&cz);
19382        unsafe {
19383            b.launch(cfg)?;
19384        }
19385        Ok(())
19386    }
19387
19388    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
19389    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
19390    /// of dynamic shared memory: identical loop structure and reduction order, so the output
19391    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
19392    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
19393    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
19394    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
19395    /// T==T_kv caller cannot silently allocate tens of GB.
19396    #[allow(clippy::too_many_arguments)]
19397    pub fn sdpa_naive_gmem(
19398        &self,
19399        q: &CudaSlice<f32>,
19400        k: &CudaSlice<f32>,
19401        v: &CudaSlice<f32>,
19402        o: &mut CudaSlice<f32>,
19403        head_dim: usize,
19404        n_head: usize,
19405        n_head_kv: usize,
19406        t: usize,
19407        t_kv: usize,
19408        scale: f32,
19409        causal: bool,
19410    ) -> Result<(), Box<dyn std::error::Error>> {
19411        let ws_len = n_head
19412            .checked_mul(t)
19413            .and_then(|x| x.checked_mul(t_kv))
19414            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19415        let ws_bytes = ws_len
19416            .checked_mul(std::mem::size_of::<f32>())
19417            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19418        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19419            return Err(format!(
19420                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19421                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19422                 needs a tiled/flash kernel, not the naive oracle"
19423            )
19424            .into());
19425        }
19426        let mut scores = self.uninit(ws_len)?;
19427        let f = self.func("sdpa_naive_gmem_f32");
19428        let cfg = LaunchConfig {
19429            grid_dim: (n_head as u32, t as u32, 1),
19430            block_dim: (128, 1, 1),
19431            shared_mem_bytes: 0,
19432        };
19433        let (hd, nh, nhkv, ti, tkvi, cz) = (
19434            head_dim as i32,
19435            n_head as i32,
19436            n_head_kv as i32,
19437            t as i32,
19438            t_kv as i32,
19439            causal as i32,
19440        );
19441        let __s_b = self.gpu.stream();
19442        let mut b = __s_b.launch_builder(&f);
19443        b.arg(q)
19444            .arg(k)
19445            .arg(v)
19446            .arg(o)
19447            .arg(&mut scores)
19448            .arg(&hd)
19449            .arg(&nh)
19450            .arg(&nhkv)
19451            .arg(&ti)
19452            .arg(&tkvi)
19453            .arg(&scale)
19454            .arg(&cz);
19455        unsafe {
19456            b.launch(cfg)?;
19457        }
19458        Ok(())
19459    }
19460
19461    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
19462    /// bidirectional image islands. `span_id` labels each absolute kv position
19463    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
19464    /// reproducing the reference's non-causal image batch. window 0 = no window.
19465    #[allow(clippy::too_many_arguments)]
19466    pub fn sdpa_naive_island(
19467        &self,
19468        q: &CudaSlice<f32>,
19469        k: &CudaSlice<f32>,
19470        v: &CudaSlice<f32>,
19471        o: &mut CudaSlice<f32>,
19472        span_id: &CudaSlice<i32>,
19473        head_dim: usize,
19474        n_head: usize,
19475        n_head_kv: usize,
19476        t: usize,
19477        t_kv: usize,
19478        scale: f32,
19479        window: usize,
19480    ) -> Result<(), Box<dyn std::error::Error>> {
19481        let f = self.func("sdpa_naive_island_f32");
19482        let cfg = LaunchConfig {
19483            grid_dim: (n_head as u32, t as u32, 1),
19484            block_dim: (128, 1, 1),
19485            shared_mem_bytes: (t_kv * 4) as u32,
19486        };
19487        let (hd, nh, nhkv, ti, tkvi, wi) = (
19488            head_dim as i32,
19489            n_head as i32,
19490            n_head_kv as i32,
19491            t as i32,
19492            t_kv as i32,
19493            window as i32,
19494        );
19495        let __s_b = self.gpu.stream();
19496        let mut b = __s_b.launch_builder(&f);
19497        b.arg(q)
19498            .arg(k)
19499            .arg(v)
19500            .arg(o)
19501            .arg(span_id)
19502            .arg(&hd)
19503            .arg(&nh)
19504            .arg(&nhkv)
19505            .arg(&ti)
19506            .arg(&tkvi)
19507            .arg(&scale)
19508            .arg(&wi);
19509        unsafe {
19510            b.launch(cfg)?;
19511        }
19512        Ok(())
19513    }
19514
19515    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
19516    #[allow(clippy::too_many_arguments)]
19517    pub fn sdpa_naive_w(
19518        &self,
19519        q: &CudaSlice<f32>,
19520        k: &CudaSlice<f32>,
19521        v: &CudaSlice<f32>,
19522        o: &mut CudaSlice<f32>,
19523        head_dim: usize,
19524        n_head: usize,
19525        n_head_kv: usize,
19526        t: usize,
19527        t_kv: usize,
19528        scale: f32,
19529        causal: bool,
19530        window: usize,
19531    ) -> Result<(), Box<dyn std::error::Error>> {
19532        let f = self.func("sdpa_naive_w_f32");
19533        let cfg = LaunchConfig {
19534            grid_dim: (n_head as u32, t as u32, 1),
19535            block_dim: (128, 1, 1),
19536            shared_mem_bytes: (t_kv * 4) as u32,
19537        };
19538        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19539            head_dim as i32,
19540            n_head as i32,
19541            n_head_kv as i32,
19542            t as i32,
19543            t_kv as i32,
19544            causal as i32,
19545            window as i32,
19546        );
19547        let __s_b = self.gpu.stream();
19548        let mut b = __s_b.launch_builder(&f);
19549        b.arg(q)
19550            .arg(k)
19551            .arg(v)
19552            .arg(o)
19553            .arg(&hd)
19554            .arg(&nh)
19555            .arg(&nhkv)
19556            .arg(&ti)
19557            .arg(&tkvi)
19558            .arg(&scale)
19559            .arg(&cz)
19560            .arg(&wi);
19561        unsafe {
19562            b.launch(cfg)?;
19563        }
19564        Ok(())
19565    }
19566
19567    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
19568    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
19569    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
19570    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
19571    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
19572    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
19573    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
19574    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
19575    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
19576    #[allow(clippy::too_many_arguments)]
19577    pub fn sdpa_naive_w_lo(
19578        &self,
19579        q: &CudaSlice<f32>,
19580        k: &CudaSlice<f32>,
19581        v: &CudaSlice<f32>,
19582        o: &mut CudaSlice<f32>,
19583        head_dim: usize,
19584        n_head: usize,
19585        n_head_kv: usize,
19586        t: usize,
19587        t_kv: usize,
19588        scale: f32,
19589        causal: bool,
19590        window: usize,
19591    ) -> Result<(), Box<dyn std::error::Error>> {
19592        let kv_lo = if window > 0 {
19593            (t_kv - t + 1).saturating_sub(window)
19594        } else {
19595            0
19596        };
19597        let smem = (t_kv - kv_lo) * 4;
19598        if smem > 48 * 1024 {
19599            return Err(format!(
19600                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19601                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19602                 a window this wide needs the multi-pass long-ctx kernel"
19603            )
19604            .into());
19605        }
19606        let f = self.func("sdpa_naive_w_lo_f32");
19607        let cfg = LaunchConfig {
19608            grid_dim: (n_head as u32, t as u32, 1),
19609            block_dim: (128, 1, 1),
19610            shared_mem_bytes: smem as u32,
19611        };
19612        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19613            head_dim as i32,
19614            n_head as i32,
19615            n_head_kv as i32,
19616            t as i32,
19617            t_kv as i32,
19618            causal as i32,
19619            window as i32,
19620            kv_lo as i32,
19621        );
19622        let __s_b = self.gpu.stream();
19623        let mut b = __s_b.launch_builder(&f);
19624        b.arg(q)
19625            .arg(k)
19626            .arg(v)
19627            .arg(o)
19628            .arg(&hd)
19629            .arg(&nh)
19630            .arg(&nhkv)
19631            .arg(&ti)
19632            .arg(&tkvi)
19633            .arg(&scale)
19634            .arg(&cz)
19635            .arg(&wi)
19636            .arg(&lo);
19637        unsafe {
19638            b.launch(cfg)?;
19639        }
19640        Ok(())
19641    }
19642
19643    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
19644    pub fn sdpa_naive_view(
19645        &self,
19646        q: &CudaSlice<f32>,
19647        k: &cudarc::driver::CudaView<f32>,
19648        v: &cudarc::driver::CudaView<f32>,
19649        o: &mut CudaSlice<f32>,
19650        head_dim: usize,
19651        n_head: usize,
19652        n_head_kv: usize,
19653        t: usize,
19654        t_kv: usize,
19655        scale: f32,
19656        causal: bool,
19657    ) -> Result<(), Box<dyn std::error::Error>> {
19658        let f = self.func("sdpa_naive_f32");
19659        let cfg = LaunchConfig {
19660            grid_dim: (n_head as u32, t as u32, 1),
19661            block_dim: (128, 1, 1),
19662            shared_mem_bytes: (t_kv * 4) as u32,
19663        };
19664        let (hd, nh, nhkv, ti, tkvi, cz) = (
19665            head_dim as i32,
19666            n_head as i32,
19667            n_head_kv as i32,
19668            t as i32,
19669            t_kv as i32,
19670            causal as i32,
19671        );
19672        let __s_b = self.gpu.stream();
19673        let mut b = __s_b.launch_builder(&f);
19674        b.arg(q)
19675            .arg(k)
19676            .arg(v)
19677            .arg(o)
19678            .arg(&hd)
19679            .arg(&nh)
19680            .arg(&nhkv)
19681            .arg(&ti)
19682            .arg(&tkvi)
19683            .arg(&scale)
19684            .arg(&cz);
19685        unsafe {
19686            b.launch(cfg)?;
19687        }
19688        Ok(())
19689    }
19690
19691    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
19692    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
19693    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
19694    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
19695    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
19696    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
19697    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
19698    #[allow(clippy::too_many_arguments)]
19699    pub fn fa_dequant_kv_view_f32(
19700        &self,
19701        k: &cudarc::driver::CudaView<u8>,
19702        v: &cudarc::driver::CudaView<u8>,
19703        kf: &mut CudaSlice<f32>,
19704        vf: &mut CudaSlice<f32>,
19705        kv_dim_k: usize,
19706        kv_dim_v: usize,
19707        t_kv: usize,
19708        k_tok_bytes: usize,
19709        v_tok_bytes: usize,
19710        g: bool,
19711    ) -> Result<(), Box<dyn std::error::Error>> {
19712        let f = if g {
19713            self.func_g("fa_dequant_kv_ws_f32")
19714        } else {
19715            self.func("fa_dequant_kv_ws_f32")
19716        };
19717        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19718        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19719        let cfg = LaunchConfig {
19720            grid_dim: (nblk.max(1), 1, 1),
19721            block_dim: (256, 1, 1),
19722            shared_mem_bytes: 0,
19723        };
19724        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19725        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19726        let __s_b = self.gpu.stream();
19727        let mut b = __s_b.launch_builder(&f);
19728        b.arg(k)
19729            .arg(v)
19730            .arg(&mut *kf)
19731            .arg(&mut *vf)
19732            .arg(&kdk)
19733            .arg(&kdv)
19734            .arg(&tkvi)
19735            .arg(&ktb)
19736            .arg(&vtb);
19737        unsafe {
19738            b.launch(cfg)?;
19739        }
19740        Ok(())
19741    }
19742
19743    #[allow(clippy::too_many_arguments)]
19744    pub fn sdpa_naive_quantized_view(
19745        &self,
19746        q: &CudaSlice<f32>,
19747        k: &cudarc::driver::CudaView<u8>,
19748        v: &cudarc::driver::CudaView<u8>,
19749        o: &mut CudaSlice<f32>,
19750        head_dim: usize,
19751        n_head: usize,
19752        n_head_kv: usize,
19753        t: usize,
19754        t_kv: usize,
19755        scale: f32,
19756        causal: bool,
19757        k_tok_bytes: usize,
19758        v_tok_bytes: usize,
19759    ) -> Result<(), Box<dyn std::error::Error>> {
19760        let kv_dim = n_head_kv * head_dim;
19761        let mut kf = self.uninit(t_kv * kv_dim)?;
19762        let mut vf = self.uninit(t_kv * kv_dim)?;
19763        let f = self.func("fa_dequant_kv_ws_f32");
19764        let total = (2 * t_kv * kv_dim) as u64;
19765        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19766        let cfg = LaunchConfig {
19767            grid_dim: (nblk.max(1), 1, 1),
19768            block_dim: (256, 1, 1),
19769            shared_mem_bytes: 0,
19770        };
19771        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19772        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19773        let __s_b = self.gpu.stream();
19774        let mut b = __s_b.launch_builder(&f);
19775        b.arg(k)
19776            .arg(v)
19777            .arg(&mut kf)
19778            .arg(&mut vf)
19779            .arg(&kv_dim_i)
19780            .arg(&kv_dim_i)
19781            .arg(&t_kv_i)
19782            .arg(&k_tok_bytes_i)
19783            .arg(&v_tok_bytes_i);
19784        unsafe { b.launch(cfg)? };
19785        self.sdpa_naive(
19786            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19787        )
19788    }
19789
19790    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
19791    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
19792    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
19793    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
19794    /// unwindowed function above and produces bit-identical output at window == 0.
19795    ///
19796    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
19797    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
19798    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
19799    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
19800    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
19801    #[allow(clippy::too_many_arguments)]
19802    pub fn sdpa_naive_w_quantized_view(
19803        &self,
19804        q: &CudaSlice<f32>,
19805        k: &cudarc::driver::CudaView<u8>,
19806        v: &cudarc::driver::CudaView<u8>,
19807        o: &mut CudaSlice<f32>,
19808        head_dim: usize,
19809        n_head: usize,
19810        n_head_kv: usize,
19811        t: usize,
19812        t_kv: usize,
19813        scale: f32,
19814        causal: bool,
19815        window: usize,
19816        k_tok_bytes: usize,
19817        v_tok_bytes: usize,
19818    ) -> Result<(), Box<dyn std::error::Error>> {
19819        let kv_dim = n_head_kv * head_dim;
19820        let mut kf = self.uninit(t_kv * kv_dim)?;
19821        let mut vf = self.uninit(t_kv * kv_dim)?;
19822        let f = self.func("fa_dequant_kv_ws_f32");
19823        let total = (2 * t_kv * kv_dim) as u64;
19824        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19825        let cfg = LaunchConfig {
19826            grid_dim: (nblk.max(1), 1, 1),
19827            block_dim: (256, 1, 1),
19828            shared_mem_bytes: 0,
19829        };
19830        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19831        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19832        let __s_b = self.gpu.stream();
19833        let mut b = __s_b.launch_builder(&f);
19834        b.arg(k)
19835            .arg(v)
19836            .arg(&mut kf)
19837            .arg(&mut vf)
19838            .arg(&kv_dim_i)
19839            .arg(&kv_dim_i)
19840            .arg(&t_kv_i)
19841            .arg(&k_tok_bytes_i)
19842            .arg(&v_tok_bytes_i);
19843        unsafe { b.launch(cfg)? };
19844        self.sdpa_naive_w(
19845            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19846        )
19847    }
19848
19849    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
19850    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
19851    /// Q/K/V/O [head_dim, n_head(_kv), T].
19852    pub fn fa_prefill(
19853        &self,
19854        q: &CudaSlice<f32>,
19855        k: &CudaSlice<f32>,
19856        v: &CudaSlice<f32>,
19857        o: &mut CudaSlice<f32>,
19858        head_dim: usize,
19859        n_head: usize,
19860        n_head_kv: usize,
19861        t: usize,
19862        t_kv: usize,
19863        scale: f32,
19864        causal: bool,
19865    ) -> Result<(), Box<dyn std::error::Error>> {
19866        if portable_mma_gated() {
19867            return self.sdpa_naive(
19868                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19869            );
19870        }
19871        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
19872        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
19873        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
19874        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
19875        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
19876        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
19877        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
19878        let fa3_on = head_dim == 256
19879            && causal
19880            && t == t_kv
19881            && match std::env::var("MEMRA_FA3").as_deref() {
19882                Ok("0") => false,
19883                // The force arm consults the arch now: the bf16 stage below calls
19884                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
19885                // a portable build. Refuse at the switch, not at the lookup.
19886                Ok("1") => {
19887                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19888                    true
19889                }
19890                _ => cfg!(memra_hopper_mma),
19891            };
19892        if fa3_on {
19893            let n = t * n_head * head_dim;
19894            let nkv = t * n_head_kv * head_dim;
19895            let mut q16 = self.alloc_u8_uninit(n * 2)?;
19896            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19897            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19898            self.f32_to_bf16_into(q, &mut q16, n)?;
19899            self.f32_to_bf16_into(k, &mut k16, nkv)?;
19900            self.f32_to_bf16_into(v, &mut v16, nkv)?;
19901            let rc = {
19902                use cudarc::driver::{DevicePtr, DevicePtrMut};
19903                let stream = self.gpu.stream();
19904                let (qp, _g1) = q16.device_ptr(&stream);
19905                let (kp, _g2) = k16.device_ptr(&stream);
19906                let (vp, _g3) = v16.device_ptr(&stream);
19907                let (op, _g4) = o.device_ptr_mut(&stream);
19908                unsafe {
19909                    memra_fa3_prefill(
19910                        qp as *const core::ffi::c_void,
19911                        kp as *const core::ffi::c_void,
19912                        vp as *const core::ffi::c_void,
19913                        op as *mut f32,
19914                        t as i32,
19915                        n_head as i32,
19916                        n_head_kv as i32,
19917                        head_dim as i32,
19918                        scale,
19919                        stream.cu_stream() as *mut core::ffi::c_void,
19920                    )
19921                }
19922            };
19923            if rc != 0 {
19924                return Err(format!("memra_fa3_prefill rc={rc}").into());
19925            }
19926            return Ok(());
19927        }
19928        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
19929        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
19930        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
19931        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
19932        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19933        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19934        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19935            const BLOCK_Q: usize = 64;
19936            const BKX: usize = 32;
19937            let f = self.func("fa_prefill_bf16_p1");
19938            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19939                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19940            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19941            f.set_attribute(
19942                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19943                shmem as i32,
19944            )?;
19945            let cfg = LaunchConfig {
19946                grid_dim: (
19947                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19948                    n_head as u32,
19949                    1,
19950                ),
19951                block_dim: (32, 4, 1),
19952                shared_mem_bytes: shmem,
19953            };
19954            let (hd, nh, nhkv, ti, tkvi, cz) = (
19955                head_dim as i32,
19956                n_head as i32,
19957                n_head_kv as i32,
19958                t as i32,
19959                t_kv as i32,
19960                causal as i32,
19961            );
19962            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19963            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19964            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19965            let __s_b = self.gpu.stream();
19966            let mut b = __s_b.launch_builder(&f);
19967            b.arg(&qb)
19968                .arg(&kb)
19969                .arg(&vb)
19970                .arg(o)
19971                .arg(&hd)
19972                .arg(&nh)
19973                .arg(&nhkv)
19974                .arg(&ti)
19975                .arg(&tkvi)
19976                .arg(&scale)
19977                .arg(&cz);
19978            unsafe {
19979                b.launch(cfg)?;
19980            }
19981            return Ok(());
19982        }
19983        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
19984        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
19985        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
19986        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
19987        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
19988        const BK: usize = 32;
19989        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
19990        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
19991        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19992        let (block_q, warps, w2_sfx): (usize, u32, &str) =
19993            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19994        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
19995        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
19996        // other head_dims to sdpa_naive before reaching here.
19997        let hd_sfx = fa_hd_suffix(head_dim)?;
19998        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19999        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
20000        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
20001        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
20002        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
20003        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20004        let (kb16, vb16) = if bf16kv {
20005            let n = t_kv * n_head_kv * head_dim;
20006            let mut kb = self.alloc_u8_uninit(n * 2)?;
20007            let mut vb = self.alloc_u8_uninit(n * 2)?;
20008            let fcv = self.func("f32_to_bf16_bulk");
20009            let ni = n as i64;
20010            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20011            let __s_b = self.gpu.stream();
20012            let mut b = __s_b.launch_builder(&fcv);
20013            b.arg(k).arg(&mut kb).arg(&ni);
20014            unsafe {
20015                b.launch(cfgc)?;
20016            }
20017            let __s_b = self.gpu.stream();
20018            let mut b = __s_b.launch_builder(&fcv);
20019            b.arg(v).arg(&mut vb).arg(&ni);
20020            unsafe {
20021                b.launch(cfgc)?;
20022            }
20023            (Some(kb), Some(vb))
20024        } else {
20025            (None, None)
20026        };
20027        let f = self.func(&if bf16kv {
20028            format!("fa_prefill_bf16kv_pp{hd_sfx}")
20029        } else {
20030            format!(
20031                "fa_prefill_f32{}{}{hd_sfx}",
20032                if floor { "" } else { "_pp" },
20033                if floor { "" } else { w2_sfx }
20034            )
20035        });
20036        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
20037        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
20038        let kv_stages = if bf16kv { 2 } else { 1 };
20039        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20040            + 4 * (block_q * BK + 2 * block_q)) as u32;
20041        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20042        f.set_attribute(
20043            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20044            shmem as i32,
20045        )?;
20046        let cfg = LaunchConfig {
20047            grid_dim: (
20048                (t as u32 + block_q as u32 - 1) / block_q as u32,
20049                n_head as u32,
20050                1,
20051            ),
20052            block_dim: (32, warps, 1),
20053            shared_mem_bytes: shmem,
20054        };
20055        let (hd, nh, nhkv, ti, tkvi, cz) = (
20056            head_dim as i32,
20057            n_head as i32,
20058            n_head_kv as i32,
20059            t as i32,
20060            t_kv as i32,
20061            causal as i32,
20062        );
20063        let __s_b = self.gpu.stream();
20064        let mut b = __s_b.launch_builder(&f);
20065        b.arg(q);
20066        match (&kb16, &vb16) {
20067            (Some(kb), Some(vb)) => {
20068                b.arg(kb).arg(vb);
20069            }
20070            _ => {
20071                b.arg(k).arg(v);
20072            }
20073        }
20074        b.arg(o)
20075            .arg(&hd)
20076            .arg(&nh)
20077            .arg(&nhkv)
20078            .arg(&ti)
20079            .arg(&tkvi)
20080            .arg(&scale)
20081            .arg(&cz);
20082        unsafe {
20083            b.launch(cfg)?;
20084        }
20085        Ok(())
20086    }
20087
20088    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
20089    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
20090    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
20091    #[allow(clippy::too_many_arguments)]
20092    pub fn fa_prefill_w(
20093        &self,
20094        q: &CudaSlice<f32>,
20095        k: &CudaSlice<f32>,
20096        v: &CudaSlice<f32>,
20097        o: &mut CudaSlice<f32>,
20098        head_dim: usize,
20099        n_head: usize,
20100        n_head_kv: usize,
20101        t: usize,
20102        t_kv: usize,
20103        scale: f32,
20104        causal: bool,
20105        window: usize,
20106    ) -> Result<(), Box<dyn std::error::Error>> {
20107        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
20108        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
20109        if portable_mma_gated() {
20110            return self.sdpa_naive_w(
20111                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20112            );
20113        }
20114        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
20115        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
20116        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
20117        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20118        let faw_f32 =
20119            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20120        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20121        self.fa_prefill_w_arm(
20122            q,
20123            k,
20124            v,
20125            o,
20126            head_dim,
20127            n_head,
20128            n_head_kv,
20129            t,
20130            t_kv,
20131            scale,
20132            causal,
20133            window,
20134            floor || faw_f32,
20135            floor,
20136        )
20137    }
20138
20139    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
20140    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
20141    #[allow(clippy::too_many_arguments)]
20142    pub fn fa_prefill_w_pre(
20143        &self,
20144        qb: &CudaSlice<u8>,
20145        kb: &CudaSlice<u8>,
20146        vb: &CudaSlice<u8>,
20147        o: &mut CudaSlice<f32>,
20148        head_dim: usize,
20149        n_head: usize,
20150        n_head_kv: usize,
20151        t: usize,
20152        t_kv: usize,
20153        scale: f32,
20154        causal: bool,
20155        window: usize,
20156        v_f16: bool,
20157    ) -> Result<(), Box<dyn std::error::Error>> {
20158        const BLOCK_Q: usize = 64;
20159        const BK: usize = 32;
20160        debug_assert_eq!(head_dim, 256);
20161        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20162        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20163        if hp {
20164            const BLOCK_QH: usize = 32;
20165            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
20166            // else re-encode through the pooled scratch (stream-ordered reuse).
20167            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20168            let vh: &CudaSlice<u8> = if v_f16 {
20169                vb
20170            } else {
20171                let n = t_kv * n_head_kv * head_dim;
20172                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20173                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20174                }
20175                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20176                vguard.as_ref().unwrap()
20177            };
20178            let f = self.func("fa_prefill_w_bf16_p1h2");
20179            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20180            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20181            f.set_attribute(
20182                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20183                shmem as i32,
20184            )?;
20185            let cfg = LaunchConfig {
20186                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20187                block_dim: (32, 4, 1),
20188                shared_mem_bytes: shmem,
20189            };
20190            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20191                head_dim as i32,
20192                n_head as i32,
20193                n_head_kv as i32,
20194                t as i32,
20195                t_kv as i32,
20196                causal as i32,
20197                window as i32,
20198            );
20199            let __s_b = self.gpu.stream();
20200            let mut b = __s_b.launch_builder(&f);
20201            b.arg(qb)
20202                .arg(kb)
20203                .arg(vh)
20204                .arg(o)
20205                .arg(&hd)
20206                .arg(&nh)
20207                .arg(&nhkv)
20208                .arg(&ti)
20209                .arg(&tkvi)
20210                .arg(&scale)
20211                .arg(&cz)
20212                .arg(&wi);
20213            unsafe {
20214                b.launch(cfg)?;
20215            }
20216            return Ok(());
20217        }
20218        let f = self.func("fa_prefill_w_bf16_p1");
20219        let shmem =
20220            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) 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: (
20228                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20229                n_head as u32,
20230                1,
20231            ),
20232            block_dim: (32, 4, 1),
20233            shared_mem_bytes: shmem,
20234        };
20235        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20236            head_dim as i32,
20237            n_head as i32,
20238            n_head_kv as i32,
20239            t as i32,
20240            t_kv as i32,
20241            causal as i32,
20242            window as i32,
20243        );
20244        let __s_b = self.gpu.stream();
20245        let mut b = __s_b.launch_builder(&f);
20246        b.arg(qb)
20247            .arg(kb)
20248            .arg(vb)
20249            .arg(o)
20250            .arg(&hd)
20251            .arg(&nh)
20252            .arg(&nhkv)
20253            .arg(&ti)
20254            .arg(&tkvi)
20255            .arg(&scale)
20256            .arg(&cz)
20257            .arg(&wi);
20258        unsafe {
20259            b.launch(cfg)?;
20260        }
20261        Ok(())
20262    }
20263
20264    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
20265    #[allow(clippy::too_many_arguments)]
20266    pub fn fa_prefill_w_arm(
20267        &self,
20268        q: &CudaSlice<f32>,
20269        k: &CudaSlice<f32>,
20270        v: &CudaSlice<f32>,
20271        o: &mut CudaSlice<f32>,
20272        head_dim: usize,
20273        n_head: usize,
20274        n_head_kv: usize,
20275        t: usize,
20276        t_kv: usize,
20277        scale: f32,
20278        causal: bool,
20279        window: usize,
20280        f32_stage: bool,
20281        floor: bool,
20282    ) -> Result<(), Box<dyn std::error::Error>> {
20283        const BLOCK_Q: usize = 64;
20284        const BK: usize = 32;
20285        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20286        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
20287        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
20288        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
20289        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20290        let p1 = !floor
20291            && !f32_stage
20292            && *P1_ON.get_or_init(|| {
20293                std::env::var("MEMRA_FAW_P1")
20294                    .map(|v| v != "0")
20295                    .unwrap_or(true)
20296            });
20297        let hp =
20298            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20299        if hp {
20300            const BLOCK_QH: usize = 32;
20301            let f = self.func("fa_prefill_w_bf16_p1h2");
20302            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20303            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20304            f.set_attribute(
20305                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20306                shmem as i32,
20307            )?;
20308            let cfg = LaunchConfig {
20309                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20310                block_dim: (32, 4, 1),
20311                shared_mem_bytes: shmem,
20312            };
20313            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20314                head_dim as i32,
20315                n_head as i32,
20316                n_head_kv as i32,
20317                t as i32,
20318                t_kv as i32,
20319                causal as i32,
20320                window as i32,
20321            );
20322            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20323            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20324            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20325            let __s_b = self.gpu.stream();
20326            let mut b = __s_b.launch_builder(&f);
20327            b.arg(&qb)
20328                .arg(&kb)
20329                .arg(&vh)
20330                .arg(o)
20331                .arg(&hd)
20332                .arg(&nh)
20333                .arg(&nhkv)
20334                .arg(&ti)
20335                .arg(&tkvi)
20336                .arg(&scale)
20337                .arg(&cz)
20338                .arg(&wi);
20339            unsafe {
20340                b.launch(cfg)?;
20341            }
20342            return Ok(());
20343        }
20344        if p1 {
20345            let f = self.func("fa_prefill_w_bf16_p1");
20346            let shmem =
20347                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20348            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20349            f.set_attribute(
20350                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20351                shmem as i32,
20352            )?;
20353            let cfg = LaunchConfig {
20354                grid_dim: (
20355                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20356                    n_head as u32,
20357                    1,
20358                ),
20359                block_dim: (32, 4, 1),
20360                shared_mem_bytes: shmem,
20361            };
20362            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20363                head_dim as i32,
20364                n_head as i32,
20365                n_head_kv as i32,
20366                t as i32,
20367                t_kv as i32,
20368                causal as i32,
20369                window as i32,
20370            );
20371            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20372            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20373            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20374            let __s_b = self.gpu.stream();
20375            let mut b = __s_b.launch_builder(&f);
20376            b.arg(&qb)
20377                .arg(&kb)
20378                .arg(&vb)
20379                .arg(o)
20380                .arg(&hd)
20381                .arg(&nh)
20382                .arg(&nhkv)
20383                .arg(&ti)
20384                .arg(&tkvi)
20385                .arg(&scale)
20386                .arg(&cz)
20387                .arg(&wi);
20388            unsafe {
20389                b.launch(cfg)?;
20390            }
20391            return Ok(());
20392        }
20393        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
20394        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
20395        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20396        let g4 = !floor
20397            && !f32_stage
20398            && n_head_kv == 1
20399            && n_head % 4 == 0
20400            && *G4_ON.get_or_init(|| {
20401                std::env::var("MEMRA_FAW_G4")
20402                    .map(|v| v != "0")
20403                    .unwrap_or(true)
20404            });
20405        if g4 {
20406            const SP_M: usize = 16;
20407            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
20408            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
20409            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20410            let o2 = *O2_ON.get_or_init(|| {
20411                std::env::var("MEMRA_FAW_O2")
20412                    .map(|v| v != "0")
20413                    .unwrap_or(true)
20414            });
20415            let f = self.func(if o2 {
20416                "fa_prefill_w_bf16_g4o2"
20417            } else {
20418                "fa_prefill_w_bf16_g4"
20419            });
20420            let shmem = if o2 {
20421                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20422            } else {
20423                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20424                    as u32
20425            };
20426            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20427            f.set_attribute(
20428                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20429                shmem as i32,
20430            )?;
20431            let cfg = LaunchConfig {
20432                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20433                block_dim: (32, 4, 1),
20434                shared_mem_bytes: shmem,
20435            };
20436            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20437                head_dim as i32,
20438                n_head as i32,
20439                n_head_kv as i32,
20440                t as i32,
20441                t_kv as i32,
20442                causal as i32,
20443                window as i32,
20444            );
20445            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20446            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20447            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20448            let __s_b = self.gpu.stream();
20449            let mut b = __s_b.launch_builder(&f);
20450            b.arg(&qb)
20451                .arg(&kb)
20452                .arg(&vb)
20453                .arg(o)
20454                .arg(&hd)
20455                .arg(&nh)
20456                .arg(&nhkv)
20457                .arg(&ti)
20458                .arg(&tkvi)
20459                .arg(&scale)
20460                .arg(&cz)
20461                .arg(&wi);
20462            unsafe {
20463                b.launch(cfg)?;
20464            }
20465            return Ok(());
20466        }
20467        let f = self.func(if floor {
20468            "fa_prefill_w_f32"
20469        } else if f32_stage {
20470            "fa_prefill_w_f32_pp"
20471        } else {
20472            "fa_prefill_w_bf16_pp"
20473        });
20474        let shmem =
20475            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20476        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20477        f.set_attribute(
20478            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20479            shmem as i32,
20480        )?;
20481        let cfg = LaunchConfig {
20482            grid_dim: (
20483                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20484                n_head as u32,
20485                1,
20486            ),
20487            block_dim: (32, 4, 1),
20488            shared_mem_bytes: shmem,
20489        };
20490        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20491            head_dim as i32,
20492            n_head as i32,
20493            n_head_kv as i32,
20494            t as i32,
20495            t_kv as i32,
20496            causal as i32,
20497            window as i32,
20498        );
20499        if f32_stage {
20500            let __s_b = self.gpu.stream();
20501            let mut b = __s_b.launch_builder(&f);
20502            b.arg(q)
20503                .arg(k)
20504                .arg(v)
20505                .arg(o)
20506                .arg(&hd)
20507                .arg(&nh)
20508                .arg(&nhkv)
20509                .arg(&ti)
20510                .arg(&tkvi)
20511                .arg(&scale)
20512                .arg(&cz)
20513                .arg(&wi);
20514            unsafe {
20515                b.launch(cfg)?;
20516            }
20517        } else {
20518            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20519            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20520            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20521            let __s_b = self.gpu.stream();
20522            let mut b = __s_b.launch_builder(&f);
20523            b.arg(&qb)
20524                .arg(&kb)
20525                .arg(&vb)
20526                .arg(o)
20527                .arg(&hd)
20528                .arg(&nh)
20529                .arg(&nhkv)
20530                .arg(&ti)
20531                .arg(&tkvi)
20532                .arg(&scale)
20533                .arg(&cz)
20534                .arg(&wi);
20535            unsafe {
20536                b.launch(cfg)?;
20537            }
20538        }
20539        Ok(())
20540    }
20541
20542    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
20543    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
20544    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
20545    #[allow(clippy::too_many_arguments)]
20546    pub fn fa_prefill_hd512(
20547        &self,
20548        q: &CudaSlice<f32>,
20549        k: &CudaSlice<f32>,
20550        v: &CudaSlice<f32>,
20551        o: &mut CudaSlice<f32>,
20552        head_dim: usize,
20553        n_head: usize,
20554        n_head_kv: usize,
20555        t: usize,
20556        t_kv: usize,
20557        scale: f32,
20558        causal: bool,
20559    ) -> Result<(), Box<dyn std::error::Error>> {
20560        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
20561        if portable_mma_gated() {
20562            return self.sdpa_naive(
20563                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20564            );
20565        }
20566        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
20567        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
20568        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
20569        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
20570        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
20571        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20572        let f32_stage =
20573            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20574        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
20575        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
20576        // Own numeric config (partial-sum order) — battery-gated.
20577        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20578        let sp = !f32_stage
20579            && *SP_ON.get_or_init(|| {
20580                std::env::var("MEMRA_FA512_SP")
20581                    .map(|v| v != "0")
20582                    .unwrap_or(true)
20583            });
20584        self.fa_prefill_hd512_arm(
20585            q,
20586            k,
20587            v,
20588            o,
20589            head_dim,
20590            n_head,
20591            n_head_kv,
20592            t,
20593            t_kv,
20594            scale,
20595            causal,
20596            f32_stage,
20597            sp,
20598            sp && fa_f16pv_on(),
20599        )
20600    }
20601
20602    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
20603    #[allow(clippy::too_many_arguments)]
20604    pub fn fa_prefill_hd512_pre(
20605        &self,
20606        qb: &CudaSlice<u8>,
20607        kb: &CudaSlice<u8>,
20608        vb: &CudaSlice<u8>,
20609        o: &mut CudaSlice<f32>,
20610        head_dim: usize,
20611        n_head: usize,
20612        n_head_kv: usize,
20613        t: usize,
20614        t_kv: usize,
20615        scale: f32,
20616        causal: bool,
20617        v_f16: bool,
20618    ) -> Result<(), Box<dyn std::error::Error>> {
20619        debug_assert_eq!(head_dim, 512);
20620        const SP_M: usize = 16;
20621        const BKS: usize = 32;
20622        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
20623        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
20624        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
20625        let f16pv = fa_f16pv_on();
20626        let nw = if f16pv { fa512_wide_warps() } else { 2 };
20627        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20628        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20629        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20630        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20631            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
20632            let n = t_kv * n_head_kv * head_dim;
20633            let need = n * 2;
20634            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20635                *vguard = Some(self.alloc_uninit::<u8>(need)?);
20636            }
20637            let dst = vguard.as_mut().unwrap();
20638            self.bf16_to_f16_into(vb, n, dst)?;
20639            vguard.as_ref().unwrap()
20640        } else {
20641            vb
20642        };
20643        let f = self.func(if hp {
20644            "fa_prefill_bf16_hd512_sp16h2"
20645        } else {
20646            match (f16pv, nw) {
20647                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20648                (true, _) => "fa_prefill_bf16_hd512_sp16",
20649                _ => "fa_prefill_bf16_hd512_sp",
20650            }
20651        });
20652        let (nwarp, npart) = if hp {
20653            (4usize, 4usize)
20654        } else if nw > 2 {
20655            (nw, nw)
20656        } else {
20657            (2, 1)
20658        };
20659        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
20660        let shmem = if hp {
20661            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20662                as u32
20663        } else {
20664            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20665                + 4 * (npart * SP_M * BKS + SP_M)) as u32
20666        };
20667        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20668        f.set_attribute(
20669            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20670            shmem as i32,
20671        )?;
20672        let grid_y = if hp {
20673            (n_head / 2) as u32
20674        } else {
20675            n_head as u32
20676        };
20677        let cfg = LaunchConfig {
20678            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20679            block_dim: (32, nwarp as u32, 1),
20680            shared_mem_bytes: shmem,
20681        };
20682        let (hd, nh, nhkv, ti, tkvi, cz) = (
20683            head_dim as i32,
20684            n_head as i32,
20685            n_head_kv as i32,
20686            t as i32,
20687            t_kv as i32,
20688            causal as i32,
20689        );
20690        let __s_b = self.gpu.stream();
20691        let mut b = __s_b.launch_builder(&f);
20692        b.arg(qb)
20693            .arg(kb)
20694            .arg(vref)
20695            .arg(o)
20696            .arg(&hd)
20697            .arg(&nh)
20698            .arg(&nhkv)
20699            .arg(&ti)
20700            .arg(&tkvi)
20701            .arg(&scale)
20702            .arg(&cz);
20703        unsafe {
20704            b.launch(cfg)?;
20705        }
20706        Ok(())
20707    }
20708
20709    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
20710    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
20711    #[allow(clippy::too_many_arguments)]
20712    pub fn fa_prefill_hd512_arm(
20713        &self,
20714        q: &CudaSlice<f32>,
20715        k: &CudaSlice<f32>,
20716        v: &CudaSlice<f32>,
20717        o: &mut CudaSlice<f32>,
20718        head_dim: usize,
20719        n_head: usize,
20720        n_head_kv: usize,
20721        t: usize,
20722        t_kv: usize,
20723        scale: f32,
20724        causal: bool,
20725        f32_stage: bool,
20726        sp: bool,
20727        f16pv: bool,
20728    ) -> Result<(), Box<dyn std::error::Error>> {
20729        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20730        if sp && !f32_stage {
20731            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
20732            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
20733            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
20734            const SP_M: usize = 16;
20735            const BKS: usize = 32;
20736            let nw = if f16pv { fa512_wide_warps() } else { 2 };
20737            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20738            let f = self.func(if hp {
20739                "fa_prefill_bf16_hd512_sp16h2"
20740            } else {
20741                match (f16pv, nw) {
20742                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20743                    (true, _) => "fa_prefill_bf16_hd512_sp16",
20744                    _ => "fa_prefill_bf16_hd512_sp",
20745                }
20746            });
20747            let (nwarp, npart) = if hp {
20748                (4usize, 4usize)
20749            } else if nw > 2 {
20750                (nw, nw)
20751            } else {
20752                (2, 1)
20753            };
20754            let shmem = if hp {
20755                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20756                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20757            } else {
20758                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20759                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
20760            };
20761            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20762            f.set_attribute(
20763                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20764                shmem as i32,
20765            )?;
20766            let grid_y = if hp {
20767                (n_head / 2) as u32
20768            } else {
20769                n_head as u32
20770            };
20771            let cfg = LaunchConfig {
20772                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20773                block_dim: (32, nwarp as u32, 1),
20774                shared_mem_bytes: shmem,
20775            };
20776            let (hd, nh, nhkv, ti, tkvi, cz) = (
20777                head_dim as i32,
20778                n_head as i32,
20779                n_head_kv as i32,
20780                t as i32,
20781                t_kv as i32,
20782                causal as i32,
20783            );
20784            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20785            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20786            let vb = if f16pv {
20787                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20788            } else {
20789                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20790            };
20791            let __s_b = self.gpu.stream();
20792            let mut b = __s_b.launch_builder(&f);
20793            b.arg(&qb)
20794                .arg(&kb)
20795                .arg(&vb)
20796                .arg(o)
20797                .arg(&hd)
20798                .arg(&nh)
20799                .arg(&nhkv)
20800                .arg(&ti)
20801                .arg(&tkvi)
20802                .arg(&scale)
20803                .arg(&cz);
20804            unsafe {
20805                b.launch(cfg)?;
20806            }
20807            return Ok(());
20808        }
20809        const BLOCK_Q: usize = 32;
20810        const BK: usize = 32;
20811        const HALF: usize = 256;
20812        let f = self.func(if f32_stage {
20813            "fa_prefill_f32_hd512"
20814        } else {
20815            "fa_prefill_bf16_hd512"
20816        });
20817        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
20818        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20819            + 4 * BLOCK_Q) as u32;
20820        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20821        f.set_attribute(
20822            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20823            shmem as i32,
20824        )?;
20825        let cfg = LaunchConfig {
20826            grid_dim: (
20827                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20828                n_head as u32,
20829                2,
20830            ),
20831            block_dim: (32, 2, 1),
20832            shared_mem_bytes: shmem,
20833        };
20834        let (hd, nh, nhkv, ti, tkvi, cz) = (
20835            head_dim as i32,
20836            n_head as i32,
20837            n_head_kv as i32,
20838            t as i32,
20839            t_kv as i32,
20840            causal as i32,
20841        );
20842        if f32_stage {
20843            let __s_b = self.gpu.stream();
20844            let mut b = __s_b.launch_builder(&f);
20845            b.arg(q)
20846                .arg(k)
20847                .arg(v)
20848                .arg(o)
20849                .arg(&hd)
20850                .arg(&nh)
20851                .arg(&nhkv)
20852                .arg(&ti)
20853                .arg(&tkvi)
20854                .arg(&scale)
20855                .arg(&cz);
20856            unsafe {
20857                b.launch(cfg)?;
20858            }
20859        } else {
20860            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20861            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20862            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20863            let __s_b = self.gpu.stream();
20864            let mut b = __s_b.launch_builder(&f);
20865            b.arg(&qb)
20866                .arg(&kb)
20867                .arg(&vb)
20868                .arg(o)
20869                .arg(&hd)
20870                .arg(&nh)
20871                .arg(&nhkv)
20872                .arg(&ti)
20873                .arg(&tkvi)
20874                .arg(&scale)
20875                .arg(&cz);
20876            unsafe {
20877                b.launch(cfg)?;
20878            }
20879        }
20880        Ok(())
20881    }
20882
20883    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
20884    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
20885    /// separate f32_to_bf16 the FA entries would run).
20886    #[allow(clippy::too_many_arguments)]
20887    pub fn rope_neox2_bf16e(
20888        &self,
20889        q: &mut CudaSlice<f32>,
20890        k: &mut CudaSlice<f32>,
20891        qb: &mut CudaSlice<u8>,
20892        kb: &mut CudaSlice<u8>,
20893        pos: &CudaSlice<i32>,
20894        head_dim: usize,
20895        n_dims: usize,
20896        nh_q: usize,
20897        nh_k: usize,
20898        n_tokens: usize,
20899        base: f32,
20900        freq_scale: f32,
20901        ff: Option<&CudaSlice<f32>>,
20902    ) -> Result<(), Box<dyn std::error::Error>> {
20903        let f = self.func("rope_neox2_bf16e_f32");
20904        let rows = ((nh_q + nh_k) * n_tokens) as u32;
20905        let cfg = LaunchConfig {
20906            grid_dim: (rows, 1, 1),
20907            block_dim: ((head_dim / 2) as u32, 1, 1),
20908            shared_mem_bytes: 0,
20909        };
20910        let theta_scale = base.powf(-2.0 / n_dims as f32);
20911        let (hd, nd, nhq, nhk, nt) = (
20912            head_dim as i32,
20913            n_dims as i32,
20914            nh_q as i32,
20915            nh_k as i32,
20916            n_tokens as i32,
20917        );
20918        let __s_b = self.gpu.stream();
20919        let mut b = __s_b.launch_builder(&f);
20920        match ff {
20921            Some(t) => {
20922                b.arg(&mut *q)
20923                    .arg(&mut *k)
20924                    .arg(&mut *qb)
20925                    .arg(&mut *kb)
20926                    .arg(pos)
20927                    .arg(&hd)
20928                    .arg(&nd)
20929                    .arg(&nhq)
20930                    .arg(&nhk)
20931                    .arg(&nt)
20932                    .arg(&theta_scale)
20933                    .arg(&freq_scale)
20934                    .arg(t);
20935                unsafe {
20936                    b.launch(cfg)?;
20937                }
20938            }
20939            None => {
20940                let null: u64 = 0;
20941                b.arg(&mut *q)
20942                    .arg(&mut *k)
20943                    .arg(&mut *qb)
20944                    .arg(&mut *kb)
20945                    .arg(pos)
20946                    .arg(&hd)
20947                    .arg(&nd)
20948                    .arg(&nhq)
20949                    .arg(&nhk)
20950                    .arg(&nt)
20951                    .arg(&theta_scale)
20952                    .arg(&freq_scale)
20953                    .arg(&null);
20954                unsafe {
20955                    b.launch(cfg)?;
20956                }
20957            }
20958        }
20959        Ok(())
20960    }
20961
20962    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
20963    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
20964    pub fn f32_to_bf16(
20965        &self,
20966        x: &CudaSlice<f32>,
20967        n: usize,
20968    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20969        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
20970        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20971        let f = self.func("f32_to_bf16_flat");
20972        let n_i = n as i64;
20973        let cfg = LaunchConfig {
20974            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20975            block_dim: (256, 1, 1),
20976            shared_mem_bytes: 0,
20977        };
20978        let __s_b = self.gpu.stream();
20979        let mut b = __s_b.launch_builder(&f);
20980        b.arg(x).arg(&mut y).arg(&n_i);
20981        unsafe {
20982            b.launch(cfg)?;
20983        }
20984        Ok(y)
20985    }
20986
20987    pub fn f32_to_f16(
20988        &self,
20989        x: &CudaSlice<f32>,
20990        n: usize,
20991    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20992        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20993        let mut y = self.alloc_uninit::<u8>(n * 2)?;
20994        let f = self.func("f32_to_f16_flat");
20995        let n_i = n as i64;
20996        let cfg = LaunchConfig {
20997            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20998            block_dim: (256, 1, 1),
20999            shared_mem_bytes: 0,
21000        };
21001        let __s_b = self.gpu.stream();
21002        let mut b = __s_b.launch_builder(&f);
21003        b.arg(x).arg(&mut y).arg(&n_i);
21004        unsafe {
21005            b.launch(cfg)?;
21006        }
21007        Ok(y)
21008    }
21009
21010    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
21011    pub fn bf16_to_f16(
21012        &self,
21013        xb: &CudaSlice<u8>,
21014        n: usize,
21015    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21016        let mut y = self.alloc_uninit::<u8>(n * 2)?;
21017        self.bf16_to_f16_into(xb, n, &mut y)?;
21018        Ok(y)
21019    }
21020
21021    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
21022    pub fn bf16_to_f16_into(
21023        &self,
21024        xb: &CudaSlice<u8>,
21025        n: usize,
21026        y: &mut CudaSlice<u8>,
21027    ) -> Result<(), Box<dyn std::error::Error>> {
21028        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21029        assert!(y.len() >= n * 2);
21030        let f = self.func("bf16_to_f16_flat");
21031        let n2 = (n / 2) as i64;
21032        let cfg = LaunchConfig {
21033            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21034            block_dim: (256, 1, 1),
21035            shared_mem_bytes: 0,
21036        };
21037        let __s_b = self.gpu.stream();
21038        let mut b = __s_b.launch_builder(&f);
21039        b.arg(xb).arg(y).arg(&n2);
21040        unsafe {
21041            b.launch(cfg)?;
21042        }
21043        Ok(())
21044    }
21045
21046    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
21047    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
21048    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
21049    /// head_dim in {256, 128}, bf16kv lane on.
21050    #[allow(clippy::too_many_arguments)]
21051    pub fn fa_prefill_vl8(
21052        &self,
21053        seqs: &[FaSeqVl],
21054        head_dim: usize,
21055        n_head: usize,
21056        n_head_kv: usize,
21057        scale: f32,
21058    ) -> Result<(), Box<dyn std::error::Error>> {
21059        const BK: usize = 32;
21060        let b = seqs.len();
21061        assert!(b >= 1 && b <= 8);
21062        let mut packed = [FaSeqVl::default(); 8];
21063        packed[..b].copy_from_slice(seqs);
21064        let v = FaVl8(packed);
21065        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21066        let ept = (n_head_kv * head_dim) as i32;
21067        {
21068            let f = self.func("fa_mirror_vl");
21069            let max_n = (max_t as i64) * ept as i64;
21070            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21071            for which in 0..2i32 {
21072                let cfg = LaunchConfig {
21073                    grid_dim: (blocks, 1, b as u32),
21074                    block_dim: (256, 1, 1),
21075                    shared_mem_bytes: 0,
21076                };
21077                let __s_lb = self.gpu.stream();
21078                let mut lb = __s_lb.launch_builder(&f);
21079                lb.arg(&v).arg(&ept).arg(&which);
21080                unsafe {
21081                    lb.launch(cfg)?;
21082                }
21083            }
21084        }
21085        let hd_sfx = fa_hd_suffix(head_dim)?;
21086        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21087        let block_q = 64usize;
21088        let kv_stages = 2usize;
21089        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21090            + 4 * (block_q * BK + 2 * block_q)) as u32;
21091        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21092        f.set_attribute(
21093            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21094            shmem as i32,
21095        )?;
21096        let cfg = LaunchConfig {
21097            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21098            block_dim: (32, 4, 1),
21099            shared_mem_bytes: shmem,
21100        };
21101        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21102        let __s_lb = self.gpu.stream();
21103        let mut lb = __s_lb.launch_builder(&f);
21104        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21105        unsafe {
21106            lb.launch(cfg)?;
21107        }
21108        Ok(())
21109    }
21110
21111    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
21112    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
21113    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
21114    #[allow(clippy::too_many_arguments)]
21115    pub fn attn_pre_vl8(
21116        &self,
21117        seqs: &[AttnPreVl],
21118        wq: &CudaSlice<f32>,
21119        wk: &CudaSlice<f32>,
21120        head_dim: usize,
21121        rope_dims: usize,
21122        n_head: usize,
21123        n_head_kv: usize,
21124        eps: f32,
21125        freq_base: f32,
21126        freq_scale: f32,
21127        kv_dim_k: usize,
21128        kv_dim_v: usize,
21129        k_tok_bytes: usize,
21130        v_tok_bytes: usize,
21131    ) -> Result<(), Box<dyn std::error::Error>> {
21132        let b = seqs.len();
21133        assert!(b >= 1 && b <= 8);
21134        let mut packed = [AttnPreVl::default(); 8];
21135        packed[..b].copy_from_slice(seqs);
21136        let v = AttnPreVl8(packed);
21137        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21138        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21139        {
21140            let f = self.func("q_gate_split_vl");
21141            let n = max_t * (n_head * head_dim) as u32;
21142            let cfg = LaunchConfig {
21143                grid_dim: (n.div_ceil(256), 1, b as u32),
21144                block_dim: (256, 1, 1),
21145                shared_mem_bytes: 0,
21146            };
21147            let __s_lb = self.gpu.stream();
21148            let mut lb = __s_lb.launch_builder(&f);
21149            lb.arg(&v).arg(&hd).arg(&nh);
21150            unsafe {
21151                lb.launch(cfg)?;
21152            }
21153        }
21154        {
21155            let f = self.func("attn_rms_vl");
21156            let cfg = LaunchConfig {
21157                grid_dim: (max_t * n_head as u32, 2, b as u32),
21158                block_dim: (rms_block(), 1, 1),
21159                shared_mem_bytes: 0,
21160            };
21161            let __s_lb = self.gpu.stream();
21162            let mut lb = __s_lb.launch_builder(&f);
21163            lb.arg(&v)
21164                .arg(wq)
21165                .arg(wk)
21166                .arg(&hd)
21167                .arg(&nh)
21168                .arg(&nhkv)
21169                .arg(&eps);
21170            unsafe {
21171                lb.launch(cfg)?;
21172            }
21173        }
21174        {
21175            let f = self.func("attn_rope_vl");
21176            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21177            let nd = rope_dims as i32;
21178            let cfg = LaunchConfig {
21179                grid_dim: (max_t * n_head as u32, 2, b as u32),
21180                block_dim: ((head_dim / 2) as u32, 1, 1),
21181                shared_mem_bytes: 0,
21182            };
21183            let __s_lb = self.gpu.stream();
21184            let mut lb = __s_lb.launch_builder(&f);
21185            lb.arg(&v)
21186                .arg(&hd)
21187                .arg(&nd)
21188                .arg(&nh)
21189                .arg(&nhkv)
21190                .arg(&theta_scale)
21191                .arg(&freq_scale);
21192            unsafe {
21193                lb.launch(cfg)?;
21194            }
21195        }
21196        {
21197            let f = self.func("append_kv_vl");
21198            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21199            let cfg = LaunchConfig {
21200                grid_dim: (nblk, max_t, b as u32),
21201                block_dim: (32, 1, 1),
21202                shared_mem_bytes: 0,
21203            };
21204            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21205            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21206            let __s_lb = self.gpu.stream();
21207            let mut lb = __s_lb.launch_builder(&f);
21208            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21209            unsafe {
21210                lb.launch(cfg)?;
21211            }
21212        }
21213        Ok(())
21214    }
21215
21216    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
21217    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
21218    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
21219    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
21220    pub fn fa_prefill_view(
21221        &self,
21222        q: &CudaSlice<f32>,
21223        k: &cudarc::driver::CudaView<u8>,
21224        v: &cudarc::driver::CudaView<u8>,
21225        o: &mut CudaSlice<f32>,
21226        head_dim: usize,
21227        n_head: usize,
21228        n_head_kv: usize,
21229        t: usize,
21230        t_kv: usize,
21231        scale: f32,
21232        causal: bool,
21233        k_tok_bytes: usize,
21234        v_tok_bytes: usize,
21235        g: bool,
21236    ) -> Result<(), Box<dyn std::error::Error>> {
21237        if portable_mma_gated() {
21238            return self.sdpa_naive_quantized_view(
21239                q,
21240                k,
21241                v,
21242                o,
21243                head_dim,
21244                n_head,
21245                n_head_kv,
21246                t,
21247                t_kv,
21248                scale,
21249                causal,
21250                k_tok_bytes,
21251                v_tok_bytes,
21252            );
21253        }
21254        const BLOCK_Q: usize = 64;
21255        const BK: usize = 32;
21256        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
21257        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
21258        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21259        let f = if g {
21260            self.func_g(&name)
21261        } else {
21262            self.func(&name)
21263        };
21264        let shmem =
21265            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21266        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21267        f.set_attribute(
21268            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21269            shmem as i32,
21270        )?;
21271        let cfg = LaunchConfig {
21272            grid_dim: (
21273                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21274                n_head as u32,
21275                1,
21276            ),
21277            block_dim: (32, 4, 1),
21278            shared_mem_bytes: shmem,
21279        };
21280        let (hd, nh, nhkv, ti, tkvi, cz) = (
21281            head_dim as i32,
21282            n_head as i32,
21283            n_head_kv as i32,
21284            t as i32,
21285            t_kv as i32,
21286            causal as i32,
21287        );
21288        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21289        let __s_b = self.gpu.stream();
21290        let mut b = __s_b.launch_builder(&f);
21291        b.arg(q)
21292            .arg(k)
21293            .arg(v)
21294            .arg(o)
21295            .arg(&hd)
21296            .arg(&nh)
21297            .arg(&nhkv)
21298            .arg(&ti)
21299            .arg(&tkvi)
21300            .arg(&scale)
21301            .arg(&cz)
21302            .arg(&ktb)
21303            .arg(&vtb);
21304        unsafe {
21305            b.launch(cfg)?;
21306        }
21307        Ok(())
21308    }
21309
21310    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
21311    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
21312    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
21313    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
21314    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
21315    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
21316    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
21317    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
21318    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
21319    #[allow(clippy::too_many_arguments)]
21320    pub fn fa_prefill_view_ws(
21321        &self,
21322        q: &CudaSlice<f32>,
21323        k: &cudarc::driver::CudaView<u8>,
21324        v: &cudarc::driver::CudaView<u8>,
21325        o: &mut CudaSlice<f32>,
21326        head_dim: usize,
21327        n_head: usize,
21328        n_head_kv: usize,
21329        t: usize,
21330        t_kv: usize,
21331        scale: f32,
21332        causal: bool,
21333        k_tok_bytes: usize,
21334        v_tok_bytes: usize,
21335        g: bool,
21336    ) -> Result<(), Box<dyn std::error::Error>> {
21337        if portable_mma_gated() {
21338            return self.sdpa_naive_quantized_view(
21339                q,
21340                k,
21341                v,
21342                o,
21343                head_dim,
21344                n_head,
21345                n_head_kv,
21346                t,
21347                t_kv,
21348                scale,
21349                causal,
21350                k_tok_bytes,
21351                v_tok_bytes,
21352            );
21353        }
21354        const BLOCK_Q: usize = 64;
21355        const BK: usize = 32;
21356        let kv_dim_k = n_head_kv * head_dim;
21357        let kv_dim_v = n_head_kv * head_dim;
21358        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
21359        let v_ws_bytes = t_kv * kv_dim_v * 2;
21360        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
21361        let mut guard = self.prime_deqw_ws.lock().unwrap();
21362        let need_grow = match guard.as_ref() {
21363            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21364            None => true,
21365        };
21366        if need_grow {
21367            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21368            let (ck, cv) = guard
21369                .as_ref()
21370                .map(|(a, b)| (a.len(), b.len()))
21371                .unwrap_or((0, 0));
21372            *guard = Some((
21373                self.alloc_u8(grow(ck, k_ws_bytes))?,
21374                self.alloc_u8(grow(cv, v_ws_bytes))?,
21375            ));
21376        }
21377        let (kw, vw) = guard.as_mut().unwrap();
21378        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
21379        {
21380            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
21381            let f = if g {
21382                self.func_g("fa_dequant_kv_ws_bf16")
21383            } else {
21384                self.func("fa_dequant_kv_ws_bf16")
21385            };
21386            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21387            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21388            let cfg = LaunchConfig {
21389                grid_dim: (nblk.max(1), 1, 1),
21390                block_dim: (256, 1, 1),
21391                shared_mem_bytes: 0,
21392            };
21393            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21394            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21395            let __s_b = self.gpu.stream();
21396            let mut b = __s_b.launch_builder(&f);
21397            b.arg(k)
21398                .arg(v)
21399                .arg(&mut *kw)
21400                .arg(&mut *vw)
21401                .arg(&kdk)
21402                .arg(&kdv)
21403                .arg(&tkvi)
21404                .arg(&ktb)
21405                .arg(&vtb);
21406            unsafe {
21407                b.launch(cfg)?;
21408            }
21409        }
21410        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
21411        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
21412        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
21413        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
21414        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
21415        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
21416        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
21417        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21418            .map(|v| v != "0")
21419            .unwrap_or(true);
21420        {
21421            let hd_sfx = fa_hd_suffix(head_dim)?;
21422            let f = self.func(&format!(
21423                "fa_prefill_qw{}{hd_sfx}",
21424                if db { "_db" } else { "" }
21425            ));
21426            let shmem = if db {
21427                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
21428                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21429            } else {
21430                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21431            };
21432            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21433            f.set_attribute(
21434                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21435                shmem as i32,
21436            )?;
21437            let cfg = LaunchConfig {
21438                grid_dim: (
21439                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21440                    n_head as u32,
21441                    1,
21442                ),
21443                block_dim: (32, 4, 1),
21444                shared_mem_bytes: shmem,
21445            };
21446            let (hd, nh, nhkv, ti, tkvi, cz) = (
21447                head_dim as i32,
21448                n_head as i32,
21449                n_head_kv as i32,
21450                t as i32,
21451                t_kv as i32,
21452                causal as i32,
21453            );
21454            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21455            let __s_b = self.gpu.stream();
21456            let mut b = __s_b.launch_builder(&f);
21457            b.arg(q)
21458                .arg(&*kw)
21459                .arg(&*vw)
21460                .arg(o)
21461                .arg(&hd)
21462                .arg(&nh)
21463                .arg(&nhkv)
21464                .arg(&ti)
21465                .arg(&tkvi)
21466                .arg(&scale)
21467                .arg(&cz)
21468                .arg(&kdk)
21469                .arg(&kdv);
21470            unsafe {
21471                b.launch(cfg)?;
21472            }
21473        }
21474        Ok(())
21475    }
21476
21477    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
21478    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
21479    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
21480    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
21481    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
21482    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
21483    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
21484    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
21485    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
21486    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
21487    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
21488    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
21489    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
21490    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
21491    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
21492    #[allow(clippy::too_many_arguments)]
21493    pub fn fa_prefill_view_ws_w_hd128(
21494        &self,
21495        q: &CudaSlice<f32>,
21496        k: &cudarc::driver::CudaView<u8>,
21497        v: &cudarc::driver::CudaView<u8>,
21498        o: &mut CudaSlice<f32>,
21499        head_dim: usize,
21500        n_head: usize,
21501        n_head_kv: usize,
21502        t: usize,
21503        t_kv: usize,
21504        scale: f32,
21505        causal: bool,
21506        window: usize,
21507        k_tok_bytes: usize,
21508        v_tok_bytes: usize,
21509    ) -> Result<(), Box<dyn std::error::Error>> {
21510        assert_eq!(
21511            head_dim, 128,
21512            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21513        );
21514        if portable_mma_gated() {
21515            return self.sdpa_naive_w_quantized_view(
21516                q,
21517                k,
21518                v,
21519                o,
21520                head_dim,
21521                n_head,
21522                n_head_kv,
21523                t,
21524                t_kv,
21525                scale,
21526                causal,
21527                window,
21528                k_tok_bytes,
21529                v_tok_bytes,
21530            );
21531        }
21532        const BLOCK_Q: usize = 64;
21533        const BK: usize = 32;
21534        let kv_dim_k = n_head_kv * head_dim;
21535        let kv_dim_v = n_head_kv * head_dim;
21536        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
21537        let v_ws_bytes = t_kv * kv_dim_v * 2;
21538        let mut guard = self.prime_deqw_ws.lock().unwrap();
21539        let need_grow = match guard.as_ref() {
21540            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21541            None => true,
21542        };
21543        if need_grow {
21544            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21545            let (ck, cv) = guard
21546                .as_ref()
21547                .map(|(a, b)| (a.len(), b.len()))
21548                .unwrap_or((0, 0));
21549            *guard = Some((
21550                self.alloc_u8(grow(ck, k_ws_bytes))?,
21551                self.alloc_u8(grow(cv, v_ws_bytes))?,
21552            ));
21553        }
21554        let (kw, vw) = guard.as_mut().unwrap();
21555        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
21556        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
21557        {
21558            let f = self.func("fa_dequant_kv_ws_bf16");
21559            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21560            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21561            let cfg = LaunchConfig {
21562                grid_dim: (nblk.max(1), 1, 1),
21563                block_dim: (256, 1, 1),
21564                shared_mem_bytes: 0,
21565            };
21566            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21567            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21568            let __s_b = self.gpu.stream();
21569            let mut b = __s_b.launch_builder(&f);
21570            b.arg(k)
21571                .arg(v)
21572                .arg(&mut *kw)
21573                .arg(&mut *vw)
21574                .arg(&kdk)
21575                .arg(&kdv)
21576                .arg(&tkvi)
21577                .arg(&ktb)
21578                .arg(&vtb);
21579            unsafe {
21580                b.launch(cfg)?;
21581            }
21582        }
21583        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
21584        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21585            .map(|v| v != "0")
21586            .unwrap_or(true);
21587        {
21588            let f = self.func(if db {
21589                "fa_prefill_qw_db_w_hd128"
21590            } else {
21591                "fa_prefill_qw_w_hd128"
21592            });
21593            let shmem = if db {
21594                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21595            } else {
21596                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21597            };
21598            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21599            f.set_attribute(
21600                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21601                shmem as i32,
21602            )?;
21603            let cfg = LaunchConfig {
21604                grid_dim: (
21605                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21606                    n_head as u32,
21607                    1,
21608                ),
21609                block_dim: (32, 4, 1),
21610                shared_mem_bytes: shmem,
21611            };
21612            let (hd, nh, nhkv, ti, tkvi, cz) = (
21613                head_dim as i32,
21614                n_head as i32,
21615                n_head_kv as i32,
21616                t as i32,
21617                t_kv as i32,
21618                causal as i32,
21619            );
21620            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21621            let __s_b = self.gpu.stream();
21622            let mut b = __s_b.launch_builder(&f);
21623            b.arg(q)
21624                .arg(&*kw)
21625                .arg(&*vw)
21626                .arg(o)
21627                .arg(&hd)
21628                .arg(&nh)
21629                .arg(&nhkv)
21630                .arg(&ti)
21631                .arg(&tkvi)
21632                .arg(&scale)
21633                .arg(&cz)
21634                .arg(&kdk)
21635                .arg(&kdv)
21636                .arg(&wnd);
21637            unsafe {
21638                b.launch(cfg)?;
21639            }
21640        }
21641        Ok(())
21642    }
21643
21644    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
21645    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
21646    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
21647    pub fn fa_decode(
21648        &self,
21649        q: &CudaSlice<f32>,
21650        k: &cudarc::driver::CudaView<u8>,
21651        v: &cudarc::driver::CudaView<u8>,
21652        o: &mut CudaSlice<f32>,
21653        head_dim: usize,
21654        n_head: usize,
21655        n_head_kv: usize,
21656        t_kv: usize,
21657        scale: f32,
21658        k_tok_bytes: usize,
21659        v_tok_bytes: usize,
21660    ) -> Result<(), Box<dyn std::error::Error>> {
21661        self.fa_decode_kvmod(
21662            q,
21663            k,
21664            v,
21665            o,
21666            head_dim,
21667            n_head,
21668            n_head_kv,
21669            t_kv,
21670            scale,
21671            k_tok_bytes,
21672            v_tok_bytes,
21673            false,
21674        )
21675    }
21676
21677    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
21678    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
21679    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
21680    #[allow(clippy::too_many_arguments)]
21681    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
21682    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
21683    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
21684    #[allow(clippy::too_many_arguments)]
21685    #[allow(clippy::too_many_arguments)]
21686    fn fa_decode_scalar_unified(
21687        &self,
21688        q: &cudarc::driver::CudaView<f32>,
21689        k: &cudarc::driver::CudaView<u8>,
21690        v: &cudarc::driver::CudaView<u8>,
21691        o: &mut cudarc::driver::CudaViewMut<f32>,
21692        head_dim: usize,
21693        n_head: usize,
21694        n_head_kv: usize,
21695        t_kv_host: usize,
21696        t_kv_dev: Option<&CudaSlice<i32>>,
21697        scale: f32,
21698        n_splits: usize,
21699        split_keys: usize,
21700        k_tok_bytes: usize,
21701        v_tok_bytes: usize,
21702        g: bool,
21703        part_o: &mut CudaSlice<f32>,
21704        part_m: &mut CudaSlice<f32>,
21705        part_l: &mut CudaSlice<f32>,
21706        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21707    ) -> Result<(), Box<dyn std::error::Error>> {
21708        let f = if g {
21709            self.func_g("fa_decode_f32")
21710        } else {
21711            self.fa_func("fa_decode_f32", head_dim)
21712        };
21713        let cfg = LaunchConfig {
21714            grid_dim: (n_head as u32, n_splits as u32, 1),
21715            block_dim: (head_dim as u32, 1, 1),
21716            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21717        };
21718        let (hd, nh, nhkv, nsp) = (
21719            head_dim as i32,
21720            n_head as i32,
21721            n_head_kv as i32,
21722            n_splits as i32,
21723        );
21724        let (ktb, vtb, tkvi, ski) = (
21725            k_tok_bytes as i64,
21726            v_tok_bytes as i64,
21727            t_kv_host as i32,
21728            split_keys as i32,
21729        );
21730        let __s_b = self.gpu.stream();
21731        let mut b = __s_b.launch_builder(&f);
21732        match t_kv_dev {
21733            Some(d) => {
21734                b.arg(q)
21735                    .arg(k)
21736                    .arg(v)
21737                    .arg(&mut *part_o)
21738                    .arg(&mut *part_m)
21739                    .arg(&mut *part_l)
21740                    .arg(&hd)
21741                    .arg(&nh)
21742                    .arg(&nhkv)
21743                    .arg(&tkvi)
21744                    .arg(d)
21745                    .arg(&scale)
21746                    .arg(&nsp)
21747                    .arg(&ski)
21748                    .arg(&ktb)
21749                    .arg(&vtb);
21750                unsafe {
21751                    b.launch(cfg)?;
21752                }
21753            }
21754            None => {
21755                let null: u64 = 0;
21756                b.arg(q)
21757                    .arg(k)
21758                    .arg(v)
21759                    .arg(&mut *part_o)
21760                    .arg(&mut *part_m)
21761                    .arg(&mut *part_l)
21762                    .arg(&hd)
21763                    .arg(&nh)
21764                    .arg(&nhkv)
21765                    .arg(&tkvi)
21766                    .arg(&null)
21767                    .arg(&scale)
21768                    .arg(&nsp)
21769                    .arg(&ski)
21770                    .arg(&ktb)
21771                    .arg(&vtb);
21772                unsafe {
21773                    b.launch(cfg)?;
21774                }
21775            }
21776        }
21777        let cfg2 = LaunchConfig {
21778            grid_dim: (n_head as u32, 1, 1),
21779            block_dim: (head_dim as u32, 1, 1),
21780            shared_mem_bytes: 0,
21781        };
21782        if let Some((oq, od)) = q8_out {
21783            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
21784            let fc = if g {
21785                self.func_g("fa_decode_combine_q8_1")
21786            } else {
21787                self.fa_func("fa_decode_combine_q8_1", head_dim)
21788            };
21789            let __s_b2 = self.gpu.stream();
21790            let mut b2 = __s_b2.launch_builder(&fc);
21791            b2.arg(&*part_o)
21792                .arg(&*part_m)
21793                .arg(&*part_l)
21794                .arg(oq)
21795                .arg(od)
21796                .arg(&hd)
21797                .arg(&nh)
21798                .arg(&nsp);
21799            unsafe {
21800                b2.launch(cfg2)?;
21801            }
21802            return Ok(());
21803        }
21804        let fc = if g {
21805            self.func_g("fa_decode_combine_f32")
21806        } else {
21807            self.fa_func("fa_decode_combine_f32", head_dim)
21808        };
21809        let __s_b2 = self.gpu.stream();
21810        let mut b2 = __s_b2.launch_builder(&fc);
21811        b2.arg(&*part_o)
21812            .arg(&*part_m)
21813            .arg(&*part_l)
21814            .arg(o)
21815            .arg(&hd)
21816            .arg(&nh)
21817            .arg(&nsp);
21818        unsafe {
21819            b2.launch(cfg2)?;
21820        }
21821        Ok(())
21822    }
21823
21824    pub fn fa_decode_kvmod(
21825        &self,
21826        q: &CudaSlice<f32>,
21827        k: &cudarc::driver::CudaView<u8>,
21828        v: &cudarc::driver::CudaView<u8>,
21829        o: &mut CudaSlice<f32>,
21830        head_dim: usize,
21831        n_head: usize,
21832        n_head_kv: usize,
21833        t_kv: usize,
21834        scale: f32,
21835        k_tok_bytes: usize,
21836        v_tok_bytes: usize,
21837        g: bool,
21838    ) -> Result<(), Box<dyn std::error::Error>> {
21839        let q_view = q.as_view();
21840        let mut o_view = o.as_view_mut();
21841        self.fa_decode_kvmod_view(
21842            &q_view,
21843            k,
21844            v,
21845            &mut o_view,
21846            head_dim,
21847            n_head,
21848            n_head_kv,
21849            t_kv,
21850            scale,
21851            k_tok_bytes,
21852            v_tok_bytes,
21853            g,
21854        )
21855    }
21856
21857    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
21858    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
21859    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
21860    /// per-session KV view and FA launch.
21861    #[allow(clippy::too_many_arguments)]
21862    pub fn fa_decode_kvmod_view(
21863        &self,
21864        q: &cudarc::driver::CudaView<f32>,
21865        k: &cudarc::driver::CudaView<u8>,
21866        v: &cudarc::driver::CudaView<u8>,
21867        o: &mut cudarc::driver::CudaViewMut<f32>,
21868        head_dim: usize,
21869        n_head: usize,
21870        n_head_kv: usize,
21871        t_kv: usize,
21872        scale: f32,
21873        k_tok_bytes: usize,
21874        v_tok_bytes: usize,
21875        g: bool,
21876    ) -> Result<(), Box<dyn std::error::Error>> {
21877        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
21878        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
21879        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
21880        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
21881        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
21882        //
21883        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
21884        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
21885        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
21886        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
21887        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
21888        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
21889        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
21890        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
21891        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
21892        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
21893        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
21894        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
21895        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
21896        // fall to the exact scalar there instead of the broken register arm.
21897        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21898        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
21899        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
21900        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
21901        if g && head_dim == 256 && !fa_v4_at(t_kv) {
21902            fa_vec = false;
21903        }
21904        let sp = fa_split_keys(t_kv, n_head_kv);
21905        let n_splits = if fa_vec {
21906            ((t_kv + sp - 1) / sp).max(1)
21907        } else {
21908            ((t_kv + 255) / 256).max(1)
21909        };
21910        let o_len = n_head * n_splits * head_dim;
21911        let ml_len = n_head * n_splits;
21912        let mut part_guard = self.fa_part_pool.lock().unwrap();
21913        if part_guard
21914            .as_ref()
21915            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21916            .unwrap_or(true)
21917        {
21918            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21919            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21920            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21921            // later live allocations land at those addresses, and the next graph REPLAY writes
21922            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21923            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21924            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21925            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21926            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21927            // (total retired < final size).
21928            let old = part_guard.take();
21929            let (co, cm) = old
21930                .as_ref()
21931                .map(|pp| (pp.0.len(), pp.1.len()))
21932                .unwrap_or((0, 0));
21933            if let Some(old) = old {
21934                self.fa_part_retired.lock().unwrap().push(old);
21935            }
21936            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21937                eprintln!(
21938                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21939                    co, o_len, cm, ml_len
21940                );
21941            }
21942            *part_guard = Some((
21943                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21944                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21945                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21946            ));
21947        }
21948        let pg = part_guard.as_mut().unwrap();
21949        self.gpu
21950            .stream()
21951            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21952        self.gpu
21953            .stream()
21954            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21955        self.gpu
21956            .stream()
21957            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21958        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21959        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21960        let (hd, nh, nhkv, tkvi, nsp) = (
21961            head_dim as i32,
21962            n_head as i32,
21963            n_head_kv as i32,
21964            t_kv as i32,
21965            n_splits as i32,
21966        );
21967        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21968        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
21969        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
21970        // silently truncating the accumulator.
21971        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
21972        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
21973        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
21974        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
21975        // 178.4 -> 173.7 when 512 rode vec unconditionally).
21976        let fa512_min = fa512_min_tkv();
21977        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
21978        // g-module keeps the v4 pick (its class is not the depth-decay class).
21979        let deep = fa_vec
21980            && head_dim == 256
21981            && fa_v4_at(t_kv)
21982            && !g
21983            && fa_deep_at(t_kv)
21984            && !matches!(fa_v4_mode(), "noB3" | "stage");
21985        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21986            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
21987            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
21988            let gqa = (n_head / n_head_kv).max(1) as u32;
21989            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21990            (
21991                fv,
21992                LaunchConfig {
21993                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21994                    block_dim: (32, gqa, 1),
21995                    shared_mem_bytes: 0,
21996                },
21997            )
21998        } else if fa_vec && head_dim <= 256 {
21999            let gqa = (n_head / n_head_kv).max(1) as u32;
22000            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
22001            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
22002            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
22003            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
22004            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
22005            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
22006            // dequant each tile ONCE per block.
22007            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
22008            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
22009            // there by 12x — latency, not bandwidth, rules small KV).
22010            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22011            let smem_tkv = *SMEM_TKV.get_or_init(|| {
22012                std::env::var("MEMRA_FA_SMEM_TKV")
22013                    .ok()
22014                    .and_then(|v| v.parse().ok())
22015                    .unwrap_or_else(|| {
22016                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22017                    })
22018            });
22019            if fa_v4_at(t_kv) && head_dim == 256 {
22020                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
22021                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
22022                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
22023                let v4name = match fa_v4_mode() {
22024                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
22025                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
22026                    _ if deep => "fa_decode_vec_q_v4_deep",
22027                    _ => "fa_decode_vec_q_v4",
22028                };
22029                let fv = if g {
22030                    self.func_g(v4name)
22031                } else {
22032                    self.func(v4name)
22033                };
22034                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
22035                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
22036                let shmem = (if deep { 12160 } else { 11520 }
22037                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22038                use cudarc::driver::sys::CUfunction_attribute_enum as A;
22039                fv.set_attribute(
22040                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22041                    shmem as i32,
22042                )?;
22043                (
22044                    fv,
22045                    LaunchConfig {
22046                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22047                        block_dim: (32, gqa, 1),
22048                        shared_mem_bytes: shmem,
22049                    },
22050                )
22051            } else if fa_v3_active(head_dim) {
22052                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
22053                // smem = sV only (half of v2's).
22054                let fv = if g {
22055                    self.func_g("fa_decode_vec_q_v3")
22056                } else {
22057                    self.func("fa_decode_vec_q_v3")
22058                };
22059                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
22060                (
22061                    fv,
22062                    LaunchConfig {
22063                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22064                        block_dim: (32, gqa, 1),
22065                        shared_mem_bytes: shmem,
22066                    },
22067                )
22068            } else if fa_v2_on() {
22069                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
22070                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
22071                // partials; same 32KB sK+sV tile as the smem twin.
22072                let fv = if g {
22073                    self.func_g("fa_decode_vec_q_v2")
22074                } else {
22075                    self.func("fa_decode_vec_q_v2")
22076                };
22077                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22078                (
22079                    fv,
22080                    LaunchConfig {
22081                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22082                        block_dim: (32, gqa, 1),
22083                        shared_mem_bytes: shmem,
22084                    },
22085                )
22086            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22087            {
22088                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
22089                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
22090                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
22091                let fv = if g {
22092                    self.func_g("fa_decode_vec_q_smem")
22093                } else {
22094                    self.func("fa_decode_vec_q_smem")
22095                };
22096                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22097                use cudarc::driver::sys::CUfunction_attribute_enum as A;
22098                fv.set_attribute(
22099                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22100                    shmem as i32,
22101                )?;
22102                (
22103                    fv,
22104                    LaunchConfig {
22105                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22106                        block_dim: (32, gqa, 1),
22107                        shared_mem_bytes: shmem,
22108                    },
22109                )
22110            } else {
22111                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
22112                // dequant, zero dynamic shared memory.
22113                let fv = if g {
22114                    self.func_g("fa_decode_vec_q")
22115                } else {
22116                    self.func("fa_decode_vec_q")
22117                };
22118                (
22119                    fv,
22120                    LaunchConfig {
22121                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22122                        block_dim: (32, gqa, 1),
22123                        shared_mem_bytes: 0,
22124                    },
22125                )
22126            }
22127        } else {
22128            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
22129            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
22130            return self.fa_decode_scalar_unified(
22131                q,
22132                k,
22133                v,
22134                o,
22135                head_dim,
22136                n_head,
22137                n_head_kv,
22138                t_kv,
22139                None,
22140                scale,
22141                n_splits,
22142                if fa_vec { sp } else { 256 },
22143                k_tok_bytes,
22144                v_tok_bytes,
22145                g,
22146                part_o,
22147                part_m,
22148                part_l,
22149                None,
22150            );
22151        };
22152        let __s_b = self.gpu.stream();
22153        let mut b = __s_b.launch_builder(&f);
22154        b.arg(q)
22155            .arg(k)
22156            .arg(v)
22157            .arg(&mut *part_o)
22158            .arg(&mut *part_m)
22159            .arg(&mut *part_l)
22160            .arg(&hd)
22161            .arg(&nh)
22162            .arg(&nhkv)
22163            .arg(&tkvi)
22164            .arg(&scale)
22165            .arg(&nsp)
22166            .arg(&ktb)
22167            .arg(&vtb);
22168        unsafe {
22169            b.launch(cfg)?;
22170        }
22171        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
22172        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
22173        let (fc, cfg2) = (
22174            if g {
22175                self.func_g("fa_decode_combine_f32")
22176            } else {
22177                self.fa_func("fa_decode_combine_f32", head_dim)
22178            },
22179            LaunchConfig {
22180                grid_dim: (n_head as u32, 1, 1),
22181                block_dim: (head_dim as u32, 1, 1),
22182                shared_mem_bytes: 0,
22183            },
22184        );
22185        let __s_b2 = self.gpu.stream();
22186        let mut b2 = __s_b2.launch_builder(&fc);
22187        b2.arg(&*part_o)
22188            .arg(&*part_m)
22189            .arg(&*part_l)
22190            .arg(o)
22191            .arg(&hd)
22192            .arg(&nh)
22193            .arg(&nsp);
22194        unsafe {
22195            b2.launch(cfg2)?;
22196        }
22197        Ok(())
22198    }
22199
22200    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
22201    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
22202    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
22203    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
22204    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
22205    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
22206    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
22207    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
22208    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
22209    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
22210    #[allow(clippy::too_many_arguments)]
22211    pub fn fa_decode_batch_seqs_v4(
22212        &self,
22213        q: &CudaSlice<f32>,
22214        kv_ptrs: &cudarc::driver::CudaView<u64>,
22215        pos_seq: &CudaSlice<i32>,
22216        o: &mut CudaSlice<f32>,
22217        head_dim: usize,
22218        n_head: usize,
22219        n_head_kv: usize,
22220        b_n: usize,
22221        t_kv_max: usize,
22222        scale: f32,
22223        split_keys: usize,
22224        k_tok_bytes: usize,
22225        v_tok_bytes: usize,
22226    ) -> Result<(), Box<dyn std::error::Error>> {
22227        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22228        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22229        let o_len = b_n * n_head * n_splits_max * head_dim;
22230        let ml_len = b_n * n_head * n_splits_max;
22231        let mut part_guard = self.fa_part_pool.lock().unwrap();
22232        if part_guard
22233            .as_ref()
22234            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22235            .unwrap_or(true)
22236        {
22237            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22238            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22239            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22240            // later live allocations land at those addresses, and the next graph REPLAY writes
22241            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22242            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22243            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22244            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22245            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22246            // (total retired < final size).
22247            let old = part_guard.take();
22248            let (co, cm) = old
22249                .as_ref()
22250                .map(|pp| (pp.0.len(), pp.1.len()))
22251                .unwrap_or((0, 0));
22252            if let Some(old) = old {
22253                self.fa_part_retired.lock().unwrap().push(old);
22254            }
22255            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22256                eprintln!(
22257                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22258                    co, o_len, cm, ml_len
22259                );
22260            }
22261            *part_guard = Some((
22262                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22263                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22264                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22265            ));
22266        }
22267        let pg = part_guard.as_mut().unwrap();
22268        self.gpu
22269            .stream()
22270            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22271        self.gpu
22272            .stream()
22273            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22274        self.gpu
22275            .stream()
22276            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22277        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22278        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22279        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22280        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22281        let gqa = (n_head / n_head_kv).max(1) as u32;
22282        let f = self.func("fa_decode_vec_q_seqs_v4");
22283        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
22284        let shmem = (11520 + 32 * head_dim * 2) as u32;
22285        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22286        f.set_attribute(
22287            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22288            shmem as i32,
22289        )?;
22290        let cfg = LaunchConfig {
22291            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22292            block_dim: (32, gqa, 1),
22293            shared_mem_bytes: shmem,
22294        };
22295        {
22296            let __s_b = self.gpu.stream();
22297            let mut b = __s_b.launch_builder(&f);
22298            b.arg(q)
22299                .arg(kv_ptrs)
22300                .arg(pos_seq)
22301                .arg(&mut *part_o)
22302                .arg(&mut *part_m)
22303                .arg(&mut *part_l)
22304                .arg(&hd)
22305                .arg(&nh)
22306                .arg(&nhkv)
22307                .arg(&scale)
22308                .arg(&nspm)
22309                .arg(&spk)
22310                .arg(&ktb)
22311                .arg(&vtb);
22312            unsafe {
22313                b.launch(cfg)?;
22314            }
22315        }
22316        let fc = self.func("fa_decode_combine_seqs");
22317        let cfg2 = LaunchConfig {
22318            grid_dim: (n_head as u32, b_n as u32, 1),
22319            block_dim: (head_dim as u32, 1, 1),
22320            shared_mem_bytes: 0,
22321        };
22322        let __s_b2 = self.gpu.stream();
22323        let mut b2 = __s_b2.launch_builder(&fc);
22324        b2.arg(&*part_o)
22325            .arg(&*part_m)
22326            .arg(&*part_l)
22327            .arg(o)
22328            .arg(&hd)
22329            .arg(&nh)
22330            .arg(pos_seq)
22331            .arg(&nspm)
22332            .arg(&spk);
22333        unsafe {
22334            b2.launch(cfg2)?;
22335        }
22336        Ok(())
22337    }
22338
22339    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
22340    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
22341    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
22342    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
22343    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
22344    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
22345    #[allow(clippy::too_many_arguments)]
22346    pub fn append_kv_quantized_seqs(
22347        &self,
22348        k_rows: &CudaSlice<f32>,
22349        v_rows: &CudaSlice<f32>,
22350        kv_ptrs: &cudarc::driver::CudaView<u64>,
22351        pos_seq: &CudaSlice<i32>,
22352        b_n: usize,
22353        kv_dim_k: usize,
22354        kv_dim_v: usize,
22355        k_tok_bytes: usize,
22356        v_tok_bytes: usize,
22357    ) -> Result<(), Box<dyn std::error::Error>> {
22358        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22359        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22360        let cfg = LaunchConfig {
22361            grid_dim: (nblk, b_n as u32, 1),
22362            block_dim: (32, 1, 1),
22363            shared_mem_bytes: 0,
22364        };
22365        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22366        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22367        let __s_b = self.gpu.stream();
22368        let mut b = __s_b.launch_builder(&f);
22369        b.arg(k_rows)
22370            .arg(v_rows)
22371            .arg(kv_ptrs)
22372            .arg(pos_seq)
22373            .arg(&kdk)
22374            .arg(&kdv)
22375            .arg(&ktb)
22376            .arg(&vtb);
22377        unsafe {
22378            b.launch(cfg)?;
22379        }
22380        Ok(())
22381    }
22382
22383    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
22384    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
22385    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
22386    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
22387    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
22388    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22389        std::env::var("MEMRA_NO_FA_VEC").is_err()
22390            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22391            && base_len + 1 >= fa_vec_min_tkv()
22392            && head_dim <= 256
22393            && head_dim % 32 == 0
22394    }
22395
22396    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
22397    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
22398    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
22399    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
22400    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
22401    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
22402    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
22403    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
22404    #[allow(clippy::too_many_arguments)]
22405    pub fn fa_decode_rows(
22406        &self,
22407        q: &CudaSlice<f32>,
22408        k: &cudarc::driver::CudaView<u8>,
22409        v: &cudarc::driver::CudaView<u8>,
22410        o: &mut CudaSlice<f32>,
22411        head_dim: usize,
22412        n_head: usize,
22413        n_head_kv: usize,
22414        base_len: usize,
22415        t: usize,
22416        scale: f32,
22417        k_tok_bytes: usize,
22418        v_tok_bytes: usize,
22419        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
22420        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
22421        // keep the host arg. None is a bug for hd512 (asserted below).
22422        base_dev: Option<(&CudaSlice<i32>, i32)>,
22423        // K and V planes hold the same values (gemma globals, wv:=wk): pick
22424        // the _kv twin — V plane never read, value rides the q8_0 key dq.
22425        kv_shared: bool,
22426        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
22427        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
22428        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
22429        g: bool,
22430        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
22431        // (hd512 path) — the standalone quantize launch folds away.
22432        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22433    ) -> Result<(), Box<dyn std::error::Error>> {
22434        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22435        let t_kv_max = base_len + t; // LAST row's key bound
22436        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
22437        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
22438        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
22439        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
22440        // (parity law), so the partition is freely tunable — verify and decode move together.
22441        if head_dim == 512 {
22442            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22443            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
22444            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
22445            let v = *SP512.get_or_init(|| {
22446                std::env::var("MEMRA_FA_SP512")
22447                    .ok()
22448                    .and_then(|x| x.parse().ok())
22449                    .unwrap_or(0)
22450            });
22451            sp = if v >= 8 {
22452                v
22453            } else {
22454                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22455            };
22456        }
22457        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22458        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22459        let gqa = (n_head / n_head_kv).max(1) as u32;
22460        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
22461        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
22462        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
22463        // the different partition changes the combine's FP order (greedy tie flips at depth;
22464        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
22465        // consecutive rows by their OWN ladder value and launch once per group — each row then
22466        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
22467        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
22468        // sp override is t_kv-independent by construction).
22469        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
22470        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22471            groups.push((0, t, sp));
22472        } else {
22473            let mut r0 = 0usize;
22474            while r0 < t {
22475                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22476                let mut r1 = r0 + 1;
22477                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22478                    r1 += 1;
22479                }
22480                groups.push((r0, r1 - r0, sp_g));
22481                r0 = r1;
22482            }
22483        }
22484        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
22485        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
22486        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
22487        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22488        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22489            std::env::var("MEMRA_FA_SMEM_TKV")
22490                .ok()
22491                .and_then(|v| v.parse().ok())
22492                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22493        });
22494        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22495        let v3 = fa_v3_active(head_dim);
22496        let smem_rows =
22497            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22498        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
22499        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
22500        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
22501        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
22502        let _ = kv_shared;
22503        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
22504        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
22505        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22506        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
22507        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
22508        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
22509        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
22510        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
22511        // (kv_head, split) stages its tile once and loops the rows over it — kills the
22512        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
22513        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
22514        // shared by every hd512 caller through this wrapper (decode+verify flip together;
22515        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
22516        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
22517        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
22518        // not unpack-bound; jsonl 2026-07-14.
22519        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22520        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
22521        let tb512 = head_dim == 512
22522            && sp <= 32
22523            && n_head / n_head_kv.max(1) <= 16
22524            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22525        let fname = if tb512 {
22526            "fa_decode_vec_q_rows_v4_512_tb"
22527        } else if i2 {
22528            "fa_decode_vec_q_rows_dpl16_i2"
22529        } else if head_dim == 512 {
22530            "fa_decode_vec_q_rows_dpl16"
22531        }
22532        // gemma globals (parity law)
22533        else if v4 {
22534            "fa_decode_vec_q_rows_v4"
22535        } else if v3 {
22536            "fa_decode_vec_q_rows_v3"
22537        } else if fa_v2_on() {
22538            "fa_decode_vec_q_rows_v2"
22539        } else if smem_rows {
22540            "fa_decode_vec_q_rows_smem"
22541        } else {
22542            "fa_decode_vec_q_rows"
22543        };
22544        let f = if head_dim == 512 {
22545            self.fa_func(fname, head_dim)
22546        } else if g {
22547            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
22548            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
22549            // g-module rows against decode's g-module v4 — different programs, short-VG
22550            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
22551            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
22552            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
22553            // dq macros are format-aware.
22554            self.func_g(if smem_rows {
22555                "fa_decode_vec_q_rows"
22556            } else {
22557                fname
22558            })
22559        } else {
22560            self.func(fname)
22561        };
22562        let shmem = if tb512 {
22563            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
22564            let gk = Self::gkv_on();
22565            let sh =
22566                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22567            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22568            f.set_attribute(
22569                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22570                sh as i32,
22571            )?;
22572            sh
22573        } else if v4 || v3 || smem_rows || fa_v2_on() {
22574            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
22575            let sh = (if v4 {
22576                11520 + 32 * head_dim * if g { 1 } else { 2 }
22577            } else if v3 {
22578                32 * head_dim * 2
22579            } else {
22580                2 * 32 * head_dim * 2
22581            }) as u32;
22582            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22583            f.set_attribute(
22584                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22585                sh as i32,
22586            )?;
22587            sh
22588        } else {
22589            0
22590        };
22591        // Per-GROUP launches (single group in the common case — identical to the pre-fix
22592        // single launch there): each group gets its own partials (the rows kernel indexes
22593        // partials by its LOCAL grid.z row) and q/o row-offset views.
22594        for &(r0, t_g, sp_g) in &groups {
22595            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22596            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22597            let base_i = (base_len + r0) as i32;
22598            let o_len = t_g * n_head * n_splits_g * head_dim;
22599            let ml_len = t_g * n_head * n_splits_g;
22600            let mut part_guard = self.fa_part_pool.lock().unwrap();
22601            if part_guard
22602                .as_ref()
22603                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22604                .unwrap_or(true)
22605            {
22606                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22607                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22608                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22609                // later live allocations land at those addresses, and the next graph REPLAY writes
22610                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22611                // output corruption began the burst after the trunk's t_kv growth first realloc'd
22612                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22613                // the baked addresses alive (single-stream: eager writes the new buffers, replays
22614                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22615                // (total retired < final size).
22616                let old = part_guard.take();
22617                let (co, cm) = old
22618                    .as_ref()
22619                    .map(|pp| (pp.0.len(), pp.1.len()))
22620                    .unwrap_or((0, 0));
22621                if let Some(old) = old {
22622                    self.fa_part_retired.lock().unwrap().push(old);
22623                }
22624                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22625                    eprintln!(
22626                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22627                        co, o_len, cm, ml_len
22628                    );
22629                }
22630                *part_guard = Some((
22631                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22632                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22633                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22634                ));
22635            }
22636            let pg = part_guard.as_mut().unwrap();
22637            self.gpu
22638                .stream()
22639                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22640            self.gpu
22641                .stream()
22642                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22643            self.gpu
22644                .stream()
22645                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22646            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22647            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22648            let qv = self.view(q, t * n_head * head_dim);
22649            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22650            let cfg = LaunchConfig {
22651                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22652                block_dim: (32, gqa, 1),
22653                shared_mem_bytes: shmem,
22654            };
22655            {
22656                let __s_b = self.gpu.stream();
22657                let mut b = __s_b.launch_builder(&f);
22658                if tb512 {
22659                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
22660                    let (bd, plus) =
22661                        base_dev.expect("hd512 rows twin requires a device base counter");
22662                    let plus_g = plus + r0 as i32;
22663                    let nr = t_g as i32;
22664                    if Self::pdl_on() && Self::pdl_wb_on() {
22665                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
22666                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22667                        let s = &self.gpu.stream();
22668                        let (pq, _b0) = q_g.device_ptr(s);
22669                        let (pk, _b1) = k.device_ptr(s);
22670                        let (pv, _b2) = v.device_ptr(s);
22671                        let (po, _b3) = part_o.device_ptr_mut(s);
22672                        let (pm, _b4) = part_m.device_ptr_mut(s);
22673                        let (pl, _b5) = part_l.device_ptr_mut(s);
22674                        let (pb, _b6) = bd.device_ptr(s);
22675                        let mut ps = [
22676                            &pq as *const _ as *mut std::ffi::c_void,
22677                            &pk as *const _ as *mut _,
22678                            &pv as *const _ as *mut _,
22679                            &po as *const _ as *mut _,
22680                            &pm as *const _ as *mut _,
22681                            &pl as *const _ as *mut _,
22682                            &hd as *const _ as *mut _,
22683                            &nh as *const _ as *mut _,
22684                            &nhkv as *const _ as *mut _,
22685                            &pb as *const _ as *mut _,
22686                            &plus_g as *const _ as *mut _,
22687                            &scale as *const _ as *mut _,
22688                            &nspm as *const _ as *mut _,
22689                            &spk as *const _ as *mut _,
22690                            &ktb as *const _ as *mut _,
22691                            &vtb as *const _ as *mut _,
22692                            &nr as *const _ as *mut _,
22693                        ];
22694                        unsafe {
22695                            self.launch_pdl_flash(
22696                                Self::gkv_on(),
22697                                "fa_decode_vec_q_rows_v4_512_tb",
22698                                (n_head_kv as u32, n_splits_g as u32, 1),
22699                                (32, gqa, 1),
22700                                shmem,
22701                                &mut ps,
22702                            )?;
22703                        }
22704                    } else {
22705                        let cfg_tb = LaunchConfig {
22706                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22707                            block_dim: (32, gqa, 1),
22708                            shared_mem_bytes: shmem,
22709                        };
22710                        b.arg(&q_g)
22711                            .arg(k)
22712                            .arg(v)
22713                            .arg(&mut *part_o)
22714                            .arg(&mut *part_m)
22715                            .arg(&mut *part_l)
22716                            .arg(&hd)
22717                            .arg(&nh)
22718                            .arg(&nhkv)
22719                            .arg(bd)
22720                            .arg(&plus_g)
22721                            .arg(&scale)
22722                            .arg(&nspm)
22723                            .arg(&spk)
22724                            .arg(&ktb)
22725                            .arg(&vtb)
22726                            .arg(&nr);
22727                        unsafe {
22728                            b.launch(cfg_tb)?;
22729                        }
22730                    }
22731                } else if head_dim == 512 {
22732                    let (bd, plus) =
22733                        base_dev.expect("hd512 rows twin requires a device base counter");
22734                    let plus_g = plus + r0 as i32;
22735                    b.arg(&q_g)
22736                        .arg(k)
22737                        .arg(v)
22738                        .arg(&mut *part_o)
22739                        .arg(&mut *part_m)
22740                        .arg(&mut *part_l)
22741                        .arg(&hd)
22742                        .arg(&nh)
22743                        .arg(&nhkv)
22744                        .arg(bd)
22745                        .arg(&plus_g)
22746                        .arg(&scale)
22747                        .arg(&nspm)
22748                        .arg(&spk)
22749                        .arg(&ktb)
22750                        .arg(&vtb);
22751                    unsafe {
22752                        b.launch(cfg)?;
22753                    }
22754                } else {
22755                    b.arg(&q_g)
22756                        .arg(k)
22757                        .arg(v)
22758                        .arg(&mut *part_o)
22759                        .arg(&mut *part_m)
22760                        .arg(&mut *part_l)
22761                        .arg(&hd)
22762                        .arg(&nh)
22763                        .arg(&nhkv)
22764                        .arg(&base_i)
22765                        .arg(&scale)
22766                        .arg(&nspm)
22767                        .arg(&spk)
22768                        .arg(&ktb)
22769                        .arg(&vtb);
22770                    unsafe {
22771                        b.launch(cfg)?;
22772                    }
22773                }
22774            }
22775            let cfg2 = LaunchConfig {
22776                grid_dim: (n_head as u32, t_g as u32, 1),
22777                block_dim: (head_dim as u32, 1, 1),
22778                shared_mem_bytes: 0,
22779            };
22780            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22781            if head_dim == 512 {
22782                // device-len combine (shared by verify/eager/graph — parity by symbol): the
22783                // per-row n_splits derives from the SAME counter the rows kernel read.
22784                let (bd, plus) = base_dev.unwrap();
22785                let plus_g = plus + r0 as i32;
22786                if let Some((oq, od)) = q8_out.as_mut() {
22787                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
22788                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22789                    if Self::pdl_on() && Self::pdl_wb_on() {
22790                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
22791                        use cudarc::driver::{DevicePtr, DevicePtrMut};
22792                        let s = &self.gpu.stream();
22793                        let (po, _g0) = part_o.device_ptr(s);
22794                        let (pm, _g1) = part_m.device_ptr(s);
22795                        let (pl, _g2) = part_l.device_ptr(s);
22796                        let (pq, _g3) = oq.device_ptr_mut(s);
22797                        let (pd, _g4) = od.device_ptr_mut(s);
22798                        let (pb, _g5) = bd.device_ptr(s);
22799                        let mut ps = [
22800                            &po as *const _ as *mut std::ffi::c_void,
22801                            &pm as *const _ as *mut _,
22802                            &pl as *const _ as *mut _,
22803                            &pq as *const _ as *mut _,
22804                            &pd as *const _ as *mut _,
22805                            &hd as *const _ as *mut _,
22806                            &nh as *const _ as *mut _,
22807                            &pb as *const _ as *mut _,
22808                            &plus_g as *const _ as *mut _,
22809                            &nspm as *const _ as *mut _,
22810                            &spk as *const _ as *mut _,
22811                        ];
22812                        unsafe {
22813                            self.launch_pdl_flash(
22814                                Self::gkv_on(),
22815                                "fa_decode_combine_rows_dc_q8_1",
22816                                cfg2.grid_dim,
22817                                cfg2.block_dim,
22818                                0,
22819                                &mut ps,
22820                            )?;
22821                        }
22822                        continue;
22823                    }
22824                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22825                    let __s_b2 = self.gpu.stream();
22826                    let mut b2 = __s_b2.launch_builder(&fc);
22827                    b2.arg(&*part_o)
22828                        .arg(&*part_m)
22829                        .arg(&*part_l)
22830                        .arg(&mut **oq)
22831                        .arg(&mut **od)
22832                        .arg(&hd)
22833                        .arg(&nh)
22834                        .arg(bd)
22835                        .arg(&plus_g)
22836                        .arg(&nspm)
22837                        .arg(&spk);
22838                    unsafe {
22839                        b2.launch(cfg2)?;
22840                    }
22841                    continue;
22842                }
22843                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22844                let __s_b2 = self.gpu.stream();
22845                let mut b2 = __s_b2.launch_builder(&fc);
22846                b2.arg(&*part_o)
22847                    .arg(&*part_m)
22848                    .arg(&*part_l)
22849                    .arg(&mut o_g)
22850                    .arg(&hd)
22851                    .arg(&nh)
22852                    .arg(bd)
22853                    .arg(&plus_g)
22854                    .arg(&nspm)
22855                    .arg(&spk);
22856                unsafe {
22857                    b2.launch(cfg2)?;
22858                }
22859            } else {
22860                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
22861                // leave the caller's pair unwritten (consumer would read garbage).
22862                assert!(
22863                    q8_out.is_none(),
22864                    "rows q8 emit requires the hd512 dc combine"
22865                );
22866                let fc = self.func("fa_decode_combine_rows");
22867                let __s_b2 = self.gpu.stream();
22868                let mut b2 = __s_b2.launch_builder(&fc);
22869                b2.arg(&*part_o)
22870                    .arg(&*part_m)
22871                    .arg(&*part_l)
22872                    .arg(&mut o_g)
22873                    .arg(&hd)
22874                    .arg(&nh)
22875                    .arg(&base_i)
22876                    .arg(&nspm)
22877                    .arg(&spk);
22878                unsafe {
22879                    b2.launch(cfg2)?;
22880                }
22881            }
22882        }
22883        Ok(())
22884    }
22885
22886    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
22887    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
22888    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
22889    #[allow(clippy::too_many_arguments)]
22890    pub fn fa_decode_rows_w(
22891        &self,
22892        q: &CudaSlice<f32>,
22893        k: &cudarc::driver::CudaView<u8>,
22894        v: &cudarc::driver::CudaView<u8>,
22895        o: &mut CudaSlice<f32>,
22896        head_dim: usize,
22897        n_head: usize,
22898        n_head_kv: usize,
22899        base_dev: &CudaSlice<i32>,
22900        base_plus: i32,
22901        t: usize,
22902        scale: f32,
22903        window: usize,
22904        k_tok_bytes: usize,
22905        v_tok_bytes: usize,
22906        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22907    ) -> Result<(), Box<dyn std::error::Error>> {
22908        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
22909        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
22910        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
22911        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
22912        debug_assert!(head_dim == 256);
22913        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
22914        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
22915        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
22916        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
22917        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
22918        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
22919        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
22920        let sp = {
22921            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22922            let v = *SPW.get_or_init(|| {
22923                std::env::var("MEMRA_FA_SPW")
22924                    .ok()
22925                    .and_then(|x| x.parse().ok())
22926                    .unwrap_or(0)
22927            });
22928            if v >= 8 {
22929                v
22930            } else {
22931                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22932            }
22933        };
22934        let n_splits_max = (window + sp - 1) / sp;
22935        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22936        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22937        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22938        let gqa = (n_head / n_head_kv).max(1) as u32;
22939        let o_len = t * n_head * n_splits_max * head_dim;
22940        let ml_len = t * n_head * n_splits_max;
22941        let mut part_guard = self.fa_part_pool.lock().unwrap();
22942        if part_guard
22943            .as_ref()
22944            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22945            .unwrap_or(true)
22946        {
22947            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22948            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22949            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22950            // later live allocations land at those addresses, and the next graph REPLAY writes
22951            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22952            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22953            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22954            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22955            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22956            // (total retired < final size).
22957            let old = part_guard.take();
22958            let (co, cm) = old
22959                .as_ref()
22960                .map(|pp| (pp.0.len(), pp.1.len()))
22961                .unwrap_or((0, 0));
22962            if let Some(old) = old {
22963                self.fa_part_retired.lock().unwrap().push(old);
22964            }
22965            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22966                eprintln!(
22967                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22968                    co, o_len, cm, ml_len
22969                );
22970            }
22971            *part_guard = Some((
22972                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22973                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22974                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22975            ));
22976        }
22977        let pg = part_guard.as_mut().unwrap();
22978        self.gpu
22979            .stream()
22980            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22981        self.gpu
22982            .stream()
22983            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22984        self.gpu
22985            .stream()
22986            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22987        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22988        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
22989        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
22990        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
22991        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
22992        // floor (deep-ctx broadcast win); register twin between.
22993        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22994        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22995            std::env::var("MEMRA_FA_SMEM_TKV")
22996                .ok()
22997                .and_then(|v| v.parse().ok())
22998                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22999        });
23000        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
23001        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
23002        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
23003        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
23004        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
23005        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23006        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
23007        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
23008        // per (lane, format-module) keeps parity structural; the old register-i2 detour
23009        // (-33%) is retired.
23010        let wg = Self::wkv_on();
23011        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
23012        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
23013        let sp2 =
23014            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23015        if sp2 {
23016            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23017            if Self::pdl_on() && Self::pdl_wb_on() {
23018                // wave-B2b: flavor mirrors wg.
23019                use cudarc::driver::{DevicePtr, DevicePtrMut};
23020                let s = &self.gpu.stream();
23021                let (pq, _b0) = q.device_ptr(s);
23022                let (pk, _b1) = k.device_ptr(s);
23023                let (pv, _b2) = v.device_ptr(s);
23024                let (po, _b3) = part_o.device_ptr_mut(s);
23025                let (pm, _b4) = part_m.device_ptr_mut(s);
23026                let (pl, _b5) = part_l.device_ptr_mut(s);
23027                let (pb, _b6) = base_dev.device_ptr(s);
23028                let mut ps = [
23029                    &pq as *const _ as *mut std::ffi::c_void,
23030                    &pk as *const _ as *mut _,
23031                    &pv as *const _ as *mut _,
23032                    &po as *const _ as *mut _,
23033                    &pm as *const _ as *mut _,
23034                    &pl as *const _ as *mut _,
23035                    &hd as *const _ as *mut _,
23036                    &nh as *const _ as *mut _,
23037                    &nhkv as *const _ as *mut _,
23038                    &pb as *const _ as *mut _,
23039                    &base_plus as *const _ as *mut _,
23040                    &scale as *const _ as *mut _,
23041                    &nspm as *const _ as *mut _,
23042                    &spk as *const _ as *mut _,
23043                    &ktb as *const _ as *mut _,
23044                    &vtb as *const _ as *mut _,
23045                    &wini as *const _ as *mut _,
23046                ];
23047                unsafe {
23048                    self.launch_pdl_flash(
23049                        wg,
23050                        "fa_decode_vec_q_rows_v4_w_sp",
23051                        (n_head_kv as u32, n_splits_max as u32, t as u32),
23052                        (32, gqa + 1, 1),
23053                        sh,
23054                        &mut ps,
23055                    )?;
23056                }
23057            } else {
23058                let f = if wg {
23059                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23060                } else {
23061                    self.func("fa_decode_vec_q_rows_v4_w_sp")
23062                };
23063                f.set_attribute(
23064                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23065                    sh as i32,
23066                )?;
23067                let cfg = LaunchConfig {
23068                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23069                    block_dim: (32, gqa + 1, 1),
23070                    shared_mem_bytes: sh,
23071                };
23072                let __s_b = self.gpu.stream();
23073                let mut b = __s_b.launch_builder(&f);
23074                b.arg(q)
23075                    .arg(k)
23076                    .arg(v)
23077                    .arg(&mut *part_o)
23078                    .arg(&mut *part_m)
23079                    .arg(&mut *part_l)
23080                    .arg(&hd)
23081                    .arg(&nh)
23082                    .arg(&nhkv)
23083                    .arg(base_dev)
23084                    .arg(&base_plus)
23085                    .arg(&scale)
23086                    .arg(&nspm)
23087                    .arg(&spk)
23088                    .arg(&ktb)
23089                    .arg(&vtb)
23090                    .arg(&wini);
23091                unsafe {
23092                    b.launch(cfg)?;
23093                }
23094            }
23095        } else {
23096            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23097                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
23098                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23099                use cudarc::driver::{DevicePtr, DevicePtrMut};
23100                let s = &self.gpu.stream();
23101                let (pq, _b0) = q.device_ptr(s);
23102                let (pk, _b1) = k.device_ptr(s);
23103                let (pv, _b2) = v.device_ptr(s);
23104                let (po, _b3) = part_o.device_ptr_mut(s);
23105                let (pm, _b4) = part_m.device_ptr_mut(s);
23106                let (pl, _b5) = part_l.device_ptr_mut(s);
23107                let (pb, _b6) = base_dev.device_ptr(s);
23108                let mut ps = [
23109                    &pq as *const _ as *mut std::ffi::c_void,
23110                    &pk as *const _ as *mut _,
23111                    &pv as *const _ as *mut _,
23112                    &po as *const _ as *mut _,
23113                    &pm as *const _ as *mut _,
23114                    &pl as *const _ as *mut _,
23115                    &hd as *const _ as *mut _,
23116                    &nh as *const _ as *mut _,
23117                    &nhkv as *const _ as *mut _,
23118                    &pb as *const _ as *mut _,
23119                    &base_plus as *const _ as *mut _,
23120                    &scale as *const _ as *mut _,
23121                    &nspm as *const _ as *mut _,
23122                    &spk as *const _ as *mut _,
23123                    &ktb as *const _ as *mut _,
23124                    &vtb as *const _ as *mut _,
23125                    &wini as *const _ as *mut _,
23126                ];
23127                unsafe {
23128                    self.launch_pdl_flash(
23129                        wg,
23130                        "fa_decode_vec_q_rows_v4_w",
23131                        (n_head_kv as u32, n_splits_max as u32, t as u32),
23132                        (32, gqa, 1),
23133                        sh,
23134                        &mut ps,
23135                    )?;
23136                }
23137            } else {
23138                let pick = |name: &str| {
23139                    if wg {
23140                        self.func_g(name)
23141                    } else {
23142                        self.func(name)
23143                    }
23144                };
23145                let (f, sh) = if fa_v4_at(window) {
23146                    let f = pick("fa_decode_vec_q_rows_v4_w");
23147                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23148                } else if smem_tkv > 0 && window >= smem_tkv {
23149                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
23150                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
23151                    (
23152                        pick("fa_decode_vec_q_rows_smem_w"),
23153                        (2 * 32 * head_dim * 2) as u32,
23154                    )
23155                } else {
23156                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23157                };
23158                f.set_attribute(
23159                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23160                    sh as i32,
23161                )?;
23162                let cfg = LaunchConfig {
23163                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23164                    block_dim: (32, gqa, 1),
23165                    shared_mem_bytes: sh,
23166                };
23167                let __s_b = self.gpu.stream();
23168                let mut b = __s_b.launch_builder(&f);
23169                b.arg(q)
23170                    .arg(k)
23171                    .arg(v)
23172                    .arg(&mut *part_o)
23173                    .arg(&mut *part_m)
23174                    .arg(&mut *part_l)
23175                    .arg(&hd)
23176                    .arg(&nh)
23177                    .arg(&nhkv)
23178                    .arg(base_dev)
23179                    .arg(&base_plus)
23180                    .arg(&scale)
23181                    .arg(&nspm)
23182                    .arg(&spk)
23183                    .arg(&ktb)
23184                    .arg(&vtb)
23185                    .arg(&wini);
23186                unsafe {
23187                    b.launch(cfg)?;
23188                }
23189            }
23190        }
23191        let cfg2 = LaunchConfig {
23192            grid_dim: (n_head as u32, t as u32, 1),
23193            block_dim: (head_dim as u32, 1, 1),
23194            shared_mem_bytes: 0,
23195        };
23196        if let Some((oq, od)) = q8_out {
23197            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
23198            // consumes the pair directly; the standalone quantize launch folds away.
23199            if Self::pdl_on() && Self::pdl_wb_on() {
23200                // wave-B2: flavor mirrors the builder's wg choice.
23201                use cudarc::driver::{DevicePtr, DevicePtrMut};
23202                let s = &self.gpu.stream();
23203                let (po, _g0) = part_o.device_ptr(s);
23204                let (pm, _g1) = part_m.device_ptr(s);
23205                let (pl, _g2) = part_l.device_ptr(s);
23206                let (pq, _g3) = oq.device_ptr_mut(s);
23207                let (pd, _g4) = od.device_ptr_mut(s);
23208                let mut ps = [
23209                    &po as *const _ as *mut std::ffi::c_void,
23210                    &pm as *const _ as *mut _,
23211                    &pl as *const _ as *mut _,
23212                    &pq as *const _ as *mut _,
23213                    &pd as *const _ as *mut _,
23214                    &hd as *const _ as *mut _,
23215                    &nh as *const _ as *mut _,
23216                    &nspm as *const _ as *mut _,
23217                    &spk as *const _ as *mut _,
23218                    &wini as *const _ as *mut _,
23219                ];
23220                unsafe {
23221                    self.launch_pdl_flash(
23222                        wg,
23223                        "fa_decode_combine_rows_w_q8_1",
23224                        cfg2.grid_dim,
23225                        cfg2.block_dim,
23226                        0,
23227                        &mut ps,
23228                    )?;
23229                }
23230                return Ok(());
23231            }
23232            let fc = if wg {
23233                self.func_g("fa_decode_combine_rows_w_q8_1")
23234            } else {
23235                self.func("fa_decode_combine_rows_w_q8_1")
23236            };
23237            let __s_b2 = self.gpu.stream();
23238            let mut b2 = __s_b2.launch_builder(&fc);
23239            b2.arg(&*part_o)
23240                .arg(&*part_m)
23241                .arg(&*part_l)
23242                .arg(oq)
23243                .arg(od)
23244                .arg(&hd)
23245                .arg(&nh)
23246                .arg(&nspm)
23247                .arg(&spk)
23248                .arg(&wini);
23249            unsafe {
23250                b2.launch(cfg2)?;
23251            }
23252            return Ok(());
23253        }
23254        let fc = if wg {
23255            self.func_g("fa_decode_combine_rows_w")
23256        } else {
23257            self.func("fa_decode_combine_rows_w")
23258        };
23259        let __s_b2 = self.gpu.stream();
23260        let mut b2 = __s_b2.launch_builder(&fc);
23261        b2.arg(&*part_o)
23262            .arg(&*part_m)
23263            .arg(&*part_l)
23264            .arg(o)
23265            .arg(&hd)
23266            .arg(&nh)
23267            .arg(&nspm)
23268            .arg(&spk)
23269            .arg(&wini);
23270        unsafe {
23271            b2.launch(cfg2)?;
23272        }
23273        Ok(())
23274    }
23275
23276    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
23277    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
23278    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
23279    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
23280    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
23281    #[allow(clippy::too_many_arguments)]
23282    pub fn fa_decode_rows_dc(
23283        &self,
23284        q: &CudaSlice<f32>,
23285        k: &cudarc::driver::CudaView<u8>,
23286        v: &cudarc::driver::CudaView<u8>,
23287        o: &mut CudaSlice<f32>,
23288        head_dim: usize,
23289        n_head: usize,
23290        n_head_kv: usize,
23291        base_dev: &CudaSlice<i32>,
23292        t_kv_upper: usize,
23293        t: usize,
23294        scale: f32,
23295        k_tok_bytes: usize,
23296        v_tok_bytes: usize,
23297        base_plus: i32,
23298        g: bool,
23299    ) -> Result<(), Box<dyn std::error::Error>> {
23300        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23301        assert!(
23302            v4 || fa_v3_active(head_dim),
23303            "stream fa rows requires the v3 or v4 lane"
23304        );
23305        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23306        if v4 {
23307            let sp = fa_split_keys(t_kv_upper, n_head_kv);
23308            let n_splits_max = (t_kv_upper + sp - 1) / sp;
23309            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23310            let (nspm, spk) = (n_splits_max as i32, sp as i32);
23311            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23312            let gqa = (n_head / n_head_kv).max(1) as u32;
23313            let o_len = t * n_head * n_splits_max * head_dim;
23314            let ml_len = t * n_head * n_splits_max;
23315            let mut part_guard = self.fa_part_pool.lock().unwrap();
23316            if part_guard
23317                .as_ref()
23318                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23319                .unwrap_or(true)
23320            {
23321                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23322                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23323                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23324                // later live allocations land at those addresses, and the next graph REPLAY writes
23325                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23326                // output corruption began the burst after the trunk's t_kv growth first realloc'd
23327                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23328                // the baked addresses alive (single-stream: eager writes the new buffers, replays
23329                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23330                // (total retired < final size).
23331                let old = part_guard.take();
23332                let (co, cm) = old
23333                    .as_ref()
23334                    .map(|pp| (pp.0.len(), pp.1.len()))
23335                    .unwrap_or((0, 0));
23336                if let Some(old) = old {
23337                    self.fa_part_retired.lock().unwrap().push(old);
23338                }
23339                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23340                    eprintln!(
23341                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23342                        co, o_len, cm, ml_len
23343                    );
23344                }
23345                *part_guard = Some((
23346                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23347                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23348                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23349                ));
23350            }
23351            let pg = part_guard.as_mut().unwrap();
23352            self.gpu
23353                .stream()
23354                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23355            self.gpu
23356                .stream()
23357                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23358            self.gpu
23359                .stream()
23360                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23361            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23362            let f = if g {
23363                self.func_g("fa_decode_vec_q_rows_v4_dc")
23364            } else {
23365                self.func("fa_decode_vec_q_rows_v4_dc")
23366            };
23367            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23368            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23369            f.set_attribute(
23370                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23371                sh as i32,
23372            )?;
23373            let cfg = LaunchConfig {
23374                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23375                block_dim: (32, gqa, 1),
23376                shared_mem_bytes: sh,
23377            };
23378            let __s_b = self.gpu.stream();
23379            let mut b = __s_b.launch_builder(&f);
23380            b.arg(q)
23381                .arg(k)
23382                .arg(v)
23383                .arg(&mut *part_o)
23384                .arg(&mut *part_m)
23385                .arg(&mut *part_l)
23386                .arg(&hd)
23387                .arg(&nh)
23388                .arg(&nhkv)
23389                .arg(base_dev)
23390                .arg(&base_plus)
23391                .arg(&scale)
23392                .arg(&nspm)
23393                .arg(&spk)
23394                .arg(&ktb)
23395                .arg(&vtb);
23396            unsafe {
23397                b.launch(cfg)?;
23398            }
23399            let fc = self.func("fa_decode_combine_rows_dc");
23400            let cfg2 = LaunchConfig {
23401                grid_dim: (n_head as u32, t as u32, 1),
23402                block_dim: (head_dim as u32, 1, 1),
23403                shared_mem_bytes: 0,
23404            };
23405            let __s_b2 = self.gpu.stream();
23406            let mut b2 = __s_b2.launch_builder(&fc);
23407            b2.arg(&*part_o)
23408                .arg(&*part_m)
23409                .arg(&*part_l)
23410                .arg(o)
23411                .arg(&hd)
23412                .arg(&nh)
23413                .arg(base_dev)
23414                .arg(&base_plus)
23415                .arg(&nspm)
23416                .arg(&spk);
23417            unsafe {
23418                b2.launch(cfg2)?;
23419            }
23420            return Ok(());
23421        }
23422        let sp = fa_split_keys(t_kv_upper, n_head_kv);
23423        let n_splits_max = (t_kv_upper + sp - 1) / sp;
23424        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23425        let (nspm, spk) = (n_splits_max as i32, sp as i32);
23426        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23427        let gqa = (n_head / n_head_kv).max(1) as u32;
23428        let o_len = t * n_head * n_splits_max * head_dim;
23429        let ml_len = t * n_head * n_splits_max;
23430        let mut part_guard = self.fa_part_pool.lock().unwrap();
23431        if part_guard
23432            .as_ref()
23433            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23434            .unwrap_or(true)
23435        {
23436            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23437            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23438            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23439            // later live allocations land at those addresses, and the next graph REPLAY writes
23440            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23441            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23442            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23443            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23444            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23445            // (total retired < final size).
23446            let old = part_guard.take();
23447            let (co, cm) = old
23448                .as_ref()
23449                .map(|pp| (pp.0.len(), pp.1.len()))
23450                .unwrap_or((0, 0));
23451            if let Some(old) = old {
23452                self.fa_part_retired.lock().unwrap().push(old);
23453            }
23454            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23455                eprintln!(
23456                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23457                    co, o_len, cm, ml_len
23458                );
23459            }
23460            *part_guard = Some((
23461                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23462                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23463                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23464            ));
23465        }
23466        let pg = part_guard.as_mut().unwrap();
23467        self.gpu
23468            .stream()
23469            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23470        self.gpu
23471            .stream()
23472            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23473        self.gpu
23474            .stream()
23475            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23476        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23477        let f = self.func("fa_decode_vec_q_rows_v3_dc");
23478        let sh = (32 * head_dim * 2) as u32;
23479        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23480        f.set_attribute(
23481            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23482            sh as i32,
23483        )?;
23484        let cfg = LaunchConfig {
23485            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23486            block_dim: (32, gqa, 1),
23487            shared_mem_bytes: sh,
23488        };
23489        let __s_b = self.gpu.stream();
23490        let mut b = __s_b.launch_builder(&f);
23491        b.arg(q)
23492            .arg(k)
23493            .arg(v)
23494            .arg(&mut *part_o)
23495            .arg(&mut *part_m)
23496            .arg(&mut *part_l)
23497            .arg(&hd)
23498            .arg(&nh)
23499            .arg(&nhkv)
23500            .arg(base_dev)
23501            .arg(&scale)
23502            .arg(&nspm)
23503            .arg(&spk)
23504            .arg(&ktb)
23505            .arg(&vtb);
23506        unsafe {
23507            b.launch(cfg)?;
23508        }
23509        let fc = self.func("fa_decode_combine_rows_dc");
23510        let cfg2 = LaunchConfig {
23511            grid_dim: (n_head as u32, t as u32, 1),
23512            block_dim: (head_dim as u32, 1, 1),
23513            shared_mem_bytes: 0,
23514        };
23515        let plus0 = 0i32;
23516        let __s_b2 = self.gpu.stream();
23517        let mut b2 = __s_b2.launch_builder(&fc);
23518        b2.arg(&*part_o)
23519            .arg(&*part_m)
23520            .arg(&*part_l)
23521            .arg(o)
23522            .arg(&hd)
23523            .arg(&nh)
23524            .arg(base_dev)
23525            .arg(&plus0)
23526            .arg(&nspm)
23527            .arg(&spk);
23528        unsafe {
23529            b2.launch(cfg2)?;
23530        }
23531        Ok(())
23532    }
23533
23534    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
23535    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
23536    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
23537    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
23538    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
23539    ///
23540    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
23541    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
23542    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
23543    /// grouping (different but mathematically-equal log-sum-exp merge).
23544    pub fn fa_decode_dc(
23545        &self,
23546        q: &CudaSlice<f32>,
23547        k: &cudarc::driver::CudaView<u8>,
23548        v: &cudarc::driver::CudaView<u8>,
23549        o: &mut CudaSlice<f32>,
23550        head_dim: usize,
23551        n_head: usize,
23552        n_head_kv: usize,
23553        t_kv_dev: &CudaSlice<i32>,
23554        bucket_max: usize,
23555        scale: f32,
23556        k_tok_bytes: usize,
23557        v_tok_bytes: usize,
23558        g: bool,
23559    ) -> Result<(), Box<dyn std::error::Error>> {
23560        self.fa_decode_dc_q8(
23561            q,
23562            k,
23563            v,
23564            o,
23565            head_dim,
23566            n_head,
23567            n_head_kv,
23568            t_kv_dev,
23569            bucket_max,
23570            scale,
23571            k_tok_bytes,
23572            v_tok_bytes,
23573            g,
23574            None,
23575        )
23576    }
23577
23578    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
23579    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
23580    #[allow(clippy::too_many_arguments)]
23581    pub fn fa_decode_dc_q8(
23582        &self,
23583        q: &CudaSlice<f32>,
23584        k: &cudarc::driver::CudaView<u8>,
23585        v: &cudarc::driver::CudaView<u8>,
23586        o: &mut CudaSlice<f32>,
23587        head_dim: usize,
23588        n_head: usize,
23589        n_head_kv: usize,
23590        t_kv_dev: &CudaSlice<i32>,
23591        bucket_max: usize,
23592        scale: f32,
23593        k_tok_bytes: usize,
23594        v_tok_bytes: usize,
23595        g: bool,
23596        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23597    ) -> Result<(), Box<dyn std::error::Error>> {
23598        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
23599        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
23600        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
23601        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
23602        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
23603        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
23604        // 2026-07-12).
23605        let mut fa_vec =
23606            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23607        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23608            fa_vec = false;
23609        } // mirror kvmod/geom
23610        let sp = fa_split_keys(bucket_max, n_head_kv);
23611        let n_splits = if fa_vec {
23612            ((bucket_max + sp - 1) / sp).max(1)
23613        } else {
23614            ((bucket_max + 255) / 256).max(1)
23615        };
23616        let o_len = n_head * n_splits * head_dim;
23617        let ml_len = n_head * n_splits;
23618        let mut part_guard = self.fa_part_pool.lock().unwrap();
23619        if part_guard
23620            .as_ref()
23621            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23622            .unwrap_or(true)
23623        {
23624            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
23625            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
23626            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
23627            // later live allocations land at those addresses, and the next graph REPLAY writes
23628            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
23629            // output corruption began the burst after the trunk's t_kv growth first realloc'd
23630            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
23631            // the baked addresses alive (single-stream: eager writes the new buffers, replays
23632            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
23633            // (total retired < final size).
23634            let old = part_guard.take();
23635            let (co, cm) = old
23636                .as_ref()
23637                .map(|pp| (pp.0.len(), pp.1.len()))
23638                .unwrap_or((0, 0));
23639            if let Some(old) = old {
23640                self.fa_part_retired.lock().unwrap().push(old);
23641            }
23642            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23643                eprintln!(
23644                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23645                    co, o_len, cm, ml_len
23646                );
23647            }
23648            *part_guard = Some((
23649                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23650                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23651                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23652            ));
23653        }
23654        let pg = part_guard.as_mut().unwrap();
23655        self.gpu
23656            .stream()
23657            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23658        self.gpu
23659            .stream()
23660            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23661        self.gpu
23662            .stream()
23663            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23664        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23665        let (hd, nh, nhkv, nsp) = (
23666            head_dim as i32,
23667            n_head as i32,
23668            n_head_kv as i32,
23669            n_splits as i32,
23670        );
23671        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23672        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23673        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
23674        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
23675        let deep = fa_vec
23676            && head_dim == 256
23677            && fa_v4_at(bucket_max)
23678            && !g
23679            && fa_deep_at(bucket_max)
23680            && !matches!(fa_v4_mode(), "noB3" | "stage");
23681        let (f, cfg) = if fa_vec
23682            && head_dim == 512
23683            && bucket_max >= {
23684                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23685                *FA512_MIN_DC.get_or_init(|| {
23686                    std::env::var("MEMRA_FA512_MIN")
23687                        .ok()
23688                        .and_then(|v| v.parse().ok())
23689                        .unwrap_or(512)
23690                })
23691            } {
23692            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
23693            let gqa = (n_head / n_head_kv).max(1) as u32;
23694            (
23695                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23696                LaunchConfig {
23697                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23698                    block_dim: (32, gqa, 1),
23699                    shared_mem_bytes: 0,
23700                },
23701            )
23702        } else if fa_vec && head_dim == 512 {
23703            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
23704            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
23705            let q_view = q.as_view();
23706            let mut o_view = o.as_view_mut();
23707            return self.fa_decode_scalar_unified(
23708                &q_view,
23709                k,
23710                v,
23711                &mut o_view,
23712                head_dim,
23713                n_head,
23714                n_head_kv,
23715                0,
23716                Some(t_kv_dev),
23717                scale,
23718                n_splits,
23719                sp,
23720                k_tok_bytes,
23721                v_tok_bytes,
23722                g,
23723                &mut *part_o,
23724                &mut *part_m,
23725                &mut *part_l,
23726                q8_out,
23727            );
23728        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23729            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
23730            // incl the g-module route + raw-e4m3 sV sizing.
23731            let gqa = (n_head / n_head_kv).max(1) as u32;
23732            let fv = if g {
23733                self.func_g("fa_decode_vec_q_v4_dc")
23734            } else if deep {
23735                self.func("fa_decode_vec_q_v4_deep_dc")
23736            } else {
23737                self.func("fa_decode_vec_q_v4_dc")
23738            };
23739            let shmem =
23740                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23741            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23742            fv.set_attribute(
23743                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23744                shmem as i32,
23745            )?;
23746            (
23747                fv,
23748                LaunchConfig {
23749                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23750                    block_dim: (32, gqa, 1),
23751                    shared_mem_bytes: shmem,
23752                },
23753            )
23754        } else if fa_vec && fa_v3_active(head_dim) {
23755            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
23756            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
23757            let gqa = (n_head / n_head_kv).max(1) as u32;
23758            let fv = if g {
23759                self.func_g("fa_decode_vec_q_v3_dc")
23760            } else {
23761                self.func("fa_decode_vec_q_v3_dc")
23762            };
23763            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
23764            (
23765                fv,
23766                LaunchConfig {
23767                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23768                    block_dim: (32, gqa, 1),
23769                    shared_mem_bytes: shmem,
23770                },
23771            )
23772        } else if fa_vec && fa_v2_on() {
23773            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
23774            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
23775            // a numeric config; eager, rows-verify and graph all switch together).
23776            let gqa = (n_head / n_head_kv).max(1) as u32;
23777            let fv = if g {
23778                self.func_g("fa_decode_vec_q_v2_dc")
23779            } else {
23780                self.func("fa_decode_vec_q_v2_dc")
23781            };
23782            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
23783            (
23784                fv,
23785                LaunchConfig {
23786                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23787                    block_dim: (32, gqa, 1),
23788                    shared_mem_bytes: shmem,
23789                },
23790            )
23791        } else if fa_vec {
23792            let gqa = (n_head / n_head_kv).max(1) as u32;
23793            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
23794            let fv = if g {
23795                self.func_g("fa_decode_vec_q_dc")
23796            } else {
23797                self.func("fa_decode_vec_q_dc")
23798            };
23799            (
23800                fv,
23801                LaunchConfig {
23802                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23803                    block_dim: (32, gqa, 1),
23804                    shared_mem_bytes: 0,
23805                },
23806            )
23807        } else {
23808            let q_view = q.as_view();
23809            let mut o_view = o.as_view_mut();
23810            return self.fa_decode_scalar_unified(
23811                &q_view,
23812                k,
23813                v,
23814                &mut o_view,
23815                head_dim,
23816                n_head,
23817                n_head_kv,
23818                0,
23819                Some(t_kv_dev),
23820                scale,
23821                n_splits,
23822                if fa_vec { sp } else { 256 },
23823                k_tok_bytes,
23824                v_tok_bytes,
23825                g,
23826                &mut *part_o,
23827                &mut *part_m,
23828                &mut *part_l,
23829                q8_out,
23830            );
23831        };
23832        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
23833        let __s_b = self.gpu.stream();
23834        let mut b = __s_b.launch_builder(&f);
23835        b.arg(q)
23836            .arg(k)
23837            .arg(v)
23838            .arg(&mut *part_o)
23839            .arg(&mut *part_m)
23840            .arg(&mut *part_l)
23841            .arg(&hd)
23842            .arg(&nh)
23843            .arg(&nhkv)
23844            .arg(t_kv_dev)
23845            .arg(&scale)
23846            .arg(&nsp)
23847            .arg(&ski)
23848            .arg(&ktb)
23849            .arg(&vtb);
23850        unsafe {
23851            b.launch(cfg)?;
23852        }
23853        let cfg2 = LaunchConfig {
23854            grid_dim: (n_head as u32, 1, 1),
23855            block_dim: (head_dim as u32, 1, 1),
23856            shared_mem_bytes: 0,
23857        };
23858        if let Some((oq, od)) = q8_out {
23859            let fc = if g {
23860                self.func_g("fa_decode_combine_q8_1")
23861            } else {
23862                self.fa_func("fa_decode_combine_q8_1", head_dim)
23863            };
23864            let __s_b2 = self.gpu.stream();
23865            let mut b2 = __s_b2.launch_builder(&fc);
23866            b2.arg(&*part_o)
23867                .arg(&*part_m)
23868                .arg(&*part_l)
23869                .arg(oq)
23870                .arg(od)
23871                .arg(&hd)
23872                .arg(&nh)
23873                .arg(&nsp);
23874            unsafe {
23875                b2.launch(cfg2)?;
23876            }
23877            return Ok(());
23878        }
23879        let fc = if g {
23880            self.func_g("fa_decode_combine_f32")
23881        } else {
23882            self.fa_func("fa_decode_combine_f32", head_dim)
23883        };
23884        let __s_b2 = self.gpu.stream();
23885        let mut b2 = __s_b2.launch_builder(&fc);
23886        b2.arg(&*part_o)
23887            .arg(&*part_m)
23888            .arg(&*part_l)
23889            .arg(o)
23890            .arg(&hd)
23891            .arg(&nh)
23892            .arg(&nsp);
23893        unsafe {
23894            b2.launch(cfg2)?;
23895        }
23896        Ok(())
23897    }
23898
23899    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
23900    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
23901    /// at equal rows.
23902    #[allow(clippy::too_many_arguments)]
23903    pub fn append_kv_quantized_dcw(
23904        &self,
23905        k_row: &CudaSlice<f32>,
23906        v_row: &CudaSlice<f32>,
23907        kc: &mut CudaSlice<u8>,
23908        vc: &mut CudaSlice<u8>,
23909        len_dev: &CudaSlice<i32>,
23910        base_dev: Option<&CudaSlice<i32>>,
23911        kv_dim_k: usize,
23912        kv_dim_v: usize,
23913        k_tok_bytes: usize,
23914        v_tok_bytes: usize,
23915    ) -> Result<(), Box<dyn std::error::Error>> {
23916        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23917        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23918        let cfg = LaunchConfig {
23919            grid_dim: (nblk, 1, 1),
23920            block_dim: (32, 1, 1),
23921            shared_mem_bytes: 0,
23922        };
23923        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23924        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23925        let null: u64 = 0;
23926        let __s_b = self.gpu.stream();
23927        let mut b = __s_b.launch_builder(&f);
23928        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23929        match base_dev {
23930            Some(base) => {
23931                b.arg(base);
23932            }
23933            None => {
23934                b.arg(&null);
23935            }
23936        }
23937        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23938        unsafe {
23939            b.launch(cfg)?;
23940        }
23941        Ok(())
23942    }
23943
23944    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
23945    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23946        let f = self.func("inc_i32");
23947        let cfg = LaunchConfig {
23948            grid_dim: (1, 1, 1),
23949            block_dim: (1, 1, 1),
23950            shared_mem_bytes: 0,
23951        };
23952        let __s_b = self.gpu.stream();
23953        let mut b = __s_b.launch_builder(&f);
23954        b.arg(counter);
23955        unsafe {
23956            b.launch(cfg)?;
23957        }
23958        Ok(())
23959    }
23960
23961    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
23962    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
23963    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
23964    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
23965    /// kernel class on this lane); callers keep eager below the vec floor and for any other
23966    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
23967    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
23968    /// alive across bucket growth.
23969    #[allow(clippy::too_many_arguments)]
23970    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
23971    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
23972    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
23973    fn fa_part_pool_grow(
23974        &self,
23975        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
23976        o_len: usize,
23977        ml_len: usize,
23978    ) -> Result<(), Box<dyn std::error::Error>> {
23979        if part_guard
23980            .as_ref()
23981            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23982            .unwrap_or(true)
23983        {
23984            let old = part_guard.take();
23985            let (co, cm) = old
23986                .as_ref()
23987                .map(|pp| (pp.0.len(), pp.1.len()))
23988                .unwrap_or((0, 0));
23989            if let Some(old) = old {
23990                self.fa_part_retired.lock().unwrap().push(old);
23991            }
23992            *part_guard = Some((
23993                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23994                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23995                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23996            ));
23997        }
23998        Ok(())
23999    }
24000
24001    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
24002    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
24003    pub fn fa_dcw_pool_ensure(
24004        &self,
24005        head_dim: usize,
24006        n_head: usize,
24007        n_head_kv: usize,
24008        bucket_max: usize,
24009    ) -> Result<(), Box<dyn std::error::Error>> {
24010        let sp = fa_split_keys(bucket_max, n_head_kv);
24011        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24012        let o_len = n_head * n_splits * head_dim;
24013        let ml_len = n_head * n_splits;
24014        let mut part_guard = self.fa_part_pool.lock().unwrap();
24015        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24016    }
24017
24018    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
24019    /// appended; one launch walks the KV stream once with two query rows (per-row causal
24020    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
24021    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
24022    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
24023    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
24024    /// outputs (the head gate fuses into the combine as in the t=1 path).
24025    #[allow(clippy::too_many_arguments)]
24026    pub fn fa_decode_dcw2(
24027        &self,
24028        q2: &CudaSlice<f32>,
24029        k_ring: &cudarc::driver::CudaView<u8>,
24030        v_ring: &cudarc::driver::CudaView<u8>,
24031        o2: &mut CudaSlice<f32>,
24032        head_dim: usize,
24033        n_head: usize,
24034        n_head_kv: usize,
24035        len_dev: &CudaSlice<i32>,
24036        base_dev: Option<&CudaSlice<i32>>,
24037        window: usize,
24038        bucket_max: usize,
24039        scale: f32,
24040        k_tok_bytes: usize,
24041        v_tok_bytes: usize,
24042        gate2: &CudaSlice<f32>,
24043    ) -> Result<(), Box<dyn std::error::Error>> {
24044        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24045        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24046            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24047        }
24048        let sp = fa_split_keys(bucket_max, n_head_kv);
24049        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24050        // Partials for BOTH rows: row-major halves.
24051        let o_len = 2 * n_head * n_splits * head_dim;
24052        let ml_len = 2 * n_head * n_splits;
24053        let mut part_guard = self.fa_part_pool.lock().unwrap();
24054        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24055        let pg = part_guard.as_mut().unwrap();
24056        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24057        let (hd, nh, nhkv, nsp) = (
24058            head_dim as i32,
24059            n_head as i32,
24060            n_head_kv as i32,
24061            n_splits as i32,
24062        );
24063        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24064        let (ski, win) = (sp as i32, window as i32);
24065        let gqa = (n_head / n_head_kv).max(1) as u32;
24066        let smem = (32 * head_dim * 2) as u32;
24067        let f = self.func("fa_decode_vec_q_v3_dcw2");
24068        let cfg = LaunchConfig {
24069            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24070            block_dim: (32, gqa, 1),
24071            shared_mem_bytes: smem,
24072        };
24073        let null: u64 = 0;
24074        {
24075            let __s_b = self.gpu.stream();
24076            let mut b = __s_b.launch_builder(&f);
24077            b.arg(q2)
24078                .arg(k_ring)
24079                .arg(v_ring)
24080                .arg(&mut *part_o)
24081                .arg(&mut *part_m)
24082                .arg(&mut *part_l)
24083                .arg(&hd)
24084                .arg(&nh)
24085                .arg(&nhkv)
24086                .arg(len_dev);
24087            match base_dev {
24088                Some(base) => {
24089                    b.arg(base);
24090                }
24091                None => {
24092                    b.arg(&null);
24093                }
24094            }
24095            b.arg(&win)
24096                .arg(&scale)
24097                .arg(&nsp)
24098                .arg(&ski)
24099                .arg(&ktb)
24100                .arg(&vtb);
24101            unsafe {
24102                b.launch(cfg)?;
24103            }
24104        }
24105        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
24106        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
24107        // one launch covers both rows with the exact t=1 program per (row, head).
24108        let fc = {
24109            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24110            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24111                self.func("fa_decode_combine_gate_f32_s")
24112            } else {
24113                self.func("fa_decode_combine_gate_f32")
24114            }
24115        };
24116        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24117        let nh2 = (2 * n_head) as i32;
24118        let cfg2 = LaunchConfig {
24119            grid_dim: ((2 * n_head) as u32, 1, 1),
24120            block_dim: (head_dim as u32, 1, 1),
24121            shared_mem_bytes: if combine_shared {
24122                (2 * n_splits * 4) as u32
24123            } else {
24124                0
24125            },
24126        };
24127        let __s_b2 = self.gpu.stream();
24128        let mut b2 = __s_b2.launch_builder(&fc);
24129        b2.arg(&*part_o)
24130            .arg(&*part_m)
24131            .arg(&*part_l)
24132            .arg(gate2)
24133            .arg(o2)
24134            .arg(&hd)
24135            .arg(&nh2)
24136            .arg(&nsp);
24137        unsafe {
24138            b2.launch(cfg2)?;
24139        }
24140        Ok(())
24141    }
24142
24143    /// T-ROW dcw decode attention over a per-row session table (the per-session
24144    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
24145    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
24146    /// program verbatim with that row's ring/len/base and its own split geometry, so each
24147    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
24148    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
24149    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
24150    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
24151    #[allow(clippy::too_many_arguments)]
24152    pub fn fa_decode_dcw_rows(
24153        &self,
24154        q_rows: &CudaSlice<f32>,
24155        tab: &CudaSlice<u64>,
24156        o_rows: &mut CudaSlice<f32>,
24157        t: usize,
24158        head_dim: usize,
24159        n_head: usize,
24160        n_head_kv: usize,
24161        window: usize,
24162        max_ns: usize,
24163        scale: f32,
24164        k_tok_bytes: usize,
24165        v_tok_bytes: usize,
24166        gate_rows: &CudaSlice<f32>,
24167    ) -> Result<(), Box<dyn std::error::Error>> {
24168        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24169            || head_dim > 256
24170            || head_dim % 32 != 0
24171            || !fa_v3_on()
24172        {
24173            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24174        }
24175        if fa_sm_count() < 128
24176            || std::env::var("MEMRA_FA_SPLIT").is_ok()
24177            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24178            || std::env::var("MEMRA_FA_SP16").is_ok()
24179        {
24180            return Err(
24181                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24182                 (or a <128-SM rig) keep the per-row path"
24183                    .into(),
24184            );
24185        }
24186        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24187            return Err("fa_decode_dcw_rows geometry".into());
24188        }
24189        let o_len = t * n_head * max_ns * head_dim;
24190        let ml_len = t * n_head * max_ns;
24191        let mut part_guard = self.fa_part_pool.lock().unwrap();
24192        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24193        let pg = part_guard.as_mut().unwrap();
24194        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24195        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24196        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24197        let (win, mns) = (window as i32, max_ns as i32);
24198        let gqa = (n_head / n_head_kv).max(1) as u32;
24199        let smem = (32 * head_dim * 2) as u32;
24200        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24201        let cfg = LaunchConfig {
24202            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24203            block_dim: (32, gqa, 1),
24204            shared_mem_bytes: smem,
24205        };
24206        {
24207            let __s_b = self.gpu.stream();
24208            let mut b = __s_b.launch_builder(&f);
24209            b.arg(q_rows)
24210                .arg(tab)
24211                .arg(&mut *part_o)
24212                .arg(&mut *part_m)
24213                .arg(&mut *part_l)
24214                .arg(&hd)
24215                .arg(&nh)
24216                .arg(&nhkv)
24217                .arg(&win)
24218                .arg(&scale)
24219                .arg(&mns)
24220                .arg(&ktb)
24221                .arg(&vtb);
24222            unsafe {
24223                b.launch(cfg)?;
24224            }
24225        }
24226        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
24227        // row r head h reads its own partial bank; splits past a row's ns_eff carry
24228        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
24229        let fc = {
24230            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24231            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24232                self.func("fa_decode_combine_gate_f32_s")
24233            } else {
24234                self.func("fa_decode_combine_gate_f32")
24235            }
24236        };
24237        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24238        let nht = (t * n_head) as i32;
24239        let cfg2 = LaunchConfig {
24240            grid_dim: ((t * n_head) as u32, 1, 1),
24241            block_dim: (head_dim as u32, 1, 1),
24242            shared_mem_bytes: if combine_shared {
24243                (2 * max_ns * 4) as u32
24244            } else {
24245                0
24246            },
24247        };
24248        let __s_b2 = self.gpu.stream();
24249        let mut b2 = __s_b2.launch_builder(&fc);
24250        b2.arg(&*part_o)
24251            .arg(&*part_m)
24252            .arg(&*part_l)
24253            .arg(gate_rows)
24254            .arg(o_rows)
24255            .arg(&hd)
24256            .arg(&nht)
24257            .arg(&mns);
24258        unsafe {
24259            b2.launch(cfg2)?;
24260        }
24261        Ok(())
24262    }
24263
24264    pub fn fa_decode_dcw(
24265        &self,
24266        q: &CudaSlice<f32>,
24267        k_ring: &cudarc::driver::CudaView<u8>,
24268        v_ring: &cudarc::driver::CudaView<u8>,
24269        o: &mut CudaSlice<f32>,
24270        head_dim: usize,
24271        n_head: usize,
24272        n_head_kv: usize,
24273        len_dev: &CudaSlice<i32>,
24274        base_dev: Option<&CudaSlice<i32>>,
24275        window: usize,
24276        bucket_max: usize,
24277        scale: f32,
24278        k_tok_bytes: usize,
24279        v_tok_bytes: usize,
24280        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
24281        // one launch saved); `o` then receives the GATED output and the caller skips its
24282        // attn_head_gate call.
24283        fused_gate: Option<&CudaSlice<f32>>,
24284    ) -> Result<(), Box<dyn std::error::Error>> {
24285        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24286        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24287            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"
24288                .into());
24289        }
24290        let sp = fa_split_keys(bucket_max, n_head_kv);
24291        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24292        let o_len = n_head * n_splits * head_dim;
24293        let ml_len = n_head * n_splits;
24294        let mut part_guard = self.fa_part_pool.lock().unwrap();
24295        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24296        let pg = part_guard.as_mut().unwrap();
24297        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
24298        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
24299        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
24300        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
24301        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24302        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
24303        // finds the attention children BY their three-memset signature and updates the
24304        // memset widths per bucket — capturing without them silently kills retargeting
24305        // (battery-v8 token drift, 2026-08-21).
24306        let memset_on = *MEMSET_ON
24307            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24308            || crate::tp::token_graph_building();
24309        if memset_on {
24310            self.gpu
24311                .stream()
24312                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24313            self.gpu
24314                .stream()
24315                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24316            self.gpu
24317                .stream()
24318                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24319        }
24320        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24321        let (hd, nh, nhkv, nsp) = (
24322            head_dim as i32,
24323            n_head as i32,
24324            n_head_kv as i32,
24325            n_splits as i32,
24326        );
24327        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24328        let (ski, win) = (sp as i32, window as i32);
24329        let gqa = (n_head / n_head_kv).max(1) as u32;
24330        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
24331        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
24332        // see fa_dec_v3_walk_u). Same launch geometry.
24333        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24334        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24335        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24336            Ok("2") => 2,
24337            Ok("1") => 1,
24338            _ => 0,
24339        });
24340        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
24341        // permission-blocked in this container and the module params are not exposed, so this
24342        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
24343        // prints cumulative cycle shares every 430 launches.
24344        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24345        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24346        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24347            std::sync::Mutex::new(None);
24348        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
24349        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
24350        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
24351        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24352        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24353            && (n_head / n_head_kv) % 2 == 0
24354            && (n_head / n_head_kv) >= 2;
24355        let f = if fprof {
24356            self.func("fa_decode_vec_q_v3_dcw_prof")
24357        } else if hs2 {
24358            self.func("fa_decode_vec_q_v3_dcw_hs2")
24359        } else if hoist == 2 {
24360            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
24361            self.func("fa_decode_vec_q_v3_dcw_hc")
24362        } else if hoist == 1 {
24363            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
24364            self.func("fa_decode_vec_q_v3_dcw_h")
24365        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24366            self.func("fa_decode_vec_q_v3_dcw_u8")
24367        } else {
24368            self.func("fa_decode_vec_q_v3_dcw")
24369        };
24370        let cfg = LaunchConfig {
24371            grid_dim: if hs2 {
24372                ((2 * n_head_kv) as u32, n_splits as u32, 1)
24373            } else {
24374                (n_head_kv as u32, n_splits as u32, 1)
24375            },
24376            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24377            shared_mem_bytes: smem,
24378        };
24379        let null: u64 = 0;
24380        let __s_b = self.gpu.stream();
24381        let mut b = __s_b.launch_builder(&f);
24382        b.arg(q)
24383            .arg(k_ring)
24384            .arg(v_ring)
24385            .arg(&mut *part_o)
24386            .arg(&mut *part_m)
24387            .arg(&mut *part_l)
24388            .arg(&hd)
24389            .arg(&nh)
24390            .arg(&nhkv)
24391            .arg(len_dev);
24392        match base_dev {
24393            Some(base) => {
24394                b.arg(base);
24395            }
24396            None => {
24397                b.arg(&null);
24398            }
24399        }
24400        b.arg(&win)
24401            .arg(&scale)
24402            .arg(&nsp)
24403            .arg(&ski)
24404            .arg(&ktb)
24405            .arg(&vtb);
24406        if fprof {
24407            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24408            if guard
24409                .as_ref()
24410                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24411            {
24412                *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24413            }
24414            let (_, buf) = guard.as_mut().expect("armed above");
24415            b.arg(&*buf);
24416            unsafe {
24417                b.launch(cfg)?;
24418            }
24419            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24420            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24421            if n % 430 == 0 {
24422                self.stream().synchronize()?;
24423                let h = self.dtoh_u64(buf)?;
24424                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24425                let tot: u64 = h[..6].iter().sum();
24426                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24427                for (i, name) in phases.iter().enumerate() {
24428                    let pct = if tot > 0 {
24429                        h[i] as f64 / tot as f64 * 100.0
24430                    } else {
24431                        0.0
24432                    };
24433                    line.push_str(&format!(" {name}={pct:.1}%"));
24434                }
24435                if h[6] > 0 {
24436                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24437                }
24438                eprintln!("{line}");
24439            }
24440        } else {
24441            unsafe {
24442                b.launch(cfg)?;
24443            }
24444        }
24445        let mut combine_shared = false;
24446        let fc = if fused_gate.is_some() {
24447            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
24448            // n_splits-deep dependent global load chain every thread used to walk twice).
24449            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24450            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24451                combine_shared = true;
24452                self.func("fa_decode_combine_gate_f32_s")
24453            } else {
24454                self.func("fa_decode_combine_gate_f32")
24455            }
24456        } else {
24457            self.fa_func("fa_decode_combine_f32", head_dim)
24458        };
24459        let cfg2 = LaunchConfig {
24460            grid_dim: (n_head as u32, 1, 1),
24461            block_dim: (head_dim as u32, 1, 1),
24462            shared_mem_bytes: if combine_shared {
24463                (2 * n_splits * 4) as u32
24464            } else {
24465                0
24466            },
24467        };
24468        let __s_b2 = self.gpu.stream();
24469        let mut b2 = __s_b2.launch_builder(&fc);
24470        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24471        if let Some(gate_row) = fused_gate {
24472            b2.arg(gate_row);
24473        }
24474        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24475        unsafe {
24476            b2.launch(cfg2)?;
24477        }
24478        Ok(())
24479    }
24480
24481    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
24482    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
24483    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
24484    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
24485    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
24486    pub fn fa_geom_eager(
24487        &self,
24488        t_kv: usize,
24489        head_dim: usize,
24490        n_head_kv: usize,
24491        g: bool,
24492    ) -> (bool, usize) {
24493        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
24494        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
24495        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
24496        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24497        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
24498        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
24499        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
24500        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
24501        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
24502        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24503        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24504        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
24505        // family; everything else falls to the g-module scalar.
24506        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
24507        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
24508        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
24509        if g && head_dim == 256 && !fa_v4_at(t_kv) {
24510            fa_vec = false;
24511        }
24512        let sp = fa_split_keys(t_kv, n_head_kv);
24513        let n_splits = if fa_vec {
24514            ((t_kv + sp - 1) / sp).max(1)
24515        } else {
24516            ((t_kv + 255) / 256).max(1)
24517        };
24518        (fa_vec, n_splits)
24519    }
24520
24521    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
24522    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
24523    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
24524    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
24525    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
24526    pub fn fa_bucket_key(
24527        &self,
24528        t_kv: usize,
24529        head_dim: usize,
24530        n_head_kv: usize,
24531        g: bool,
24532    ) -> (bool, usize) {
24533        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24534    }
24535
24536    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
24537    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
24538    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
24539    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
24540    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
24541    /// device data) — every per-step varying scalar must come from a device counter. Returns the
24542    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
24543    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
24544    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
24545    /// replays (transients returning to the pool get reused by unrelated work and corrupt
24546    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
24547    pub fn capture_graph_retained<F>(
24548        &self,
24549        step: F,
24550    ) -> Result<
24551        (
24552            cudarc::driver::CudaGraph,
24553            Vec<Box<dyn std::any::Any + Send>>,
24554        ),
24555        Box<dyn std::error::Error>,
24556    >
24557    where
24558        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24559    {
24560        use cudarc::driver::sys::CUgraphInstantiate_flags;
24561        self.capture_graph_retained_flags(
24562            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24563            step,
24564        )
24565    }
24566
24567    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
24568    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
24569    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
24570    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
24571    pub fn capture_graph_retained_flags<F>(
24572        &self,
24573        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24574        mut step: F,
24575    ) -> Result<
24576        (
24577            cudarc::driver::CudaGraph,
24578            Vec<Box<dyn std::any::Any + Send>>,
24579        ),
24580        Box<dyn std::error::Error>,
24581    >
24582    where
24583        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24584    {
24585        use cudarc::driver::sys::CUstreamCaptureMode;
24586        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
24587        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
24588        // while the capture region is open become dead copy NODES replayed every launch
24589        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
24590        // warmup runs allocate the same transient sequence at the same pool addresses, so
24591        // retaining the warmup clones preserves the draft-graph fix without polluting the
24592        // captured graph.
24593        self.capture_keep.lock().unwrap().clear();
24594        let was_tracking = self.gpu.ctx.is_event_tracking();
24595        if was_tracking {
24596            unsafe {
24597                self.gpu.ctx.disable_event_tracking();
24598            }
24599        }
24600        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24601            self.capture_keep_on
24602                .store(true, std::sync::atomic::Ordering::Relaxed);
24603            let w = (|| {
24604                step(self)?;
24605                step(self)
24606            })();
24607            self.capture_keep_on
24608                .store(false, std::sync::atomic::Ordering::Relaxed);
24609            w?;
24610            self.gpu.stream().synchronize()?;
24611            self.gpu
24612                .stream()
24613                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24614            let r = step(self);
24615            let g = self.gpu.stream().end_capture(flags);
24616            r?;
24617            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24618            graph.upload()?;
24619            Ok(graph)
24620        };
24621        let result = run();
24622        self.capture_keep_on
24623            .store(false, std::sync::atomic::Ordering::Relaxed);
24624        if was_tracking {
24625            unsafe {
24626                self.gpu.ctx.enable_event_tracking();
24627            }
24628        }
24629        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24630        Ok((result?, keeper))
24631    }
24632
24633    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
24634    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
24635    /// alloc-free with persistent operands, and their bodies carry device side effects
24636    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
24637    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
24638    pub fn capture_graph_retained_nowarm<F>(
24639        &self,
24640        mut step: F,
24641    ) -> Result<
24642        (
24643            cudarc::driver::CudaGraph,
24644            Vec<Box<dyn std::any::Any + Send>>,
24645        ),
24646        Box<dyn std::error::Error>,
24647    >
24648    where
24649        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24650    {
24651        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24652        let was_tracking = self.gpu.ctx.is_event_tracking();
24653        if was_tracking {
24654            unsafe {
24655                self.gpu.ctx.disable_event_tracking();
24656            }
24657        }
24658        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24659            self.gpu.stream().synchronize()?;
24660            self.gpu
24661                .stream()
24662                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24663            let r = step(self);
24664            let g = self.gpu.stream().end_capture(
24665                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24666            );
24667            r?;
24668            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24669            graph.upload()?;
24670            Ok(graph)
24671        };
24672        let result = run();
24673        if was_tracking {
24674            unsafe {
24675                self.gpu.ctx.enable_event_tracking();
24676            }
24677        }
24678        Ok((result?, Vec::new()))
24679    }
24680
24681    pub fn capture_graph<F>(
24682        &self,
24683        mut step: F,
24684    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24685    where
24686        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24687    {
24688        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24689        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
24690        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
24691        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
24692        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
24693        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
24694        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
24695        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
24696        let was_tracking = self.gpu.ctx.is_event_tracking();
24697        if was_tracking {
24698            unsafe {
24699                self.gpu.ctx.disable_event_tracking();
24700            }
24701        }
24702        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
24703        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
24704        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
24705        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
24706        // measure that scan's real cost on the generic path. Diagnostic door only; the
24707        // default stays AUTO_FREE until a measured A/B justifies moving it.
24708        let iflag = {
24709            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24710            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24711                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
24712                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
24713                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24714                Ok("priority") => {
24715                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24716                }
24717                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24718            })
24719        };
24720        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
24721        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
24722        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
24723        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
24724        // eager step executions and are node-count-invariant. Printing the split bounds the
24725        // refactor's ceiling instead of assuming it.
24726        let ct = {
24727            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24728            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24729        };
24730        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
24731        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
24732        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
24733        // chased, and node-count-invariant, so no capture-body refactor could touch it.
24734        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
24735        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
24736        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
24737        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
24738        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
24739        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
24740        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
24741        // grow and never frees, resident counters/scratch, cache set in place), and the
24742        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
24743        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
24744        // settling and pool mapping. Arbitrated adversarially, not by taste:
24745        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
24746        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
24747        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
24748        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
24749        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
24750        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
24751        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
24752        let warmups = {
24753            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24754            *W.get_or_init(|| {
24755                std::env::var("MEMRA_GRAPH_WARMUPS")
24756                    .ok()
24757                    .and_then(|v| v.parse().ok())
24758                    .filter(|n| *n >= 1)
24759                    .unwrap_or(1)
24760            })
24761        };
24762        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24763            let t_w = std::time::Instant::now();
24764            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
24765            for _ in 0..warmups {
24766                step(self)?;
24767            }
24768            self.gpu.stream().synchronize()?;
24769            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24770            // capture the third run.
24771            let t_c = std::time::Instant::now();
24772            self.gpu
24773                .stream()
24774                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24775            // If the body errors mid-capture, end the capture before propagating so the stream isn't
24776            // left in a capturing state.
24777            let r = step(self);
24778            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24779            let t_i = std::time::Instant::now();
24780            let g = self.gpu.stream().end_capture(iflag);
24781            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24782            r?;
24783            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24784            let t_u = std::time::Instant::now();
24785            graph.upload()?;
24786            if ct {
24787                println!(
24788                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
24789                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
24790                    t_u.elapsed().as_secs_f64() * 1e3
24791                );
24792            }
24793            Ok(graph)
24794        };
24795        let result = run();
24796        if was_tracking {
24797            unsafe {
24798                self.gpu.ctx.enable_event_tracking();
24799            }
24800        }
24801        result
24802    }
24803
24804    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
24805    pub fn gdn_scan_s128_view(
24806        &self,
24807        q: &CudaSlice<f32>,
24808        k: &CudaSlice<f32>,
24809        v: &CudaSlice<f32>,
24810        g: &CudaSlice<f32>,
24811        beta: &CudaSlice<f32>,
24812        state_in: &cudarc::driver::CudaView<f32>,
24813        state_out: &mut cudarc::driver::CudaViewMut<f32>,
24814        o: &mut CudaSlice<f32>,
24815        n_head: usize,
24816        t: usize,
24817        scale: f32,
24818    ) -> Result<(), Box<dyn std::error::Error>> {
24819        let f = self.func("gdn_scan_s128");
24820        const S_V: u32 = 128;
24821        const WARP: u32 = 32;
24822        const COLS: u32 = 4;
24823        let cfg = LaunchConfig {
24824            grid_dim: (n_head as u32, 1, S_V / COLS),
24825            block_dim: (WARP, COLS, 1),
24826            shared_mem_bytes: 0,
24827        };
24828        let (h, ti) = (n_head as i32, t as i32);
24829        let __s_b = self.gpu.stream();
24830        let mut b = __s_b.launch_builder(&f);
24831        b.arg(q)
24832            .arg(k)
24833            .arg(v)
24834            .arg(g)
24835            .arg(beta)
24836            .arg(state_in)
24837            .arg(state_out)
24838            .arg(o)
24839            .arg(&h)
24840            .arg(&ti)
24841            .arg(&scale);
24842        unsafe {
24843            b.launch(cfg)?;
24844        }
24845        Ok(())
24846    }
24847
24848    /// conv1d where the input is a CudaView (resident conv state assembled in place).
24849    pub fn ssm_conv1d_view(
24850        &self,
24851        x: &cudarc::driver::CudaView<f32>,
24852        w: &CudaSlice<f32>,
24853        y: &mut CudaSlice<f32>,
24854        conv_dim: usize,
24855        t: usize,
24856        d_conv: usize,
24857        silu: bool,
24858    ) -> Result<(), Box<dyn std::error::Error>> {
24859        let f = self.func("ssm_conv1d_silu_f32");
24860        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
24861        let cfg = LaunchConfig {
24862            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24863            block_dim: (256, 1, 1),
24864            shared_mem_bytes: 0,
24865        };
24866        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24867        let __s_b = self.gpu.stream();
24868        let mut b = __s_b.launch_builder(&f);
24869        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24870        unsafe {
24871            b.launch(cfg)?;
24872        }
24873        Ok(())
24874    }
24875
24876    /// Depthwise causal conv1d + optional SiLU.
24877    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
24878    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
24879    /// FUSED prefill conv (token-major input, zero left-state): replaces
24880    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
24881    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
24882    pub fn ssm_conv1d_tm(
24883        &self,
24884        qkv_tm: &CudaSlice<f32>,
24885        w: &CudaSlice<f32>,
24886        y: &mut CudaSlice<f32>,
24887        conv_dim: usize,
24888        t: usize,
24889        d_conv: usize,
24890    ) -> Result<(), Box<dyn std::error::Error>> {
24891        let f = self.func("ssm_conv1d_tm_f32");
24892        let cfg = LaunchConfig {
24893            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24894            block_dim: (256, 1, 1),
24895            shared_mem_bytes: 0,
24896        };
24897        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24898        let __s_b = self.gpu.stream();
24899        let mut b = __s_b.launch_builder(&f);
24900        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24901        unsafe {
24902            b.launch(cfg)?;
24903        }
24904        Ok(())
24905    }
24906
24907    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
24908    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
24909    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
24910    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
24911    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
24912    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
24913    /// columns; the final ring == what T sequential decode ring rolls leave).
24914    pub fn ssm_conv1d_tm_state(
24915        &self,
24916        qkv_tm: &CudaSlice<f32>,
24917        conv_state: &mut CudaSlice<f32>,
24918        w: &CudaSlice<f32>,
24919        y: &mut CudaSlice<f32>,
24920        conv_dim: usize,
24921        t: usize,
24922        d_conv: usize,
24923    ) -> Result<(), Box<dyn std::error::Error>> {
24924        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24925    }
24926
24927    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
24928    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
24929    #[allow(clippy::too_many_arguments)]
24930    pub fn ssm_conv1d_tm_state_pad(
24931        &self,
24932        qkv_tm: &CudaSlice<f32>,
24933        conv_state: &mut CudaSlice<f32>,
24934        w: &CudaSlice<f32>,
24935        y: &mut CudaSlice<f32>,
24936        conv_dim: usize,
24937        t: usize,
24938        d_conv: usize,
24939        pad_len: Option<&CudaSlice<i32>>,
24940    ) -> Result<(), Box<dyn std::error::Error>> {
24941        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24942        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
24943        // the window kernel both read the pre-roll ring; the roll launches after both) — but
24944        // cloning first keeps the ordering trivially correct under any future stream split.
24945        let ring_old = if t < d_conv - 1 {
24946            Some(self.clone_dtod(conv_state)?)
24947        } else {
24948            None
24949        };
24950        {
24951            let f = self.func("ssm_conv1d_tm_state_f32");
24952            let cfg = LaunchConfig {
24953                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24954                block_dim: (256, 1, 1),
24955                shared_mem_bytes: 0,
24956            };
24957            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24958            let __s_b = self.gpu.stream();
24959            let mut b = __s_b.launch_builder(&f);
24960            b.arg(qkv_tm)
24961                .arg(&*conv_state)
24962                .arg(w)
24963                .arg(y)
24964                .arg(&cd)
24965                .arg(&ti)
24966                .arg(&dc);
24967            unsafe {
24968                b.launch(cfg)?;
24969            }
24970        }
24971        match (ring_old, pad_len) {
24972            (None, Some(len_d)) => {
24973                let f = self.func("ssm_conv_ring_update_dev_f32");
24974                let n = conv_dim * (d_conv - 1);
24975                let cfg = LaunchConfig::for_num_elems(n as u32);
24976                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24977                let __s_b = self.gpu.stream();
24978                let mut b = __s_b.launch_builder(&f);
24979                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24980                unsafe {
24981                    b.launch(cfg)?;
24982                }
24983            }
24984            (None, None) => {
24985                let f = self.func("ssm_conv_ring_update_f32");
24986                let n = conv_dim * (d_conv - 1);
24987                let cfg = LaunchConfig::for_num_elems(n as u32);
24988                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24989                let __s_b = self.gpu.stream();
24990                let mut b = __s_b.launch_builder(&f);
24991                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24992                unsafe {
24993                    b.launch(cfg)?;
24994                }
24995            }
24996            (Some(old), _) => {
24997                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
24998            }
24999        }
25000        Ok(())
25001    }
25002
25003    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
25004    pub fn ssm_conv1d_tm_state_pad_v(
25005        &self,
25006        qkv_tm: &cudarc::driver::CudaView<f32>,
25007        conv_state: &mut CudaSlice<f32>,
25008        w: &CudaSlice<f32>,
25009        y: &mut CudaSlice<f32>,
25010        conv_dim: usize,
25011        t: usize,
25012        d_conv: usize,
25013        pad_len: Option<&CudaSlice<i32>>,
25014    ) -> Result<(), Box<dyn std::error::Error>> {
25015        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25016        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
25017        // the window kernel both read the pre-roll ring; the roll launches after both) — but
25018        // cloning first keeps the ordering trivially correct under any future stream split.
25019        let ring_old = if t < d_conv - 1 {
25020            Some(self.clone_dtod(conv_state)?)
25021        } else {
25022            None
25023        };
25024        {
25025            let f = self.func("ssm_conv1d_tm_state_f32");
25026            let cfg = LaunchConfig {
25027                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25028                block_dim: (256, 1, 1),
25029                shared_mem_bytes: 0,
25030            };
25031            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25032            let __s_b = self.gpu.stream();
25033            let mut b = __s_b.launch_builder(&f);
25034            b.arg(qkv_tm)
25035                .arg(&*conv_state)
25036                .arg(w)
25037                .arg(y)
25038                .arg(&cd)
25039                .arg(&ti)
25040                .arg(&dc);
25041            unsafe {
25042                b.launch(cfg)?;
25043            }
25044        }
25045        match (ring_old, pad_len) {
25046            (None, Some(len_d)) => {
25047                let f = self.func("ssm_conv_ring_update_dev_f32");
25048                let n = conv_dim * (d_conv - 1);
25049                let cfg = LaunchConfig::for_num_elems(n as u32);
25050                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25051                let __s_b = self.gpu.stream();
25052                let mut b = __s_b.launch_builder(&f);
25053                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25054                unsafe {
25055                    b.launch(cfg)?;
25056                }
25057            }
25058            (None, None) => {
25059                let f = self.func("ssm_conv_ring_update_f32");
25060                let n = conv_dim * (d_conv - 1);
25061                let cfg = LaunchConfig::for_num_elems(n as u32);
25062                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25063                let __s_b = self.gpu.stream();
25064                let mut b = __s_b.launch_builder(&f);
25065                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25066                unsafe {
25067                    b.launch(cfg)?;
25068                }
25069            }
25070            (Some(_), _) => unreachable!(
25071                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25072            ),
25073        }
25074        Ok(())
25075    }
25076
25077    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
25078    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
25079    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
25080    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
25081    pub fn ssm_conv_ring_rebuild(
25082        &self,
25083        qkv_tm: &CudaSlice<f32>,
25084        ring_old: &CudaSlice<f32>,
25085        conv_state: &mut CudaSlice<f32>,
25086        conv_dim: usize,
25087        tc: usize,
25088        d_conv: usize,
25089    ) -> Result<(), Box<dyn std::error::Error>> {
25090        let f = self.func("ssm_conv_ring_rebuild_f32");
25091        let n = conv_dim * (d_conv - 1);
25092        let cfg = LaunchConfig::for_num_elems(n as u32);
25093        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25094        let __s_b = self.gpu.stream();
25095        let mut b = __s_b.launch_builder(&f);
25096        b.arg(qkv_tm)
25097            .arg(ring_old)
25098            .arg(conv_state)
25099            .arg(&cd)
25100            .arg(&ti)
25101            .arg(&dc);
25102        unsafe {
25103            b.launch(cfg)?;
25104        }
25105        Ok(())
25106    }
25107
25108    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
25109    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
25110    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
25111    /// the argmax + run-spec gates are the authority.
25112    #[allow(clippy::too_many_arguments)]
25113    pub fn gdn_prep_decode(
25114        &self,
25115        conv_out: &CudaSlice<f32>,
25116        beta_raw: &CudaSlice<f32>,
25117        alpha: &CudaSlice<f32>,
25118        dt_bias: &CudaSlice<f32>,
25119        a: &CudaSlice<f32>,
25120        q_l2: &mut CudaSlice<f32>,
25121        k_l2: &mut CudaSlice<f32>,
25122        v_g: &mut CudaSlice<f32>,
25123        beta: &mut CudaSlice<f32>,
25124        g_log: &mut CudaSlice<f32>,
25125        d_state: usize,
25126        num_v: usize,
25127        num_k: usize,
25128        key_dim: usize,
25129        eps: f32,
25130    ) -> Result<(), Box<dyn std::error::Error>> {
25131        let f = self.func("gdn_prep_decode_f32");
25132        let cfg = LaunchConfig {
25133            grid_dim: (num_v as u32, 1, 1),
25134            block_dim: (32, 4, 1),
25135            shared_mem_bytes: 0,
25136        };
25137        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25138        let __s_b = self.gpu.stream();
25139        let mut b = __s_b.launch_builder(&f);
25140        b.arg(conv_out)
25141            .arg(beta_raw)
25142            .arg(alpha)
25143            .arg(dt_bias)
25144            .arg(a)
25145            .arg(q_l2)
25146            .arg(k_l2)
25147            .arg(v_g)
25148            .arg(beta)
25149            .arg(g_log)
25150            .arg(&ds)
25151            .arg(&nv)
25152            .arg(&nk)
25153            .arg(&kd)
25154            .arg(&eps);
25155        unsafe {
25156            b.launch(cfg)?;
25157        }
25158        Ok(())
25159    }
25160
25161    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
25162    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
25163    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
25164    #[allow(clippy::too_many_arguments)]
25165    pub fn ssm_conv1d_gdn(
25166        &self,
25167        qkv_tm: &CudaSlice<f32>,
25168        w: &CudaSlice<f32>,
25169        q_g: &mut CudaSlice<f32>,
25170        k_g: &mut CudaSlice<f32>,
25171        v_g: &mut CudaSlice<f32>,
25172        conv_dim: usize,
25173        t: usize,
25174        d_conv: usize,
25175        d_state: usize,
25176        num_v: usize,
25177        num_k: usize,
25178        key_dim: usize,
25179    ) -> Result<(), Box<dyn std::error::Error>> {
25180        let f = self.func("ssm_conv1d_gdn_f32");
25181        let cfg = LaunchConfig {
25182            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25183            block_dim: (256, 1, 1),
25184            shared_mem_bytes: 0,
25185        };
25186        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25187        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25188        let __s_b = self.gpu.stream();
25189        let mut b = __s_b.launch_builder(&f);
25190        b.arg(qkv_tm)
25191            .arg(w)
25192            .arg(q_g)
25193            .arg(k_g)
25194            .arg(v_g)
25195            .arg(&cd)
25196            .arg(&ti)
25197            .arg(&dc)
25198            .arg(&ds)
25199            .arg(&nv)
25200            .arg(&nk)
25201            .arg(&kd);
25202        unsafe {
25203            b.launch(cfg)?;
25204        }
25205        Ok(())
25206    }
25207
25208    pub fn ssm_conv1d(
25209        &self,
25210        x: &CudaSlice<f32>,
25211        w: &CudaSlice<f32>,
25212        y: &mut CudaSlice<f32>,
25213        conv_dim: usize,
25214        t: usize,
25215        d_conv: usize,
25216        silu: bool,
25217    ) -> Result<(), Box<dyn std::error::Error>> {
25218        let f = self.func("ssm_conv1d_silu_f32");
25219        let cfg = LaunchConfig {
25220            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25221            block_dim: (256, 1, 1),
25222            shared_mem_bytes: 0,
25223        };
25224        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25225        let __s_b = self.gpu.stream();
25226        let mut b = __s_b.launch_builder(&f);
25227        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25228        unsafe {
25229            b.launch(cfg)?;
25230        }
25231        Ok(())
25232    }
25233
25234    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
25235    /// o:[128,H,T]. Single sequence.
25236    pub fn gdn_scan_s128(
25237        &self,
25238        q: &CudaSlice<f32>,
25239        k: &CudaSlice<f32>,
25240        v: &CudaSlice<f32>,
25241        g: &CudaSlice<f32>,
25242        beta: &CudaSlice<f32>,
25243        state_in: &CudaSlice<f32>,
25244        state_out: &mut CudaSlice<f32>,
25245        o: &mut CudaSlice<f32>,
25246        n_head: usize,
25247        t: usize,
25248        scale: f32,
25249    ) -> Result<(), Box<dyn std::error::Error>> {
25250        let f = self.func("gdn_scan_s128");
25251        const S_V: u32 = 128;
25252        const WARP: u32 = 32;
25253        const COLS_PER_BLOCK: u32 = 4;
25254        let cfg = LaunchConfig {
25255            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25256            block_dim: (WARP, COLS_PER_BLOCK, 1),
25257            shared_mem_bytes: 0,
25258        };
25259        let (h, ti) = (n_head as i32, t as i32);
25260        let __s_b = self.gpu.stream();
25261        let mut b = __s_b.launch_builder(&f);
25262        b.arg(q)
25263            .arg(k)
25264            .arg(v)
25265            .arg(g)
25266            .arg(beta)
25267            .arg(state_in)
25268            .arg(state_out)
25269            .arg(o)
25270            .arg(&h)
25271            .arg(&ti)
25272            .arg(&scale);
25273        unsafe {
25274            b.launch(cfg)?;
25275        }
25276        Ok(())
25277    }
25278
25279    // ==== B2' batched decode state ops (decode_batch.rs) ====
25280    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
25281    // Bodies are the single-seq kernels per sequence — bit-identical per row.
25282
25283    #[allow(clippy::too_many_arguments)]
25284    pub fn ssm_conv1d_fused_decode_b(
25285        &self,
25286        qkv_cols: &CudaSlice<f32>,
25287        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25288        w: &CudaSlice<f32>,
25289        conv_outs: &mut CudaSlice<f32>,
25290        conv_dim: usize,
25291        d_conv: usize,
25292        b_n: usize,
25293    ) -> Result<(), Box<dyn std::error::Error>> {
25294        let f = self.func("ssm_conv1d_fused_decode_b_f32");
25295        let cfg = LaunchConfig {
25296            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25297            block_dim: (256, 1, 1),
25298            shared_mem_bytes: 0,
25299        };
25300        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25301        let __s_b = self.gpu.stream();
25302        let mut b = __s_b.launch_builder(&f);
25303        b.arg(qkv_cols)
25304            .arg(conv_state_ptrs)
25305            .arg(w)
25306            .arg(conv_outs)
25307            .arg(&cd)
25308            .arg(&dc);
25309        unsafe {
25310            b.launch(cfg)?;
25311        }
25312        Ok(())
25313    }
25314
25315    #[allow(clippy::too_many_arguments)]
25316    pub fn gdn_prep_decode_b(
25317        &self,
25318        conv_outs: &CudaSlice<f32>,
25319        beta_raws: &CudaSlice<f32>,
25320        alphas: &CudaSlice<f32>,
25321        dt_bias: &CudaSlice<f32>,
25322        a: &CudaSlice<f32>,
25323        q_l2: &mut CudaSlice<f32>,
25324        k_l2: &mut CudaSlice<f32>,
25325        v_g: &mut CudaSlice<f32>,
25326        beta: &mut CudaSlice<f32>,
25327        g_log: &mut CudaSlice<f32>,
25328        d_state: usize,
25329        num_v: usize,
25330        num_k: usize,
25331        key_dim: usize,
25332        eps: f32,
25333        conv_dim: usize,
25334        b_n: usize,
25335    ) -> Result<(), Box<dyn std::error::Error>> {
25336        let f = self.func("gdn_prep_decode_b_f32");
25337        let cfg = LaunchConfig {
25338            grid_dim: (num_v as u32, 1, b_n as u32),
25339            block_dim: (32, 4, 1),
25340            shared_mem_bytes: 0,
25341        };
25342        let (ds, nv, nk, kd, cd) = (
25343            d_state as i32,
25344            num_v as i32,
25345            num_k as i32,
25346            key_dim as i32,
25347            conv_dim as i32,
25348        );
25349        let __s_b = self.gpu.stream();
25350        let mut b = __s_b.launch_builder(&f);
25351        b.arg(conv_outs)
25352            .arg(beta_raws)
25353            .arg(alphas)
25354            .arg(dt_bias)
25355            .arg(a)
25356            .arg(q_l2)
25357            .arg(k_l2)
25358            .arg(v_g)
25359            .arg(beta)
25360            .arg(g_log)
25361            .arg(&ds)
25362            .arg(&nv)
25363            .arg(&nk)
25364            .arg(&kd)
25365            .arg(&eps)
25366            .arg(&cd);
25367        unsafe {
25368            b.launch(cfg)?;
25369        }
25370        Ok(())
25371    }
25372
25373    #[allow(clippy::too_many_arguments)]
25374    pub fn gdn_scan_s128_batched(
25375        &self,
25376        q: &CudaSlice<f32>,
25377        k: &CudaSlice<f32>,
25378        v: &CudaSlice<f32>,
25379        g: &CudaSlice<f32>,
25380        beta: &CudaSlice<f32>,
25381        state_in_ptrs: &cudarc::driver::CudaView<u64>,
25382        state_out_ptrs: &cudarc::driver::CudaView<u64>,
25383        o: &mut CudaSlice<f32>,
25384        n_head: usize,
25385        b_n: usize,
25386        scale: f32,
25387    ) -> Result<(), Box<dyn std::error::Error>> {
25388        let f = self.func("gdn_scan_s128_b");
25389        const S_V: u32 = 128;
25390        const WARP: u32 = 32;
25391        const COLS_PER_BLOCK: u32 = 4;
25392        let cfg = LaunchConfig {
25393            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25394            block_dim: (WARP, COLS_PER_BLOCK, 1),
25395            shared_mem_bytes: 0,
25396        };
25397        let h = n_head as i32;
25398        let __s_b = self.gpu.stream();
25399        let mut b = __s_b.launch_builder(&f);
25400        b.arg(q)
25401            .arg(k)
25402            .arg(v)
25403            .arg(g)
25404            .arg(beta)
25405            .arg(state_in_ptrs)
25406            .arg(state_out_ptrs)
25407            .arg(o)
25408            .arg(&h)
25409            .arg(&scale);
25410        unsafe {
25411            b.launch(cfg)?;
25412        }
25413        Ok(())
25414    }
25415
25416    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
25417    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
25418    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
25419    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
25420    /// numeric class; only the pointer arithmetic moved host-side.
25421    #[allow(clippy::too_many_arguments)]
25422    pub fn ssm_conv1d_fused_decode_b_view(
25423        &self,
25424        qkv_cols: &cudarc::driver::CudaView<f32>,
25425        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25426        w: &CudaSlice<f32>,
25427        conv_outs: &mut CudaSlice<f32>,
25428        conv_dim: usize,
25429        d_conv: usize,
25430        b_n: usize,
25431    ) -> Result<(), Box<dyn std::error::Error>> {
25432        let f = self.func("ssm_conv1d_fused_decode_b_f32");
25433        let cfg = LaunchConfig {
25434            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25435            block_dim: (256, 1, 1),
25436            shared_mem_bytes: 0,
25437        };
25438        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25439        let __s_b = self.gpu.stream();
25440        let mut b = __s_b.launch_builder(&f);
25441        b.arg(qkv_cols)
25442            .arg(conv_state_ptrs)
25443            .arg(w)
25444            .arg(conv_outs)
25445            .arg(&cd)
25446            .arg(&dc);
25447        unsafe {
25448            b.launch(cfg)?;
25449        }
25450        Ok(())
25451    }
25452
25453    #[allow(clippy::too_many_arguments)]
25454    pub fn gdn_prep_decode_b_view(
25455        &self,
25456        conv_outs: &CudaSlice<f32>,
25457        beta_raws: &cudarc::driver::CudaView<f32>,
25458        alphas: &cudarc::driver::CudaView<f32>,
25459        dt_bias: &CudaSlice<f32>,
25460        a: &CudaSlice<f32>,
25461        q_l2: &mut CudaSlice<f32>,
25462        k_l2: &mut CudaSlice<f32>,
25463        v_g: &mut CudaSlice<f32>,
25464        beta: &mut CudaSlice<f32>,
25465        g_log: &mut CudaSlice<f32>,
25466        d_state: usize,
25467        num_v: usize,
25468        num_k: usize,
25469        key_dim: usize,
25470        eps: f32,
25471        conv_dim: usize,
25472        b_n: usize,
25473    ) -> Result<(), Box<dyn std::error::Error>> {
25474        let f = self.func("gdn_prep_decode_b_f32");
25475        let cfg = LaunchConfig {
25476            grid_dim: (num_v as u32, 1, b_n as u32),
25477            block_dim: (32, 4, 1),
25478            shared_mem_bytes: 0,
25479        };
25480        let (ds, nv, nk, kd, cd) = (
25481            d_state as i32,
25482            num_v as i32,
25483            num_k as i32,
25484            key_dim as i32,
25485            conv_dim as i32,
25486        );
25487        let __s_b = self.gpu.stream();
25488        let mut b = __s_b.launch_builder(&f);
25489        b.arg(conv_outs)
25490            .arg(beta_raws)
25491            .arg(alphas)
25492            .arg(dt_bias)
25493            .arg(a)
25494            .arg(q_l2)
25495            .arg(k_l2)
25496            .arg(v_g)
25497            .arg(beta)
25498            .arg(g_log)
25499            .arg(&ds)
25500            .arg(&nv)
25501            .arg(&nk)
25502            .arg(&kd)
25503            .arg(&eps)
25504            .arg(&cd);
25505        unsafe {
25506            b.launch(cfg)?;
25507        }
25508        Ok(())
25509    }
25510
25511    #[allow(clippy::too_many_arguments)]
25512    pub fn gdn_scan_s128_batched_view(
25513        &self,
25514        q: &CudaSlice<f32>,
25515        k: &CudaSlice<f32>,
25516        v: &CudaSlice<f32>,
25517        g: &CudaSlice<f32>,
25518        beta: &CudaSlice<f32>,
25519        state_in_ptrs: &cudarc::driver::CudaView<u64>,
25520        state_out_ptrs: &cudarc::driver::CudaView<u64>,
25521        o: &mut cudarc::driver::CudaViewMut<f32>,
25522        n_head: usize,
25523        b_n: usize,
25524        scale: f32,
25525    ) -> Result<(), Box<dyn std::error::Error>> {
25526        let f = self.func("gdn_scan_s128_b");
25527        const S_V: u32 = 128;
25528        const WARP: u32 = 32;
25529        const COLS_PER_BLOCK: u32 = 4;
25530        let cfg = LaunchConfig {
25531            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25532            block_dim: (WARP, COLS_PER_BLOCK, 1),
25533            shared_mem_bytes: 0,
25534        };
25535        let h = n_head as i32;
25536        let __s_b = self.gpu.stream();
25537        let mut b = __s_b.launch_builder(&f);
25538        b.arg(q)
25539            .arg(k)
25540            .arg(v)
25541            .arg(g)
25542            .arg(beta)
25543            .arg(state_in_ptrs)
25544            .arg(state_out_ptrs)
25545            .arg(o)
25546            .arg(&h)
25547            .arg(&scale);
25548        unsafe {
25549            b.launch(cfg)?;
25550        }
25551        Ok(())
25552    }
25553
25554    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
25555    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
25556    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
25557    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
25558    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
25559    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
25560    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
25561    /// identity law); prime_cache/forward/forward_last are the only callers.
25562    pub fn gdn_chunked_enabled() -> bool {
25563        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25564        *E.get_or_init(|| {
25565            std::env::var("MEMRA_GDN_CHUNKED")
25566                .map(|v| v != "0")
25567                .unwrap_or(true)
25568        })
25569    }
25570
25571    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
25572    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
25573    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
25574    /// of 32 in [32, 128] (kernel row mappings require it).
25575    pub fn gdn_chunk_size() -> usize {
25576        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25577        *C.get_or_init(|| {
25578            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25579                .ok()
25580                .and_then(|v| v.parse().ok())
25581                .unwrap_or(32);
25582            c.clamp(32, 128) / 32 * 32
25583        })
25584    }
25585
25586    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
25587    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
25588    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
25589    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
25590    #[allow(clippy::too_many_arguments)]
25591    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
25592    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
25593    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25594    #[allow(clippy::too_many_arguments)]
25595    pub fn gdn_chunk_k123(
25596        &self,
25597        q: &CudaSlice<f32>,
25598        k: &CudaSlice<f32>,
25599        v: &CudaSlice<f32>,
25600        g: &CudaSlice<f32>,
25601        beta: &CudaSlice<f32>,
25602        wb16: Option<&mut CudaSlice<u8>>,
25603        n_head: usize,
25604        t: usize,
25605        c: usize,
25606        hk: usize,
25607        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25608    ) -> Result<
25609        (
25610            CudaSlice<f32>,
25611            CudaSlice<f32>,
25612            CudaSlice<f32>,
25613            CudaSlice<f32>,
25614        ),
25615        Box<dyn std::error::Error>,
25616    > {
25617        const D: usize = 128;
25618        let h = n_head;
25619        let nc = (t + c - 1) / c;
25620        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25621        let mut gcum = self.uninit(t * h)?;
25622        let mut a = self.uninit(nc * h * c * c)?;
25623        let mut p = self.uninit(nc * h * c * c)?;
25624        let mut u = self.uninit(nc * h * c * D)?;
25625        let mut w = self.uninit(nc * h * c * D)?;
25626        {
25627            // K1
25628            let f = self.func("gdn_chunk_cumgate_f32");
25629            let cfg = LaunchConfig {
25630                grid_dim: (nc as u32, h as u32, 1),
25631                block_dim: (32, 1, 1),
25632                shared_mem_bytes: 0,
25633            };
25634            let __s_b = self.gpu.stream();
25635            let mut b = __s_b.launch_builder(&f);
25636            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25637            unsafe {
25638                b.launch(cfg)?;
25639            }
25640        }
25641        if let Some((qb, kb, pb)) = k2w {
25642            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
25643            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
25644            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25645            let f = self.func("gdn_k2_wgmma");
25646            let cfg = LaunchConfig {
25647                grid_dim: (nc as u32, h as u32, 1),
25648                block_dim: (128, 1, 1),
25649                shared_mem_bytes: 0,
25650            };
25651            let hki = hk as i32;
25652            let __s_b = self.gpu.stream();
25653            let mut b = __s_b.launch_builder(&f);
25654            b.arg(qb)
25655                .arg(kb)
25656                .arg(&gcum)
25657                .arg(beta)
25658                .arg(&mut a)
25659                .arg(&mut *pb)
25660                .arg(&hi)
25661                .arg(&ti)
25662                .arg(&ci)
25663                .arg(&hki);
25664            unsafe {
25665                b.launch(cfg)?;
25666            }
25667        } else if c <= 64 && !portable_mma_gated() {
25668            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
25669            let f = self.func("gdn_chunk_attn_f32");
25670            f.set_attribute(
25671                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25672                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25673            )?;
25674            let jt = ((c + 31) / 32) as u32;
25675            let cfg = LaunchConfig {
25676                grid_dim: (nc as u32, h as u32, jt),
25677                block_dim: (256, 1, 1),
25678                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25679            };
25680            let hki = hk as i32;
25681            let __s_b = self.gpu.stream();
25682            let mut b = __s_b.launch_builder(&f);
25683            b.arg(q)
25684                .arg(k)
25685                .arg(&gcum)
25686                .arg(beta)
25687                .arg(&mut a)
25688                .arg(&mut p)
25689                .arg(&hi)
25690                .arg(&ti)
25691                .arg(&ci)
25692                .arg(&hki);
25693            unsafe {
25694                b.launch(cfg)?;
25695            }
25696        } else {
25697            // K2 generic (C = 128, or the portable target's low-smem fallback)
25698            assert!(
25699                hk == h,
25700                "generic K2 is broadcast-only (de-broadcast rides C==32)"
25701            );
25702            let f = self.func("gdn_chunk_attn_g_f32");
25703            let cfg = LaunchConfig {
25704                grid_dim: (nc as u32, h as u32, 1),
25705                block_dim: (32, 8, 1),
25706                shared_mem_bytes: 0,
25707            };
25708            let __s_b = self.gpu.stream();
25709            let mut b = __s_b.launch_builder(&f);
25710            b.arg(q)
25711                .arg(k)
25712                .arg(&gcum)
25713                .arg(beta)
25714                .arg(&mut a)
25715                .arg(&mut p)
25716                .arg(&hi)
25717                .arg(&ti)
25718                .arg(&ci);
25719            unsafe {
25720                b.launch(cfg)?;
25721            }
25722        }
25723        {
25724            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
25725            let cfg = LaunchConfig {
25726                grid_dim: (nc as u32, h as u32, 1),
25727                block_dim: (256, 1, 1),
25728                shared_mem_bytes: 0,
25729            };
25730            match c {
25731                32 | 64 => {
25732                    let f = self.func(if c == 32 {
25733                        "gdn_chunk_solve32_f32"
25734                    } else {
25735                        "gdn_chunk_solve64_f32"
25736                    });
25737                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
25738                    let wb: u64 = match wb16 {
25739                        Some(d) => self.addr_u8(d),
25740                        None => 0,
25741                    };
25742                    let hki = hk as i32;
25743                    let __s_b = self.gpu.stream();
25744                    let mut b = __s_b.launch_builder(&f);
25745                    b.arg(v)
25746                        .arg(k)
25747                        .arg(&a)
25748                        .arg(&gcum)
25749                        .arg(&mut u)
25750                        .arg(&mut w)
25751                        .arg(&wb)
25752                        .arg(&hi)
25753                        .arg(&ti)
25754                        .arg(&hki);
25755                    unsafe {
25756                        b.launch(cfg)?;
25757                    }
25758                }
25759                _ => {
25760                    assert!(hk == h, "generic K3 is broadcast-only");
25761                    let f = self.func("gdn_chunk_solve_f32");
25762                    let __s_b = self.gpu.stream();
25763                    let mut b = __s_b.launch_builder(&f);
25764                    b.arg(v)
25765                        .arg(k)
25766                        .arg(&a)
25767                        .arg(&gcum)
25768                        .arg(&mut u)
25769                        .arg(&mut w)
25770                        .arg(&hi)
25771                        .arg(&ti)
25772                        .arg(&ci);
25773                    unsafe {
25774                        b.launch(cfg)?;
25775                    }
25776                }
25777            }
25778        }
25779        Ok((gcum, p, u, w))
25780    }
25781
25782    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
25783    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
25784    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
25785    pub fn gdn_db_on() -> bool {
25786        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25787    }
25788
25789    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
25790    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
25791    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
25792    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
25793    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
25794    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
25795    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
25796    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
25797    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25798        !portable_mma_gated()
25799            && c == 32
25800            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25801                Ok("1") => true,
25802                Ok("0") => false,
25803                _ => gdn_mma_default_on(),
25804            }
25805    }
25806
25807    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
25808    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
25809    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
25810    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
25811    /// force would silently produce garbage. Required since the sm_120a mma default
25812    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
25813    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25814        cfg!(memra_hopper_mma)
25815            && self.gdn_mma_enabled(c)
25816            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25817    }
25818
25819    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
25820    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
25821    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
25822    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
25823    #[allow(clippy::too_many_arguments)]
25824    pub fn ssm_conv1d_gdn_state_pad(
25825        &self,
25826        qkv_tm: &cudarc::driver::CudaView<f32>,
25827        conv_state: &mut CudaSlice<f32>,
25828        w: &CudaSlice<f32>,
25829        q_g: &mut CudaSlice<f32>,
25830        k_g: &mut CudaSlice<f32>,
25831        v_g: &mut CudaSlice<f32>,
25832        conv_dim: usize,
25833        t: usize,
25834        d_conv: usize,
25835        d_state: usize,
25836        num_v: usize,
25837        num_k: usize,
25838        key_dim: usize,
25839        hk: usize,
25840        pad_len: Option<&CudaSlice<i32>>,
25841    ) -> Result<(), Box<dyn std::error::Error>> {
25842        assert!(
25843            t >= d_conv - 1,
25844            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25845        );
25846        {
25847            let f = self.func("ssm_conv1d_gdn_state_f32");
25848            let cfg = LaunchConfig {
25849                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25850                block_dim: (256, 1, 1),
25851                shared_mem_bytes: 0,
25852            };
25853            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25854            let (ds, nv, nk, kd, hki) = (
25855                d_state as i32,
25856                num_v as i32,
25857                num_k as i32,
25858                key_dim as i32,
25859                hk as i32,
25860            );
25861            let __s_b = self.gpu.stream();
25862            let mut b = __s_b.launch_builder(&f);
25863            b.arg(qkv_tm)
25864                .arg(&*conv_state)
25865                .arg(w)
25866                .arg(q_g)
25867                .arg(k_g)
25868                .arg(v_g)
25869                .arg(&cd)
25870                .arg(&ti)
25871                .arg(&dc)
25872                .arg(&ds)
25873                .arg(&nv)
25874                .arg(&nk)
25875                .arg(&kd)
25876                .arg(&hki);
25877            unsafe {
25878                b.launch(cfg)?;
25879            }
25880        }
25881        match pad_len {
25882            Some(len_d) => {
25883                let f = self.func("ssm_conv_ring_update_dev_f32");
25884                let n = conv_dim * (d_conv - 1);
25885                let cfg = LaunchConfig::for_num_elems(n as u32);
25886                let (cd, dc) = (conv_dim as i32, d_conv as i32);
25887                let __s_b = self.gpu.stream();
25888                let mut b = __s_b.launch_builder(&f);
25889                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25890                unsafe {
25891                    b.launch(cfg)?;
25892                }
25893            }
25894            None => {
25895                let f = self.func("ssm_conv_ring_update_f32");
25896                let n = conv_dim * (d_conv - 1);
25897                let cfg = LaunchConfig::for_num_elems(n as u32);
25898                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25899                let __s_b = self.gpu.stream();
25900                let mut b = __s_b.launch_builder(&f);
25901                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25902                unsafe {
25903                    b.launch(cfg)?;
25904                }
25905            }
25906        }
25907        Ok(())
25908    }
25909
25910    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
25911    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
25912    /// K2/K3 can write them.
25913    pub fn gdn_chunk_alloc(
25914        &self,
25915        n_head: usize,
25916        t: usize,
25917        c: usize,
25918        hk: usize,
25919    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25920        const D: usize = 128;
25921        assert!(
25922            c == 32,
25923            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25924        );
25925        let h = n_head;
25926        let nc = (t + c - 1) / c;
25927        Ok(GdnChunkBufs {
25928            gcum: self.uninit(t * h)?,
25929            a: self.uninit(nc * h * c * c)?,
25930            p: self.uninit(nc * h * c * c)?,
25931            u: self.uninit(nc * h * c * D)?,
25932            w: self.uninit(nc * h * c * D)?,
25933            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25934            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25935            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25936            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25937            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25938            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25939            o: self.uninit(D * h * t)?,
25940            t,
25941            nc,
25942        })
25943    }
25944
25945    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
25946    pub fn f32_to_bf16_v(
25947        &self,
25948        x: &cudarc::driver::CudaView<f32>,
25949        dst: &mut CudaSlice<u8>,
25950        n: usize,
25951    ) -> Result<(), Box<dyn std::error::Error>> {
25952        let f = self.func("f32_to_bf16_bulk");
25953        let ni = n as i64;
25954        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25955        let __s_b = self.gpu.stream();
25956        let mut b = __s_b.launch_builder(&f);
25957        b.arg(x).arg(dst).arg(&ni);
25958        unsafe {
25959            b.launch(cfg)?;
25960        }
25961        Ok(())
25962    }
25963
25964    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
25965    pub fn f32_to_bf16_into(
25966        &self,
25967        x: &CudaSlice<f32>,
25968        dst: &mut CudaSlice<u8>,
25969        n: usize,
25970    ) -> Result<(), Box<dyn std::error::Error>> {
25971        let f = self.func("f32_to_bf16_bulk");
25972        let ni = n as i64;
25973        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25974        let __s_b = self.gpu.stream();
25975        let mut b = __s_b.launch_builder(&f);
25976        b.arg(x).arg(dst).arg(&ni);
25977        unsafe {
25978            b.launch(cfg)?;
25979        }
25980        Ok(())
25981    }
25982
25983    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
25984    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
25985    pub fn gdn_chunk_k123_vl8(
25986        &self,
25987        seqs: &[GdnSeqVl],
25988        n_head: usize,
25989        hk: usize,
25990        wq: Option<&GdnWVl8>,
25991    ) -> Result<(), Box<dyn std::error::Error>> {
25992        let b = seqs.len();
25993        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
25994        let mut packed = [GdnSeqVl::default(); 8];
25995        packed[..b].copy_from_slice(seqs);
25996        let v = GdnVl8(packed);
25997        let (hi, ci) = (n_head as i32, 32i32);
25998        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25999        {
26000            let f = self.func("gdn_chunk_cumgate_vl");
26001            let cfg = LaunchConfig {
26002                grid_dim: (max_nc, n_head as u32, b as u32),
26003                block_dim: (32, 1, 1),
26004                shared_mem_bytes: 0,
26005            };
26006            let __s_lb = self.gpu.stream();
26007            let mut lb = __s_lb.launch_builder(&f);
26008            lb.arg(&v).arg(&hi).arg(&ci);
26009            unsafe {
26010                lb.launch(cfg)?;
26011            }
26012        }
26013        let hki = hk as i32;
26014        if let Some(w) = wq {
26015            // K2-wgmma vl twin (writes A + pre-masked Pb16)
26016            let f = self.func("gdn_k2_wgmma_vl");
26017            let cfg = LaunchConfig {
26018                grid_dim: (max_nc, n_head as u32, b as u32),
26019                block_dim: (128, 1, 1),
26020                shared_mem_bytes: 0,
26021            };
26022            let __s_lb = self.gpu.stream();
26023            let mut lb = __s_lb.launch_builder(&f);
26024            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26025            unsafe {
26026                lb.launch(cfg)?;
26027            }
26028        } else {
26029            let f = self.func("gdn_chunk_attn_vl");
26030            f.set_attribute(
26031                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26032                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26033            )?;
26034            let cfg = LaunchConfig {
26035                grid_dim: (max_nc, n_head as u32, b as u32),
26036                block_dim: (256, 1, 1),
26037                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26038            };
26039            let __s_lb = self.gpu.stream();
26040            let mut lb = __s_lb.launch_builder(&f);
26041            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26042            unsafe {
26043                lb.launch(cfg)?;
26044            }
26045        }
26046        {
26047            let f = self.func("gdn_chunk_solve32_vl");
26048            let cfg = LaunchConfig {
26049                grid_dim: (max_nc, n_head as u32, b as u32),
26050                block_dim: (256, 1, 1),
26051                shared_mem_bytes: 0,
26052            };
26053            let __s_lb = self.gpu.stream();
26054            let mut lb = __s_lb.launch_builder(&f);
26055            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26056            unsafe {
26057                lb.launch(cfg)?;
26058            }
26059        }
26060        Ok(())
26061    }
26062
26063    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
26064    /// fused gate-prep, 5 launches for every sequence (per-element math identical
26065    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
26066    #[allow(clippy::too_many_arguments)]
26067    pub fn gdn_prep_vl8(
26068        &self,
26069        seqs: &[GdnPrepVl],
26070        conv_w: &CudaSlice<f32>,
26071        dt_bias: &CudaSlice<f32>,
26072        a: &CudaSlice<f32>,
26073        conv_dim: usize,
26074        d_conv: usize,
26075        d_state: usize,
26076        num_v: usize,
26077        num_k: usize,
26078        key_dim: usize,
26079        hk: usize,
26080        eps: f32,
26081    ) -> Result<(), Box<dyn std::error::Error>> {
26082        let b = seqs.len();
26083        assert!(b >= 1 && b <= 8);
26084        let mut packed = [GdnPrepVl::default(); 8];
26085        packed[..b].copy_from_slice(seqs);
26086        let v = GdnPrepVl8(packed);
26087        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26088        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26089        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26090        assert!(
26091            conv_fuse || hk == num_v,
26092            "de-broadcast requires the fused conv"
26093        );
26094        if conv_fuse {
26095            let f = self.func("ssm_conv1d_gdn_state_vl");
26096            let cfg = LaunchConfig {
26097                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26098                block_dim: (256, 1, 1),
26099                shared_mem_bytes: 0,
26100            };
26101            let (dsi, nvi, nki, kdi, hki) = (
26102                d_state as i32,
26103                num_v as i32,
26104                num_k as i32,
26105                key_dim as i32,
26106                hk as i32,
26107            );
26108            let __s_lb = self.gpu.stream();
26109            let mut lb = __s_lb.launch_builder(&f);
26110            lb.arg(&v)
26111                .arg(conv_w)
26112                .arg(&cdi)
26113                .arg(&dci)
26114                .arg(&dsi)
26115                .arg(&nvi)
26116                .arg(&nki)
26117                .arg(&kdi)
26118                .arg(&hki);
26119            unsafe {
26120                lb.launch(cfg)?;
26121            }
26122        } else {
26123            let f = self.func("ssm_conv1d_tm_state_vl");
26124            let cfg = LaunchConfig {
26125                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26126                block_dim: (256, 1, 1),
26127                shared_mem_bytes: 0,
26128            };
26129            let __s_lb = self.gpu.stream();
26130            let mut lb = __s_lb.launch_builder(&f);
26131            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26132            unsafe {
26133                lb.launch(cfg)?;
26134            }
26135        }
26136        {
26137            let f = self.func("ssm_conv_ring_update_vl");
26138            let n = (conv_dim * (d_conv - 1)) as u32;
26139            let cfg = LaunchConfig {
26140                grid_dim: (n.div_ceil(256), 1, b as u32),
26141                block_dim: (256, 1, 1),
26142                shared_mem_bytes: 0,
26143            };
26144            let __s_lb = self.gpu.stream();
26145            let mut lb = __s_lb.launch_builder(&f);
26146            lb.arg(&v).arg(&cdi).arg(&dci);
26147            unsafe {
26148                lb.launch(cfg)?;
26149            }
26150        }
26151        if !conv_fuse {
26152            let f = self.func("qkv_to_gdn_repack_vl");
26153            let n = max_t * (num_v * d_state) as u32;
26154            let cfg = LaunchConfig {
26155                grid_dim: (n.div_ceil(256), 1, b as u32),
26156                block_dim: (256, 1, 1),
26157                shared_mem_bytes: 0,
26158            };
26159            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26160            let __s_lb = self.gpu.stream();
26161            let mut lb = __s_lb.launch_builder(&f);
26162            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26163            unsafe {
26164                lb.launch(cfg)?;
26165            }
26166        }
26167        if Self::l2_v2_on(d_state) {
26168            let f = self.func("gdn_l2_v2_vl");
26169            let cfg = LaunchConfig {
26170                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26171                block_dim: (256, 1, 1),
26172                shared_mem_bytes: 0,
26173            };
26174            let (dsi, nvi) = (d_state as i32, hk as i32);
26175            let __s_lb = self.gpu.stream();
26176            let mut lb = __s_lb.launch_builder(&f);
26177            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26178            unsafe {
26179                lb.launch(cfg)?;
26180            }
26181        } else {
26182            let f = self.func("gdn_l2_vl");
26183            let cfg = LaunchConfig {
26184                grid_dim: (max_t * hk as u32, 2, b as u32),
26185                block_dim: (256, 1, 1),
26186                shared_mem_bytes: 0,
26187            };
26188            let (dsi, nvi) = (d_state as i32, hk as i32);
26189            let __s_lb = self.gpu.stream();
26190            let mut lb = __s_lb.launch_builder(&f);
26191            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26192            unsafe {
26193                lb.launch(cfg)?;
26194            }
26195        }
26196        {
26197            let f = self.func("gdn_gate_prep_vl");
26198            let n = max_t * num_v as u32;
26199            let cfg = LaunchConfig {
26200                grid_dim: (n.div_ceil(256), 1, b as u32),
26201                block_dim: (256, 1, 1),
26202                shared_mem_bytes: 0,
26203            };
26204            let nvi = num_v as i32;
26205            let __s_lb = self.gpu.stream();
26206            let mut lb = __s_lb.launch_builder(&f);
26207            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26208            unsafe {
26209                lb.launch(cfg)?;
26210            }
26211        }
26212        Ok(())
26213    }
26214
26215    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
26216    pub fn gdn_mirror_vl8(
26217        &self,
26218        seqs: &[GdnSeqVl],
26219        n_head: usize,
26220        which: i32,
26221        hk: usize,
26222    ) -> Result<(), Box<dyn std::error::Error>> {
26223        let b = seqs.len();
26224        assert!(b >= 1 && b <= 8);
26225        let mut packed = [GdnSeqVl::default(); 8];
26226        packed[..b].copy_from_slice(seqs);
26227        let v = GdnVl8(packed);
26228        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26229        let max_n = seqs
26230            .iter()
26231            .map(|s| {
26232                if which == 0 {
26233                    s.t as i64 * ept as i64
26234                } else {
26235                    s.nc as i64 * ept as i64 * 32
26236                }
26237            })
26238            .max()
26239            .unwrap();
26240        let f = self.func("gdn_mirror_vl");
26241        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26242        let cfg = LaunchConfig {
26243            grid_dim: (blocks, 1, b as u32),
26244            block_dim: (256, 1, 1),
26245            shared_mem_bytes: 0,
26246        };
26247        let __s_lb = self.gpu.stream();
26248        let mut lb = __s_lb.launch_builder(&f);
26249        lb.arg(&v).arg(&ept).arg(&which);
26250        unsafe {
26251            lb.launch(cfg)?;
26252        }
26253        Ok(())
26254    }
26255
26256    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
26257    pub fn gdn_tail_vl8(
26258        &self,
26259        seqs: &[GdnPrepVl],
26260        norm_w: &CudaSlice<f32>,
26261        d_state: usize,
26262        num_v: usize,
26263        eps: f32,
26264    ) -> Result<(), Box<dyn std::error::Error>> {
26265        let b = seqs.len();
26266        assert!(b >= 1 && b <= 8);
26267        let mut packed = [GdnPrepVl::default(); 8];
26268        packed[..b].copy_from_slice(seqs);
26269        let v = GdnPrepVl8(packed);
26270        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26271        let f = self.func("gated_rmsnorm_f16out_vl");
26272        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
26273        let cfg = LaunchConfig {
26274            grid_dim: (max_t * num_v as u32, 1, b as u32),
26275            block_dim: (128, 1, 1),
26276            shared_mem_bytes: 0,
26277        };
26278        let (dsi, nvi) = (d_state as i32, num_v as i32);
26279        let __s_lb = self.gpu.stream();
26280        let mut lb = __s_lb.launch_builder(&f);
26281        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26282        unsafe {
26283            lb.launch(cfg)?;
26284        }
26285        Ok(())
26286    }
26287
26288    /// Raw device address helpers for the varlen by-value arg struct (single-stream
26289    /// launches; every buffer outlives the call — the f16 FFI discipline).
26290    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26291        use cudarc::driver::DevicePtr;
26292        let s = self.gpu.stream();
26293        let (p, _g) = x.device_ptr(&s);
26294        p as u64
26295    }
26296    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26297        use cudarc::driver::DevicePtrMut;
26298        let s = self.gpu.stream();
26299        let (p, _g) = x.device_ptr_mut(&s);
26300        p as u64
26301    }
26302    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26303        use cudarc::driver::DevicePtr;
26304        let s = self.gpu.stream();
26305        let (p, _g) = x.device_ptr(&s);
26306        p as u64
26307    }
26308    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26309        use cudarc::driver::DevicePtr;
26310        let s = self.gpu.stream();
26311        let (p, _g) = x.device_ptr(&s);
26312        p as u64
26313    }
26314
26315    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
26316    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
26317    /// launches, so this is strictly bit-gateable against them).
26318    pub fn gdn_chunk_vl8(
26319        &self,
26320        seqs: &[GdnSeqVl],
26321        n_head: usize,
26322        scale: f32,
26323        hk: usize,
26324        wq: Option<&GdnWVl8>,
26325    ) -> Result<(), Box<dyn std::error::Error>> {
26326        const NSPLIT: u32 = 4;
26327        let b = seqs.len();
26328        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26329        let mut packed = [GdnSeqVl::default(); 8];
26330        packed[..b].copy_from_slice(seqs);
26331        let v = GdnVl8(packed);
26332        let (hi, ci) = (n_head as i32, 32i32);
26333        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26334        let hki = hk as i32;
26335        if let Some(w) = wq {
26336            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
26337            let f = self.func("gdn_k45_wgmma_vl");
26338            let cfg = LaunchConfig {
26339                grid_dim: (n_head as u32, NSPLIT, b as u32),
26340                block_dim: (256, 1, 1),
26341                shared_mem_bytes: 0,
26342            };
26343            let __s_lb = self.gpu.stream();
26344            let mut lb = __s_lb.launch_builder(&f);
26345            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26346            unsafe {
26347                lb.launch(cfg)?;
26348            }
26349            let _ = max_nc;
26350            return Ok(());
26351        }
26352        {
26353            let f = self.func("gdn_chunk_state_mma_vl");
26354            let cfg = LaunchConfig {
26355                grid_dim: (n_head as u32, NSPLIT, b as u32),
26356                block_dim: (256, 1, 1),
26357                shared_mem_bytes: 0,
26358            };
26359            let __s_lb = self.gpu.stream();
26360            let mut lb = __s_lb.launch_builder(&f);
26361            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26362            unsafe {
26363                lb.launch(cfg)?;
26364            }
26365        }
26366        {
26367            let f = self.func("gdn_chunk_output_mma_vl");
26368            let cfg = LaunchConfig {
26369                grid_dim: (max_nc, n_head as u32, b as u32),
26370                block_dim: (256, 1, 1),
26371                shared_mem_bytes: 0,
26372            };
26373            let __s_lb = self.gpu.stream();
26374            let mut lb = __s_lb.launch_builder(&f);
26375            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26376            unsafe {
26377                lb.launch(cfg)?;
26378            }
26379        }
26380        Ok(())
26381    }
26382    pub fn gdn_scan_chunked(
26383        &self,
26384        q: &CudaSlice<f32>,
26385        k: &CudaSlice<f32>,
26386        v: &CudaSlice<f32>,
26387        g: &CudaSlice<f32>,
26388        beta: &CudaSlice<f32>,
26389        kb16_pre: Option<&CudaSlice<u8>>,
26390        qb16_pre: Option<&CudaSlice<u8>>,
26391        state_in: &CudaSlice<f32>,
26392        state_out: &mut CudaSlice<f32>,
26393        o: &mut CudaSlice<f32>,
26394        n_head: usize,
26395        t: usize,
26396        scale: f32,
26397        c: usize,
26398        hk: usize,
26399    ) -> Result<(), Box<dyn std::error::Error>> {
26400        const D: usize = 128;
26401        const NSPLIT: u32 = 4;
26402        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26403        let h = n_head;
26404        let nc = (t + c - 1) / c;
26405        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26406        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
26407        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
26408        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
26409        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
26410        let gdn_mma_pre = !portable_mma_gated()
26411            && c == 32
26412            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26413                Ok("1") => true,
26414                Ok("0") => false,
26415                _ => gdn_mma_default_on(),
26416            };
26417        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26418            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26419        } else {
26420            None
26421        };
26422        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
26423        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
26424        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
26425        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
26426        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26427            && gdn_mma_pre
26428            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26429        let nk = t * hk * D;
26430        let mut kb16_local: Option<CudaSlice<u8>> = None;
26431        if gdn_mma_pre && kb16_pre.is_none() {
26432            let mut kb = self.alloc_u8_uninit(nk * 2)?;
26433            let f = self.func("f32_to_bf16_bulk");
26434            let n2 = nk as i64;
26435            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26436            let __s_b = self.gpu.stream();
26437            let mut b = __s_b.launch_builder(&f);
26438            b.arg(k).arg(&mut kb).arg(&n2);
26439            unsafe {
26440                b.launch(cfg2)?;
26441            }
26442            kb16_local = Some(kb);
26443        }
26444        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26445        if let Some(kb) = kb16_pre {
26446            assert!(kb.len() >= nk * 2, "kb16_pre too small");
26447        }
26448        let mut qb16: Option<CudaSlice<u8>> = None;
26449        let mut pb16: Option<CudaSlice<u8>> = None;
26450        if gdn_wgmma_pre {
26451            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
26452            // the standalone bulk cvt only serves callers without the prep mirror.
26453            if qb16_pre.is_none() {
26454                let mut qb = self.alloc_u8_uninit(nk * 2)?;
26455                let f = self.func("f32_to_bf16_bulk");
26456                let n2 = nk as i64;
26457                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26458                let __s_b = self.gpu.stream();
26459                let mut b = __s_b.launch_builder(&f);
26460                b.arg(q).arg(&mut qb).arg(&n2);
26461                unsafe {
26462                    b.launch(cfg2)?;
26463                }
26464                qb16 = Some(qb);
26465            } else if let Some(qb) = qb16_pre {
26466                assert!(qb.len() >= nk * 2, "qb16_pre too small");
26467            }
26468            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26469        }
26470        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26471        let k2w = if gdn_wgmma_pre {
26472            Some((
26473                *qb16_ref0.as_ref().unwrap(),
26474                *kb16_ref0.as_ref().unwrap(),
26475                pb16.as_mut().unwrap(),
26476            ))
26477        } else {
26478            None
26479        };
26480        let (gcum, p, u, w) =
26481            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26482        let _ = &w;
26483        let mut y = self.uninit(nc * h * c * D)?;
26484        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
26485        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
26486        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
26487        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
26488        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
26489        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
26490        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
26491        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
26492        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
26493        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
26494        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
26495        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
26496        // sites must agree or the pre-work arms while the scan takes the scalar route.
26497        let gdn_mma = !portable_mma_gated()
26498            && c == 32
26499            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26500                Ok("1") => true,
26501                Ok("0") => false,
26502                _ => gdn_mma_default_on(),
26503            };
26504        if gdn_mma {
26505            let wb16 = wb16_pre
26506                .take()
26507                .expect("mma path pre-allocates wb16 (K3 store fold)");
26508            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26509            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
26510            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
26511            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
26512            // pass runs inside the persistent-M kernel; Y and Ssnap are never
26513            // materialized. New numeric class (gk folds into k^T instead of ys) —
26514            // explicit opt-in until the state-carry battery promotes it. Env read per
26515            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
26516            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
26517            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
26518            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
26519            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
26520            if gdn_wgmma_pre {
26521                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
26522                let qb16 = qb16_ref0.unwrap();
26523                let pb16 = pb16.as_ref().unwrap();
26524                {
26525                    let f = self.func("gdn_k45_wgmma");
26526                    let cfg = LaunchConfig {
26527                        grid_dim: (h as u32, 4, 1),
26528                        block_dim: (256, 1, 1),
26529                        shared_mem_bytes: 0,
26530                    };
26531                    let hki = hk as i32;
26532                    let __s_b = self.gpu.stream();
26533                    let mut b = __s_b.launch_builder(&f);
26534                    b.arg(kb16_ref)
26535                        .arg(&gcum)
26536                        .arg(beta)
26537                        .arg(&u)
26538                        .arg(&wb16)
26539                        .arg(qb16)
26540                        .arg(pb16)
26541                        .arg(o)
26542                        .arg(&scale)
26543                        .arg(state_in)
26544                        .arg(&mut *state_out)
26545                        .arg(&hi)
26546                        .arg(&ti)
26547                        .arg(&ci)
26548                        .arg(&hki);
26549                    unsafe {
26550                        b.launch(cfg)?;
26551                    }
26552                }
26553                return Ok(());
26554            }
26555            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
26556            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
26557            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
26558            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26559            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26560            {
26561                let f = self.func("gdn_chunk_state_mma");
26562                let cfg = LaunchConfig {
26563                    grid_dim: (h as u32, NSPLIT, 1),
26564                    block_dim: (256, 1, 1),
26565                    shared_mem_bytes: 0,
26566                };
26567                let hki = hk as i32;
26568                let __s_b = self.gpu.stream();
26569                let mut b = __s_b.launch_builder(&f);
26570                b.arg(kb16_ref)
26571                    .arg(&gcum)
26572                    .arg(beta)
26573                    .arg(&u)
26574                    .arg(&wb16)
26575                    .arg(&mut y16)
26576                    .arg(&mut ssnap16)
26577                    .arg(state_in)
26578                    .arg(&mut *state_out)
26579                    .arg(&hi)
26580                    .arg(&ti)
26581                    .arg(&ci)
26582                    .arg(&hki);
26583                unsafe {
26584                    b.launch(cfg)?;
26585                }
26586            }
26587            {
26588                // K5-mma (bf16 St/Y consumers)
26589                let f = self.func("gdn_chunk_output_mma");
26590                let jt = ((c + 31) / 32) as u32;
26591                let cfg = LaunchConfig {
26592                    grid_dim: (nc as u32, h as u32, jt),
26593                    block_dim: (256, 1, 1),
26594                    shared_mem_bytes: 0,
26595                };
26596                let hki = hk as i32;
26597                let __s_b = self.gpu.stream();
26598                let mut b = __s_b.launch_builder(&f);
26599                b.arg(q)
26600                    .arg(&gcum)
26601                    .arg(&p)
26602                    .arg(&y16)
26603                    .arg(&ssnap16)
26604                    .arg(o)
26605                    .arg(&hi)
26606                    .arg(&ti)
26607                    .arg(&ci)
26608                    .arg(&scale)
26609                    .arg(&hki);
26610                unsafe {
26611                    b.launch(cfg)?;
26612                }
26613            }
26614            return Ok(());
26615        }
26616        {
26617            // K4 (sequential over chunks inside; blocks col-partition the state)
26618            let f = self.func("gdn_chunk_state_f32");
26619            let cfg = LaunchConfig {
26620                grid_dim: (h as u32, NSPLIT, 1),
26621                block_dim: (256, 1, 1),
26622                shared_mem_bytes: 0,
26623            };
26624            let __s_b = self.gpu.stream();
26625            let mut b = __s_b.launch_builder(&f);
26626            b.arg(k)
26627                .arg(&gcum)
26628                .arg(beta)
26629                .arg(&u)
26630                .arg(&w)
26631                .arg(&mut y)
26632                .arg(&mut ssnap)
26633                .arg(state_in)
26634                .arg(&mut *state_out)
26635                .arg(&hi)
26636                .arg(&ti)
26637                .arg(&ci);
26638            unsafe {
26639                b.launch(cfg)?;
26640            }
26641        }
26642        {
26643            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
26644            let f = self.func("gdn_chunk_output_f32");
26645            let jt = ((c + 31) / 32) as u32;
26646            let cfg = LaunchConfig {
26647                grid_dim: (nc as u32, h as u32, jt),
26648                block_dim: (256, 1, 1),
26649                shared_mem_bytes: 0,
26650            };
26651            let __s_b = self.gpu.stream();
26652            let mut b = __s_b.launch_builder(&f);
26653            b.arg(q)
26654                .arg(&gcum)
26655                .arg(&p)
26656                .arg(&y)
26657                .arg(&ssnap)
26658                .arg(o)
26659                .arg(&hi)
26660                .arg(&ti)
26661                .arg(&ci)
26662                .arg(&scale);
26663            unsafe {
26664                b.launch(cfg)?;
26665            }
26666        }
26667        Ok(())
26668    }
26669
26670    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
26671    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
26672    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
26673    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
26674    ///
26675    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
26676    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
26677    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
26678    #[allow(clippy::too_many_arguments)]
26679    #[allow(clippy::too_many_arguments)]
26680    pub fn gdn_scan_prefill(
26681        &self,
26682        q: &CudaSlice<f32>,
26683        k: &CudaSlice<f32>,
26684        v: &CudaSlice<f32>,
26685        g: &CudaSlice<f32>,
26686        beta: &CudaSlice<f32>,
26687        kb16_pre: Option<&CudaSlice<u8>>,
26688        qb16_pre: Option<&CudaSlice<u8>>,
26689        state_in: &CudaSlice<f32>,
26690        state_out: &mut CudaSlice<f32>,
26691        o: &mut CudaSlice<f32>,
26692        n_head: usize,
26693        t: usize,
26694        scale: f32,
26695        hk: usize,
26696    ) -> Result<(), Box<dyn std::error::Error>> {
26697        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26698            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26699            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26700        }
26701        if Self::gdn_chunked_enabled() && t >= 16 {
26702            self.gdn_scan_chunked(
26703                q,
26704                k,
26705                v,
26706                g,
26707                beta,
26708                kb16_pre,
26709                qb16_pre,
26710                state_in,
26711                state_out,
26712                o,
26713                n_head,
26714                t,
26715                scale,
26716                Self::gdn_chunk_size(),
26717                hk,
26718            )
26719        } else {
26720            assert!(
26721                hk == n_head,
26722                "s128 scan is broadcast-only (prep guarantees by predicate)"
26723            );
26724            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26725        }
26726    }
26727
26728    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
26729    #[allow(clippy::too_many_arguments)]
26730    fn gdn_scan_diff(
26731        &self,
26732        q: &CudaSlice<f32>,
26733        k: &CudaSlice<f32>,
26734        v: &CudaSlice<f32>,
26735        g: &CudaSlice<f32>,
26736        beta: &CudaSlice<f32>,
26737        state_in: &CudaSlice<f32>,
26738        state_out: &mut CudaSlice<f32>,
26739        o: &mut CudaSlice<f32>,
26740        n_head: usize,
26741        t: usize,
26742        scale: f32,
26743    ) -> Result<(), Box<dyn std::error::Error>> {
26744        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26745        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26746        let mut o_c = self.uninit(o.len())?;
26747        let mut st_c = self.uninit(state_out.len())?;
26748        self.gdn_scan_chunked(
26749            q,
26750            k,
26751            v,
26752            g,
26753            beta,
26754            None,
26755            None,
26756            state_in,
26757            &mut st_c,
26758            &mut o_c,
26759            n_head,
26760            t,
26761            scale,
26762            Self::gdn_chunk_size(),
26763            n_head,
26764        )?;
26765        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26766        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26767        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26768        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26769            let mut max_abs = 0f32;
26770            let mut max_rel = 0f32;
26771            let mut sum_rel = 0f64;
26772            for (x, y) in a.iter().zip(b) {
26773                let ad = (x - y).abs();
26774                let rel = ad / x.abs().max(y.abs()).max(1e-3);
26775                if ad > max_abs {
26776                    max_abs = ad;
26777                }
26778                if rel > max_rel {
26779                    max_rel = rel;
26780                }
26781                sum_rel += rel as f64;
26782            }
26783            (max_abs, max_rel, sum_rel / a.len() as f64)
26784        };
26785        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26786        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26787        println!(
26788            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26789                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26790            Self::gdn_chunk_size()
26791        );
26792        Ok(())
26793    }
26794
26795    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
26796    pub fn gdn_glog(
26797        &self,
26798        alpha: &CudaSlice<f32>,
26799        dt_bias: &CudaSlice<f32>,
26800        a: &CudaSlice<f32>,
26801        g_log: &mut CudaSlice<f32>,
26802        n_head: usize,
26803        t: usize,
26804    ) -> Result<(), Box<dyn std::error::Error>> {
26805        let f = self.func("gdn_glog_f32");
26806        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26807        let (h, ti) = (n_head as i32, t as i32);
26808        let __s_b = self.gpu.stream();
26809        let mut b = __s_b.launch_builder(&f);
26810        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26811        unsafe {
26812            b.launch(cfg)?;
26813        }
26814        Ok(())
26815    }
26816
26817    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
26818    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
26819    pub fn sigmoid_v(
26820        &self,
26821        x: &cudarc::driver::CudaView<f32>,
26822        y: &mut CudaSlice<f32>,
26823        n: usize,
26824    ) -> Result<(), Box<dyn std::error::Error>> {
26825        let f = self.func("sigmoid_f32");
26826        let cfg = LaunchConfig::for_num_elems(n as u32);
26827        let ni = n as i32;
26828        let __s_b = self.gpu.stream();
26829        let mut b = __s_b.launch_builder(&f);
26830        b.arg(x).arg(y).arg(&ni);
26831        unsafe {
26832            b.launch(cfg)?;
26833        }
26834        Ok(())
26835    }
26836
26837    pub fn gdn_glog_v(
26838        &self,
26839        alpha: &cudarc::driver::CudaView<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    pub fn sigmoid(
26859        &self,
26860        x: &CudaSlice<f32>,
26861        y: &mut CudaSlice<f32>,
26862        n: usize,
26863    ) -> Result<(), Box<dyn std::error::Error>> {
26864        let f = self.func("sigmoid_f32");
26865        let cfg = LaunchConfig::for_num_elems(n as u32);
26866        let ni = n as i32;
26867        let __s_b = self.gpu.stream();
26868        let mut b = __s_b.launch_builder(&f);
26869        b.arg(x).arg(y).arg(&ni);
26870        unsafe {
26871            b.launch(cfg)?;
26872        }
26873        Ok(())
26874    }
26875
26876    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
26877    /// (replaces sigmoid + mul + convert). Bit-identical class.
26878    pub fn sig_mul_f16out(
26879        &self,
26880        a: &CudaSlice<f32>,
26881        g: &CudaSlice<f32>,
26882        dst: &mut CudaSlice<f32>,
26883        dst16: &mut CudaSlice<u8>,
26884        n: usize,
26885    ) -> Result<(), Box<dyn std::error::Error>> {
26886        let f = self.func("sig_mul_f16out_f32");
26887        let cfg = LaunchConfig::for_num_elems(n as u32);
26888        let ni = n as i32;
26889        let __s_b = self.gpu.stream();
26890        let mut b = __s_b.launch_builder(&f);
26891        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26892        unsafe {
26893            b.launch(cfg)?;
26894        }
26895        Ok(())
26896    }
26897
26898    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
26899    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
26900    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
26901    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
26902    ///
26903    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
26904    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
26905    /// applies the wrong number of distinct gate values.
26906    #[allow(clippy::too_many_arguments)]
26907    pub fn attn_head_gate(
26908        &self,
26909        a: &CudaSlice<f32>,
26910        g: &CudaSlice<f32>,
26911        dst: &mut CudaSlice<f32>,
26912        dst16: Option<&mut CudaSlice<u8>>,
26913        head_dim: usize,
26914        n_head: usize,
26915        t: usize,
26916    ) -> Result<(), Box<dyn std::error::Error>> {
26917        let f = self.func("attn_head_gate_f32");
26918        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26919        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26920        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
26921        let d16: u64 = match dst16 {
26922            Some(d) => self.addr_u8(d),
26923            None => 0,
26924        };
26925        let __s_b = self.gpu.stream();
26926        let mut b = __s_b.launch_builder(&f);
26927        b.arg(a)
26928            .arg(g)
26929            .arg(dst)
26930            .arg(&d16)
26931            .arg(&hd)
26932            .arg(&nh)
26933            .arg(&ti);
26934        unsafe {
26935            b.launch(cfg)?;
26936        }
26937        Ok(())
26938    }
26939
26940    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
26941    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
26942    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
26943    ///
26944    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
26945    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
26946    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
26947    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
26948    #[allow(clippy::too_many_arguments)]
26949    pub fn swiglu_clamped_mul_scaled(
26950        &self,
26951        gate: &CudaSlice<f32>,
26952        up: &CudaSlice<f32>,
26953        gs: f32,
26954        us: f32,
26955        limit: f32,
26956        dst: &mut CudaSlice<f32>,
26957        n: usize,
26958    ) -> Result<(), Box<dyn std::error::Error>> {
26959        debug_assert!(
26960            limit > 1e-6,
26961            "swiglu_clamped needs a live limit; use silu_mul_scaled"
26962        );
26963        let f = self.func("swiglu_clamped_mul_scaled_f32");
26964        let cfg = LaunchConfig::for_num_elems(n as u32);
26965        let ni = n as i32;
26966        let __s_b = self.gpu.stream();
26967        let mut b = __s_b.launch_builder(&f);
26968        b.arg(gate)
26969            .arg(up)
26970            .arg(&gs)
26971            .arg(&us)
26972            .arg(&limit)
26973            .arg(dst)
26974            .arg(&ni);
26975        unsafe {
26976            b.launch(cfg)?;
26977        }
26978        Ok(())
26979    }
26980
26981    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
26982    pub fn gated_rmsnorm(
26983        &self,
26984        o: &CudaSlice<f32>,
26985        w: &CudaSlice<f32>,
26986        z: &CudaSlice<f32>,
26987        dst: &mut CudaSlice<f32>,
26988        ncols: usize,
26989        nrows: usize,
26990        eps: f32,
26991    ) -> Result<(), Box<dyn std::error::Error>> {
26992        let f = self.func("gated_rmsnorm_f32");
26993        let cfg = LaunchConfig {
26994            grid_dim: (nrows as u32, 1, 1),
26995            block_dim: (128, 1, 1),
26996            shared_mem_bytes: 0,
26997        };
26998        let (nc, e) = (ncols as i32, eps);
26999        let __s_b = self.gpu.stream();
27000        let mut b = __s_b.launch_builder(&f);
27001        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27002        unsafe {
27003            b.launch(cfg)?;
27004        }
27005        Ok(())
27006    }
27007
27008    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
27009    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
27010    pub fn gated_rmsnorm_f16out(
27011        &self,
27012        o: &CudaSlice<f32>,
27013        w: &CudaSlice<f32>,
27014        z: &CudaSlice<f32>,
27015        dst: &mut CudaSlice<f32>,
27016        dst16: &mut CudaSlice<u8>,
27017        ncols: usize,
27018        nrows: usize,
27019        eps: f32,
27020    ) -> Result<(), Box<dyn std::error::Error>> {
27021        let f = self.func("gated_rmsnorm_f16out_f32");
27022        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
27023        let cfg = LaunchConfig {
27024            grid_dim: (nrows as u32, 1, 1),
27025            block_dim: (128, 1, 1),
27026            shared_mem_bytes: 0,
27027        };
27028        let (nc, e) = (ncols as i32, eps);
27029        let __s_b = self.gpu.stream();
27030        let mut b = __s_b.launch_builder(&f);
27031        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27032        unsafe {
27033            b.launch(cfg)?;
27034        }
27035        Ok(())
27036    }
27037
27038    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
27039    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
27040    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
27041    #[allow(clippy::too_many_arguments)]
27042    pub fn add_rms_norm_zq8(
27043        &self,
27044        a: &CudaSlice<f32>,
27045        b_in: &CudaSlice<f32>,
27046        w: &CudaSlice<f32>,
27047        res: &mut CudaSlice<f32>,
27048        z: &mut CudaSlice<f32>,
27049        ncols: usize,
27050        nrows: usize,
27051        eps: f32,
27052    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27053        assert!(ncols % 32 == 0);
27054        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27055        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27056        let f = self.func("add_rms_norm_zq8");
27057        let cfg = LaunchConfig {
27058            grid_dim: (nrows as u32, 1, 1),
27059            block_dim: (1024, 1, 1),
27060            shared_mem_bytes: 0,
27061        };
27062        let (nc, ep) = (ncols as i32, eps);
27063        let __s_b = self.gpu.stream();
27064        let mut b = __s_b.launch_builder(&f);
27065        b.arg(a)
27066            .arg(b_in)
27067            .arg(w)
27068            .arg(res)
27069            .arg(z)
27070            .arg(&mut q)
27071            .arg(&mut d)
27072            .arg(&nc)
27073            .arg(&ep);
27074        unsafe {
27075            b.launch(cfg)?;
27076        }
27077        Ok((q, d))
27078    }
27079
27080    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
27081    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
27082    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
27083    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
27084    pub fn gated_rmsnorm_zv(
27085        &self,
27086        o: &CudaSlice<f32>,
27087        w: &CudaSlice<f32>,
27088        z: &cudarc::driver::CudaView<f32>,
27089        dst: &mut CudaSlice<f32>,
27090        ncols: usize,
27091        nrows: usize,
27092        eps: f32,
27093    ) -> Result<(), Box<dyn std::error::Error>> {
27094        let f = self.func("gated_rmsnorm_f32");
27095        let cfg = LaunchConfig {
27096            grid_dim: (nrows as u32, 1, 1),
27097            block_dim: (128, 1, 1),
27098            shared_mem_bytes: 0,
27099        };
27100        let (nc, e) = (ncols as i32, eps);
27101        let __s_b = self.gpu.stream();
27102        let mut b = __s_b.launch_builder(&f);
27103        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27104        unsafe {
27105            b.launch(cfg)?;
27106        }
27107        Ok(())
27108    }
27109
27110    pub fn gated_rmsnorm_f16out_zv(
27111        &self,
27112        o: &CudaSlice<f32>,
27113        w: &CudaSlice<f32>,
27114        z: &cudarc::driver::CudaView<f32>,
27115        dst: &mut CudaSlice<f32>,
27116        dst16: &mut CudaSlice<u8>,
27117        ncols: usize,
27118        nrows: usize,
27119        eps: f32,
27120    ) -> Result<(), Box<dyn std::error::Error>> {
27121        let f = self.func("gated_rmsnorm_f16out_f32");
27122        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
27123        let cfg = LaunchConfig {
27124            grid_dim: (nrows as u32, 1, 1),
27125            block_dim: (128, 1, 1),
27126            shared_mem_bytes: 0,
27127        };
27128        let (nc, e) = (ncols as i32, eps);
27129        let __s_b = self.gpu.stream();
27130        let mut b = __s_b.launch_builder(&f);
27131        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27132        unsafe {
27133            b.launch(cfg)?;
27134        }
27135        Ok(())
27136    }
27137
27138    pub fn gated_rmsnorm_q8_1(
27139        &self,
27140        o: &CudaSlice<f32>,
27141        w: &CudaSlice<f32>,
27142        z: &CudaSlice<f32>,
27143        ncols: usize,
27144        nrows: usize,
27145        eps: f32,
27146    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27147        assert!(ncols % 32 == 0);
27148        let f = self.func("gated_rmsnorm_q8_1");
27149        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27150        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27151        let cfg = LaunchConfig {
27152            grid_dim: (nrows as u32, 1, 1),
27153            block_dim: (128, 1, 1),
27154            shared_mem_bytes: 0,
27155        };
27156        let (nc, ep) = (ncols as i32, eps);
27157        let __s_b = self.gpu.stream();
27158        let mut b = __s_b.launch_builder(&f);
27159        b.arg(o)
27160            .arg(w)
27161            .arg(z)
27162            .arg(&mut out_q)
27163            .arg(&mut out_d)
27164            .arg(&nc)
27165            .arg(&ep);
27166        unsafe {
27167            b.launch(cfg)?;
27168        }
27169        Ok((out_q, out_d))
27170    }
27171
27172    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
27173    pub fn transpose(
27174        &self,
27175        inp: &CudaSlice<f32>,
27176        rows: usize,
27177        cols: usize,
27178    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27179        let f = self.func("transpose_f32");
27180        let mut out = self.zeros(rows * cols)?;
27181        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27182        let (r, c) = (rows as i32, cols as i32);
27183        let __s_b = self.gpu.stream();
27184        let mut b = __s_b.launch_builder(&f);
27185        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27186        unsafe {
27187            b.launch(cfg)?;
27188        }
27189        Ok(out)
27190    }
27191
27192    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
27193    pub fn repeat_heads(
27194        &self,
27195        inp: &CudaSlice<f32>,
27196        out: &mut CudaSlice<f32>,
27197        head_dim: usize,
27198        n_in: usize,
27199        n_out: usize,
27200        t: usize,
27201    ) -> Result<(), Box<dyn std::error::Error>> {
27202        let f = self.func("repeat_heads_f32");
27203        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27204        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27205        let __s_b = self.gpu.stream();
27206        let mut b = __s_b.launch_builder(&f);
27207        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27208        unsafe {
27209            b.launch(cfg)?;
27210        }
27211        Ok(())
27212    }
27213
27214    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
27215    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
27216    ///
27217    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
27218    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
27219    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
27220    pub fn q_gate_split(
27221        &self,
27222        qf: &CudaSlice<f32>,
27223        q_out: &mut CudaSlice<f32>,
27224        gate_out: &mut CudaSlice<f32>,
27225        head_dim: usize,
27226        n_head: usize,
27227        t: usize,
27228    ) -> Result<(), Box<dyn std::error::Error>> {
27229        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27230        let out_need = head_dim * n_head * t;
27231        if q_out.len() < out_need || gate_out.len() < out_need {
27232            return Err(format!(
27233                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27234                q_out.len(),
27235                gate_out.len()
27236            )
27237            .into());
27238        }
27239        let f = self.func("q_gate_split_f32");
27240        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27241        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27242        let __s_b = self.gpu.stream();
27243        let mut b = __s_b.launch_builder(&f);
27244        b.arg(qf)
27245            .arg(q_out)
27246            .arg(gate_out)
27247            .arg(&hd)
27248            .arg(&nh)
27249            .arg(&ti);
27250        unsafe {
27251            b.launch(cfg)?;
27252        }
27253        Ok(())
27254    }
27255
27256    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
27257    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
27258    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
27259    pub fn qkv_to_gdn_repack(
27260        &self,
27261        conv_out: &CudaSlice<f32>,
27262        q_g: &mut CudaSlice<f32>,
27263        k_g: &mut CudaSlice<f32>,
27264        v_g: &mut CudaSlice<f32>,
27265        d_state: usize,
27266        num_v: usize,
27267        num_k: usize,
27268        key_dim: usize,
27269        t: usize,
27270    ) -> Result<(), Box<dyn std::error::Error>> {
27271        let f = self.func("qkv_to_gdn_repack_f32");
27272        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27273        let (ds, nv, nk, kd, ti) = (
27274            d_state as i32,
27275            num_v as i32,
27276            num_k as i32,
27277            key_dim as i32,
27278            t as i32,
27279        );
27280        let __s_b = self.gpu.stream();
27281        let mut b = __s_b.launch_builder(&f);
27282        b.arg(conv_out)
27283            .arg(q_g)
27284            .arg(k_g)
27285            .arg(v_g)
27286            .arg(&ds)
27287            .arg(&nv)
27288            .arg(&nk)
27289            .arg(&kd)
27290            .arg(&ti);
27291        unsafe {
27292            b.launch(cfg)?;
27293        }
27294        Ok(())
27295    }
27296
27297    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
27298    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
27299    pub fn conv_left_pad(
27300        &self,
27301        src: &CudaSlice<f32>,
27302        dst: &mut CudaSlice<f32>,
27303        conv_dim: usize,
27304        t: usize,
27305        pad: usize,
27306    ) -> Result<(), Box<dyn std::error::Error>> {
27307        let f = self.func("conv_left_pad_f32");
27308        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27309        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27310        let __s_b = self.gpu.stream();
27311        let mut b = __s_b.launch_builder(&f);
27312        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27313        unsafe {
27314            b.launch(cfg)?;
27315        }
27316        Ok(())
27317    }
27318
27319    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
27320    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
27321    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
27322    pub fn conv_assemble_and_roll(
27323        &self,
27324        qkv_col: &CudaSlice<f32>,
27325        conv_state: &mut CudaSlice<f32>,
27326        conv_in: &mut CudaSlice<f32>,
27327        conv_dim: usize,
27328        pad: usize,
27329    ) -> Result<(), Box<dyn std::error::Error>> {
27330        let f = self.func("conv_assemble_and_roll_f32");
27331        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27332        let (cd, p) = (conv_dim as i32, pad as i32);
27333        let __s_b = self.gpu.stream();
27334        let mut b = __s_b.launch_builder(&f);
27335        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27336        unsafe {
27337            b.launch(cfg)?;
27338        }
27339        Ok(())
27340    }
27341
27342    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
27343    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
27344    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
27345    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
27346    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
27347    pub fn ssm_conv1d_fused_decode(
27348        &self,
27349        qkv_col: &CudaSlice<f32>,
27350        conv_state: &mut CudaSlice<f32>,
27351        w: &CudaSlice<f32>,
27352        conv_out: &mut CudaSlice<f32>,
27353        conv_dim: usize,
27354        d_conv: usize,
27355    ) -> Result<(), Box<dyn std::error::Error>> {
27356        let f = self.func("ssm_conv1d_fused_decode_f32");
27357        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27358        let (cd, dc) = (conv_dim as i32, d_conv as i32);
27359        let __s_b = self.gpu.stream();
27360        let mut b = __s_b.launch_builder(&f);
27361        b.arg(qkv_col)
27362            .arg(conv_state)
27363            .arg(w)
27364            .arg(conv_out)
27365            .arg(&cd)
27366            .arg(&dc);
27367        unsafe {
27368            b.launch(cfg)?;
27369        }
27370        Ok(())
27371    }
27372
27373    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
27374    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
27375    pub fn slice_range(
27376        &self,
27377        src: &CudaSlice<f32>,
27378        start: usize,
27379        len: usize,
27380    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27381        let host = self.gpu.stream().clone_dtoh(src)?;
27382        self.gpu.stream().synchronize()?;
27383        Ok(self.htod(&host[start..start + len])?)
27384    }
27385}
27386
27387#[cfg(test)]
27388mod target_dispatch_tests {
27389    use super::legacy_quant_gemm_allowed;
27390
27391    #[test]
27392    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27393        // sm_120a native lane
27394        assert!(legacy_quant_gemm_allowed(false, false, false));
27395        assert!(!legacy_quant_gemm_allowed(false, false, true));
27396        // pure portable lane (sm_89): gated
27397        assert!(!legacy_quant_gemm_allowed(true, false, false));
27398        assert!(!legacy_quant_gemm_allowed(true, false, true));
27399        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
27400        assert!(legacy_quant_gemm_allowed(true, true, false));
27401        assert!(!legacy_quant_gemm_allowed(true, true, true));
27402    }
27403
27404    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27405    #[test]
27406    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27407        assert!(!legacy_quant_gemm_allowed(
27408            cfg!(memra_portable_cuda),
27409            cfg!(memra_hopper_mma),
27410            false
27411        ));
27412    }
27413
27414    #[cfg(memra_hopper_mma)]
27415    #[test]
27416    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27417        assert!(legacy_quant_gemm_allowed(
27418            cfg!(memra_portable_cuda),
27419            cfg!(memra_hopper_mma),
27420            false
27421        ));
27422        assert!(super::portable_mma_gated() == false);
27423    }
27424}
27425
27426/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
27427/// inherent methods (inherent methods win name resolution, so no recursion).
27428impl memra_kv::KvDev for Engine {
27429    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27430        Engine::zeros(self, n)
27431    }
27432    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27433        Engine::uninit(self, n)
27434    }
27435    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27436        Engine::alloc_u8(self, n)
27437    }
27438    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27439        Engine::htod_i32(self, v)
27440    }
27441    fn clone_dtod(
27442        &self,
27443        src: &CudaSlice<f32>,
27444    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27445        Engine::clone_dtod(self, src)
27446    }
27447    fn copy_into(
27448        &self,
27449        dst: &mut CudaSlice<f32>,
27450        off: usize,
27451        src: &CudaSlice<f32>,
27452        len: usize,
27453    ) -> Result<(), Box<dyn std::error::Error>> {
27454        Engine::copy_into(self, dst, off, src, len)
27455    }
27456    fn set_i32_one(
27457        &self,
27458        d: &mut CudaSlice<i32>,
27459        v: i32,
27460    ) -> Result<(), Box<dyn std::error::Error>> {
27461        Engine::set_i32_one(self, d, v)
27462    }
27463}
27464
27465#[cfg(test)]
27466mod fused_gate_bounds_tests {
27467    use super::*;
27468
27469    /// The fused `[q|gate]` split's read-site guard, on the device.
27470    ///
27471    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
27472    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
27473    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
27474    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
27475    /// `FusedQGateExtent` before the launch.
27476    ///
27477    /// Catch demonstration for this test (guard temporarily removed, then restored):
27478    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
27479    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
27480    /// the call returns `Err`. Receipt in the lane report.
27481    #[test]
27482    #[ignore = "requires a CUDA GPU"]
27483    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27484        let e = Engine::new(0).unwrap();
27485        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27486        let fused = 2 * head_dim * n_head * t;
27487        let out_n = head_dim * n_head * t;
27488
27489        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
27490        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27491        let mut q = e.uninit(out_n).unwrap();
27492        let mut gate = e.uninit(out_n).unwrap();
27493        let err = e
27494            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27495            .expect_err("half-width wq must be refused, not read past")
27496            .to_string();
27497        assert!(err.contains("NO fused gate"), "{err}");
27498        assert!(err.contains(&format!("{fused}")), "{err}");
27499
27500        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
27501        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
27502        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27503        let wide = e.htod(&host).unwrap();
27504        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27505            .expect("full-width wq splits");
27506        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27507        for tok in 0..t {
27508            for hh in 0..n_head {
27509                for d in 0..head_dim {
27510                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27511                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27512                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27513                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27514                }
27515            }
27516        }
27517
27518        // undersized destinations are refused too (the other half of the extent contract)
27519        let mut small = e.uninit(out_n - 1).unwrap();
27520        assert!(
27521            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27522                .is_err()
27523        );
27524    }
27525}
27526
27527/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
27528/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
27529/// any launch, so the refusal is testable without a device.
27530#[cfg(test)]
27531mod fused_rope_width_tests {
27532    use super::Engine;
27533
27534    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
27535    /// safetensors route derives the same), which is why the fusion is legal there today.
27536    #[test]
27537    fn full_width_is_accepted() {
27538        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27539        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27540        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27541    }
27542
27543    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
27544    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
27545    ///
27546    /// ```text
27547    /// attention.key_length     512   rope.dimension_count     512   (global class)
27548    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
27549    /// ```
27550    ///
27551    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
27552    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
27553    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
27554    /// instead of a silently over-rotated head.
27555    #[test]
27556    fn gemma4_official_artifact_widths_pass() {
27557        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27558        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27559    }
27560
27561    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
27562    /// with no `n_dims`, silently rotating the pass-through band.
27563    #[test]
27564    fn partial_rotary_is_refused_with_the_geometry_named() {
27565        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
27566        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27567            .expect_err("partial rotary must refuse");
27568        let msg = err.to_string();
27569        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27570        assert!(msg.contains("n_rot 64"), "{msg}");
27571        assert!(msg.contains("head_dim 256"), "{msg}");
27572        assert!(
27573            msg.contains("64..256"),
27574            "names the band it would corrupt: {msg}"
27575        );
27576        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
27577        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27578        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
27579        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27580    }
27581}