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 hyper;
62pub mod model;
63pub mod sigrouter_contract;
64pub mod vision;
65pub mod vision_gemma;
66pub mod vision_glm5;
67pub mod vision_pre;
68pub mod vision_step;
69/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
70/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
71pub mod cache {
72    pub use memra_kv::*;
73}
74pub mod decode;
75pub mod decode_batch;
76pub mod dflash;
77pub mod eagle;
78/// Measured expert-placement map (`MEMRA_EP_MAP`; glm5 alias honored) — the fail-closed
79/// `memra-ep-map-v1` reader every family's EP shard builders consume (fleet-shared by
80/// design; glm5 is the first consumer). (LAW:coactivation-expert-placement; maps are
81/// minted by the shared fleet tool from `MEMRA_MOE_WEIGHT_TRACE` traces). No CUDA deps.
82pub mod ep_map;
83pub mod gemma_spec;
84pub mod glm5_tp;
85/// glm5_next T-parallel speculative verify: the rows-walk verify, per-step KDA state-column
86/// rollback, latent/kpool truncation, and the MEMRA_GLM5_SPEC-gated draft->verify->rollback
87/// loop over the native MTP head (lane/glm5-tparallel-verify).
88pub mod glm_spec;
89pub mod graph_update;
90pub mod kda;
91/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
92/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
93/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
94pub mod mla;
95pub mod mla_ffi;
96pub mod moesd;
97pub mod parallel;
98pub mod plan_backend;
99pub mod pp;
100/// qwen4_exp (Qwen3.8-Flash-Next) GPU eager forward — onboarding phase 7, correctness arm
101/// gated against memra-reference (research/qwen4exp-bringup-20260829/GPU-EAGER.md).
102pub mod qwen4exp_gpu;
103pub mod round_stream;
104pub mod spec;
105/// Per-burst spec-round phase attribution (`MEMRA_SPEC_TRACE`; glm5 alias honored) —
106/// the draft/verify/accept/rollback/maintenance split every spec family owns, with
107/// caller-tagged emit lines so banked receipts keep their grep shape. No CUDA deps
108/// beyond the stream drains at phase boundaries.
109pub mod spec_phase;
110pub mod tp;
111pub mod tp_transport;
112pub use memra_sampling as sampler;
113
114/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
115/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
116/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
117/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
118/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
119///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
120///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
121///                     stream sync per projection (round-47 ledgered defect).
122///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
123///                     construction, zero syncs, f32 C with the act row-scale folded in.
124/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
125/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
126/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
127/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
128/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
129/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
130///
131/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
132/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
133/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
134/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
135/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
136/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
137/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
138/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
139///
140/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
141/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
142/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
143/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
144/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
145/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
146/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
147///
148/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
149/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
150/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
151/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
152/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
153/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
154/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
155/// the k-quant-only admission survives as the rollback seam, not the default.
156/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
157pub fn moe_f16g_mode() -> u8 {
158    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
159    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
160        Ok("0") => 0,
161        Ok("2") => 2,
162        Ok("3") => 3,
163        Ok(_) => 1,
164        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
165        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
166        Err(_) => 2,
167    })
168}
169/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
170/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
171/// (shape_sel, cross) for the FFI:
172///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
173///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
174///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
175///                         back to 32x64 in-launcher when the device/in_f can't take it).
176///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
177///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
178///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
179///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
180///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
181///                         verdict was stale).
182pub fn moe_f16g_sk_params() -> (i32, i32) {
183    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
184    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
185        Ok("0") => (-1, 0),
186        Ok("32") => (0, i32::MAX),
187        Ok("128") => (0, 1),
188        _ => {
189            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
190                .ok()
191                .and_then(|v| v.parse().ok())
192                .unwrap_or(64);
193            (0, cross)
194        }
195    })
196}
197/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
198/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
199/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
200/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
201/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
202/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
203/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
204/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
205/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
206/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
207pub fn moe_f16g_direct_on(qtype: i32) -> bool {
208    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
209    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
210        Ok("0") => 0,
211        Ok("kq") => 1,
212        _ => 2,
213    });
214    match m {
215        0 => false,
216        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
217        _ => true,
218    }
219}
220/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
221/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
222/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
223/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
224/// stage under q35's routing skew. Bit-identical to every other sk form by construction
225/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
226/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
227/// tail. in_f % 64 != 0 falls back in-launcher.
228pub fn moe_f16g_tail_on() -> bool {
229    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
230    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
231}
232
233/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
234/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
235/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
236/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
237/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
238/// still opens this door for A/B.
239pub fn moe_f16g_gemma_on() -> bool {
240    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
241    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
242}
243
244/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
245/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
246/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
247pub fn moe_fuse_actq_on() -> bool {
248    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
249    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
250}
251
252/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
253/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
254/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
255/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
256/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
257/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
258/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
259/// verify already use (dispatch parity, one router kernel for every t).
260/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
261pub fn router_prefill_exact_on() -> bool {
262    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
263    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
264}
265
266pub fn router_kernel_on() -> bool {
267    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
268    *ON.get_or_init(|| {
269        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
270        if !on {
271            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
272        }
273        on
274    })
275}
276
277/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
278/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
279/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
280/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
281/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
282/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
283/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
284/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
285/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
286/// seam, perf-only: bits are equal by the kernel-check gate).
287/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
288/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
289/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
290/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
291pub const ROUTER_BATCH_MIN_T: usize = 8;
292pub fn router_batch_on() -> bool {
293    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
294    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
295}
296mod cpu_experts;
297#[cfg(memra_cutlass)]
298pub mod cutlass_ffi;
299pub mod dsv4_ffi;
300pub mod dsv4_gpu;
301pub mod f16_ffi;
302pub mod fp8_ffi;
303pub mod mmq_ffi;
304pub mod moe_cache;
305pub mod prime_graph;
306pub mod spill;
307mod spill_pread;
308
309// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
310// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
311// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
312// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
313// broke every machine that wasn't the build machine. Same bytes, same module image;
314// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
315const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
316const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
317/// kda.cu: the glm5_next Kimi Delta Attention mixer (per-channel-decay delta rule).
318const KDA_FATBIN: &[u8] = include_bytes!(env!("MEMRA_KDA_FATBIN"));
319const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
320const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
321const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
322const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
323/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
324const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
325
326/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
327/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
328/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
329/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
330/// compile-time default (zero behavior change).
331fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
332    assert!(
333        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
334        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
335    );
336    match std::env::var("MEMRA_GEMM_FATBIN") {
337        Ok(path) => std::borrow::Cow::Owned(
338            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
339        ),
340        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
341    }
342}
343
344/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
345/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
346/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
347/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
348/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
349/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
350pub(crate) const fn portable_mma_gated() -> bool {
351    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
352}
353
354/// Refuse an env force that would reach a kernel THIS BUILD DOES NOT CONTAIN.
355///
356/// Doors of the shape `MEMRA_X=1 => true` are arch-blind: they were written so an operator could
357/// force a promoted path on, and the default arm (`cfg!(memra_hopper_mma)` or similar) is the only
358/// thing that consulted the arch. On a portable build the forced path then reaches
359/// `Engine::func`, which resolves lazily and ends in `panic!("kernel {name} not in any fatbin")` —
360/// a confusing crash naming a kernel the operator never heard of, several frames from the switch
361/// they actually flipped.
362///
363/// Found 2026-08-23 by tools/fatbin-lookup-census.py, which listed 20 looked-up kernels absent
364/// from the sm_89 fatbins. 18 of those turned out to be correctly unreachable (the GDN varlen
365/// chain is gated through `gdn_mma_enabled`, which starts with `!portable_mma_gated()`); these
366/// env doors were the two that were genuinely reachable, and only by explicit operator action.
367///
368/// Same shape and same message style as `gemm_fatbin_bytes`'s refusal above — one idiom for
369/// "this switch cannot work on this build", so it fails at the switch instead of at the lookup.
370#[track_caller]
371pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
372    assert!(
373        !portable_mma_gated(),
374        "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
375         portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
376    );
377}
378
379/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
380/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
381/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
382/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
383/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
384/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
385/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
386/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
387pub(crate) const fn gdn_mma_default_on() -> bool {
388    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
389}
390
391/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
392const fn konst_eq(a: &str, b: &str) -> bool {
393    let (a, b) = (a.as_bytes(), b.as_bytes());
394    if a.len() != b.len() {
395        return false;
396    }
397    let mut i = 0;
398    while i < a.len() {
399        if a[i] != b[i] {
400            return false;
401        }
402        i += 1;
403    }
404    true
405}
406
407/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
408/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
409/// in a pure helper so the dispatch guard can be regression-tested without constructing an
410/// Engine or allocating a GPU tensor.
411const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
412    (!portable_cuda || hopper_mma) && !no_gemm
413}
414
415// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
416// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
417// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
418// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
419// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
420// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
421// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
422const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
423const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
424const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
425const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
426const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
427
428/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
429/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
430pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
431
432/// The flash_attn fatbin matching the selected KV formats.
433fn flash_fatbin_bytes() -> &'static [u8] {
434    match kv_cache_formats() {
435        ("q8_0", "q5_1") => FLASH_FATBIN,
436        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
437        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
438        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
439        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
440        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
441        other => unreachable!("kv_cache_formats returned {other:?}"),
442    }
443}
444
445/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
446/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
447/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
448/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
449/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
450/// defaults (zero behavior change).
451fn k1_launch_override() -> Option<(u32, u32, u32)> {
452    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
453    *K1.get_or_init(|| {
454        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
455        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
456        match p.as_slice() {
457            [bm, bn, w] => Some((*bm, *bn, *w)),
458            _ => None,
459        }
460    })
461}
462
463/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
464/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
465/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
466/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
467/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
468/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
469pub(crate) fn wgmma_gemm_enabled() -> bool {
470    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
471    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
472}
473
474/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
475/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
476/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
477/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
478/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
479/// the split count changes the combine's FP summation order, and the spec verify's batched forward
480/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
481/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
482/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
483/// adaptive retries (any retry MUST pass run-spec self-consistency first).
484/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
485/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
486/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
487/// between eager decode and the verify (the spec-exactness law).
488pub const FA_VEC_MIN_TKV: usize = 96;
489/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
490/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
491/// which moves the crossover — sweep per model, adopt per the battery.
492pub fn fa_vec_min_tkv() -> usize {
493    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
494    *V.get_or_init(|| {
495        std::env::var("MEMRA_FA_VEC_MIN")
496            .ok()
497            .and_then(|v| v.parse().ok())
498            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
499    })
500}
501
502/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
503/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
504/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
505///
506/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
507/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
508/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
509/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
510/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
511/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
512pub fn fa_f16pv_on() -> bool {
513    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
514    *ON.get_or_init(|| {
515        std::env::var("MEMRA_FA_F16PV")
516            .map(|v| v != "0")
517            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
518    })
519}
520
521/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
522/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
523/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
524pub fn fa512_hp_on() -> bool {
525    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
526    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
527}
528
529/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
530/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
531/// accumulation. Even n_head and even GQA group required (guarded per call).
532pub fn faw_hp_on() -> bool {
533    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
534    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
535}
536
537/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
538/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
539/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
540pub fn fa512_wide_warps() -> usize {
541    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
542    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
543        Ok("1") => 4,
544        _ => 2,
545    })
546}
547
548/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
549/// and the gemma global-layer rows/parity call sites.
550pub fn fa512_min_tkv() -> usize {
551    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
552    *FA512_MIN.get_or_init(|| {
553        std::env::var("MEMRA_FA512_MIN")
554            .ok()
555            .and_then(|v| v.parse().ok())
556            .unwrap_or(512)
557    })
558}
559/// Per-model crossover default, set at model load BEFORE the first decode (per-model
560/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
561/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
562pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
563    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
564/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
565/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
566/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
567pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
568/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
569/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
570/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
571/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
572/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
573pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
574    std::sync::atomic::AtomicBool::new(false);
575/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
576/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
577/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
578/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
579/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
580/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
581pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
582    std::sync::atomic::AtomicBool::new(true);
583pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
584    std::sync::atomic::AtomicUsize::new(16);
585/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
586/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
587/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
588/// latency-bound at 256 threads — 7us/launch measured).
589pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
590/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
591pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
592/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
593/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
594/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
595/// explicit numerical-form seam. mmq_ffi reads this before the env.
596pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
597/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
598/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
599pub use memra_kv::KV_FP8_FORCE;
600/// bf16 matvec family block size (MEMRA_MMV_BLOCK, default 128, clamped to [64, 256] and a
601/// multiple of 32 — the f32acc twin's shared reduce caps at 256). NUMERIC-CLASS knob: the
602/// per-thread stride and reduction order change with the block, same acceptance class as
603/// MEMRA_RMS_BLOCK (fresh-tape identity + battery at the pinned value).
604pub(crate) fn mmv_block() -> u32 {
605    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
606    *V.get_or_init(|| {
607        std::env::var("MEMRA_MMV_BLOCK")
608            .ok()
609            .and_then(|v| v.parse().ok())
610            .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
611            .unwrap_or(128)
612    })
613}
614
615/// MEMRA_STEP_TP_W8=1: q8_0 mirror of the step TP attention projections for DECODE.
616///
617/// NUMERIC-CLASS door, same class and acceptance as `MEMRA_STEP_TP_QKV_FUSED` /
618/// `MEMRA_BF16_MMV`: the per-row arithmetic becomes an int8 dp4a dot
619/// with per-32 scales instead of a bf16xf32 fma chain, so a bit-tape cannot apply and the
620/// acceptance is the argmax gate plus the boot battery. Motivation is measured, not assumed
621/// (`decode-kernel-census`, 2026-08-25): the fused qkv shape runs 23.0 us in bf16 at
622/// 1.83 TB/s and 14.0 us in q8_0 at 1.60, and o_proj 24.2 -> 11.7 us — together
623/// ~-1.0 ms of a 13.16 ms token. Default OFF.
624/// MEMRA_W8_HYBRID=1 opts the door's HYBRID half in (LM head, shared expert, dense FFN).
625/// Default OFF on measurement AND on residency: it moved decode +0.1% (the W8 trace showed it
626/// only ever mirrored the shexp down rows, which SHEXP_OVERLAP already hides), while costing
627/// ~1.7 GB per card on top of the attention mirrors' ~0.9 GB — and at the model's NATURAL
628/// 262144-token context the full set does not fit: `MEMRA_STEP_TP_W8=1` there dies in
629/// CUDA_ERROR_OUT_OF_MEMORY while plain decode runs 76.03 tok/s.
630/// STEP37 SERVING DEFAULTS (owner flip, 2026-08-27). The step37 serving shape — the t-row walk,
631/// the q8 W8 doors, the SWA ring, the NVFP4 draft heads, and this lane's three verify fixes —
632/// was gated door by door (byte tape == plain, acceptance unchanged, run-spec K=1..8 PASS,
633/// interleaved x5 wall, vendor-default sampled cell with engagement receipts: greedy 93.18 vs
634/// 81.95 plain, sampled 81.79 vs 78.50) and the owner ordered the defaults ON. The doors' call
635/// sites are not all family-scoped (the W8 mirror routing sits inside generic matmul paths), so
636/// the default arms AT MODEL LOAD when the plan compiles to the SlidingGatedMoe program, never
637/// globally. Every door keeps a per-flag env override: `=1` forces ON for any family, `=0` is
638/// the kill switch — the rollback seam the FLAGS rows name. Per-process: a process that loads a
639/// step37-class model arms the defaults for its lifetime.
640static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
641    std::sync::atomic::AtomicBool::new(false);
642
643pub fn arm_step37_serving_defaults() {
644    STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
645    crate::cache::set_swa_ring_default(true);
646    eprintln!(
647        "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
648         (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
649    );
650}
651
652pub(crate) fn step37_defaults_armed() -> bool {
653    STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
654}
655
656/// Tri-state door: `=1` ON, `=0` OFF, unset = the family default (ON once a step37-class model
657/// armed it, OFF otherwise). The env parse is cached; the family default is read live because
658/// arming happens at model load, possibly after another door's first read.
659pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
660    match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
661        Some("1") => Some(true),
662        Some("0") => Some(false),
663        _ => None,
664    }) {
665        Some(forced) => forced,
666        None => step37_defaults_armed(),
667    }
668}
669
670pub(crate) fn w8_hybrid_on() -> bool {
671    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
672    step37_door(&ENV, "MEMRA_W8_HYBRID")
673}
674
675pub(crate) fn step_tp_w8_on() -> bool {
676    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
677    step37_door(&ENV, "MEMRA_STEP_TP_W8")
678}
679
680/// MEMRA_W8_VIEW=1: extend the W8 hybrid half to the ROW-RANGE-VIEW GEMVs, i.e. the lo halves
681/// that `MEMRA_HEAD_SPLIT` and `MEMRA_SHEXP_OVERLAP` keep on rank 0. NOT a step37 family door
682/// and NOT armed by `arm_step37_serving_defaults`: it stays off until it carries its own
683/// interleaved speed rows and its own argmax gate. Unset or `=0` is the rollback seam.
684pub(crate) fn w8_view_on() -> bool {
685    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
686    *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
687}
688
689/// MEMRA_Q8T_WONCE=1: the q8 t-column verify kernels take their weight-once `_tw` twins — one
690/// row grid, each weight int4 loaded once and dotted against all t columns — instead of the `_t`
691/// forms, whose column grid axis plus __ldcs (streaming, evict-first) re-reads the fully-shared
692/// weights from DRAM once per column (nsys 2026-08-27: qkv_rp_t 1.67x, b4_rp_t 1.43x a
693/// single-column call for 2 columns, where weight-bound scaling says ~1.1x). Per-column float
694/// program unchanged (same lane-strided blk order, own accumulator chain, same reduce); default
695/// off until the byte tape says so.
696/// MEMRA_STEP_GEMM_PRIME: prime chunks (t>=16) route the routed MoE through the grouped f16 GEMM
697/// over the resident NVFP4 banks instead of the per-token device routes. FAMILY-DEFAULT ON since
698/// 2026-08-28 because on the server route it is the only prime that WORKS: measured there, walk
699/// = ERR (tail chunk missing from the distributed kv), fallback chunked prime = 29 s on a
700/// ~450-token prompt and a 90 s TIMEOUT at 4k, grouped GEMM = 3.5-4.9 s with coherent output.
701/// `=0` is the kill switch back to the fallback prime.
702pub(crate) fn step_gemm_prime_on() -> bool {
703    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
704    step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
705}
706
707pub(crate) fn q8t_wonce_on() -> bool {
708    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
709    step37_door(&ENV, "MEMRA_Q8T_WONCE")
710}
711
712/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
713/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
714/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
715/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
716pub(crate) fn sig_expf_dev_on() -> bool {
717    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
718    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
719}
720
721pub(crate) fn topk_fast_on() -> bool {
722    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
723    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
724}
725
726/// Select the sigmoid-router kernel without ever sending a shape wider than the fast
727/// kernels' fixed eight-pick scratch. The generic and dexp kernels support the full host
728/// contract; both `_fast` twins index `[warp][8]` storage and would write out of bounds for
729/// `n_used > 8` (Hermes `0d220d8c9a3eb634`).
730fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
731    match (sig_expf, fast && n_used <= 8) {
732        (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
733        (true, false) => "moe_router_sigmoid_topk_f32_dexp",
734        (false, true) => "moe_router_sigmoid_topk_f32_fast",
735        (false, false) => "moe_router_sigmoid_topk_f32",
736    }
737}
738
739#[cfg(test)]
740mod sigmoid_topk_dispatch_tests {
741    #[test]
742    fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
743        use super::sigmoid_topk_kernel;
744
745        assert_eq!(
746            sigmoid_topk_kernel(false, true, 8),
747            "moe_router_sigmoid_topk_f32_fast"
748        );
749        assert_eq!(
750            sigmoid_topk_kernel(true, true, 8),
751            "moe_router_sigmoid_topk_f32_dexp_fast"
752        );
753        assert_eq!(
754            sigmoid_topk_kernel(false, true, 9),
755            "moe_router_sigmoid_topk_f32"
756        );
757        assert_eq!(
758            sigmoid_topk_kernel(true, true, 9),
759            "moe_router_sigmoid_topk_f32_dexp"
760        );
761    }
762}
763
764pub(crate) fn rms_block() -> u32 {
765    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
766    *V.get_or_init(|| {
767        std::env::var("MEMRA_RMS_BLOCK")
768            .ok()
769            .and_then(|v| v.parse().ok())
770            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
771    })
772}
773
774pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
775    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
776    if let Some(forced) = *S.get_or_init(|| {
777        std::env::var("MEMRA_FA_SPLIT")
778            .ok()
779            .and_then(|v| v.parse().ok())
780            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
781    }) {
782        return forced;
783    }
784    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
785    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
786    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
787    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
788    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
789    //
790    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
791    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
792    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
793    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
794    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
795    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
796    // rig-divergence law: this branch is measured on 188 SMs only).
797    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
798    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
799    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
800    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
801    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
802        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
803    {
804        return if t_kv <= 8192 {
805            16
806        } else if t_kv <= 16384 {
807            64
808        } else {
809            128
810        };
811    }
812    let big_rig = fa_sm_count() >= 128;
813    if big_rig {
814        let _ = n_head_kv;
815        if t_kv <= 2048 {
816            // MEMRA_FA_SP_SHORT=N: the SHORT rung only (the SWA layers' capped t_kv lands
817            // here on step37: 33 of 45 layers at t_kv=512). At 16 the tile loop runs
818            // HALF-EMPTY (FA_DEC_TILE=32 -> nt=16 per split), so the V staging pass moves a
819            // half tile per iteration and the combine carries 2x the partials; 32 makes each
820            // split exactly one full tile. A global MEMRA_FA_SPLIT cannot isolate this — it
821            // moves the deep-ctx rung too, where more splits measured worse.
822            // NUMERIC-CLASS door (key partition -> different per-split partials/combine):
823            // new tape + battery, exactly like every other split-ladder change.
824            static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
825            if let Some(sp) = *SHORT.get_or_init(|| {
826                std::env::var("MEMRA_FA_SP_SHORT")
827                    .ok()
828                    .and_then(|v| v.parse().ok())
829                    .filter(|&s: &usize| s >= 8 && s % 8 == 0)
830            }) {
831                return sp;
832            }
833            16
834        } else if t_kv <= 16384 {
835            64
836        } else {
837            128
838        }
839    } else if n_head_kv <= 4 {
840        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
841        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
842        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
843        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
844        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
845        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
846        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
847        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
848        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
849        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
850        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
851        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
852        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
853        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
854        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
855        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
856        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
857        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
858        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
859        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
860        if t_kv <= 512 {
861            8
862        } else if t_kv <= 16384 {
863            64
864        } else {
865            128
866        }
867    } else {
868        if t_kv <= 8192 {
869            32
870        } else if t_kv <= 16384 {
871            64
872        } else {
873            128
874        }
875    }
876}
877
878/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
879/// same attribute Engine::batched_variant reads).
880pub(crate) fn fa_sm_count() -> i32 {
881    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
882    *N.get_or_init(|| {
883        cudarc::driver::result::init().ok();
884        cudarc::driver::result::device::get(0)
885            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
886                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
887            .unwrap_or(82)
888    })
889}
890
891/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
892/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
893/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
894#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
895fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
896    match head_dim {
897        256 => Ok(""),
898        128 => Ok("_hd128"),
899        d => Err(format!(
900            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
901                          callers must gate to sdpa_naive"
902        )
903        .into()),
904    }
905}
906
907/// Quant type codes matching qmatvec.cu QType enum.
908pub const QT_Q8_0: i32 = 0;
909pub const QT_Q4_K: i32 = 1;
910pub const QT_Q6_K: i32 = 2;
911pub const QT_Q5_K: i32 = 3;
912pub const QT_Q3_K: i32 = 4;
913pub const QT_IQ4_XS: i32 = 5;
914pub const QT_IQ3_S: i32 = 6;
915pub const QT_NVFP4: i32 = 7;
916/// Slot-major v2 bank permutation of `QT_NVFP4` (see tp.rs `nvfp4_matrix_v2_permute`) — only the
917/// grouped-prefill dequant consumes this tag; every direct/dp4a lane must keep refusing it.
918pub const QT_NVFP4_V2: i32 = 107;
919/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
920/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
921/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
922/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
923/// — ONE weight copy total, no Q8_0 re-encode duplicate.
924pub const QT_F8_E4M3: i32 = 10;
925/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
926/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
927pub const QT_NVFP4_RP: i32 = 9;
928/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
929pub const QT_F32: i32 = 8;
930pub const QT_BF16: i32 = 11;
931pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
932/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
933/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
934/// dp4a/MMQ implementation exists.
935pub const QT_Q2_K: i32 = 13;
936/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
937/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
938/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
939/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
940/// scalar `scale` field is 1.0 by the layout contract.
941///
942/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
943/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
944/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
945/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
946/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
947/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
948/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
949/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
950/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
951pub const QT_F8_E4M3_BLK: i32 = 14;
952
953/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
954pub struct Engine {
955    pub gpu: memra_runtime::Gpu,
956    module: Arc<CudaModule>,
957    hybrid: Arc<CudaModule>,
958    /// Kimi Delta Attention kernels (cu/kda.cu) — separate fatbin, resolved through `func`.
959    kda: Arc<CudaModule>,
960    qmatvec: Arc<CudaModule>,
961    flash: Arc<CudaModule>,
962    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
963    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
964    /// Lazy: loaded on first global-format use; None until then.
965    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
966    gemm: Arc<CudaModule>,
967    router: Arc<CudaModule>,
968    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
969    sample: Arc<CudaModule>,
970    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
971    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
972    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
973    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
974    /// MEMRA_STEP_TP_W8, hybrid half: q8_0 mirrors of bf16 GEMV weights that do NOT live in a
975    /// TP resident bank (the LM head, the shared expert, the dense-FFN layers), keyed by the
976    /// bf16 slab's device pointer and built on first decode use. The mirror is 1.0625 B/w
977    /// against bf16's 2, and the raw slab stays resident, so prefill keeps its arithmetic.
978    /// KEYED ON (pointer, in_f, out_f), not on the pointer alone: a row-range VIEW of a slab
979    /// carries the PARENT's base pointer when the range starts at row 0, so a pointer-only key
980    /// would hand the head-split lo half (4096 x 64448) the full head's mirror (4096 x 128896)
981    /// and read 2x past the rows it owns. The shape is part of the identity of a mirror.
982    w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
983    /// Per-`in_f` q8_1 activation scratch for those mirrors (allocating per call would cost
984    /// more than the door saves).
985    #[allow(clippy::type_complexity)]
986    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
987    w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
988    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
989    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
990    /// the single largest block. The cache still owns every address for its full lifetime.
991    moe_cache_layout: Mutex<Option<Vec<usize>>>,
992    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
993    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
994    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
995    /// verify between replays) reuse their addresses and the replay reads/writes live memory
996    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
997    capture_keep_on: std::sync::atomic::AtomicBool,
998    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
999    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
1000    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
1001    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
1002    verify_exact: std::sync::atomic::AtomicBool,
1003    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
1004    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
1005    pub copy_stream: Arc<CudaStream>,
1006    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
1007    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
1008    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
1009    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
1010    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
1011    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
1012    #[cfg(memra_cutlass)]
1013    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
1014    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
1015    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
1016    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
1017    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
1018    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
1019    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
1020    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
1021    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
1022    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
1023    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
1024    #[allow(clippy::type_complexity)]
1025    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1026    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1027    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
1028    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
1029    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
1030    #[allow(clippy::type_complexity)]
1031    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1032    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1033    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
1034    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1035    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1036    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
1037    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
1038    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
1039    /// before capture under the generate_graph tracking-off window so it carries no events).
1040    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1041    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
1042    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
1043    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
1044    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
1045    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1046    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
1047    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
1048    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
1049    router_stage: Mutex<Option<PinnedStage>>,
1050    /// Persistent hc-glue decode workspace (MEMRA_HC_DECODE_WS, lane/glm5-decode-diet lever 2).
1051    /// Pooled per engine like `fa_part_pool`: the buffers are pure per-step scratch (every
1052    /// element fully overwritten before read each step), so one slot per engine is correct
1053    /// even across sessions; the walk TAKES it for the step and puts it back, and a second
1054    /// concurrent walk on the same engine simply falls back to fresh allocations.
1055    hyper_decode_ws: Mutex<Option<crate::hyper::HyperDecodeWs>>,
1056    /// Verify-walk allocation workspace (MEMRA_VERIFY_WS — glm5-alias
1057    /// MEMRA_GLM5_VERIFY_WS honored, OFF-wins; lane/glm5-matvec door W, generalized
1058    /// lane/glm5-extract-general — the pool is family-agnostic by content): the
1059    /// `MEMRA_HC_DECODE_WS` pattern extended to the spec verify walk, whose ~1380
1060    /// `cuMemAllocAsync`+Free pairs/token the t=1 workspace door structurally never reaches
1061    /// (spec decodes through the t=K+1 walk — diet-battery WINDOW.md). Size-keyed free-lists;
1062    /// verify-only call sites (the rows-exact matmul class, the KDA rows arm, the MoE vrows
1063    /// staging) draw from and recycle into it. Reuse is byte-identical by the same contract
1064    /// that makes `uninit` legal at those sites: every element is fully overwritten before
1065    /// any read, by the SAME unchanged kernels. Per-engine = per-stream, so stream ordering
1066    /// makes recycle-then-reuse safe exactly like free-then-alloc on the async pool.
1067    verify_ws: Mutex<VerifyWs>,
1068    /// Resident device mirrors of the per-expert NVFP4 `weight_scale_2` macro planes, keyed by
1069    /// `(layer, plane)` with plane 0/1/2 = gate/up/down (MEMRA_MOE_VROWS_DEV_TABLES, door D).
1070    /// The device table build needs `macro_scale(ex)` where the selection lives; the host plane
1071    /// is an immutable `Vec<f32>` of n_expert entries for the process lifetime, so ONE upload
1072    /// per (layer, plane) serves every subsequent layer-call — 3 x n_expert x 4 B (3.5 KB at
1073    /// 288 experts), 126 buffers = ~145 KB for a 42-MoE-layer model. Uploading per call instead
1074    /// would ADD three HtoD to a door whose whole purpose is removing two.
1075    vrows_macro_dev: Mutex<std::collections::HashMap<(u16, u8), CudaSlice<f32>>>,
1076    /// Resident all-ones f32 vector for the UNGATED shared-expert add (`MEMRA_HTOD_DIET`,
1077    /// door H). A family whose plan carries no `ffn_gate_inp_shexp` (GLM-5.3-Flash is the
1078    /// first) makes `moe_shexp_add` take the `g = 1.0` arm, which re-uploaded a freshly
1079    /// allocated `vec![1.0f32; t]` on EVERY MoE layer-call — 42 pageable HtoD per ship round
1080    /// to move a constant on the glm5 serving geometry. Grown to the largest t
1081    /// seen; the buffer may be LONGER than t because `add_scaled_rows_f32` reads only
1082    /// `scale[0..nrows]`.
1083    shexp_ones: Mutex<Option<CudaSlice<f32>>>,
1084}
1085
1086/// Size-keyed device-buffer free-lists for the verify walk (door W — see the field doc on
1087/// [`Engine::verify_ws`]). Exact-length keying: the walk's shapes quantize to a few
1088/// classes per round (t in 2..=8 times fixed widths), so hit rates are structural, and an
1089/// exact-size buffer keeps every `debug_assert_eq!(len, ...)` at the launchers intact.
1090#[derive(Default)]
1091pub struct VerifyWs {
1092    f32_pool: std::collections::HashMap<usize, Vec<CudaSlice<f32>>>,
1093    i8_pool: std::collections::HashMap<usize, Vec<CudaSlice<i8>>>,
1094    u64_pool: std::collections::HashMap<usize, Vec<CudaSlice<u64>>>,
1095    held_bytes: usize,
1096}
1097
1098/// Per-size-class retention cap: enough for every live shape class of one round plus the
1099/// stash generation, small enough that a shape drift cannot hoard VRAM.
1100const VWS_PER_CLASS_CAP: usize = 16;
1101/// Total retention cap (bytes). The round's recurring buffers are t*8192-f32-class and MoE
1102/// staging (<= ~1 MiB each); 256 MiB holds every class with an order of magnitude of slack.
1103const VWS_HELD_BYTES_CAP: usize = 256 << 20;
1104
1105impl VerifyWs {
1106    fn take<T>(
1107        pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1108        held: &mut usize,
1109        n: usize,
1110    ) -> Option<CudaSlice<T>> {
1111        let s = pool.get_mut(&n)?.pop()?;
1112        *held -= n * std::mem::size_of::<T>();
1113        Some(s)
1114    }
1115    fn put<T>(
1116        pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1117        held: &mut usize,
1118        s: CudaSlice<T>,
1119    ) {
1120        let n = s.len();
1121        let bytes = n * std::mem::size_of::<T>();
1122        if *held + bytes > VWS_HELD_BYTES_CAP {
1123            return; // drop: falls to the ordinary async free
1124        }
1125        let v = pool.entry(n).or_default();
1126        if v.len() >= VWS_PER_CLASS_CAP {
1127            return;
1128        }
1129        v.push(s);
1130        *held += bytes;
1131    }
1132}
1133
1134/// Device-scratch allocation census (lane/glm5-decode-diet): bumped by every `alloc_uninit`
1135/// and `zeros` call — the class the launch-diet census measured at 2,358
1136/// `cuMemAllocAsync+Free` calls/token. The decode-workspace gate reads deltas per step; the
1137/// cost axis is the CALL COUNT (the box's measured ~1.06 us/driver call), which is exactly
1138/// what this counts. Relaxed atomic: one increment per allocation, noise-level.
1139pub static SCRATCH_ALLOC_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1140
1141/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
1142/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
1143/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
1144/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
1145/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
1146/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
1147/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
1148/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
1149/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
1150fn fa_v2_on() -> bool {
1151    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
1152    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
1153    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
1154    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
1155    // + graph bit-identity green on all three models.
1156    std::env::var("MEMRA_FA_V2")
1157        .map(|v| v != "0")
1158        .unwrap_or(true)
1159}
1160
1161/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
1162/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
1163/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
1164/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
1165/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
1166/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
1167/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
1168/// `MEMRA_FA_PART_ZERO=1`: zero every freshly grown fa partial bank. DEFAULT OFF,
1169/// diagnostic only. See `fa_part_alloc` for what it discriminates and why it is not a fix.
1170pub(crate) fn fa_part_zero_on() -> bool {
1171    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1172    *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1173}
1174
1175pub(crate) fn fa_v3_on() -> bool {
1176    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
1177    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
1178    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
1179    std::env::var("MEMRA_FA_V3")
1180        .map(|v| v != "0")
1181        .unwrap_or(true)
1182}
1183
1184/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
1185/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
1186/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
1187/// predicate so the twins can never diverge.
1188fn fa_v4_mode() -> &'static str {
1189    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
1190    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
1191}
1192fn fa_v4_on() -> bool {
1193    fa_v4_mode() != "0"
1194} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
1195/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
1196/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
1197/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
1198/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
1199/// stays kernel-family-identical to decode at the same t_kv.
1200/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
1201/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
1202pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1203    std::sync::atomic::AtomicUsize::new(1024);
1204pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1205    std::sync::atomic::AtomicUsize::new(usize::MAX);
1206pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1207    fa_v4_at(t_kv)
1208}
1209fn fa_v4_at(t_kv: usize) -> bool {
1210    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1211    let mx = *M.get_or_init(|| {
1212        std::env::var("MEMRA_FA_V4_MAX")
1213            .ok()
1214            .and_then(|v| v.parse().ok())
1215            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1216    });
1217    fa_v4_on() && t_kv < mx
1218}
1219/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
1220/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
1221/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
1222/// (same split partition, same softmax/accumulation order, same partials/combine) and only
1223/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
1224/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
1225/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
1226/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
1227/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
1228/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
1229/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
1230/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
1231/// within one process (the v2/v3 pattern).
1232pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1233fn fa_deep_at(t_kv: usize) -> bool {
1234    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1235        return false;
1236    }
1237    let min = std::env::var("MEMRA_FA_DEEP_MIN")
1238        .ok()
1239        .and_then(|v| v.parse().ok())
1240        .unwrap_or(FA_DEEP_MIN_DEFAULT);
1241    t_kv >= min
1242}
1243/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
1244pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1245    fa_deep_at(t_kv)
1246}
1247
1248fn fa_v3_active(head_dim: usize) -> bool {
1249    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
1250    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
1251    fa_v3_on()
1252        && head_dim.is_multiple_of(128)
1253        && kv_cache_formats() == ("q8_0", "q5_1")
1254        && !Engine::kv_fp8_on()
1255}
1256
1257/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
1258/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
1259/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
1260/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
1261/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
1262/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
1263/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
1264pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1265    std::env::var("MEMRA_NO_FA_VEC").is_err()
1266        && t_kv >= fa_vec_min_tkv()
1267        && head_dim == 256
1268        && fa_v4_at(t_kv)
1269        && !matches!(fa_v4_mode(), "noB3" | "stage")
1270        && !Engine::kv_fp8_on()
1271}
1272/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
1273pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1274    fa_split_keys(t_kv, n_head_kv)
1275}
1276
1277/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
1278/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
1279/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
1280/// so we allocate through `result::malloc_host` with flags=0 directly.
1281struct PinnedStage {
1282    ptr: *mut u8,
1283    cap: usize,
1284}
1285unsafe impl Send for PinnedStage {}
1286impl PinnedStage {
1287    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1288        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1289        Ok(PinnedStage { ptr, cap })
1290    }
1291}
1292impl Drop for PinnedStage {
1293    fn drop(&mut self) {
1294        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1295    }
1296}
1297
1298/// Owned page-locked CACHEABLE host buffer (flags=0, deliberately NOT write-combined) for the
1299/// prefix-cache host tier (lane/kv-host-spill-20260830). Same allocation class as `PinnedStage`
1300/// above and for the same reason: `ctx().alloc_pinned` is CU_MEMHOSTALLOC_WRITECOMBINED, which
1301/// is right for H2D-only staging but pathologically slow for host READS (see the HostBuf CAVEAT
1302/// in model.rs), and these bytes are CPU-read by the MEMRA_KV_HOST_VERIFY digest arm. Public
1303/// because the server's host-tier cache owns these buffers across requests.
1304pub struct PinnedHostBuf {
1305    ptr: *mut u8,
1306    len: usize,
1307}
1308// Safety: the allocation is process-wide page-locked host memory; the raw pointer is owned by
1309// this struct alone and freed exactly once in Drop (identical justification to PinnedStage).
1310unsafe impl Send for PinnedHostBuf {}
1311impl PinnedHostBuf {
1312    /// Allocate `len` pinned cacheable bytes (a zero-length request still pins one byte so the
1313    /// pointer stays valid, mirroring the device planes' `alloc_u8(kb.max(1))` convention).
1314    pub fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1315        let ptr = unsafe { cudarc::driver::result::malloc_host(len.max(1), 0)? } as *mut u8;
1316        Ok(PinnedHostBuf { ptr, len })
1317    }
1318    pub fn len(&self) -> usize {
1319        self.len
1320    }
1321    pub fn is_empty(&self) -> bool {
1322        self.len == 0
1323    }
1324    pub fn as_slice(&self) -> &[u8] {
1325        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1326    }
1327    pub fn as_mut_slice(&mut self) -> &mut [u8] {
1328        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1329    }
1330}
1331impl Drop for PinnedHostBuf {
1332    fn drop(&mut self) {
1333        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1334    }
1335}
1336
1337/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
1338/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
1339pub const ARGMAX_NB: usize = 256;
1340
1341/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
1342pub(crate) use memra_fa3_vl as fa3_vl_raw;
1343
1344unsafe extern "C" {
1345    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
1346    fn memra_fa3_prefill(
1347        q16: *const core::ffi::c_void,
1348        k16: *const core::ffi::c_void,
1349        v16: *const core::ffi::c_void,
1350        o: *mut f32,
1351        t: i32,
1352        h: i32,
1353        hkv: i32,
1354        d: i32,
1355        scale: f32,
1356        stream: *mut core::ffi::c_void,
1357    ) -> i32;
1358    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
1359    pub(crate) fn memra_fa3_vl(
1360        q16s: *const *const core::ffi::c_void,
1361        k16s: *const *const core::ffi::c_void,
1362        v16s: *const *const core::ffi::c_void,
1363        os: *const *mut f32,
1364        ts: *const i32,
1365        b: i32,
1366        h: i32,
1367        hkv: i32,
1368        d: i32,
1369        scale: f32,
1370        stream: *mut core::ffi::c_void,
1371    ) -> i32;
1372}
1373
1374/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1375/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1376/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1377/// (slots are never re-allocated), so passing raw values is stable across the launch.
1378#[repr(C)]
1379#[derive(Clone, Copy)]
1380pub struct WPtr8(pub [u64; 8]);
1381unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1382
1383/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1384/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1385/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1386/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1387#[repr(C)]
1388#[derive(Clone, Copy, Default)]
1389pub struct GdnSeqVl {
1390    pub kb16: u64,
1391    pub gcum: u64,
1392    pub beta: u64,
1393    pub u: u64,
1394    pub wb16: u64,
1395    pub y: u64,
1396    pub ssnap: u64,
1397    pub state_in: u64,
1398    pub state_out: u64,
1399    pub q: u64,
1400    pub p: u64,
1401    pub o: u64,
1402    pub k: u64,
1403    pub v: u64,
1404    pub g: u64,
1405    pub a: u64,
1406    pub w: u64,
1407    pub t: i32,
1408    pub nc: i32,
1409}
1410unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1411#[repr(C)]
1412#[derive(Clone, Copy)]
1413pub struct GdnVl8(pub [GdnSeqVl; 8]);
1414unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1415
1416/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1417/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1418#[repr(C)]
1419#[derive(Clone, Copy, Default)]
1420pub struct GdnWVl {
1421    pub qb16: u64,
1422    pub pb16: u64,
1423}
1424unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1425#[repr(C)]
1426#[derive(Clone, Copy)]
1427pub struct GdnWVl8(pub [GdnWVl; 8]);
1428unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1429
1430/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1431#[repr(C)]
1432#[derive(Clone, Copy, Default)]
1433pub struct GdnPrepVl {
1434    pub qkv: u64,
1435    pub conv_state: u64,
1436    pub conv_out: u64,
1437    pub q_g: u64,
1438    pub k_g: u64,
1439    pub v_g: u64,
1440    pub q_l2: u64,
1441    pub k_l2: u64,
1442    pub beta_raw: u64,
1443    pub alpha: u64,
1444    pub beta: u64,
1445    pub g_log: u64,
1446    pub o: u64,
1447    pub z: u64,
1448    pub gn: u64,
1449    pub gn16: u64,
1450    pub kb16: u64,
1451    pub qb16: u64,
1452    pub t: i32,
1453    pub pad: i32,
1454}
1455unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1456#[repr(C)]
1457#[derive(Clone, Copy)]
1458pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1459unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1460
1461/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1462#[repr(C)]
1463#[derive(Clone, Copy, Default)]
1464pub struct FaSeqVl {
1465    pub q: u64,
1466    pub k16: u64,
1467    pub v16: u64,
1468    pub o: u64,
1469    pub kf: u64,
1470    pub vf: u64,
1471    pub t: i32,
1472    pub pad: i32,
1473}
1474unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1475#[repr(C)]
1476#[derive(Clone, Copy)]
1477pub struct FaVl8(pub [FaSeqVl; 8]);
1478unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1479
1480/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1481#[repr(C)]
1482#[derive(Clone, Copy, Default)]
1483pub struct AttnPreVl {
1484    pub qf: u64,
1485    pub kf: u64,
1486    pub vf: u64,
1487    pub q: u64,
1488    pub gate: u64,
1489    pub qn: u64,
1490    pub kn: u64,
1491    pub kc: u64,
1492    pub vc: u64,
1493    pub t: i32,
1494    pub pad: i32,
1495}
1496unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1497#[repr(C)]
1498#[derive(Clone, Copy)]
1499pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1500unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1501
1502/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1503/// varlen K1-K5 chain fills them).
1504pub struct GdnChunkBufs {
1505    pub gcum: CudaSlice<f32>,
1506    pub a: CudaSlice<f32>,
1507    pub p: CudaSlice<f32>,
1508    pub u: CudaSlice<f32>,
1509    pub w: CudaSlice<f32>,
1510    pub kb16: CudaSlice<u8>,
1511    pub wb16: CudaSlice<u8>,
1512    pub y16: CudaSlice<u8>,
1513    pub ssnap16: CudaSlice<u8>,
1514    pub qb16: CudaSlice<u8>,
1515    pub pb16: CudaSlice<u8>,
1516    pub o: CudaSlice<f32>,
1517    pub t: usize,
1518    pub nc: usize,
1519}
1520
1521/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1522#[repr(C)]
1523#[derive(Clone, Copy)]
1524pub struct F32x8(pub [f32; 8]);
1525unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1526
1527/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1528/// process. Bench binaries read it right after the call to print gen-only throughput without the
1529/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1530pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1531
1532/// Fused MoE-epilogue dispatches taken since process start (`MEMRA_MOE_FUSED_EPI`), incremented
1533/// once per (token, layer) that actually runs `moe_fused_epi_token_q8`.
1534///
1535/// This exists because the arm cannot be observed any other way: setting `MEMRA_MOE_STATS` /
1536/// `MEMRA_MOE_TRACE` / `MEMRA_MOE_WEIGHT_TRACE` / `MEMRA_MOE_INPUT_TRACE_DIR` sets
1537/// `observe_routes` in `moe_ffn_inner`, which DIVERTS dispatch to the host-routed path — so a
1538/// gate that tried to prove the fused arm ran by tracing would prove it about a different
1539/// program. Read it via [`moe_fused_epilogue_dispatches`] around a workload.
1540pub static MOE_FUSED_EPI_DISPATCHES: std::sync::atomic::AtomicU64 =
1541    std::sync::atomic::AtomicU64::new(0);
1542
1543/// Snapshot of [`MOE_FUSED_EPI_DISPATCHES`]. Gates take a before/after pair around a workload and
1544/// assert on the delta, anchoring on the arm's own invocation rather than on a flag being set
1545/// (LAW:wiring-assertions-match-prose).
1546pub fn moe_fused_epilogue_dispatches() -> u64 {
1547    MOE_FUSED_EPI_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1548}
1549
1550/// Verify-rows batched MoE dispatches taken since process start (lane/glm5-vrest): incremented
1551/// once per (layer, verify-call) that runs the pairs-shaped routed-expert program
1552/// (`moe_gate_up_preclamp8_q8_rows` + `moe_down8_fma_q8_rows`) instead of the per-(token,expert)
1553/// sequential loop. Rides `MEMRA_GLM5_VERIFY_BATCH`'s arm — no flag of its own. Same rationale
1554/// as [`MOE_FUSED_EPI_DISPATCHES`]: the observation envs divert dispatch, so gates anchor on the
1555/// arm's own invocation (LAW:wiring-assertions-match-prose).
1556pub static MOE_VROWS_DISPATCHES: std::sync::atomic::AtomicU64 =
1557    std::sync::atomic::AtomicU64::new(0);
1558
1559/// Snapshot of [`MOE_VROWS_DISPATCHES`] — gates take a before/after delta around a workload.
1560pub fn moe_vrows_dispatches() -> u64 {
1561    MOE_VROWS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1562}
1563
1564/// `MEMRA_BF16_TCOLS_WIDE` (lane/glm5-matvec door T, default ON since the 2026-08-31 mv-battery
1565/// flip; `=0` is the rollback seam): FloatBf16 rows calls at
1566/// t=2..=16 ride the weight-once t-column twins (`matvec_bf16_f32acc_x4_tcols` for t<=8, the
1567/// NEW `..._tcols16` for 9..=16) instead of the grid.y=t weight-rereading `_rows` kernel. The
1568/// motivating call is the DFlash2 drafter's block head: `eh.matmul(head, rows, 15)` re-read
1569/// the 1.269 GB lm head 15x per spec round (diet-battery c8-ship census, 5.31 ms/round —
1570/// 13% of decode GPU). Bit-identical per (row, token) by the tcols class's standing
1571/// construction; gated by `glm5_matvec_doors_gpu`. Read per call — the rollback seam.
1572fn bf16_tcols_wide_on() -> bool {
1573    std::env::var("MEMRA_BF16_TCOLS_WIDE").as_deref() != Ok("0")
1574}
1575
1576/// Engagement counter for the wide-t tcols door (`MEMRA_BF16_TCOLS_WIDE`), incremented at the
1577/// door's own dispatch (LAW:wiring-assertions-match-prose). Read via
1578/// [`bf16_tcols_wide_dispatches`].
1579pub static BF16_TCOLS_WIDE_DISPATCHES: std::sync::atomic::AtomicU64 =
1580    std::sync::atomic::AtomicU64::new(0);
1581
1582/// Snapshot of [`BF16_TCOLS_WIDE_DISPATCHES`] — gates take a before/after delta.
1583pub fn bf16_tcols_wide_dispatches() -> u64 {
1584    BF16_TCOLS_WIDE_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1585}
1586
1587/// `MEMRA_BF16_TCOLS_X1` (lane/glm5-matvec door X, default ON since the 2026-08-31 mv-battery
1588/// flip; `=0` is the rollback seam): the tcols dispatch takes
1589/// the one-row-per-block grid twin (`matvec_bf16_f32acc_x1_tcols`, grid.x = out_f) instead of
1590/// the 4-rows-per-block form. WHY: the trunk kda shapes (out_f 4096/8192) launch 1024/2048
1591/// blocks — ~one resident wave, and the census pins them at 1.05 TB/s (59% of peak) while the
1592/// SAME kernel at the lm head's 38720-block grid runs 1.43 TB/s (80%). Per-row program and
1593/// tree verbatim — bit-identical. Gated by `glm5_matvec_doors_gpu`. Read per call.
1594fn bf16_tcols_x1_on() -> bool {
1595    std::env::var("MEMRA_BF16_TCOLS_X1").as_deref() != Ok("0")
1596}
1597
1598/// Engagement counter for the x1-grid tcols door (`MEMRA_BF16_TCOLS_X1`).
1599pub static BF16_TCOLS_X1_DISPATCHES: std::sync::atomic::AtomicU64 =
1600    std::sync::atomic::AtomicU64::new(0);
1601
1602/// Snapshot of [`BF16_TCOLS_X1_DISPATCHES`] — gates take a before/after delta.
1603pub fn bf16_tcols_x1_dispatches() -> u64 {
1604    BF16_TCOLS_X1_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1605}
1606
1607/// `MEMRA_BF16_TCOLS_RED_FUSED=1` (lane/glm5-door-r door R, default OFF): the tcols
1608/// dispatches take the `_rf` fused-reduce-tail twins (`matvec_bf16_f32acc_x1_tcols_rf` /
1609/// `..._x4_tcols_rf` / `..._x4_tcols16_rf`). WHY (moe-loc LANE.md §2.2): after door X the
1610/// kda trunk's tcols calls sit at 67.0% of peak because the reduce tail runs t SEPARATE
1611/// strided trees — ~30 block-wide barriers at t=3.34 (135 at the drafter head's t=15)
1612/// against a 4-iteration main loop; the kernel is barrier/tail-bound. The twins share ONE
1613/// barrier sequence across the t columns (`red[t*blockDim]`, dynamic shared) and run levels
1614/// s<=16 as a `__shfl_down_sync` chain at the IDENTICAL pairing and operand order — 9t -> 3
1615/// barriers per block, bit-identical by pairing preservation (gated with a shifted-pairing
1616/// red in `glm5_matvec_doors_gpu`). Engages only when `MEMRA_MMV_BLOCK` is a power of two
1617/// (the fused tail's block-wide loop must pass exactly through s=32; the default 128 is).
1618/// Read per call — unset or `=0` is byte-for-byte the standing tcols program.
1619fn bf16_tcols_red_fused_on() -> bool {
1620    std::env::var("MEMRA_BF16_TCOLS_RED_FUSED").as_deref() == Ok("1")
1621}
1622
1623/// Engagement counter for the fused-reduce-tail tcols door (`MEMRA_BF16_TCOLS_RED_FUSED`),
1624/// incremented at the door's own dispatch (LAW:wiring-assertions-match-prose).
1625pub static BF16_TCOLS_RED_FUSED_DISPATCHES: std::sync::atomic::AtomicU64 =
1626    std::sync::atomic::AtomicU64::new(0);
1627
1628/// Snapshot of [`BF16_TCOLS_RED_FUSED_DISPATCHES`] — gates take a before/after delta.
1629pub fn bf16_tcols_red_fused_dispatches() -> u64 {
1630    BF16_TCOLS_RED_FUSED_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1631}
1632
1633/// `MEMRA_MOE_VROWS_PACK=1` (lane/glm5-matvec door M, default OFF): the verify-rows MoE pair
1634/// launches its `_w4` warp-packed twins — MEMRA_MMVQ_ROWS = 4 warps per block on threadIdx.y
1635/// (the qmatvec mmvq family's standing shape) instead of one warp per block. The unpacked
1636/// launch caps residency at the blocks/SM limit (<=67% of warp slots) and schedules ~65k
1637/// one-warp blocks per launch; per-warp body verbatim, bit-identical per (row, pair). Gated
1638/// by `glm5_matvec_doors_gpu`. Read per call.
1639fn moe_vrows_pack_on() -> bool {
1640    std::env::var("MEMRA_MOE_VROWS_PACK").as_deref() == Ok("1")
1641}
1642
1643/// Engagement counter for the warp-packed verify-rows MoE door (`MEMRA_MOE_VROWS_PACK`).
1644pub static MOE_VROWS_PACK_DISPATCHES: std::sync::atomic::AtomicU64 =
1645    std::sync::atomic::AtomicU64::new(0);
1646
1647/// Snapshot of [`MOE_VROWS_PACK_DISPATCHES`] — gates take a before/after delta.
1648pub fn moe_vrows_pack_dispatches() -> u64 {
1649    MOE_VROWS_PACK_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1650}
1651
1652/// `MEMRA_MOE_VROWS_DEV_TABLES=1` (lane/glm5-moe-loc door D, default OFF): the verify-rows MoE
1653/// pair builds its `ptrs`/`scl` tables ON DEVICE from the router's own `sel`/`w` device output
1654/// (`moe_vrows_tables_from_sel`) instead of on the host, and the layer routes through the
1655/// readback-free `moe_router_sigmoid_topk` rather than `..._host`. WHY: the host table build is
1656/// the ONLY consumer of the selection on the serving shape, and it costs a full
1657/// `cuStreamSynchronize` + 2 DtoH + 2 pageable HtoD + 2 host Vec allocations per MoE layer-call
1658/// = 42 device-wide drains + 84 DtoH + 84 HtoD per ship round. Bit-identical: same integer
1659/// `base + ex*stride`, same macro-plane lookups, same single `w * macro_down` product. Read per
1660/// call; fails closed to the host path whenever any host-visible route consumer is armed.
1661fn moe_vrows_dev_tables_on() -> bool {
1662    std::env::var("MEMRA_MOE_VROWS_DEV_TABLES").as_deref() == Ok("1")
1663}
1664
1665/// Engagement counter for the device-side vrows table build (`MEMRA_MOE_VROWS_DEV_TABLES`).
1666pub static MOE_VROWS_DEV_TABLES_DISPATCHES: std::sync::atomic::AtomicU64 =
1667    std::sync::atomic::AtomicU64::new(0);
1668
1669/// Snapshot of [`MOE_VROWS_DEV_TABLES_DISPATCHES`] — gates take a before/after delta.
1670pub fn moe_vrows_dev_tables_dispatches() -> u64 {
1671    MOE_VROWS_DEV_TABLES_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1672}
1673
1674/// Router readbacks (one full `cuStreamSynchronize` + 2 DtoH each) that door D skipped. The
1675/// count receipt for the host seam: a gate asserts it moves 1:1 with
1676/// [`MOE_VROWS_DEV_TABLES_DISPATCHES`] on the ON arm and stays flat on the OFF arm.
1677pub static MOE_VROWS_ROUTER_SYNCS_AVOIDED: std::sync::atomic::AtomicU64 =
1678    std::sync::atomic::AtomicU64::new(0);
1679
1680/// Snapshot of [`MOE_VROWS_ROUTER_SYNCS_AVOIDED`].
1681pub fn moe_vrows_router_syncs_avoided() -> u64 {
1682    MOE_VROWS_ROUTER_SYNCS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
1683}
1684
1685/// `MEMRA_MOE_VROWS_DEDUP_STAT=1` (lane/glm5-moe-loc, default OFF — a MEASUREMENT instrument,
1686/// not a serving door): on the host table-build arm, count the pair union's expert VISITS and
1687/// DISTINCT experts per layer-call into [`MOE_VROWS_PAIR_VISITS`] /
1688/// [`MOE_VROWS_PAIR_DISTINCT`]. WHY IT EXISTS: the pair runs at ~90% of this card class's
1689/// theoretical DRAM peak (moe-loc LANE.md §1), so cross-row expert-slab dedup is the ONLY
1690/// remaining byte lever, and its size is exactly `1 - distinct/visits` — an unmeasured routing
1691/// property whose independent-routing bound is 3.2% but whose structural ceiling is 70%. This
1692/// counter turns a speculative kernel campaign into a priced decision for the cost of a host
1693/// bitset. Requires `MEMRA_MOE_VROWS_DEV_TABLES=0` (door D removes the host selection).
1694fn moe_vrows_dedup_stat_on() -> bool {
1695    std::env::var("MEMRA_MOE_VROWS_DEDUP_STAT").as_deref() == Ok("1")
1696}
1697
1698/// Expert VISITS (t x n_used) summed over vrows layer-calls under `MEMRA_MOE_VROWS_DEDUP_STAT`.
1699pub static MOE_VROWS_PAIR_VISITS: std::sync::atomic::AtomicU64 =
1700    std::sync::atomic::AtomicU64::new(0);
1701
1702/// DISTINCT experts in the pair union, summed over the same layer-calls. The dedup lever is
1703/// `1 - distinct/visits`; equal counters mean routing is disjoint across the verify rows and
1704/// there is no byte to save.
1705pub static MOE_VROWS_PAIR_DISTINCT: std::sync::atomic::AtomicU64 =
1706    std::sync::atomic::AtomicU64::new(0);
1707
1708/// `(visits, distinct)` for one layer-call's pair union — the dedup lever's whole arithmetic.
1709/// `visits` is `t * n_used`, the slab reads the pair performs today; `distinct` is how many of
1710/// them are to a DIFFERENT expert. `1 - distinct/visits` is the share of the pair's 9.86 ms/round
1711/// that a dedup kernel could remove, and nothing else about the pair is removable (it already
1712/// runs at ~90% of theoretical DRAM peak). Split out from the call site so the counting itself is
1713/// unit-testable on planted overlaps rather than inferred from a live routing tape.
1714pub(crate) fn vrows_overlap_counts(sel_all: &[u32]) -> (u64, u64) {
1715    let mut seen = std::collections::HashSet::with_capacity(sel_all.len());
1716    for &ex in sel_all {
1717        seen.insert(ex);
1718    }
1719    (sel_all.len() as u64, seen.len() as u64)
1720}
1721
1722/// vrows layer-calls the dedup instrument has observed — the reporting cadence's clock.
1723static MOE_VROWS_DEDUP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1724
1725/// AN INSTRUMENT HAS TO SPEAK. A box window greps a server log; it cannot read a Rust atomic, so
1726/// the dedup counters emit their cumulative ratio on the first vrows layer-call and every 42
1727/// after (42 = the MoE layer count, i.e. about one line per decode round). The reported
1728/// `repeat` IS the dedup lever's ceiling: the share of the pair's 9.86 ms/round that reading a
1729/// shared expert slab once could remove, and the only removable share that exists (LANE.md §1 —
1730/// the pair already runs at ~90% of theoretical DRAM peak).
1731fn moe_vrows_dedup_report() {
1732    let n = MOE_VROWS_DEDUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1733    if n != 0 && !n.is_multiple_of(42) {
1734        return;
1735    }
1736    let (visits, distinct) = moe_vrows_pair_overlap();
1737    if visits == 0 {
1738        return;
1739    }
1740    let repeat = 100.0 * (1.0 - distinct as f64 / visits as f64);
1741    eprintln!(
1742        "[moe-vrows-dedup] layer-calls={} visits={visits} distinct={distinct} \
1743         repeat={repeat:.2}% = the cross-row expert-slab dedup ceiling on the vrows pair \
1744         (MEMRA_MOE_VROWS_DEDUP_STAT=1)",
1745        n + 1
1746    );
1747}
1748
1749/// Gate hook for [`vrows_overlap_counts`] — the counting is the whole instrument, so it is gated
1750/// on planted overlaps (disjoint / partial / identical) rather than inferred from a live tape.
1751pub fn vrows_overlap_counts_for_test(sel_all: &[u32]) -> (u64, u64) {
1752    vrows_overlap_counts(sel_all)
1753}
1754
1755/// Snapshot of the dedup instrument as `(visits, distinct)`.
1756pub fn moe_vrows_pair_overlap() -> (u64, u64) {
1757    (
1758        MOE_VROWS_PAIR_VISITS.load(std::sync::atomic::Ordering::Relaxed),
1759        MOE_VROWS_PAIR_DISTINCT.load(std::sync::atomic::Ordering::Relaxed),
1760    )
1761}
1762
1763/// `MEMRA_MOE_VROWS_DEDUP_ORDER=1` (lane/glm5-dedup door E, default OFF): the verify-rows
1764/// gate/up launch takes the `_ord` twin — grid TRANSPOSED so the pair index is the fastest
1765/// dimension, walking an EXPERT-MAJOR order plane appended to the pointer table. WHY: the
1766/// struct-battery instrument measured a **21.96% repeat fraction** across the pair's expert
1767/// visits (2.55M visits, 6.9x the 3.21% independent-routing bound), and the pair is already at
1768/// 90.2% of theoretical DRAM peak, so the only lever left is not re-reading a slab a sibling
1769/// verify row already read — which requires the repeat visit to be SCHEDULED inside the reuse
1770/// window. Bit-identical by construction: every output is a pure function of its `(o, pr)`
1771/// coordinate and no block communicates, so re-indexing which block computes which output moves
1772/// no bits (`glm5_dedup_sched_gpu`). The WIN is a scheduling property, unpriceable on an
1773/// exactness-only rig — hence default OFF with the box pricing the flip.
1774///
1775/// Refused by name, falling closed to the shipped schedule: door M (`MEMRA_MOE_VROWS_PACK`, the
1776/// refuted 4-warp pack) takes precedence in the launcher, and the door engages only when the
1777/// order plane is actually present (`ptrs.len() >= 4*n_pairs`), so a direct launcher call with a
1778/// 3-plane table keeps the shipped program.
1779fn moe_vrows_dedup_order_on() -> bool {
1780    std::env::var("MEMRA_MOE_VROWS_DEDUP_ORDER").as_deref() == Ok("1")
1781}
1782
1783/// `MEMRA_MOE_VROWS_DOWN_TMAJ=1` (lane/glm5-dedup door E-down, default OFF): the verify-rows down
1784/// launch takes the `_tmaj` twin — grid transposed to `(t, out_f)` so the t verify rows at one
1785/// output row are adjacent blocks and a repeated expert's down row is read once for every token
1786/// sharing it. The down chain's slot-ordered `__fmaf_rn` accumulation is INSIDE the block and is
1787/// untouched (it keeps its original slot order — the vrest gate-4 bit bar); only the grid moves.
1788/// Split from [`moe_vrows_dedup_order_on`] as its own flag so the box can attribute the two
1789/// halves of the lever separately (gate/up is 2/3 of the pair's bytes, down 1/3). Same refusals:
1790/// door M wins, and `out_f > 65535` falls closed (a grid.y bound, not a serving shape).
1791fn moe_vrows_down_tmaj_on() -> bool {
1792    std::env::var("MEMRA_MOE_VROWS_DOWN_TMAJ").as_deref() == Ok("1")
1793}
1794
1795/// Engagement counter for the expert-major gate/up schedule (`MEMRA_MOE_VROWS_DEDUP_ORDER`).
1796pub static MOE_VROWS_DEDUP_ORDER_DISPATCHES: std::sync::atomic::AtomicU64 =
1797    std::sync::atomic::AtomicU64::new(0);
1798
1799/// Snapshot of [`MOE_VROWS_DEDUP_ORDER_DISPATCHES`] — gates take a before/after delta.
1800pub fn moe_vrows_dedup_order_dispatches() -> u64 {
1801    MOE_VROWS_DEDUP_ORDER_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1802}
1803
1804/// Engagement counter for the token-major down schedule (`MEMRA_MOE_VROWS_DOWN_TMAJ`).
1805pub static MOE_VROWS_DOWN_TMAJ_DISPATCHES: std::sync::atomic::AtomicU64 =
1806    std::sync::atomic::AtomicU64::new(0);
1807
1808/// Snapshot of [`MOE_VROWS_DOWN_TMAJ_DISPATCHES`] — gates take a before/after delta.
1809pub fn moe_vrows_down_tmaj_dispatches() -> u64 {
1810    MOE_VROWS_DOWN_TMAJ_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1811}
1812
1813/// AVOIDED SLAB READS — the box receipt for door E. Every layer-call adds `visits - distinct`,
1814/// i.e. the expert-slab reads whose repeat visit the expert-major schedule places inside the
1815/// reuse window. Multiply by the per-visit slab bytes (gate+up 9.4372 MB, down 4.7186 MB at the
1816/// serving geometry) for the bytes the schedule makes avoidable; that product is the CEILING of
1817/// the win, not the win (the realized share is a cache/scheduling property the box prices).
1818///
1819/// HOST-ARM ONLY, by construction: with door D on there is no host-side selection to count and a
1820/// 4-byte readback would reintroduce the very `cuStreamSynchronize` door D removed. The counting
1821/// boot is therefore `MEMRA_MOE_VROWS_DEV_TABLES=0`, exactly like the dedup instrument — while
1822/// [`MOE_VROWS_DEDUP_ORDER_DISPATCHES`] moves in BOTH table arms.
1823pub static MOE_VROWS_SLAB_READS_AVOIDED: std::sync::atomic::AtomicU64 =
1824    std::sync::atomic::AtomicU64::new(0);
1825
1826/// Snapshot of [`MOE_VROWS_SLAB_READS_AVOIDED`].
1827pub fn moe_vrows_slab_reads_avoided() -> u64 {
1828    MOE_VROWS_SLAB_READS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
1829}
1830
1831/// The EXPERT-MAJOR order plane, host build — the stable sort by `(expert id, pair index)` whose
1832/// bit-for-bit twin is the `moe_vrows_order_from_sel` counting rank. Returned as the `[n_pairs]`
1833/// tail plane the pointer table carries at `[3*n_pairs ..)`, and split out from the call site so
1834/// the device kernel can be gated against it directly.
1835pub(crate) fn vrows_expert_major_order(sel_all: &[u32]) -> Vec<u64> {
1836    let mut ord: Vec<u64> = (0..sel_all.len() as u64).collect();
1837    // Stable by construction: `sort_by_key` on the expert id keeps ascending pair order inside
1838    // each expert's run, so per-token slot order survives within a shared expert.
1839    ord.sort_by_key(|&p| sel_all[p as usize]);
1840    ord
1841}
1842
1843/// Gate hook for [`vrows_expert_major_order`] — the permutation is the whole door, so it is gated
1844/// against the device build and on planted selections rather than inferred from a live tape.
1845pub fn vrows_expert_major_order_for_test(sel_all: &[u32]) -> Vec<u64> {
1846    vrows_expert_major_order(sel_all)
1847}
1848
1849// ---- THE FLAG-ALIAS LAW for boolean doors (lane/glm5-extract2, phase 2) ------------------
1850//
1851// A door extracted from a family name to its general name keeps the FAMILY NAME HONORED:
1852// every banked gate script, box battery and in-flight lane sets the old name today, so
1853// refusing it would break receipts mid-bank for the price of one extra env read. Phase 1
1854// established the pattern for the two default-ON doors it moved (`MEMRA_VERIFY_WS`
1855// OFF-wins; `MEMRA_SPEC_TRACE` general-wins-loudly) and for the one VALUED door
1856// (`MEMRA_EP_MAP`, [`ep_map::resolve_ep_map_env`], which refuses a disagreeing pair at load).
1857// [`alias_door_from`] is the same law for a DEFAULT-OFF BOOLEAN door read PER CALL.
1858
1859/// Pure two-name resolution for a default-OFF boolean door (unit-tested without env
1860/// mutation — the phase-1 co-refusal-test pattern). Returns `(armed, the name the operator
1861/// actually set)` so every downstream refusal names the flag they typed, exactly as
1862/// [`ep_map::resolve_ep_map_env`] does for the valued seam.
1863///
1864/// * either name `=1` arms the door; anything else (including `=0`) is a deliberate pin;
1865/// * both set to the SAME value resolves to the general name;
1866/// * both set to DISAGREEING values is an operator error and is refused — `Err` carries the
1867///   message naming BOTH flags. The CALLER falls closed to the shipped program rather than
1868///   picking a precedence winner.
1869pub(crate) fn alias_door_from(
1870    general: (&'static str, Option<&str>),
1871    alias: (&'static str, Option<&str>),
1872) -> Result<(bool, &'static str), String> {
1873    match (general.1, alias.1) {
1874        (Some(g), Some(a)) if g != a => Err(format!(
1875            "{}={g:?} and {}={a:?} disagree — the alias and the general flag name ONE door \
1876             (unset one); refused rather than silently picking a precedence winner, and the \
1877             door falls closed to the shipped program",
1878            general.0, alias.0
1879        )),
1880        (Some(g), _) => Ok((g == "1", general.0)),
1881        (None, Some(a)) => Ok((a == "1", alias.0)),
1882        (None, None) => Ok((false, general.0)),
1883    }
1884}
1885
1886/// Env-reading wrapper over [`alias_door_from`]. A disagreeing pair FALLS CLOSED (door not
1887/// armed = the shipped program) and prints the refusal ONCE PER PROCESS through `latch`.
1888///
1889/// COST, stated because "read-site only" is true of the ARITHMETIC and not of the lookups:
1890/// honoring two names doubles the `env::var` calls on a per-call door (door H goes from ~64 to
1891/// ~128 lookups per ship round across `i32_mirror_store` and the shexp add), and `env::var`
1892/// takes the process environ lock. That is the price of not breaking every banked script, it
1893/// is paid only on doors whose call sites are already per-layer rather than per-token, and it
1894/// is unmeasured on a rig that cannot time host effects (LAW:rig-exactness-only). If a door
1895/// ever moves to a per-token site, resolve it once behind a `OnceLock` and give up the
1896/// in-process arm flipping the gates use today — that is the trade, named in advance.
1897///
1898/// It does not panic and it does not return `Result`: this is read per call inside the round,
1899/// and an abort in the GPU worker thread exits the process and kills every live session
1900/// (engine panics are fleet-fatal). A per-call door refuses by NOT ARMING; the loud line is
1901/// the operator's receipt that neither value won.
1902fn alias_door(
1903    general: &'static str,
1904    alias: &'static str,
1905    latch: &'static std::sync::atomic::AtomicBool,
1906) -> (bool, &'static str) {
1907    let g = std::env::var(general).ok();
1908    let a = std::env::var(alias).ok();
1909    match alias_door_from((general, g.as_deref()), (alias, a.as_deref())) {
1910        Ok(resolved) => resolved,
1911        Err(msg) => {
1912            if !latch.swap(true, std::sync::atomic::Ordering::Relaxed) {
1913                eprintln!("[flag-alias] {msg}");
1914            }
1915            (false, general)
1916        }
1917    }
1918}
1919
1920/// `MEMRA_HTOD_DIET=1` (default OFF; generalized from `MEMRA_GLM5_HTOD_DIET`, which stays
1921/// honored per the flag-alias law above — door H, lane/glm5-moe-loc): ENGINE-GENERIC HtoD
1922/// hygiene. Nothing in either class is family knowledge; both are "the host uploaded bytes
1923/// the device already had".
1924///
1925/// 1. The UNGATED shared-expert add re-uploaded a fresh `vec![1.0f32; t]` per MoE layer-call
1926///    (42 pageable HtoD/round to move a CONSTANT) — it now reads a resident ones buffer
1927///    ([`Engine::shexp_ones`]). Applies to every MoE family whose plan carries no
1928///    `ffn_gate_inp_shexp`.
1929/// 2. The latent-plane `len_d` i32 mirror took `memcpy_htod(&[v], ..)`, a SYNCHRONIZING
1930///    pageable copy, at 11 walk sites + 11 rollback sites per round. It now takes
1931///    [`Engine::i32_set_k`], the existing async twin whose value rides the kernel argument —
1932///    whose own doc already says the copy form is "fine at stream-idle boundaries, poison
1933///    mid-round". Applies to every latent-KV consumer ([`Engine::i32_mirror_store`] is an
1934///    Engine method, not a family method).
1935///
1936/// Both write identical values to identical buffers and both are stream-ordered, so the arms are
1937/// bit-identical by construction. Default OFF because no box timing receipt exists (rig is
1938/// exactness-only): 64 driver calls/round of measured count, UNPRICED wall. Read per call.
1939pub fn htod_diet_on() -> bool {
1940    htod_diet_armed().0
1941}
1942
1943/// Once-per-process latch for door H's disagreeing-pair line.
1944static HTOD_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
1945    std::sync::atomic::AtomicBool::new(false);
1946
1947/// Resolve door H, returning the armed flag name for refusals/announces.
1948pub(crate) fn htod_diet_armed() -> (bool, &'static str) {
1949    alias_door(
1950        "MEMRA_HTOD_DIET",
1951        "MEMRA_GLM5_HTOD_DIET",
1952        &HTOD_DIET_ALIAS_WARNED,
1953    )
1954}
1955
1956/// HtoD calls avoided by door H (`MEMRA_HTOD_DIET`): the count receipt. A gate asserts it
1957/// tracks the layer-call count on the ON arm and stays flat on the OFF arm.
1958pub static HTOD_DIET_AVOIDED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1959
1960/// Snapshot of [`HTOD_DIET_AVOIDED`] — gates take a before/after delta.
1961pub fn htod_diet_avoided() -> u64 {
1962    HTOD_DIET_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
1963}
1964
1965/// `MEMRA_EP_DIET=1` (default OFF; generalized from `MEMRA_GLM5_EP_DIET`, which stays honored
1966/// per the flag-alias law — lane/glm5-ep-diet): the EP DISPATCH DIET door, general to any
1967/// expert-parallel MoE walk. What the door names is a movement CLASS, not a family: one bulk
1968/// peer activation fan-out per layer-call instead of per-token uploads, compact peer staging
1969/// with one bulk return instead of a per-slot round-trip dribble, and one scatter launch
1970/// instead of the `t*n_used` sequential axpy chain. The glm5 TP-2 walk is today's CONSUMER
1971/// (its kernels, its combine order, its counters in `glm5_tp.rs`); hy3/step EP walks arm the
1972/// same door for their own walks.
1973///
1974/// The glm5 consumer's contract, unchanged: same per-slot expert kernels, same slot-ordered combine chain, restructured
1975/// data movement: ONE bulk peer z fan-out per layer-call (skipped entirely when no peer-owned
1976/// expert routed), zero per-slot host round-trips (peer rows stage compact on the peer and
1977/// return in ONE bulk DtoH+HtoD), and the t*n_used sequential `axpy_f32` combine launches
1978/// collapse into ONE `moe_pairs_scatter` launch — whose kernel header carries the
1979/// byte-identity contract vs the zeros+sequential-axpy chain. Decode stays BYTE-identical to
1980/// the v1 walk (and therefore to plain) by construction; `glm5-tp-gate` re-proves it with the
1981/// door pinned ON. Default OFF: the rig is exactness-only and the door changes the round's
1982/// SYNC STRUCTURE (the class the diet window warned does not always transfer from counts to
1983/// wall) — it ships with count receipts and the box window prices the wall. Read per call;
1984/// `=0`/unset restores the v1 per-slot walk byte-for-byte.
1985pub fn ep_diet_on() -> bool {
1986    ep_diet_armed().0
1987}
1988
1989/// Once-per-process latch for the EP-diet door's disagreeing-pair line.
1990static EP_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
1991    std::sync::atomic::AtomicBool::new(false);
1992
1993/// Resolve the EP-diet door, returning the armed flag name — the co-refusal in `hybrid.rs`
1994/// names the flag the operator actually set.
1995pub(crate) fn ep_diet_armed() -> (bool, &'static str) {
1996    alias_door("MEMRA_EP_DIET", "MEMRA_GLM5_EP_DIET", &EP_DIET_ALIAS_WARNED)
1997}
1998
1999/// `MEMRA_EP_GROUPED_PRIME=1` (default OFF; generalized from `MEMRA_GLM5_EP_GROUPED_PRIME`,
2000/// which stays honored per the flag-alias law — lane/glm5-ep-diet): the EP GROUPED-PRIME door,
2001/// general to any expert-parallel MoE walk — "run the family's own chunked grouped MoE prefill
2002/// program per rank over each rank's resident expert slab, then add the peer's bulk-returned
2003/// partial". The glm5 TP-2 walk is today's consumer.
2004///
2005/// The glm5 consumer's contract, unchanged: port the chunked
2006/// grouped MoE prefill (`MEMRA_MOE_GROUPED_PREFILL`, the plain walk's default-ON 85->616-639
2007/// tok/s prefill program) through the glm5 TP-2 EP walk: the SAME sigmoid host-oracle
2008/// routing, per-rank expert-major CSR restricted to each rank's owned experts, one grouped
2009/// f16 GEMM per projection PER RANK over the rank's resident EP slab (pointer tables minted
2010/// at arm time), per-rank slot-ordered scatter, then root adds the peer's bulk-returned
2011/// partial. Fires only where the plain grouped arm would (f16g-eligible qtypes, PRE-clamp,
2012/// n_used<=8); everything else — including the rig fixture's Q8_0 bank — falls closed to the
2013/// (dieted) sequential EP walk. Numeric class: per-expert GEMMs are the plain grouped arm's;
2014/// the ONE reassociation is the per-token root+peer partial add (band-gated, never claimed
2015/// byte). Read per call.
2016pub fn ep_grouped_prime_on() -> bool {
2017    ep_grouped_prime_armed().0
2018}
2019
2020/// Once-per-process latch for the EP grouped-prime door's disagreeing-pair line.
2021static EP_GROUPED_PRIME_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2022    std::sync::atomic::AtomicBool::new(false);
2023
2024/// Resolve the EP grouped-prime door, returning the armed flag name for the co-refusal.
2025pub(crate) fn ep_grouped_prime_armed() -> (bool, &'static str) {
2026    alias_door(
2027        "MEMRA_EP_GROUPED_PRIME",
2028        "MEMRA_GLM5_EP_GROUPED_PRIME",
2029        &EP_GROUPED_PRIME_ALIAS_WARNED,
2030    )
2031}
2032
2033/// `MEMRA_TOPK_SHARDS` (lane/glm5-matvec door K, default ON since the 2026-08-31 mv-battery
2034/// flip; `=0` is the rollback seam): `topk_rows` runs the exact
2035/// two-launch shard split (per-(row,shard) partial top-k + per-row shard merge) instead of the
2036/// one-block-per-row kernel. The standing kernel puts n_rows blocks on the card (the DFlash2
2037/// selector: 15 blocks on 188 SMs, 9.3 MB read in 1.31 ms = 7 GB/s). Top-k under the total
2038/// order (value desc, column asc) is a discrete selection, so the shard split is
2039/// OUTPUT-IDENTICAL by construction (same insertion comparisons, same tie rules in both
2040/// stages); gated by `glm5_matvec_doors_gpu` incl. planted-tie fixtures. Read per call.
2041fn topk_shards_on() -> bool {
2042    std::env::var("MEMRA_TOPK_SHARDS").as_deref() != Ok("0")
2043}
2044
2045/// Engagement counter for the sharded top-k door (`MEMRA_TOPK_SHARDS`).
2046pub static TOPK_SHARDS_DISPATCHES: std::sync::atomic::AtomicU64 =
2047    std::sync::atomic::AtomicU64::new(0);
2048
2049/// Snapshot of [`TOPK_SHARDS_DISPATCHES`] — gates take a before/after delta.
2050pub fn topk_shards_dispatches() -> u64 {
2051    TOPK_SHARDS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2052}
2053
2054/// `MEMRA_VERIFY_WS` (lane/glm5-matvec door W, default ON since the 2026-08-31 mv-battery
2055/// flip; `=0` is the rollback seam; generalized from `MEMRA_GLM5_VERIFY_WS`, which stays
2056/// honored as the family alias — OFF-WINS composition: either name `=0` disables, so every
2057/// banked gate arm and box script pinning the old name keeps its exact semantics, and the
2058/// old name is never silently dead): the verify walk's
2059/// recurring buffers draw from the engine's size-keyed free-lists and recycle back instead
2060/// of one `cuMemAllocAsync`+Free pair per buffer (~1380+1370 driver calls/token on the ship
2061/// shape — diet-battery apisum; `MEMRA_HC_DECODE_WS` owns only the t=1 walk and never
2062/// reaches the spec serving shape). Byte-identical by the sites' own full-overwrite uninit
2063/// contract; gated by `glm5_matvec_doors_gpu` (multi-call byte identity + the
2064/// `SCRATCH_ALLOC_CALLS` delta receipt). Read per call — the rollback seam.
2065fn verify_ws_on() -> bool {
2066    verify_ws_on_from(
2067        std::env::var("MEMRA_VERIFY_WS").ok().as_deref(),
2068        std::env::var("MEMRA_GLM5_VERIFY_WS").ok().as_deref(),
2069    )
2070}
2071
2072/// The pure OFF-wins composition over the general name and the glm5 alias (unit-tested
2073/// without env mutation; default ON, either name `=0` disables).
2074fn verify_ws_on_from(general: Option<&str>, glm5_alias: Option<&str>) -> bool {
2075    general != Some("0") && glm5_alias != Some("0")
2076}
2077
2078/// Engagement counter for the verify-walk workspace (`MEMRA_VERIFY_WS`): incremented
2079/// once per POOL HIT (a reused buffer = one avoided alloc + one avoided free). Gates anchor
2080/// on the delta; `SCRATCH_ALLOC_CALLS` carries the complementary real-alloc count.
2081pub static VERIFY_WS_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2082
2083/// Snapshot of [`VERIFY_WS_HITS`] — gates take a before/after delta.
2084pub fn verify_ws_hits() -> u64 {
2085    VERIFY_WS_HITS.load(std::sync::atomic::Ordering::Relaxed)
2086}
2087
2088/// Engagement counter for the glm5_next tensor-core MLA prefill chain
2089/// (`MEMRA_MLA_TC_PREFILL`), incremented once per (layer, chunk) dispatch at the chain's own
2090/// invocation, AFTER the strided-batched GEMM decline check — a declined shape does not count.
2091/// A gate that must prove "the TC arm ran N times for this workload" reads this delta; the
2092/// once-per-boot announce line dedups and cannot carry a count
2093/// (LAW:wiring-assertions-match-prose).
2094pub static MLA_TC_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2095    std::sync::atomic::AtomicU64::new(0);
2096
2097/// Snapshot of [`MLA_TC_PREFILL_DISPATCHES`]. Gates take a before/after pair around a workload
2098/// and assert on the delta — including the DECODE byte-identity gate, whose assertion is that
2099/// this stays FLAT across t=1 steps with the flag on.
2100pub fn mla_tc_prefill_dispatches() -> u64 {
2101    MLA_TC_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2102}
2103
2104/// Engagement counter for the glm5_next expert-grouped MoE PREFILL arm
2105/// (`MEMRA_MOE_GROUPED_PREFILL`), incremented once per (layer, chunk) dispatch at the arm's own
2106/// call site. Same reason the fused-epilogue counter exists: the observation env vars divert
2107/// dispatch, so a counter at the invocation is the only honest engagement receipt
2108/// (LAW:wiring-assertions-match-prose). Read via [`moe_grouped_prefill_dispatches`].
2109pub static MOE_GROUPED_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2110    std::sync::atomic::AtomicU64::new(0);
2111
2112/// Snapshot of [`MOE_GROUPED_PREFILL_DISPATCHES`]. Gates take a before/after pair around a
2113/// workload and assert on the delta.
2114pub fn moe_grouped_prefill_dispatches() -> u64 {
2115    MOE_GROUPED_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2116}
2117
2118/// RAII guard from `Engine::exact_scope`: restores the pre-scope `verify_exact` value on
2119/// drop, so error propagation (`?`) can never leave the engine latched in the
2120/// decode-exact matmul program (hermes finding, fixed 2026-08-23). Holds the flag, not
2121/// the Engine, so the restoration contract is unit-testable without a GPU.
2122#[must_use = "dropping immediately ends the exact scope"]
2123pub struct ExactScope<'a> {
2124    flag: &'a std::sync::atomic::AtomicBool,
2125    prev: bool,
2126}
2127
2128impl<'a> ExactScope<'a> {
2129    pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
2130        let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
2131        flag.store(on, std::sync::atomic::Ordering::Relaxed);
2132        ExactScope { flag, prev }
2133    }
2134}
2135
2136impl Drop for ExactScope<'_> {
2137    fn drop(&mut self) {
2138        self.flag
2139            .store(self.prev, std::sync::atomic::Ordering::Relaxed);
2140    }
2141}
2142
2143#[cfg(test)]
2144mod verify_ws_flag_tests {
2145    use super::verify_ws_on_from;
2146
2147    #[test]
2148    fn off_wins_across_general_and_alias() {
2149        // default ON
2150        assert!(verify_ws_on_from(None, None));
2151        // either name =0 disables (the banked gate arms pin the ALIAS =0; the general
2152        // name must be exactly as loud)
2153        assert!(!verify_ws_on_from(Some("0"), None));
2154        assert!(!verify_ws_on_from(None, Some("0")));
2155        assert!(!verify_ws_on_from(Some("1"), Some("0")));
2156        assert!(!verify_ws_on_from(Some("0"), Some("1")));
2157        // explicit ON on either name keeps the default
2158        assert!(verify_ws_on_from(Some("1"), None));
2159        assert!(verify_ws_on_from(None, Some("1")));
2160    }
2161}
2162
2163#[cfg(test)]
2164mod alias_door_tests {
2165    use super::alias_door_from;
2166
2167    const G: &str = "MEMRA_EP_DIET";
2168    const A: &str = "MEMRA_GLM5_EP_DIET";
2169
2170    fn r(g: Option<&str>, a: Option<&str>) -> Result<(bool, &'static str), String> {
2171        alias_door_from((G, g), (A, a))
2172    }
2173
2174    #[test]
2175    fn default_off_and_either_name_arms() {
2176        // unset/unset: the door is OFF and the general name is what a refusal would cite
2177        assert_eq!(r(None, None).unwrap(), (false, G));
2178        // either name =1 arms it, and the ARMED NAME is the one the operator set
2179        assert_eq!(r(Some("1"), None).unwrap(), (true, G));
2180        assert_eq!(r(None, Some("1")).unwrap(), (true, A));
2181        // =0 is a deliberate pin on either name, never an arming
2182        assert_eq!(r(Some("0"), None).unwrap(), (false, G));
2183        assert_eq!(r(None, Some("0")).unwrap(), (false, A));
2184        // anything that is not "1" is not an arming (no truthiness guessing)
2185        assert_eq!(r(None, Some("on")).unwrap(), (false, A));
2186        assert_eq!(r(Some(""), None).unwrap(), (false, G));
2187    }
2188
2189    #[test]
2190    fn agreeing_pair_resolves_to_the_general_name() {
2191        assert_eq!(r(Some("1"), Some("1")).unwrap(), (true, G));
2192        assert_eq!(r(Some("0"), Some("0")).unwrap(), (false, G));
2193    }
2194
2195    #[test]
2196    fn disagreeing_pair_refuses_and_names_both() {
2197        for (g, a) in [("1", "0"), ("0", "1")] {
2198            let err = r(Some(g), Some(a)).expect_err("a disagreeing pair must refuse");
2199            assert!(
2200                err.contains(G),
2201                "the refusal must name the general flag: {err}"
2202            );
2203            assert!(err.contains(A), "the refusal must name the alias: {err}");
2204            // and it must say which way it falls, so an operator reading the line knows the
2205            // door is CLOSED rather than guessing a precedence winner
2206            assert!(err.contains("falls closed"), "{err}");
2207        }
2208    }
2209}
2210
2211#[cfg(test)]
2212mod exact_scope_tests {
2213    use std::sync::atomic::{AtomicBool, Ordering};
2214
2215    #[test]
2216    fn error_path_restores_verify_exact() {
2217        // TOOTH (hermes finding, fixed 2026-08-23): dspark_spec_session_burst called
2218        // set_verify_exact(true)/(false) manually with `?`s in between — any error left
2219        // the engine latched in the decode-exact matmul program for every later request.
2220        // The RAII scope must restore across an error propagation.
2221        let flag = AtomicBool::new(false);
2222        let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
2223            let _scope = super::ExactScope::set(flag, true);
2224            assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
2225            Err("draft forward failed")? // the `?` exit the manual pair leaked on
2226        };
2227        assert!(failing(&flag).is_err());
2228        assert!(
2229            !flag.load(Ordering::Relaxed),
2230            "error propagation must restore the pre-scope value"
2231        );
2232        // Nested/previous-value contract: a scope entered while already ON restores ON.
2233        let flag = AtomicBool::new(true);
2234        {
2235            let _scope = super::ExactScope::set(&flag, true);
2236        }
2237        assert!(flag.load(Ordering::Relaxed));
2238        // Early drop ends the scope exactly where the manual `false` used to sit.
2239        let flag = AtomicBool::new(false);
2240        let scope = super::ExactScope::set(&flag, true);
2241        drop(scope);
2242        assert!(!flag.load(Ordering::Relaxed));
2243    }
2244}
2245
2246impl Engine {
2247    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
2248        let gpu = memra_runtime::Gpu::new(ordinal)?;
2249        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
2250        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
2251        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
2252        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
2253            use cudarc::driver::sys::CUdevice_attribute_enum as A;
2254            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
2255                .and_then(|d| unsafe {
2256                    Ok((
2257                        cudarc::driver::result::device::get_attribute(
2258                            d,
2259                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
2260                        )?,
2261                        cudarc::driver::result::device::get_attribute(
2262                            d,
2263                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
2264                        )?,
2265                    ))
2266                })
2267                .unwrap_or((0, 0));
2268            let built = env!("MEMRA_BUILT_CUDA_ARCH");
2269            let ok = matches!(
2270                (built, maj, min),
2271                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
2272            );
2273            if !ok {
2274                return Err(format!(
2275                    "memra was built for sm_{built} but device {ordinal} reports compute \
2276                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
2277                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
2278                )
2279                .into());
2280            }
2281        }
2282        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
2283        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
2284        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
2285        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
2286        unsafe {
2287            use cudarc::driver::sys;
2288            let dev: sys::CUdevice = ordinal as sys::CUdevice;
2289            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
2290            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
2291                let mut thresh: u64 = u64::MAX;
2292                let _ = sys::cuMemPoolSetAttribute(
2293                    pool,
2294                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
2295                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
2296                );
2297            }
2298        }
2299        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
2300        let hybrid = gpu
2301            .ctx
2302            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
2303        let kda = gpu.ctx.load_module(Ptx::from_binary(KDA_FATBIN.to_vec()))?;
2304        let qmatvec = gpu
2305            .ctx
2306            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
2307        let flash = gpu
2308            .ctx
2309            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
2310        let gemm = gpu
2311            .ctx
2312            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
2313        let router = gpu
2314            .ctx
2315            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
2316        let sample = gpu
2317            .ctx
2318            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
2319        let copy_stream = gpu.ctx.new_stream()?;
2320        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
2321        // cudarc is in multi-stream mode (main stream +
2322        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
2323        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
2324        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
2325        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
2326        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
2327        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
2328        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
2329        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
2330        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
2331        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
2332        // implicit event tracking.
2333        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
2334        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
2335        if std::env::var("MEMRA_EVT")
2336            .map(|v| v == "1")
2337            .unwrap_or(false)
2338        {
2339            // escape hatch: keep cudarc's implicit cross-stream event tracking.
2340        } else {
2341            unsafe {
2342                gpu.ctx.disable_event_tracking();
2343            }
2344        }
2345        Ok(Self {
2346            gpu,
2347            module,
2348            hybrid,
2349            kda,
2350            qmatvec,
2351            flash,
2352            flash_g: std::sync::OnceLock::new(),
2353            gemm,
2354            router,
2355            sample,
2356            moe_cache: Mutex::new(None),
2357            w8_mirrors: Mutex::new(std::collections::HashMap::new()),
2358            w8_act: Mutex::new(std::collections::HashMap::new()),
2359            moe_cache_layout: Mutex::new(None),
2360            copy_stream,
2361            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
2362            verify_exact: std::sync::atomic::AtomicBool::new(false),
2363            capture_keep: Mutex::new(Vec::new()),
2364            argmax_partials: Mutex::new(None),
2365            prime_deqw_ws: Mutex::new(None),
2366            router_stage: Mutex::new(None),
2367            hyper_decode_ws: Mutex::new(None),
2368            verify_ws: Mutex::new(VerifyWs::default()),
2369            vrows_macro_dev: Mutex::new(std::collections::HashMap::new()),
2370            shexp_ones: Mutex::new(None),
2371            fp8_scratch: Mutex::new(None),
2372            fa_vf16_scratch: Mutex::new(None),
2373            fa_part_pool: Mutex::new(None),
2374            fa_part_retired: Mutex::new(Vec::new()),
2375            fn_cache: Mutex::new(Default::default()),
2376            f16_scratch: Mutex::new(None),
2377            #[cfg(memra_cutlass)]
2378            cutlass_scratch: Mutex::new(None),
2379        })
2380    }
2381
2382    pub fn ctx(&self) -> &Arc<CudaContext> {
2383        &self.gpu.ctx
2384    }
2385
2386    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
2387    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
2388    ///
2389    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
2390    /// they are mapped to this process, so `free` counts them as gone, yet the very next
2391    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
2392    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
2393    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
2394    ///
2395    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
2396    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
2397    /// under-count headroom does not belong in a gate that queues real work, but the honest
2398    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
2399    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
2400    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
2401    ///
2402    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
2403    pub fn pool_cached_bytes(&self) -> usize {
2404        let (reserved, used) = self.pool_reserved_used();
2405        reserved.saturating_sub(used)
2406    }
2407
2408    /// Bytes the driver's per-device CUDA GRAPH memory pool currently holds RESERVED
2409    /// (`cuDeviceGetGraphMemAttribute` RESERVED_MEM_CURRENT) — the backing store of every
2410    /// captured alloc node, which on this engine means the dspark verify-graph pool
2411    /// (decode/step graphs bake pre-allocated buffers and own no alloc nodes). This memory
2412    /// is DISTINCT from the async pool above: `mem_get_info`'s `free` already excludes it,
2413    /// it is never released back (the vgraph pool has no eviction by design), and it GROWS
2414    /// as new (segment, vt)/(vt, rung, hi) keys capture — the growth is what
2415    /// `dspark_vg_admission_debt` charges at admission. Returns 0 if the attribute cannot
2416    /// be queried (never a false headroom claim, matching `pool_cached_bytes`).
2417    pub fn device_graph_mem_reserved(&self) -> usize {
2418        use cudarc::driver::sys as cus;
2419        let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
2420            return 0;
2421        };
2422        let mut bytes: u64 = 0;
2423        let rc = unsafe {
2424            cus::cuDeviceGetGraphMemAttribute(
2425                dev,
2426                cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
2427                &mut bytes as *mut u64 as *mut std::ffi::c_void,
2428            )
2429        };
2430        if rc == cus::cudaError_enum::CUDA_SUCCESS {
2431            bytes as usize
2432        } else {
2433            0
2434        }
2435    }
2436
2437    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
2438    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
2439    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
2440    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
2441    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
2442    /// (0, 0) if the pool cannot be queried.
2443    /// Release every CACHED (freed-but-retained) block of the default async mempool
2444    /// back to the driver (deploy-headroom lane, 2026-08-27). The boot-time
2445    /// RELEASE_THRESHOLD=u64::MAX pin keeps freed blocks cached for graph-launch speed,
2446    /// which is right for steady serving and wrong at a blue/green overlap: a green
2447    /// PROCESS cannot use blue's cached pool. cuMemPoolTrimTo(0) frees only unused
2448    /// blocks — live allocations are untouched; later allocs re-map once. Returns the
2449    /// bytes released (reserved delta), 0 if the pool cannot be queried.
2450    pub fn pool_trim_to_zero(&self) -> usize {
2451        use cudarc::driver::sys;
2452        let (before, _) = self.pool_reserved_used();
2453        unsafe {
2454            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
2455            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
2456                != sys::CUresult::CUDA_SUCCESS
2457            {
2458                return 0;
2459            }
2460            let _ = sys::cuMemPoolTrimTo(pool, 0);
2461        }
2462        let (after, _) = self.pool_reserved_used();
2463        before.saturating_sub(after)
2464    }
2465
2466    pub fn pool_reserved_used(&self) -> (usize, usize) {
2467        use cudarc::driver::sys;
2468        unsafe {
2469            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
2470            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
2471                != sys::CUresult::CUDA_SUCCESS
2472            {
2473                return (0, 0);
2474            }
2475            let (mut reserved, mut used) = (0u64, 0u64);
2476            if sys::cuMemPoolGetAttribute(
2477                pool,
2478                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
2479                &mut reserved as *mut u64 as *mut core::ffi::c_void,
2480            ) != sys::CUresult::CUDA_SUCCESS
2481            {
2482                return (0, 0);
2483            }
2484            if sys::cuMemPoolGetAttribute(
2485                pool,
2486                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
2487                &mut used as *mut u64 as *mut core::ffi::c_void,
2488            ) != sys::CUresult::CUDA_SUCCESS
2489            {
2490                return (0, 0);
2491            }
2492            (reserved as usize, used as usize)
2493        }
2494    }
2495
2496    /// Async-pool HIGH-WATER pair since the last reset: (RESERVED_MEM_HIGH, USED_MEM_HIGH)
2497    /// in bytes, then reset both watermarks to their CURRENT values
2498    /// (lane/step37-vram-admission-20260830). This is the instrument the boot admission
2499    /// calibration reads: engine transients are allocated and freed INSIDE one step, so any
2500    /// tick-boundary sampling of `mem_get_info`/pool-current sees nothing of the peak — the
2501    /// driver-kept watermark is the only honest record of how deep a burst actually dipped.
2502    /// (0, 0) if the pool cannot be queried (never a false claim, matching
2503    /// `pool_cached_bytes`).
2504    pub fn pool_high_water_reset(&self) -> (usize, usize) {
2505        use cudarc::driver::sys;
2506        unsafe {
2507            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
2508            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
2509                != sys::CUresult::CUDA_SUCCESS
2510            {
2511                return (0, 0);
2512            }
2513            let (mut reserved, mut used) = (0u64, 0u64);
2514            if sys::cuMemPoolGetAttribute(
2515                pool,
2516                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
2517                &mut reserved as *mut u64 as *mut core::ffi::c_void,
2518            ) != sys::CUresult::CUDA_SUCCESS
2519            {
2520                return (0, 0);
2521            }
2522            if sys::cuMemPoolGetAttribute(
2523                pool,
2524                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
2525                &mut used as *mut u64 as *mut core::ffi::c_void,
2526            ) != sys::CUresult::CUDA_SUCCESS
2527            {
2528                return (0, 0);
2529            }
2530            // Setting a *_HIGH attribute resets the watermark to the pool's current value
2531            // (the value argument must be 0 per the driver contract).
2532            let mut zero: u64 = 0;
2533            let _ = sys::cuMemPoolSetAttribute(
2534                pool,
2535                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
2536                &mut zero as *mut u64 as *mut core::ffi::c_void,
2537            );
2538            let mut zero2: u64 = 0;
2539            let _ = sys::cuMemPoolSetAttribute(
2540                pool,
2541                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
2542                &mut zero2 as *mut u64 as *mut core::ffi::c_void,
2543            );
2544            (reserved as usize, used as usize)
2545        }
2546    }
2547
2548    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
2549    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
2550    pub fn stream(&self) -> Arc<CudaStream> {
2551        self.gpu.stream()
2552    }
2553    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
2554    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
2555    pub fn gkv_on() -> bool {
2556        memra_kv::gkv_on()
2557    }
2558
2559    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
2560    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
2561    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
2562    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
2563    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
2564    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
2565    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
2566    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
2567    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
2568    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
2569    /// ON for both — no acceptance cost measured.
2570    pub fn wkv_on() -> bool {
2571        memra_kv::wkv_on()
2572    }
2573
2574    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
2575    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
2576    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
2577    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
2578    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
2579    pub fn kv_fp8_on() -> bool {
2580        memra_kv::kv_fp8_on()
2581    }
2582
2583    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
2584    /// when the fp8-globals arm is on; everything else from the default flash module.
2585    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
2586        if head_dim == 512 && Self::gkv_on() {
2587            self.func_g(name)
2588        } else {
2589            self.func(name)
2590        }
2591    }
2592
2593    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
2594    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
2595    /// per-format fatbins; fall back to the base modules for those.
2596    fn func_g(&self, name: &str) -> CudaFunction {
2597        let m = self.flash_g.get_or_init(|| {
2598            self.gpu
2599                .ctx
2600                .load_module(cudarc::nvrtc::Ptx::from_binary(
2601                    FLASH_FATBIN_KF8VF8.to_vec(),
2602                ))
2603                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
2604        });
2605        let key = format!("g:{name}");
2606        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
2607            return f.clone();
2608        }
2609        let f = match m.load_function(name) {
2610            Ok(f) => f,
2611            Err(_) => self.func(name),
2612        };
2613        self.fn_cache.lock().unwrap().insert(key, f.clone());
2614        f
2615    }
2616
2617    fn func(&self, name: &str) -> CudaFunction {
2618        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
2619        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
2620        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
2621            return f.clone();
2622        }
2623        let f = self
2624            .module
2625            .load_function(name)
2626            .or_else(|_| self.hybrid.load_function(name))
2627            .or_else(|_| self.kda.load_function(name))
2628            .or_else(|_| self.qmatvec.load_function(name))
2629            .or_else(|_| self.flash.load_function(name))
2630            .or_else(|_| self.gemm.load_function(name))
2631            .or_else(|_| self.router.load_function(name))
2632            .or_else(|_| self.sample.load_function(name))
2633            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
2634        self.fn_cache
2635            .lock()
2636            .unwrap()
2637            .insert(name.to_string(), f.clone());
2638        f
2639    }
2640
2641    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
2642    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
2643    pub fn scatter_trim_logits(
2644        &self,
2645        src: &CudaSlice<f32>,
2646        d2t: &CudaSlice<u32>,
2647        dst: &mut CudaSlice<f32>,
2648        d_vocab: usize,
2649        n_vocab: usize,
2650    ) -> Result<(), Box<dyn std::error::Error>> {
2651        let f1 = self.func("scatter_trim_logits_f32");
2652        let f2 = self.func("scatter_trim_logits_pass2_f32");
2653        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
2654        let cfg1 = LaunchConfig {
2655            grid_dim: (256, 1, 1),
2656            block_dim: (256, 1, 1),
2657            shared_mem_bytes: 0,
2658        };
2659        let __s_b1 = self.gpu.stream();
2660        let mut b1 = __s_b1.launch_builder(&f1);
2661        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
2662        unsafe {
2663            b1.launch(cfg1)?;
2664        }
2665        let cfg2 = LaunchConfig {
2666            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
2667            block_dim: (256, 1, 1),
2668            shared_mem_bytes: 0,
2669        };
2670        let __s_b2 = self.gpu.stream();
2671        let mut b2 = __s_b2.launch_builder(&f2);
2672        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
2673        unsafe {
2674            b2.launch(cfg2)?;
2675        }
2676        Ok(())
2677    }
2678
2679    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
2680    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
2681
2682    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
2683    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
2684    #[allow(clippy::too_many_arguments)]
2685    pub fn filter_stats(
2686        &self,
2687        x: &CudaSlice<f32>,
2688        row_stride: usize,
2689        rows: &CudaSlice<i32>,
2690        out_th: &mut CudaSlice<f32>,
2691        out_z: &mut CudaSlice<f32>,
2692        out_max: &mut CudaSlice<f32>,
2693        n: usize,
2694        nrow: usize,
2695        temp: f32,
2696        top_k: i32,
2697        top_p: f32,
2698        min_p: f32,
2699    ) -> Result<(), Box<dyn std::error::Error>> {
2700        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
2701        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
2702        // L2-resident, so the extra passes are near-free while the per-thread selection list
2703        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
2704        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
2705        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
2706        //
2707        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
2708        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
2709        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
2710        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
2711        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
2712        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
2713        //
2714        // DETERMINISTIC KEYING (hermes finding, fixed 2026-08-23): the old admission
2715        // `16*nrow <= sm_count` fell back to the single-block program PER CALL when a tick
2716        // carried too many rows — and the two programs are NOT bit-identical (measured
2717        // ~1e-7 rel on the renorm mass: different f32 partial-sum order), so a request's
2718        // sampling threshold arithmetic depended on how many rows shared its serve tick.
2719        // Coop is now THE program on every coop-capable device: rows are CHUNKED to the
2720        // co-residency cap (sm_count/16 rows per cooperative launch) and each row's
2721        // arithmetic uses only its own 16 slices + its own ws region, so the per-row bits
2722        // are independent of batch width by construction — the kernel-check
2723        // FILTER-COOP-CHUNK cell pins exactly that. The single-block program remains only
2724        // behind the deployment-keyed seams: MEMRA_FILTER_COOP=0, or a device with
2725        // sm_count < 16 (fixed per device class, never per call).
2726        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2727        let coop_on =
2728            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
2729        if coop_on && self.sm_count() >= 16 {
2730            let cap = self.sm_count() as usize / 16;
2731            let mut done = 0usize;
2732            while done < nrow {
2733                let chunk = cap.min(nrow - done);
2734                self.filter_stats_coop_chunk(
2735                    x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
2736                    top_p, min_p,
2737                )?;
2738                done += chunk;
2739            }
2740            return Ok(());
2741        }
2742        self.filter_stats_plain_program(
2743            x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
2744        )
2745    }
2746
2747    /// One cooperative `filter_stats` launch over rows `row0..row0+chunk` (pub so the
2748    /// kernel-check FILTER-COOP-CHUNK cell can pin batch-width independence directly).
2749    /// The kernel indexes `rows`/outputs by blockIdx.y, so the chunk is expressed as
2750    /// sub-views at `row0` — per-row arithmetic is untouched by the offset.
2751    #[allow(clippy::too_many_arguments)]
2752    pub fn filter_stats_coop_chunk(
2753        &self,
2754        x: &CudaSlice<f32>,
2755        row_stride: usize,
2756        rows: &CudaSlice<i32>,
2757        row0: usize,
2758        out_th: &mut CudaSlice<f32>,
2759        out_z: &mut CudaSlice<f32>,
2760        out_max: &mut CudaSlice<f32>,
2761        n: usize,
2762        chunk: usize,
2763        temp: f32,
2764        top_k: i32,
2765        top_p: f32,
2766        min_p: f32,
2767    ) -> Result<(), Box<dyn std::error::Error>> {
2768        let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
2769        let f = self.func("filter_stats_coop_f32");
2770        let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
2771        let cfg = LaunchConfig {
2772            grid_dim: (16, chunk as u32, 1),
2773            block_dim: (512, 1, 1),
2774            shared_mem_bytes: 0,
2775        };
2776        let rows_v = rows.slice(row0..row0 + chunk);
2777        let mut th_v = out_th.slice_mut(row0..row0 + chunk);
2778        let mut z_v = out_z.slice_mut(row0..row0 + chunk);
2779        let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
2780        let __s_b = self.gpu.stream();
2781        let mut b = __s_b.launch_builder(&f);
2782        b.arg(x)
2783            .arg(&rs)
2784            .arg(&rows_v)
2785            .arg(&mut th_v)
2786            .arg(&mut z_v)
2787            .arg(&mut mx_v)
2788            .arg(&mut ws)
2789            .arg(&ni)
2790            .arg(&nr)
2791            .arg(&temp)
2792            .arg(&top_k)
2793            .arg(&top_p)
2794            .arg(&min_p);
2795        unsafe {
2796            b.launch_cooperative(cfg)?;
2797        }
2798        Ok(())
2799    }
2800
2801    /// The single-block-per-row `filter_stats` program (the pre-coop form; the
2802    /// MEMRA_FILTER_COOP=0 rollback and the occupancy fallback). Gate-callable twin of
2803    /// `filter_stats_coop_program`.
2804    #[allow(clippy::too_many_arguments)]
2805    pub fn filter_stats_plain_program(
2806        &self,
2807        x: &CudaSlice<f32>,
2808        row_stride: usize,
2809        rows: &CudaSlice<i32>,
2810        out_th: &mut CudaSlice<f32>,
2811        out_z: &mut CudaSlice<f32>,
2812        out_max: &mut CudaSlice<f32>,
2813        n: usize,
2814        nrow: usize,
2815        temp: f32,
2816        top_k: i32,
2817        top_p: f32,
2818        min_p: f32,
2819    ) -> Result<(), Box<dyn std::error::Error>> {
2820        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
2821        let f = self.func("filter_stats_f32");
2822        let cfg = LaunchConfig {
2823            grid_dim: (nrow as u32, 1, 1),
2824            block_dim: (1024, 1, 1),
2825            shared_mem_bytes: 0,
2826        };
2827        let __s_b = self.gpu.stream();
2828        let mut b = __s_b.launch_builder(&f);
2829        b.arg(x)
2830            .arg(&rs)
2831            .arg(rows)
2832            .arg(&mut *out_th)
2833            .arg(&mut *out_z)
2834            .arg(&mut *out_max)
2835            .arg(&ni)
2836            .arg(&nr)
2837            .arg(&temp)
2838            .arg(&top_k)
2839            .arg(&top_p)
2840            .arg(&min_p);
2841        unsafe {
2842            b.launch(cfg)?;
2843        }
2844        Ok(())
2845    }
2846
2847    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
2848    #[allow(clippy::too_many_arguments)]
2849    pub fn softmax_gather_filtered(
2850        &self,
2851        x: &CudaSlice<f32>,
2852        row_stride: usize,
2853        ids: &CudaSlice<u32>,
2854        rows: &CudaSlice<i32>,
2855        th: &CudaSlice<f32>,
2856        z: &CudaSlice<f32>,
2857        out: &mut CudaSlice<f32>,
2858        n: usize,
2859        npair: usize,
2860        temp: f32,
2861    ) -> Result<(), Box<dyn std::error::Error>> {
2862        let f = self.func("softmax_gather_filtered_f32");
2863        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
2864        let cfg = LaunchConfig {
2865            grid_dim: (npair as u32, 1, 1),
2866            block_dim: (256, 1, 1),
2867            shared_mem_bytes: 0,
2868        };
2869        let __s_b = self.gpu.stream();
2870        let mut b = __s_b.launch_builder(&f);
2871        b.arg(x)
2872            .arg(&rs)
2873            .arg(ids)
2874            .arg(rows)
2875            .arg(th)
2876            .arg(z)
2877            .arg(&mut *out)
2878            .arg(&ni)
2879            .arg(&np)
2880            .arg(&temp);
2881        unsafe {
2882            b.launch(cfg)?;
2883        }
2884        Ok(())
2885    }
2886
2887    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
2888    #[allow(clippy::too_many_arguments)]
2889    pub fn residual_sample_filtered(
2890        &self,
2891        p: &CudaSlice<f32>,
2892        q: Option<&CudaSlice<f32>>,
2893        n: usize,
2894        temp: f32,
2895        seed: u64,
2896        stream_pos: u32,
2897        p_stats: (f32, f32, f32),
2898        q_stats: (f32, f32, f32),
2899        out_tok: &mut CudaSlice<u32>,
2900    ) -> Result<(), Box<dyn std::error::Error>> {
2901        let f = self.func("residual_sample_filtered_f32");
2902        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2903        let has_q: i32 = q.is_some() as i32;
2904        let qbuf = q.unwrap_or(p);
2905        let (pm, pth, pz) = p_stats;
2906        let (qm, qth, qz) = q_stats;
2907        let cfg = LaunchConfig {
2908            grid_dim: (1, 1, 1),
2909            block_dim: (1024, 1, 1),
2910            shared_mem_bytes: 0,
2911        };
2912        let __s_b = self.gpu.stream();
2913        let mut b = __s_b.launch_builder(&f);
2914        b.arg(p)
2915            .arg(qbuf)
2916            .arg(&has_q)
2917            .arg(&ni)
2918            .arg(&temp)
2919            .arg(&slo)
2920            .arg(&shi)
2921            .arg(&stream_pos)
2922            .arg(&pm)
2923            .arg(&pth)
2924            .arg(&pz)
2925            .arg(&qm)
2926            .arg(&qth)
2927            .arg(&qz)
2928            .arg(&mut *out_tok);
2929        unsafe {
2930            b.launch(cfg)?;
2931        }
2932        Ok(())
2933    }
2934
2935    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
2936    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
2937    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
2938    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
2939    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
2940    #[allow(clippy::too_many_arguments)]
2941    pub fn residual_sample_sparse_q(
2942        &self,
2943        p: &CudaSlice<f32>,
2944        cand_ids: &CudaSlice<u32>,
2945        q_probs: &CudaSlice<f32>,
2946        n_cand: usize,
2947        n: usize,
2948        temp: f32,
2949        seed: u64,
2950        stream_pos: u32,
2951        p_stats: (f32, f32, f32),
2952        out_tok: &mut CudaSlice<u32>,
2953    ) -> Result<(), Box<dyn std::error::Error>> {
2954        assert!(
2955            (1..=32).contains(&n_cand),
2956            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
2957        );
2958        let f = self.func("residual_sample_sparse_q_f32");
2959        let (ni, nc) = (n as i32, n_cand as i32);
2960        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2961        let (pm, pth, pz) = p_stats;
2962        let cfg = LaunchConfig {
2963            grid_dim: (1, 1, 1),
2964            block_dim: (1024, 1, 1),
2965            shared_mem_bytes: 0,
2966        };
2967        let __s_b = self.gpu.stream();
2968        let mut b = __s_b.launch_builder(&f);
2969        b.arg(p)
2970            .arg(cand_ids)
2971            .arg(q_probs)
2972            .arg(&nc)
2973            .arg(&ni)
2974            .arg(&temp)
2975            .arg(&slo)
2976            .arg(&shi)
2977            .arg(&stream_pos)
2978            .arg(&pm)
2979            .arg(&pth)
2980            .arg(&pz)
2981            .arg(&mut *out_tok);
2982        unsafe {
2983            b.launch(cfg)?;
2984        }
2985        Ok(())
2986    }
2987
2988    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
2989    #[allow(clippy::too_many_arguments)]
2990    pub fn gumbel_perturb_filtered(
2991        &self,
2992        x: &CudaSlice<f32>,
2993        y: &mut CudaSlice<f32>,
2994        n: usize,
2995        seed: u64,
2996        stream_pos: u32,
2997        temp: f32,
2998        row_max: f32,
2999        th: f32,
3000    ) -> Result<(), Box<dyn std::error::Error>> {
3001        let f = self.func("gumbel_perturb_filtered_f32");
3002        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3003        let cfg = LaunchConfig {
3004            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3005            block_dim: (256, 1, 1),
3006            shared_mem_bytes: 0,
3007        };
3008        let __s_b = self.gpu.stream();
3009        let mut b = __s_b.launch_builder(&f);
3010        b.arg(x)
3011            .arg(&mut *y)
3012            .arg(&ni)
3013            .arg(&slo)
3014            .arg(&shi)
3015            .arg(&stream_pos)
3016            .arg(&temp)
3017            .arg(&row_max)
3018            .arg(&th);
3019        unsafe {
3020            b.launch(cfg)?;
3021        }
3022        Ok(())
3023    }
3024
3025    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
3026    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
3027    /// filtered rejection sampling exact for the penalized target.
3028    #[allow(clippy::too_many_arguments)]
3029    pub fn penalize_logits(
3030        &self,
3031        x: &mut CudaSlice<f32>,
3032        hist: &CudaSlice<u32>,
3033        n_hist: usize,
3034        rep: f32,
3035        freq: f32,
3036        present: f32,
3037        n: usize,
3038    ) -> Result<(), Box<dyn std::error::Error>> {
3039        if n_hist == 0 {
3040            return Ok(());
3041        }
3042        let f = self.func("penalize_logits_f32");
3043        let (nh, ni) = (n_hist as i32, n as i32);
3044        let cfg = LaunchConfig {
3045            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
3046            block_dim: (128, 1, 1),
3047            shared_mem_bytes: 0,
3048        };
3049        let __s_b = self.gpu.stream();
3050        let mut b = __s_b.launch_builder(&f);
3051        b.arg(&mut *x)
3052            .arg(hist)
3053            .arg(&nh)
3054            .arg(&rep)
3055            .arg(&freq)
3056            .arg(&present)
3057            .arg(&ni);
3058        unsafe {
3059            b.launch(cfg)?;
3060        }
3061        Ok(())
3062    }
3063
3064    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
3065    #[allow(clippy::too_many_arguments)]
3066    pub fn penalize_logits_rows(
3067        &self,
3068        x: &mut CudaSlice<f32>,
3069        hist: &CudaSlice<u32>,
3070        n_hist: usize,
3071        rep: f32,
3072        freq: f32,
3073        present: f32,
3074        n: usize,
3075        nrow: usize,
3076    ) -> Result<(), Box<dyn std::error::Error>> {
3077        if n_hist == 0 || nrow == 0 {
3078            return Ok(());
3079        }
3080        let f = self.func("penalize_logits_rows_f32");
3081        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
3082        let cfg = LaunchConfig {
3083            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
3084            block_dim: (128, 1, 1),
3085            shared_mem_bytes: 0,
3086        };
3087        let __s_b = self.gpu.stream();
3088        let mut b = __s_b.launch_builder(&f);
3089        b.arg(&mut *x)
3090            .arg(hist)
3091            .arg(&nh)
3092            .arg(&rep)
3093            .arg(&freq)
3094            .arg(&present)
3095            .arg(&ni)
3096            .arg(&nr);
3097        unsafe {
3098            b.launch(cfg)?;
3099        }
3100        Ok(())
3101    }
3102
3103    /// Heterogeneous serving-batch penalties over host-maintained sparse window counts.
3104    /// `offsets[r]..offsets[r+1]` indexes the unique positive-count `(id,count)` entries for logits row
3105    /// `rows[r]`; each row may carry independent repetition/frequency/presence coefficients.
3106    /// One thread owns one distinct logit, so the kernel needs neither atomics nor the
3107    /// history-squared dedup scan used by the speculative raw-history oracle.
3108    #[allow(clippy::too_many_arguments)]
3109    pub fn penalize_logits_sparse_rows(
3110        &self,
3111        x: &mut CudaSlice<f32>,
3112        ids: &[u32],
3113        counts: &[u32],
3114        offsets: &[i32],
3115        rows: &[i32],
3116        reps: &[f32],
3117        freqs: &[f32],
3118        presents: &[f32],
3119        n: usize,
3120    ) -> Result<(), Box<dyn std::error::Error>> {
3121        let nrow = rows.len();
3122        if nrow == 0 {
3123            return Ok(());
3124        }
3125        let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
3126        let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
3127        let entry_count =
3128            i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
3129        if ids.len() != counts.len()
3130            || offsets.len() != nrow + 1
3131            || reps.len() != nrow
3132            || freqs.len() != nrow
3133            || presents.len() != nrow
3134            || offsets.first().copied() != Some(0)
3135            || offsets.last().copied() != Some(entry_count)
3136        {
3137            return Err("sparse penalty row metadata shape mismatch".into());
3138        }
3139        if counts.contains(&0) {
3140            return Err("sparse penalty counts must be positive".into());
3141        }
3142        let mut max_len = 0usize;
3143        for pair in offsets.windows(2) {
3144            if pair[0] < 0 || pair[1] < pair[0] {
3145                return Err("sparse penalty offsets must be monotonic".into());
3146            }
3147            max_len = max_len.max((pair[1] - pair[0]) as usize);
3148        }
3149        if max_len == 0 {
3150            return Ok(());
3151        }
3152
3153        let mut seen = std::collections::HashSet::with_capacity(ids.len());
3154        for (r, &row) in rows.iter().enumerate() {
3155            if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
3156                return Err("sparse penalty row index exceeds logits shape".into());
3157            }
3158            let begin = offsets[r] as usize;
3159            let end = offsets[r + 1] as usize;
3160            for &id in &ids[begin..end] {
3161                if id as usize >= n {
3162                    return Err("sparse penalty token id exceeds logits row".into());
3163                }
3164                if !seen.insert((row, id)) {
3165                    return Err("sparse penalty entries must be unique per logits row".into());
3166                }
3167            }
3168        }
3169
3170        // SAFETY: the checks above establish every invariant of the launch-only helper.
3171        unsafe {
3172            self.penalize_logits_sparse_rows_unchecked(
3173                x, ids, counts, offsets, rows, reps, freqs, presents, n,
3174            )
3175        }
3176    }
3177
3178    /// Launch-only form for the serving hot path, whose `HashMap`-backed producer already
3179    /// guarantees unique ids and whose rows are enumerated from the live batch.
3180    ///
3181    /// # Safety
3182    ///
3183    /// Shapes must match the safe wrapper, offsets must be monotonic and in bounds, every row
3184    /// must index `x`, and each `(row,id)` pair must occur at most once. Token ids outside the
3185    /// logits row are safe no-ops because the kernel bounds-checks them before computing `x`.
3186    #[allow(clippy::too_many_arguments)]
3187    pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
3188        &self,
3189        x: &mut CudaSlice<f32>,
3190        ids: &[u32],
3191        counts: &[u32],
3192        offsets: &[i32],
3193        rows: &[i32],
3194        reps: &[f32],
3195        freqs: &[f32],
3196        presents: &[f32],
3197        n: usize,
3198    ) -> Result<(), Box<dyn std::error::Error>> {
3199        let nrow = rows.len();
3200        if nrow == 0 {
3201            return Ok(());
3202        }
3203        let max_len = offsets
3204            .windows(2)
3205            .map(|pair| (pair[1] - pair[0]) as usize)
3206            .max()
3207            .unwrap_or(0);
3208        if max_len == 0 {
3209            return Ok(());
3210        }
3211        let ids_d = self.htod_u32_v(ids)?;
3212        let counts_d = self.htod_u32_v(counts)?;
3213        let offsets_d = self.htod_i32(offsets)?;
3214        let rows_d = self.htod_i32(rows)?;
3215        let reps_d = self.htod(reps)?;
3216        let freqs_d = self.htod(freqs)?;
3217        let presents_d = self.htod(presents)?;
3218        let f = self.func("penalize_logits_sparse_rows_f32");
3219        let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
3220        let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
3221        let cfg = LaunchConfig {
3222            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
3223            block_dim: (128, 1, 1),
3224            shared_mem_bytes: 0,
3225        };
3226        let __s_b = self.gpu.stream();
3227        let mut b = __s_b.launch_builder(&f);
3228        b.arg(&mut *x)
3229            .arg(&ids_d)
3230            .arg(&counts_d)
3231            .arg(&offsets_d)
3232            .arg(&rows_d)
3233            .arg(&reps_d)
3234            .arg(&freqs_d)
3235            .arg(&presents_d)
3236            .arg(&ni)
3237            .arg(&nr);
3238        unsafe {
3239            b.launch(cfg)?;
3240        }
3241        Ok(())
3242    }
3243
3244    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
3245    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
3246    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
3247    /// is the within-round evolving penalty state block drafting needs: verify row r's
3248    /// target is penalized by every token committed before it INCLUDING same-round
3249    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
3250    /// approximation this exists to replace on the dspark route.
3251    #[allow(clippy::too_many_arguments)]
3252    pub fn penalize_logits_rows_inc(
3253        &self,
3254        x: &mut CudaSlice<f32>,
3255        hist: &CudaSlice<u32>,
3256        n_hist0: usize,
3257        rep: f32,
3258        freq: f32,
3259        present: f32,
3260        n: usize,
3261        nrow: usize,
3262        win: usize,
3263    ) -> Result<(), Box<dyn std::error::Error>> {
3264        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
3265            return Ok(());
3266        }
3267        debug_assert!(
3268            hist.len() >= n_hist0 + nrow - 1,
3269            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
3270        );
3271        let f = self.func("penalize_logits_rows_inc_f32");
3272        let max_len = win.min(n_hist0 + nrow - 1).max(1);
3273        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
3274        let cfg = LaunchConfig {
3275            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
3276            block_dim: (128, 1, 1),
3277            shared_mem_bytes: 0,
3278        };
3279        let __s_b = self.gpu.stream();
3280        let mut b = __s_b.launch_builder(&f);
3281        b.arg(&mut *x)
3282            .arg(hist)
3283            .arg(&nh)
3284            .arg(&rep)
3285            .arg(&freq)
3286            .arg(&present)
3287            .arg(&ni)
3288            .arg(&nr)
3289            .arg(&wi);
3290        unsafe {
3291            b.launch(cfg)?;
3292        }
3293        Ok(())
3294    }
3295
3296    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
3297    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
3298    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
3299    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
3300    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
3301    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
3302    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
3303    pub fn wpf_level() -> u32 {
3304        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
3305        *ON.get_or_init(|| {
3306            std::env::var("MEMRA_WPF")
3307                .ok()
3308                .and_then(|v| v.parse().ok())
3309                .unwrap_or(1)
3310        })
3311    }
3312
3313    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
3314    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
3315    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
3316    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
3317    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
3318    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
3319    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
3320    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
3321    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
3322    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
3323    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
3324    /// Prefer `exact_scope` — the RAII form — anywhere a `?` can exit the scope: a manual
3325    /// true/false pair leaves the flag LATCHED engine-wide when an error propagates
3326    /// between the two calls (hermes finding on dspark_spec_session_burst, fixed
3327    /// 2026-08-23), and every later request then runs the exact-GEMM program.
3328    pub fn set_verify_exact(&self, on: bool) {
3329        self.verify_exact
3330            .store(on, std::sync::atomic::Ordering::Relaxed);
3331    }
3332    pub(crate) fn verify_exact_on(&self) -> bool {
3333        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
3334    }
3335
3336    /// RAII scope over `verify_exact`: sets the flag to `on` now and restores the
3337    /// PREVIOUS value on drop — unwind, early `return`, and every `?` exit included.
3338    /// This is the required form for any scope an error can leave (see
3339    /// `set_verify_exact`); dropping the guard early (`drop(scope)`) ends the scope
3340    /// exactly where the manual `set_verify_exact(false)` used to sit.
3341    pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
3342        ExactScope::set(&self.verify_exact, on)
3343    }
3344
3345    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
3346    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
3347    pub fn qkv_append_on() -> bool {
3348        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3349        *ON.get_or_init(|| {
3350            std::env::var("MEMRA_QKV_APPEND")
3351                .map(|v| v != "0")
3352                .unwrap_or(true)
3353        })
3354    }
3355
3356    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
3357    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
3358    pub fn pdl_wb_on() -> bool {
3359        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3360        *ON.get_or_init(|| {
3361            std::env::var("MEMRA_PDL_WB")
3362                .map(|v| v != "0")
3363                .unwrap_or(true)
3364        })
3365    }
3366
3367    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
3368    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
3369    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
3370    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
3371    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
3372    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
3373    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
3374    pub fn norm_ilp_on() -> bool {
3375        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3376        *ON.get_or_init(|| {
3377            std::env::var("MEMRA_NORM_ILP")
3378                .map(|v| v != "0")
3379                .unwrap_or(true)
3380        })
3381    }
3382
3383    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
3384    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
3385    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
3386    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
3387    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
3388    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
3389    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
3390    pub fn tk_ffn_dual_on() -> bool {
3391        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3392        *ON.get_or_init(|| {
3393            std::env::var("MEMRA_TK_FFN_DUAL")
3394                .map(|v| v != "0")
3395                .unwrap_or(true)
3396        })
3397    }
3398
3399    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
3400    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
3401    /// per-model no-harm bisect knob.
3402    pub fn pdl_mmvq_on() -> bool {
3403        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3404        *ON.get_or_init(|| {
3405            std::env::var("MEMRA_PDL_MMVQ")
3406                .map(|v| v != "0")
3407                .unwrap_or(true)
3408        })
3409    }
3410
3411    pub fn pdl_on() -> bool {
3412        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3413        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
3414    }
3415
3416    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
3417    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
3418    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
3419    /// on the producer before any read), bit-identical by construction.
3420    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
3421    pub fn pdl_nvfp4q8_on() -> bool {
3422        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3423        *ON.get_or_init(|| {
3424            std::env::var("MEMRA_PDL_NVFP4")
3425                .map(|v| v != "0")
3426                .unwrap_or(true)
3427        })
3428    }
3429
3430    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
3431    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
3432    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
3433    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
3434    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
3435    fn q40_mr1_on() -> bool {
3436        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
3437        match *Q40MR.get_or_init(|| {
3438            std::env::var("MEMRA_Q40_MR")
3439                .ok()
3440                .and_then(|v| v.parse().ok())
3441        }) {
3442            Some(v) => v == 1,
3443            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3444        }
3445    }
3446
3447    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
3448    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
3449    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
3450    /// writes wrong bytes silently.
3451    fn pdl_func_flash(
3452        &self,
3453        g: bool,
3454        name: &'static str,
3455    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
3456        use cudarc::driver::sys as cu;
3457        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
3458        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
3459        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
3460        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
3461        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
3462        // this engine's CUcontext; single-context runs behave exactly as before.
3463        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
3464            std::sync::Mutex::new(None);
3465        #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3466        static FNS: std::sync::Mutex<
3467            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
3468        > = std::sync::Mutex::new(None);
3469        let ctx_key = self.ctx().cu_ctx() as usize;
3470        if let Some(&f) = FNS
3471            .lock()
3472            .unwrap()
3473            .get_or_insert_with(Default::default)
3474            .get(&(ctx_key, g, name))
3475        {
3476            return Ok(f as cu::CUfunction);
3477        }
3478        let module = {
3479            let mut mods = MODS.lock().unwrap();
3480            let map = mods.get_or_insert_with(Default::default);
3481            match map.get(&(ctx_key, g)) {
3482                Some(&m) => m,
3483                None => {
3484                    let m = self.pdl_load_module_in_ctx(if g {
3485                        FLASH_FATBIN_KF8VF8
3486                    } else {
3487                        FLASH_FATBIN
3488                    })?;
3489                    map.insert((ctx_key, g), m);
3490                    m
3491                }
3492            }
3493        };
3494        let cname = std::ffi::CString::new(name)?;
3495        let mut f: cu::CUfunction = std::ptr::null_mut();
3496        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
3497        if r != cu::CUresult::CUDA_SUCCESS {
3498            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
3499        }
3500        FNS.lock()
3501            .unwrap()
3502            .get_or_insert_with(Default::default)
3503            .insert((ctx_key, g, name), f as usize);
3504        Ok(f)
3505    }
3506
3507    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
3508    /// the module to the thread's CURRENT context — a remote-stage engine must not
3509    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
3510    /// current context before returning.
3511    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
3512        use cudarc::driver::sys as cu;
3513        let mut prev: cu::CUcontext = std::ptr::null_mut();
3514        unsafe {
3515            cu::cuCtxGetCurrent(&mut prev).result()?;
3516        }
3517        self.ctx().bind_to_thread()?;
3518        let mut m: cu::CUmodule = std::ptr::null_mut();
3519        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
3520        let restore = if prev.is_null() {
3521            cu::CUresult::CUDA_SUCCESS
3522        } else {
3523            unsafe { cu::cuCtxSetCurrent(prev) }
3524        };
3525        if r != cu::CUresult::CUDA_SUCCESS {
3526            return Err(format!("pdl module load: {r:?}").into());
3527        }
3528        if restore != cu::CUresult::CUDA_SUCCESS {
3529            return Err(format!("pdl module load: ctx restore {restore:?}").into());
3530        }
3531        Ok(m as usize)
3532    }
3533
3534    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
3535    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
3536    pub fn raw_kernel_function(
3537        &self,
3538        name: &'static str,
3539    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
3540        self.pdl_func(name)
3541    }
3542
3543    fn pdl_func(
3544        &self,
3545        name: &'static str,
3546    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
3547        use cudarc::driver::sys as cu;
3548        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
3549        // are context-scoped; key everything by this engine's CUcontext).
3550        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
3551            std::sync::Mutex::new(None);
3552        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
3553        // duplicate module, loaded lazily on the first kernels-module miss.
3554        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
3555            std::sync::Mutex::new(None);
3556        static FNS: std::sync::Mutex<
3557            Option<std::collections::HashMap<(usize, &'static str), usize>>,
3558        > = std::sync::Mutex::new(None);
3559        let ctx_key = self.ctx().cu_ctx() as usize;
3560        if let Some(&f) = FNS
3561            .lock()
3562            .unwrap()
3563            .get_or_insert_with(Default::default)
3564            .get(&(ctx_key, name))
3565        {
3566            return Ok(f as cu::CUfunction);
3567        }
3568        let module = {
3569            let mut mods = MODULES.lock().unwrap();
3570            let map = mods.get_or_insert_with(Default::default);
3571            match map.get(&ctx_key) {
3572                Some(&m) => m,
3573                None => {
3574                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
3575                    map.insert(ctx_key, m);
3576                    m
3577                }
3578            }
3579        };
3580        let cname = std::ffi::CString::new(name)?;
3581        let mut f: cu::CUfunction = std::ptr::null_mut();
3582        let mut r =
3583            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
3584        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
3585            let qmodule = {
3586                let mut mods = QMODULES.lock().unwrap();
3587                let map = mods.get_or_insert_with(Default::default);
3588                match map.get(&ctx_key) {
3589                    Some(&m) => m,
3590                    None => {
3591                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
3592                        map.insert(ctx_key, m);
3593                        m
3594                    }
3595                }
3596            };
3597            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
3598        }
3599        if r != cu::CUresult::CUDA_SUCCESS {
3600            return Err(format!("pdl_func {name}: {r:?}").into());
3601        }
3602        FNS.lock()
3603            .unwrap()
3604            .get_or_insert_with(Default::default)
3605            .insert((ctx_key, name), f as usize);
3606        Ok(f)
3607    }
3608
3609    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
3610    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
3611    ///
3612    /// # Safety
3613    /// `params` must match the kernel's exact parameter list (order, types, count) —
3614    /// a mismatch corrupts the launch silently.
3615    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
3616    /// builder path's fa_func/func_g choice exactly).
3617    ///
3618    /// # Safety
3619    /// Same contract as `launch_pdl`.
3620    unsafe fn launch_pdl_flash(
3621        &self,
3622        g: bool,
3623        name: &'static str,
3624        grid: (u32, u32, u32),
3625        block: (u32, u32, u32),
3626        smem: u32,
3627        params: &mut [*mut std::ffi::c_void],
3628    ) -> Result<(), Box<dyn std::error::Error>> {
3629        use cudarc::driver::sys as cu;
3630        let f = self.pdl_func_flash(g, name)?;
3631        if smem > 0 {
3632            // mirror the builder path's opt-in ceiling (idempotent host-side set).
3633            let r =
3634                unsafe {
3635                    cu::cuFuncSetAttribute(f,
3636                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
3637                smem as i32)
3638                };
3639            if r != cu::CUresult::CUDA_SUCCESS {
3640                return Err(format!("pdl smem attr {name}: {r:?}").into());
3641            }
3642        }
3643        let mut attr = cu::CUlaunchAttribute {
3644            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
3645            pad: [0; 4],
3646            value: cu::CUlaunchAttributeValue {
3647                programmaticStreamSerializationAllowed: 1,
3648            },
3649        };
3650        let cfg = cu::CUlaunchConfig {
3651            gridDimX: grid.0,
3652            gridDimY: grid.1,
3653            gridDimZ: grid.2,
3654            blockDimX: block.0,
3655            blockDimY: block.1,
3656            blockDimZ: block.2,
3657            sharedMemBytes: smem,
3658            hStream: self.gpu.stream().cu_stream(),
3659            attrs: &mut attr,
3660            numAttrs: 1,
3661        };
3662        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
3663        if r != cu::CUresult::CUDA_SUCCESS {
3664            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
3665        }
3666        Ok(())
3667    }
3668
3669    unsafe fn launch_pdl(
3670        &self,
3671        name: &'static str,
3672        grid: (u32, u32, u32),
3673        block: (u32, u32, u32),
3674        params: &mut [*mut std::ffi::c_void],
3675    ) -> Result<(), Box<dyn std::error::Error>> {
3676        use cudarc::driver::sys as cu;
3677        let f = self.pdl_func(name)?;
3678        let mut attr = cu::CUlaunchAttribute {
3679            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
3680            pad: [0; 4],
3681            value: cu::CUlaunchAttributeValue {
3682                programmaticStreamSerializationAllowed: 1,
3683            },
3684        };
3685        let cfg = cu::CUlaunchConfig {
3686            gridDimX: grid.0,
3687            gridDimY: grid.1,
3688            gridDimZ: grid.2,
3689            blockDimX: block.0,
3690            blockDimY: block.1,
3691            blockDimZ: block.2,
3692            sharedMemBytes: 0,
3693            hStream: self.gpu.stream().cu_stream(),
3694            attrs: &mut attr,
3695            numAttrs: 1,
3696        };
3697        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
3698        if r != cu::CUresult::CUDA_SUCCESS {
3699            return Err(format!("launch_pdl {name}: {r:?}").into());
3700        }
3701        Ok(())
3702    }
3703
3704    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
3705    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
3706    pub fn prefetch_weight_l2(
3707        &self,
3708        w: &crate::model::GpuTensor,
3709    ) -> Result<(), Box<dyn std::error::Error>> {
3710        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
3711            let p = rp4.as_ref().unwrap_or(bytes);
3712            self.prefetch_l2(p, p.len())?;
3713        }
3714        Ok(())
3715    }
3716
3717    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
3718    /// by the DEVICE token id at tok[idx] into f32.
3719    pub fn gather_row_bf16(
3720        &self,
3721        table: &CudaSlice<u8>,
3722        tok: &CudaSlice<u32>,
3723        idx: usize,
3724        dst: &mut CudaSlice<f32>,
3725        ncols: usize,
3726    ) -> Result<(), Box<dyn std::error::Error>> {
3727        let f = self.func("gather_row_bf16_f32");
3728        let cfg = LaunchConfig {
3729            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
3730            block_dim: (256, 1, 1),
3731            shared_mem_bytes: 0,
3732        };
3733        let (nc, ix) = (ncols as i32, idx as i32);
3734        let __s_b = self.gpu.stream();
3735        let mut b = __s_b.launch_builder(&f);
3736        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
3737        unsafe {
3738            b.launch(cfg)?;
3739        }
3740        Ok(())
3741    }
3742
3743    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
3744    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
3745    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
3746    ///   `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
3747    ///   finish(1).
3748    #[allow(clippy::too_many_arguments)]
3749    pub fn dflash2_dynconv(
3750        &self,
3751        x: &CudaSlice<f32>,
3752        dyn_: &CudaSlice<f32>,
3753        base: &CudaSlice<f32>,
3754        out: &mut CudaSlice<f32>,
3755        rows: usize,
3756        hidden: usize,
3757        group_size: usize,
3758        ksize: usize,
3759        half: usize,
3760    ) -> Result<(), Box<dyn std::error::Error>> {
3761        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
3762        let f = self.func("dflash2_dynconv_f32");
3763        let n = rows * hidden;
3764        let cfg = LaunchConfig {
3765            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3766            block_dim: (256, 1, 1),
3767            shared_mem_bytes: 0,
3768        };
3769        let (ri, hi, gi, ki, hf) = (
3770            rows as i32,
3771            hidden as i32,
3772            group_size as i32,
3773            ksize as i32,
3774            half as i32,
3775        );
3776        let __s_b = self.gpu.stream();
3777        let mut b = __s_b.launch_builder(&f);
3778        b.arg(x)
3779            .arg(dyn_)
3780            .arg(base)
3781            .arg(out)
3782            .arg(&ri)
3783            .arg(&hi)
3784            .arg(&gi)
3785            .arg(&ki)
3786            .arg(&hf);
3787        unsafe {
3788            b.launch(cfg)?;
3789        }
3790        Ok(())
3791    }
3792
3793    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
3794    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
3795    /// value-descending, ties to the lower index.
3796    pub fn topk_rows(
3797        &self,
3798        logits: &CudaSlice<f32>,
3799        n_rows: usize,
3800        n_cols: usize,
3801        k: usize,
3802    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
3803        assert!((1..=32).contains(&k), "topk_rows supports 1..=32, got {k}");
3804        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
3805        // MEMRA_TOPK_SHARDS (lane/glm5-matvec door K, default ON since 2026-08-31): the exact two-launch
3806        // shard split — n_rows*16 partial blocks + a per-row merge — instead of n_rows
3807        // blocks total (the DFlash2 selector: 15 blocks on the whole card, 7 GB/s). Top-k
3808        // under (value desc, column asc) is discrete selection: output-identical by
3809        // construction, gated by glm5_matvec_doors_gpu. Small columns fall through (the
3810        // shard overhead would dominate and the standing grid is already wide enough).
3811        if topk_shards_on() && n_cols >= 16 * 1024 && k <= n_cols / 16 {
3812            if TOPK_SHARDS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
3813                eprintln!(
3814                    "[topk-shards] engaged: rows={n_rows} cols={n_cols} k={k} shards=16 \
3815                     (MEMRA_TOPK_SHARDS=1)"
3816                );
3817            }
3818            return self.topk_rows_sharded(logits, n_rows, n_cols, k, 16);
3819        }
3820        let f = self.func("topk_rows_f32");
3821        let nth = 256usize;
3822        let mut vals = self.uninit(n_rows * k)?;
3823        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
3824        let cfg = LaunchConfig {
3825            grid_dim: (n_rows as u32, 1, 1),
3826            block_dim: (nth as u32, 1, 1),
3827            shared_mem_bytes: (nth * k * 8) as u32,
3828        };
3829        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
3830        let __s_b = self.gpu.stream();
3831        let mut b = __s_b.launch_builder(&f);
3832        b.arg(logits)
3833            .arg(&nr)
3834            .arg(&nc)
3835            .arg(&ki)
3836            .arg(&mut vals)
3837            .arg(&mut idxs);
3838        unsafe {
3839            b.launch(cfg)?;
3840        }
3841        Ok((vals, idxs))
3842    }
3843
3844    /// The exact two-launch shard split behind `MEMRA_TOPK_SHARDS` (see [`Self::topk_rows`]):
3845    /// per-(row, shard) partial top-k with the standing kernel's insertion/tie rules on
3846    /// global column indices, then a per-row k-way merge with the standing kernel's merge
3847    /// rules. Output-identical to `topk_rows_f32` by construction (discrete selection under
3848    /// the total order value-desc/index-asc); gated by `glm5_matvec_doors_gpu`.
3849    fn topk_rows_sharded(
3850        &self,
3851        logits: &CudaSlice<f32>,
3852        n_rows: usize,
3853        n_cols: usize,
3854        k: usize,
3855        n_shards: usize,
3856    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
3857        assert!((1..=64).contains(&n_shards), "shard merge head cap is 64");
3858        let nth = 256usize;
3859        let mut pvals = self.uninit(n_rows * n_shards * k)?;
3860        let mut pidxs = self.alloc_uninit::<u32>(n_rows * n_shards * k)?;
3861        let f1 = self.func("topk_rows_shard_f32");
3862        let cfg1 = LaunchConfig {
3863            grid_dim: (n_rows as u32, n_shards as u32, 1),
3864            block_dim: (nth as u32, 1, 1),
3865            shared_mem_bytes: (nth * k * 8) as u32,
3866        };
3867        let (nr, nc, ki, ns) = (n_rows as i32, n_cols as i32, k as i32, n_shards as i32);
3868        {
3869            let __s_b = self.gpu.stream();
3870            let mut b = __s_b.launch_builder(&f1);
3871            b.arg(logits)
3872                .arg(&nr)
3873                .arg(&nc)
3874                .arg(&ki)
3875                .arg(&ns)
3876                .arg(&mut pvals)
3877                .arg(&mut pidxs);
3878            unsafe {
3879                b.launch(cfg1)?;
3880            }
3881        }
3882        let mut vals = self.uninit(n_rows * k)?;
3883        let mut idxs = self.alloc_uninit::<u32>(n_rows * k)?;
3884        let f2 = self.func("topk_rows_shard_merge_f32");
3885        let cfg2 = LaunchConfig {
3886            grid_dim: (n_rows as u32, 1, 1),
3887            block_dim: (32, 1, 1),
3888            shared_mem_bytes: 0,
3889        };
3890        let __s_b = self.gpu.stream();
3891        let mut b = __s_b.launch_builder(&f2);
3892        b.arg(&pvals)
3893            .arg(&pidxs)
3894            .arg(&nr)
3895            .arg(&ns)
3896            .arg(&ki)
3897            .arg(&mut vals)
3898            .arg(&mut idxs);
3899        unsafe {
3900            b.launch(cfg2)?;
3901        }
3902        Ok((vals, idxs))
3903    }
3904
3905    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
3906    pub fn add_row_inplace(
3907        &self,
3908        logits: &mut CudaSlice<f32>,
3909        bias: &CudaSlice<f32>,
3910        n: usize,
3911        row_off: usize,
3912    ) -> Result<(), Box<dyn std::error::Error>> {
3913        let f = self.func("add_row_inplace_f32");
3914        let cfg = LaunchConfig {
3915            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3916            block_dim: (256, 1, 1),
3917            shared_mem_bytes: 0,
3918        };
3919        let (ni, off) = (n as i32, row_off as i64);
3920        let __s_b = self.gpu.stream();
3921        let mut b = __s_b.launch_builder(&f);
3922        b.arg(logits).arg(bias).arg(&ni).arg(&off);
3923        unsafe {
3924            b.launch(cfg)?;
3925        }
3926        Ok(())
3927    }
3928
3929    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
3930    pub fn prefetch_l2(
3931        &self,
3932        p: &CudaSlice<u8>,
3933        n: usize,
3934    ) -> Result<(), Box<dyn std::error::Error>> {
3935        let f = self.func("prefetch_l2_bytes");
3936        let lines = n.div_ceil(128);
3937        let ni = n as i64;
3938        let cfg = LaunchConfig {
3939            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
3940            block_dim: (256, 1, 1),
3941            shared_mem_bytes: 0,
3942        };
3943        let __s_b = self.gpu.stream();
3944        let mut b = __s_b.launch_builder(&f);
3945        b.arg(p).arg(&ni);
3946        unsafe {
3947            b.launch(cfg)?;
3948        }
3949        Ok(())
3950    }
3951
3952    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
3953    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
3954    pub fn router_gemv(
3955        &self,
3956        w: &CudaSlice<f32>,
3957        x: &CudaSlice<f32>,
3958        n_embd: usize,
3959        n_experts: usize,
3960        t: usize,
3961    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3962        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
3963        // stream differs) — too small to justify a numeric config change; deleted.
3964        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
3965        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
3966        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
3967        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
3968            Ok("0") => false,
3969            Ok(_) => true,
3970            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3971        };
3972        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
3973        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
3974        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
3975        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
3976        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
3977        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
3978        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
3979        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
3980        // (perf-only, bits equal).
3981        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
3982        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
3983    }
3984
3985    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
3986    /// force both forms; `batch` requires `w8`).
3987    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
3988    pub fn router_gemv_form(
3989        &self,
3990        w: &CudaSlice<f32>,
3991        x: &CudaSlice<f32>,
3992        n_embd: usize,
3993        n_experts: usize,
3994        t: usize,
3995        w8: bool,
3996        batch: bool,
3997    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3998        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
3999        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
4000        let f = if batch {
4001            self.func("router_gemv_f32_w8_batch")
4002        } else if w8 {
4003            self.func("router_gemv_f32_w8")
4004        } else {
4005            self.func("router_gemv_f32")
4006        };
4007        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4008        let cfg = if batch {
4009            LaunchConfig {
4010                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
4011                block_dim: (32, 8, 1),
4012                shared_mem_bytes: 0,
4013            }
4014        } else {
4015            LaunchConfig {
4016                grid_dim: (n_experts as u32, t as u32, 1),
4017                block_dim: (32, if w8 { 8 } else { 1 }, 1),
4018                shared_mem_bytes: 0,
4019            }
4020        };
4021        let __s_b = self.gpu.stream();
4022        let mut b = __s_b.launch_builder(&f);
4023        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
4024        unsafe {
4025            b.launch(cfg)?;
4026        }
4027        Ok(y)
4028    }
4029
4030    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
4031    /// buffer — token-graph alloc-free.
4032    pub fn router_gemv_into(
4033        &self,
4034        w: &CudaSlice<f32>,
4035        x: &CudaSlice<f32>,
4036        y: &mut CudaSlice<f32>,
4037        n_embd: usize,
4038        n_experts: usize,
4039        t: usize,
4040    ) -> Result<(), Box<dyn std::error::Error>> {
4041        if y.len() < t * n_experts {
4042            return Err("router_gemv_into output too small".into());
4043        }
4044        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4045            Ok("0") => false,
4046            Ok(_) => true,
4047            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4048        };
4049        let f = if w8 {
4050            self.func("router_gemv_f32_w8")
4051        } else {
4052            self.func("router_gemv_f32")
4053        };
4054        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4055        let cfg = LaunchConfig {
4056            grid_dim: (n_experts as u32, t as u32, 1),
4057            block_dim: (32, if w8 { 8 } else { 1 }, 1),
4058            shared_mem_bytes: 0,
4059        };
4060        let __s_b = self.gpu.stream();
4061        let mut b = __s_b.launch_builder(&f);
4062        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
4063        unsafe {
4064            b.launch(cfg)?;
4065        }
4066        Ok(())
4067    }
4068
4069    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
4070    pub fn rows_permute(
4071        &self,
4072        src: &CudaSlice<f32>,
4073        idx: &CudaSlice<i32>,
4074        nrows: usize,
4075        ncols: usize,
4076    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4077        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
4078        let f = self.func("rows_permute_f32");
4079        let (nc, nr) = (ncols as i32, nrows as i32);
4080        let cfg = LaunchConfig {
4081            grid_dim: (nrows as u32, 1, 1),
4082            block_dim: (256, 1, 1),
4083            shared_mem_bytes: 0,
4084        };
4085        let __s_b = self.gpu.stream();
4086        let mut b = __s_b.launch_builder(&f);
4087        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
4088        unsafe {
4089            b.launch(cfg)?;
4090        }
4091        Ok(dst)
4092    }
4093
4094    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
4095    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
4096    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
4097    /// decode chain and the small-t spec-verify chain match per row by construction.
4098    pub fn sigmoid_dot_rows(
4099        &self,
4100        x: &CudaSlice<f32>,
4101        w: &CudaSlice<f32>,
4102        n_embd: usize,
4103        t: usize,
4104    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4105        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
4106        // config; same class as MEMRA_ROUTER_V2).
4107        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4108        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
4109            let gs = self.linear(x, w, t, n_embd, 1)?;
4110            let mut g = self.uninit(t)?;
4111            self.sigmoid(&gs, &mut g, t)?;
4112            return Ok(g);
4113        }
4114        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
4115        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
4116        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
4117        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
4118        // flags doctrine; this per-token form serves every t.
4119        let mut g = self.alloc_uninit::<f32>(t)?;
4120        let f = self.func("sigmoid_dot_rows_f32");
4121        let (ne, ti) = (n_embd as i32, t as i32);
4122        let cfg = LaunchConfig {
4123            grid_dim: (t as u32, 1, 1),
4124            block_dim: (32, 8, 1),
4125            shared_mem_bytes: 0,
4126        };
4127        let __s_b = self.gpu.stream();
4128        let mut b = __s_b.launch_builder(&f);
4129        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
4130        unsafe {
4131            b.launch(cfg)?;
4132        }
4133        Ok(g)
4134    }
4135
4136    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
4137    pub fn sigmoid_dot_rows_into(
4138        &self,
4139        x: &CudaSlice<f32>,
4140        w: &CudaSlice<f32>,
4141        g: &mut CudaSlice<f32>,
4142        n_embd: usize,
4143        t: usize,
4144    ) -> Result<(), Box<dyn std::error::Error>> {
4145        if g.len() < t {
4146            return Err("sigmoid_dot_rows_into output too small".into());
4147        }
4148        let f = self.func("sigmoid_dot_rows_f32");
4149        let (ne, ti) = (n_embd as i32, t as i32);
4150        let cfg = LaunchConfig {
4151            grid_dim: (t as u32, 1, 1),
4152            block_dim: (32, 8, 1),
4153            shared_mem_bytes: 0,
4154        };
4155        let __s_b = self.gpu.stream();
4156        let mut b = __s_b.launch_builder(&f);
4157        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
4158        unsafe {
4159            b.launch(cfg)?;
4160        }
4161        Ok(())
4162    }
4163
4164    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
4165    pub fn spec_rollback_stream(
4166        &self,
4167        len_ptrs: &CudaSlice<u64>,
4168        pos_start: &CudaSlice<i32>,
4169        acc: &CudaSlice<u32>,
4170        base: usize,
4171        n_rows: usize,
4172    ) -> Result<(), Box<dyn std::error::Error>> {
4173        let f = self.func("spec_rollback_stream");
4174        let (b, nr) = (base as i32, n_rows as i32);
4175        let cfg = LaunchConfig {
4176            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
4177            block_dim: (64, 1, 1),
4178            shared_mem_bytes: 0,
4179        };
4180        let __s_bl = self.gpu.stream();
4181        let mut bl = __s_bl.launch_builder(&f);
4182        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
4183        unsafe {
4184            bl.launch(cfg)?;
4185        }
4186        Ok(())
4187    }
4188
4189    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
4190    pub fn plain_tok_ring(
4191        &self,
4192        vam: &CudaSlice<u32>,
4193        pos_start: &CudaSlice<i32>,
4194        base: usize,
4195        ring: &mut CudaSlice<u32>,
4196    ) -> Result<(), Box<dyn std::error::Error>> {
4197        let f = self.func("plain_tok_ring");
4198        let (b, cap) = (base as i32, ring.len() as i32);
4199        let cfg = LaunchConfig {
4200            grid_dim: (1, 1, 1),
4201            block_dim: (32, 1, 1),
4202            shared_mem_bytes: 0,
4203        };
4204        let __s_bl = self.gpu.stream();
4205        let mut bl = __s_bl.launch_builder(&f);
4206        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
4207        unsafe {
4208            bl.launch(cfg)?;
4209        }
4210        Ok(())
4211    }
4212
4213    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
4214    pub fn spec_ring_commit(
4215        &self,
4216        vtok: &CudaSlice<u32>,
4217        acc: &CudaSlice<u32>,
4218        brk: &CudaSlice<u32>,
4219        ring: &mut CudaSlice<u32>,
4220        pend: &mut CudaSlice<u32>,
4221    ) -> Result<(), Box<dyn std::error::Error>> {
4222        let f = self.func("spec_ring_commit");
4223        let cfg = LaunchConfig {
4224            grid_dim: (1, 1, 1),
4225            block_dim: (32, 1, 1),
4226            shared_mem_bytes: 0,
4227        };
4228        let __s_b = self.gpu.stream();
4229        let mut b = __s_b.launch_builder(&f);
4230        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
4231        unsafe {
4232            b.launch(cfg)?;
4233        }
4234        Ok(())
4235    }
4236    pub fn i32_copy_add(
4237        &self,
4238        src: &CudaSlice<i32>,
4239        dst: &mut CudaSlice<i32>,
4240        delta: i32,
4241    ) -> Result<(), Box<dyn std::error::Error>> {
4242        let f = self.func("i32_copy_add");
4243        let cfg = LaunchConfig {
4244            grid_dim: (1, 1, 1),
4245            block_dim: (32, 1, 1),
4246            shared_mem_bytes: 0,
4247        };
4248        let __s_b = self.gpu.stream();
4249        let mut b = __s_b.launch_builder(&f);
4250        b.arg(src).arg(dst).arg(&delta);
4251        unsafe {
4252            b.launch(cfg)?;
4253        }
4254        Ok(())
4255    }
4256    pub fn u32_copy(
4257        &self,
4258        src: &CudaSlice<u32>,
4259        dst: &mut CudaSlice<u32>,
4260    ) -> Result<(), Box<dyn std::error::Error>> {
4261        let f = self.func("u32_copy");
4262        let cfg = LaunchConfig {
4263            grid_dim: (1, 1, 1),
4264            block_dim: (32, 1, 1),
4265            shared_mem_bytes: 0,
4266        };
4267        let __s_b = self.gpu.stream();
4268        let mut b = __s_b.launch_builder(&f);
4269        b.arg(src).arg(dst);
4270        unsafe {
4271            b.launch(cfg)?;
4272        }
4273        Ok(())
4274    }
4275
4276    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
4277    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
4278    /// caps acceptance exactly like drafting fewer tokens).
4279    pub fn spec_adapt_k(
4280        &self,
4281        acc: &CudaSlice<u32>,
4282        brk: &mut CudaSlice<u32>,
4283        floor: usize,
4284        cap: usize,
4285    ) -> Result<(), Box<dyn std::error::Error>> {
4286        let f = self.func("spec_adapt_k");
4287        let (fl, cp) = (floor as i32, cap as i32);
4288        let cfg = LaunchConfig {
4289            grid_dim: (1, 1, 1),
4290            block_dim: (32, 1, 1),
4291            shared_mem_bytes: 0,
4292        };
4293        let __s_b = self.gpu.stream();
4294        let mut b = __s_b.launch_builder(&f);
4295        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
4296        unsafe {
4297            b.launch(cfg)?;
4298        }
4299        Ok(())
4300    }
4301
4302    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
4303    pub fn spec_accept_greedy_dc(
4304        &self,
4305        preds: &CudaSlice<u32>,
4306        vtok: &CudaSlice<u32>,
4307        last_pred: &CudaSlice<u32>,
4308        brk: &CudaSlice<u32>,
4309        out: &mut CudaSlice<u32>,
4310    ) -> Result<(), Box<dyn std::error::Error>> {
4311        let f = self.func("spec_accept_greedy_dc");
4312        let cfg = LaunchConfig {
4313            grid_dim: (1, 1, 1),
4314            block_dim: (32, 1, 1),
4315            shared_mem_bytes: 0,
4316        };
4317        let __s_b = self.gpu.stream();
4318        let mut b = __s_b.launch_builder(&f);
4319        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
4320        unsafe {
4321            b.launch(cfg)?;
4322        }
4323        Ok(())
4324    }
4325
4326    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
4327    pub fn pos_iota(
4328        &self,
4329        pos0: &CudaSlice<i32>,
4330        out: &mut CudaSlice<i32>,
4331        t: usize,
4332    ) -> Result<(), Box<dyn std::error::Error>> {
4333        let f = self.func("pos_iota_i32");
4334        let ti = t as i32;
4335        let cfg = LaunchConfig {
4336            grid_dim: (1, 1, 1),
4337            block_dim: (t.max(1) as u32, 1, 1),
4338            shared_mem_bytes: 0,
4339        };
4340        let __s_b = self.gpu.stream();
4341        let mut b = __s_b.launch_builder(&f);
4342        b.arg(pos0).arg(out).arg(&ti);
4343        unsafe {
4344            b.launch(cfg)?;
4345        }
4346        Ok(())
4347    }
4348    #[allow(clippy::too_many_arguments)]
4349    pub fn append_kv_quantized_rows_dc(
4350        &self,
4351        k_rows: &CudaSlice<f32>,
4352        v_rows: &CudaSlice<f32>,
4353        kc: &mut CudaSlice<u8>,
4354        vc: &mut CudaSlice<u8>,
4355        t0_dev: &CudaSlice<i32>,
4356        t: usize,
4357        kv_dim_k: usize,
4358        kv_dim_v: usize,
4359        k_tok_bytes: usize,
4360        v_tok_bytes: usize,
4361        g: bool,
4362    ) -> Result<(), Box<dyn std::error::Error>> {
4363        let f = if g {
4364            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
4365        } else {
4366            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
4367        };
4368        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4369        let cfg = LaunchConfig {
4370            grid_dim: (nblk, t as u32, 1),
4371            block_dim: (32, 1, 1),
4372            shared_mem_bytes: 0,
4373        };
4374        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4375        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4376        let __s_b = self.gpu.stream();
4377        let mut b = __s_b.launch_builder(&f);
4378        b.arg(k_rows)
4379            .arg(v_rows)
4380            .arg(kc)
4381            .arg(vc)
4382            .arg(t0_dev)
4383            .arg(&kdk)
4384            .arg(&kdv)
4385            .arg(&ktb)
4386            .arg(&vtb);
4387        unsafe {
4388            b.launch(cfg)?;
4389        }
4390        Ok(())
4391    }
4392
4393    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
4394    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
4395    #[allow(clippy::too_many_arguments)]
4396    pub fn append_kv_quantized_row_dc_inc(
4397        &self,
4398        k_row: &CudaSlice<f32>,
4399        v_row: &CudaSlice<f32>,
4400        kc: &mut CudaSlice<u8>,
4401        vc: &mut CudaSlice<u8>,
4402        t0_dev: &mut CudaSlice<i32>,
4403        kv_dim_k: usize,
4404        kv_dim_v: usize,
4405        k_tok_bytes: usize,
4406        v_tok_bytes: usize,
4407        g: bool,
4408    ) -> Result<(), Box<dyn std::error::Error>> {
4409        let f = if g {
4410            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
4411        } else {
4412            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
4413        };
4414        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
4415        let cfg = LaunchConfig {
4416            grid_dim: (1, 1, 1),
4417            block_dim: (nthreads, 1, 1),
4418            shared_mem_bytes: 0,
4419        };
4420        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4421        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4422        let __s_b = self.gpu.stream();
4423        let mut b = __s_b.launch_builder(&f);
4424        b.arg(k_row)
4425            .arg(v_row)
4426            .arg(kc)
4427            .arg(vc)
4428            .arg(t0_dev)
4429            .arg(&kdk)
4430            .arg(&kdv)
4431            .arg(&ktb)
4432            .arg(&vtb);
4433        unsafe {
4434            b.launch(cfg)?;
4435        }
4436        Ok(())
4437    }
4438
4439    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
4440    pub fn pack_tok_p(
4441        &self,
4442        tok: &CudaSlice<u32>,
4443        p: &CudaSlice<f32>,
4444        out: &mut CudaSlice<u32>,
4445        slot: usize,
4446    ) -> Result<(), Box<dyn std::error::Error>> {
4447        let f = self.func("pack_tok_p");
4448        let sl = slot as i32;
4449        let cfg = LaunchConfig {
4450            grid_dim: (1, 1, 1),
4451            block_dim: (32, 1, 1),
4452            shared_mem_bytes: 0,
4453        };
4454        let __s_b = self.gpu.stream();
4455        let mut b = __s_b.launch_builder(&f);
4456        b.arg(tok).arg(p).arg(out).arg(&sl);
4457        unsafe {
4458            b.launch(cfg)?;
4459        }
4460        Ok(())
4461    }
4462    pub fn tok_map_u32(
4463        &self,
4464        tok: &mut CudaSlice<u32>,
4465        map: &CudaSlice<u32>,
4466    ) -> Result<(), Box<dyn std::error::Error>> {
4467        let f = self.func("tok_map_u32");
4468        let cfg = LaunchConfig {
4469            grid_dim: (1, 1, 1),
4470            block_dim: (32, 1, 1),
4471            shared_mem_bytes: 0,
4472        };
4473        let __s_b = self.gpu.stream();
4474        let mut b = __s_b.launch_builder(&f);
4475        b.arg(tok).arg(map);
4476        unsafe {
4477            b.launch(cfg)?;
4478        }
4479        Ok(())
4480    }
4481
4482    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
4483    #[allow(clippy::too_many_arguments)]
4484    pub fn spec_assemble_verify(
4485        &self,
4486        tokp: &CudaSlice<u32>,
4487        pend: &CudaSlice<u32>,
4488        d2t: Option<&CudaSlice<u32>>,
4489        vtok: &mut CudaSlice<u32>,
4490        brk: &mut CudaSlice<u32>,
4491        p_min: f32,
4492        k: usize,
4493        pmin0: bool,
4494    ) -> Result<(), Box<dyn std::error::Error>> {
4495        let f = self.func("spec_assemble_verify");
4496        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
4497        let cfg = LaunchConfig {
4498            grid_dim: (1, 1, 1),
4499            block_dim: (32, 1, 1),
4500            shared_mem_bytes: 0,
4501        };
4502        let __s_b = self.gpu.stream();
4503        let mut b = __s_b.launch_builder(&f);
4504        match d2t {
4505            Some(m) => {
4506                b.arg(tokp)
4507                    .arg(pend)
4508                    .arg(m)
4509                    .arg(vtok)
4510                    .arg(brk)
4511                    .arg(&p_min)
4512                    .arg(&ki)
4513                    .arg(&pm);
4514                unsafe {
4515                    b.launch(cfg)?;
4516                }
4517            }
4518            None => {
4519                let null: u64 = 0;
4520                b.arg(tokp)
4521                    .arg(pend)
4522                    .arg(&null)
4523                    .arg(vtok)
4524                    .arg(brk)
4525                    .arg(&p_min)
4526                    .arg(&ki)
4527                    .arg(&pm);
4528                unsafe {
4529                    b.launch(cfg)?;
4530                }
4531            }
4532        }
4533        Ok(())
4534    }
4535
4536    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
4537    #[allow(clippy::too_many_arguments)]
4538    pub fn ssm_conv_ring_rebuild_dc(
4539        &self,
4540        qkv_tm: &CudaSlice<f32>,
4541        ring_old: &CudaSlice<f32>,
4542        conv_state: &mut CudaSlice<f32>,
4543        conv_dim: usize,
4544        acc: &CudaSlice<u32>,
4545        base: usize,
4546        t_v: usize,
4547        d_conv: usize,
4548    ) -> Result<(), Box<dyn std::error::Error>> {
4549        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
4550        let n = conv_dim * (d_conv - 1);
4551        let cfg = LaunchConfig::for_num_elems(n as u32);
4552        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
4553        let __s_b = self.gpu.stream();
4554        let mut b = __s_b.launch_builder(&f);
4555        b.arg(qkv_tm)
4556            .arg(ring_old)
4557            .arg(conv_state)
4558            .arg(&cd)
4559            .arg(acc)
4560            .arg(&b0)
4561            .arg(&tv)
4562            .arg(&dc);
4563        unsafe {
4564            b.launch(cfg)?;
4565        }
4566        Ok(())
4567    }
4568    #[allow(clippy::too_many_arguments)]
4569    pub fn gdn_scan_s128_dc(
4570        &self,
4571        q: &CudaSlice<f32>,
4572        k: &CudaSlice<f32>,
4573        v: &CudaSlice<f32>,
4574        g: &CudaSlice<f32>,
4575        beta: &CudaSlice<f32>,
4576        state_in: &CudaSlice<f32>,
4577        state_out: &mut CudaSlice<f32>,
4578        o: &mut CudaSlice<f32>,
4579        n_head: usize,
4580        acc: &CudaSlice<u32>,
4581        base: usize,
4582        t_v: usize,
4583        scale: f32,
4584    ) -> Result<(), Box<dyn std::error::Error>> {
4585        let f = self.func("gdn_scan_s128_dc");
4586        const S_V: u32 = 128;
4587        const WARP: u32 = 32;
4588        const COLS_PER_BLOCK: u32 = 4;
4589        let cfg = LaunchConfig {
4590            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
4591            block_dim: (WARP, COLS_PER_BLOCK, 1),
4592            shared_mem_bytes: 0,
4593        };
4594        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
4595        let __s_b = self.gpu.stream();
4596        let mut b = __s_b.launch_builder(&f);
4597        b.arg(q)
4598            .arg(k)
4599            .arg(v)
4600            .arg(g)
4601            .arg(beta)
4602            .arg(state_in)
4603            .arg(state_out)
4604            .arg(o)
4605            .arg(&h)
4606            .arg(acc)
4607            .arg(&b0)
4608            .arg(&tv)
4609            .arg(&scale);
4610        unsafe {
4611            b.launch(cfg)?;
4612        }
4613        Ok(())
4614    }
4615
4616    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
4617    pub fn spec_rollback_kv(
4618        &self,
4619        len_ptrs: &CudaSlice<u64>,
4620        saved: &CudaSlice<i32>,
4621        acc: &CudaSlice<u32>,
4622        base: usize,
4623        n_layer: usize,
4624    ) -> Result<(), Box<dyn std::error::Error>> {
4625        let f = self.func("spec_rollback_kv");
4626        let (b, nl) = (base as i32, n_layer as i32);
4627        let cfg = LaunchConfig {
4628            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
4629            block_dim: (64, 1, 1),
4630            shared_mem_bytes: 0,
4631        };
4632        let __s_bl = self.gpu.stream();
4633        let mut bl = __s_bl.launch_builder(&f);
4634        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
4635        unsafe {
4636            bl.launch(cfg)?;
4637        }
4638        Ok(())
4639    }
4640
4641    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
4642    pub fn spec_fork_valid(
4643        &self,
4644        acc: &CudaSlice<u32>,
4645        optimistic_pending: u32,
4646        valid: &mut CudaSlice<u32>,
4647    ) -> Result<(), Box<dyn std::error::Error>> {
4648        let f = self.func("spec_fork_valid");
4649        let cfg = LaunchConfig {
4650            grid_dim: (1, 1, 1),
4651            block_dim: (1, 1, 1),
4652            shared_mem_bytes: 0,
4653        };
4654        let __s_bl = self.gpu.stream();
4655        let mut bl = __s_bl.launch_builder(&f);
4656        bl.arg(acc).arg(&optimistic_pending).arg(valid);
4657        unsafe {
4658            bl.launch(cfg)?;
4659        }
4660        Ok(())
4661    }
4662
4663    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
4664    pub fn spec_fork_reconcile_kv(
4665        &self,
4666        len_ptrs: &CudaSlice<u64>,
4667        saved: &CudaSlice<i32>,
4668        acc: &CudaSlice<u32>,
4669        valid: &CudaSlice<u32>,
4670        base: usize,
4671        n_layer: usize,
4672    ) -> Result<(), Box<dyn std::error::Error>> {
4673        let f = self.func("spec_fork_reconcile_kv");
4674        let (b, nl) = (base as i32, n_layer as i32);
4675        let cfg = LaunchConfig {
4676            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
4677            block_dim: (64, 1, 1),
4678            shared_mem_bytes: 0,
4679        };
4680        let __s_bl = self.gpu.stream();
4681        let mut bl = __s_bl.launch_builder(&f);
4682        bl.arg(len_ptrs)
4683            .arg(saved)
4684            .arg(acc)
4685            .arg(valid)
4686            .arg(&b)
4687            .arg(&nl);
4688        unsafe {
4689            bl.launch(cfg)?;
4690        }
4691        Ok(())
4692    }
4693
4694    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
4695    pub fn spec_fork_restore_f32(
4696        &self,
4697        snapshot: &CudaSlice<f32>,
4698        state: &mut CudaSlice<f32>,
4699        valid: &CudaSlice<u32>,
4700    ) -> Result<(), Box<dyn std::error::Error>> {
4701        assert_eq!(
4702            snapshot.len(),
4703            state.len(),
4704            "fork recurrent snapshot shape mismatch"
4705        );
4706        let f = self.func("spec_fork_restore_f32");
4707        let n = state.len() as i32;
4708        #[allow(clippy::manual_clamp)]
4709        // allow: the min/max chain mirrors the reference arithmetic order in pinned sizing/quant math
4710        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
4711        let cfg = LaunchConfig {
4712            grid_dim: (blocks, 1, 1),
4713            block_dim: (256, 1, 1),
4714            shared_mem_bytes: 0,
4715        };
4716        let __s_bl = self.gpu.stream();
4717        let mut bl = __s_bl.launch_builder(&f);
4718        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
4719        unsafe {
4720            bl.launch(cfg)?;
4721        }
4722        Ok(())
4723    }
4724
4725    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
4726    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
4727    pub fn spec_seed_gather(
4728        &self,
4729        vx: &CudaSlice<f32>,
4730        fill_prev: &CudaSlice<f32>,
4731        acc: &CudaSlice<u32>,
4732        h_seed: &mut CudaSlice<f32>,
4733        base: usize,
4734        n_embd: usize,
4735    ) -> Result<(), Box<dyn std::error::Error>> {
4736        let f = self.func("spec_seed_gather");
4737        let (b, ne) = (base as i32, n_embd as i32);
4738        let cfg = LaunchConfig {
4739            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
4740            block_dim: (256, 1, 1),
4741            shared_mem_bytes: 0,
4742        };
4743        let __s_bl = self.gpu.stream();
4744        let mut bl = __s_bl.launch_builder(&f);
4745        bl.arg(vx)
4746            .arg(fill_prev)
4747            .arg(acc)
4748            .arg(h_seed)
4749            .arg(&b)
4750            .arg(&ne);
4751        unsafe {
4752            bl.launch(cfg)?;
4753        }
4754        Ok(())
4755    }
4756
4757    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
4758    pub fn spec_accept_greedy(
4759        &self,
4760        preds: &CudaSlice<u32>,
4761        draft: &CudaSlice<u32>,
4762        last_pred: u32,
4763        base: usize,
4764        k_round: usize,
4765        out: &mut CudaSlice<u32>,
4766    ) -> Result<(), Box<dyn std::error::Error>> {
4767        let f = self.func("spec_accept_greedy");
4768        let (b, k) = (base as i32, k_round as i32);
4769        let cfg = LaunchConfig {
4770            grid_dim: (1, 1, 1),
4771            block_dim: (32, 1, 1),
4772            shared_mem_bytes: 0,
4773        };
4774        let __s_bl = self.gpu.stream();
4775        let mut bl = __s_bl.launch_builder(&f);
4776        bl.arg(preds)
4777            .arg(draft)
4778            .arg(&last_pred)
4779            .arg(&b)
4780            .arg(&k)
4781            .arg(out);
4782        unsafe {
4783            bl.launch(cfg)?;
4784        }
4785        Ok(())
4786    }
4787
4788    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
4789    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
4790    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
4791
4792    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
4793    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
4794    pub fn gumbel_perturb(
4795        &self,
4796        x: &CudaSlice<f32>,
4797        y: &mut CudaSlice<f32>,
4798        n: usize,
4799        seed: u64,
4800        stream_pos: u32,
4801        temp: f32,
4802    ) -> Result<(), Box<dyn std::error::Error>> {
4803        let f = self.func("gumbel_perturb_f32");
4804        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4805        let cfg = LaunchConfig {
4806            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4807            block_dim: (256, 1, 1),
4808            shared_mem_bytes: 0,
4809        };
4810        let __s_b = self.gpu.stream();
4811        let mut b = __s_b.launch_builder(&f);
4812        b.arg(x)
4813            .arg(&mut *y)
4814            .arg(&ni)
4815            .arg(&slo)
4816            .arg(&shi)
4817            .arg(&stream_pos)
4818            .arg(&temp);
4819        unsafe {
4820            b.launch(cfg)?;
4821        }
4822        Ok(())
4823    }
4824
4825    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
4826    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
4827    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
4828    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
4829    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
4830    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
4831    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
4832    pub fn mask_logits_col(
4833        &self,
4834        logits: &mut CudaSlice<f32>,
4835        mask: &CudaSlice<u32>,
4836        col: usize,
4837        n: usize,
4838        mask_words: usize,
4839    ) -> Result<(), Box<dyn std::error::Error>> {
4840        let f = self.func("mask_logits_f32");
4841        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
4842        let cfg = LaunchConfig {
4843            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
4844            block_dim: (256, 1, 1),
4845            shared_mem_bytes: 0,
4846        };
4847        let __s_b = self.gpu.stream();
4848        let mut b = __s_b.launch_builder(&f);
4849        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
4850        unsafe {
4851            b.launch(cfg)?;
4852        }
4853        Ok(())
4854    }
4855
4856    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
4857    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
4858    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
4859    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
4860    /// (the lane index is the in-row position; `col` only moves the input pointer). That
4861    /// pointer-invariance IS the serving isolation contract for sampled rows.
4862    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4863    pub fn gumbel_perturb_col(
4864        &self,
4865        x: &CudaSlice<f32>,
4866        col: usize,
4867        y: &mut CudaSlice<f32>,
4868        n: usize,
4869        seed: u64,
4870        stream_pos: u32,
4871        temp: f32,
4872    ) -> Result<(), Box<dyn std::error::Error>> {
4873        let f = self.func("gumbel_perturb_f32");
4874        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4875        let col_view = x.slice(col * n..(col + 1) * n);
4876        let cfg = LaunchConfig {
4877            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4878            block_dim: (256, 1, 1),
4879            shared_mem_bytes: 0,
4880        };
4881        let __s_b = self.gpu.stream();
4882        let mut b = __s_b.launch_builder(&f);
4883        b.arg(&col_view)
4884            .arg(&mut *y)
4885            .arg(&ni)
4886            .arg(&slo)
4887            .arg(&shi)
4888            .arg(&stream_pos)
4889            .arg(&temp);
4890        unsafe {
4891            b.launch(cfg)?;
4892        }
4893        Ok(())
4894    }
4895
4896    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
4897    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
4898    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
4899    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
4900    /// the serving isolation contract for sampled rows).
4901    #[allow(clippy::too_many_arguments)]
4902    pub fn gumbel_perturb_filtered_col(
4903        &self,
4904        x: &CudaSlice<f32>,
4905        col: usize,
4906        y: &mut CudaSlice<f32>,
4907        n: usize,
4908        seed: u64,
4909        stream_pos: u32,
4910        temp: f32,
4911        stat_max: &CudaSlice<f32>,
4912        stat_th: &CudaSlice<f32>,
4913        stat_idx: usize,
4914    ) -> Result<(), Box<dyn std::error::Error>> {
4915        let f = self.func("gumbel_perturb_filtered_col_f32");
4916        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4917        let (ci, si) = (col as i32, stat_idx as i32);
4918        let cfg = LaunchConfig {
4919            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4920            block_dim: (256, 1, 1),
4921            shared_mem_bytes: 0,
4922        };
4923        let __s_b = self.gpu.stream();
4924        let mut b = __s_b.launch_builder(&f);
4925        b.arg(x)
4926            .arg(&ci)
4927            .arg(&mut *y)
4928            .arg(&ni)
4929            .arg(&slo)
4930            .arg(&shi)
4931            .arg(&stream_pos)
4932            .arg(&temp)
4933            .arg(stat_max)
4934            .arg(stat_th)
4935            .arg(&si);
4936        unsafe {
4937            b.launch(cfg)?;
4938        }
4939        Ok(())
4940    }
4941
4942    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
4943    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
4944    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
4945    /// reads it (counter is data, not state — graph-replay-safe).
4946    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
4947        let f = self.func("memra_sctr_inc");
4948        let cfg = LaunchConfig {
4949            grid_dim: (1, 1, 1),
4950            block_dim: (1, 1, 1),
4951            shared_mem_bytes: 0,
4952        };
4953        let __s_b = self.gpu.stream();
4954        let mut b = __s_b.launch_builder(&f);
4955        b.arg(&mut *ctr);
4956        unsafe {
4957            b.launch(cfg)?;
4958        }
4959        Ok(())
4960    }
4961
4962    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
4963    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
4964    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
4965    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
4966    pub fn gumbel_perturb_ctr(
4967        &self,
4968        x: &CudaSlice<f32>,
4969        y: &mut CudaSlice<f32>,
4970        n: usize,
4971        seed: u64,
4972        ctr: &CudaSlice<u32>,
4973        temp: f32,
4974    ) -> Result<(), Box<dyn std::error::Error>> {
4975        let f = self.func("gumbel_perturb_ctr_f32");
4976        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4977        let cfg = LaunchConfig {
4978            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4979            block_dim: (256, 1, 1),
4980            shared_mem_bytes: 0,
4981        };
4982        let __s_b = self.gpu.stream();
4983        let mut b = __s_b.launch_builder(&f);
4984        b.arg(x)
4985            .arg(&mut *y)
4986            .arg(&ni)
4987            .arg(&slo)
4988            .arg(&shi)
4989            .arg(ctr)
4990            .arg(&temp);
4991        unsafe {
4992            b.launch(cfg)?;
4993        }
4994        Ok(())
4995    }
4996
4997    /// Graph-capturable `gumbel_perturb_filtered` (lane/step37-draft-graph-serving): the
4998    /// sampling-event counter comes from DEVICE memory (`ctr[0]`) and the filter stats
4999    /// (row_max, th) from DEVICE slots — the `filter_stats` outputs of the same captured
5000    /// body. Identical math (same Philox call, same lane mapping, same e0 filter test) to
5001    /// `gumbel_perturb_filtered` at stream_pos == ctr[0], row_max == mx[0], th == th_d[0]:
5002    /// the eager and graph FILTERED sampled chains produce bit-identical perturbations for
5003    /// the same (seed, counter, stats). Launch geometry mirrors the host-scalar wrapper.
5004    #[allow(clippy::too_many_arguments)]
5005    pub fn gumbel_perturb_filtered_ctr(
5006        &self,
5007        x: &CudaSlice<f32>,
5008        y: &mut CudaSlice<f32>,
5009        n: usize,
5010        seed: u64,
5011        ctr: &CudaSlice<u32>,
5012        temp: f32,
5013        stat_max: &CudaSlice<f32>,
5014        stat_th: &CudaSlice<f32>,
5015    ) -> Result<(), Box<dyn std::error::Error>> {
5016        let f = self.func("gumbel_perturb_filtered_ctr_f32");
5017        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5018        let cfg = LaunchConfig {
5019            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5020            block_dim: (256, 1, 1),
5021            shared_mem_bytes: 0,
5022        };
5023        let __s_b = self.gpu.stream();
5024        let mut b = __s_b.launch_builder(&f);
5025        b.arg(x)
5026            .arg(&mut *y)
5027            .arg(&ni)
5028            .arg(&slo)
5029            .arg(&shi)
5030            .arg(ctr)
5031            .arg(&temp)
5032            .arg(stat_max)
5033            .arg(stat_th);
5034        unsafe {
5035            b.launch(cfg)?;
5036        }
5037        Ok(())
5038    }
5039
5040    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
5041    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
5042    /// (smallest-index tie-break — matches the argmax-gate contract).
5043    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5044    pub fn softmax_gather(
5045        &self,
5046        x: &CudaSlice<f32>,
5047        row_stride: usize,
5048        ids: &CudaSlice<u32>,
5049        rows: &CudaSlice<i32>,
5050        out: &mut CudaSlice<f32>,
5051        n: usize,
5052        npair: usize,
5053        temp: f32,
5054    ) -> Result<(), Box<dyn std::error::Error>> {
5055        let f = self.func("softmax_gather_f32");
5056        let (ni, rs) = (n as i32, row_stride as i64);
5057        let np = npair as i32;
5058        let cfg = LaunchConfig {
5059            grid_dim: (npair as u32, 1, 1),
5060            block_dim: (256, 1, 1),
5061            shared_mem_bytes: 0,
5062        };
5063        let __s_b = self.gpu.stream();
5064        let mut b = __s_b.launch_builder(&f);
5065        b.arg(x)
5066            .arg(&rs)
5067            .arg(ids)
5068            .arg(rows)
5069            .arg(&mut *out)
5070            .arg(&ni)
5071            .arg(&np)
5072            .arg(&temp);
5073        unsafe {
5074            b.launch(cfg)?;
5075        }
5076        Ok(())
5077    }
5078
5079    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
5080    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
5081    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
5082    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5083    pub fn residual_sample(
5084        &self,
5085        p: &CudaSlice<f32>,
5086        q: Option<&CudaSlice<f32>>,
5087        n: usize,
5088        temp: f32,
5089        seed: u64,
5090        stream_pos: u32,
5091        out_tok: &mut CudaSlice<u32>,
5092    ) -> Result<(), Box<dyn std::error::Error>> {
5093        let f = self.func("residual_sample_f32");
5094        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5095        let nth = 1024u32;
5096        let cfg = LaunchConfig {
5097            grid_dim: (1, 1, 1),
5098            block_dim: (nth, 1, 1),
5099            shared_mem_bytes: 0,
5100        };
5101        let has_q: i32 = q.is_some() as i32;
5102        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
5103        let __s_b = self.gpu.stream();
5104        let mut b = __s_b.launch_builder(&f);
5105        b.arg(p)
5106            .arg(qbuf)
5107            .arg(&has_q)
5108            .arg(&ni)
5109            .arg(&temp)
5110            .arg(&slo)
5111            .arg(&shi)
5112            .arg(&stream_pos)
5113            .arg(&mut *out_tok);
5114        unsafe {
5115            b.launch(cfg)?;
5116        }
5117        Ok(())
5118    }
5119
5120    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
5121    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
5122    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
5123    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
5124    pub fn with_moe_cache<R>(
5125        &self,
5126        max_block_bytes: usize,
5127        f: impl FnOnce(
5128            &mut crate::moe_cache::MoeSlotCache,
5129            &Engine,
5130        ) -> Result<R, Box<dyn std::error::Error>>,
5131    ) -> Result<R, Box<dyn std::error::Error>> {
5132        let mut guard = self.moe_cache.lock().unwrap();
5133        if guard.is_none() {
5134            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
5135        }
5136        let cache = guard.as_mut().unwrap();
5137        f(cache, self)
5138    }
5139
5140    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
5141    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
5142    pub fn freeze_moe_cache(&self) {
5143        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
5144            cache.freeze();
5145        }
5146    }
5147
5148    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
5149    /// Never constructs a cache.
5150    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
5151        self.moe_cache
5152            .lock()
5153            .unwrap()
5154            .as_ref()
5155            .map(crate::moe_cache::MoeSlotCache::export_residency)
5156    }
5157
5158    pub(crate) fn moe_cache_frozen(&self) -> bool {
5159        self.moe_cache
5160            .lock()
5161            .unwrap()
5162            .as_ref()
5163            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
5164    }
5165
5166    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
5167    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
5168    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
5169    /// while leaving the profiling warmup's established batched behavior untouched.
5170    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
5171    /// tokenwise arm anyway.)
5172    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
5173        crate::cpu_experts::configured()
5174            && self.moe_cache_frozen()
5175            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
5176    }
5177
5178    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
5179    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
5180        assert!(
5181            self.moe_cache.lock().unwrap().is_none(),
5182            "MoE cache layout configured after cache construction"
5183        );
5184        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
5185    }
5186
5187    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
5188        self.moe_cache_layout.lock().unwrap().clone()
5189    }
5190
5191    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
5192    pub fn moe_cache_enabled() -> bool {
5193        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
5194    }
5195
5196    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
5197    /// Returns None if the cache was never built (disabled or no MoE forward ran).
5198    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
5199        let guard = self.moe_cache.lock().unwrap();
5200        guard
5201            .as_ref()
5202            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
5203    }
5204
5205    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
5206    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
5207    /// callers compare a before/after snapshot around a decode window.
5208    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5209    pub fn cpu_expert_stats(
5210        &self,
5211    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
5212        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
5213    }
5214
5215    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
5216    /// the backend tail that resident-GPU expert work did not hide.
5217    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
5218        crate::cpu_experts::predictor_stats()
5219    }
5220
5221    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
5222        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
5223    }
5224
5225    /// CPU-routed expert selections grouped by how many of their three projections were already
5226    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
5227    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
5228        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
5229    }
5230
5231    /// Positioned-read proof-backend counters:
5232    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
5233    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
5234        let guard = self.moe_cache.lock().unwrap();
5235        guard
5236            .as_ref()
5237            .and_then(|cache| cache.pread_stats())
5238            .map(|stats| {
5239                (
5240                    stats.reads,
5241                    stats.bytes,
5242                    stats.read_errors,
5243                    stats.short_reads,
5244                    stats.fallbacks,
5245                    stats.buffer_waits,
5246                    stats.ring_full,
5247                )
5248            })
5249    }
5250
5251    /// Spill configuration values that warned and substituted their documented defaults.
5252    pub fn spill_config_fallbacks(&self) -> u64 {
5253        crate::spill_pread::config_fallbacks()
5254    }
5255
5256    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
5257    pub fn moe_cache_reset_counters(&self) {
5258        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
5259            c.reset_counters();
5260        }
5261    }
5262
5263    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5264        Ok(self.gpu.stream().clone_htod(v)?)
5265    }
5266
5267    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
5268    /// past the final q4_0 block through their aligned window — the bytes never reach a
5269    /// result (funnelshift discards them) but must be mapped memory.
5270    pub fn htod_bytes_padded(
5271        &self,
5272        v: &[u8],
5273        pad: usize,
5274    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5275        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
5276        {
5277            let mut view = d.slice_mut(0..v.len());
5278            self.gpu.stream().memcpy_htod(v, &mut view)?;
5279        }
5280        Ok(d)
5281    }
5282
5283    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
5284    pub fn copy_into(
5285        &self,
5286        dst: &mut CudaSlice<f32>,
5287        off: usize,
5288        src: &CudaSlice<f32>,
5289        len: usize,
5290    ) -> Result<(), Box<dyn std::error::Error>> {
5291        let mut view = dst.slice_mut(off..off + len);
5292        self.gpu
5293            .stream()
5294            .memcpy_dtod(&src.slice(0..len), &mut view)?;
5295        Ok(())
5296    }
5297
5298    /// D2D copy with an offset on BOTH sides. `copy_into` always reads the source from 0,
5299    /// which cannot express "copy the TAIL of this buffer" — the shape a sliding-window draft
5300    /// KV export needs (lane/dspark-draft-plane-20260827).
5301    pub fn copy_range_into(
5302        &self,
5303        dst: &mut CudaSlice<f32>,
5304        dst_off: usize,
5305        src: &CudaSlice<f32>,
5306        src_off: usize,
5307        len: usize,
5308    ) -> Result<(), Box<dyn std::error::Error>> {
5309        let mut view = dst.slice_mut(dst_off..dst_off + len);
5310        self.gpu
5311            .stream()
5312            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
5313        Ok(())
5314    }
5315
5316    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
5317    /// u8 twin of copy_into (D2D byte-range copy at an offset).
5318    pub fn copy_u8_into(
5319        &self,
5320        dst: &mut CudaSlice<u8>,
5321        off: usize,
5322        src: &CudaSlice<u8>,
5323        len: usize,
5324    ) -> Result<(), Box<dyn std::error::Error>> {
5325        // try_slice_mut, not slice_mut: an out-of-bounds range here panics the GPU worker
5326        // thread and takes the whole server with it (2026-08-29 warm-turn-at-40k incident).
5327        // A bounds miss is a caller bug, but it must fail the request, not the fleet.
5328        let cap = dst.len();
5329        let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
5330            format!(
5331                "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
5332                off + len,
5333            )
5334        })?;
5335        self.gpu
5336            .stream()
5337            .memcpy_dtod(&src.slice(0..len), &mut view)?;
5338        Ok(())
5339    }
5340
5341    /// D2D byte-range copy with explicit source and destination offsets.
5342    pub fn copy_u8_range_into(
5343        &self,
5344        dst: &mut CudaSlice<u8>,
5345        dst_off: usize,
5346        src: &CudaSlice<u8>,
5347        src_off: usize,
5348        len: usize,
5349    ) -> Result<(), Box<dyn std::error::Error>> {
5350        // try_slice_mut for the same reason as copy_u8_into: bounds misses fail the request,
5351        // never panic the worker.
5352        let cap = dst.len();
5353        let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
5354            format!(
5355                "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
5356                dst_off + len,
5357            )
5358        })?;
5359        self.gpu
5360            .stream()
5361            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
5362        Ok(())
5363    }
5364
5365    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
5366    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
5367    /// keeping the audited attention range contiguous without changing its absolute start.
5368    /// #[track_caller]: every ring-backed append that REBASES sets the plane's `base`, and a
5369    /// later append or rewind that needs a lower row is then refused. Three attempts at the
5370    /// SWA-ring lap failed because the writer that actually moved `base` was never the site being
5371    /// patched — the bare "SWA ring lapped required rows" message named neither the caller nor
5372    /// what it retained. Cost of the annotation is nothing; cost of not having it was two wrong
5373    /// fixes on hardware.
5374    #[track_caller]
5375    pub fn prepare_kv_append(
5376        &self,
5377        kv: &mut crate::cache::KvLayer,
5378        retain_from: usize,
5379        append_rows: usize,
5380    ) -> Result<usize, Box<dyn std::error::Error>> {
5381        let caller = std::panic::Location::caller();
5382        let base_before = kv.ring.as_ref().map(|r| r.base());
5383        let Some(plan) = kv
5384            .ring
5385            .as_ref()
5386            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
5387            .transpose()
5388            .map_err(|err| -> Box<dyn std::error::Error> {
5389                format!(
5390                    "{err} [append len={} retain_from={retain_from} append_rows={append_rows}                      base={base_before:?} called from {caller}]",
5391                    kv.len
5392                )
5393                .into()
5394            })?
5395        else {
5396            return Ok(kv.len);
5397        };
5398        match plan {
5399            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
5400            crate::cache::KvRingAppend::Rebase {
5401                src_row,
5402                keep_rows,
5403                new_base,
5404                write_row,
5405            } => {
5406                if keep_rows > 0 {
5407                    let k_len = keep_rows * kv.k_tok_bytes;
5408                    let v_len = keep_rows * kv.v_tok_bytes;
5409                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
5410                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
5411                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
5412                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
5413                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
5414                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
5415                }
5416                // One line per distinct (caller, new_base) so the writers that move `base` are
5417                // enumerable from a single run instead of inferred from which error fires.
5418                if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
5419                    eprintln!(
5420                        "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
5421                         retain_from={retain_from} called from {caller}",
5422                        kv.len
5423                    );
5424                }
5425                kv.ring.as_mut().unwrap().apply_rebase(new_base);
5426                // The dcw draft arm's device mirror of the ring base (see KvLayer::base_d).
5427                // Rebase is the ONLY writer of `base`, and rebases run host-side outside any
5428                // captured region, so this one line keeps the device view exact.
5429                if let Some(base_d) = kv.base_d.as_mut() {
5430                    self.set_i32_one(base_d, new_base as i32)?;
5431                }
5432                Ok(write_row)
5433            }
5434        }
5435    }
5436
5437    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
5438    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
5439    pub fn htod_u8_into(
5440        &self,
5441        dst: &mut CudaSlice<u8>,
5442        off: usize,
5443        src: &[u8],
5444    ) -> Result<(), Box<dyn std::error::Error>> {
5445        let mut view = dst.slice_mut(off..off + src.len());
5446        self.gpu.stream().memcpy_htod(src, &mut view)?;
5447        Ok(())
5448    }
5449
5450    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
5451        b.slice(0..len)
5452    }
5453
5454    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
5455    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
5456    pub fn view_u8_range<'a>(
5457        &self,
5458        b: &'a CudaSlice<u8>,
5459        start: usize,
5460        end: usize,
5461    ) -> cudarc::driver::CudaView<'a, u8> {
5462        b.slice(start..end)
5463    }
5464    pub fn view_u8<'a>(
5465        &self,
5466        b: &'a CudaSlice<u8>,
5467        len: usize,
5468    ) -> cudarc::driver::CudaView<'a, u8> {
5469        b.slice(0..len)
5470    }
5471
5472    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
5473    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
5474    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
5475    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5476    pub fn append_kv_quantized(
5477        &self,
5478        k_row: &CudaSlice<f32>,
5479        v_row: &CudaSlice<f32>,
5480        kc: &mut CudaSlice<u8>,
5481        vc: &mut CudaSlice<u8>,
5482        t: usize,
5483        kv_dim_k: usize,
5484        kv_dim_v: usize,
5485        k_tok_bytes: usize,
5486        v_tok_bytes: usize,
5487        g: bool,
5488    ) -> Result<(), Box<dyn std::error::Error>> {
5489        let f = if g {
5490            self.func_g("append_quantize_kv_q8_0_q5_1")
5491        } else {
5492            self.func("append_quantize_kv_q8_0_q5_1")
5493        };
5494        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5495        let cfg = LaunchConfig {
5496            grid_dim: (nblk, 1, 1),
5497            block_dim: (32, 1, 1),
5498            shared_mem_bytes: 0,
5499        };
5500        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
5501        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5502        let __s_b = self.gpu.stream();
5503        let mut b = __s_b.launch_builder(&f);
5504        b.arg(k_row)
5505            .arg(v_row)
5506            .arg(kc)
5507            .arg(vc)
5508            .arg(&ti)
5509            .arg(&kdk)
5510            .arg(&kdv)
5511            .arg(&ktb)
5512            .arg(&vtb);
5513        unsafe {
5514            b.launch(cfg)?;
5515        }
5516        Ok(())
5517    }
5518
5519    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
5520    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
5521    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
5522    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5523    pub fn append_kv_quantized_dc(
5524        &self,
5525        k_row: &CudaSlice<f32>,
5526        v_row: &CudaSlice<f32>,
5527        kc: &mut CudaSlice<u8>,
5528        vc: &mut CudaSlice<u8>,
5529        t_dev: &CudaSlice<i32>,
5530        kv_dim_k: usize,
5531        kv_dim_v: usize,
5532        k_tok_bytes: usize,
5533        v_tok_bytes: usize,
5534        g: bool,
5535    ) -> Result<(), Box<dyn std::error::Error>> {
5536        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5537        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5538        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5539        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
5540        if Self::pdl_on() && Self::pdl_wb_on() {
5541            use cudarc::driver::{DevicePtr, DevicePtrMut};
5542            let s = &self.gpu.stream();
5543            let (pk, _g0) = k_row.device_ptr(s);
5544            let (pv, _g1) = v_row.device_ptr(s);
5545            let (pkc, _g2) = kc.device_ptr_mut(s);
5546            let (pvc, _g3) = vc.device_ptr_mut(s);
5547            let (pt, _g4) = t_dev.device_ptr(s);
5548            let mut ps = [
5549                &pk as *const _ as *mut std::ffi::c_void,
5550                &pv as *const _ as *mut _,
5551                &pkc as *const _ as *mut _,
5552                &pvc as *const _ as *mut _,
5553                &pt as *const _ as *mut _,
5554                &kdk as *const _ as *mut _,
5555                &kdv as *const _ as *mut _,
5556                &ktb as *const _ as *mut _,
5557                &vtb as *const _ as *mut _,
5558            ];
5559            unsafe {
5560                self.launch_pdl_flash(
5561                    g,
5562                    "append_quantize_kv_q8_0_q5_1_dc",
5563                    (nblk, 1, 1),
5564                    (32, 1, 1),
5565                    0,
5566                    &mut ps,
5567                )?;
5568            }
5569            return Ok(());
5570        }
5571        let f = if g {
5572            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
5573        } else {
5574            self.func("append_quantize_kv_q8_0_q5_1_dc")
5575        };
5576        let cfg = LaunchConfig {
5577            grid_dim: (nblk, 1, 1),
5578            block_dim: (32, 1, 1),
5579            shared_mem_bytes: 0,
5580        };
5581        let __s_b = self.gpu.stream();
5582        let mut b = __s_b.launch_builder(&f);
5583        b.arg(k_row)
5584            .arg(v_row)
5585            .arg(kc)
5586            .arg(vc)
5587            .arg(t_dev)
5588            .arg(&kdk)
5589            .arg(&kdv)
5590            .arg(&ktb)
5591            .arg(&vtb);
5592        unsafe {
5593            b.launch(cfg)?;
5594        }
5595        Ok(())
5596    }
5597
5598    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
5599    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
5600    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
5601    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
5602    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
5603    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
5604    #[allow(clippy::too_many_arguments)]
5605    pub fn append_kv_quantized_rows(
5606        &self,
5607        k_rows: &CudaSlice<f32>,
5608        v_rows: &CudaSlice<f32>,
5609        kc: &mut CudaSlice<u8>,
5610        vc: &mut CudaSlice<u8>,
5611        t0: usize,
5612        t: usize,
5613        kv_dim_k: usize,
5614        kv_dim_v: usize,
5615        k_tok_bytes: usize,
5616        v_tok_bytes: usize,
5617        g: bool,
5618    ) -> Result<(), Box<dyn std::error::Error>> {
5619        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
5620            for i in 0..t {
5621                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
5622                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
5623                self.append_kv_quantized_view(
5624                    &k_row,
5625                    &v_row,
5626                    kc,
5627                    vc,
5628                    t0 + i,
5629                    kv_dim_k,
5630                    kv_dim_v,
5631                    k_tok_bytes,
5632                    v_tok_bytes,
5633                    g,
5634                )?;
5635            }
5636            return Ok(());
5637        }
5638        let f = if g {
5639            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
5640        } else {
5641            self.func("append_quantize_kv_q8_0_q5_1_rows")
5642        };
5643        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5644        let cfg = LaunchConfig {
5645            grid_dim: (nblk, t as u32, 1),
5646            block_dim: (32, 1, 1),
5647            shared_mem_bytes: 0,
5648        };
5649        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
5650        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5651        let __s_b = self.gpu.stream();
5652        let mut b = __s_b.launch_builder(&f);
5653        b.arg(k_rows)
5654            .arg(v_rows)
5655            .arg(kc)
5656            .arg(vc)
5657            .arg(&t0i)
5658            .arg(&kdk)
5659            .arg(&kdv)
5660            .arg(&ktb)
5661            .arg(&vtb);
5662        unsafe {
5663            b.launch(cfg)?;
5664        }
5665        Ok(())
5666    }
5667
5668    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
5669    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
5670    /// later, inside a captured graph) without a host round-trip.
5671    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
5672        let f = self.func("inc_i32");
5673        let cfg = LaunchConfig {
5674            grid_dim: (1, 1, 1),
5675            block_dim: (1, 1, 1),
5676            shared_mem_bytes: 0,
5677        };
5678        let __s_b = self.gpu.stream();
5679        let mut b = __s_b.launch_builder(&f);
5680        b.arg(p);
5681        unsafe {
5682            b.launch(cfg)?;
5683        }
5684        Ok(())
5685    }
5686
5687    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
5688    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
5689    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5690    pub fn append_kv_quantized_view(
5691        &self,
5692        k_row: &cudarc::driver::CudaView<f32>,
5693        v_row: &cudarc::driver::CudaView<f32>,
5694        kc: &mut CudaSlice<u8>,
5695        vc: &mut CudaSlice<u8>,
5696        t: usize,
5697        kv_dim_k: usize,
5698        kv_dim_v: usize,
5699        k_tok_bytes: usize,
5700        v_tok_bytes: usize,
5701        g: bool,
5702    ) -> Result<(), Box<dyn std::error::Error>> {
5703        let stream = self.gpu.stream();
5704        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
5705        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
5706        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
5707        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
5708        let f = if g {
5709            self.func_g("append_quantize_kv_q8_0_q5_1")
5710        } else {
5711            self.func("append_quantize_kv_q8_0_q5_1")
5712        };
5713        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5714        let cfg = LaunchConfig {
5715            grid_dim: (nblk, 1, 1),
5716            block_dim: (32, 1, 1),
5717            shared_mem_bytes: 0,
5718        };
5719        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
5720        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5721        let mut b = stream.launch_builder(&f);
5722        b.arg(k_row)
5723            .arg(v_row)
5724            .arg(kc)
5725            .arg(vc)
5726            .arg(&ti)
5727            .arg(&kdk)
5728            .arg(&kdv)
5729            .arg(&ktb)
5730            .arg(&vtb);
5731        unsafe {
5732            b.launch(cfg)?;
5733        }
5734        Ok(())
5735    }
5736
5737    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
5738    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
5739    pub fn copy_view_into(
5740        &self,
5741        dst: &mut CudaSlice<f32>,
5742        off: usize,
5743        src: &cudarc::driver::CudaView<f32>,
5744        len: usize,
5745    ) -> Result<(), Box<dyn std::error::Error>> {
5746        let mut view = dst.slice_mut(off..off + len);
5747        self.gpu
5748            .stream()
5749            .memcpy_dtod(&src.slice(0..len), &mut view)?;
5750        Ok(())
5751    }
5752
5753    /// Real device-to-device COPY of `src` into a freshly allocated buffer. Used for cache
5754    /// snapshots (MTP-PLAN §D.4), where a snapshot must not alias the live buffer.
5755    ///
5756    /// CORRECTION (memra-next#23, verified against the LOCKED cudarc 0.19.8): this comment used to say
5757    /// "`CudaSlice::clone()` only bumps a refcount and would alias the live buffer". That is
5758    /// FALSE and it propagated — `impl Clone for CudaSlice` is `try_clone().unwrap()`, and
5759    /// `try_clone` is `self.stream.clone_dtod(self)`, so a plain `.clone()` already allocates and
5760    /// copies. Code that wants real aliasing needs an `Arc<CudaSlice<T>>` (see
5761    /// `vision::EmbedOverlay::rows`).
5762    ///
5763    /// THE TWO ARE NOT INTERCHANGEABLE, AND THE DIFFERENCE IS NOT ONLY FALLIBILITY — a second
5764    /// correction, from the peer review of that first one, because getting this backwards is how
5765    /// a residency bug gets written. `CudaSlice::clone()` allocates on the SLICE's own stream, so
5766    /// the copy lands in the SOURCE's context. This method allocates on `self.gpu.stream()`,
5767    /// which is the thread-local pp stage stream whenever a stage scope is active — so under a
5768    /// stage scope THIS method is the one that lands in a foreign context. Choose by what you
5769    /// need: `try_clone()` for a fallible copy that stays with the source, this method for a copy
5770    /// deliberately placed on the calling engine's current stream (and check the landing context
5771    /// if residency matters). Minor: cudarc's path uses an uninitialized alloc, this one
5772    /// `alloc_zeros`, i.e. an extra full memset.
5773    pub fn clone_dtod(
5774        &self,
5775        src: &CudaSlice<f32>,
5776    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5777        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
5778        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
5779        Ok(dst)
5780    }
5781
5782    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
5783    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
5784    pub fn dtod_copy_view(
5785        &self,
5786        src: &cudarc::driver::CudaView<f32>,
5787        dst: &mut CudaSlice<f32>,
5788    ) -> Result<(), Box<dyn std::error::Error>> {
5789        self.gpu.stream().memcpy_dtod(src, dst)?;
5790        Ok(())
5791    }
5792
5793    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
5794    pub fn dtod_copy_view_i8(
5795        &self,
5796        src: &cudarc::driver::CudaView<i8>,
5797        dst: &mut CudaSlice<i8>,
5798    ) -> Result<(), Box<dyn std::error::Error>> {
5799        self.gpu.stream().memcpy_dtod(src, dst)?;
5800        Ok(())
5801    }
5802
5803    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
5804    pub fn dtod_copy_into(
5805        &self,
5806        src: &CudaSlice<f32>,
5807        dst: &mut CudaSlice<f32>,
5808        offset: usize,
5809    ) -> Result<(), Box<dyn std::error::Error>> {
5810        let n = src.len();
5811        let mut dv = dst.slice_mut(offset..offset + n);
5812        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
5813        Ok(())
5814    }
5815
5816    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
5817    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
5818    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
5819    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
5820    /// Bytes and stream order are identical to the memcpy sequence it replaces.
5821    pub fn copy_batch_uniform_f32(
5822        &self,
5823        table: &CudaSlice<u64>,
5824        n: usize,
5825        words: usize,
5826    ) -> Result<(), Box<dyn std::error::Error>> {
5827        if n == 0 || words == 0 {
5828            return Ok(());
5829        }
5830        debug_assert!(
5831            table.len() >= 2 * n,
5832            "pointer table must hold n srcs + n dsts"
5833        );
5834        let f = self.func("copy_batch_uniform_f32");
5835        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
5836        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
5837        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
5838        let (ni, wi) = (n as i32, words as i32);
5839        let cfg = LaunchConfig {
5840            grid_dim: (chunks, n as u32, 1),
5841            block_dim: (256, 1, 1),
5842            shared_mem_bytes: 0,
5843        };
5844        let __s = self.gpu.stream();
5845        let mut b = __s.launch_builder(&f);
5846        b.arg(table).arg(&ni).arg(&wi);
5847        unsafe {
5848            b.launch(cfg)?;
5849        }
5850        Ok(())
5851    }
5852
5853    /// Copy one uniform quantized K/V row range for each layer in `table` and publish every
5854    /// layer's device length in the same launch. Table layout is five pointer planes:
5855    /// K source, V source, K destination, V destination, and i32 length destination.
5856    #[allow(clippy::too_many_arguments)] // allow: row bytes and source strides are independent K/V geometry and collapsing them would hide the peer-layout contract
5857    pub fn copy_batch_uniform_kv_u8_set_len(
5858        &self,
5859        table: &CudaSlice<u64>,
5860        n: usize,
5861        rows: usize,
5862        k_row_bytes: usize,
5863        v_row_bytes: usize,
5864        k_src_stride: usize,
5865        v_src_stride: usize,
5866        logical_len: usize,
5867    ) -> Result<(), Box<dyn std::error::Error>> {
5868        if n == 0 || rows == 0 || (k_row_bytes == 0 && v_row_bytes == 0) {
5869            return Ok(());
5870        }
5871        if table.len() < 5 * n {
5872            return Err(format!(
5873                "TP KV repair table has {} words, expected at least {}",
5874                table.len(),
5875                5 * n
5876            )
5877            .into());
5878        }
5879        let ni = i32::try_from(n).map_err(|_| "TP KV repair layer count exceeds i32")?;
5880        let rows = i32::try_from(rows).map_err(|_| "TP KV repair rows exceed i32")?;
5881        let kb = i32::try_from(k_row_bytes).map_err(|_| "TP KV repair K bytes exceed i32")?;
5882        let vb = i32::try_from(v_row_bytes).map_err(|_| "TP KV repair V bytes exceed i32")?;
5883        let ks = i32::try_from(k_src_stride).map_err(|_| "TP KV repair K stride exceeds i32")?;
5884        let vs = i32::try_from(v_src_stride).map_err(|_| "TP KV repair V stride exceeds i32")?;
5885        let len = i32::try_from(logical_len).map_err(|_| "TP KV repair length exceeds i32")?;
5886        let f = self.func("copy_batch_uniform_kv_u8_set_len");
5887        let cfg = LaunchConfig {
5888            grid_dim: (n as u32, 1, 1),
5889            block_dim: (256, 1, 1),
5890            shared_mem_bytes: 0,
5891        };
5892        let stream = self.gpu.stream();
5893        let mut builder = stream.launch_builder(&f);
5894        builder
5895            .arg(table)
5896            .arg(&ni)
5897            .arg(&rows)
5898            .arg(&kb)
5899            .arg(&vb)
5900            .arg(&ks)
5901            .arg(&vs)
5902            .arg(&len);
5903        unsafe {
5904            builder.launch(cfg)?;
5905        }
5906        Ok(())
5907    }
5908
5909    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
5910    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
5911    pub fn htod_u64_into(
5912        &self,
5913        v: &[u64],
5914        dst: &mut CudaSlice<u64>,
5915    ) -> Result<(), Box<dyn std::error::Error>> {
5916        let mut view = dst.slice_mut(0..v.len());
5917        self.gpu.stream().memcpy_htod(v, &mut view)?;
5918        Ok(())
5919    }
5920
5921    /// f32 twin of [`Self::htod_u64_into`] (the MoE vrows scale tables through the
5922    /// verify-walk workspace, door W).
5923    pub fn htod_f32_into(
5924        &self,
5925        v: &[f32],
5926        dst: &mut CudaSlice<f32>,
5927    ) -> Result<(), Box<dyn std::error::Error>> {
5928        let mut view = dst.slice_mut(0..v.len());
5929        self.gpu.stream().memcpy_htod(v, &mut view)?;
5930        Ok(())
5931    }
5932
5933    /// `htod_f32_into` landing at an element offset: `dst[off..off+v.len()] = v`. The EP
5934    /// dispatch-diet's bulk peer-row return lands the peer's compact block directly into the
5935    /// pair-slab tail with ONE upload instead of a per-row scatter.
5936    pub fn htod_f32_into_at(
5937        &self,
5938        v: &[f32],
5939        dst: &mut CudaSlice<f32>,
5940        off: usize,
5941    ) -> Result<(), Box<dyn std::error::Error>> {
5942        if off + v.len() > dst.len() {
5943            return Err(format!(
5944                "htod_f32_into_at range {}..{} exceeds dst {}",
5945                off,
5946                off + v.len(),
5947                dst.len()
5948            )
5949            .into());
5950        }
5951        let mut view = dst.slice_mut(off..off + v.len());
5952        self.gpu.stream().memcpy_htod(v, &mut view)?;
5953        Ok(())
5954    }
5955
5956    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
5957    /// device pointer-table entry at run time, so a captured graph follows the gdn
5958    /// ping-pong through the same table its scan kernels read — a baked memcpy node
5959    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
5960    pub fn copy_indirect_src_f32(
5961        &self,
5962        src_entry: &cudarc::driver::CudaView<u64>,
5963        dst: &mut CudaSlice<f32>,
5964        dst_off: usize,
5965        words: usize,
5966    ) -> Result<(), Box<dyn std::error::Error>> {
5967        let f = self.func("copy_indirect_src_f32");
5968        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
5969        let wi = words as i32;
5970        let cfg = LaunchConfig {
5971            grid_dim: (chunks, 1, 1),
5972            block_dim: (256, 1, 1),
5973            shared_mem_bytes: 0,
5974        };
5975        let mut dv = dst.slice_mut(dst_off..dst_off + words);
5976        let __s = self.gpu.stream();
5977        let mut b = __s.launch_builder(&f);
5978        b.arg(src_entry).arg(&mut dv).arg(&wi);
5979        unsafe {
5980            b.launch(cfg)?;
5981        }
5982        Ok(())
5983    }
5984
5985    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
5986    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
5987        self.alloc_uninit::<i8>(n)
5988    }
5989
5990    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
5991    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5992    pub fn qmatvec(
5993        &self,
5994        w: &CudaSlice<u8>,
5995        x: &CudaSlice<f32>,
5996        m: usize,
5997        in_f: usize,
5998        out_f: usize,
5999        qtype: i32,
6000        row_bytes: usize,
6001    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6002        let f = self.func("qmatvec_f32");
6003        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6004        let cfg = LaunchConfig {
6005            grid_dim: (out_f as u32, m as u32, 1),
6006            block_dim: (256, 1, 1),
6007            shared_mem_bytes: 0,
6008        };
6009        let (inf, outf, mi, qt, rb) =
6010            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6011        let __s_b = self.gpu.stream();
6012        let mut b = __s_b.launch_builder(&f);
6013        b.arg(w)
6014            .arg(x)
6015            .arg(&mut y)
6016            .arg(&inf)
6017            .arg(&outf)
6018            .arg(&mi)
6019            .arg(&qt)
6020            .arg(&rb);
6021        unsafe {
6022            b.launch(cfg)?;
6023        }
6024        Ok(y)
6025    }
6026
6027    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
6028    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6029        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
6030        self.keep_if_capturing(&s);
6031        Ok(s)
6032    }
6033
6034    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
6035    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
6036    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
6037    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6038        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
6039        self.keep_if_capturing(&s);
6040        Ok(s)
6041    }
6042
6043    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
6044    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
6045    pub fn memset_zeros_view(
6046        &self,
6047        dst: &mut cudarc::driver::CudaViewMut<f32>,
6048    ) -> Result<(), Box<dyn std::error::Error>> {
6049        self.gpu.stream().memset_zeros(dst)?;
6050        Ok(())
6051    }
6052
6053    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
6054    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
6055    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
6056    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
6057    /// stream would require an event).
6058    pub fn stage_expert(
6059        &self,
6060        host_bytes: &[u8],
6061        scratch: &mut CudaSlice<u8>,
6062        off: usize,
6063    ) -> Result<(), Box<dyn std::error::Error>> {
6064        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
6065        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
6066        Ok(())
6067    }
6068
6069    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
6070    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
6071    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
6072    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
6073    /// One CTA per token row, 256 threads (one per expert).
6074    pub fn moe_router_topk(
6075        &self,
6076        logits: &CudaSlice<f32>,
6077        t: usize,
6078        n_expert: usize,
6079        n_used: usize,
6080    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6081        let f = self.func("moe_router_topk_f32");
6082        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
6083        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
6084        let cfg = LaunchConfig {
6085            grid_dim: (t as u32, 1, 1),
6086            block_dim: (n_expert as u32, 1, 1),
6087            shared_mem_bytes: 0,
6088        };
6089        let (ne, nu) = (n_expert as i32, n_used as i32);
6090        let __s_b = self.gpu.stream();
6091        let mut b = __s_b.launch_builder(&f);
6092        b.arg(logits)
6093            .arg(&mut sel_idx)
6094            .arg(&mut sel_w)
6095            .arg(&ne)
6096            .arg(&nu);
6097        unsafe {
6098            b.launch(cfg)?;
6099        }
6100        Ok((sel_idx, sel_w))
6101    }
6102
6103    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
6104    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
6105    pub fn moe_router_topk_scaled(
6106        &self,
6107        logits: &CudaSlice<f32>,
6108        t: usize,
6109        n_expert: usize,
6110        n_used: usize,
6111        ex_scale: &CudaSlice<f32>,
6112    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6113        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
6114        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
6115        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
6116        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
6117        let f = self.func("moe_router_topk_scaled_f32");
6118        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
6119        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
6120        let cfg = LaunchConfig {
6121            grid_dim: (t as u32, 1, 1),
6122            block_dim: (n_expert as u32, 1, 1),
6123            shared_mem_bytes: 0,
6124        };
6125        let (ne, nu) = (n_expert as i32, n_used as i32);
6126        let __s_b = self.gpu.stream();
6127        let mut b = __s_b.launch_builder(&f);
6128        b.arg(logits)
6129            .arg(&mut sel_idx)
6130            .arg(&mut sel_w)
6131            .arg(&ne)
6132            .arg(&nu)
6133            .arg(ex_scale);
6134        unsafe {
6135            b.launch(cfg)?;
6136        }
6137        Ok((sel_idx, sel_w))
6138    }
6139
6140    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
6141    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
6142    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
6143    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
6144    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
6145    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
6146    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
6147    pub fn moe_router_topk_host(
6148        &self,
6149        logits: &CudaSlice<f32>,
6150        t: usize,
6151        n_expert: usize,
6152        n_used: usize,
6153    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6154        let f = self.func("moe_router_topk_f32");
6155        let n = t * n_used;
6156        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
6157        let mut sel_w = self.alloc_uninit::<f32>(n)?;
6158        let cfg = LaunchConfig {
6159            grid_dim: (t as u32, 1, 1),
6160            block_dim: (n_expert as u32, 1, 1),
6161            shared_mem_bytes: 0,
6162        };
6163        let (ne, nu) = (n_expert as i32, n_used as i32);
6164        let __s_b = self.gpu.stream();
6165        let mut b = __s_b.launch_builder(&f);
6166        b.arg(logits)
6167            .arg(&mut sel_idx)
6168            .arg(&mut sel_w)
6169            .arg(&ne)
6170            .arg(&nu);
6171        unsafe {
6172            b.launch(cfg)?;
6173        }
6174        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
6175        let bytes = n * 8;
6176        let mut guard = self.router_stage.lock().unwrap();
6177        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
6178            *guard = Some(PinnedStage::new(bytes.max(4096))?);
6179        }
6180        let stage = guard.as_mut().unwrap();
6181        let (si, sw) = unsafe {
6182            (
6183                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
6184                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
6185            )
6186        };
6187        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
6188        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
6189        self.gpu.stream().synchronize()?; // ONE sync for both
6190        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
6191    }
6192
6193    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
6194    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
6195    /// original expert ids before top-k. Exact key ties choose the smaller original id.
6196    #[allow(clippy::too_many_arguments)]
6197    pub fn moe_router_sigmoid_topk(
6198        &self,
6199        logits: &CudaSlice<f32>,
6200        t: usize,
6201        n_expert: usize,
6202        n_used: usize,
6203        active_count: usize,
6204        correction_bias: &CudaSlice<f32>,
6205        active: &CudaSlice<u8>,
6206        scaling_factor: f32,
6207        route_norm: bool,
6208    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6209        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6210        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
6211            return Err(format!(
6212                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
6213            )
6214            .into());
6215        }
6216        if logits.len() < t * n_expert
6217            || correction_bias.len() != n_expert
6218            || active.len() != n_expert
6219        {
6220            return Err(format!(
6221                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
6222                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
6223            ).into());
6224        }
6225        let f = self.func(crate::sigmoid_topk_kernel(
6226            crate::sig_expf_dev_on(),
6227            crate::topk_fast_on(),
6228            n_used,
6229        ));
6230        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
6231        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
6232        let threads = n_expert.div_ceil(32) * 32;
6233        let cfg = LaunchConfig {
6234            grid_dim: (t as u32, 1, 1),
6235            block_dim: (threads as u32, 1, 1),
6236            shared_mem_bytes: 0,
6237        };
6238        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
6239        let __s_b = self.gpu.stream();
6240        let mut b = __s_b.launch_builder(&f);
6241        b.arg(logits)
6242            .arg(correction_bias)
6243            .arg(active)
6244            .arg(&mut sel_idx)
6245            .arg(&mut sel_w)
6246            .arg(&ne)
6247            .arg(&nu)
6248            .arg(&scaling_factor)
6249            .arg(&rn);
6250        unsafe {
6251            b.launch(cfg)?;
6252        }
6253        Ok((sel_idx, sel_w))
6254    }
6255
6256    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
6257    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
6258    #[allow(clippy::too_many_arguments)]
6259    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
6260    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
6261    /// the model engine can wait on it with a same-device stream memop.
6262    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
6263        if ptr == 0 {
6264            return Err("ring_flag_raw: unarmed flag".into());
6265        }
6266        let f = self.func("memra_ring_flag");
6267        let cfg = LaunchConfig {
6268            grid_dim: (1, 1, 1),
6269            block_dim: (32, 1, 1),
6270            shared_mem_bytes: 0,
6271        };
6272        let __s_b = self.gpu.stream();
6273        let mut b = __s_b.launch_builder(&f);
6274        b.arg(&ptr).arg(&value);
6275        unsafe {
6276            b.launch(cfg)?;
6277        }
6278        Ok(())
6279    }
6280
6281    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
6282    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
6283    pub fn moe_sel_w_mirror(
6284        &self,
6285        sel_src: &CudaSlice<i32>,
6286        w_src: &CudaSlice<f32>,
6287        sel_dst: &mut CudaSlice<i32>,
6288        w_dst: &mut CudaSlice<f32>,
6289        n: usize,
6290    ) -> Result<(), Box<dyn std::error::Error>> {
6291        if n == 0
6292            || n > i32::MAX as usize
6293            || sel_src.len() < n
6294            || w_src.len() < n
6295            || sel_dst.len() < n
6296            || w_dst.len() < n
6297        {
6298            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
6299        }
6300        let f = self.func("moe_sel_w_mirror");
6301        let threads = if n <= 32 { 32 } else { 128 };
6302        let cfg = LaunchConfig {
6303            grid_dim: ((n as u32).div_ceil(threads), 1, 1),
6304            block_dim: (threads, 1, 1),
6305            shared_mem_bytes: 0,
6306        };
6307        let ni = n as i32;
6308        let __s_b = self.gpu.stream();
6309        let mut b = __s_b.launch_builder(&f);
6310        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
6311        unsafe {
6312            b.launch(cfg)?;
6313        }
6314        Ok(())
6315    }
6316
6317    /// One-launch W4A16 EP staging: peer-read the active f32 input plus routed ids/weights from
6318    /// the root device, round the input directly into the rank-local BF16 buffer, and mirror the
6319    /// fixed route metadata. The caller orders root production with an entry event.
6320    #[allow(clippy::too_many_arguments)]
6321    pub fn nvfp4_ep_stage_inputs(
6322        &self,
6323        input_src: &CudaSlice<f32>,
6324        sel_src: &CudaSlice<i32>,
6325        w_src: &CudaSlice<f32>,
6326        input_bf16_dst: &mut CudaSlice<u8>,
6327        sel_dst: &mut CudaSlice<i32>,
6328        w_dst: &mut CudaSlice<f32>,
6329        input_values: usize,
6330        pairs: usize,
6331        copy_weights: bool,
6332    ) -> Result<(), Box<dyn std::error::Error>> {
6333        if input_values == 0
6334            || pairs == 0
6335            || input_src.len() < input_values
6336            || sel_src.len() < pairs
6337            || w_src.len() < pairs
6338            || input_bf16_dst.len() < 2 * input_values
6339            || sel_dst.len() < pairs
6340            || w_dst.len() < pairs
6341        {
6342            return Err(format!(
6343                "W4A16 EP stage geometry input={} sel={} weights={} input_bf16={} \
6344                 sel_dst={} weights_dst={} active={input_values} pairs={pairs}",
6345                input_src.len(),
6346                sel_src.len(),
6347                w_src.len(),
6348                input_bf16_dst.len(),
6349                sel_dst.len(),
6350                w_dst.len(),
6351            )
6352            .into());
6353        }
6354        let f = self.func("nvfp4_ep_stage_inputs");
6355        let n = input_values.max(pairs);
6356        let cfg = LaunchConfig::for_num_elems(n as u32);
6357        let (input_values, pairs, copy_weights) =
6358            (input_values as i32, pairs as i32, i32::from(copy_weights));
6359        let __s_b = self.gpu.stream();
6360        let mut b = __s_b.launch_builder(&f);
6361        b.arg(input_src)
6362            .arg(sel_src)
6363            .arg(w_src)
6364            .arg(input_bf16_dst)
6365            .arg(sel_dst)
6366            .arg(w_dst)
6367            .arg(&input_values)
6368            .arg(&pairs)
6369            .arg(&copy_weights);
6370        unsafe {
6371            b.launch(cfg)?;
6372        }
6373        Ok(())
6374    }
6375
6376    /// Capture-safe twin of `nvfp4_ep_stage_inputs`: the three sources are persistent raw
6377    /// device addresses owned by the root engine. Destinations remain rank-local typed slices.
6378    #[allow(clippy::too_many_arguments)]
6379    pub fn nvfp4_ep_stage_inputs_raw(
6380        &self,
6381        input_src: u64,
6382        sel_src: u64,
6383        w_src: u64,
6384        input_bf16_dst: &mut CudaSlice<u8>,
6385        sel_dst: &mut CudaSlice<i32>,
6386        w_dst: &mut CudaSlice<f32>,
6387        input_values: usize,
6388        pairs: usize,
6389        copy_weights: bool,
6390    ) -> Result<(), Box<dyn std::error::Error>> {
6391        if input_src == 0
6392            || sel_src == 0
6393            || w_src == 0
6394            || input_values == 0
6395            || pairs == 0
6396            || input_bf16_dst.len() < 2 * input_values
6397            || sel_dst.len() < pairs
6398            || w_dst.len() < pairs
6399        {
6400            return Err(format!(
6401                "W4A16 EP raw stage geometry input={input_src:#x} sel={sel_src:#x} \
6402                 weights={w_src:#x} input_bf16={} sel_dst={} weights_dst={} \
6403                 active={input_values} pairs={pairs}",
6404                input_bf16_dst.len(),
6405                sel_dst.len(),
6406                w_dst.len(),
6407            )
6408            .into());
6409        }
6410        let f = self.func("nvfp4_ep_stage_inputs");
6411        let n = input_values.max(pairs);
6412        let cfg = LaunchConfig::for_num_elems(n as u32);
6413        let (input_values, pairs, copy_weights) =
6414            (input_values as i32, pairs as i32, i32::from(copy_weights));
6415        let __s_b = self.gpu.stream();
6416        let mut b = __s_b.launch_builder(&f);
6417        b.arg(&input_src)
6418            .arg(&sel_src)
6419            .arg(&w_src)
6420            .arg(input_bf16_dst)
6421            .arg(sel_dst)
6422            .arg(w_dst)
6423            .arg(&input_values)
6424            .arg(&pairs)
6425            .arg(&copy_weights);
6426        unsafe {
6427            b.launch(cfg)?;
6428        }
6429        Ok(())
6430    }
6431
6432    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6433    pub fn moe_router_sigmoid_topk_into(
6434        &self,
6435        logits: &CudaSlice<f32>,
6436        t: usize,
6437        n_expert: usize,
6438        n_used: usize,
6439        active_count: usize,
6440        correction_bias: &CudaSlice<f32>,
6441        active: &CudaSlice<u8>,
6442        scaling_factor: f32,
6443        route_norm: bool,
6444        sel_idx: &mut CudaSlice<i32>,
6445        sel_w: &mut CudaSlice<f32>,
6446    ) -> Result<(), Box<dyn std::error::Error>> {
6447        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6448        if n_expert == 0
6449            || n_expert > 1024
6450            || n_used == 0
6451            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
6452            || n_used > n_expert
6453            || logits.len() < t * n_expert
6454            || correction_bias.len() != n_expert
6455            || active.len() != n_expert
6456            || sel_idx.len() < t * n_used
6457            || sel_w.len() < t * n_used
6458        {
6459            return Err("sigmoid router _into geometry mismatch".into());
6460        }
6461        let f = self.func(crate::sigmoid_topk_kernel(
6462            crate::sig_expf_dev_on(),
6463            crate::topk_fast_on(),
6464            n_used,
6465        ));
6466        let threads = n_expert.div_ceil(32) * 32;
6467        let cfg = LaunchConfig {
6468            grid_dim: (t as u32, 1, 1),
6469            block_dim: (threads as u32, 1, 1),
6470            shared_mem_bytes: 0,
6471        };
6472        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
6473        let __s_b = self.gpu.stream();
6474        let mut b = __s_b.launch_builder(&f);
6475        b.arg(logits)
6476            .arg(correction_bias)
6477            .arg(active)
6478            .arg(&mut *sel_idx)
6479            .arg(&mut *sel_w)
6480            .arg(&ne)
6481            .arg(&nu)
6482            .arg(&scaling_factor)
6483            .arg(&rn);
6484        unsafe {
6485            b.launch(cfg)?;
6486        }
6487        Ok(())
6488    }
6489
6490    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
6491    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
6492    #[allow(clippy::too_many_arguments)]
6493    pub fn moe_router_sigmoid_topk_host(
6494        &self,
6495        logits: &CudaSlice<f32>,
6496        t: usize,
6497        n_expert: usize,
6498        n_used: usize,
6499        active_count: usize,
6500        correction_bias: &CudaSlice<f32>,
6501        active: &CudaSlice<u8>,
6502        scaling_factor: f32,
6503        route_norm: bool,
6504    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6505        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
6506            logits,
6507            t,
6508            n_expert,
6509            n_used,
6510            active_count,
6511            correction_bias,
6512            active,
6513            scaling_factor,
6514            route_norm,
6515        )?;
6516        let n = t * n_used;
6517        let bytes = n * 8;
6518        let mut guard = self.router_stage.lock().unwrap();
6519        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
6520            *guard = Some(PinnedStage::new(bytes.max(4096))?);
6521        }
6522        let stage = guard.as_mut().unwrap();
6523        let (si, sw) = unsafe {
6524            (
6525                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
6526                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
6527            )
6528        };
6529        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
6530        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
6531        self.gpu.stream().synchronize()?;
6532        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
6533    }
6534
6535    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
6536    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
6537    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
6538    pub fn stage_expert_async(
6539        &self,
6540        host_bytes: &[u8],
6541        scratch: &mut CudaSlice<u8>,
6542        off: usize,
6543    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
6544        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
6545        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
6546        Ok(self.copy_stream.record_event(None)?)
6547    }
6548
6549    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
6550    pub fn compute_wait(
6551        &self,
6552        ev: &cudarc::driver::CudaEvent,
6553    ) -> Result<(), Box<dyn std::error::Error>> {
6554        self.gpu.stream().wait(ev)?;
6555        Ok(())
6556    }
6557
6558    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
6559    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
6560    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
6561    /// CudaView base+offset pointer is honored by the launch arg.
6562    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6563    pub fn qmatvec_view(
6564        &self,
6565        w: &CudaSlice<u8>,
6566        range: std::ops::Range<usize>,
6567        x: &cudarc::driver::CudaView<f32>,
6568        m: usize,
6569        in_f: usize,
6570        out_f: usize,
6571        qtype: i32,
6572        row_bytes: usize,
6573    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6574        self.qmatvec_view_inner(w, range, x, m, in_f, out_f, qtype, row_bytes)
6575    }
6576
6577    /// W4A16 expert matvec: round the floating activation to checkpoint BF16 before the
6578    /// existing f32-dequant weight dot. The output remains f32. This is selected per model by
6579    /// `MoeWeights`; it is not a process-global NVFP4 policy.
6580    #[allow(clippy::too_many_arguments)]
6581    pub fn qmatvec_view_bf16_activation(
6582        &self,
6583        w: &CudaSlice<u8>,
6584        range: std::ops::Range<usize>,
6585        x: &cudarc::driver::CudaView<f32>,
6586        m: usize,
6587        in_f: usize,
6588        out_f: usize,
6589        qtype: i32,
6590        row_bytes: usize,
6591    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6592        let n = m * in_f;
6593        if x.len() != n {
6594            return Err(format!(
6595                "W4A16 BF16 activation input length {} != {m}x{in_f}",
6596                x.len()
6597            )
6598            .into());
6599        }
6600        let mut x_bf16 = self.alloc_u8_uninit(n * 2)?;
6601        self.f32_to_bf16_v(x, &mut x_bf16, n)?;
6602        let x_f32 = self.bf16_to_f32(&x_bf16.slice(0..n * 2), n)?;
6603        self.qmatvec_view_inner(
6604            w,
6605            range,
6606            &x_f32.slice(0..n),
6607            m,
6608            in_f,
6609            out_f,
6610            qtype,
6611            row_bytes,
6612        )
6613    }
6614
6615    #[allow(clippy::too_many_arguments)]
6616    fn qmatvec_view_inner(
6617        &self,
6618        w: &CudaSlice<u8>,
6619        range: std::ops::Range<usize>,
6620        x: &cudarc::driver::CudaView<f32>,
6621        m: usize,
6622        in_f: usize,
6623        out_f: usize,
6624        qtype: i32,
6625        row_bytes: usize,
6626    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6627        let f = self.func("qmatvec_f32");
6628        let wv = w.slice(range); // CudaView<u8>, offset honored
6629        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6630        let cfg = LaunchConfig {
6631            grid_dim: (out_f as u32, m as u32, 1),
6632            block_dim: (256, 1, 1),
6633            shared_mem_bytes: 0,
6634        };
6635        let (inf, outf, mi, qt, rb) =
6636            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6637        let __s_b = self.gpu.stream();
6638        let mut b = __s_b.launch_builder(&f);
6639        b.arg(&wv)
6640            .arg(x)
6641            .arg(&mut y)
6642            .arg(&inf)
6643            .arg(&outf)
6644            .arg(&mi)
6645            .arg(&qt)
6646            .arg(&rb);
6647        unsafe {
6648            b.launch(cfg)?;
6649        }
6650        Ok(y)
6651    }
6652
6653    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
6654    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
6655    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
6656    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
6657    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
6658    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
6659    #[allow(clippy::too_many_arguments)]
6660    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
6661    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
6662    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
6663    pub fn moe_gate_up_silu8_q8(
6664        &self,
6665        gp: WPtr8,
6666        up: WPtr8,
6667        aq: &CudaSlice<i8>,
6668        ad: &CudaSlice<f32>,
6669        in_f: usize,
6670        n_ff: usize,
6671        n_used: usize,
6672        qt_g: i32,
6673        qt_u: i32,
6674        rb_g: usize,
6675        rb_u: usize,
6676    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6677        let f = self.func("moe_gate_up_silu8_q8");
6678        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6679        let cfg = LaunchConfig {
6680            grid_dim: (n_ff as u32, n_used as u32, 1),
6681            block_dim: (32, 1, 1),
6682            shared_mem_bytes: 0,
6683        };
6684        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
6685        let __s_b = self.gpu.stream();
6686        let mut b = __s_b.launch_builder(&f);
6687        b.arg(&gp)
6688            .arg(&up)
6689            .arg(aq)
6690            .arg(ad)
6691            .arg(&mut act)
6692            .arg(&inf)
6693            .arg(&nff)
6694            .arg(&qt_g)
6695            .arg(&qt_u)
6696            .arg(&rbg)
6697            .arg(&rbu);
6698        unsafe {
6699            b.launch(cfg)?;
6700        }
6701        Ok(act)
6702    }
6703
6704    /// The PRE-clamped, macro-folding twin of [`Engine::moe_gate_up_silu8_q8`] — the kernel
6705    /// class for any MoE family whose activation clamps the gate BEFORE the silu (glm5_next is
6706    /// the first such family; the door names the arithmetic, not the family).
6707    ///
6708    /// Same grid/block/dots/warp reduction; the epilogue is
6709    /// `silu(min(gate*gs, limit)) * clamp(up*us, ±limit)` — `swiglu_preclamped_mul_scaled_f32`'s
6710    /// expression verbatim — and `gs`/`us` carry the SELECTED experts' NVFP4 `weight_scale_2`
6711    /// macro scales in router slot order (1.0 for a macro-free bank).
6712    ///
6713    /// `limit` must be live: at `limit == 0` every gate collapses to `silu(0) == 0`, so a caller
6714    /// with no clamp belongs on the plain-SiLU sibling, not here. Same contract as
6715    /// [`Engine::swiglu_preclamped_mul_scaled`].
6716    #[allow(clippy::too_many_arguments)]
6717    pub fn moe_gate_up_preclamp8_q8(
6718        &self,
6719        gp: WPtr8,
6720        up: WPtr8,
6721        aq: &CudaSlice<i8>,
6722        ad: &CudaSlice<f32>,
6723        gs: F32x8,
6724        us: F32x8,
6725        limit: f32,
6726        in_f: usize,
6727        n_ff: usize,
6728        n_used: usize,
6729        qt_g: i32,
6730        qt_u: i32,
6731        rb_g: usize,
6732        rb_u: usize,
6733    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6734        debug_assert!(
6735            limit > 1e-6,
6736            "moe_gate_up_preclamp8_q8 needs a live limit; use moe_gate_up_silu8_q8"
6737        );
6738        let f = self.func("moe_gate_up_preclamp8_q8");
6739        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6740        let cfg = LaunchConfig {
6741            grid_dim: (n_ff as u32, n_used as u32, 1),
6742            block_dim: (32, 1, 1),
6743            shared_mem_bytes: 0,
6744        };
6745        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
6746        let __s_b = self.gpu.stream();
6747        let mut b = __s_b.launch_builder(&f);
6748        b.arg(&gp)
6749            .arg(&up)
6750            .arg(aq)
6751            .arg(ad)
6752            .arg(&gs)
6753            .arg(&us)
6754            .arg(&limit)
6755            .arg(&mut act)
6756            .arg(&inf)
6757            .arg(&nff)
6758            .arg(&qt_g)
6759            .arg(&qt_u)
6760            .arg(&rbg)
6761            .arg(&rbu);
6762        unsafe {
6763            b.launch(cfg)?;
6764        }
6765        Ok(act)
6766    }
6767
6768    #[allow(clippy::too_many_arguments)]
6769    pub fn moe_down8_fma_q8(
6770        &self,
6771        dp: WPtr8,
6772        w: F32x8,
6773        aq2: &CudaSlice<i8>,
6774        ad2: &CudaSlice<f32>,
6775        dst: &mut cudarc::driver::CudaViewMut<f32>,
6776        in_f: usize,
6777        out_f: usize,
6778        n_used: usize,
6779        qt: i32,
6780        rb: usize,
6781    ) -> Result<(), Box<dyn std::error::Error>> {
6782        let f = self.func("moe_down8_fma_q8");
6783        let cfg = LaunchConfig {
6784            grid_dim: (out_f as u32, 1, 1),
6785            block_dim: (32, 1, 1),
6786            shared_mem_bytes: 0,
6787        };
6788        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
6789        let __s_b = self.gpu.stream();
6790        let mut b = __s_b.launch_builder(&f);
6791        b.arg(&dp)
6792            .arg(&w)
6793            .arg(aq2)
6794            .arg(ad2)
6795            .arg(dst)
6796            .arg(&inf)
6797            .arg(&outf)
6798            .arg(&nu)
6799            .arg(&qt)
6800            .arg(&rbi);
6801        unsafe {
6802            b.launch(cfg)?;
6803        }
6804        Ok(())
6805    }
6806
6807    /// DEVICE-SIDE build of the verify-rows pair's pointer/scale tables (door D,
6808    /// `MEMRA_MOE_VROWS_DEV_TABLES`) from the router's own device selection. Replaces the host
6809    /// loop plus its two pageable HtoD, and lets the caller skip the router's pinned readback
6810    /// and its full `cuStreamSynchronize` entirely. Arithmetic is term-for-term the host loop's
6811    /// (see the kernel comment in `qmatvec.cu`), so the tables — and therefore every downstream
6812    /// byte — are identical.
6813    ///
6814    /// `macros` is the model's immutable `(gate, up, down)` `weight_scale_2` host planes, or
6815    /// `None` for a non-macro bank (the kernel then takes 1.0f, `macro_scale`'s own answer).
6816    /// The planes get a resident device mirror keyed by `(il, plane)` on first use — uploading
6817    /// them per call would ADD three HtoD to a door whose purpose is removing two.
6818    #[allow(clippy::too_many_arguments)]
6819    // allow: the parameter list mirrors the kernel/FFI/call contract
6820    pub fn moe_vrows_tables_from_sel(
6821        &self,
6822        sel: &CudaSlice<i32>,
6823        selw: &CudaSlice<f32>,
6824        il: u16,
6825        macros: Option<(&[f32], &[f32], &[f32])>,
6826        (pg, pu, pd): (u64, u64, u64),
6827        (sg, su, sd): (usize, usize, usize),
6828        n_pairs: usize,
6829        ptrs: &mut CudaSlice<u64>,
6830        scl: &mut CudaSlice<f32>,
6831    ) -> Result<(), Box<dyn std::error::Error>> {
6832        debug_assert!(sel.len() >= n_pairs && selw.len() >= n_pairs);
6833        // `>=` not `==`: door E appends a fourth (expert-major order) plane to the same table.
6834        debug_assert!(ptrs.len() >= 3 * n_pairs);
6835        debug_assert_eq!(scl.len(), 3 * n_pairs);
6836        // Resident macro mirrors, uploaded once per (layer, plane). The guard is held across the
6837        // launch because `CudaSlice` is not clonable — the same shape as the w8-mirror sites.
6838        let mut mac = self
6839            .vrows_macro_dev
6840            .lock()
6841            .map_err(|_| "vrows macro mirror map is poisoned")?;
6842        if let Some((hg, hu, hd)) = macros {
6843            for (plane, host) in [(0u8, hg), (1u8, hu), (2u8, hd)] {
6844                // `entry` rather than contains_key+insert: the upload is fallible, so it lands in
6845                // the Vacant arm instead of an `or_insert_with` closure.
6846                if let std::collections::hash_map::Entry::Vacant(slot) = mac.entry((il, plane)) {
6847                    slot.insert(self.htod(host)?);
6848                }
6849            }
6850        }
6851        // Absent macro planes: the three kernel pointers must still be legal device addresses,
6852        // so the call aliases the selection weights and never dereferences them (have_macros=0).
6853        let (mg, mu, md, have) = match macros {
6854            Some(_) => (
6855                mac.get(&(il, 0)).expect("gate macro mirror built above"),
6856                mac.get(&(il, 1)).expect("up macro mirror built above"),
6857                mac.get(&(il, 2)).expect("down macro mirror built above"),
6858                1i32,
6859            ),
6860            None => (selw, selw, selw, 0i32),
6861        };
6862        let f = self.func("moe_vrows_tables_from_sel");
6863        let threads = 128u32;
6864        let cfg = LaunchConfig {
6865            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
6866            block_dim: (threads, 1, 1),
6867            shared_mem_bytes: 0,
6868        };
6869        let (sgi, sui, sdi) = (sg as i64, su as i64, sd as i64);
6870        let (npi, havei) = (n_pairs as i32, have);
6871        let __s_b = self.gpu.stream();
6872        let mut b = __s_b.launch_builder(&f);
6873        b.arg(sel)
6874            .arg(selw)
6875            .arg(mg)
6876            .arg(mu)
6877            .arg(md)
6878            .arg(&mut *ptrs)
6879            .arg(&mut *scl)
6880            .arg(&pg)
6881            .arg(&pu)
6882            .arg(&pd)
6883            .arg(&sgi)
6884            .arg(&sui)
6885            .arg(&sdi)
6886            .arg(&npi)
6887            .arg(&havei);
6888        unsafe {
6889            b.launch(cfg)?;
6890        }
6891        Ok(())
6892    }
6893
6894    /// DEVICE-SIDE build of the verify-rows pair's EXPERT-MAJOR order plane (door E,
6895    /// `MEMRA_MOE_VROWS_DEDUP_ORDER`) from the router's own device selection, written into the
6896    /// pointer table's fourth plane `ptrs[3*n_pairs ..)`. Bit-identical to
6897    /// [`crate::vrows_expert_major_order`]: both are a stable order on `(expert id, pair index)`,
6898    /// the kernel by counting rank (see its comment in `qmatvec.cu`), the host by a stable sort.
6899    ///
6900    /// This launch exists ONLY in the door-D (device tables) arm — the host arm appends the plane
6901    /// to the vector it already uploads, so it costs zero extra transfers there. Cost in the
6902    /// device arm: 42 launches/round = ~0.093 ms at the box's 2.216 us eager-launch constant,
6903    /// against a predicted -2.17 ms/round; folding it into `moe_vrows_tables_from_sel` (same
6904    /// inputs, same one-thread-per-pair grid) is the named follow-up that recovers it.
6905    pub fn moe_vrows_order_from_sel(
6906        &self,
6907        sel: &CudaSlice<i32>,
6908        n_pairs: usize,
6909        ptrs: &mut CudaSlice<u64>,
6910    ) -> Result<(), Box<dyn std::error::Error>> {
6911        debug_assert!(sel.len() >= n_pairs);
6912        debug_assert!(
6913            ptrs.len() >= 4 * n_pairs,
6914            "the order plane lives at ptrs[3*n_pairs .. 4*n_pairs)"
6915        );
6916        let f = self.func("moe_vrows_order_from_sel");
6917        let threads = 128u32;
6918        let cfg = LaunchConfig {
6919            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
6920            block_dim: (threads, 1, 1),
6921            shared_mem_bytes: 0,
6922        };
6923        let np = n_pairs as i32;
6924        let __s_b = self.gpu.stream();
6925        let mut b = __s_b.launch_builder(&f);
6926        b.arg(sel).arg(&mut *ptrs).arg(&np);
6927        unsafe {
6928            b.launch(cfg)?;
6929        }
6930        Ok(())
6931    }
6932
6933    /// Verify-rows twin of [`Self::moe_gate_up_preclamp8_q8`] (lane/glm5-vrest): one launch
6934    /// covers ALL `n_pairs = t * n_used` routed pairs of a spec-verify batch. `ptrs` /
6935    /// `scl` are the `[3 * n_pairs]` plane-major (gate | up | down) expert-pointer and
6936    /// scale tables (gs | us | w*macro_down); per pair the kernel body is the t=1 fused
6937    /// epilogue's verbatim, bit-gated per row vs the sequential chain.
6938    #[allow(clippy::too_many_arguments)]
6939    // allow: the parameter list mirrors the kernel/FFI/call contract
6940    pub fn moe_gate_up_preclamp8_q8_rows(
6941        &self,
6942        ptrs: &CudaSlice<u64>,
6943        scl: &CudaSlice<f32>,
6944        aq: &CudaSlice<i8>,
6945        ad: &CudaSlice<f32>,
6946        limit: f32,
6947        in_f: usize,
6948        n_ff: usize,
6949        n_used: usize,
6950        n_pairs: usize,
6951        qt_g: i32,
6952        qt_u: i32,
6953        rb_g: usize,
6954        rb_u: usize,
6955    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6956        debug_assert!(
6957            limit > 1e-6,
6958            "moe_gate_up_preclamp8_q8_rows needs a live limit; the kernel collapses every gate \
6959             to silu(0) at limit 0"
6960        );
6961        debug_assert!(ptrs.len() >= 3 * n_pairs);
6962        debug_assert_eq!(scl.len(), 3 * n_pairs);
6963        // MEMRA_MOE_VROWS_DEDUP_ORDER (lane/glm5-dedup door E, default OFF): the `_ord` twin —
6964        // pair index the FASTEST grid dimension, walked in expert-major order from the table's
6965        // fourth plane, so two verify rows sharing an expert read the identical gate/up rows in
6966        // adjacent blocks. `ptrs.len() >= 4*n_pairs` is a REQUIREMENT not a hint: the door engages
6967        // only when the caller actually built the order plane, so a direct launcher call with the
6968        // shipped 3-plane table (every standing gate) keeps the shipped program. Door M wins the
6969        // tie by being tested first — the two are refused together rather than crossed.
6970        let packed = moe_vrows_pack_on();
6971        let ordered =
6972            !packed && moe_vrows_dedup_order_on() && ptrs.len() >= 4 * n_pairs && n_ff <= 65535;
6973        let (f, cfg) = if packed {
6974            if MOE_VROWS_PACK_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
6975                eprintln!(
6976                    "[moe-vrows-pack] engaged: 4-warp blocks on the verify-rows MoE pair \
6977                     (MEMRA_MOE_VROWS_PACK=1)"
6978                );
6979            }
6980            (
6981                self.func("moe_gate_up_preclamp8_q8_rows_w4"),
6982                LaunchConfig {
6983                    grid_dim: ((n_ff as u32).div_ceil(4), n_pairs as u32, 1),
6984                    block_dim: (32, 4, 1),
6985                    shared_mem_bytes: 0,
6986                },
6987            )
6988        } else if ordered {
6989            if MOE_VROWS_DEDUP_ORDER_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
6990                == 0
6991            {
6992                eprintln!(
6993                    "[moe-vrows-dedup-order] engaged: verify-rows gate/up walks the pair union \
6994                     EXPERT-MAJOR with the pair index as the fastest grid dimension, so the \
6995                     21.96%-measured repeat visits read a shared expert slab's rows in adjacent \
6996                     blocks (MEMRA_MOE_VROWS_DEDUP_ORDER=1)"
6997                );
6998            }
6999            (
7000                self.func("moe_gate_up_preclamp8_q8_rows_ord"),
7001                LaunchConfig {
7002                    grid_dim: (n_pairs as u32, n_ff as u32, 1),
7003                    block_dim: (32, 1, 1),
7004                    shared_mem_bytes: 0,
7005                },
7006            )
7007        } else {
7008            (
7009                self.func("moe_gate_up_preclamp8_q8_rows"),
7010                LaunchConfig {
7011                    grid_dim: (n_ff as u32, n_pairs as u32, 1),
7012                    block_dim: (32, 1, 1),
7013                    shared_mem_bytes: 0,
7014                },
7015            )
7016        };
7017        // Door W: the vrows launcher is verify-walk-only; act is a pooled draw.
7018        let mut act = self.vws_uninit(n_pairs * n_ff)?;
7019        let (inf, nff, nu, np) = (in_f as i32, n_ff as i32, n_used as i32, n_pairs as i32);
7020        let (rbg, rbu) = (rb_g as i64, rb_u as i64);
7021        let __s_b = self.gpu.stream();
7022        let mut b = __s_b.launch_builder(&f);
7023        b.arg(ptrs)
7024            .arg(scl)
7025            .arg(aq)
7026            .arg(ad)
7027            .arg(&limit)
7028            .arg(&mut act)
7029            .arg(&inf)
7030            .arg(&nff)
7031            .arg(&nu)
7032            .arg(&np)
7033            .arg(&qt_g)
7034            .arg(&qt_u)
7035            .arg(&rbg)
7036            .arg(&rbu);
7037        unsafe {
7038            b.launch(cfg)?;
7039        }
7040        Ok(act)
7041    }
7042
7043    /// Verify-rows twin of [`Self::moe_down8_fma_q8`] (lane/glm5-vrest): every verify row's
7044    /// slot-ordered down+FMA chain in one launch. `dst` is `[t, out_f]`, fully overwritten;
7045    /// `ptrs`/`scl` are the same tables the gate/up rows launch consumed (down plane).
7046    #[allow(clippy::too_many_arguments)]
7047    // allow: the parameter list mirrors the kernel/FFI/call contract
7048    pub fn moe_down8_fma_q8_rows(
7049        &self,
7050        ptrs: &CudaSlice<u64>,
7051        scl: &CudaSlice<f32>,
7052        aq2: &CudaSlice<i8>,
7053        ad2: &CudaSlice<f32>,
7054        dst: &mut CudaSlice<f32>,
7055        in_f: usize,
7056        out_f: usize,
7057        n_used: usize,
7058        n_pairs: usize,
7059        qt: i32,
7060        rb: usize,
7061    ) -> Result<(), Box<dyn std::error::Error>> {
7062        debug_assert!(ptrs.len() >= 3 * n_pairs);
7063        debug_assert_eq!(scl.len(), 3 * n_pairs);
7064        debug_assert_eq!(n_pairs % n_used, 0, "pairs are dense slot-major");
7065        let t = n_pairs / n_used;
7066        debug_assert!(dst.len() >= t * out_f);
7067        // MEMRA_MOE_VROWS_PACK (door M): the _w4 twin, same packing as the gate/up launch.
7068        let packed = moe_vrows_pack_on();
7069        // MEMRA_MOE_VROWS_DOWN_TMAJ (door E-down): grid transposed to (t, out_f) — token fastest —
7070        // so the t verify rows at one output row are adjacent blocks and a repeated expert's down
7071        // row is read once for every token that shares it. The slot-ordered __fmaf_rn chain is
7072        // inside the block and keeps its ORIGINAL slot order; only the grid moves. Needs no table
7073        // plane (the down chain cannot be permuted), so it composes with either table provenance.
7074        let tmaj = !packed && moe_vrows_down_tmaj_on() && out_f <= 65535;
7075        let (f, cfg) = if packed {
7076            (
7077                self.func("moe_down8_fma_q8_rows_w4"),
7078                LaunchConfig {
7079                    grid_dim: ((out_f as u32).div_ceil(4), t as u32, 1),
7080                    block_dim: (32, 4, 1),
7081                    shared_mem_bytes: 0,
7082                },
7083            )
7084        } else if tmaj {
7085            if MOE_VROWS_DOWN_TMAJ_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
7086                == 0
7087            {
7088                eprintln!(
7089                    "[moe-vrows-down-tmaj] engaged: verify-rows down/FMA grid transposed to \
7090                     (t, out_f) so the verify rows at one output row are adjacent blocks; the \
7091                     slot-ordered FMA chain is unchanged (MEMRA_MOE_VROWS_DOWN_TMAJ=1)"
7092                );
7093            }
7094            (
7095                self.func("moe_down8_fma_q8_rows_tmaj"),
7096                LaunchConfig {
7097                    grid_dim: (t as u32, out_f as u32, 1),
7098                    block_dim: (32, 1, 1),
7099                    shared_mem_bytes: 0,
7100                },
7101            )
7102        } else {
7103            (
7104                self.func("moe_down8_fma_q8_rows"),
7105                LaunchConfig {
7106                    grid_dim: (out_f as u32, t as u32, 1),
7107                    block_dim: (32, 1, 1),
7108                    shared_mem_bytes: 0,
7109                },
7110            )
7111        };
7112        let (inf, outf, nu, np, rbi) = (
7113            in_f as i32,
7114            out_f as i32,
7115            n_used as i32,
7116            n_pairs as i32,
7117            rb as i64,
7118        );
7119        let __s_b = self.gpu.stream();
7120        let mut b = __s_b.launch_builder(&f);
7121        b.arg(ptrs)
7122            .arg(scl)
7123            .arg(aq2)
7124            .arg(ad2)
7125            .arg(dst)
7126            .arg(&inf)
7127            .arg(&outf)
7128            .arg(&nu)
7129            .arg(&np)
7130            .arg(&qt)
7131            .arg(&rbi);
7132        unsafe {
7133            b.launch(cfg)?;
7134        }
7135        Ok(())
7136    }
7137
7138    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
7139    #[allow(clippy::too_many_arguments)]
7140    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7141    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
7142    pub fn qmatvec_expert_q8(
7143        &self,
7144        w: &CudaSlice<u8>,
7145        range: std::ops::Range<usize>,
7146        aq: &CudaSlice<i8>,
7147        ad: &CudaSlice<f32>,
7148        m: usize,
7149        in_f: usize,
7150        out_f: usize,
7151        qtype: i32,
7152        row_bytes: usize,
7153    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7154        let f = self.func("qmatvec_expert_q8");
7155        let wv = w.slice(range);
7156        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7157        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
7158        let cfg = LaunchConfig {
7159            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
7160            block_dim: (32, ROWS, 1),
7161            shared_mem_bytes: 0,
7162        };
7163        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7164        let __s_b = self.gpu.stream();
7165        let mut b = __s_b.launch_builder(&f);
7166        b.arg(&wv)
7167            .arg(aq)
7168            .arg(ad)
7169            .arg(&mut y)
7170            .arg(&inf)
7171            .arg(&outf)
7172            .arg(&mi)
7173            .arg(&qtype)
7174            .arg(&rbi);
7175        unsafe {
7176            b.launch(cfg)?;
7177        }
7178        Ok(y)
7179    }
7180
7181    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7182    pub fn moe_gate_up_silu8(
7183        &self,
7184        gp: WPtr8,
7185        up: WPtr8,
7186        x: &cudarc::driver::CudaView<f32>,
7187        in_f: usize,
7188        n_ff: usize,
7189        n_used: usize,
7190        qt_g: i32,
7191        qt_u: i32,
7192        rb_g: usize,
7193        rb_u: usize,
7194    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7195        let f = self.func("moe_gate_up_silu8_f32");
7196        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
7197        let cfg = LaunchConfig {
7198            grid_dim: (n_ff as u32, n_used as u32, 1),
7199            block_dim: (256, 1, 1),
7200            shared_mem_bytes: 0,
7201        };
7202        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7203        let __s_b = self.gpu.stream();
7204        let mut b = __s_b.launch_builder(&f);
7205        b.arg(&gp)
7206            .arg(&up)
7207            .arg(x)
7208            .arg(&mut act)
7209            .arg(&inf)
7210            .arg(&nff)
7211            .arg(&qt_g)
7212            .arg(&qt_u)
7213            .arg(&rbg)
7214            .arg(&rbu);
7215        unsafe {
7216            b.launch(cfg)?;
7217        }
7218        Ok(act)
7219    }
7220
7221    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
7222    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
7223    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
7224    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
7225    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
7226    #[allow(clippy::too_many_arguments)]
7227    pub fn moe_down8_fma_into(
7228        &self,
7229        dp: WPtr8,
7230        w: F32x8,
7231        act: &CudaSlice<f32>,
7232        dst: &mut cudarc::driver::CudaViewMut<f32>,
7233        in_f: usize,
7234        out_f: usize,
7235        n_used: usize,
7236        qt: i32,
7237        rb: usize,
7238    ) -> Result<(), Box<dyn std::error::Error>> {
7239        let f = self.func("moe_down8_fma_f32");
7240        let cfg = LaunchConfig {
7241            grid_dim: (out_f as u32, 1, 1),
7242            block_dim: (256, 1, 1),
7243            shared_mem_bytes: 0,
7244        };
7245        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7246        let __s_b = self.gpu.stream();
7247        let mut b = __s_b.launch_builder(&f);
7248        b.arg(&dp)
7249            .arg(&w)
7250            .arg(act)
7251            .arg(dst)
7252            .arg(&inf)
7253            .arg(&outf)
7254            .arg(&nu)
7255            .arg(&qt)
7256            .arg(&rbv);
7257        unsafe {
7258            b.launch(cfg)?;
7259        }
7260        Ok(())
7261    }
7262
7263    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
7264    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
7265    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
7266    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
7267    #[allow(clippy::too_many_arguments)]
7268    /// dp4a q8 twin of the _dev pair (resident-experts arc).
7269    ///
7270    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
7271    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
7272    /// down's FMA chain stays slot-ordered serial). Seams:
7273    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
7274    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
7275    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
7276    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
7277    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
7278    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
7279    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
7280    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
7281    ///                       only) | w8h2 (h2 x slot-parallel)
7282    #[allow(clippy::too_many_arguments)]
7283    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
7284    #[allow(clippy::too_many_arguments)]
7285    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
7286    pub fn moe_pairs_matvec_q8(
7287        &self,
7288        table: &CudaSlice<u64>,
7289        proj: i32,
7290        pair_tok: &CudaSlice<i32>,
7291        pair_ex: &CudaSlice<i32>,
7292        aq: &CudaSlice<i8>,
7293        ad: &CudaSlice<f32>,
7294        in_f: usize,
7295        out_f: usize,
7296        n_expert: usize,
7297        n_pairs: usize,
7298        qtype: i32,
7299        row_bytes: usize,
7300    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7301        let f = self.func("moe_pairs_matvec_q8");
7302        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7303        const ROWS: u32 = 4;
7304        let cfg = LaunchConfig {
7305            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
7306            block_dim: (32, ROWS, 1),
7307            shared_mem_bytes: 0,
7308        };
7309        let (inf, outf, ne, np, rbi) = (
7310            in_f as i32,
7311            out_f as i32,
7312            n_expert as i32,
7313            n_pairs as i32,
7314            row_bytes as i64,
7315        );
7316        let __s_b = self.gpu.stream();
7317        let mut b = __s_b.launch_builder(&f);
7318        b.arg(table)
7319            .arg(&proj)
7320            .arg(pair_tok)
7321            .arg(pair_ex)
7322            .arg(aq)
7323            .arg(ad)
7324            .arg(&mut y)
7325            .arg(&inf)
7326            .arg(&outf)
7327            .arg(&ne)
7328            .arg(&np)
7329            .arg(&qtype)
7330            .arg(&rbi);
7331        unsafe {
7332            b.launch(cfg)?;
7333        }
7334        Ok(y)
7335    }
7336
7337    /// Expert-major pair matvec (weight-reuse across each expert's token group).
7338    #[allow(clippy::too_many_arguments)]
7339    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
7340    pub fn moe_pairs_matvec_q8_em(
7341        &self,
7342        table: &CudaSlice<u64>,
7343        proj: i32,
7344        ex_ids: &CudaSlice<i32>,
7345        ex_off: &CudaSlice<i32>,
7346        ex_pairs: &CudaSlice<i32>,
7347        pair_tok: &CudaSlice<i32>,
7348        aq: &CudaSlice<i8>,
7349        ad: &CudaSlice<f32>,
7350        in_f: usize,
7351        out_f: usize,
7352        n_expert: usize,
7353        n_active: usize,
7354        n_pairs: usize,
7355        qtype: i32,
7356        row_bytes: usize,
7357    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7358        let f = self.func("moe_pairs_matvec_q8_em");
7359        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7360        const ROWS: u32 = 4;
7361        let cfg = LaunchConfig {
7362            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
7363            block_dim: (32, ROWS, 1),
7364            shared_mem_bytes: 0,
7365        };
7366        let (inf, outf, ne, na, rbi) = (
7367            in_f as i32,
7368            out_f as i32,
7369            n_expert as i32,
7370            n_active as i32,
7371            row_bytes as i64,
7372        );
7373        let __s_b = self.gpu.stream();
7374        let mut b = __s_b.launch_builder(&f);
7375        b.arg(table)
7376            .arg(&proj)
7377            .arg(ex_ids)
7378            .arg(ex_off)
7379            .arg(ex_pairs)
7380            .arg(pair_tok)
7381            .arg(aq)
7382            .arg(ad)
7383            .arg(&mut y)
7384            .arg(&inf)
7385            .arg(&outf)
7386            .arg(&ne)
7387            .arg(&na)
7388            .arg(&qtype)
7389            .arg(&rbi);
7390        unsafe {
7391            b.launch(cfg)?;
7392        }
7393        Ok(y)
7394    }
7395
7396    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
7397    // weight group once per (row,group) then dp4a's across the expert's token group.
7398    #[allow(clippy::too_many_arguments)]
7399    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
7400    pub fn moe_pairs_matvec_q8_dec(
7401        &self,
7402        table: &CudaSlice<u64>,
7403        proj: i32,
7404        ex_ids: &CudaSlice<i32>,
7405        ex_off: &CudaSlice<i32>,
7406        ex_pairs: &CudaSlice<i32>,
7407        pair_tok: &CudaSlice<i32>,
7408        aq: &CudaSlice<i8>,
7409        ad: &CudaSlice<f32>,
7410        in_f: usize,
7411        out_f: usize,
7412        n_expert: usize,
7413        n_active: usize,
7414        n_pairs: usize,
7415        qtype: i32,
7416        row_bytes: usize,
7417    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7418        let f = self.func("moe_pairs_matvec_q8_dec");
7419        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7420        const ROWS: u32 = 4;
7421        let cfg = LaunchConfig {
7422            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
7423            block_dim: (32, ROWS, 1),
7424            shared_mem_bytes: 0,
7425        };
7426        let (inf, outf, ne, na, rbi) = (
7427            in_f as i32,
7428            out_f as i32,
7429            n_expert as i32,
7430            n_active as i32,
7431            row_bytes as i64,
7432        );
7433        let __s_b = self.gpu.stream();
7434        let mut b = __s_b.launch_builder(&f);
7435        b.arg(table)
7436            .arg(&proj)
7437            .arg(ex_ids)
7438            .arg(ex_off)
7439            .arg(ex_pairs)
7440            .arg(pair_tok)
7441            .arg(aq)
7442            .arg(ad)
7443            .arg(&mut y)
7444            .arg(&inf)
7445            .arg(&outf)
7446            .arg(&ne)
7447            .arg(&na)
7448            .arg(&qtype)
7449            .arg(&rbi);
7450        unsafe {
7451            b.launch(cfg)?;
7452        }
7453        Ok(y)
7454    }
7455
7456    pub fn moe_pairs_gelu_mul(
7457        &self,
7458        gate: &CudaSlice<f32>,
7459        up: &CudaSlice<f32>,
7460        n: usize,
7461    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7462        let f = self.func("moe_pairs_gelu_mul");
7463        let mut act = self.alloc_uninit::<f32>(n)?;
7464        let cfg = LaunchConfig::for_num_elems(n as u32);
7465        let nl = n as i64;
7466        let __s_b = self.gpu.stream();
7467        let mut b = __s_b.launch_builder(&f);
7468        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
7469        unsafe {
7470            b.launch(cfg)?;
7471        }
7472        Ok(act)
7473    }
7474
7475    pub fn moe_pairs_silu_mul(
7476        &self,
7477        gate: &CudaSlice<f32>,
7478        up: &CudaSlice<f32>,
7479        n: usize,
7480    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7481        let f = self.func("moe_pairs_silu_mul");
7482        let mut act = self.alloc_uninit::<f32>(n)?;
7483        let cfg = LaunchConfig::for_num_elems(n as u32);
7484        let nl = n as i64;
7485        let __s_b = self.gpu.stream();
7486        let mut b = __s_b.launch_builder(&f);
7487        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
7488        unsafe {
7489            b.launch(cfg)?;
7490        }
7491        Ok(act)
7492    }
7493
7494    #[allow(clippy::too_many_arguments)]
7495    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
7496    pub fn moe_pairs_scatter(
7497        &self,
7498        y_down: &CudaSlice<f32>,
7499        pair_w: &CudaSlice<f32>,
7500        tok_pair_off: &CudaSlice<i32>,
7501        tok_pair_ids: &CudaSlice<i32>,
7502        moe_out: &mut CudaSlice<f32>,
7503        t: usize,
7504        n_embd: usize,
7505    ) -> Result<(), Box<dyn std::error::Error>> {
7506        let f = self.func("moe_pairs_scatter");
7507        let cfg = LaunchConfig {
7508            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
7509            block_dim: (256, 1, 1),
7510            shared_mem_bytes: 0,
7511        };
7512        let ne = n_embd as i32;
7513        let __s_b = self.gpu.stream();
7514        let mut b = __s_b.launch_builder(&f);
7515        b.arg(y_down)
7516            .arg(pair_w)
7517            .arg(tok_pair_off)
7518            .arg(tok_pair_ids)
7519            .arg(moe_out)
7520            .arg(&ne);
7521        unsafe {
7522            b.launch(cfg)?;
7523        }
7524        Ok(())
7525    }
7526
7527    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
7528    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
7529    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
7530    #[allow(clippy::too_many_arguments)]
7531    pub fn moe_gate_up_gelu8_dev_q8(
7532        &self,
7533        table: &CudaSlice<u64>,
7534        sel: &cudarc::driver::CudaView<i32>,
7535        aq: &CudaSlice<i8>,
7536        ad: &CudaSlice<f32>,
7537        in_f: usize,
7538        n_ff: usize,
7539        n_used: usize,
7540        n_expert: usize,
7541        qt_g: i32,
7542        qt_u: i32,
7543        rb_g: usize,
7544        rb_u: usize,
7545    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7546        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7547        let (inf, nff, ne, rbg, rbu) = (
7548            in_f as i32,
7549            n_ff as i32,
7550            n_expert as i32,
7551            rb_g as i64,
7552            rb_u as i64,
7553        );
7554        let f = self.func("moe_gate_up_gelu8_dev_q8");
7555        let cfg = LaunchConfig {
7556            grid_dim: (n_ff as u32, n_used as u32, 1),
7557            block_dim: (32, 1, 1),
7558            shared_mem_bytes: 0,
7559        };
7560        let __s_b = self.gpu.stream();
7561        let mut b = __s_b.launch_builder(&f);
7562        b.arg(table)
7563            .arg(sel)
7564            .arg(aq)
7565            .arg(ad)
7566            .arg(&mut act)
7567            .arg(&inf)
7568            .arg(&nff)
7569            .arg(&ne)
7570            .arg(&qt_g)
7571            .arg(&qt_u)
7572            .arg(&rbg)
7573            .arg(&rbu);
7574        unsafe {
7575            b.launch(cfg)?;
7576        }
7577        Ok(act)
7578    }
7579
7580    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
7581    #[allow(clippy::too_many_arguments)]
7582    pub fn moe_gate_up_gelu8_dev_q8_rows(
7583        &self,
7584        table: &CudaSlice<u64>,
7585        sel: &CudaSlice<i32>,
7586        aq: &CudaSlice<i8>,
7587        ad: &CudaSlice<f32>,
7588        t: usize,
7589        in_f: usize,
7590        n_ff: usize,
7591        n_used: usize,
7592        n_expert: usize,
7593        qt_g: i32,
7594        qt_u: i32,
7595        rb_g: usize,
7596        rb_u: usize,
7597    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7598        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
7599        let (inf, nff, ne, rbg, rbu, nu) = (
7600            in_f as i32,
7601            n_ff as i32,
7602            n_expert as i32,
7603            rb_g as i64,
7604            rb_u as i64,
7605            n_used as i32,
7606        );
7607        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
7608        let cfg = LaunchConfig {
7609            grid_dim: (n_ff as u32, n_used as u32, t as u32),
7610            block_dim: (32, 1, 1),
7611            shared_mem_bytes: 0,
7612        };
7613        let __s_b = self.gpu.stream();
7614        let mut b = __s_b.launch_builder(&f);
7615        b.arg(table)
7616            .arg(sel)
7617            .arg(aq)
7618            .arg(ad)
7619            .arg(&mut act)
7620            .arg(&inf)
7621            .arg(&nff)
7622            .arg(&ne)
7623            .arg(&qt_g)
7624            .arg(&qt_u)
7625            .arg(&rbg)
7626            .arg(&rbu)
7627            .arg(&nu);
7628        unsafe {
7629            b.launch(cfg)?;
7630        }
7631        Ok(act)
7632    }
7633
7634    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
7635    #[allow(clippy::too_many_arguments)]
7636    pub fn moe_gate_up_gelu8_dev_q8_csr(
7637        &self,
7638        table: &CudaSlice<u64>,
7639        sel: &CudaSlice<i32>,
7640        aq: &CudaSlice<i8>,
7641        ad: &CudaSlice<f32>,
7642        n_pairs: usize,
7643        in_f: usize,
7644        n_ff: usize,
7645        n_used: usize,
7646        n_expert: usize,
7647        qt_g: i32,
7648        qt_u: i32,
7649        rb_g: usize,
7650        rb_u: usize,
7651    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7652        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7653        let (inf, nff, ne, rbg, rbu, nu, npi) = (
7654            in_f as i32,
7655            n_ff as i32,
7656            n_expert as i32,
7657            rb_g as i64,
7658            rb_u as i64,
7659            n_used as i32,
7660            n_pairs as i32,
7661        );
7662        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
7663        let cfg = LaunchConfig {
7664            grid_dim: (n_ff as u32, n_pairs as u32, 1),
7665            block_dim: (32, 1, 1),
7666            shared_mem_bytes: 0,
7667        };
7668        let __s_b = self.gpu.stream();
7669        let mut b = __s_b.launch_builder(&f);
7670        b.arg(table)
7671            .arg(sel)
7672            .arg(aq)
7673            .arg(ad)
7674            .arg(&mut act)
7675            .arg(&inf)
7676            .arg(&nff)
7677            .arg(&ne)
7678            .arg(&qt_g)
7679            .arg(&qt_u)
7680            .arg(&rbg)
7681            .arg(&rbu)
7682            .arg(&nu)
7683            .arg(&npi);
7684        unsafe {
7685            b.launch(cfg)?;
7686        }
7687        Ok(act)
7688    }
7689
7690    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
7691    #[allow(clippy::too_many_arguments)]
7692    pub fn moe_down8_fma_dev_q8_rows_g(
7693        &self,
7694        table: &CudaSlice<u64>,
7695        sel: &CudaSlice<i32>,
7696        w: &CudaSlice<f32>,
7697        aq2: &CudaSlice<i8>,
7698        ad2: &CudaSlice<f32>,
7699        dst: &mut CudaSlice<f32>,
7700        t: usize,
7701        in_f: usize,
7702        out_f: usize,
7703        n_used: usize,
7704        n_expert: usize,
7705        qt: i32,
7706        rb: usize,
7707    ) -> Result<(), Box<dyn std::error::Error>> {
7708        let (inf, outf, nu, ne, rbi) = (
7709            in_f as i32,
7710            out_f as i32,
7711            n_used as i32,
7712            n_expert as i32,
7713            rb as i64,
7714        );
7715        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
7716        // eight warps, then replay the original slot-ordered FMA chain. Every
7717        // other shape retains the generic one-warp rows kernel.
7718        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
7719        let f = self.func(if step_b1_w8 {
7720            "moe_down8_fma_dev_q8_rows_w8"
7721        } else {
7722            "moe_down8_fma_dev_q8_rows_g"
7723        });
7724        let cfg = LaunchConfig {
7725            grid_dim: (out_f as u32, 1, t as u32),
7726            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
7727            shared_mem_bytes: 0,
7728        };
7729        let __s_b = self.gpu.stream();
7730        let mut b = __s_b.launch_builder(&f);
7731        b.arg(table)
7732            .arg(sel)
7733            .arg(w)
7734            .arg(aq2)
7735            .arg(ad2)
7736            .arg(dst)
7737            .arg(&inf)
7738            .arg(&outf)
7739            .arg(&nu)
7740            .arg(&ne)
7741            .arg(&qt)
7742            .arg(&rbi);
7743        unsafe {
7744            b.launch(cfg)?;
7745        }
7746        Ok(())
7747    }
7748
7749    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
7750    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
7751    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
7752    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
7753        let (out_f, in_f) = (2048usize, 2816usize);
7754        let nblk = in_f / 32;
7755        let mut seed = 0x9E3779B97F4A7C15u64;
7756        let mut rng = move || {
7757            seed = seed
7758                .wrapping_mul(6364136223846793005)
7759                .wrapping_add(1442695040888963407);
7760            (seed >> 33) as u8
7761        };
7762        let mut w = vec![0u8; out_f * nblk * 18];
7763        for b in w.iter_mut() {
7764            *b = rng();
7765        }
7766        for r in 0..out_f {
7767            for g in 0..nblk {
7768                let off = (r * nblk + g) * 18;
7769                w[off] = 0x00;
7770                w[off + 1] = 0x2C; // sane half d
7771            }
7772        }
7773        let qplane = out_f * nblk * 16;
7774        let mut wrp = vec![0u8; w.len()];
7775        for r in 0..out_f {
7776            for g in 0..nblk {
7777                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
7778                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
7779                    .copy_from_slice(&src[0..2]);
7780                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
7781            }
7782        }
7783        let w_d = self.htod_bytes(&w)?;
7784        let wrp_d = self.htod_bytes(&wrp)?;
7785        let mut aq = vec![0i8; m * in_f];
7786        for v in aq.iter_mut() {
7787            *v = rng() as i8;
7788        }
7789        let aq_d = self.htod_i8(&aq)?;
7790        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
7791        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
7792        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
7793        const RPB: u32 = 4;
7794        let cfg = LaunchConfig {
7795            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
7796            block_dim: (32, RPB, 1),
7797            shared_mem_bytes: 0,
7798        };
7799        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
7800        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
7801        let fb = self.func("qmatvec_q4_0_mmvq_b4");
7802        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
7803        {
7804            let __s_b = self.gpu.stream();
7805            let mut b = __s_b.launch_builder(&fb);
7806            b.arg(&w_d)
7807                .arg(&aq_d)
7808                .arg(&ad_d)
7809                .arg(&mut y0)
7810                .arg(&inf)
7811                .arg(&outf)
7812                .arg(&mi)
7813                .arg(&rb);
7814            unsafe {
7815                b.launch(cfg)?;
7816            }
7817            let __s_b = self.gpu.stream();
7818            let mut b = __s_b.launch_builder(&fr);
7819            b.arg(&wrp_d)
7820                .arg(&aq_d)
7821                .arg(&ad_d)
7822                .arg(&mut y1)
7823                .arg(&inf)
7824                .arg(&outf)
7825                .arg(&mi)
7826                .arg(&qp);
7827            unsafe {
7828                b.launch(cfg)?;
7829            }
7830        }
7831        self.gpu.stream().synchronize()?;
7832        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
7833        let nd = h0
7834            .iter()
7835            .zip(&h1)
7836            .filter(|(a, b)| a.to_bits() != b.to_bits())
7837            .count();
7838        if nd != 0 {
7839            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
7840        }
7841        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
7842            self.gpu.stream().synchronize()?;
7843            let t0 = std::time::Instant::now();
7844            for _ in 0..500 {
7845                if rp {
7846                    let __s_b = self.gpu.stream();
7847                    let mut b = __s_b.launch_builder(&fr);
7848                    b.arg(&wrp_d)
7849                        .arg(&aq_d)
7850                        .arg(&ad_d)
7851                        .arg(&mut y1)
7852                        .arg(&inf)
7853                        .arg(&outf)
7854                        .arg(&mi)
7855                        .arg(&qp);
7856                    unsafe {
7857                        b.launch(cfg)?;
7858                    }
7859                } else {
7860                    let __s_b = self.gpu.stream();
7861                    let mut b = __s_b.launch_builder(&fb);
7862                    b.arg(&w_d)
7863                        .arg(&aq_d)
7864                        .arg(&ad_d)
7865                        .arg(&mut y0)
7866                        .arg(&inf)
7867                        .arg(&outf)
7868                        .arg(&mi)
7869                        .arg(&rb);
7870                    unsafe {
7871                        b.launch(cfg)?;
7872                    }
7873                }
7874            }
7875            self.gpu.stream().synchronize()?;
7876            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
7877        };
7878        let _ = time(false)?;
7879        let _ = time(true)?; // warm
7880        Ok((time(false)?, time(true)?))
7881    }
7882
7883    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
7884    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
7885    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
7886    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
7887    pub fn build_q4_rp4(
7888        &self,
7889        t: &mut crate::model::GpuTensor,
7890    ) -> Result<(), Box<dyn std::error::Error>> {
7891        use crate::model::GpuTensor;
7892        let GpuTensor::Quant {
7893            bytes,
7894            qtype,
7895            row_bytes,
7896            ne,
7897            rp4,
7898            ..
7899        } = t
7900        else {
7901            return Ok(());
7902        };
7903        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
7904            return Ok(());
7905        }
7906        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7907        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
7908            return Ok(());
7909        }
7910        let nblk = in_f / 32;
7911        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
7912        let f = self.func("q4_0_split_rp_build");
7913        let n = (out_f * nblk) as i32;
7914        let cfg = LaunchConfig {
7915            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
7916            block_dim: (256, 1, 1),
7917            shared_mem_bytes: 0,
7918        };
7919        let (of, nb) = (out_f as i32, nblk as i32);
7920        let _ = n;
7921        let __s_b = self.gpu.stream();
7922        let mut b = __s_b.launch_builder(&f);
7923        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
7924        unsafe {
7925            b.launch(cfg)?;
7926        }
7927        *rp4 = Some(dst);
7928        Ok(())
7929    }
7930
7931    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
7932    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
7933    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
7934    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
7935    pub fn build_q8_rp4(
7936        &self,
7937        t: &mut crate::model::GpuTensor,
7938    ) -> Result<(), Box<dyn std::error::Error>> {
7939        use crate::model::GpuTensor;
7940        let GpuTensor::Quant {
7941            bytes,
7942            qtype,
7943            row_bytes,
7944            ne,
7945            rp4,
7946            ..
7947        } = t
7948        else {
7949            return Ok(());
7950        };
7951        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
7952            return Ok(());
7953        }
7954        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7955        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
7956            return Ok(());
7957        }
7958        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
7959        Ok(())
7960    }
7961
7962    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
7963    /// mirror without a GpuTensor (same kernel the loader path above uses).
7964    pub fn build_q8_rp4_raw(
7965        &self,
7966        bytes: &CudaSlice<u8>,
7967        in_f: usize,
7968        out_f: usize,
7969    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7970        assert!(in_f.is_multiple_of(32));
7971        let nblk = in_f / 32;
7972        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
7973        let f = self.func("q8_0_split_rp_build");
7974        let cfg = LaunchConfig {
7975            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
7976            block_dim: (256, 1, 1),
7977            shared_mem_bytes: 0,
7978        };
7979        let (of, nb) = (out_f as i32, nblk as i32);
7980        let __s_b = self.gpu.stream();
7981        let mut b = __s_b.launch_builder(&f);
7982        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
7983        unsafe {
7984            b.launch(cfg)?;
7985        }
7986        Ok(dst)
7987    }
7988
7989    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
7990    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
7991    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
7992    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
7993    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
7994    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
7995    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
7996    pub fn build_q4k_rp4(
7997        &self,
7998        t: &mut crate::model::GpuTensor,
7999    ) -> Result<(), Box<dyn std::error::Error>> {
8000        use crate::model::GpuTensor;
8001        let GpuTensor::Quant {
8002            bytes,
8003            qtype,
8004            row_bytes,
8005            ne,
8006            rp4,
8007            ..
8008        } = t
8009        else {
8010            return Ok(());
8011        };
8012        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
8013            return Ok(());
8014        }
8015        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
8016        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
8017            return Ok(());
8018        }
8019        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
8020        Ok(())
8021    }
8022
8023    pub fn build_q6k_rp4(
8024        &self,
8025        t: &mut crate::model::GpuTensor,
8026    ) -> Result<(), Box<dyn std::error::Error>> {
8027        use crate::model::GpuTensor;
8028        let GpuTensor::Quant {
8029            bytes,
8030            qtype,
8031            row_bytes,
8032            ne,
8033            rp4,
8034            ..
8035        } = t
8036        else {
8037            return Ok(());
8038        };
8039        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
8040            return Ok(());
8041        }
8042        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
8043        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
8044            return Ok(());
8045        }
8046        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
8047        Ok(())
8048    }
8049
8050    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
8051    pub fn build_kq_rp4_raw(
8052        &self,
8053        bytes: &CudaSlice<u8>,
8054        in_f: usize,
8055        out_f: usize,
8056        qtype: i32,
8057    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8058        assert!(in_f.is_multiple_of(256));
8059        let nsbk = in_f / 256;
8060        let (sb_bytes, kname) = match qtype {
8061            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
8062            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
8063            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
8064        };
8065        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
8066        let f = self.func(kname);
8067        let cfg = LaunchConfig {
8068            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
8069            block_dim: (256, 1, 1),
8070            shared_mem_bytes: 0,
8071        };
8072        let (of, nb) = (out_f as i32, nsbk as i32);
8073        let __s_b = self.gpu.stream();
8074        let mut b = __s_b.launch_builder(&f);
8075        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
8076        unsafe {
8077            b.launch(cfg)?;
8078        }
8079        Ok(dst)
8080    }
8081
8082    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
8083    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
8084    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
8085    pub fn kqrp_enabled() -> bool {
8086        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8087        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
8088            Ok("0") => false,
8089            Ok(_) => true,
8090            Err(_) => cfg!(memra_hopper_mma),
8091        })
8092    }
8093
8094    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
8095    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
8096    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
8097    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
8098    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
8099    pub fn build_q4_rp_swap(
8100        &self,
8101        t: &mut crate::model::GpuTensor,
8102    ) -> Result<bool, Box<dyn std::error::Error>> {
8103        use crate::model::GpuTensor;
8104        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
8105        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
8106        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
8107        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
8108        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
8109        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
8110        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
8111        // this fn's OWN builder serves may ever be swapped; everything else refuses
8112        // here, regardless of walk ordering.
8113        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
8114            return Ok(false);
8115        }
8116        self.build_q4_rp4(t)?;
8117        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
8118        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
8119            return Ok(false);
8120        };
8121        match rp4.take() {
8122            Some(split) => {
8123                *bytes = split; // the GGUF-layout buffer drops here
8124                *rp = true;
8125                Ok(true)
8126            }
8127            None => Ok(false),
8128        }
8129    }
8130
8131    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
8132    pub fn q4rp_enabled() -> bool {
8133        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8134        *ON.get_or_init(|| {
8135            std::env::var("MEMRA_Q4RP")
8136                .map(|v| v != "0")
8137                .unwrap_or(true)
8138        })
8139    }
8140
8141    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
8142    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
8143    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8144    pub fn copy_rows_strided(
8145        &self,
8146        src: &CudaSlice<f32>,
8147        dst: &mut CudaSlice<f32>,
8148        row_elems: usize,
8149        n_rows: usize,
8150        src_stride: usize,
8151        src_off: usize,
8152    ) -> Result<(), Box<dyn std::error::Error>> {
8153        let f = self.func("copy_rows_strided_f32");
8154        let cfg = LaunchConfig {
8155            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
8156            block_dim: (256, 1, 1),
8157            shared_mem_bytes: 0,
8158        };
8159        let (re, nr) = (row_elems as i32, n_rows as i32);
8160        let (st, off) = (src_stride as i64, src_off as i64);
8161        let __s_b = self.gpu.stream();
8162        let mut b = __s_b.launch_builder(&f);
8163        b.arg(src)
8164            .arg(&mut *dst)
8165            .arg(&re)
8166            .arg(&nr)
8167            .arg(&st)
8168            .arg(&off);
8169        unsafe {
8170            b.launch(cfg)?;
8171        }
8172        Ok(())
8173    }
8174
8175    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
8176    ///
8177    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
8178    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
8179    /// one peer copy per token.
8180    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8181    pub fn place_rows_strided(
8182        &self,
8183        src: &CudaSlice<f32>,
8184        dst: &mut CudaSlice<f32>,
8185        row_elems: usize,
8186        n_rows: usize,
8187        dst_stride: usize,
8188        dst_off: usize,
8189    ) -> Result<(), Box<dyn std::error::Error>> {
8190        if row_elems == 0 || n_rows == 0 {
8191            return Err("strided row placement requires nonzero rows and row width".into());
8192        }
8193        let src_len = n_rows
8194            .checked_mul(row_elems)
8195            .ok_or("strided row placement source size overflow")?;
8196        let dst_len = n_rows
8197            .checked_sub(1)
8198            .and_then(|rows| rows.checked_mul(dst_stride))
8199            .and_then(|base| base.checked_add(dst_off))
8200            .and_then(|base| base.checked_add(row_elems))
8201            .ok_or("strided row placement destination size overflow")?;
8202        let row_end = dst_off
8203            .checked_add(row_elems)
8204            .ok_or("strided row placement row size overflow")?;
8205        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
8206            return Err(format!(
8207                "strided row placement geometry mismatch: src={} need_src={src_len} \
8208                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
8209                 dst_stride={dst_stride} dst_off={dst_off}",
8210                src.len(),
8211                dst.len(),
8212            )
8213            .into());
8214        }
8215        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
8216            return Err("strided row placement exceeds CUDA kernel geometry".into());
8217        }
8218        let f = self.func("place_rows_strided_f32");
8219        let cfg = LaunchConfig {
8220            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
8221            block_dim: (256, 1, 1),
8222            shared_mem_bytes: 0,
8223        };
8224        let (re, nr) = (row_elems as i32, n_rows as i32);
8225        let (st, off) = (dst_stride as i64, dst_off as i64);
8226        let __s_b = self.gpu.stream();
8227        let mut b = __s_b.launch_builder(&f);
8228        b.arg(src)
8229            .arg(&mut *dst)
8230            .arg(&re)
8231            .arg(&nr)
8232            .arg(&st)
8233            .arg(&off);
8234        unsafe {
8235            b.launch(cfg)?;
8236        }
8237        Ok(())
8238    }
8239
8240    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
8241    pub fn u32_set_k(
8242        &self,
8243        dst: &mut CudaSlice<u32>,
8244        v: u32,
8245        idx: usize,
8246    ) -> Result<(), Box<dyn std::error::Error>> {
8247        let f = self.func("u32_set_k");
8248        let cfg = LaunchConfig {
8249            grid_dim: (1, 1, 1),
8250            block_dim: (1, 1, 1),
8251            shared_mem_bytes: 0,
8252        };
8253        let ii = idx as i32;
8254        let __s_b = self.gpu.stream();
8255        let mut b = __s_b.launch_builder(&f);
8256        b.arg(dst).arg(&v).arg(&ii);
8257        unsafe {
8258            b.launch(cfg)?;
8259        }
8260        Ok(())
8261    }
8262
8263    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
8264    pub fn i32_add_k(
8265        &self,
8266        d: &mut CudaSlice<i32>,
8267        v: i32,
8268    ) -> Result<(), Box<dyn std::error::Error>> {
8269        let f = self.func("i32_add_k");
8270        let cfg = LaunchConfig {
8271            grid_dim: (1, 1, 1),
8272            block_dim: (32, 1, 1),
8273            shared_mem_bytes: 0,
8274        };
8275        let __s_b = self.gpu.stream();
8276        let mut b = __s_b.launch_builder(&f);
8277        b.arg(d).arg(&v);
8278        unsafe {
8279            b.launch(cfg)?;
8280        }
8281        Ok(())
8282    }
8283
8284    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
8285    pub fn i32_iota_from(
8286        &self,
8287        ctr: &CudaSlice<i32>,
8288        dst: &mut CudaSlice<i32>,
8289        n: usize,
8290    ) -> Result<(), Box<dyn std::error::Error>> {
8291        let f = self.func("i32_iota_from");
8292        let cfg = LaunchConfig::for_num_elems(n as u32);
8293        let ni = n as i32;
8294        let __s_b = self.gpu.stream();
8295        let mut b = __s_b.launch_builder(&f);
8296        b.arg(ctr).arg(dst).arg(&ni);
8297        unsafe {
8298            b.launch(cfg)?;
8299        }
8300        Ok(())
8301    }
8302
8303    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
8304    pub fn u32_map_k(
8305        &self,
8306        buf: &mut CudaSlice<u32>,
8307        map: &CudaSlice<u32>,
8308        idx: usize,
8309    ) -> Result<(), Box<dyn std::error::Error>> {
8310        let f = self.func("u32_map_k");
8311        let cfg = LaunchConfig {
8312            grid_dim: (1, 1, 1),
8313            block_dim: (1, 1, 1),
8314            shared_mem_bytes: 0,
8315        };
8316        let ii = idx as i32;
8317        let __s_b = self.gpu.stream();
8318        let mut b = __s_b.launch_builder(&f);
8319        b.arg(buf).arg(map).arg(&ii);
8320        unsafe {
8321            b.launch(cfg)?;
8322        }
8323        Ok(())
8324    }
8325
8326    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
8327    #[allow(clippy::too_many_arguments)]
8328    pub fn u32_pack2(
8329        &self,
8330        a: &CudaSlice<u32>,
8331        off_a: usize,
8332        n1: usize,
8333        b_in: &CudaSlice<u32>,
8334        n2: usize,
8335        out: &mut CudaSlice<u32>,
8336    ) -> Result<(), Box<dyn std::error::Error>> {
8337        let f = self.func("u32_pack2");
8338        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
8339        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
8340        let __s_b = self.gpu.stream();
8341        let mut b = __s_b.launch_builder(&f);
8342        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
8343        unsafe {
8344            b.launch(cfg)?;
8345        }
8346        Ok(())
8347    }
8348
8349    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
8350    pub fn moe_w_exscale(
8351        &self,
8352        w: &mut CudaSlice<f32>,
8353        sel: &CudaSlice<i32>,
8354        s: &CudaSlice<f32>,
8355        n: usize,
8356    ) -> Result<(), Box<dyn std::error::Error>> {
8357        let f = self.func("moe_w_exscale");
8358        let cfg = LaunchConfig::for_num_elems(n as u32);
8359        let ni = n as i32;
8360        let __s_b = self.gpu.stream();
8361        let mut b = __s_b.launch_builder(&f);
8362        b.arg(w).arg(sel).arg(s).arg(&ni);
8363        unsafe {
8364            b.launch(cfg)?;
8365        }
8366        Ok(())
8367    }
8368
8369    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
8370    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
8371    pub fn moe_w_scale_by_expert(
8372        &self,
8373        w: &mut CudaSlice<f32>,
8374        sel: &CudaSlice<i32>,
8375        macros: &CudaSlice<f32>,
8376        n_expert: usize,
8377        n: usize,
8378    ) -> Result<(), Box<dyn std::error::Error>> {
8379        let f = self.func("moe_w_scale_by_expert");
8380        let cfg = LaunchConfig {
8381            grid_dim: (n.div_ceil(64) as u32, 1, 1),
8382            block_dim: (64, 1, 1),
8383            shared_mem_bytes: 0,
8384        };
8385        let (ne, nn) = (n_expert as i32, n as i32);
8386        let __s_b = self.gpu.stream();
8387        let mut b = __s_b.launch_builder(&f);
8388        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
8389        unsafe {
8390            b.launch(cfg)?;
8391        }
8392        Ok(())
8393    }
8394
8395    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
8396    pub fn moe_gate_up_silu8_dev_q8(
8397        &self,
8398        table: &CudaSlice<u64>,
8399        sel: &cudarc::driver::CudaView<i32>,
8400        aq: &CudaSlice<i8>,
8401        ad: &CudaSlice<f32>,
8402        in_f: usize,
8403        n_ff: usize,
8404        n_used: usize,
8405        n_expert: usize,
8406        qt_g: i32,
8407        qt_u: i32,
8408        rb_g: usize,
8409        rb_u: usize,
8410        macros: &CudaSlice<f32>,
8411    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8412        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
8413        let (mode, wpb) = GU.get_or_init(|| {
8414            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
8415            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
8416                .ok()
8417                .and_then(|v| v.parse().ok())
8418                .unwrap_or(4u32)
8419                .clamp(1, 16);
8420            (mode, wpb)
8421        });
8422        let (mode, wpb) = (mode.as_str(), *wpb);
8423        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8424        let (inf, nff, ne, rbg, rbu) = (
8425            in_f as i32,
8426            n_ff as i32,
8427            n_expert as i32,
8428            rb_g as i64,
8429            rb_u as i64,
8430        );
8431        let (f, cfg) = match mode {
8432            "1" | "2" | "4" => {
8433                let rpw: u32 = mode.parse().unwrap();
8434                let f = self.func(match rpw {
8435                    1 => "moe_gate_up_silu8_dev_q8_r1",
8436                    2 => "moe_gate_up_silu8_dev_q8_r2",
8437                    _ => "moe_gate_up_silu8_dev_q8_r4",
8438                });
8439                let rows_per_block = (rpw * wpb) as usize;
8440                let gx = n_ff.div_ceil(rows_per_block) as u32;
8441                (
8442                    f,
8443                    LaunchConfig {
8444                        grid_dim: (gx, n_used as u32, 1),
8445                        block_dim: (32, wpb, 1),
8446                        shared_mem_bytes: 0,
8447                    },
8448                )
8449            }
8450            "j8" if n_used <= 32 => (
8451                self.func("moe_gate_up_silu8_dev_q8_j8"),
8452                LaunchConfig {
8453                    grid_dim: (n_ff as u32, 1, 1),
8454                    block_dim: (32, n_used as u32, 1),
8455                    shared_mem_bytes: 0,
8456                },
8457            ),
8458            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
8459            "vsm2" => {
8460                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
8461                let sh = (rb_g + rb_u) as u32;
8462                use cudarc::driver::sys::CUfunction_attribute_enum as A;
8463                f.set_attribute(
8464                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
8465                    sh as i32,
8466                )?;
8467                (
8468                    f,
8469                    LaunchConfig {
8470                        grid_dim: (n_ff as u32, n_used as u32, 1),
8471                        block_dim: (32, 1, 1),
8472                        shared_mem_bytes: sh,
8473                    },
8474                )
8475            }
8476            "vsm" => {
8477                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
8478                let sh = (rb_g + rb_u) as u32;
8479                use cudarc::driver::sys::CUfunction_attribute_enum as A;
8480                f.set_attribute(
8481                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
8482                    sh as i32,
8483                )?;
8484                (
8485                    f,
8486                    LaunchConfig {
8487                        grid_dim: (n_ff as u32, n_used as u32, 1),
8488                        block_dim: (32, 1, 1),
8489                        shared_mem_bytes: sh,
8490                    },
8491                )
8492            }
8493            "sg" => (
8494                self.func("moe_gate_up_silu8_dev_q8_sg"),
8495                LaunchConfig {
8496                    grid_dim: (n_ff as u32, n_used as u32, 1),
8497                    block_dim: (32, 1, 1),
8498                    shared_mem_bytes: 0,
8499                },
8500            ),
8501            "j8sg" if n_used <= 32 => (
8502                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
8503                LaunchConfig {
8504                    grid_dim: (n_ff as u32, 1, 1),
8505                    block_dim: (32, n_used as u32, 1),
8506                    shared_mem_bytes: 0,
8507                },
8508            ),
8509            "u64" if in_f == 2048 => (
8510                self.func("moe_gate_up_silu8_dev_q8_u64"),
8511                LaunchConfig {
8512                    grid_dim: (n_ff as u32, n_used as u32, 1),
8513                    block_dim: (32, 1, 1),
8514                    shared_mem_bytes: 0,
8515                },
8516            ),
8517            "gs4" if in_f == 2048 => (
8518                self.func("moe_gate_up_silu8_dev_q8_gs4"),
8519                LaunchConfig {
8520                    grid_dim: (n_ff as u32, n_used as u32, 1),
8521                    block_dim: (32, 4, 1),
8522                    shared_mem_bytes: 0,
8523                },
8524            ),
8525            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
8526            "v" | "" => (
8527                self.func("moe_gate_up_silu8_dev_q8_v"),
8528                LaunchConfig {
8529                    grid_dim: (n_ff as u32, n_used as u32, 1),
8530                    block_dim: (32, 1, 1),
8531                    shared_mem_bytes: 0,
8532                },
8533            ),
8534            "s2" => (
8535                self.func("moe_gate_up_silu8_dev_q8_s2"),
8536                LaunchConfig {
8537                    grid_dim: (n_ff as u32, n_used as u32, 1),
8538                    block_dim: (32, 2, 1),
8539                    shared_mem_bytes: 0,
8540                },
8541            ),
8542            "s2z" => {
8543                let rz = wpb.min(16); // s2z smem tile is [16][2]
8544                (
8545                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
8546                    LaunchConfig {
8547                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
8548                        block_dim: (32, 2, rz),
8549                        shared_mem_bytes: 0,
8550                    },
8551                )
8552            }
8553            _ => (
8554                self.func("moe_gate_up_silu8_dev_q8"),
8555                LaunchConfig {
8556                    grid_dim: (n_ff as u32, n_used as u32, 1),
8557                    block_dim: (32, 1, 1),
8558                    shared_mem_bytes: 0,
8559                },
8560            ),
8561        };
8562        let __s_b = self.gpu.stream();
8563        let mut b = __s_b.launch_builder(&f);
8564        b.arg(table)
8565            .arg(sel)
8566            .arg(aq)
8567            .arg(ad)
8568            .arg(&mut act)
8569            .arg(&inf)
8570            .arg(&nff)
8571            .arg(&ne)
8572            .arg(&qt_g)
8573            .arg(&qt_u)
8574            .arg(&rbg)
8575            .arg(&rbu)
8576            .arg(macros);
8577        unsafe {
8578            b.launch(cfg)?;
8579        }
8580        Ok(act)
8581    }
8582
8583    #[allow(clippy::too_many_arguments)]
8584    pub fn moe_down8_fma_dev_q8(
8585        &self,
8586        table: &CudaSlice<u64>,
8587        sel: &cudarc::driver::CudaView<i32>,
8588        w: &cudarc::driver::CudaView<f32>,
8589        aq2: &CudaSlice<i8>,
8590        ad2: &CudaSlice<f32>,
8591        dst: &mut cudarc::driver::CudaViewMut<f32>,
8592        in_f: usize,
8593        out_f: usize,
8594        n_used: usize,
8595        n_expert: usize,
8596        qt: i32,
8597        rb: usize,
8598    ) -> Result<(), Box<dyn std::error::Error>> {
8599        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
8600        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
8601        let (inf, outf, nu, ne, rbi) = (
8602            in_f as i32,
8603            out_f as i32,
8604            n_used as i32,
8605            n_expert as i32,
8606            rb as i64,
8607        );
8608        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
8609        // the h2 twins are nsb==16 (in_f==512) shape-gated.
8610        let (f, cfg) = match mode.as_str() {
8611            m @ ("1" | "2" | "4") if n_used <= 8 => {
8612                let rpw: usize = m.parse().unwrap();
8613                let f = self.func(match rpw {
8614                    1 => "moe_down8_fma_dev_q8_w8r1",
8615                    2 => "moe_down8_fma_dev_q8_w8r2",
8616                    _ => "moe_down8_fma_dev_q8_w8r4",
8617                });
8618                (
8619                    f,
8620                    LaunchConfig {
8621                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
8622                        block_dim: (32, n_used as u32, 1),
8623                        shared_mem_bytes: 0,
8624                    },
8625                )
8626            }
8627            "h2" if in_f == 512 => (
8628                self.func("moe_down8_fma_dev_q8_h2"),
8629                LaunchConfig {
8630                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8631                    block_dim: (32, 1, 1),
8632                    shared_mem_bytes: 0,
8633                },
8634            ),
8635            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
8636            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
8637            "" if in_f == 704 && n_used <= 8 => (
8638                self.func("moe_down8_fma_dev_q8_w8r2"),
8639                LaunchConfig {
8640                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8641                    block_dim: (32, n_used as u32, 1),
8642                    shared_mem_bytes: 0,
8643                },
8644            ),
8645            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
8646            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
8647            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
8648            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
8649                self.func("moe_down8_fma_dev_q8_w8h2v"),
8650                LaunchConfig {
8651                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8652                    block_dim: (32, n_used as u32, 1),
8653                    shared_mem_bytes: 0,
8654                },
8655            ),
8656            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
8657                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
8658                LaunchConfig {
8659                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8660                    block_dim: (32, n_used as u32, 1),
8661                    shared_mem_bytes: 0,
8662                },
8663            ),
8664            "w8h2r2" if in_f == 512 && n_used <= 8 => (
8665                self.func("moe_down8_fma_dev_q8_w8h2r2"),
8666                LaunchConfig {
8667                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8668                    block_dim: (32, n_used as u32, 1),
8669                    shared_mem_bytes: 0,
8670                },
8671            ),
8672            "w8h2" if in_f == 512 && n_used <= 8 => (
8673                self.func("moe_down8_fma_dev_q8_w8h2"),
8674                LaunchConfig {
8675                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8676                    block_dim: (32, n_used as u32, 1),
8677                    shared_mem_bytes: 0,
8678                },
8679            ),
8680            _ => (
8681                self.func("moe_down8_fma_dev_q8"),
8682                LaunchConfig {
8683                    grid_dim: (out_f as u32, 1, 1),
8684                    block_dim: (32, 1, 1),
8685                    shared_mem_bytes: 0,
8686                },
8687            ),
8688        };
8689        let __s_b = self.gpu.stream();
8690        let mut b = __s_b.launch_builder(&f);
8691        b.arg(table)
8692            .arg(sel)
8693            .arg(w)
8694            .arg(aq2)
8695            .arg(ad2)
8696            .arg(dst)
8697            .arg(&inf)
8698            .arg(&outf)
8699            .arg(&nu)
8700            .arg(&ne)
8701            .arg(&qt)
8702            .arg(&rbi);
8703        unsafe {
8704            b.launch(cfg)?;
8705        }
8706        Ok(())
8707    }
8708
8709    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
8710    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
8711    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
8712    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
8713    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
8714    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
8715    #[allow(clippy::too_many_arguments)]
8716    pub fn moe_gate_up_silu8_dev_q8_rows(
8717        &self,
8718        table: &CudaSlice<u64>,
8719        sel: &CudaSlice<i32>,
8720        aq: &CudaSlice<i8>,
8721        ad: &CudaSlice<f32>,
8722        t: usize,
8723        in_f: usize,
8724        n_ff: usize,
8725        n_used: usize,
8726        n_expert: usize,
8727        qt_g: i32,
8728        qt_u: i32,
8729        rb_g: usize,
8730        rb_u: usize,
8731        macros: &CudaSlice<f32>,
8732    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8733        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
8734        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
8735        let cfg = LaunchConfig {
8736            grid_dim: (n_ff as u32, n_used as u32, t as u32),
8737            block_dim: (32, 1, 1),
8738            shared_mem_bytes: 0,
8739        };
8740        let (inf, nff, ne, nu, rbg, rbu) = (
8741            in_f as i32,
8742            n_ff as i32,
8743            n_expert as i32,
8744            n_used as i32,
8745            rb_g as i64,
8746            rb_u as i64,
8747        );
8748        let __s_b = self.gpu.stream();
8749        let mut b = __s_b.launch_builder(&f);
8750        b.arg(table)
8751            .arg(sel)
8752            .arg(aq)
8753            .arg(ad)
8754            .arg(&mut act)
8755            .arg(&inf)
8756            .arg(&nff)
8757            .arg(&ne)
8758            .arg(&qt_g)
8759            .arg(&qt_u)
8760            .arg(&rbg)
8761            .arg(&rbu)
8762            .arg(&nu)
8763            .arg(macros);
8764        unsafe {
8765            b.launch(cfg)?;
8766        }
8767        Ok(act)
8768    }
8769
8770    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
8771    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
8772    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
8773    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
8774    #[allow(clippy::too_many_arguments)]
8775    pub fn moe_down8_fma_dev_q8_rows(
8776        &self,
8777        table: &CudaSlice<u64>,
8778        sel: &CudaSlice<i32>,
8779        w: &CudaSlice<f32>,
8780        aq2: &CudaSlice<i8>,
8781        ad2: &CudaSlice<f32>,
8782        dst: &mut CudaSlice<f32>,
8783        t: usize,
8784        in_f: usize,
8785        out_f: usize,
8786        n_used: usize,
8787        n_expert: usize,
8788        qt: i32,
8789        rb: usize,
8790    ) -> Result<(), Box<dyn std::error::Error>> {
8791        assert!(
8792            in_f == 512 && n_used <= 8,
8793            "down rows twin is w8h2v shape-gated"
8794        );
8795        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
8796        let cfg = LaunchConfig {
8797            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
8798            block_dim: (32, n_used as u32, 1),
8799            shared_mem_bytes: 0,
8800        };
8801        let (inf, outf, nu, ne, rbi) = (
8802            in_f as i32,
8803            out_f as i32,
8804            n_used as i32,
8805            n_expert as i32,
8806            rb as i64,
8807        );
8808        let __s_b = self.gpu.stream();
8809        let mut b = __s_b.launch_builder(&f);
8810        b.arg(table)
8811            .arg(sel)
8812            .arg(w)
8813            .arg(aq2)
8814            .arg(ad2)
8815            .arg(dst)
8816            .arg(&inf)
8817            .arg(&outf)
8818            .arg(&nu)
8819            .arg(&ne)
8820            .arg(&qt)
8821            .arg(&rbi);
8822        unsafe {
8823            b.launch(cfg)?;
8824        }
8825        Ok(())
8826    }
8827
8828    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
8829    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
8830    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
8831    #[allow(clippy::too_many_arguments)]
8832    pub fn moe_gate_up_silu8_dev_q8_csr(
8833        &self,
8834        table: &CudaSlice<u64>,
8835        sel: &CudaSlice<i32>,
8836        aq: &CudaSlice<i8>,
8837        ad: &CudaSlice<f32>,
8838        n_pairs: usize,
8839        in_f: usize,
8840        n_ff: usize,
8841        n_used: usize,
8842        n_expert: usize,
8843        qt_g: i32,
8844        qt_u: i32,
8845        rb_g: usize,
8846        rb_u: usize,
8847    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8848        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
8849        // host gate guarantees qt_g == qt_u within a supported class.
8850        let f = if qt_g == crate::QT_NVFP4 {
8851            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
8852        } else {
8853            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
8854        };
8855        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
8856        let cfg = LaunchConfig {
8857            grid_dim: (n_ff as u32, n_pairs as u32, 1),
8858            block_dim: (32, 1, 1),
8859            shared_mem_bytes: 0,
8860        };
8861        let (inf, nff, ne, nu, npi, rbg, rbu) = (
8862            in_f as i32,
8863            n_ff as i32,
8864            n_expert as i32,
8865            n_used as i32,
8866            n_pairs as i32,
8867            rb_g as i64,
8868            rb_u as i64,
8869        );
8870        let __s_b = self.gpu.stream();
8871        let mut b = __s_b.launch_builder(&f);
8872        b.arg(table)
8873            .arg(sel)
8874            .arg(aq)
8875            .arg(ad)
8876            .arg(&mut act)
8877            .arg(&inf)
8878            .arg(&nff)
8879            .arg(&ne)
8880            .arg(&qt_g)
8881            .arg(&qt_u)
8882            .arg(&rbg)
8883            .arg(&rbu)
8884            .arg(&nu)
8885            .arg(&npi);
8886        unsafe {
8887            b.launch(cfg)?;
8888        }
8889        Ok(act)
8890    }
8891
8892    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
8893    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
8894    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
8895    #[allow(clippy::too_many_arguments)]
8896    pub fn moe_down8_fma_dev_q8_variant(
8897        &self,
8898        variant: &str,
8899        table: &CudaSlice<u64>,
8900        sel: &cudarc::driver::CudaView<i32>,
8901        w: &cudarc::driver::CudaView<f32>,
8902        aq2: &CudaSlice<i8>,
8903        ad2: &CudaSlice<f32>,
8904        dst: &mut cudarc::driver::CudaViewMut<f32>,
8905        in_f: usize,
8906        out_f: usize,
8907        n_used: usize,
8908        n_expert: usize,
8909        qt: i32,
8910        rb: usize,
8911    ) -> Result<(), Box<dyn std::error::Error>> {
8912        let (inf, outf, nu, ne, rbi) = (
8913            in_f as i32,
8914            out_f as i32,
8915            n_used as i32,
8916            n_expert as i32,
8917            rb as i64,
8918        );
8919        let (f, cfg) = match variant {
8920            "w8h2" | "w8h2v" => (
8921                self.func(if variant == "w8h2" {
8922                    "moe_down8_fma_dev_q8_w8h2"
8923                } else {
8924                    "moe_down8_fma_dev_q8_w8h2v"
8925                }),
8926                LaunchConfig {
8927                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8928                    block_dim: (32, n_used as u32, 1),
8929                    shared_mem_bytes: 0,
8930                },
8931            ),
8932            "w8h2r2" | "w8h2r2v" => (
8933                self.func(if variant == "w8h2r2" {
8934                    "moe_down8_fma_dev_q8_w8h2r2"
8935                } else {
8936                    "moe_down8_fma_dev_q8_w8h2r2v"
8937                }),
8938                LaunchConfig {
8939                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8940                    block_dim: (32, n_used as u32, 1),
8941                    shared_mem_bytes: 0,
8942                },
8943            ),
8944            _ => (
8945                self.func("moe_down8_fma_dev_q8"),
8946                LaunchConfig {
8947                    grid_dim: (out_f as u32, 1, 1),
8948                    block_dim: (32, 1, 1),
8949                    shared_mem_bytes: 0,
8950                },
8951            ),
8952        };
8953        let __s_b = self.gpu.stream();
8954        let mut b = __s_b.launch_builder(&f);
8955        b.arg(table)
8956            .arg(sel)
8957            .arg(w)
8958            .arg(aq2)
8959            .arg(ad2)
8960            .arg(dst)
8961            .arg(&inf)
8962            .arg(&outf)
8963            .arg(&nu)
8964            .arg(&ne)
8965            .arg(&qt)
8966            .arg(&rbi);
8967        unsafe {
8968            b.launch(cfg)?;
8969        }
8970        Ok(())
8971    }
8972
8973    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
8974    #[allow(clippy::too_many_arguments)]
8975    pub fn moe_gate_up_silu8_dev_q8_variant(
8976        &self,
8977        variant: &str,
8978        table: &CudaSlice<u64>,
8979        sel: &cudarc::driver::CudaView<i32>,
8980        aq: &CudaSlice<i8>,
8981        ad: &CudaSlice<f32>,
8982        in_f: usize,
8983        n_ff: usize,
8984        n_used: usize,
8985        n_expert: usize,
8986        qt_g: i32,
8987        qt_u: i32,
8988        rb_g: usize,
8989        rb_u: usize,
8990    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8991        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8992        let (inf, nff, ne, rbg, rbu) = (
8993            in_f as i32,
8994            n_ff as i32,
8995            n_expert as i32,
8996            rb_g as i64,
8997            rb_u as i64,
8998        );
8999        let f = self.func(if variant == "v" {
9000            "moe_gate_up_silu8_dev_q8_v"
9001        } else {
9002            "moe_gate_up_silu8_dev_q8"
9003        });
9004        let cfg = LaunchConfig {
9005            grid_dim: (n_ff as u32, n_used as u32, 1),
9006            block_dim: (32, 1, 1),
9007            shared_mem_bytes: 0,
9008        };
9009        let __s_b = self.gpu.stream();
9010        let mut b = __s_b.launch_builder(&f);
9011        b.arg(table)
9012            .arg(sel)
9013            .arg(aq)
9014            .arg(ad)
9015            .arg(&mut act)
9016            .arg(&inf)
9017            .arg(&nff)
9018            .arg(&ne)
9019            .arg(&qt_g)
9020            .arg(&qt_u)
9021            .arg(&rbg)
9022            .arg(&rbu);
9023        unsafe {
9024            b.launch(cfg)?;
9025        }
9026        Ok(act)
9027    }
9028
9029    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9030    pub fn moe_gate_up_silu8_dev(
9031        &self,
9032        table: &CudaSlice<u64>,
9033        sel: &cudarc::driver::CudaView<i32>,
9034        x: &cudarc::driver::CudaView<f32>,
9035        in_f: usize,
9036        n_ff: usize,
9037        n_used: usize,
9038        n_expert: usize,
9039        qt_g: i32,
9040        qt_u: i32,
9041        rb_g: usize,
9042        rb_u: usize,
9043        macros: &CudaSlice<f32>,
9044    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9045        let f = self.func("moe_gate_up_silu8_dev");
9046        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
9047        let cfg = LaunchConfig {
9048            grid_dim: (n_ff as u32, n_used as u32, 1),
9049            block_dim: (256, 1, 1),
9050            shared_mem_bytes: 0,
9051        };
9052        let (inf, nff, ne, rbg, rbu) = (
9053            in_f as i32,
9054            n_ff as i32,
9055            n_expert as i32,
9056            rb_g as i64,
9057            rb_u as i64,
9058        );
9059        let __s_b = self.gpu.stream();
9060        let mut b = __s_b.launch_builder(&f);
9061        b.arg(table)
9062            .arg(sel)
9063            .arg(x)
9064            .arg(&mut act)
9065            .arg(&inf)
9066            .arg(&nff)
9067            .arg(&ne)
9068            .arg(&qt_g)
9069            .arg(&qt_u)
9070            .arg(&rbg)
9071            .arg(&rbu)
9072            .arg(macros);
9073        unsafe {
9074            b.launch(cfg)?;
9075        }
9076        Ok(act)
9077    }
9078
9079    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
9080    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
9081    #[allow(clippy::too_many_arguments)]
9082    pub fn moe_down8_fma_dev(
9083        &self,
9084        table: &CudaSlice<u64>,
9085        sel: &cudarc::driver::CudaView<i32>,
9086        w: &cudarc::driver::CudaView<f32>,
9087        act: &CudaSlice<f32>,
9088        dst: &mut cudarc::driver::CudaViewMut<f32>,
9089        in_f: usize,
9090        out_f: usize,
9091        n_used: usize,
9092        n_expert: usize,
9093        qt: i32,
9094        rb: usize,
9095    ) -> Result<(), Box<dyn std::error::Error>> {
9096        let f = self.func("moe_down8_fma_dev");
9097        let cfg = LaunchConfig {
9098            grid_dim: (out_f as u32, 1, 1),
9099            block_dim: (256, 1, 1),
9100            shared_mem_bytes: 0,
9101        };
9102        let (inf, outf, nu, ne, rbv) = (
9103            in_f as i32,
9104            out_f as i32,
9105            n_used as i32,
9106            n_expert as i32,
9107            rb as i64,
9108        );
9109        let __s_b = self.gpu.stream();
9110        let mut b = __s_b.launch_builder(&f);
9111        b.arg(table)
9112            .arg(sel)
9113            .arg(w)
9114            .arg(act)
9115            .arg(dst)
9116            .arg(&inf)
9117            .arg(&outf)
9118            .arg(&nu)
9119            .arg(&ne)
9120            .arg(&qt)
9121            .arg(&rbv);
9122        unsafe {
9123            b.launch(cfg)?;
9124        }
9125        Ok(())
9126    }
9127
9128    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
9129    pub fn axpy_into(
9130        &self,
9131        src: &CudaSlice<f32>,
9132        alpha: f32,
9133        dst: &mut cudarc::driver::CudaViewMut<f32>,
9134        n: usize,
9135    ) -> Result<(), Box<dyn std::error::Error>> {
9136        let f = self.func("axpy_f32");
9137        let cfg = LaunchConfig::for_num_elems(n as u32);
9138        let (a, ni) = (alpha, n as i32);
9139        let __s_b = self.gpu.stream();
9140        let mut b = __s_b.launch_builder(&f);
9141        b.arg(src).arg(dst).arg(&a).arg(&ni);
9142        unsafe {
9143            b.launch(cfg)?;
9144        }
9145        Ok(())
9146    }
9147
9148    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
9149    pub fn axpy_host_into(
9150        &self,
9151        src: &cudarc::driver::CudaView<'_, f32>,
9152        alpha: f32,
9153        dst: &mut cudarc::driver::CudaViewMut<f32>,
9154        n: usize,
9155    ) -> Result<(), Box<dyn std::error::Error>> {
9156        let f = self.func("axpy_host_f32");
9157        let cfg = LaunchConfig::for_num_elems(n as u32);
9158        let (a, ni) = (alpha, n as i32);
9159        let __s_b = self.gpu.stream();
9160        let mut b = __s_b.launch_builder(&f);
9161        b.arg(src).arg(dst).arg(&a).arg(&ni);
9162        unsafe {
9163            b.launch(cfg)?;
9164        }
9165        Ok(())
9166    }
9167
9168    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
9169    pub fn add_scaled_rows(
9170        &self,
9171        src: &CudaSlice<f32>,
9172        scale: &CudaSlice<f32>,
9173        dst: &mut CudaSlice<f32>,
9174        ncols: usize,
9175        nrows: usize,
9176    ) -> Result<(), Box<dyn std::error::Error>> {
9177        let f = self.func("add_scaled_rows_f32");
9178        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9179        let (nc, nr) = (ncols as i32, nrows as i32);
9180        let __s_b = self.gpu.stream();
9181        let mut b = __s_b.launch_builder(&f);
9182        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
9183        unsafe {
9184            b.launch(cfg)?;
9185        }
9186        Ok(())
9187    }
9188
9189    /// `add_scaled_rows` with an all-ones scale drawn from the resident ones buffer (door H,
9190    /// `MEMRA_HTOD_DIET`) — the UNGATED shared-expert add, without re-uploading the
9191    /// constant every MoE layer-call. Same kernel, same values: the buffer may be longer than
9192    /// `nrows` because `add_scaled_rows_f32` reads only `scale[0..nrows]`.
9193    pub fn add_scaled_rows_ones(
9194        &self,
9195        src: &CudaSlice<f32>,
9196        dst: &mut CudaSlice<f32>,
9197        ncols: usize,
9198        nrows: usize,
9199    ) -> Result<(), Box<dyn std::error::Error>> {
9200        let mut guard = self
9201            .shexp_ones
9202            .lock()
9203            .map_err(|_| "shexp ones buffer is poisoned")?;
9204        if guard.as_ref().map(|b| b.len() < nrows).unwrap_or(true) {
9205            // One upload per process (or per growth step): the serving shapes are t <= 8 for the
9206            // verify walk and the prime's chunk width otherwise.
9207            *guard = Some(self.htod(&vec![1.0f32; nrows.max(64)])?);
9208        }
9209        let ones = guard.as_ref().expect("just ensured");
9210        let f = self.func("add_scaled_rows_f32");
9211        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9212        let (nc, nr) = (ncols as i32, nrows as i32);
9213        let __s_b = self.gpu.stream();
9214        let mut b = __s_b.launch_builder(&f);
9215        b.arg(src).arg(ones).arg(&mut *dst).arg(&nc).arg(&nr);
9216        unsafe {
9217            b.launch(cfg)?;
9218        }
9219        Ok(())
9220    }
9221
9222    /// The `len_d` i32 mirror store, door H aware (`MEMRA_HTOD_DIET`): the async
9223    /// [`Self::i32_set_k`] launch when the door is on, else the shipped synchronizing pageable
9224    /// `memcpy_htod`. Identical value into the identical slot, both stream-ordered.
9225    pub fn i32_mirror_store(
9226        &self,
9227        dst: &mut CudaSlice<i32>,
9228        v: i32,
9229    ) -> Result<(), Box<dyn std::error::Error>> {
9230        if crate::htod_diet_on() {
9231            HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9232            return self.i32_set_k(dst, v);
9233        }
9234        self.gpu.stream().memcpy_htod(&[v], dst)?;
9235        Ok(())
9236    }
9237
9238    /// y[r, :] *= s[r] in place (per-CSR-row macro scale for the grouped prime's gate/up —
9239    /// silu is nonlinear, so per-expert NVFP4 macros must land before it).
9240    pub fn scale_rows(
9241        &self,
9242        y: &mut CudaSlice<f32>,
9243        s: &CudaSlice<f32>,
9244        ncols: usize,
9245        nrows: usize,
9246    ) -> Result<(), Box<dyn std::error::Error>> {
9247        let f = self.func("scale_rows_f32");
9248        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9249        let (nc, nr) = (ncols as i32, nrows as i32);
9250        let __s_b = self.gpu.stream();
9251        let mut b = __s_b.launch_builder(&f);
9252        b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
9253        unsafe {
9254            b.launch(cfg)?;
9255        }
9256        Ok(())
9257    }
9258
9259    /// Fused grouped-prime tail: join both rank partials (canonical shard order), permute
9260    /// CSR->pair via `inv`, weight, and scatter to tokens in one pass — replaces
9261    /// rows_permute + add + scatter and the three large temporaries they needed.
9262    #[allow(clippy::too_many_arguments)]
9263    pub fn moe_prime_join_scatter(
9264        &self,
9265        y0: &CudaSlice<f32>,
9266        y1: &CudaSlice<f32>,
9267        inv: &CudaSlice<i32>,
9268        w: &CudaSlice<f32>,
9269        out: &mut CudaSlice<f32>,
9270        ncols: usize,
9271        n_used: usize,
9272        t: usize,
9273    ) -> Result<(), Box<dyn std::error::Error>> {
9274        let f = self.func("moe_prime_join_scatter_f32");
9275        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9276        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9277        let __s_b = self.gpu.stream();
9278        let mut b = __s_b.launch_builder(&f);
9279        b.arg(y0)
9280            .arg(y1)
9281            .arg(inv)
9282            .arg(w)
9283            .arg(&mut *out)
9284            .arg(&nc)
9285            .arg(&nu)
9286            .arg(&ti);
9287        unsafe {
9288            b.launch(cfg)?;
9289        }
9290        Ok(())
9291    }
9292
9293    /// out[t, :] += sum_j w[t*n_used+j] * y[t*n_used+j, :], the j-sum sequential per thread —
9294    /// a pinned per-token reduction order, never atomics (the grouped prime's scatter).
9295    pub fn moe_pairs_weighted_scatter(
9296        &self,
9297        y: &CudaSlice<f32>,
9298        w: &CudaSlice<f32>,
9299        out: &mut CudaSlice<f32>,
9300        ncols: usize,
9301        n_used: usize,
9302        t: usize,
9303    ) -> Result<(), Box<dyn std::error::Error>> {
9304        let f = self.func("moe_pairs_weighted_scatter_f32");
9305        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9306        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9307        let __s_b = self.gpu.stream();
9308        let mut b = __s_b.launch_builder(&f);
9309        b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
9310        unsafe {
9311            b.launch(cfg)?;
9312        }
9313        Ok(())
9314    }
9315
9316    // ======== A2 GROUPED MoE PREFILL KERNELS ========
9317
9318    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
9319    pub fn gather_rows(
9320        &self,
9321        src: &CudaSlice<f32>,
9322        idx: &CudaSlice<i32>,
9323        dst: &mut CudaSlice<f32>,
9324        ncols: usize,
9325        m_e: usize,
9326    ) -> Result<(), Box<dyn std::error::Error>> {
9327        let f = self.func("gather_rows_f32");
9328        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
9329        let (nc, me) = (ncols as i32, m_e as i32);
9330        let __s_b = self.gpu.stream();
9331        let mut b = __s_b.launch_builder(&f);
9332        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
9333        unsafe {
9334            b.launch(cfg)?;
9335        }
9336        Ok(())
9337    }
9338
9339    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
9340    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
9341    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
9342    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
9343    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9344    pub fn scatter_slot(
9345        &self,
9346        src: &CudaSlice<f32>,
9347        tok_idx: &CudaSlice<i32>,
9348        slot_idx: &CudaSlice<i32>,
9349        weight: &CudaSlice<f32>,
9350        dst: &mut CudaSlice<f32>,
9351        wbuf: &mut CudaSlice<f32>,
9352        ncols: usize,
9353        n_used: usize,
9354        m_e: usize,
9355    ) -> Result<(), Box<dyn std::error::Error>> {
9356        let f = self.func("scatter_add_slot_f32");
9357        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
9358        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
9359        let __s_b = self.gpu.stream();
9360        let mut b = __s_b.launch_builder(&f);
9361        b.arg(src)
9362            .arg(tok_idx)
9363            .arg(slot_idx)
9364            .arg(weight)
9365            .arg(dst)
9366            .arg(wbuf)
9367            .arg(&nc)
9368            .arg(&nu)
9369            .arg(&me);
9370        unsafe {
9371            b.launch(cfg)?;
9372        }
9373        Ok(())
9374    }
9375
9376    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
9377    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
9378    /// Uses FMA for bit-identity with the sequential axpy path.
9379    pub fn reduce_slots(
9380        &self,
9381        slots: &CudaSlice<f32>,
9382        wbuf: &CudaSlice<f32>,
9383        dst: &mut CudaSlice<f32>,
9384        ncols: usize,
9385        n_used: usize,
9386        t: usize,
9387    ) -> Result<(), Box<dyn std::error::Error>> {
9388        let f = self.func("reduce_slots_f32");
9389        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9390        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9391        let __s_b = self.gpu.stream();
9392        let mut b = __s_b.launch_builder(&f);
9393        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
9394        unsafe {
9395            b.launch(cfg)?;
9396        }
9397        Ok(())
9398    }
9399
9400    /// Canonical slot-order reduction with separately rounded multiply and add.
9401    ///
9402    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
9403    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
9404    pub fn reduce_slots_host(
9405        &self,
9406        slots: &CudaSlice<f32>,
9407        wbuf: &CudaSlice<f32>,
9408        dst: &mut CudaSlice<f32>,
9409        ncols: usize,
9410        n_used: usize,
9411        t: usize,
9412    ) -> Result<(), Box<dyn std::error::Error>> {
9413        let f = self.func("reduce_slots_host_f32");
9414        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9415        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9416        let __s_b = self.gpu.stream();
9417        let mut b = __s_b.launch_builder(&f);
9418        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
9419        unsafe {
9420            b.launch(cfg)?;
9421        }
9422        Ok(())
9423    }
9424
9425    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
9426    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
9427    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
9428    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
9429    /// GPU time, ~half of it redundant re-quantization of the same row.
9430    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
9431    pub fn quantize_q8_1_view(
9432        &self,
9433        x: &cudarc::driver::CudaView<f32>,
9434        m: usize,
9435        in_f: usize,
9436    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9437        let f = self.func("quantize_q8_1");
9438        let nblk = in_f / 32;
9439        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
9440        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
9441        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9442        let (inf, mi) = (in_f as i32, m as i32);
9443        let __s_b = self.gpu.stream();
9444        let mut b = __s_b.launch_builder(&f);
9445        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
9446        unsafe {
9447            b.launch(cfg)?;
9448        }
9449        Ok((q, d))
9450    }
9451
9452    pub fn quantize_q8_1(
9453        &self,
9454        x: &CudaSlice<f32>,
9455        m: usize,
9456        in_f: usize,
9457    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9458        let nblk = in_f / 32;
9459        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
9460        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
9461        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
9462        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9463        let (inf, mi) = (in_f as i32, m as i32);
9464        if Self::pdl_on() && Self::pdl_wb_on() {
9465            {
9466                use cudarc::driver::{DevicePtr, DevicePtrMut};
9467                let s = &self.gpu.stream();
9468                let (px, _g0) = x.device_ptr(s);
9469                let (pq, _g1) = q.device_ptr_mut(s);
9470                let (pd, _g2) = d.device_ptr_mut(s);
9471                let mut ps = [
9472                    &px as *const _ as *mut std::ffi::c_void,
9473                    &pq as *const _ as *mut _,
9474                    &pd as *const _ as *mut _,
9475                    &inf as *const _ as *mut _,
9476                    &mi as *const _ as *mut _,
9477                ];
9478                unsafe {
9479                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9480                }
9481            }
9482            return Ok((q, d));
9483        }
9484        let f = self.func("quantize_q8_1");
9485        let __s_b = self.gpu.stream();
9486        let mut b = __s_b.launch_builder(&f);
9487        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
9488        unsafe {
9489            b.launch(cfg)?;
9490        }
9491        Ok((q, d))
9492    }
9493
9494    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
9495    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
9496    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
9497    pub fn quantize_fp4_act(
9498        &self,
9499        x: &CudaSlice<f32>,
9500        m: usize,
9501        in_f: usize,
9502    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
9503        let f = self.func("quantize_fp4_act");
9504        let nb16 = in_f / 16;
9505        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
9506        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
9507        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
9508        let (inf, mi) = (in_f as i32, m as i32);
9509        let __s_b = self.gpu.stream();
9510        let mut b = __s_b.launch_builder(&f);
9511        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
9512        unsafe {
9513            b.launch(cfg)?;
9514        }
9515        Ok((aq4, ad4))
9516    }
9517
9518    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
9519    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
9520    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
9521    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
9522    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9523    pub fn qmatvec_gemm_nvfp4_fp4(
9524        &self,
9525        bytes: &CudaSlice<u8>,
9526        x: &CudaSlice<f32>,
9527        m: usize,
9528        in_f: usize,
9529        out_f: usize,
9530        row_bytes: usize,
9531        scale: f32,
9532    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9533        assert!(
9534            in_f.is_multiple_of(64),
9535            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
9536        );
9537        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
9538        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
9539        if scale != 1.0 {
9540            self.scale_inplace(&mut y, scale, m * out_f)?;
9541        }
9542        Ok(y)
9543    }
9544
9545    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
9546    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
9547    #[allow(clippy::too_many_arguments)]
9548    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9549    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
9550    fn fp4_gemm_launch(
9551        &self,
9552        bytes: &CudaSlice<u8>,
9553        aq4: &CudaSlice<u32>,
9554        ad4: &CudaSlice<u8>,
9555        m: usize,
9556        in_f: usize,
9557        out_f: usize,
9558        row_bytes: usize,
9559    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9560        let f = self.func("qmatvec_gemm_nvfp4_fp4");
9561        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9562        const BM: u32 = 64;
9563        const BN: u32 = 256;
9564        let cfg = LaunchConfig {
9565            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
9566            block_dim: (32, 4, 1),
9567            shared_mem_bytes: 0,
9568        };
9569        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9570        let __s_b = self.gpu.stream();
9571        let mut b = __s_b.launch_builder(&f);
9572        b.arg(bytes)
9573            .arg(aq4)
9574            .arg(ad4)
9575            .arg(&mut y)
9576            .arg(&inf)
9577            .arg(&outf)
9578            .arg(&mi)
9579            .arg(&rb);
9580        unsafe {
9581            b.launch(cfg)?;
9582        }
9583        Ok(y)
9584    }
9585
9586    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
9587    pub fn qmatvec_gemm_nvfp4_fp4_raw(
9588        &self,
9589        bytes: &CudaSlice<u8>,
9590        x: &CudaSlice<f32>,
9591        m: usize,
9592        in_f: usize,
9593        out_f: usize,
9594        row_bytes: usize,
9595    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9596        assert!(
9597            in_f.is_multiple_of(64),
9598            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
9599        );
9600        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
9601        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
9602    }
9603
9604    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
9605    pub fn qmatvec_q8_0_fast(
9606        &self,
9607        w: &CudaSlice<u8>,
9608        x: &CudaSlice<f32>,
9609        m: usize,
9610        in_f: usize,
9611        out_f: usize,
9612        row_bytes: usize,
9613    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9614        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9615        let f = self.func("qmatvec_q8_0_dp4a");
9616        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9617        let cfg = LaunchConfig {
9618            grid_dim: (out_f as u32, m as u32, 1),
9619            block_dim: (128, 1, 1),
9620            shared_mem_bytes: 0,
9621        };
9622        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9623        let __s_b = self.gpu.stream();
9624        let mut b = __s_b.launch_builder(&f);
9625        b.arg(w)
9626            .arg(&aq)
9627            .arg(&ad)
9628            .arg(&mut y)
9629            .arg(&inf)
9630            .arg(&outf)
9631            .arg(&mi)
9632            .arg(&rb);
9633        unsafe {
9634            b.launch(cfg)?;
9635        }
9636        Ok(y)
9637    }
9638
9639    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
9640    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9641    pub fn qmatvec_q4_K_fast(
9642        &self,
9643        w: &CudaSlice<u8>,
9644        x: &CudaSlice<f32>,
9645        m: usize,
9646        in_f: usize,
9647        out_f: usize,
9648        row_bytes: usize,
9649    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9650        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9651        let f = self.func("qmatvec_q4_K_dp4a");
9652        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9653        let cfg = LaunchConfig {
9654            grid_dim: (out_f as u32, m as u32, 1),
9655            block_dim: (128, 1, 1),
9656            shared_mem_bytes: 0,
9657        };
9658        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9659        let __s_b = self.gpu.stream();
9660        let mut b = __s_b.launch_builder(&f);
9661        b.arg(w)
9662            .arg(&aq)
9663            .arg(&ad)
9664            .arg(&mut y)
9665            .arg(&inf)
9666            .arg(&outf)
9667            .arg(&mi)
9668            .arg(&rb);
9669        unsafe {
9670            b.launch(cfg)?;
9671        }
9672        Ok(y)
9673    }
9674
9675    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
9676    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9677    pub fn qmatvec_q6_K_fast(
9678        &self,
9679        w: &CudaSlice<u8>,
9680        x: &CudaSlice<f32>,
9681        m: usize,
9682        in_f: usize,
9683        out_f: usize,
9684        row_bytes: usize,
9685    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9686        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9687        let f = self.func("qmatvec_q6_K_dp4a");
9688        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9689        let cfg = LaunchConfig {
9690            grid_dim: (out_f as u32, m as u32, 1),
9691            block_dim: (128, 1, 1),
9692            shared_mem_bytes: 0,
9693        };
9694        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9695        let __s_b = self.gpu.stream();
9696        let mut b = __s_b.launch_builder(&f);
9697        b.arg(w)
9698            .arg(&aq)
9699            .arg(&ad)
9700            .arg(&mut y)
9701            .arg(&inf)
9702            .arg(&outf)
9703            .arg(&mi)
9704            .arg(&rb);
9705        unsafe {
9706            b.launch(cfg)?;
9707        }
9708        Ok(y)
9709    }
9710
9711    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
9712    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9713    pub fn qmatvec_q5_K_fast(
9714        &self,
9715        w: &CudaSlice<u8>,
9716        x: &CudaSlice<f32>,
9717        m: usize,
9718        in_f: usize,
9719        out_f: usize,
9720        row_bytes: usize,
9721    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9722        self.qmatvec_dp4a_named(
9723            "qmatvec_q5_K_dp4a",
9724            &w.slice(0..w.len()),
9725            x,
9726            m,
9727            in_f,
9728            out_f,
9729            row_bytes,
9730        )
9731    }
9732    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
9733    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9734    pub fn qmatvec_q3_K_fast(
9735        &self,
9736        w: &CudaSlice<u8>,
9737        x: &CudaSlice<f32>,
9738        m: usize,
9739        in_f: usize,
9740        out_f: usize,
9741        row_bytes: usize,
9742    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9743        self.qmatvec_dp4a_named(
9744            "qmatvec_q3_K_dp4a",
9745            &w.slice(0..w.len()),
9746            x,
9747            m,
9748            in_f,
9749            out_f,
9750            row_bytes,
9751        )
9752    }
9753    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
9754    pub fn qmatvec_nvfp4_fast_rp(
9755        &self,
9756        w: &CudaSlice<u8>,
9757        x: &CudaSlice<f32>,
9758        m: usize,
9759        in_f: usize,
9760        out_f: usize,
9761        row_bytes: usize,
9762    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9763        assert!(
9764            in_f.is_multiple_of(64),
9765            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9766        );
9767        self.qmatvec_dp4a_named(
9768            "qmatvec_nvfp4_dp4a_rp",
9769            &w.slice(0..w.len()),
9770            x,
9771            m,
9772            in_f,
9773            out_f,
9774            row_bytes,
9775        )
9776    }
9777    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
9778    pub fn qmatvec_nvfp4_fast(
9779        &self,
9780        w: &cudarc::driver::CudaView<'_, u8>,
9781        x: &CudaSlice<f32>,
9782        m: usize,
9783        in_f: usize,
9784        out_f: usize,
9785        row_bytes: usize,
9786    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9787        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
9788        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
9789        assert!(
9790            in_f.is_multiple_of(64),
9791            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9792        );
9793        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
9794    }
9795    /// Slot-major-layout twin of `qmatvec_nvfp4_fast`: bit-identical per row, coalesced
9796    /// reads. Since the 2026-08-29 `MEMRA_NVFP4_BANK_V2` door removal its only in-tree
9797    /// producer of slot-major banks is the EP2 whole-expert bank build; this is EP2's
9798    /// host-canonical oracle reader (plus offline harnesses like moe_tp2_repro).
9799    pub fn qmatvec_nvfp4_fast_v2(
9800        &self,
9801        w: &cudarc::driver::CudaView<'_, u8>,
9802        x: &CudaSlice<f32>,
9803        m: usize,
9804        in_f: usize,
9805        out_f: usize,
9806        row_bytes: usize,
9807    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9808        assert!(
9809            in_f.is_multiple_of(64),
9810            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9811        );
9812        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
9813    }
9814    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
9815    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9816    pub fn qmatvec_iq4_XS_fast(
9817        &self,
9818        w: &CudaSlice<u8>,
9819        x: &CudaSlice<f32>,
9820        m: usize,
9821        in_f: usize,
9822        out_f: usize,
9823        row_bytes: usize,
9824    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9825        self.qmatvec_dp4a_named(
9826            "qmatvec_iq4_XS_dp4a",
9827            &w.slice(0..w.len()),
9828            x,
9829            m,
9830            in_f,
9831            out_f,
9832            row_bytes,
9833        )
9834    }
9835
9836    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
9837    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9838    fn qmatvec_dp4a_named(
9839        &self,
9840        name: &str,
9841        w: &cudarc::driver::CudaView<'_, u8>,
9842        x: &CudaSlice<f32>,
9843        m: usize,
9844        in_f: usize,
9845        out_f: usize,
9846        row_bytes: usize,
9847    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9848        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9849        let f = self.func(name);
9850        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9851        let cfg = LaunchConfig {
9852            grid_dim: (out_f as u32, m as u32, 1),
9853            block_dim: (128, 1, 1),
9854            shared_mem_bytes: 0,
9855        };
9856        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9857        let __s_b = self.gpu.stream();
9858        let mut b = __s_b.launch_builder(&f);
9859        b.arg(w)
9860            .arg(&aq)
9861            .arg(&ad)
9862            .arg(&mut y)
9863            .arg(&inf)
9864            .arg(&outf)
9865            .arg(&mi)
9866            .arg(&rb);
9867        unsafe {
9868            b.launch(cfg)?;
9869        }
9870        Ok(y)
9871    }
9872
9873    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
9874    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
9875    /// its output); this entry exists so a routed-expert program can quantize one activation
9876    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
9877    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
9878    #[allow(clippy::too_many_arguments)]
9879    pub fn qmatvec_nvfp4_fast_prequant_into(
9880        &self,
9881        w: &CudaSlice<u8>,
9882        aq: &CudaSlice<i8>,
9883        ad: &CudaSlice<f32>,
9884        y: &mut CudaSlice<f32>,
9885        m: usize,
9886        in_f: usize,
9887        out_f: usize,
9888        row_bytes: usize,
9889    ) -> Result<(), Box<dyn std::error::Error>> {
9890        assert!(
9891            in_f.is_multiple_of(64),
9892            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9893        );
9894        if y.len() < m * out_f {
9895            return Err(format!(
9896                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
9897                y.len()
9898            )
9899            .into());
9900        }
9901        let f = self.func("qmatvec_nvfp4_dp4a");
9902        let cfg = LaunchConfig {
9903            grid_dim: (out_f as u32, m as u32, 1),
9904            block_dim: (128, 1, 1),
9905            shared_mem_bytes: 0,
9906        };
9907        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9908        let __s_b = self.gpu.stream();
9909        let mut b = __s_b.launch_builder(&f);
9910        b.arg(w)
9911            .arg(aq)
9912            .arg(ad)
9913            .arg(y)
9914            .arg(&inf)
9915            .arg(&outf)
9916            .arg(&mi)
9917            .arg(&rb);
9918        unsafe {
9919            b.launch(cfg)?;
9920        }
9921        Ok(())
9922    }
9923
9924    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
9925    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
9926    #[allow(clippy::too_many_arguments)]
9927    pub fn matvec_f32_qkv_into(
9928        &self,
9929        wq: &CudaSlice<f32>,
9930        wk: &CudaSlice<f32>,
9931        wv: &CudaSlice<f32>,
9932        wg: &CudaSlice<f32>,
9933        x: &CudaSlice<f32>,
9934        yq: &mut CudaSlice<f32>,
9935        yk: &mut CudaSlice<f32>,
9936        yv: &mut CudaSlice<f32>,
9937        yg: &mut CudaSlice<f32>,
9938        in_f: usize,
9939        out_q: usize,
9940        out_kv: usize,
9941        out_g: usize,
9942    ) -> Result<(), Box<dyn std::error::Error>> {
9943        if !in_f.is_multiple_of(4)
9944            || wq.len() != out_q * in_f
9945            || wk.len() != out_kv * in_f
9946            || wv.len() != out_kv * in_f
9947            || wg.len() < out_g * in_f
9948            || x.len() < in_f
9949            || yq.len() < out_q
9950            || yk.len() < out_kv
9951            || yv.len() < out_kv
9952            || (out_g > 0 && yg.len() < out_g)
9953        {
9954            return Err(format!(
9955                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
9956                 wq={} wk={} wv={} wg={}",
9957                wq.len(),
9958                wk.len(),
9959                wv.len(),
9960                wg.len()
9961            )
9962            .into());
9963        }
9964        let f = self.func("matvec_f32_qkv");
9965        let cfg = LaunchConfig {
9966            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
9967            block_dim: (128, 1, 1),
9968            shared_mem_bytes: 0,
9969        };
9970        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
9971        let __s_b = self.gpu.stream();
9972        let mut b = __s_b.launch_builder(&f);
9973        b.arg(wq)
9974            .arg(wk)
9975            .arg(wv)
9976            .arg(wg)
9977            .arg(x)
9978            .arg(yq)
9979            .arg(yk)
9980            .arg(yv)
9981            .arg(yg)
9982            .arg(&inf)
9983            .arg(&oq)
9984            .arg(&okv)
9985            .arg(&og);
9986        unsafe {
9987            b.launch(cfg)?;
9988        }
9989        Ok(())
9990    }
9991
9992    /// PROGRAM 2 (`MEMRA_NVFP4_SEL_GU`, default OFF): the routed gate and up sweeps in ONE
9993    /// launch. The two sweeps share `sel`/`aq`/`ad` and have identical geometry, so blocks
9994    /// `[0,out_f)` run the exact `_sel_v2` body on the GATE bank and `[out_f,2*out_f)` on the UP
9995    /// bank — per-row BIT-IDENTICAL to two `qmatvec_nvfp4_sel_into` calls, with half the sweep
9996    /// launches and double the grid fill.
9997    ///
9998    /// SLOT-MAJOR ONLY, and the caller proves it: the kernel reads the slot-major byte map, so
9999    /// this refuses banks that do not carry it rather than trusting an env door. In the removed
10000    /// implementation this fusion auto-armed on `nvfp4_bank_v2_on()` with NO door of its own,
10001    /// which is one of the three programs that moved together behind one env var and made the
10002    /// 2026-08-29 bisect unable to name a mechanism (DIAGNOSIS.md).
10003    #[allow(clippy::too_many_arguments)]
10004    pub fn qmatvec_nvfp4_sel_gu_into(
10005        &self,
10006        gate_bank: &CudaSlice<u8>,
10007        up_bank: &CudaSlice<u8>,
10008        sel: &CudaSlice<i32>,
10009        aq: &CudaSlice<i8>,
10010        ad: &CudaSlice<f32>,
10011        yg: &mut CudaSlice<f32>,
10012        yu: &mut CudaSlice<f32>,
10013        n_sel: usize,
10014        in_f: usize,
10015        out_f: usize,
10016        row_bytes: usize,
10017        expert_stride: usize,
10018        slot_major: bool,
10019    ) -> Result<(), Box<dyn std::error::Error>> {
10020        assert!(
10021            in_f.is_multiple_of(64),
10022            "NVFP4 dp4a requires in_f % 64 == 0"
10023        );
10024        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
10025            return Err("NVFP4 gu sel geometry".into());
10026        }
10027        if !slot_major {
10028            return Err(
10029                "NVFP4 gu sel fusion reads slot-major rows: these banks are block_nvfp4 \
10030                        v1 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
10031                    .into(),
10032            );
10033        }
10034        // MEMRA_NVFP4_SEL_GU_RPW=2|4 (sub-door, default OFF, UNPRICED): multirow twin — the
10035        // activation group is read once and reused across RPW rows' gate+up dots. Per-row
10036        // accumulation order and reduce tree are the base kernel's -> bit-identical.
10037        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
10038        let rpw = *RPW.get_or_init(|| {
10039            std::env::var("MEMRA_NVFP4_SEL_GU_RPW")
10040                .ok()
10041                .and_then(|v| v.parse().ok())
10042                .filter(|r| *r == 2 || *r == 4)
10043                .unwrap_or(1)
10044        });
10045        let rpw = if out_f.is_multiple_of(rpw) { rpw } else { 1 };
10046        // MEMRA_NVFP4_SEL_GU_WPR=1 (sub-door, default OFF, UNPRICED): warp-per-row.
10047        // NUMERIC-CLASS — the per-row REDUCTION ORDER changes, so a bit tape cannot apply and
10048        // acceptance is the argmax gate plus the boot battery (the QKV_FUSED/BF16_MMV class).
10049        // It is deliberately NOT part of this lane's priced arms, which are all bit-gateable.
10050        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10051        let wpr =
10052            *WPR.get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_GU_WPR").as_deref() == Ok("1"));
10053        let f = self.func(match (wpr, rpw) {
10054            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
10055            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
10056            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
10057            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
10058        });
10059        let cfg = LaunchConfig {
10060            grid_dim: if wpr {
10061                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
10062            } else if rpw == 1 {
10063                ((2 * out_f) as u32, n_sel as u32, 1)
10064            } else {
10065                ((out_f / rpw) as u32, n_sel as u32, 1)
10066            },
10067            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
10068            shared_mem_bytes: 0,
10069        };
10070        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10071        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10072        let (ars, adrs) = (0i64, 0i64);
10073        let __s_b = self.gpu.stream();
10074        let mut b = __s_b.launch_builder(&f);
10075        b.arg(gate_bank)
10076            .arg(up_bank)
10077            .arg(sel)
10078            .arg(aq)
10079            .arg(ad)
10080            .arg(yg)
10081            .arg(yu)
10082            .arg(&inf)
10083            .arg(&outf)
10084            .arg(&ns)
10085            .arg(&rb)
10086            .arg(&es)
10087            .arg(&ars)
10088            .arg(&adrs);
10089        unsafe {
10090            b.launch(cfg)?;
10091        }
10092        Ok(())
10093    }
10094
10095    /// PROGRAM 3 (`MEMRA_NVFP4_SEL_DOWN8`, **default ON since 2026-09-01**): the DOWN sweep and
10096    /// the route-weight
10097    /// combine in ONE launch (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm
10098    /// ported to the NVFP4 banks). Block = `(32, n_sel)`: one warp per slot instead of one warp
10099    /// per (row, slot), and the `n_sel x out_f` partial buffer disappears. BIT-IDENTICAL to
10100    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce tree,
10101    /// same slot-ordered combine chain.
10102    ///
10103    /// Requires slot-major rows and `nsb <= 32` (the fit-block class the reduce identity is
10104    /// argued at). The removed implementation refused on `!nvfp4_bank_v2_on()`; it now refuses on
10105    /// the LAYOUT THE CALLER READ OFF THE BANK, so the guard cannot disagree with the bytes.
10106    #[allow(clippy::too_many_arguments)]
10107    pub fn qmatvec_nvfp4_sel_down8_into(
10108        &self,
10109        bank: &CudaSlice<u8>,
10110        sel: &CudaSlice<i32>,
10111        aq: &CudaSlice<i8>,
10112        ad: &CudaSlice<f32>,
10113        route_w: &CudaSlice<f32>,
10114        md: &CudaSlice<f32>,
10115        dst: &mut CudaSlice<f32>,
10116        n_sel: usize,
10117        in_f: usize,
10118        out_f: usize,
10119        row_bytes: usize,
10120        expert_stride: usize,
10121        act_row_stride: usize,
10122        ad_row_stride: usize,
10123        slot_major: bool,
10124    ) -> Result<(), Box<dyn std::error::Error>> {
10125        if !in_f.is_multiple_of(64)
10126            || n_sel == 0
10127            || n_sel > 8
10128            || (in_f >> 5) > 32
10129            || dst.len() < out_f
10130            || sel.len() < n_sel
10131            || route_w.len() < n_sel
10132        {
10133            return Err(format!(
10134                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
10135                dst.len()
10136            )
10137            .into());
10138        }
10139        if !slot_major {
10140            return Err(
10141                "NVFP4 sel down8 reads slot-major rows: this shard is block_nvfp4 v1 \
10142                        (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
10143                    .into(),
10144            );
10145        }
10146        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
10147        let cfg = LaunchConfig {
10148            grid_dim: (out_f as u32, 1, 1),
10149            block_dim: (32, n_sel as u32, 1),
10150            shared_mem_bytes: 0,
10151        };
10152        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10153        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10154        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
10155        let __s_b = self.gpu.stream();
10156        let mut b = __s_b.launch_builder(&f);
10157        b.arg(bank)
10158            .arg(sel)
10159            .arg(aq)
10160            .arg(ad)
10161            .arg(route_w)
10162            .arg(md)
10163            .arg(dst)
10164            .arg(&inf)
10165            .arg(&outf)
10166            .arg(&ns)
10167            .arg(&rb)
10168            .arg(&es)
10169            .arg(&ars)
10170            .arg(&adrs);
10171        unsafe {
10172            b.launch(cfg)?;
10173        }
10174        Ok(())
10175    }
10176
10177    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
10178    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
10179    #[allow(clippy::too_many_arguments)]
10180    pub fn qmatvec_nvfp4_sel_gu_ep_into(
10181        &self,
10182        gate_bank: &CudaSlice<u8>,
10183        up_bank: &CudaSlice<u8>,
10184        sel: &CudaSlice<i32>,
10185        aq: &CudaSlice<i8>,
10186        ad: &CudaSlice<f32>,
10187        yg: &mut CudaSlice<f32>,
10188        yu: &mut CudaSlice<f32>,
10189        n_sel: usize,
10190        in_f: usize,
10191        out_f: usize,
10192        row_bytes: usize,
10193        expert_stride: usize,
10194        owner: usize,
10195    ) -> Result<(), Box<dyn std::error::Error>> {
10196        assert!(
10197            in_f.is_multiple_of(64),
10198            "NVFP4 dp4a requires in_f % 64 == 0"
10199        );
10200        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
10201            return Err("NVFP4 gu ep geometry".into());
10202        }
10203        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
10204        let cfg = LaunchConfig {
10205            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
10206            block_dim: (128, 1, 1),
10207            shared_mem_bytes: 0,
10208        };
10209        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
10210        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10211        let (ars, adrs) = (0i64, 0i64);
10212        let __s_b = self.gpu.stream();
10213        let mut b = __s_b.launch_builder(&f);
10214        b.arg(gate_bank)
10215            .arg(up_bank)
10216            .arg(sel)
10217            .arg(aq)
10218            .arg(ad)
10219            .arg(yg)
10220            .arg(yu)
10221            .arg(&inf)
10222            .arg(&outf)
10223            .arg(&ns)
10224            .arg(&rb)
10225            .arg(&es)
10226            .arg(&ars)
10227            .arg(&adrs)
10228            .arg(&own);
10229        unsafe {
10230            b.launch(cfg)?;
10231        }
10232        Ok(())
10233    }
10234
10235    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
10236    #[allow(clippy::too_many_arguments)]
10237    pub fn silu_mul_scaled_q8_1_sel_ep_into(
10238        &self,
10239        gate: &CudaSlice<f32>,
10240        up: &CudaSlice<f32>,
10241        gmac: &CudaSlice<f32>,
10242        umac: &CudaSlice<f32>,
10243        sel: &CudaSlice<i32>,
10244        limit: Option<f32>,
10245        out_q: &mut CudaSlice<i8>,
10246        out_d: &mut CudaSlice<f32>,
10247        n_per: usize,
10248        n_sel: usize,
10249        owner: usize,
10250    ) -> Result<(), Box<dyn std::error::Error>> {
10251        if !n_per.is_multiple_of(32)
10252            || out_q.len() < n_sel * n_per
10253            || out_d.len() < n_sel * n_per / 32
10254        {
10255            return Err("NVFP4 silu ep geometry".into());
10256        }
10257        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
10258        let warps = n_sel * n_per / 32;
10259        let cfg = LaunchConfig {
10260            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
10261            block_dim: (128, 1, 1),
10262            shared_mem_bytes: 0,
10263        };
10264        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
10265        let (lim, has) = match limit {
10266            Some(l) => (l, 1i32),
10267            None => (0.0f32, 0i32),
10268        };
10269        let __s_b = self.gpu.stream();
10270        let mut b = __s_b.launch_builder(&f);
10271        b.arg(gate)
10272            .arg(up)
10273            .arg(gmac)
10274            .arg(umac)
10275            .arg(sel)
10276            .arg(&lim)
10277            .arg(&has)
10278            .arg(out_q)
10279            .arg(out_d)
10280            .arg(&np)
10281            .arg(&ns)
10282            .arg(&own);
10283        unsafe {
10284            b.launch(cfg)?;
10285        }
10286        Ok(())
10287    }
10288
10289    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
10290    #[allow(clippy::too_many_arguments)]
10291    pub fn qmatvec_nvfp4_sel_down8_ep_into(
10292        &self,
10293        bank: &CudaSlice<u8>,
10294        sel: &CudaSlice<i32>,
10295        aq: &CudaSlice<i8>,
10296        ad: &CudaSlice<f32>,
10297        route_w: &CudaSlice<f32>,
10298        md: &CudaSlice<f32>,
10299        dst: &mut CudaSlice<f32>,
10300        n_sel: usize,
10301        in_f: usize,
10302        out_f: usize,
10303        row_bytes: usize,
10304        expert_stride: usize,
10305        act_row_stride: usize,
10306        ad_row_stride: usize,
10307        owner: usize,
10308    ) -> Result<(), Box<dyn std::error::Error>> {
10309        if !in_f.is_multiple_of(64)
10310            || n_sel == 0
10311            || n_sel > 8
10312            || (in_f >> 5) > 64
10313            || dst.len() < out_f
10314        {
10315            return Err("NVFP4 down8 ep geometry".into());
10316        }
10317        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
10318        let cfg = LaunchConfig {
10319            grid_dim: (out_f as u32, 1, 1),
10320            block_dim: (32, n_sel as u32, 1),
10321            shared_mem_bytes: 0,
10322        };
10323        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
10324        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10325        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
10326        let __s_b = self.gpu.stream();
10327        let mut b = __s_b.launch_builder(&f);
10328        b.arg(bank)
10329            .arg(sel)
10330            .arg(aq)
10331            .arg(ad)
10332            .arg(route_w)
10333            .arg(md)
10334            .arg(dst)
10335            .arg(&inf)
10336            .arg(&outf)
10337            .arg(&ns)
10338            .arg(&rb)
10339            .arg(&es)
10340            .arg(&ars)
10341            .arg(&adrs)
10342            .arg(&own);
10343        unsafe {
10344            b.launch(cfg)?;
10345        }
10346        Ok(())
10347    }
10348
10349    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
10350    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
10351    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
10352    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
10353    /// kernel — the batching only removes host launch latency.
10354    ///
10355    /// `slot_major` names the LAYOUT OF THE BYTES AT `bank` and is REQUIRED, never defaulted:
10356    /// true routes the `_sel_v2` reader (slot g's 16 qs bytes at `g*16`, scale tail at
10357    /// `nslots*16`), false the block_nvfp4 v1 reader. The caller reads it off the resident bank
10358    /// (`ResidentNvfp4{Column,Row}BankRank::slot_major`) — never off an env door, and never with
10359    /// a default. A defaulted layout scalar in exactly this position is what produced the
10360    /// 2026-08-29 step37 corruption (`kq_fetch(..., int in_f = 0)`,
10361    /// research/step37-bankv3-20260901/DIAGNOSIS.md).
10362    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10363    pub fn qmatvec_nvfp4_sel_into(
10364        &self,
10365        bank: &CudaSlice<u8>,
10366        sel: &CudaSlice<i32>,
10367        aq: &CudaSlice<i8>,
10368        ad: &CudaSlice<f32>,
10369        y: &mut CudaSlice<f32>,
10370        n_sel: usize,
10371        in_f: usize,
10372        out_f: usize,
10373        row_bytes: usize,
10374        expert_stride: usize,
10375        act_row_stride: usize,
10376        ad_row_stride: usize,
10377        slot_major: bool,
10378    ) -> Result<(), Box<dyn std::error::Error>> {
10379        assert!(
10380            in_f.is_multiple_of(64),
10381            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10382        );
10383        if y.len() < n_sel * out_f || sel.len() < n_sel {
10384            return Err(format!(
10385                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
10386                y.len(),
10387                sel.len()
10388            )
10389            .into());
10390        }
10391        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
10392        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
10393        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
10394        // sequential-rows variant was flat). Default stays the single-row form.
10395        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
10396        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
10397        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
10398        let mode = *MR.get_or_init(|| {
10399            if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
10400                2
10401            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
10402                1
10403            } else {
10404                0
10405            }
10406        });
10407        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
10408        // Mode 3 is the SLOT-MAJOR reader, and it is chosen by the BANK's layout, not by an env
10409        // door: `MEMRA_SEL_MR`/`MEMRA_SEL_STREAM` are v1-layout probes and cannot read these
10410        // bytes at all, so the layout overrides them rather than racing them.
10411        let mode = if slot_major { 3 } else { mode };
10412        // MEMRA_NVFP4_SEL_SM_STREAM=1 (sub-door of MEMRA_NVFP4_BANK_SM, default OFF, UNPRICED):
10413        // 8 contiguous rows per block with next-row int4 prefetch. Needs 16B-aligned rows
10414        // (step37 gate/up 2304B yes, down 360B no -> single-row) and one slot per thread.
10415        static SM_STREAM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10416        let sm_stream = mode == 3
10417            && *SM_STREAM
10418                .get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_SM_STREAM").as_deref() == Ok("1"))
10419            && row_bytes.is_multiple_of(16)
10420            && in_f <= 4096;
10421        let kname = match (mode, sm_stream) {
10422            (3, true) => "qmatvec_nvfp4_dp4a_sel_v2s",
10423            (3, false) => "qmatvec_nvfp4_dp4a_sel_v2",
10424            (2, _) => "qmatvec_nvfp4_dp4a_sel_stream",
10425            (1, _) => "qmatvec_nvfp4_dp4a_sel_mr4",
10426            _ => "qmatvec_nvfp4_dp4a_sel",
10427        };
10428        // ENGAGEMENT RECEIPT for PROGRAM 1, one line per distinct (kernel, geometry) pair. The
10429        // door being SET in the environment does not prove the slot-major READER ran; only the
10430        // selected kernel name does. Without this, a pricing cell that reports a flat delta
10431        // cannot distinguish "the program is worth nothing" from "the program never ran" — the
10432        // defect the MEMRA_BF16_MMV lane hit when its engagement grep returned 0 in both arms.
10433        {
10434            static SEEN_SEL: std::sync::Mutex<Vec<(&'static str, usize, usize)>> =
10435                std::sync::Mutex::new(Vec::new());
10436            let combo = (kname, in_f, out_f);
10437            let mut seen = SEEN_SEL.lock().unwrap();
10438            if !seen.contains(&combo) {
10439                seen.push(combo);
10440                eprintln!(
10441                    "[nvfp4-sel] kernel={kname} slot_major={slot_major} in_f={in_f} \
10442                     out_f={out_f} nsb={} row_bytes={row_bytes}",
10443                    in_f >> 5
10444                );
10445            }
10446        }
10447        let f = self.func(kname);
10448        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
10449        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
10450        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
10451        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
10452        let nsb = in_f >> 5;
10453        let fit_block: u32 = if (mode == 0 || mode == 3) && !sm_stream && nsb <= 32 {
10454            32
10455        } else if mode == 1 {
10456            512
10457        } else {
10458            128
10459        };
10460        let cfg = LaunchConfig {
10461            grid_dim: (
10462                if sm_stream {
10463                    (out_f as u32).div_ceil(8)
10464                } else {
10465                    match mode {
10466                        2 => (out_f as u32).div_ceil(16),
10467                        1 => (out_f as u32).div_ceil(4),
10468                        _ => out_f as u32,
10469                    }
10470                },
10471                n_sel as u32,
10472                1,
10473            ),
10474            block_dim: (fit_block, 1, 1),
10475            shared_mem_bytes: 0,
10476        };
10477        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10478        let (rb, es, ars, adrs) = (
10479            row_bytes as i64,
10480            expert_stride as i64,
10481            act_row_stride as i64,
10482            ad_row_stride as i64,
10483        );
10484        let __s_b = self.gpu.stream();
10485        let mut b = __s_b.launch_builder(&f);
10486        b.arg(bank)
10487            .arg(sel)
10488            .arg(aq)
10489            .arg(ad)
10490            .arg(y)
10491            .arg(&inf)
10492            .arg(&outf)
10493            .arg(&ns)
10494            .arg(&rb)
10495            .arg(&es)
10496            .arg(&ars)
10497            .arg(&adrs);
10498        unsafe {
10499            b.launch(cfg)?;
10500        }
10501        Ok(())
10502    }
10503
10504    /// W4A16 selected-expert gate+up pair. `x_bf16` contains checkpoint-rounded BF16
10505    /// activations; selected ids are local to the rank's contiguous expert bank.
10506    #[allow(clippy::too_many_arguments)]
10507    pub fn qmatvec_nvfp4_bf16_sel_dual_rows_into(
10508        &self,
10509        gate_bank: &CudaSlice<u8>,
10510        up_bank: &CudaSlice<u8>,
10511        sel: &CudaSlice<i32>,
10512        token_rows: &CudaSlice<i32>,
10513        x_bf16: &CudaSlice<u8>,
10514        gate_out: &mut CudaSlice<f32>,
10515        up_out: &mut CudaSlice<f32>,
10516        n_sel: usize,
10517        in_f: usize,
10518        out_f: usize,
10519        row_bytes: usize,
10520        expert_stride: usize,
10521        tokens: usize,
10522    ) -> Result<(), Box<dyn std::error::Error>> {
10523        if !in_f.is_multiple_of(64)
10524            || sel.len() < n_sel
10525            || token_rows.len() < n_sel
10526            || gate_out.len() < n_sel * out_f
10527            || up_out.len() < n_sel * out_f
10528            || x_bf16.len() < 2 * in_f * tokens
10529        {
10530            return Err(format!(
10531                "W4A16 NVFP4 dual selected rows geometry sel={} token_rows={} x={} gate={} up={} \
10532                 n_sel={n_sel} tokens={tokens} in={in_f} out={out_f}",
10533                sel.len(),
10534                token_rows.len(),
10535                x_bf16.len(),
10536                gate_out.len(),
10537                up_out.len(),
10538            )
10539            .into());
10540        }
10541        let adjacent_rows = tokens > 1;
10542        let f = if adjacent_rows {
10543            self.func("qmatvec_nvfp4_bf16_sel_quad_rows")
10544        } else {
10545            self.func("qmatvec_nvfp4_bf16_sel_dual_rows")
10546        };
10547        let cfg = LaunchConfig {
10548            grid_dim: (
10549                if adjacent_rows {
10550                    out_f.div_ceil(2) as u32
10551                } else {
10552                    (2 * out_f) as u32
10553                },
10554                n_sel as u32,
10555                1,
10556            ),
10557            block_dim: (256, 1, 1),
10558            shared_mem_bytes: 0,
10559        };
10560        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10561        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10562        let __s_b = self.gpu.stream();
10563        let mut b = __s_b.launch_builder(&f);
10564        b.arg(gate_bank)
10565            .arg(up_bank)
10566            .arg(sel)
10567            .arg(token_rows)
10568            .arg(x_bf16)
10569            .arg(gate_out)
10570            .arg(up_out)
10571            .arg(&inf)
10572            .arg(&outf)
10573            .arg(&ns)
10574            .arg(&rb)
10575            .arg(&es);
10576        unsafe {
10577            b.launch(cfg)?;
10578        }
10579        Ok(())
10580    }
10581
10582    /// Device-routed W4A16 gate+up over fixed token/slot rows. Selection ids remain global;
10583    /// each rank rejects non-owned slots and translates owned ids into its local expert bank.
10584    #[allow(clippy::too_many_arguments)]
10585    pub fn qmatvec_nvfp4_bf16_ep_dual_slots_into(
10586        &self,
10587        gate_bank: &CudaSlice<u8>,
10588        up_bank: &CudaSlice<u8>,
10589        sel: &CudaSlice<i32>,
10590        x_bf16: &CudaSlice<u8>,
10591        gate_out: &mut CudaSlice<f32>,
10592        up_out: &mut CudaSlice<f32>,
10593        n_pairs: usize,
10594        top_k: usize,
10595        in_f: usize,
10596        out_f: usize,
10597        owner_start: usize,
10598        owner_end: usize,
10599        row_bytes: usize,
10600        expert_stride: usize,
10601    ) -> Result<(), Box<dyn std::error::Error>> {
10602        let tokens = n_pairs.div_ceil(top_k);
10603        if top_k == 0
10604            || owner_start >= owner_end
10605            || !in_f.is_multiple_of(64)
10606            || sel.len() < n_pairs
10607            || gate_out.len() < n_pairs * out_f
10608            || up_out.len() < n_pairs * out_f
10609            || x_bf16.len() < 2 * in_f * tokens
10610        {
10611            return Err(format!(
10612                "W4A16 NVFP4 device EP dual-slot geometry sel={} x={} gate={} up={} \
10613                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10614                 owner={owner_start}..{owner_end}",
10615                sel.len(),
10616                x_bf16.len(),
10617                gate_out.len(),
10618                up_out.len(),
10619            )
10620            .into());
10621        }
10622        let pair_parallel = tokens > 1;
10623        let f = if pair_parallel {
10624            self.func("qmatvec_nvfp4_bf16_ep_quad_pairs")
10625        } else {
10626            self.func("qmatvec_nvfp4_bf16_ep_dual_slots")
10627        };
10628        let cfg = LaunchConfig {
10629            grid_dim: (
10630                if pair_parallel {
10631                    out_f.div_ceil(2) as u32
10632                } else {
10633                    (2 * out_f) as u32
10634                },
10635                if pair_parallel { n_pairs as u32 } else { 1 },
10636                1,
10637            ),
10638            block_dim: (256, 1, 1),
10639            shared_mem_bytes: 0,
10640        };
10641        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10642        let (os, oe) = (owner_start as i32, owner_end as i32);
10643        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10644        let __s_b = self.gpu.stream();
10645        let mut b = __s_b.launch_builder(&f);
10646        b.arg(gate_bank)
10647            .arg(up_bank)
10648            .arg(sel)
10649            .arg(x_bf16)
10650            .arg(gate_out)
10651            .arg(up_out)
10652            .arg(&inf)
10653            .arg(&outf)
10654            .arg(&np)
10655            .arg(&tk)
10656            .arg(&os)
10657            .arg(&oe)
10658            .arg(&rb)
10659            .arg(&es);
10660        unsafe {
10661            b.launch(cfg)?;
10662        }
10663        Ok(())
10664    }
10665
10666    /// Optional A8 t=1 gate+up program over global fixed slots.
10667    #[allow(clippy::too_many_arguments)]
10668    pub fn qmatvec_nvfp4_q8_ep_dual_slots_into(
10669        &self,
10670        gate_bank: &CudaSlice<u8>,
10671        up_bank: &CudaSlice<u8>,
10672        sel: &CudaSlice<i32>,
10673        aq: &CudaSlice<i8>,
10674        ad: &CudaSlice<f32>,
10675        gate_out: &mut CudaSlice<f32>,
10676        up_out: &mut CudaSlice<f32>,
10677        n_pairs: usize,
10678        top_k: usize,
10679        in_f: usize,
10680        out_f: usize,
10681        owner_start: usize,
10682        owner_end: usize,
10683        row_bytes: usize,
10684        expert_stride: usize,
10685    ) -> Result<(), Box<dyn std::error::Error>> {
10686        let tokens = n_pairs.div_ceil(top_k);
10687        if top_k == 0
10688            || owner_start >= owner_end
10689            || !in_f.is_multiple_of(64)
10690            || sel.len() < n_pairs
10691            || aq.len() < tokens * in_f
10692            || ad.len() < tokens * (in_f / 32)
10693            || gate_out.len() < n_pairs * out_f
10694            || up_out.len() < n_pairs * out_f
10695        {
10696            return Err(format!(
10697                "W4A8 NVFP4 device EP gate/up geometry sel={} aq={} ad={} gate={} up={} \
10698                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10699                 owner={owner_start}..{owner_end}",
10700                sel.len(),
10701                aq.len(),
10702                ad.len(),
10703                gate_out.len(),
10704                up_out.len(),
10705            )
10706            .into());
10707        }
10708        let f = self.func("qmatvec_nvfp4_q8_ep_dual_slots");
10709        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
10710        let cfg = LaunchConfig {
10711            grid_dim: (out_f as u32, 1, 1),
10712            block_dim: (threads, 1, 1),
10713            shared_mem_bytes: 0,
10714        };
10715        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10716        let (os, oe) = (owner_start as i32, owner_end as i32);
10717        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10718        let __s_b = self.gpu.stream();
10719        let mut b = __s_b.launch_builder(&f);
10720        b.arg(gate_bank)
10721            .arg(up_bank)
10722            .arg(sel)
10723            .arg(aq)
10724            .arg(ad)
10725            .arg(gate_out)
10726            .arg(up_out)
10727            .arg(&inf)
10728            .arg(&outf)
10729            .arg(&np)
10730            .arg(&tk)
10731            .arg(&os)
10732            .arg(&oe)
10733            .arg(&rb)
10734            .arg(&es);
10735        unsafe {
10736            b.launch(cfg)?;
10737        }
10738        Ok(())
10739    }
10740
10741    /// Known-good paired gate+up Q8 schedule: one CTA owns the same output row in both banks,
10742    /// shares the activation bytes, and retains one independent accumulator/reduction per bank.
10743    #[allow(clippy::too_many_arguments)]
10744    pub fn qmatvec_nvfp4_q8_ep_paired_slots_into(
10745        &self,
10746        gate_bank: &CudaSlice<u8>,
10747        up_bank: &CudaSlice<u8>,
10748        sel: &CudaSlice<i32>,
10749        aq: &CudaSlice<i8>,
10750        ad: &CudaSlice<f32>,
10751        gate_out: &mut CudaSlice<f32>,
10752        up_out: &mut CudaSlice<f32>,
10753        n_pairs: usize,
10754        top_k: usize,
10755        in_f: usize,
10756        out_f: usize,
10757        owner_start: usize,
10758        owner_end: usize,
10759        row_bytes: usize,
10760        expert_stride: usize,
10761    ) -> Result<(), Box<dyn std::error::Error>> {
10762        let tokens = n_pairs.div_ceil(top_k);
10763        if top_k == 0
10764            || owner_start >= owner_end
10765            || !in_f.is_multiple_of(64)
10766            || sel.len() < n_pairs
10767            || aq.len() < tokens * in_f
10768            || ad.len() < tokens * (in_f / 32)
10769            || gate_out.len() < n_pairs * out_f
10770            || up_out.len() < n_pairs * out_f
10771        {
10772            return Err(format!(
10773                "W4A8 NVFP4 paired gate/up geometry sel={} aq={} ad={} gate={} up={} \
10774                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10775                 owner={owner_start}..{owner_end}",
10776                sel.len(),
10777                aq.len(),
10778                ad.len(),
10779                gate_out.len(),
10780                up_out.len(),
10781            )
10782            .into());
10783        }
10784        let f = self.func("qmatvec_nvfp4_q8_ep_paired_slots");
10785        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
10786        let cfg = LaunchConfig {
10787            grid_dim: (out_f as u32, 1, 1),
10788            block_dim: (threads, 1, 1),
10789            shared_mem_bytes: 0,
10790        };
10791        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10792        let (os, oe) = (owner_start as i32, owner_end as i32);
10793        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10794        let __s_b = self.gpu.stream();
10795        let mut b = __s_b.launch_builder(&f);
10796        b.arg(gate_bank)
10797            .arg(up_bank)
10798            .arg(sel)
10799            .arg(aq)
10800            .arg(ad)
10801            .arg(gate_out)
10802            .arg(up_out)
10803            .arg(&inf)
10804            .arg(&outf)
10805            .arg(&np)
10806            .arg(&tk)
10807            .arg(&os)
10808            .arg(&oe)
10809            .arg(&rb)
10810            .arg(&es);
10811        unsafe {
10812            b.launch(cfg)?;
10813        }
10814        Ok(())
10815    }
10816
10817    /// W4A16 selected down rows scattered into canonical global pair positions on the root.
10818    #[allow(clippy::too_many_arguments)]
10819    pub fn qmatvec_nvfp4_bf16_sel_down_rows_raw(
10820        &self,
10821        bank: &CudaSlice<u8>,
10822        sel: &CudaSlice<i32>,
10823        global_pairs: &CudaSlice<i32>,
10824        activation_bf16: &CudaSlice<u8>,
10825        macros_down: &CudaSlice<f32>,
10826        dst_raw: u64,
10827        n_sel: usize,
10828        in_f: usize,
10829        out_f: usize,
10830        row_bytes: usize,
10831        expert_stride: usize,
10832        total_pairs: usize,
10833    ) -> Result<(), Box<dyn std::error::Error>> {
10834        if !in_f.is_multiple_of(64)
10835            || sel.len() < n_sel
10836            || global_pairs.len() < n_sel
10837            || activation_bf16.len() < 2 * n_sel * in_f
10838            || dst_raw == 0
10839        {
10840            return Err(format!(
10841                "W4A16 NVFP4 down rows geometry sel={} pairs={} act={} dst_raw={dst_raw:#x} \
10842                 n_sel={n_sel} total_pairs={total_pairs} in={in_f} out={out_f}",
10843                sel.len(),
10844                global_pairs.len(),
10845                activation_bf16.len(),
10846            )
10847            .into());
10848        }
10849        let f = self.func("qmatvec_nvfp4_bf16_sel_down_rows");
10850        let cfg = LaunchConfig {
10851            grid_dim: (out_f.div_ceil(2) as u32, n_sel as u32, 1),
10852            block_dim: (256, 1, 1),
10853            shared_mem_bytes: 0,
10854        };
10855        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10856        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10857        let __s_b = self.gpu.stream();
10858        let mut b = __s_b.launch_builder(&f);
10859        b.arg(bank)
10860            .arg(sel)
10861            .arg(global_pairs)
10862            .arg(activation_bf16)
10863            .arg(macros_down)
10864            .arg(&dst_raw)
10865            .arg(&inf)
10866            .arg(&outf)
10867            .arg(&ns)
10868            .arg(&rb)
10869            .arg(&es);
10870        unsafe {
10871            b.launch(cfg)?;
10872        }
10873        Ok(())
10874    }
10875
10876    /// Device-routed W4A16 down rows. Exactly one owner rank writes each global token/slot row
10877    /// into the root device's peer-accessible slab.
10878    #[allow(clippy::too_many_arguments)]
10879    pub fn qmatvec_nvfp4_bf16_ep_down_slots_raw(
10880        &self,
10881        bank: &CudaSlice<u8>,
10882        sel: &CudaSlice<i32>,
10883        activation_bf16: &CudaSlice<u8>,
10884        macros_down: &CudaSlice<f32>,
10885        dst_raw: u64,
10886        n_pairs: usize,
10887        in_f: usize,
10888        out_f: usize,
10889        owner_start: usize,
10890        owner_end: usize,
10891        row_bytes: usize,
10892        expert_stride: usize,
10893    ) -> Result<(), Box<dyn std::error::Error>> {
10894        if owner_start >= owner_end
10895            || !in_f.is_multiple_of(64)
10896            || sel.len() < n_pairs
10897            || activation_bf16.len() < 2 * n_pairs * in_f
10898            || dst_raw == 0
10899        {
10900            return Err(format!(
10901                "W4A16 NVFP4 device EP down-slot geometry sel={} act={} dst_raw={dst_raw:#x} \
10902                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
10903                sel.len(),
10904                activation_bf16.len(),
10905            )
10906            .into());
10907        }
10908        let f = self.func("qmatvec_nvfp4_bf16_ep_down_slots");
10909        let cfg = LaunchConfig {
10910            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
10911            block_dim: (256, 1, 1),
10912            shared_mem_bytes: 0,
10913        };
10914        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
10915        let (os, oe) = (owner_start as i32, owner_end as i32);
10916        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10917        let __s_b = self.gpu.stream();
10918        let mut b = __s_b.launch_builder(&f);
10919        b.arg(bank)
10920            .arg(sel)
10921            .arg(activation_bf16)
10922            .arg(macros_down)
10923            .arg(&dst_raw)
10924            .arg(&inf)
10925            .arg(&outf)
10926            .arg(&np)
10927            .arg(&os)
10928            .arg(&oe)
10929            .arg(&rb)
10930            .arg(&es);
10931        unsafe {
10932            b.launch(cfg)?;
10933        }
10934        Ok(())
10935    }
10936
10937    /// Pair-parallel multi-token twin of `qmatvec_nvfp4_bf16_ep_down_slots_raw`.
10938    #[allow(clippy::too_many_arguments)]
10939    pub fn qmatvec_nvfp4_bf16_ep_down_pairs_raw(
10940        &self,
10941        bank: &CudaSlice<u8>,
10942        sel: &CudaSlice<i32>,
10943        activation_bf16: &CudaSlice<u8>,
10944        macros_down: &CudaSlice<f32>,
10945        dst_raw: u64,
10946        n_pairs: usize,
10947        in_f: usize,
10948        out_f: usize,
10949        owner_start: usize,
10950        owner_end: usize,
10951        row_bytes: usize,
10952        expert_stride: usize,
10953    ) -> Result<(), Box<dyn std::error::Error>> {
10954        if owner_start >= owner_end
10955            || !in_f.is_multiple_of(64)
10956            || sel.len() < n_pairs
10957            || activation_bf16.len() < 2 * n_pairs * in_f
10958            || dst_raw == 0
10959        {
10960            return Err(format!(
10961                "W4A16 NVFP4 device EP down-pair geometry sel={} act={} dst_raw={dst_raw:#x} \
10962                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
10963                sel.len(),
10964                activation_bf16.len(),
10965            )
10966            .into());
10967        }
10968        let f = self.func("qmatvec_nvfp4_bf16_ep_down_pairs");
10969        let cfg = LaunchConfig {
10970            grid_dim: (out_f.div_ceil(2) as u32, n_pairs as u32, 1),
10971            block_dim: (256, 1, 1),
10972            shared_mem_bytes: 0,
10973        };
10974        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
10975        let (os, oe) = (owner_start as i32, owner_end as i32);
10976        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10977        let __s_b = self.gpu.stream();
10978        let mut b = __s_b.launch_builder(&f);
10979        b.arg(bank)
10980            .arg(sel)
10981            .arg(activation_bf16)
10982            .arg(macros_down)
10983            .arg(&dst_raw)
10984            .arg(&inf)
10985            .arg(&outf)
10986            .arg(&np)
10987            .arg(&os)
10988            .arg(&oe)
10989            .arg(&rb)
10990            .arg(&es);
10991        unsafe {
10992            b.launch(cfg)?;
10993        }
10994        Ok(())
10995    }
10996
10997    /// Host-expf W4A16 SwiGLU selected rows, rounded directly to BF16 for the down projection.
10998    #[allow(clippy::too_many_arguments)]
10999    pub fn silu_mul_scaled_host_expf_bf16_sel_into(
11000        &self,
11001        gate: &CudaSlice<f32>,
11002        up: &CudaSlice<f32>,
11003        gate_macros: &CudaSlice<f32>,
11004        up_macros: &CudaSlice<f32>,
11005        sel: &CudaSlice<i32>,
11006        limit: Option<f32>,
11007        output_bf16: &mut CudaSlice<u8>,
11008        n_per: usize,
11009        n_sel: usize,
11010    ) -> Result<(), Box<dyn std::error::Error>> {
11011        let n = n_per * n_sel;
11012        if sel.len() < n_sel || gate.len() < n || up.len() < n || output_bf16.len() < 2 * n {
11013            return Err(format!(
11014                "W4A16 selected activation geometry sel={} gate={} up={} out={} \
11015                 n_per={n_per} n_sel={n_sel}",
11016                sel.len(),
11017                gate.len(),
11018                up.len(),
11019                output_bf16.len(),
11020            )
11021            .into());
11022        }
11023        let (limit, has_limit) = match limit {
11024            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11025            Some(limit) => {
11026                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
11027            }
11028            None => (0.0f32, 0i32),
11029        };
11030        let f = self.func("silu_mul_scaled_host_expf_bf16_sel");
11031        let cfg = LaunchConfig::for_num_elems(n as u32);
11032        let (np, ns) = (n_per as i32, n_sel as i32);
11033        let __s_b = self.gpu.stream();
11034        let mut b = __s_b.launch_builder(&f);
11035        b.arg(gate)
11036            .arg(up)
11037            .arg(gate_macros)
11038            .arg(up_macros)
11039            .arg(sel)
11040            .arg(&limit)
11041            .arg(&has_limit)
11042            .arg(output_bf16)
11043            .arg(&np)
11044            .arg(&ns);
11045        unsafe {
11046            b.launch(cfg)?;
11047        }
11048        Ok(())
11049    }
11050
11051    /// Device-routed fixed token/slot W4A16 activation. Global expert ids are translated into
11052    /// rank-local macro rows only on the owning rank.
11053    #[allow(clippy::too_many_arguments)]
11054    pub fn silu_mul_scaled_host_expf_bf16_ep_slots_into(
11055        &self,
11056        gate: &CudaSlice<f32>,
11057        up: &CudaSlice<f32>,
11058        gate_macros: &CudaSlice<f32>,
11059        up_macros: &CudaSlice<f32>,
11060        sel: &CudaSlice<i32>,
11061        owner_start: usize,
11062        owner_end: usize,
11063        limit: Option<f32>,
11064        output_bf16: &mut CudaSlice<u8>,
11065        n_per: usize,
11066        n_pairs: usize,
11067    ) -> Result<(), Box<dyn std::error::Error>> {
11068        let n = n_per * n_pairs;
11069        if owner_start >= owner_end
11070            || sel.len() < n_pairs
11071            || gate.len() < n
11072            || up.len() < n
11073            || output_bf16.len() < 2 * n
11074        {
11075            return Err(format!(
11076                "W4A16 device EP activation geometry sel={} gate={} up={} out={} \
11077                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11078                sel.len(),
11079                gate.len(),
11080                up.len(),
11081                output_bf16.len(),
11082            )
11083            .into());
11084        }
11085        let (limit, has_limit) = match limit {
11086            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11087            Some(limit) => {
11088                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
11089            }
11090            None => (0.0f32, 0i32),
11091        };
11092        let f = self.func("silu_mul_scaled_host_expf_bf16_ep_slots");
11093        let cfg = LaunchConfig::for_num_elems(n as u32);
11094        let (np, pairs) = (n_per as i32, n_pairs as i32);
11095        let (os, oe) = (owner_start as i32, owner_end as i32);
11096        let __s_b = self.gpu.stream();
11097        let mut b = __s_b.launch_builder(&f);
11098        b.arg(gate)
11099            .arg(up)
11100            .arg(gate_macros)
11101            .arg(up_macros)
11102            .arg(sel)
11103            .arg(&limit)
11104            .arg(&has_limit)
11105            .arg(output_bf16)
11106            .arg(&np)
11107            .arg(&pairs)
11108            .arg(&os)
11109            .arg(&oe);
11110        unsafe {
11111            b.launch(cfg)?;
11112        }
11113        Ok(())
11114    }
11115
11116    /// Optional A8 host-expf SwiGLU over global fixed slots.
11117    #[allow(clippy::too_many_arguments)]
11118    pub fn silu_mul_scaled_host_expf_q8_ep_slots_into(
11119        &self,
11120        gate: &CudaSlice<f32>,
11121        up: &CudaSlice<f32>,
11122        gate_macros: &CudaSlice<f32>,
11123        up_macros: &CudaSlice<f32>,
11124        sel: &CudaSlice<i32>,
11125        owner_start: usize,
11126        owner_end: usize,
11127        limit: Option<f32>,
11128        output_q8: &mut CudaSlice<i8>,
11129        output_scales: &mut CudaSlice<f32>,
11130        n_per: usize,
11131        n_pairs: usize,
11132    ) -> Result<(), Box<dyn std::error::Error>> {
11133        let n = n_per * n_pairs;
11134        if owner_start >= owner_end
11135            || !n_per.is_multiple_of(32)
11136            || sel.len() < n_pairs
11137            || gate.len() < n
11138            || up.len() < n
11139            || output_q8.len() < n
11140            || output_scales.len() < n / 32
11141        {
11142            return Err(format!(
11143                "W4A8 device EP activation geometry sel={} gate={} up={} q8={} scales={} \
11144                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11145                sel.len(),
11146                gate.len(),
11147                up.len(),
11148                output_q8.len(),
11149                output_scales.len(),
11150            )
11151            .into());
11152        }
11153        let (limit, has_limit) = match limit {
11154            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11155            Some(limit) => {
11156                return Err(format!("W4A8 selected activation limit {limit} is invalid").into());
11157            }
11158            None => (0.0f32, 0i32),
11159        };
11160        let f = self.func("silu_mul_scaled_host_expf_q8_ep_slots");
11161        let warps = n / 32;
11162        let cfg = LaunchConfig {
11163            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
11164            block_dim: (128, 1, 1),
11165            shared_mem_bytes: 0,
11166        };
11167        let (np, pairs) = (n_per as i32, n_pairs as i32);
11168        let (os, oe) = (owner_start as i32, owner_end as i32);
11169        let __s_b = self.gpu.stream();
11170        let mut b = __s_b.launch_builder(&f);
11171        b.arg(gate)
11172            .arg(up)
11173            .arg(gate_macros)
11174            .arg(up_macros)
11175            .arg(sel)
11176            .arg(&limit)
11177            .arg(&has_limit)
11178            .arg(output_q8)
11179            .arg(output_scales)
11180            .arg(&np)
11181            .arg(&pairs)
11182            .arg(&os)
11183            .arg(&oe);
11184        unsafe {
11185            b.launch(cfg)?;
11186        }
11187        Ok(())
11188    }
11189
11190    /// W4A16 selected-expert down projection plus owner-local route combine. The destination may
11191    /// reside in the model engine's peer-accessible root pool.
11192    #[allow(clippy::too_many_arguments)]
11193    pub fn qmatvec_nvfp4_bf16_sel_down_fma_into(
11194        &self,
11195        bank: &CudaSlice<u8>,
11196        sel: &CudaSlice<i32>,
11197        activation_bf16: &CudaSlice<u8>,
11198        route_weights: &CudaSlice<f32>,
11199        macros_down: &CudaSlice<f32>,
11200        dst: &mut cudarc::driver::CudaViewMut<f32>,
11201        n_sel: usize,
11202        in_f: usize,
11203        out_f: usize,
11204        row_bytes: usize,
11205        expert_stride: usize,
11206    ) -> Result<(), Box<dyn std::error::Error>> {
11207        if !in_f.is_multiple_of(64)
11208            || sel.len() < n_sel
11209            || route_weights.len() < n_sel
11210            || activation_bf16.len() < 2 * n_sel * in_f
11211            || dst.len() < out_f
11212        {
11213            return Err(format!(
11214                "W4A16 NVFP4 down selected geometry sel={} act={} weights={} dst={} \
11215                 n_sel={n_sel} in={in_f} out={out_f}",
11216                sel.len(),
11217                activation_bf16.len(),
11218                route_weights.len(),
11219                dst.len(),
11220            )
11221            .into());
11222        }
11223        let f = self.func("qmatvec_nvfp4_bf16_sel_down_fma");
11224        let cfg = LaunchConfig {
11225            grid_dim: (out_f as u32, 1, 1),
11226            block_dim: (256, 1, 1),
11227            shared_mem_bytes: 0,
11228        };
11229        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11230        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11231        let __s_b = self.gpu.stream();
11232        let mut b = __s_b.launch_builder(&f);
11233        b.arg(bank)
11234            .arg(sel)
11235            .arg(activation_bf16)
11236            .arg(route_weights)
11237            .arg(macros_down)
11238            .arg(dst)
11239            .arg(&inf)
11240            .arg(&outf)
11241            .arg(&ns)
11242            .arg(&rb)
11243            .arg(&es);
11244        unsafe {
11245            b.launch(cfg)?;
11246        }
11247        Ok(())
11248    }
11249
11250    /// Device-routed t=1 W4A16 down projection plus owner-local weighted combine.
11251    #[allow(clippy::too_many_arguments)]
11252    pub fn qmatvec_nvfp4_bf16_ep_down_fma_into(
11253        &self,
11254        bank: &CudaSlice<u8>,
11255        sel: &CudaSlice<i32>,
11256        activation_bf16: &CudaSlice<u8>,
11257        route_weights: &CudaSlice<f32>,
11258        macros_down: &CudaSlice<f32>,
11259        dst: &mut cudarc::driver::CudaViewMut<f32>,
11260        n_pairs: usize,
11261        in_f: usize,
11262        out_f: usize,
11263        owner_start: usize,
11264        owner_end: usize,
11265        row_bytes: usize,
11266        expert_stride: usize,
11267    ) -> Result<(), Box<dyn std::error::Error>> {
11268        if owner_start >= owner_end
11269            || !in_f.is_multiple_of(64)
11270            || sel.len() < n_pairs
11271            || route_weights.len() < n_pairs
11272            || activation_bf16.len() < 2 * n_pairs * in_f
11273            || dst.len() < out_f
11274        {
11275            return Err(format!(
11276                "W4A16 device EP down-FMA geometry sel={} act={} weights={} dst={} \
11277                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
11278                sel.len(),
11279                activation_bf16.len(),
11280                route_weights.len(),
11281                dst.len(),
11282            )
11283            .into());
11284        }
11285        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
11286        let cfg = LaunchConfig {
11287            grid_dim: (out_f as u32, 1, 1),
11288            block_dim: (256, 1, 1),
11289            shared_mem_bytes: 0,
11290        };
11291        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11292        let (os, oe) = (owner_start as i32, owner_end as i32);
11293        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11294        let __s_b = self.gpu.stream();
11295        let mut b = __s_b.launch_builder(&f);
11296        b.arg(bank)
11297            .arg(sel)
11298            .arg(activation_bf16)
11299            .arg(route_weights)
11300            .arg(macros_down)
11301            .arg(dst)
11302            .arg(&inf)
11303            .arg(&outf)
11304            .arg(&np)
11305            .arg(&os)
11306            .arg(&oe)
11307            .arg(&rb)
11308            .arg(&es);
11309        unsafe {
11310            b.launch(cfg)?;
11311        }
11312        Ok(())
11313    }
11314
11315    /// Capture-safe twin of `qmatvec_nvfp4_bf16_ep_down_fma_into`. The destination is one
11316    /// rank-owned row inside a persistent root-device slab.
11317    #[allow(clippy::too_many_arguments)]
11318    pub fn qmatvec_nvfp4_bf16_ep_down_fma_raw(
11319        &self,
11320        bank: &CudaSlice<u8>,
11321        sel: &CudaSlice<i32>,
11322        activation_bf16: &CudaSlice<u8>,
11323        route_weights: &CudaSlice<f32>,
11324        macros_down: &CudaSlice<f32>,
11325        dst_raw: u64,
11326        n_pairs: usize,
11327        in_f: usize,
11328        out_f: usize,
11329        owner_start: usize,
11330        owner_end: usize,
11331        row_bytes: usize,
11332        expert_stride: usize,
11333    ) -> Result<(), Box<dyn std::error::Error>> {
11334        if dst_raw == 0
11335            || owner_start >= owner_end
11336            || !in_f.is_multiple_of(64)
11337            || sel.len() < n_pairs
11338            || route_weights.len() < n_pairs
11339            || activation_bf16.len() < 2 * n_pairs * in_f
11340        {
11341            return Err(format!(
11342                "W4A16 device EP raw down-FMA geometry sel={} act={} weights={} \
11343                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11344                 owner={owner_start}..{owner_end}",
11345                sel.len(),
11346                activation_bf16.len(),
11347                route_weights.len(),
11348            )
11349            .into());
11350        }
11351        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
11352        let cfg = LaunchConfig {
11353            grid_dim: (out_f as u32, 1, 1),
11354            block_dim: (256, 1, 1),
11355            shared_mem_bytes: 0,
11356        };
11357        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11358        let (os, oe) = (owner_start as i32, owner_end as i32);
11359        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11360        let __s_b = self.gpu.stream();
11361        let mut b = __s_b.launch_builder(&f);
11362        b.arg(bank)
11363            .arg(sel)
11364            .arg(activation_bf16)
11365            .arg(route_weights)
11366            .arg(macros_down)
11367            .arg(&dst_raw)
11368            .arg(&inf)
11369            .arg(&outf)
11370            .arg(&np)
11371            .arg(&os)
11372            .arg(&oe)
11373            .arg(&rb)
11374            .arg(&es);
11375        unsafe {
11376            b.launch(cfg)?;
11377        }
11378        Ok(())
11379    }
11380
11381    /// Optional A8 fixed-slot down rows. Each owner rank writes its selected pair rows directly
11382    /// into the root slot slab; the root applies route weights in canonical token/slot order.
11383    #[allow(clippy::too_many_arguments)]
11384    pub fn qmatvec_nvfp4_q8_ep_down_slots_raw(
11385        &self,
11386        bank: &CudaSlice<u8>,
11387        sel: &CudaSlice<i32>,
11388        aq: &CudaSlice<i8>,
11389        ad: &CudaSlice<f32>,
11390        macros_down: &CudaSlice<f32>,
11391        dst_raw: u64,
11392        n_pairs: usize,
11393        in_f: usize,
11394        out_f: usize,
11395        owner_start: usize,
11396        owner_end: usize,
11397        row_bytes: usize,
11398        expert_stride: usize,
11399    ) -> Result<(), Box<dyn std::error::Error>> {
11400        if dst_raw == 0
11401            || owner_start >= owner_end
11402            || !in_f.is_multiple_of(64)
11403            || sel.len() < n_pairs
11404            || aq.len() < n_pairs * in_f
11405            || ad.len() < n_pairs * (in_f / 32)
11406        {
11407            return Err(format!(
11408                "W4A8 device EP raw down-slot geometry sel={} aq={} ad={} \
11409                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11410                 owner={owner_start}..{owner_end}",
11411                sel.len(),
11412                aq.len(),
11413                ad.len(),
11414            )
11415            .into());
11416        }
11417        let f = self.func("qmatvec_nvfp4_q8_ep_down_slots");
11418        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11419        let cfg = LaunchConfig {
11420            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
11421            block_dim: (threads, 1, 1),
11422            shared_mem_bytes: 0,
11423        };
11424        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11425        let (os, oe) = (owner_start as i32, owner_end as i32);
11426        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11427        let __s_b = self.gpu.stream();
11428        let mut b = __s_b.launch_builder(&f);
11429        b.arg(bank)
11430            .arg(sel)
11431            .arg(aq)
11432            .arg(ad)
11433            .arg(macros_down)
11434            .arg(&dst_raw)
11435            .arg(&inf)
11436            .arg(&outf)
11437            .arg(&np)
11438            .arg(&os)
11439            .arg(&oe)
11440            .arg(&rb)
11441            .arg(&es);
11442        unsafe {
11443            b.launch(cfg)?;
11444        }
11445        Ok(())
11446    }
11447
11448    /// Historical A8 t=1 down + owner-local route combine into a persistent root row.
11449    #[allow(clippy::too_many_arguments)]
11450    pub fn qmatvec_nvfp4_q8_ep_down_fma_raw(
11451        &self,
11452        bank: &CudaSlice<u8>,
11453        sel: &CudaSlice<i32>,
11454        aq: &CudaSlice<i8>,
11455        ad: &CudaSlice<f32>,
11456        route_weights: &CudaSlice<f32>,
11457        macros_down: &CudaSlice<f32>,
11458        dst_raw: u64,
11459        n_pairs: usize,
11460        in_f: usize,
11461        out_f: usize,
11462        owner_start: usize,
11463        owner_end: usize,
11464        row_bytes: usize,
11465        expert_stride: usize,
11466    ) -> Result<(), Box<dyn std::error::Error>> {
11467        if dst_raw == 0
11468            || owner_start >= owner_end
11469            || !in_f.is_multiple_of(64)
11470            || sel.len() < n_pairs
11471            || aq.len() < n_pairs * in_f
11472            || ad.len() < n_pairs * (in_f / 32)
11473            || route_weights.len() < n_pairs
11474        {
11475            return Err(format!(
11476                "W4A8 device EP raw down-FMA geometry sel={} aq={} ad={} weights={} \
11477                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11478                 owner={owner_start}..{owner_end}",
11479                sel.len(),
11480                aq.len(),
11481                ad.len(),
11482                route_weights.len(),
11483            )
11484            .into());
11485        }
11486        let f = self.func("qmatvec_nvfp4_q8_ep_down_fma");
11487        let cfg = LaunchConfig {
11488            grid_dim: (out_f as u32, 1, 1),
11489            block_dim: (256, 1, 1),
11490            shared_mem_bytes: 0,
11491        };
11492        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11493        let (os, oe) = (owner_start as i32, owner_end as i32);
11494        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11495        let __s_b = self.gpu.stream();
11496        let mut b = __s_b.launch_builder(&f);
11497        b.arg(bank)
11498            .arg(sel)
11499            .arg(aq)
11500            .arg(ad)
11501            .arg(route_weights)
11502            .arg(macros_down)
11503            .arg(&dst_raw)
11504            .arg(&inf)
11505            .arg(&outf)
11506            .arg(&np)
11507            .arg(&os)
11508            .arg(&oe)
11509            .arg(&rb)
11510            .arg(&es);
11511        unsafe {
11512            b.launch(cfg)?;
11513        }
11514        Ok(())
11515    }
11516
11517    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
11518    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
11519    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
11520    /// takes the plain SiLU kernel.
11521    #[allow(clippy::too_many_arguments)]
11522    pub fn silu_mul_scaled_q8_1_sel_into(
11523        &self,
11524        gate: &CudaSlice<f32>,
11525        up: &CudaSlice<f32>,
11526        gmac: &CudaSlice<f32>,
11527        umac: &CudaSlice<f32>,
11528        sel: &CudaSlice<i32>,
11529        limit: Option<f32>,
11530        out_q: &mut CudaSlice<i8>,
11531        out_d: &mut CudaSlice<f32>,
11532        n_per: usize,
11533        n_sel: usize,
11534    ) -> Result<(), Box<dyn std::error::Error>> {
11535        let n = n_per * n_sel;
11536        if !n_per.is_multiple_of(32) || out_q.len() < n || out_d.len() < n / 32 {
11537            return Err(format!(
11538                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
11539                out_q.len(),
11540                out_d.len()
11541            )
11542            .into());
11543        }
11544        if let Some(limit) = limit {
11545            if limit <= 1e-6 {
11546                return Err(format!(
11547                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
11548                )
11549                .into());
11550            }
11551            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
11552            let cfg = LaunchConfig::for_num_elems(n as u32);
11553            let (np, ns) = (n_per as i32, n_sel as i32);
11554            let __s_b = self.gpu.stream();
11555            let mut b = __s_b.launch_builder(&f);
11556            b.arg(gate)
11557                .arg(up)
11558                .arg(gmac)
11559                .arg(umac)
11560                .arg(sel)
11561                .arg(&limit)
11562                .arg(out_q)
11563                .arg(out_d)
11564                .arg(&np)
11565                .arg(&ns);
11566            unsafe {
11567                b.launch(cfg)?;
11568            }
11569            return Ok(());
11570        }
11571        let f = self.func("silu_mul_scaled_q8_1_sel");
11572        let cfg = LaunchConfig::for_num_elems(n as u32);
11573        let (np, ns) = (n_per as i32, n_sel as i32);
11574        let __s_b = self.gpu.stream();
11575        let mut b = __s_b.launch_builder(&f);
11576        b.arg(gate)
11577            .arg(up)
11578            .arg(gmac)
11579            .arg(umac)
11580            .arg(sel)
11581            .arg(out_q)
11582            .arg(out_d)
11583            .arg(&np)
11584            .arg(&ns);
11585        unsafe {
11586            b.launch(cfg)?;
11587        }
11588        Ok(())
11589    }
11590
11591    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11592        Ok(self.gpu.stream().clone_htod(v)?)
11593    }
11594    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
11595        Ok(self.gpu.stream().clone_htod(v)?)
11596    }
11597    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
11598    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
11599        Ok(self.gpu.stream().clone_htod(v)?)
11600    }
11601    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
11602        Ok(self.gpu.stream().clone_htod(v)?)
11603    }
11604    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
11605    pub fn dtoh_view(
11606        &self,
11607        d: &cudarc::driver::CudaView<f32>,
11608    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11609        let v = self.gpu.stream().clone_dtoh(d)?;
11610        self.gpu.stream().synchronize()?;
11611        Ok(v)
11612    }
11613    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11614        let v = self.gpu.stream().clone_dtoh(d)?;
11615        self.gpu.stream().synchronize()?;
11616        Ok(v)
11617    }
11618    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
11619    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
11620    /// issuing them together avoids a second stream synchronization in every trunk layer.
11621    pub fn dtoh_pair(
11622        &self,
11623        a: &CudaSlice<f32>,
11624        b: &CudaSlice<f32>,
11625    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
11626        let av = self.gpu.stream().clone_dtoh(a)?;
11627        let bv = self.gpu.stream().clone_dtoh(b)?;
11628        self.gpu.stream().synchronize()?;
11629        Ok((av, bv))
11630    }
11631    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
11632    /// cross a shape-sensitive host boundary.
11633    pub fn dtoh_pair_views(
11634        &self,
11635        a: &cudarc::driver::CudaView<f32>,
11636        b: &cudarc::driver::CudaView<f32>,
11637    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
11638        let av = self.gpu.stream().clone_dtoh(a)?;
11639        let bv = self.gpu.stream().clone_dtoh(b)?;
11640        self.gpu.stream().synchronize()?;
11641        Ok((av, bv))
11642    }
11643    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
11644    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
11645        let v = self.gpu.stream().clone_dtoh(d)?;
11646        self.gpu.stream().synchronize()?;
11647        Ok(v)
11648    }
11649    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
11650    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
11651        let v = self.gpu.stream().clone_dtoh(d)?;
11652        self.gpu.stream().synchronize()?;
11653        Ok(v)
11654    }
11655    pub fn dtoh_u8_view(
11656        &self,
11657        d: &cudarc::driver::CudaView<u8>,
11658    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
11659        let v = self.gpu.stream().clone_dtoh(d)?;
11660        self.gpu.stream().synchronize()?;
11661        Ok(v)
11662    }
11663    /// D2H copy of the first `n` bytes of `d` into a pinned CACHEABLE host buffer: the
11664    /// prefix-cache host-tier demote primitive (lane/kv-host-spill-20260830). Queued on the
11665    /// worker stream and synchronized before returning, exactly like `dtoh_u8`: v1 keeps every
11666    /// host-tier copy on the CUDA owner thread (the HY3 spill law). SEAM (named, not built): an
11667    /// overlapped copy-stream variant would queue this on a dedicated D2H stream with an event
11668    /// handshake against the compute stream; build it only with a tick-stall receipt that says
11669    /// the sync copy is the bottleneck.
11670    pub fn dtoh_u8_into_pinned(
11671        &self,
11672        d: &CudaSlice<u8>,
11673        dst: &mut PinnedHostBuf,
11674        n: usize,
11675    ) -> Result<(), Box<dyn std::error::Error>> {
11676        if n > d.len() || n > dst.len() {
11677            return Err(format!(
11678                "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
11679                d.len(),
11680                dst.len(),
11681            )
11682            .into());
11683        }
11684        if n == 0 {
11685            return Ok(());
11686        }
11687        let host = &mut dst.as_mut_slice()[..n];
11688        self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
11689        self.gpu.stream().synchronize()?;
11690        Ok(())
11691    }
11692    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11693        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11694        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
11695        self.keep_if_capturing(&s);
11696        Ok(s)
11697    }
11698
11699    /// Take the pooled hc-glue decode workspace (MEMRA_HC_DECODE_WS) for one step's walk; put
11700    /// it back with [`Self::hyper_ws_put`]. A `None` here means another walk holds it (or it
11701    /// was never built) — the caller allocates fresh, which is always correct.
11702    pub(crate) fn hyper_ws_take(&self) -> Option<crate::hyper::HyperDecodeWs> {
11703        self.hyper_decode_ws.lock().unwrap().take()
11704    }
11705
11706    pub(crate) fn hyper_ws_put(&self, ws: crate::hyper::HyperDecodeWs) {
11707        *self.hyper_decode_ws.lock().unwrap() = Some(ws);
11708    }
11709
11710    // ---- Verify-walk workspace (MEMRA_VERIFY_WS, door W — see VerifyWs). ----
11711    // take/recycle are no-ops with the door off, so every OFF-arm call site is byte-for-byte
11712    // the shipped program (fresh alloc, ordinary async free). All pooled sites are
11713    // verify-walk-only by construction (rows-exact matmuls, the KDA Rows stash arm, the MoE
11714    // vrows staging), and the pool is per-engine = per-stream: recycle-then-reuse carries the
11715    // same stream-ordering guarantee the async allocator's free-then-alloc does.
11716
11717    /// Pool-or-alloc f32 scratch for a verify-walk site (uninit contract unchanged).
11718    pub(crate) fn vws_uninit(
11719        &self,
11720        n: usize,
11721    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11722        if verify_ws_on() {
11723            let mut ws = self.verify_ws.lock().unwrap();
11724            let ws = &mut *ws;
11725            if let Some(s) = VerifyWs::take(&mut ws.f32_pool, &mut ws.held_bytes, n) {
11726                if VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
11727                    eprintln!(
11728                        "[glm5-verify-ws] engaged: verify-walk buffers recycling through \
11729                         the size-keyed pool (MEMRA_GLM5_VERIFY_WS=1)"
11730                    );
11731                }
11732                return Ok(s);
11733            }
11734        }
11735        self.alloc_uninit::<f32>(n)
11736    }
11737
11738    /// Pool-or-alloc i8 scratch (q8_1 activation planes).
11739    pub(crate) fn vws_uninit_i8(
11740        &self,
11741        n: usize,
11742    ) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
11743        if verify_ws_on() {
11744            let mut ws = self.verify_ws.lock().unwrap();
11745            let ws = &mut *ws;
11746            if let Some(s) = VerifyWs::take(&mut ws.i8_pool, &mut ws.held_bytes, n) {
11747                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11748                return Ok(s);
11749            }
11750        }
11751        self.alloc_uninit::<i8>(n)
11752    }
11753
11754    /// Pool-or-alloc u64 scratch (the MoE vrows pointer tables).
11755    pub(crate) fn vws_uninit_u64(
11756        &self,
11757        n: usize,
11758    ) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
11759        if verify_ws_on() {
11760            let mut ws = self.verify_ws.lock().unwrap();
11761            let ws = &mut *ws;
11762            if let Some(s) = VerifyWs::take(&mut ws.u64_pool, &mut ws.held_bytes, n) {
11763                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11764                return Ok(s);
11765            }
11766        }
11767        self.alloc_uninit::<u64>(n)
11768    }
11769
11770    /// Return a dead verify-walk buffer to the pool (no-op with the door off: the buffer
11771    /// drops to the ordinary async free, the shipped program).
11772    pub(crate) fn vws_recycle(&self, s: CudaSlice<f32>) {
11773        if verify_ws_on() {
11774            let mut ws = self.verify_ws.lock().unwrap();
11775            let ws = &mut *ws;
11776            VerifyWs::put(&mut ws.f32_pool, &mut ws.held_bytes, s);
11777        }
11778    }
11779
11780    /// i8 twin of [`Self::vws_recycle`].
11781    pub(crate) fn vws_recycle_i8(&self, s: CudaSlice<i8>) {
11782        if verify_ws_on() {
11783            let mut ws = self.verify_ws.lock().unwrap();
11784            let ws = &mut *ws;
11785            VerifyWs::put(&mut ws.i8_pool, &mut ws.held_bytes, s);
11786        }
11787    }
11788
11789    /// u64 twin of [`Self::vws_recycle`].
11790    pub(crate) fn vws_recycle_u64(&self, s: CudaSlice<u64>) {
11791        if verify_ws_on() {
11792            let mut ws = self.verify_ws.lock().unwrap();
11793            let ws = &mut *ws;
11794            VerifyWs::put(&mut ws.u64_pool, &mut ws.held_bytes, s);
11795        }
11796    }
11797
11798    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
11799    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
11800    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
11801    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
11802    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
11803    /// back (or kept resident for graph replay). Returns the device token buffer.
11804    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
11805    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
11806    pub fn prob_of_token_device(
11807        &self,
11808        logits: &CudaSlice<f32>,
11809        tok: &CudaSlice<u32>,
11810        n_vocab: usize,
11811    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11812        let nb = ARGMAX_NB;
11813        let mut part = self.alloc_uninit::<f32>(nb)?;
11814        let mut p = self.alloc_uninit::<f32>(1)?;
11815        let f1 = self.func("prob_of_token_partial_f32");
11816        let cfg1 = LaunchConfig {
11817            grid_dim: (nb as u32, 1, 1),
11818            block_dim: (256, 1, 1),
11819            shared_mem_bytes: 0,
11820        };
11821        let nv = n_vocab as i32;
11822        let __s_b1 = self.gpu.stream();
11823        let mut b1 = __s_b1.launch_builder(&f1);
11824        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
11825        unsafe {
11826            b1.launch(cfg1)?;
11827        }
11828        let f2 = self.func("prob_of_token_final_f32");
11829        let cfg2 = LaunchConfig {
11830            grid_dim: (1, 1, 1),
11831            block_dim: (256, 1, 1),
11832            shared_mem_bytes: 0,
11833        };
11834        let nbi = nb as i32;
11835        let __s_b2 = self.gpu.stream();
11836        let mut b2 = __s_b2.launch_builder(&f2);
11837        b2.arg(&part).arg(&mut p).arg(&nbi);
11838        unsafe {
11839            b2.launch(cfg2)?;
11840        }
11841        Ok(p)
11842    }
11843
11844    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
11845    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
11846    /// where the host reads the p-min confidence between replays. Same kernels, same math.
11847    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
11848    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
11849    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
11850    pub fn prob_of_token_device_col(
11851        &self,
11852        logits: &CudaSlice<f32>,
11853        tok_all: &CudaSlice<u32>,
11854        tok_idx: usize,
11855        p_out: &mut CudaSlice<f32>,
11856        p_idx: usize,
11857        n_vocab: usize,
11858    ) -> Result<(), Box<dyn std::error::Error>> {
11859        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
11860        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
11861        let nb = ARGMAX_NB;
11862        let mut part = self.alloc_uninit::<f32>(nb)?;
11863        let f1 = self.func("prob_of_token_partial_f32");
11864        let cfg1 = LaunchConfig {
11865            grid_dim: (nb as u32, 1, 1),
11866            block_dim: (256, 1, 1),
11867            shared_mem_bytes: 0,
11868        };
11869        let nv = n_vocab as i32;
11870        let __s_b1 = self.gpu.stream();
11871        let mut b1 = __s_b1.launch_builder(&f1);
11872        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
11873        unsafe {
11874            b1.launch(cfg1)?;
11875        }
11876        let f2 = self.func("prob_of_token_final_f32");
11877        let cfg2 = LaunchConfig {
11878            grid_dim: (1, 1, 1),
11879            block_dim: (256, 1, 1),
11880            shared_mem_bytes: 0,
11881        };
11882        let nbi = nb as i32;
11883        let __s_b2 = self.gpu.stream();
11884        let mut b2 = __s_b2.launch_builder(&f2);
11885        b2.arg(&part).arg(&mut p_v).arg(&nbi);
11886        unsafe {
11887            b2.launch(cfg2)?;
11888        }
11889        Ok(())
11890    }
11891
11892    pub fn prob_of_token_device_into(
11893        &self,
11894        logits: &CudaSlice<f32>,
11895        tok: &CudaSlice<u32>,
11896        p_out: &mut CudaSlice<f32>,
11897        n_vocab: usize,
11898    ) -> Result<(), Box<dyn std::error::Error>> {
11899        let nb = ARGMAX_NB;
11900        let mut part = self.alloc_uninit::<f32>(nb)?;
11901        let f1 = self.func("prob_of_token_partial_f32");
11902        let cfg1 = LaunchConfig {
11903            grid_dim: (nb as u32, 1, 1),
11904            block_dim: (256, 1, 1),
11905            shared_mem_bytes: 0,
11906        };
11907        let nv = n_vocab as i32;
11908        let __s_b1 = self.gpu.stream();
11909        let mut b1 = __s_b1.launch_builder(&f1);
11910        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
11911        unsafe {
11912            b1.launch(cfg1)?;
11913        }
11914        let f2 = self.func("prob_of_token_final_f32");
11915        let cfg2 = LaunchConfig {
11916            grid_dim: (1, 1, 1),
11917            block_dim: (256, 1, 1),
11918            shared_mem_bytes: 0,
11919        };
11920        let nbi = nb as i32;
11921        let __s_b2 = self.gpu.stream();
11922        let mut b2 = __s_b2.launch_builder(&f2);
11923        b2.arg(&part).arg(p_out).arg(&nbi);
11924        unsafe {
11925            b2.launch(cfg2)?;
11926        }
11927        Ok(())
11928    }
11929
11930    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
11931    /// (graph-constant params, device-varying index). Capture-safe.
11932    pub fn u32_hist_append(
11933        &self,
11934        tok: &CudaSlice<u32>,
11935        hist: &mut CudaSlice<u32>,
11936        idx: &mut CudaSlice<i32>,
11937    ) -> Result<(), Box<dyn std::error::Error>> {
11938        let f = self.func("u32_hist_append");
11939        let cfg = LaunchConfig {
11940            grid_dim: (1, 1, 1),
11941            block_dim: (32, 1, 1),
11942            shared_mem_bytes: 0,
11943        };
11944        let __s_b = self.gpu.stream();
11945        let mut b = __s_b.launch_builder(&f);
11946        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
11947        unsafe {
11948            b.launch(cfg)?;
11949        }
11950        Ok(())
11951    }
11952
11953    pub fn argmax_token_device(
11954        &self,
11955        logits: &CudaSlice<f32>,
11956        n_vocab: usize,
11957    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
11958        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
11959        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
11960        Ok(tok)
11961    }
11962    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
11963    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
11964    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
11965    /// pointer is baked once and the token id never round-trips to host inside steady state. The
11966    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
11967    /// captured passes bake fixed addresses.
11968    pub fn argmax_token_device_into(
11969        &self,
11970        logits: &CudaSlice<f32>,
11971        tok: &mut CudaSlice<u32>,
11972        n_vocab: usize,
11973    ) -> Result<(), Box<dyn std::error::Error>> {
11974        let nb = ARGMAX_NB;
11975        let f1 = self.func("argmax_partial_f32");
11976        let f2 = self.func("argmax_final_f32");
11977        let mut guard = self.argmax_partials.lock().unwrap();
11978        if guard.is_none() {
11979            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
11980            // buffers carry no cudarc events (illegal inside capture).
11981            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
11982            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
11983            *guard = Some((pv, pi));
11984        }
11985        let (part_v, part_i) = guard.as_mut().unwrap();
11986        let nv = n_vocab as i32;
11987        let nbi = nb as i32;
11988        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
11989        let cfg1 = LaunchConfig {
11990            grid_dim: (nb as u32, 1, 1),
11991            block_dim: (256, 1, 1),
11992            shared_mem_bytes: 0,
11993        };
11994        let __s_b1 = self.gpu.stream();
11995        let mut b1 = __s_b1.launch_builder(&f1);
11996        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
11997        unsafe {
11998            b1.launch(cfg1)?;
11999        }
12000        // pass 2: one block reduces NB partials -> token_out[0].
12001        let cfg2 = LaunchConfig {
12002            grid_dim: (1, 1, 1),
12003            block_dim: (256, 1, 1),
12004            shared_mem_bytes: 0,
12005        };
12006        let __s_b2 = self.gpu.stream();
12007        let mut b2 = __s_b2.launch_builder(&f2);
12008        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
12009        unsafe {
12010            b2.launch(cfg2)?;
12011        }
12012        Ok(())
12013    }
12014    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
12015    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
12016    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
12017    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
12018    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
12019    pub fn argmax_token_device_col(
12020        &self,
12021        logits: &CudaSlice<f32>,
12022        col: usize,
12023        n_vocab: usize,
12024        toks: &mut CudaSlice<u32>,
12025        out_idx: usize,
12026    ) -> Result<(), Box<dyn std::error::Error>> {
12027        let nb = ARGMAX_NB;
12028        let f1 = self.func("argmax_partial_f32");
12029        let f2 = self.func("argmax_final_f32");
12030        let mut guard = self.argmax_partials.lock().unwrap();
12031        if guard.is_none() {
12032            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
12033            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
12034            *guard = Some((pv, pi));
12035        }
12036        let (part_v, part_i) = guard.as_mut().unwrap();
12037        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
12038        let nv = n_vocab as i32;
12039        let nbi = nb as i32;
12040        let cfg1 = LaunchConfig {
12041            grid_dim: (nb as u32, 1, 1),
12042            block_dim: (256, 1, 1),
12043            shared_mem_bytes: 0,
12044        };
12045        let __s_b1 = self.gpu.stream();
12046        let mut b1 = __s_b1.launch_builder(&f1);
12047        b1.arg(&col_view)
12048            .arg(&mut *part_v)
12049            .arg(&mut *part_i)
12050            .arg(&nv);
12051        unsafe {
12052            b1.launch(cfg1)?;
12053        }
12054        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
12055        let cfg2 = LaunchConfig {
12056            grid_dim: (1, 1, 1),
12057            block_dim: (256, 1, 1),
12058            shared_mem_bytes: 0,
12059        };
12060        let __s_b2 = self.gpu.stream();
12061        let mut b2 = __s_b2.launch_builder(&f2);
12062        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
12063        unsafe {
12064            b2.launch(cfg2)?;
12065        }
12066        Ok(())
12067    }
12068    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
12069    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
12070        Ok(self.gpu.stream().clone_htod(v)?)
12071    }
12072    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
12073        let v = self.gpu.stream().clone_dtoh(d)?;
12074        self.gpu.stream().synchronize()?;
12075        Ok(v)
12076    }
12077
12078    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
12079        let v = self.gpu.stream().clone_dtoh(d)?;
12080        self.gpu.stream().synchronize()?;
12081        Ok(v)
12082    }
12083    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
12084    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
12085    /// contents change every step, the address must not, so a captured graph can read it).
12086    pub fn htod_u32_into(
12087        &self,
12088        dst: &mut CudaSlice<u32>,
12089        src: &[u32],
12090    ) -> Result<(), Box<dyn std::error::Error>> {
12091        let mut view = dst.slice_mut(0..src.len());
12092        self.gpu.stream().memcpy_htod(src, &mut view)?;
12093        Ok(())
12094    }
12095
12096    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
12097    /// table without changing the device address its reconcile kernel consumes.
12098    pub fn htod_i32_into(
12099        &self,
12100        dst: &mut CudaSlice<i32>,
12101        src: &[i32],
12102    ) -> Result<(), Box<dyn std::error::Error>> {
12103        let mut view = dst.slice_mut(0..src.len());
12104        self.gpu.stream().memcpy_htod(src, &mut view)?;
12105        Ok(())
12106    }
12107
12108    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
12109        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
12110        self.keep_if_capturing(&s);
12111        Ok(s)
12112    }
12113    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
12114    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
12115    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
12116    pub fn embed_gather_device_into(
12117        &self,
12118        embd: &CudaSlice<u8>,
12119        token_d: &CudaSlice<u32>,
12120        x_out: &mut CudaSlice<f32>,
12121        n_embd: usize,
12122        qtype: i32,
12123        row_bytes: usize,
12124    ) -> Result<(), Box<dyn std::error::Error>> {
12125        let f = self.func("embed_gather_u32");
12126        let cfg = LaunchConfig {
12127            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
12128            block_dim: (256, 1, 1),
12129            shared_mem_bytes: 0,
12130        };
12131        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
12132        let __s_b = self.gpu.stream();
12133        let mut b = __s_b.launch_builder(&f);
12134        b.arg(embd)
12135            .arg(token_d)
12136            .arg(x_out)
12137            .arg(&ne)
12138            .arg(&qt)
12139            .arg(&rb);
12140        unsafe {
12141            b.launch(cfg)?;
12142        }
12143        Ok(())
12144    }
12145    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
12146    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
12147        let v = self.gpu.stream().clone_dtoh(d)?;
12148        self.gpu.stream().synchronize()?;
12149        Ok(v[0])
12150    }
12151    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
12152    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
12153    /// the counter value after the throwaway capture warmups corrupt it.
12154    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
12155    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
12156    /// copy (fine at stream-idle boundaries, poison mid-round).
12157    pub fn i32_set_k(
12158        &self,
12159        dst: &mut CudaSlice<i32>,
12160        v: i32,
12161    ) -> Result<(), Box<dyn std::error::Error>> {
12162        let f = self.func("i32_set_k");
12163        let cfg = LaunchConfig {
12164            grid_dim: (1, 1, 1),
12165            block_dim: (1, 1, 1),
12166            shared_mem_bytes: 0,
12167        };
12168        let idx = 0i32;
12169        let __s_b = self.gpu.stream();
12170        let mut b = __s_b.launch_builder(&f);
12171        b.arg(dst).arg(&v).arg(&idx);
12172        unsafe {
12173            b.launch(cfg)?;
12174        }
12175        Ok(())
12176    }
12177
12178    pub fn set_i32_one(
12179        &self,
12180        d: &mut CudaSlice<i32>,
12181        v: i32,
12182    ) -> Result<(), Box<dyn std::error::Error>> {
12183        self.gpu.stream().memcpy_htod(&[v], d)?;
12184        Ok(())
12185    }
12186    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
12187    /// during priming / capture-state restore.
12188    pub fn set_u32_one(
12189        &self,
12190        d: &mut CudaSlice<u32>,
12191        v: u32,
12192    ) -> Result<(), Box<dyn std::error::Error>> {
12193        self.gpu.stream().memcpy_htod(&[v], d)?;
12194        Ok(())
12195    }
12196    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
12197    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
12198        let v = self.gpu.stream().clone_dtoh(d)?;
12199        self.gpu.stream().synchronize()?;
12200        Ok(v[0])
12201    }
12202    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
12203    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
12204        Ok(self.gpu.stream().clone_htod(bytes)?)
12205    }
12206    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
12207    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
12208    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
12209    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
12210    pub fn embed_gather_device(
12211        &self,
12212        embd: &CudaSlice<u8>,
12213        token_d: &CudaSlice<u32>,
12214        n_embd: usize,
12215        qtype: i32,
12216        row_bytes: usize,
12217    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12218        let f = self.func("embed_gather_u32");
12219        let mut x = self.alloc_uninit::<f32>(n_embd)?;
12220        let cfg = LaunchConfig {
12221            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
12222            block_dim: (256, 1, 1),
12223            shared_mem_bytes: 0,
12224        };
12225        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
12226        let __s_b = self.gpu.stream();
12227        let mut b = __s_b.launch_builder(&f);
12228        b.arg(embd)
12229            .arg(token_d)
12230            .arg(&mut x)
12231            .arg(&ne)
12232            .arg(&qt)
12233            .arg(&rb);
12234        unsafe {
12235            b.launch(cfg)?;
12236        }
12237        Ok(x)
12238    }
12239
12240    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
12241    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
12242    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
12243    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
12244    pub fn embed_gather_device_t(
12245        &self,
12246        embd: &CudaSlice<u8>,
12247        tokens: &[u32],
12248        n_embd: usize,
12249        qtype: i32,
12250        row_bytes: usize,
12251    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12252        let t = tokens.len();
12253        let tok_d = self.gpu.stream().clone_htod(tokens)?;
12254        let f = self.func("embed_gather_u32_t");
12255        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12256        let cfg = LaunchConfig {
12257            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12258            block_dim: (256, 1, 1),
12259            shared_mem_bytes: 0,
12260        };
12261        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12262        let __s_b = self.gpu.stream();
12263        let mut b = __s_b.launch_builder(&f);
12264        b.arg(embd)
12265            .arg(&tok_d)
12266            .arg(&mut x)
12267            .arg(&ne)
12268            .arg(&qt)
12269            .arg(&rb)
12270            .arg(&ti);
12271        unsafe {
12272            b.launch(cfg)?;
12273        }
12274        Ok(x)
12275    }
12276
12277    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
12278    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
12279    /// as embed_gather_device_t — bit-identical rows.
12280    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
12281    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
12282    pub fn embed_gather_device_tv(
12283        &self,
12284        embd: &CudaSlice<u8>,
12285        tok_v: &cudarc::driver::CudaView<u32>,
12286        t: usize,
12287        n_embd: usize,
12288        qtype: i32,
12289        row_bytes: usize,
12290    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12291        let f = self.func("embed_gather_u32_t");
12292        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12293        let cfg = LaunchConfig {
12294            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12295            block_dim: (256, 1, 1),
12296            shared_mem_bytes: 0,
12297        };
12298        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12299        let __s_b = self.gpu.stream();
12300        let mut b = __s_b.launch_builder(&f);
12301        b.arg(embd)
12302            .arg(tok_v)
12303            .arg(&mut x)
12304            .arg(&ne)
12305            .arg(&qt)
12306            .arg(&rb)
12307            .arg(&ti);
12308        unsafe {
12309            b.launch(cfg)?;
12310        }
12311        Ok(x)
12312    }
12313
12314    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
12315    pub fn embed_gather_device_td(
12316        &self,
12317        embd: &CudaSlice<u8>,
12318        tok_d: &CudaSlice<u32>,
12319        t: usize,
12320        n_embd: usize,
12321        qtype: i32,
12322        row_bytes: usize,
12323    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12324        let f = self.func("embed_gather_u32_t");
12325        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12326        let cfg = LaunchConfig {
12327            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12328            block_dim: (256, 1, 1),
12329            shared_mem_bytes: 0,
12330        };
12331        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12332        let __s_b = self.gpu.stream();
12333        let mut b = __s_b.launch_builder(&f);
12334        b.arg(embd)
12335            .arg(tok_d)
12336            .arg(&mut x)
12337            .arg(&ne)
12338            .arg(&qt)
12339            .arg(&rb)
12340            .arg(&ti);
12341        unsafe {
12342            b.launch(cfg)?;
12343        }
12344        Ok(x)
12345    }
12346
12347    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
12348    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
12349    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
12350    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
12351    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
12352    #[inline]
12353    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
12354    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
12355        if self
12356            .capture_keep_on
12357            .load(std::sync::atomic::Ordering::Relaxed)
12358        {
12359            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
12360        }
12361    }
12362
12363    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
12364        &self,
12365        n: usize,
12366    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
12367        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12368        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
12369        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
12370        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
12371        // not cover engine-internal buffers). Debug-only: massive launch overhead.
12372        {
12373            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12374            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
12375                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
12376                use cudarc::driver::DevicePtrMut;
12377                let n_bytes = s.len() * std::mem::size_of::<T>();
12378                let stream = self.gpu.stream();
12379                let (p_, _g) = s.device_ptr_mut(&stream);
12380                unsafe {
12381                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
12382                        .result()?;
12383                }
12384            }
12385        }
12386        self.keep_if_capturing(&s);
12387        Ok(s)
12388    }
12389
12390    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
12391    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
12392    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
12393    /// consumers alloc through this (m=1 decode arms).
12394    pub fn uninit_q8_pair(
12395        &self,
12396        n: usize,
12397    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12398        Ok((
12399            self.alloc_uninit::<i8>(n)?,
12400            self.alloc_uninit::<f32>(n / 32)?,
12401        ))
12402    }
12403
12404    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12405        self.alloc_uninit::<f32>(n)
12406    }
12407
12408    /// i8 uninitialized scratch (same contract as `uninit`).
12409    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12410        self.alloc_uninit::<i8>(n)
12411    }
12412
12413    /// i32 uninitialized scratch (same contract as `uninit`) — the DSA indexer's position lists.
12414    pub fn uninit_i32(&self, n: usize) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
12415        self.alloc_uninit::<i32>(n)
12416    }
12417
12418    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
12419    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
12420    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
12421    #[allow(clippy::too_many_arguments)]
12422    pub fn rms_norm3(
12423        &self,
12424        x: &CudaSlice<f32>,
12425        w0: &CudaSlice<f32>,
12426        w1: &CudaSlice<f32>,
12427        w2: &CudaSlice<f32>,
12428        d0: &mut CudaSlice<f32>,
12429        d1: &mut CudaSlice<f32>,
12430        d2: &mut CudaSlice<f32>,
12431        ncols: usize,
12432        nrows: usize,
12433        eps: f32,
12434    ) -> Result<(), Box<dyn std::error::Error>> {
12435        let f = self.func("rms_norm3_f32");
12436        let cfg = LaunchConfig {
12437            grid_dim: (nrows as u32, 1, 1),
12438            block_dim: (rms_block(), 1, 1),
12439            shared_mem_bytes: 0,
12440        };
12441        let (nc, e) = (ncols as i32, eps);
12442        let __s_b = self.gpu.stream();
12443        let mut b = __s_b.launch_builder(&f);
12444        b.arg(x)
12445            .arg(w0)
12446            .arg(w1)
12447            .arg(w2)
12448            .arg(d0)
12449            .arg(d1)
12450            .arg(d2)
12451            .arg(&nc)
12452            .arg(&e);
12453        unsafe {
12454            b.launch(cfg)?;
12455        }
12456        Ok(())
12457    }
12458
12459    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
12460    #[allow(clippy::too_many_arguments)]
12461    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
12462    /// piggybacks on the same conditions.
12463    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
12464        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12465        *WARP_ON.get_or_init(|| {
12466            std::env::var("MEMRA_QKVNORM_W")
12467                .map(|v| v != "0")
12468                .unwrap_or(true)
12469        }) && ncols.is_multiple_of(4)
12470            && rows >= 64
12471    }
12472
12473    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
12474    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
12475    #[allow(clippy::too_many_arguments)]
12476    pub fn rms_norm_qkv_w4b(
12477        &self,
12478        q: &CudaSlice<f32>,
12479        k: &CudaSlice<f32>,
12480        v: &CudaSlice<f32>,
12481        wq: &CudaSlice<f32>,
12482        wk: &CudaSlice<f32>,
12483        wv: &CudaSlice<f32>,
12484        dq: &mut CudaSlice<f32>,
12485        dk: &mut CudaSlice<f32>,
12486        dv: &mut CudaSlice<f32>,
12487        dvb: &mut CudaSlice<u8>,
12488        ncols: usize,
12489        rq: usize,
12490        rk: usize,
12491        eps: f32,
12492        vf16: bool,
12493    ) -> Result<(), Box<dyn std::error::Error>> {
12494        assert!(ncols.is_multiple_of(4) && rq + 2 * rk >= 64);
12495        let f = self.func("rms_norm_qkv_w4b_f32");
12496        let rows = (rq + 2 * rk) as u32;
12497        let cfg = LaunchConfig {
12498            grid_dim: (rows.div_ceil(8), 1, 1),
12499            block_dim: (256, 1, 1),
12500            shared_mem_bytes: 0,
12501        };
12502        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
12503        let vf = vf16 as i32;
12504        let __s_b = self.gpu.stream();
12505        let mut b = __s_b.launch_builder(&f);
12506        b.arg(q)
12507            .arg(k)
12508            .arg(v)
12509            .arg(wq)
12510            .arg(wk)
12511            .arg(wv)
12512            .arg(dq)
12513            .arg(dk)
12514            .arg(dv)
12515            .arg(&mut *dvb)
12516            .arg(&nc)
12517            .arg(&rqi)
12518            .arg(&rki)
12519            .arg(&rvi)
12520            .arg(&e)
12521            .arg(&vf);
12522        unsafe {
12523            b.launch(cfg)?;
12524        }
12525        Ok(())
12526    }
12527
12528    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
12529    pub fn rms_norm_qkv(
12530        &self,
12531        q: &CudaSlice<f32>,
12532        k: &CudaSlice<f32>,
12533        v: &CudaSlice<f32>,
12534        wq: &CudaSlice<f32>,
12535        wk: &CudaSlice<f32>,
12536        wv: &CudaSlice<f32>,
12537        dq: &mut CudaSlice<f32>,
12538        dk: &mut CudaSlice<f32>,
12539        dv: &mut CudaSlice<f32>,
12540        ncols: usize,
12541        rq: usize,
12542        rk: usize,
12543        eps: f32,
12544    ) -> Result<(), Box<dyn std::error::Error>> {
12545        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
12546        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
12547        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
12548        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12549        let warp_on = *WARP_ON.get_or_init(|| {
12550            std::env::var("MEMRA_QKVNORM_W")
12551                .map(|v| v != "0")
12552                .unwrap_or(true)
12553        });
12554        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
12555        // replay numerics are untouched on every model; only prefill depth takes the new config.
12556        if warp_on && ncols.is_multiple_of(4) && rq + 2 * rk >= 64 {
12557            let f = self.func("rms_norm_qkv_w4_f32");
12558            let rows = (rq + 2 * rk) as u32;
12559            let cfg = LaunchConfig {
12560                grid_dim: (rows.div_ceil(8), 1, 1),
12561                block_dim: (256, 1, 1),
12562                shared_mem_bytes: 0,
12563            };
12564            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
12565            let __s_b = self.gpu.stream();
12566            let mut b = __s_b.launch_builder(&f);
12567            b.arg(q)
12568                .arg(k)
12569                .arg(v)
12570                .arg(wq)
12571                .arg(wk)
12572                .arg(wv)
12573                .arg(dq)
12574                .arg(dk)
12575                .arg(dv)
12576                .arg(&nc)
12577                .arg(&rqi)
12578                .arg(&rki)
12579                .arg(&rvi)
12580                .arg(&e);
12581            unsafe {
12582                b.launch(cfg)?;
12583            }
12584            return Ok(());
12585        }
12586        let f = self.func("rms_norm_qkv_f32");
12587        let grid = (rq + 2 * rk) as u32;
12588        let cfg = LaunchConfig {
12589            grid_dim: (grid, 1, 1),
12590            block_dim: (rms_block(), 1, 1),
12591            shared_mem_bytes: 0,
12592        };
12593        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
12594        let __s_b = self.gpu.stream();
12595        let mut b = __s_b.launch_builder(&f);
12596        b.arg(q)
12597            .arg(k)
12598            .arg(v)
12599            .arg(wq)
12600            .arg(wk)
12601            .arg(wv)
12602            .arg(dq)
12603            .arg(dk)
12604            .arg(dv)
12605            .arg(&nc)
12606            .arg(&rqi)
12607            .arg(&rki)
12608            .arg(&e);
12609        unsafe {
12610            b.launch(cfg)?;
12611        }
12612        Ok(())
12613    }
12614
12615    /// gemma4 fused pair of rms_norms over two different inputs (same width).
12616    #[allow(clippy::too_many_arguments)]
12617    pub fn rms_norm2x(
12618        &self,
12619        a: &CudaSlice<f32>,
12620        bb: &CudaSlice<f32>,
12621        wa: &CudaSlice<f32>,
12622        wb: &CudaSlice<f32>,
12623        da: &mut CudaSlice<f32>,
12624        db: &mut CudaSlice<f32>,
12625        ncols: usize,
12626        nrows: usize,
12627        eps: f32,
12628    ) -> Result<(), Box<dyn std::error::Error>> {
12629        let f = self.func("rms_norm2x_f32");
12630        let cfg = LaunchConfig {
12631            grid_dim: (2 * nrows as u32, 1, 1),
12632            block_dim: (rms_block(), 1, 1),
12633            shared_mem_bytes: 0,
12634        };
12635        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
12636        let __s_b = self.gpu.stream();
12637        let mut b = __s_b.launch_builder(&f);
12638        b.arg(a)
12639            .arg(bb)
12640            .arg(wa)
12641            .arg(wb)
12642            .arg(da)
12643            .arg(db)
12644            .arg(&nc)
12645            .arg(&nr)
12646            .arg(&e);
12647        unsafe {
12648            b.launch(cfg)?;
12649        }
12650        Ok(())
12651    }
12652
12653    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
12654    pub fn softcap(
12655        &self,
12656        y: &mut CudaSlice<f32>,
12657        cap: f32,
12658        n: usize,
12659    ) -> Result<(), Box<dyn std::error::Error>> {
12660        let f = self.func("softcap_f32");
12661        let cfg = LaunchConfig::for_num_elems(n as u32);
12662        let ni = n as i32;
12663        let __s_b = self.gpu.stream();
12664        let mut b = __s_b.launch_builder(&f);
12665        b.arg(y).arg(&cap).arg(&ni);
12666        unsafe {
12667            b.launch(cfg)?;
12668        }
12669        Ok(())
12670    }
12671
12672    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
12673    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
12674    pub fn mask_ids_rows(
12675        &self,
12676        y: &mut CudaSlice<f32>,
12677        ids: &CudaSlice<i32>,
12678        n_ids: usize,
12679        n_vocab: usize,
12680        t: usize,
12681    ) -> Result<(), Box<dyn std::error::Error>> {
12682        let f = self.func("mask_ids_rows_f32");
12683        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
12684        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
12685        let __s_b = self.gpu.stream();
12686        let mut b = __s_b.launch_builder(&f);
12687        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
12688        unsafe {
12689            b.launch(cfg)?;
12690        }
12691        Ok(())
12692    }
12693
12694    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
12695    #[allow(clippy::too_many_arguments)]
12696    pub fn add_scale_rms_norm(
12697        &self,
12698        a: &CudaSlice<f32>,
12699        b_in: &CudaSlice<f32>,
12700        c: f32,
12701        w: &CudaSlice<f32>,
12702        res: &mut CudaSlice<f32>,
12703        dst: &mut CudaSlice<f32>,
12704        ncols: usize,
12705        nrows: usize,
12706        eps: f32,
12707    ) -> Result<(), Box<dyn std::error::Error>> {
12708        let f = self.func("add_scale_rms_norm_f32");
12709        let cfg = LaunchConfig {
12710            grid_dim: (nrows as u32, 1, 1),
12711            block_dim: (rms_block(), 1, 1),
12712            shared_mem_bytes: 0,
12713        };
12714        let (nc, e2) = (ncols as i32, eps);
12715        let __s_b = self.gpu.stream();
12716        let mut b = __s_b.launch_builder(&f);
12717        b.arg(a)
12718            .arg(b_in)
12719            .arg(&c)
12720            .arg(w)
12721            .arg(res)
12722            .arg(dst)
12723            .arg(&nc)
12724            .arg(&e2);
12725        unsafe {
12726            b.launch(cfg)?;
12727        }
12728        Ok(())
12729    }
12730
12731    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
12732    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
12733    #[allow(clippy::too_many_arguments)]
12734    pub fn add_scale_rms_norm_q8_1(
12735        &self,
12736        a: &CudaSlice<f32>,
12737        b_in: &CudaSlice<f32>,
12738        c: f32,
12739        w: &CudaSlice<f32>,
12740        res: &mut CudaSlice<f32>,
12741        ncols: usize,
12742        nrows: usize,
12743        eps: f32,
12744    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12745        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12746        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12747        let (nc, e2) = (ncols as i32, eps);
12748        if Self::pdl_on() && Self::pdl_wb_on() {
12749            {
12750                use cudarc::driver::{DevicePtr, DevicePtrMut};
12751                let s = &self.gpu.stream();
12752                let (pa, _g0) = a.device_ptr(s);
12753                let (pb, _g1) = b_in.device_ptr(s);
12754                let (pw, _g2) = w.device_ptr(s);
12755                let (pr, _g3) = res.device_ptr_mut(s);
12756                let (pq, _g4) = out_q.device_ptr_mut(s);
12757                let (pd, _g5) = out_d.device_ptr_mut(s);
12758                let mut ps = [
12759                    &pa as *const _ as *mut std::ffi::c_void,
12760                    &pb as *const _ as *mut _,
12761                    &c as *const _ as *mut _,
12762                    &pw as *const _ as *mut _,
12763                    &pr as *const _ as *mut _,
12764                    &pq as *const _ as *mut _,
12765                    &pd as *const _ as *mut _,
12766                    &nc as *const _ as *mut _,
12767                    &e2 as *const _ as *mut _,
12768                ];
12769                unsafe {
12770                    self.launch_pdl(
12771                        "add_scale_rms_norm_q8_1",
12772                        (nrows as u32, 1, 1),
12773                        (rms_block(), 1, 1),
12774                        &mut ps,
12775                    )?;
12776                }
12777            }
12778            return Ok((out_q, out_d));
12779        }
12780        let f = self.func("add_scale_rms_norm_q8_1");
12781        let cfg = LaunchConfig {
12782            grid_dim: (nrows as u32, 1, 1),
12783            block_dim: (rms_block(), 1, 1),
12784            shared_mem_bytes: 0,
12785        };
12786        let __s_b = self.gpu.stream();
12787        let mut b = __s_b.launch_builder(&f);
12788        b.arg(a)
12789            .arg(b_in)
12790            .arg(&c)
12791            .arg(w)
12792            .arg(res)
12793            .arg(&mut out_q)
12794            .arg(&mut out_d)
12795            .arg(&nc)
12796            .arg(&e2);
12797        unsafe {
12798            b.launch(cfg)?;
12799        }
12800        Ok((out_q, out_d))
12801    }
12802
12803    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
12804    #[allow(clippy::too_many_arguments)]
12805    pub fn add_scale_rms_norm_q8_1_into(
12806        &self,
12807        a: &CudaSlice<f32>,
12808        b_in: &CudaSlice<f32>,
12809        c: f32,
12810        w: &CudaSlice<f32>,
12811        res: &mut CudaSlice<f32>,
12812        ncols: usize,
12813        nrows: usize,
12814        eps: f32,
12815        out_q: &mut CudaSlice<i8>,
12816        out_d: &mut CudaSlice<f32>,
12817    ) -> Result<(), Box<dyn std::error::Error>> {
12818        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
12819        let (nc, e2) = (ncols as i32, eps);
12820        if Self::pdl_on() && Self::pdl_wb_on() {
12821            use cudarc::driver::{DevicePtr, DevicePtrMut};
12822            let s = &self.gpu.stream();
12823            let (pa, _g0) = a.device_ptr(s);
12824            let (pb, _g1) = b_in.device_ptr(s);
12825            let (pw, _g2) = w.device_ptr(s);
12826            let (pr, _g3) = res.device_ptr_mut(s);
12827            let (pq, _g4) = out_q.device_ptr_mut(s);
12828            let (pd, _g5) = out_d.device_ptr_mut(s);
12829            let mut ps = [
12830                &pa as *const _ as *mut std::ffi::c_void,
12831                &pb as *const _ as *mut _,
12832                &c as *const _ as *mut _,
12833                &pw as *const _ as *mut _,
12834                &pr as *const _ as *mut _,
12835                &pq as *const _ as *mut _,
12836                &pd as *const _ as *mut _,
12837                &nc as *const _ as *mut _,
12838                &e2 as *const _ as *mut _,
12839            ];
12840            unsafe {
12841                self.launch_pdl(
12842                    "add_scale_rms_norm_q8_1",
12843                    (nrows as u32, 1, 1),
12844                    (rms_block(), 1, 1),
12845                    &mut ps,
12846                )?;
12847            }
12848            return Ok(());
12849        }
12850        let f = self.func("add_scale_rms_norm_q8_1");
12851        let cfg = LaunchConfig {
12852            grid_dim: (nrows as u32, 1, 1),
12853            block_dim: (rms_block(), 1, 1),
12854            shared_mem_bytes: 0,
12855        };
12856        let __s_b = self.gpu.stream();
12857        let mut b = __s_b.launch_builder(&f);
12858        b.arg(a)
12859            .arg(b_in)
12860            .arg(&c)
12861            .arg(w)
12862            .arg(res)
12863            .arg(&mut *out_q)
12864            .arg(&mut *out_d)
12865            .arg(&nc)
12866            .arg(&e2);
12867        unsafe {
12868            b.launch(cfg)?;
12869        }
12870        Ok(())
12871    }
12872
12873    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
12874    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
12875    #[allow(clippy::too_many_arguments)]
12876    pub fn rms_pre_add_scale_rms_norm_q8_1(
12877        &self,
12878        a: &CudaSlice<f32>,
12879        wa: &CudaSlice<f32>,
12880        b_in: &CudaSlice<f32>,
12881        c: f32,
12882        w: &CudaSlice<f32>,
12883        res: &mut CudaSlice<f32>,
12884        ncols: usize,
12885        nrows: usize,
12886        eps: f32,
12887    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12888        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12889        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12890        let (nc, e2) = (ncols as i32, eps);
12891        if Self::pdl_on() {
12892            {
12893                use cudarc::driver::{DevicePtr, DevicePtrMut};
12894                let s = &self.gpu.stream();
12895                let (pa, _g0) = a.device_ptr(s);
12896                let (pwa, _g1) = wa.device_ptr(s);
12897                let (pb, _g2) = b_in.device_ptr(s);
12898                let (pw, _g3) = w.device_ptr(s);
12899                let (pr, _g4) = res.device_ptr_mut(s);
12900                let (pq, _g5) = out_q.device_ptr_mut(s);
12901                let (pd, _g6) = out_d.device_ptr_mut(s);
12902                let mut ps = [
12903                    &pa as *const _ as *mut std::ffi::c_void,
12904                    &pwa as *const _ as *mut _,
12905                    &pb as *const _ as *mut _,
12906                    &c as *const _ as *mut _,
12907                    &pw as *const _ as *mut _,
12908                    &pr as *const _ as *mut _,
12909                    &pq as *const _ as *mut _,
12910                    &pd as *const _ as *mut _,
12911                    &nc as *const _ as *mut _,
12912                    &e2 as *const _ as *mut _,
12913                ];
12914                unsafe {
12915                    self.launch_pdl(
12916                        "rms_pre_add_scale_rms_norm_q8_1",
12917                        (nrows as u32, 1, 1),
12918                        (rms_block(), 1, 1),
12919                        &mut ps,
12920                    )?;
12921                }
12922            }
12923            return Ok((out_q, out_d));
12924        }
12925        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
12926        let cfg = LaunchConfig {
12927            grid_dim: (nrows as u32, 1, 1),
12928            block_dim: (rms_block(), 1, 1),
12929            shared_mem_bytes: 0,
12930        };
12931        let __s_b = self.gpu.stream();
12932        let mut b = __s_b.launch_builder(&f);
12933        b.arg(a)
12934            .arg(wa)
12935            .arg(b_in)
12936            .arg(&c)
12937            .arg(w)
12938            .arg(res)
12939            .arg(&mut out_q)
12940            .arg(&mut out_d)
12941            .arg(&nc)
12942            .arg(&e2);
12943        unsafe {
12944            b.launch(cfg)?;
12945        }
12946        Ok((out_q, out_d))
12947    }
12948
12949    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
12950    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
12951    pub fn gelu_tanh_mul_q8_1(
12952        &self,
12953        gate: &CudaSlice<f32>,
12954        up: &cudarc::driver::CudaView<f32>,
12955        act: &mut CudaSlice<f32>,
12956        ncols: usize,
12957        nrows: usize,
12958    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12959        debug_assert!(ncols.is_multiple_of(128));
12960        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12961        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12962        let nc = ncols as i32;
12963        if Self::pdl_on() {
12964            {
12965                use cudarc::driver::{DevicePtr, DevicePtrMut};
12966                let s = &self.gpu.stream();
12967                let (pg, _g0) = gate.device_ptr(s);
12968                let (pu, _g1) = up.device_ptr(s);
12969                let (pact, _g2) = act.device_ptr_mut(s);
12970                let (pq, _g3) = out_q.device_ptr_mut(s);
12971                let (pd, _g4) = out_d.device_ptr_mut(s);
12972                let mut ps = [
12973                    &pg as *const _ as *mut std::ffi::c_void,
12974                    &pu as *const _ as *mut _,
12975                    &pact as *const _ as *mut _,
12976                    &pq as *const _ as *mut _,
12977                    &pd as *const _ as *mut _,
12978                    &nc as *const _ as *mut _,
12979                ];
12980                unsafe {
12981                    self.launch_pdl(
12982                        "gelu_tanh_mul_q8_1",
12983                        (nrows as u32, 1, 1),
12984                        (rms_block(), 1, 1),
12985                        &mut ps,
12986                    )?;
12987                }
12988            }
12989            return Ok((out_q, out_d));
12990        }
12991        let f = self.func("gelu_tanh_mul_q8_1");
12992        let cfg = LaunchConfig {
12993            grid_dim: (nrows as u32, 1, 1),
12994            block_dim: (rms_block(), 1, 1),
12995            shared_mem_bytes: 0,
12996        };
12997        let __s_b = self.gpu.stream();
12998        let mut b = __s_b.launch_builder(&f);
12999        b.arg(gate)
13000            .arg(up)
13001            .arg(act)
13002            .arg(&mut out_q)
13003            .arg(&mut out_d)
13004            .arg(&nc);
13005        unsafe {
13006            b.launch(cfg)?;
13007        }
13008        Ok((out_q, out_d))
13009    }
13010
13011    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
13012    #[allow(clippy::too_many_arguments)]
13013    pub fn gelu_tanh_mul_q8_1_into(
13014        &self,
13015        gate: &CudaSlice<f32>,
13016        up: &cudarc::driver::CudaView<f32>,
13017        act: &mut CudaSlice<f32>,
13018        ncols: usize,
13019        nrows: usize,
13020        out_q: &mut CudaSlice<i8>,
13021        out_d: &mut CudaSlice<f32>,
13022    ) -> Result<(), Box<dyn std::error::Error>> {
13023        debug_assert!(ncols.is_multiple_of(128));
13024        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
13025        let nc = ncols as i32;
13026        if Self::pdl_on() {
13027            use cudarc::driver::{DevicePtr, DevicePtrMut};
13028            let s = &self.gpu.stream();
13029            let (pg, _g0) = gate.device_ptr(s);
13030            let (pu, _g1) = up.device_ptr(s);
13031            let (pact, _g2) = act.device_ptr_mut(s);
13032            let (pq, _g3) = out_q.device_ptr_mut(s);
13033            let (pd, _g4) = out_d.device_ptr_mut(s);
13034            let mut ps = [
13035                &pg as *const _ as *mut std::ffi::c_void,
13036                &pu as *const _ as *mut _,
13037                &pact as *const _ as *mut _,
13038                &pq as *const _ as *mut _,
13039                &pd as *const _ as *mut _,
13040                &nc as *const _ as *mut _,
13041            ];
13042            unsafe {
13043                self.launch_pdl(
13044                    "gelu_tanh_mul_q8_1",
13045                    (nrows as u32, 1, 1),
13046                    (rms_block(), 1, 1),
13047                    &mut ps,
13048                )?;
13049            }
13050            return Ok(());
13051        }
13052        let f = self.func("gelu_tanh_mul_q8_1");
13053        let cfg = LaunchConfig {
13054            grid_dim: (nrows as u32, 1, 1),
13055            block_dim: (rms_block(), 1, 1),
13056            shared_mem_bytes: 0,
13057        };
13058        let __s_b = self.gpu.stream();
13059        let mut b = __s_b.launch_builder(&f);
13060        b.arg(gate)
13061            .arg(up)
13062            .arg(&mut *act)
13063            .arg(&mut *out_q)
13064            .arg(&mut *out_d)
13065            .arg(&nc);
13066        unsafe {
13067            b.launch(cfg)?;
13068        }
13069        Ok(())
13070    }
13071
13072    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
13073    #[allow(clippy::too_many_arguments)]
13074    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
13075    pub fn add_rms_norm3_q8z(
13076        &self,
13077        a: &CudaSlice<f32>,
13078        b_in: &CudaSlice<f32>,
13079        w0: &CudaSlice<f32>,
13080        w1: &CudaSlice<f32>,
13081        w2: &CudaSlice<f32>,
13082        res: &mut CudaSlice<f32>,
13083        out1: &mut CudaSlice<f32>,
13084        ncols: usize,
13085        nrows: usize,
13086        eps: f32,
13087    ) -> Result<
13088        (
13089            (CudaSlice<i8>, CudaSlice<f32>),
13090            (CudaSlice<i8>, CudaSlice<f32>),
13091        ),
13092        Box<dyn std::error::Error>,
13093    > {
13094        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
13095        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13096        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
13097        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13098        let f = self.func("add_rms_norm3_q8z_f32");
13099        let cfg = LaunchConfig {
13100            grid_dim: (nrows as u32, 1, 1),
13101            block_dim: (rms_block(), 1, 1),
13102            shared_mem_bytes: 0,
13103        };
13104        let (nc, e2) = (ncols as i32, eps);
13105        let __s_b = self.gpu.stream();
13106        let mut b = __s_b.launch_builder(&f);
13107        b.arg(a)
13108            .arg(b_in)
13109            .arg(w0)
13110            .arg(w1)
13111            .arg(w2)
13112            .arg(res)
13113            .arg(&mut q0)
13114            .arg(&mut d0)
13115            .arg(out1)
13116            .arg(&mut q2)
13117            .arg(&mut d2)
13118            .arg(&nc)
13119            .arg(&e2);
13120        unsafe {
13121            b.launch(cfg)?;
13122        }
13123        Ok(((q0, d0), (q2, d2)))
13124    }
13125
13126    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
13127    #[allow(clippy::too_many_arguments)]
13128    pub fn add_rms_norm3(
13129        &self,
13130        a: &CudaSlice<f32>,
13131        b_in: &CudaSlice<f32>,
13132        w0: &CudaSlice<f32>,
13133        w1: &CudaSlice<f32>,
13134        w2: &CudaSlice<f32>,
13135        res: &mut CudaSlice<f32>,
13136        d0: &mut CudaSlice<f32>,
13137        d1: &mut CudaSlice<f32>,
13138        d2: &mut CudaSlice<f32>,
13139        ncols: usize,
13140        nrows: usize,
13141        eps: f32,
13142    ) -> Result<(), Box<dyn std::error::Error>> {
13143        let f = self.func("add_rms_norm3_f32");
13144        let cfg = LaunchConfig {
13145            grid_dim: (nrows as u32, 1, 1),
13146            block_dim: (rms_block(), 1, 1),
13147            shared_mem_bytes: 0,
13148        };
13149        let (nc, e2) = (ncols as i32, eps);
13150        let __s_b = self.gpu.stream();
13151        let mut b = __s_b.launch_builder(&f);
13152        b.arg(a)
13153            .arg(b_in)
13154            .arg(w0)
13155            .arg(w1)
13156            .arg(w2)
13157            .arg(res)
13158            .arg(d0)
13159            .arg(d1)
13160            .arg(d2)
13161            .arg(&nc)
13162            .arg(&e2);
13163        unsafe {
13164            b.launch(cfg)?;
13165        }
13166        Ok(())
13167    }
13168
13169    /// dst = (a + b) * c (residual add + layer scale, one launch).
13170    pub fn add_scale(
13171        &self,
13172        a: &CudaSlice<f32>,
13173        b_in: &CudaSlice<f32>,
13174        c: f32,
13175        dst: &mut CudaSlice<f32>,
13176        n: usize,
13177    ) -> Result<(), Box<dyn std::error::Error>> {
13178        let f = self.func("add_scale_f32");
13179        let cfg = LaunchConfig::for_num_elems(n as u32);
13180        let ni = n as i32;
13181        let __s_b = self.gpu.stream();
13182        let mut b = __s_b.launch_builder(&f);
13183        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
13184        unsafe {
13185            b.launch(cfg)?;
13186        }
13187        Ok(())
13188    }
13189
13190    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
13191    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13192    pub fn layer_norm_bias(
13193        &self,
13194        x: &CudaSlice<f32>,
13195        w: &CudaSlice<f32>,
13196        b: &CudaSlice<f32>,
13197        dst: &mut CudaSlice<f32>,
13198        ncols: usize,
13199        nrows: usize,
13200        eps: f32,
13201    ) -> Result<(), Box<dyn std::error::Error>> {
13202        let f = self.func("layer_norm_bias_f32");
13203        let (nc, e) = (ncols as i32, eps);
13204        let cfg = LaunchConfig {
13205            grid_dim: (nrows as u32, 1, 1),
13206            block_dim: (256, 1, 1),
13207            shared_mem_bytes: 0,
13208        };
13209        let __s_b = self.gpu.stream();
13210        let mut lb = __s_b.launch_builder(&f);
13211        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
13212        unsafe {
13213            lb.launch(cfg)?;
13214        }
13215        Ok(())
13216    }
13217
13218    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
13219    pub fn gelu_tanh(
13220        &self,
13221        x: &CudaSlice<f32>,
13222        dst: &mut CudaSlice<f32>,
13223        n: usize,
13224    ) -> Result<(), Box<dyn std::error::Error>> {
13225        let f = self.func("gelu_tanh_f32");
13226        let ni = n as i64;
13227        let cfg = LaunchConfig {
13228            grid_dim: (n.div_ceil(256) as u32, 1, 1),
13229            block_dim: (256, 1, 1),
13230            shared_mem_bytes: 0,
13231        };
13232        let __s_b = self.gpu.stream();
13233        let mut lb = __s_b.launch_builder(&f);
13234        lb.arg(x).arg(&mut *dst).arg(&ni);
13235        unsafe {
13236            lb.launch(cfg)?;
13237        }
13238        Ok(())
13239    }
13240
13241    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
13242    pub fn row_softmax(
13243        &self,
13244        x: &mut CudaSlice<f32>,
13245        ncols: usize,
13246        nrows: usize,
13247    ) -> Result<(), Box<dyn std::error::Error>> {
13248        let f = self.func("row_softmax_f32");
13249        let nc = ncols as i32;
13250        let cfg = LaunchConfig {
13251            grid_dim: (nrows as u32, 1, 1),
13252            block_dim: (256, 1, 1),
13253            shared_mem_bytes: 0,
13254        };
13255        let __s_b = self.gpu.stream();
13256        let mut lb = __s_b.launch_builder(&f);
13257        lb.arg(&mut *x).arg(&nc);
13258        unsafe {
13259            lb.launch(cfg)?;
13260        }
13261        Ok(())
13262    }
13263
13264    pub fn rms_norm(
13265        &self,
13266        x: &CudaSlice<f32>,
13267        w: &CudaSlice<f32>,
13268        dst: &mut CudaSlice<f32>,
13269        ncols: usize,
13270        nrows: usize,
13271        eps: f32,
13272    ) -> Result<(), Box<dyn std::error::Error>> {
13273        let (nc, e) = (ncols as i32, eps);
13274        let kname = if Self::norm_ilp_on() {
13275            "rms_norm_f32_v2"
13276        } else {
13277            "rms_norm_f32"
13278        };
13279        if Self::pdl_on() && Self::pdl_wb_on() {
13280            use cudarc::driver::{DevicePtr, DevicePtrMut};
13281            let s = &self.gpu.stream();
13282            let (px, _g0) = x.device_ptr(s);
13283            let (pw, _g1) = w.device_ptr(s);
13284            let (pd, _g2) = dst.device_ptr_mut(s);
13285            let mut ps = [
13286                &px as *const _ as *mut std::ffi::c_void,
13287                &pw as *const _ as *mut _,
13288                &pd as *const _ as *mut _,
13289                &nc as *const _ as *mut _,
13290                &e as *const _ as *mut _,
13291            ];
13292            unsafe {
13293                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
13294            }
13295            return Ok(());
13296        }
13297        let f = self.func(kname);
13298        let cfg = LaunchConfig {
13299            grid_dim: (nrows as u32, 1, 1),
13300            block_dim: (rms_block(), 1, 1),
13301            shared_mem_bytes: 0,
13302        };
13303        let __s_b = self.gpu.stream();
13304        let mut b = __s_b.launch_builder(&f);
13305        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
13306        unsafe {
13307            b.launch(cfg)?;
13308        }
13309        Ok(())
13310    }
13311
13312    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
13313    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
13314    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
13315    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
13316    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
13317    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
13318    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
13319    pub fn rms_norm_decode(
13320        &self,
13321        x: &CudaSlice<f32>,
13322        w: &CudaSlice<f32>,
13323        dst: &mut CudaSlice<f32>,
13324        ncols: usize,
13325        nrows: usize,
13326        eps: f32,
13327    ) -> Result<(), Box<dyn std::error::Error>> {
13328        let f = self.func(if Self::norm_ilp_on() {
13329            "rms_norm_f32_v2"
13330        } else {
13331            "rms_norm_f32"
13332        });
13333        let cfg = LaunchConfig {
13334            grid_dim: (nrows as u32, 1, 1),
13335            block_dim: (1024, 1, 1),
13336            shared_mem_bytes: 0,
13337        };
13338        let (nc, e) = (ncols as i32, eps);
13339        let __s_b = self.gpu.stream();
13340        let mut b = __s_b.launch_builder(&f);
13341        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
13342        unsafe {
13343            b.launch(cfg)?;
13344        }
13345        Ok(())
13346    }
13347
13348    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
13349    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
13350    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
13351    pub fn rms_norm_q8_1(
13352        &self,
13353        x: &CudaSlice<f32>,
13354        w: &CudaSlice<f32>,
13355        ncols: usize,
13356        nrows: usize,
13357        eps: f32,
13358    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13359        let nblk = ncols / 32;
13360        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
13361        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
13362        let (nc, e) = (ncols as i32, eps);
13363        if Self::pdl_on() {
13364            {
13365                use cudarc::driver::{DevicePtr, DevicePtrMut};
13366                let s = &self.gpu.stream();
13367                let (px, _g0) = x.device_ptr(s);
13368                let (pw, _g1) = w.device_ptr(s);
13369                let (pq, _g2) = q.device_ptr_mut(s);
13370                let (pd, _g3) = d.device_ptr_mut(s);
13371                let mut ps = [
13372                    &px as *const _ as *mut std::ffi::c_void,
13373                    &pw as *const _ as *mut _,
13374                    &pq as *const _ as *mut _,
13375                    &pd as *const _ as *mut _,
13376                    &nc as *const _ as *mut _,
13377                    &e as *const _ as *mut _,
13378                ];
13379                unsafe {
13380                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
13381                }
13382            }
13383            return Ok((q, d));
13384        }
13385        let f = self.func("rms_norm_q8_1");
13386        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
13387        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
13388        let cfg = LaunchConfig {
13389            grid_dim: (nrows as u32, 1, 1),
13390            block_dim: (1024, 1, 1),
13391            shared_mem_bytes: 0,
13392        };
13393        let __s_b = self.gpu.stream();
13394        let mut b = __s_b.launch_builder(&f);
13395        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
13396        unsafe {
13397            b.launch(cfg)?;
13398        }
13399        Ok((q, d))
13400    }
13401
13402    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
13403    /// PDL arm), caller-owned outputs.
13404    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13405    pub fn rms_norm_q8_1_into(
13406        &self,
13407        x: &CudaSlice<f32>,
13408        w: &CudaSlice<f32>,
13409        ncols: usize,
13410        nrows: usize,
13411        eps: f32,
13412        q: &mut CudaSlice<i8>,
13413        d: &mut CudaSlice<f32>,
13414    ) -> Result<(), Box<dyn std::error::Error>> {
13415        let nblk = ncols / 32;
13416        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
13417        let (nc, e) = (ncols as i32, eps);
13418        if Self::pdl_on() {
13419            use cudarc::driver::{DevicePtr, DevicePtrMut};
13420            let s = &self.gpu.stream();
13421            let (px, _g0) = x.device_ptr(s);
13422            let (pw, _g1) = w.device_ptr(s);
13423            let (pq, _g2) = q.device_ptr_mut(s);
13424            let (pd, _g3) = d.device_ptr_mut(s);
13425            let mut ps = [
13426                &px as *const _ as *mut std::ffi::c_void,
13427                &pw as *const _ as *mut _,
13428                &pq as *const _ as *mut _,
13429                &pd as *const _ as *mut _,
13430                &nc as *const _ as *mut _,
13431                &e as *const _ as *mut _,
13432            ];
13433            unsafe {
13434                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
13435            }
13436            return Ok(());
13437        }
13438        let f = self.func("rms_norm_q8_1");
13439        let cfg = LaunchConfig {
13440            grid_dim: (nrows as u32, 1, 1),
13441            block_dim: (1024, 1, 1),
13442            shared_mem_bytes: 0,
13443        };
13444        let __s_b = self.gpu.stream();
13445        let mut b = __s_b.launch_builder(&f);
13446        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
13447        unsafe {
13448            b.launch(cfg)?;
13449        }
13450        Ok(())
13451    }
13452
13453    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
13454    pub fn quantize_q8_1_into(
13455        &self,
13456        x: &CudaSlice<f32>,
13457        m: usize,
13458        in_f: usize,
13459        q: &mut CudaSlice<i8>,
13460        d: &mut CudaSlice<f32>,
13461    ) -> Result<(), Box<dyn std::error::Error>> {
13462        let nblk = in_f / 32;
13463        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
13464        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
13465        let (inf, mi) = (in_f as i32, m as i32);
13466        if Self::pdl_on() && Self::pdl_wb_on() {
13467            use cudarc::driver::{DevicePtr, DevicePtrMut};
13468            let s = &self.gpu.stream();
13469            let (px, _g0) = x.device_ptr(s);
13470            let (pq, _g1) = q.device_ptr_mut(s);
13471            let (pd, _g2) = d.device_ptr_mut(s);
13472            let mut ps = [
13473                &px as *const _ as *mut std::ffi::c_void,
13474                &pq as *const _ as *mut _,
13475                &pd as *const _ as *mut _,
13476                &inf as *const _ as *mut _,
13477                &mi as *const _ as *mut _,
13478            ];
13479            unsafe {
13480                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
13481            }
13482            return Ok(());
13483        }
13484        let f = self.func("quantize_q8_1");
13485        let __s_b = self.gpu.stream();
13486        let mut b = __s_b.launch_builder(&f);
13487        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
13488        unsafe {
13489            b.launch(cfg)?;
13490        }
13491        Ok(())
13492    }
13493
13494    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
13495    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
13496    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
13497    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13498    pub fn add_rms_norm_q8_1(
13499        &self,
13500        a: &CudaSlice<f32>,
13501        b_in: &CudaSlice<f32>,
13502        w: &CudaSlice<f32>,
13503        res: &mut CudaSlice<f32>,
13504        ncols: usize,
13505        nrows: usize,
13506        eps: f32,
13507    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13508        let nblk = ncols / 32;
13509        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
13510        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
13511        let f = self.func("add_rms_norm_q8_1");
13512        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
13513        let cfg = LaunchConfig {
13514            grid_dim: (nrows as u32, 1, 1),
13515            block_dim: (1024, 1, 1),
13516            shared_mem_bytes: 0,
13517        };
13518        let (nc, e) = (ncols as i32, eps);
13519        let __s_bld = self.gpu.stream();
13520        let mut bld = __s_bld.launch_builder(&f);
13521        bld.arg(a)
13522            .arg(b_in)
13523            .arg(w)
13524            .arg(res)
13525            .arg(&mut q)
13526            .arg(&mut d)
13527            .arg(&nc)
13528            .arg(&e);
13529        unsafe {
13530            bld.launch(cfg)?;
13531        }
13532        Ok((q, d))
13533    }
13534
13535    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
13536    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
13537    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
13538    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
13539    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
13540    #[allow(clippy::too_many_arguments)]
13541    pub fn join_add_rms_norm_raw(
13542        &self,
13543        a0_raw: u64,
13544        a1_raw: u64,
13545        x: &CudaSlice<f32>,
13546        w: &CudaSlice<f32>,
13547        res: &mut CudaSlice<f32>,
13548        dst: &mut CudaSlice<f32>,
13549        ncols: usize,
13550        eps: f32,
13551    ) -> Result<(), Box<dyn std::error::Error>> {
13552        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
13553            return Err("join_add_rms_norm geometry".into());
13554        }
13555        let f = self.func("join_add_rms_norm_f32");
13556        let cfg = LaunchConfig {
13557            grid_dim: (1, 1, 1),
13558            block_dim: (rms_block(), 1, 1),
13559            shared_mem_bytes: 0,
13560        };
13561        let (nc, e) = (ncols as i32, eps);
13562        let __s_b = self.gpu.stream();
13563        let mut b = __s_b.launch_builder(&f);
13564        b.arg(&a0_raw)
13565            .arg(&a1_raw)
13566            .arg(x)
13567            .arg(w)
13568            .arg(&mut *res)
13569            .arg(&mut *dst)
13570            .arg(&nc)
13571            .arg(&e);
13572        unsafe {
13573            b.launch(cfg)?;
13574        }
13575        Ok(())
13576    }
13577
13578    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13579    pub fn add_rms_norm(
13580        &self,
13581        a: &CudaSlice<f32>,
13582        b: &CudaSlice<f32>,
13583        w: &CudaSlice<f32>,
13584        res: &mut CudaSlice<f32>,
13585        dst: &mut CudaSlice<f32>,
13586        ncols: usize,
13587        nrows: usize,
13588        eps: f32,
13589    ) -> Result<(), Box<dyn std::error::Error>> {
13590        let (nc, e) = (ncols as i32, eps);
13591        let kname = if Self::norm_ilp_on() {
13592            "add_rms_norm_f32_v2"
13593        } else {
13594            "add_rms_norm_f32"
13595        };
13596        if Self::pdl_on() && Self::pdl_wb_on() {
13597            use cudarc::driver::{DevicePtr, DevicePtrMut};
13598            let s = &self.gpu.stream();
13599            let (pa, _g0) = a.device_ptr(s);
13600            let (pb, _g1) = b.device_ptr(s);
13601            let (pw, _g2) = w.device_ptr(s);
13602            let (pr, _g3) = res.device_ptr_mut(s);
13603            let (pd, _g4) = dst.device_ptr_mut(s);
13604            let mut ps = [
13605                &pa as *const _ as *mut std::ffi::c_void,
13606                &pb as *const _ as *mut _,
13607                &pw as *const _ as *mut _,
13608                &pr as *const _ as *mut _,
13609                &pd as *const _ as *mut _,
13610                &nc as *const _ as *mut _,
13611                &e as *const _ as *mut _,
13612            ];
13613            unsafe {
13614                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
13615            }
13616            return Ok(());
13617        }
13618        let f = self.func(kname);
13619        let cfg = LaunchConfig {
13620            grid_dim: (nrows as u32, 1, 1),
13621            block_dim: (rms_block(), 1, 1),
13622            shared_mem_bytes: 0,
13623        };
13624        let __s_b2 = self.gpu.stream();
13625        let mut b2 = __s_b2.launch_builder(&f);
13626        b2.arg(a)
13627            .arg(b)
13628            .arg(w)
13629            .arg(&mut *res)
13630            .arg(&mut *dst)
13631            .arg(&nc)
13632            .arg(&e);
13633        unsafe {
13634            b2.launch(cfg)?;
13635        }
13636        Ok(())
13637    }
13638
13639    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
13640    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
13641    #[allow(clippy::too_many_arguments)]
13642    pub fn rms_pre_add_rms_norm(
13643        &self,
13644        a: &CudaSlice<f32>,
13645        wa: &CudaSlice<f32>,
13646        b: &CudaSlice<f32>,
13647        w: &CudaSlice<f32>,
13648        res: &mut CudaSlice<f32>,
13649        dst: &mut CudaSlice<f32>,
13650        ncols: usize,
13651        nrows: usize,
13652        eps: f32,
13653    ) -> Result<(), Box<dyn std::error::Error>> {
13654        let f = self.func("rms_pre_add_rms_norm_f32");
13655        let cfg = LaunchConfig {
13656            grid_dim: (nrows as u32, 1, 1),
13657            block_dim: (rms_block(), 1, 1),
13658            shared_mem_bytes: 0,
13659        };
13660        let (nc, e) = (ncols as i32, eps);
13661        let __s_b2 = self.gpu.stream();
13662        let mut b2 = __s_b2.launch_builder(&f);
13663        b2.arg(a)
13664            .arg(wa)
13665            .arg(b)
13666            .arg(w)
13667            .arg(&mut *res)
13668            .arg(&mut *dst)
13669            .arg(&nc)
13670            .arg(&e);
13671        unsafe {
13672            b2.launch(cfg)?;
13673        }
13674        Ok(())
13675    }
13676
13677    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
13678    #[allow(clippy::too_many_arguments)]
13679    pub fn rms_pre_add_rms_norm_q8z(
13680        &self,
13681        a: &CudaSlice<f32>,
13682        wa: &CudaSlice<f32>,
13683        b: &CudaSlice<f32>,
13684        w: &CudaSlice<f32>,
13685        res: &mut CudaSlice<f32>,
13686        dst: &mut CudaSlice<f32>,
13687        ncols: usize,
13688        nrows: usize,
13689        eps: f32,
13690    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13691        debug_assert!(ncols.is_multiple_of(128));
13692        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
13693        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13694        let (nc, e) = (ncols as i32, eps);
13695        if Self::pdl_on() {
13696            {
13697                use cudarc::driver::{DevicePtr, DevicePtrMut};
13698                let s = &self.gpu.stream();
13699                let (pa, _g0) = a.device_ptr(s);
13700                let (pwa, _g1) = wa.device_ptr(s);
13701                let (pb, _g2) = b.device_ptr(s);
13702                let (pw, _g3) = w.device_ptr(s);
13703                let (pr, _g4) = res.device_ptr_mut(s);
13704                let (pdst, _g5) = dst.device_ptr_mut(s);
13705                let (pq, _g6) = out_q.device_ptr_mut(s);
13706                let (pd, _g7) = out_d.device_ptr_mut(s);
13707                let mut ps = [
13708                    &pa as *const _ as *mut std::ffi::c_void,
13709                    &pwa as *const _ as *mut _,
13710                    &pb as *const _ as *mut _,
13711                    &pw as *const _ as *mut _,
13712                    &pr as *const _ as *mut _,
13713                    &pdst as *const _ as *mut _,
13714                    &pq as *const _ as *mut _,
13715                    &pd as *const _ as *mut _,
13716                    &nc as *const _ as *mut _,
13717                    &e as *const _ as *mut _,
13718                ];
13719                unsafe {
13720                    self.launch_pdl(
13721                        "rms_pre_add_rms_norm_q8z_f32",
13722                        (nrows as u32, 1, 1),
13723                        (rms_block(), 1, 1),
13724                        &mut ps,
13725                    )?;
13726                }
13727            }
13728            return Ok((out_q, out_d));
13729        }
13730        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
13731        let cfg = LaunchConfig {
13732            grid_dim: (nrows as u32, 1, 1),
13733            block_dim: (rms_block(), 1, 1),
13734            shared_mem_bytes: 0,
13735        };
13736        let __s_b2 = self.gpu.stream();
13737        let mut b2 = __s_b2.launch_builder(&f);
13738        b2.arg(a)
13739            .arg(wa)
13740            .arg(b)
13741            .arg(w)
13742            .arg(&mut *res)
13743            .arg(&mut *dst)
13744            .arg(&mut out_q)
13745            .arg(&mut out_d)
13746            .arg(&nc)
13747            .arg(&e);
13748        unsafe {
13749            b2.launch(cfg)?;
13750        }
13751        Ok((out_q, out_d))
13752    }
13753
13754    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
13755    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
13756    /// body must stay attribute-free (the fused2_into precedent).
13757    #[allow(clippy::too_many_arguments)]
13758    pub fn rms_pre_add_rms_norm_q8z_into(
13759        &self,
13760        a: &CudaSlice<f32>,
13761        wa: &CudaSlice<f32>,
13762        b: &CudaSlice<f32>,
13763        w: &CudaSlice<f32>,
13764        res: &mut CudaSlice<f32>,
13765        dst: &mut CudaSlice<f32>,
13766        ncols: usize,
13767        nrows: usize,
13768        eps: f32,
13769        out_q: &mut CudaSlice<i8>,
13770        out_d: &mut CudaSlice<f32>,
13771    ) -> Result<(), Box<dyn std::error::Error>> {
13772        debug_assert!(ncols.is_multiple_of(128));
13773        let (nc, e) = (ncols as i32, eps);
13774        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
13775        let cfg = LaunchConfig {
13776            grid_dim: (nrows as u32, 1, 1),
13777            block_dim: (rms_block(), 1, 1),
13778            shared_mem_bytes: 0,
13779        };
13780        let __s_b = self.gpu.stream();
13781        let mut b2 = __s_b.launch_builder(&f);
13782        b2.arg(a)
13783            .arg(wa)
13784            .arg(b)
13785            .arg(w)
13786            .arg(&mut *res)
13787            .arg(&mut *dst)
13788            .arg(&mut *out_q)
13789            .arg(&mut *out_d)
13790            .arg(&nc)
13791            .arg(&e);
13792        unsafe {
13793            b2.launch(cfg)?;
13794        }
13795        Ok(())
13796    }
13797
13798    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
13799    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
13800    #[allow(clippy::too_many_arguments)]
13801    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
13802        &self,
13803        a: &CudaSlice<f32>,
13804        wa: &CudaSlice<f32>,
13805        b_in: &CudaSlice<f32>,
13806        c: f32,
13807        w: &CudaSlice<f32>,
13808        res: &mut CudaSlice<f32>,
13809        ncols: usize,
13810        nrows: usize,
13811        eps: f32,
13812        out_q: &mut CudaSlice<i8>,
13813        out_d: &mut CudaSlice<f32>,
13814    ) -> Result<(), Box<dyn std::error::Error>> {
13815        debug_assert!(ncols.is_multiple_of(128));
13816        let (nc, e2) = (ncols as i32, eps);
13817        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
13818        let cfg = LaunchConfig {
13819            grid_dim: (nrows as u32, 1, 1),
13820            block_dim: (rms_block(), 1, 1),
13821            shared_mem_bytes: 0,
13822        };
13823        let __s_b = self.gpu.stream();
13824        let mut b2 = __s_b.launch_builder(&f);
13825        b2.arg(a)
13826            .arg(wa)
13827            .arg(b_in)
13828            .arg(&c)
13829            .arg(w)
13830            .arg(&mut *res)
13831            .arg(&mut *out_q)
13832            .arg(&mut *out_d)
13833            .arg(&nc)
13834            .arg(&e2);
13835        unsafe {
13836            b2.launch(cfg)?;
13837        }
13838        Ok(())
13839    }
13840
13841    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
13842    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
13843    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
13844    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
13845    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
13846    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
13847    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
13848    pub fn g4_pnfold_on() -> bool {
13849        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13850        *ON.get_or_init(|| {
13851            std::env::var("MEMRA_G4_PNFOLD")
13852                .map(|v| v != "0")
13853                .unwrap_or(true)
13854        })
13855    }
13856
13857    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
13858    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
13859    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
13860    pub fn build_q4_out_concat3(
13861        &self,
13862        w0: &crate::model::GpuTensor,
13863        w1: &crate::model::GpuTensor,
13864        w2: &crate::model::GpuTensor,
13865    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
13866        use crate::model::GpuTensor;
13867        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
13868            match w {
13869                GpuTensor::Quant {
13870                    qtype,
13871                    row_bytes,
13872                    rp,
13873                    ..
13874                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
13875                _ => None,
13876            }
13877        };
13878        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
13879        else {
13880            return Ok(None);
13881        };
13882        if rb0 != rb1
13883            || rb0 != rb2
13884            || w0.in_features() != w1.in_features()
13885            || w0.in_features() != w2.in_features()
13886        {
13887            return Ok(None);
13888        }
13889        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
13890            match w {
13891                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
13892                _ => unreachable!(),
13893            }
13894        }
13895        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
13896        let total = rb0 * (o0 + o1 + o2);
13897        let mut cat = self.alloc_u8(total)?;
13898        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
13899        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
13900        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
13901        Ok(Some(GpuTensor::Quant {
13902            bytes: cat,
13903            qtype: QT_Q4_0,
13904            row_bytes: rb0,
13905            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
13906            scale: 1.0,
13907            rp: false,
13908            #[cfg(memra_cutlass)]
13909            cutlass: None,
13910            fp8: None,
13911            blk: None,
13912            rp4: None,
13913            f16: None,
13914        }))
13915    }
13916
13917    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
13918    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
13919    ///
13920    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
13921    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
13922    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
13923    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
13924    ///
13925    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
13926    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
13927    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
13928    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
13929    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
13930    ///
13931    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
13932    /// width. A future partial-rotary caller fails at its first launch with the geometry named
13933    /// instead of serving quietly wrong logits.
13934    fn full_width_rope_only(
13935        kernel: &str,
13936        n_rot: usize,
13937        head_dim: usize,
13938    ) -> Result<(), Box<dyn std::error::Error>> {
13939        if n_rot == head_dim {
13940            return Ok(());
13941        }
13942        Err(format!(
13943            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
13944             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
13945             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
13946             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
13947             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
13948        )
13949        .into())
13950    }
13951
13952    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
13953    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
13954    /// ([`Engine::full_width_rope_only`]).
13955    #[allow(clippy::too_many_arguments)]
13956    pub fn rms_norm_qkv_rope_cat(
13957        &self,
13958        qkv: &CudaSlice<f32>,
13959        wq: &CudaSlice<f32>,
13960        wk: &CudaSlice<f32>,
13961        wv: &CudaSlice<f32>,
13962        q: &mut CudaSlice<f32>,
13963        k: &mut CudaSlice<f32>,
13964        v: &mut CudaSlice<f32>,
13965        head_dim: usize,
13966        n_rot: usize,
13967        rq: usize,
13968        rk: usize,
13969        pos: &CudaSlice<i32>,
13970        nh_q: usize,
13971        nh_k: usize,
13972        base: f32,
13973        freq_scale: f32,
13974        ff: Option<&CudaSlice<f32>>,
13975        eps: f32,
13976    ) -> Result<(), Box<dyn std::error::Error>> {
13977        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
13978        let rows = rq + rk + rk;
13979        let theta_scale = base.powf(-2.0 / head_dim as f32);
13980        let (nc, rqi, rki, nhq, nhk) = (
13981            head_dim as i32,
13982            rq as i32,
13983            rk as i32,
13984            nh_q as i32,
13985            nh_k as i32,
13986        );
13987        if Self::pdl_on() {
13988            use cudarc::driver::{DevicePtr, DevicePtrMut};
13989            let s = &self.gpu.stream();
13990            let (pqkv, _g0) = qkv.device_ptr(s);
13991            let (pwq, _g1) = wq.device_ptr(s);
13992            let (pwk, _g2) = wk.device_ptr(s);
13993            let (pwv, _g3) = wv.device_ptr(s);
13994            let (pq, _g4) = q.device_ptr_mut(s);
13995            let (pk, _g5) = k.device_ptr_mut(s);
13996            let (pv, _g6) = v.device_ptr_mut(s);
13997            let (ppos, _g7) = pos.device_ptr(s);
13998            let (pff, _g8) = match ff {
13999                Some(t) => {
14000                    let (p, g) = t.device_ptr(s);
14001                    (p, Some(g))
14002                }
14003                None => (0, None),
14004            };
14005            let mut ps = [
14006                &pqkv as *const _ as *mut std::ffi::c_void,
14007                &pwq as *const _ as *mut _,
14008                &pwk as *const _ as *mut _,
14009                &pwv as *const _ as *mut _,
14010                &pq as *const _ as *mut _,
14011                &pk as *const _ as *mut _,
14012                &pv as *const _ as *mut _,
14013                &nc as *const _ as *mut _,
14014                &rqi as *const _ as *mut _,
14015                &rki as *const _ as *mut _,
14016                &ppos as *const _ as *mut _,
14017                &nhq as *const _ as *mut _,
14018                &nhk as *const _ as *mut _,
14019                &theta_scale as *const _ as *mut _,
14020                &freq_scale as *const _ as *mut _,
14021                &pff as *const _ as *mut _,
14022                &eps as *const _ as *mut _,
14023            ];
14024            unsafe {
14025                self.launch_pdl(
14026                    "rms_norm_qkv_rope_cat_f32",
14027                    (rows as u32, 1, 1),
14028                    (rms_block(), 1, 1),
14029                    &mut ps,
14030                )?;
14031            }
14032            return Ok(());
14033        }
14034        let f = self.func("rms_norm_qkv_rope_cat_f32");
14035        let cfg = LaunchConfig {
14036            grid_dim: (rows as u32, 1, 1),
14037            block_dim: (rms_block(), 1, 1),
14038            shared_mem_bytes: 0,
14039        };
14040        let __s_b = self.gpu.stream();
14041        let mut b = __s_b.launch_builder(&f);
14042        match ff {
14043            Some(t) => {
14044                b.arg(qkv)
14045                    .arg(wq)
14046                    .arg(wk)
14047                    .arg(wv)
14048                    .arg(&mut *q)
14049                    .arg(&mut *k)
14050                    .arg(&mut *v)
14051                    .arg(&nc)
14052                    .arg(&rqi)
14053                    .arg(&rki)
14054                    .arg(pos)
14055                    .arg(&nhq)
14056                    .arg(&nhk)
14057                    .arg(&theta_scale)
14058                    .arg(&freq_scale)
14059                    .arg(t)
14060                    .arg(&eps);
14061                unsafe {
14062                    b.launch(cfg)?;
14063                }
14064            }
14065            None => {
14066                let null: u64 = 0;
14067                b.arg(qkv)
14068                    .arg(wq)
14069                    .arg(wk)
14070                    .arg(wv)
14071                    .arg(&mut *q)
14072                    .arg(&mut *k)
14073                    .arg(&mut *v)
14074                    .arg(&nc)
14075                    .arg(&rqi)
14076                    .arg(&rki)
14077                    .arg(pos)
14078                    .arg(&nhq)
14079                    .arg(&nhk)
14080                    .arg(&theta_scale)
14081                    .arg(&freq_scale)
14082                    .arg(&null)
14083                    .arg(&eps);
14084                unsafe {
14085                    b.launch(cfg)?;
14086                }
14087            }
14088        }
14089        Ok(())
14090    }
14091
14092    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
14093    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
14094    /// ([`Engine::full_width_rope_only`]).
14095    #[allow(clippy::too_many_arguments)]
14096    pub fn rms_norm_qkv_rope(
14097        &self,
14098        q0: &CudaSlice<f32>,
14099        k0: &CudaSlice<f32>,
14100        v0: &CudaSlice<f32>,
14101        wq: &CudaSlice<f32>,
14102        wk: &CudaSlice<f32>,
14103        wv: &CudaSlice<f32>,
14104        q: &mut CudaSlice<f32>,
14105        k: &mut CudaSlice<f32>,
14106        v: &mut CudaSlice<f32>,
14107        head_dim: usize,
14108        n_rot: usize,
14109        rq: usize,
14110        rk: usize,
14111        pos: &CudaSlice<i32>,
14112        nh_q: usize,
14113        nh_k: usize,
14114        base: f32,
14115        freq_scale: f32,
14116        ff: Option<&CudaSlice<f32>>,
14117        eps: f32,
14118    ) -> Result<(), Box<dyn std::error::Error>> {
14119        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
14120        let f = self.func("rms_norm_qkv_rope_f32");
14121        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
14122        let cfg = LaunchConfig {
14123            grid_dim: (rows as u32, 1, 1),
14124            block_dim: (rms_block(), 1, 1),
14125            shared_mem_bytes: 0,
14126        };
14127        let theta_scale = base.powf(-2.0 / head_dim as f32);
14128        let (nc, rqi, rki, nhq, nhk) = (
14129            head_dim as i32,
14130            rq as i32,
14131            rk as i32,
14132            nh_q as i32,
14133            nh_k as i32,
14134        );
14135        let __s_b = self.gpu.stream();
14136        let mut b = __s_b.launch_builder(&f);
14137        match ff {
14138            Some(t) => {
14139                b.arg(q0)
14140                    .arg(k0)
14141                    .arg(v0)
14142                    .arg(wq)
14143                    .arg(wk)
14144                    .arg(wv)
14145                    .arg(&mut *q)
14146                    .arg(&mut *k)
14147                    .arg(&mut *v)
14148                    .arg(&nc)
14149                    .arg(&rqi)
14150                    .arg(&rki)
14151                    .arg(pos)
14152                    .arg(&nhq)
14153                    .arg(&nhk)
14154                    .arg(&theta_scale)
14155                    .arg(&freq_scale)
14156                    .arg(t)
14157                    .arg(&eps);
14158                unsafe {
14159                    b.launch(cfg)?;
14160                }
14161            }
14162            None => {
14163                let null: u64 = 0;
14164                b.arg(q0)
14165                    .arg(k0)
14166                    .arg(v0)
14167                    .arg(wq)
14168                    .arg(wk)
14169                    .arg(wv)
14170                    .arg(&mut *q)
14171                    .arg(&mut *k)
14172                    .arg(&mut *v)
14173                    .arg(&nc)
14174                    .arg(&rqi)
14175                    .arg(&rki)
14176                    .arg(pos)
14177                    .arg(&nhq)
14178                    .arg(&nhk)
14179                    .arg(&theta_scale)
14180                    .arg(&freq_scale)
14181                    .arg(&null)
14182                    .arg(&eps);
14183                unsafe {
14184                    b.launch(cfg)?;
14185                }
14186            }
14187        }
14188        Ok(())
14189    }
14190
14191    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
14192    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
14193    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
14194    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
14195    /// ([`Engine::full_width_rope_only`]).
14196    #[allow(clippy::too_many_arguments)]
14197    pub fn rms_norm_qkv_rope_append_dc(
14198        &self,
14199        q0: &CudaSlice<f32>,
14200        k0: &CudaSlice<f32>,
14201        v0: &CudaSlice<f32>,
14202        wq: &CudaSlice<f32>,
14203        wk: &CudaSlice<f32>,
14204        wv: &CudaSlice<f32>,
14205        q: &mut CudaSlice<f32>,
14206        k: &mut CudaSlice<f32>,
14207        v: &mut CudaSlice<f32>,
14208        head_dim: usize,
14209        n_rot: usize,
14210        rq: usize,
14211        rk: usize,
14212        pos: &CudaSlice<i32>,
14213        nh_q: usize,
14214        nh_k: usize,
14215        base: f32,
14216        freq_scale: f32,
14217        ff: Option<&CudaSlice<f32>>,
14218        eps: f32,
14219        kc: &mut CudaSlice<u8>,
14220        vc: &mut CudaSlice<u8>,
14221        t_dev: &CudaSlice<i32>,
14222        k_tok_bytes: usize,
14223        v_tok_bytes: usize,
14224        g: bool,
14225    ) -> Result<(), Box<dyn std::error::Error>> {
14226        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
14227        let rows = rq + rk + rk;
14228        let theta_scale = base.powf(-2.0 / head_dim as f32);
14229        let (nc, rqi, rki, nhq, nhk) = (
14230            head_dim as i32,
14231            rq as i32,
14232            rk as i32,
14233            nh_q as i32,
14234            nh_k as i32,
14235        );
14236        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
14237        if Self::pdl_on() && Self::pdl_wb_on() {
14238            use cudarc::driver::{DevicePtr, DevicePtrMut};
14239            let s = &self.gpu.stream();
14240            let (p0, _a0) = q0.device_ptr(s);
14241            let (p1, _a1) = k0.device_ptr(s);
14242            let (p2, _a2) = v0.device_ptr(s);
14243            let (pwq, _a3) = wq.device_ptr(s);
14244            let (pwk, _a4) = wk.device_ptr(s);
14245            let (pwv, _a5) = wv.device_ptr(s);
14246            let (pq, _a6) = q.device_ptr_mut(s);
14247            let (pk, _a7) = k.device_ptr_mut(s);
14248            let (pv, _a8) = v.device_ptr_mut(s);
14249            let (pp, _a9) = pos.device_ptr(s);
14250            let pff: u64 = match ff {
14251                Some(t) => {
14252                    let (p, _gg) = t.device_ptr(s);
14253                    p
14254                }
14255                None => 0,
14256            };
14257            let (pkc, _a10) = kc.device_ptr_mut(s);
14258            let (pvc, _a11) = vc.device_ptr_mut(s);
14259            let (pt, _a12) = t_dev.device_ptr(s);
14260            let mut ps = [
14261                &p0 as *const _ as *mut std::ffi::c_void,
14262                &p1 as *const _ as *mut _,
14263                &p2 as *const _ as *mut _,
14264                &pwq as *const _ as *mut _,
14265                &pwk as *const _ as *mut _,
14266                &pwv as *const _ as *mut _,
14267                &pq as *const _ as *mut _,
14268                &pk as *const _ as *mut _,
14269                &pv as *const _ as *mut _,
14270                &nc as *const _ as *mut _,
14271                &rqi as *const _ as *mut _,
14272                &rki as *const _ as *mut _,
14273                &pp as *const _ as *mut _,
14274                &nhq as *const _ as *mut _,
14275                &nhk as *const _ as *mut _,
14276                &theta_scale as *const _ as *mut _,
14277                &freq_scale as *const _ as *mut _,
14278                &pff as *const _ as *mut _,
14279                &eps as *const _ as *mut _,
14280                &pkc as *const _ as *mut _,
14281                &pvc as *const _ as *mut _,
14282                &pt as *const _ as *mut _,
14283                &ktb as *const _ as *mut _,
14284                &vtb as *const _ as *mut _,
14285            ];
14286            unsafe {
14287                self.launch_pdl_flash(
14288                    g,
14289                    "rms_norm_qkv_rope_append_dc_f32",
14290                    (rows as u32, 1, 1),
14291                    (rms_block(), 1, 1),
14292                    0,
14293                    &mut ps,
14294                )?;
14295            }
14296            return Ok(());
14297        }
14298        let f = if g {
14299            self.func_g("rms_norm_qkv_rope_append_dc_f32")
14300        } else {
14301            self.func("rms_norm_qkv_rope_append_dc_f32")
14302        };
14303        let cfg = LaunchConfig {
14304            grid_dim: (rows as u32, 1, 1),
14305            block_dim: (rms_block(), 1, 1),
14306            shared_mem_bytes: 0,
14307        };
14308        let __s_b = self.gpu.stream();
14309        let mut b = __s_b.launch_builder(&f);
14310        match ff {
14311            Some(t) => {
14312                b.arg(q0)
14313                    .arg(k0)
14314                    .arg(v0)
14315                    .arg(wq)
14316                    .arg(wk)
14317                    .arg(wv)
14318                    .arg(&mut *q)
14319                    .arg(&mut *k)
14320                    .arg(&mut *v)
14321                    .arg(&nc)
14322                    .arg(&rqi)
14323                    .arg(&rki)
14324                    .arg(pos)
14325                    .arg(&nhq)
14326                    .arg(&nhk)
14327                    .arg(&theta_scale)
14328                    .arg(&freq_scale)
14329                    .arg(t)
14330                    .arg(&eps)
14331                    .arg(&mut *kc)
14332                    .arg(&mut *vc)
14333                    .arg(t_dev)
14334                    .arg(&ktb)
14335                    .arg(&vtb);
14336                unsafe {
14337                    b.launch(cfg)?;
14338                }
14339            }
14340            None => {
14341                let null: u64 = 0;
14342                b.arg(q0)
14343                    .arg(k0)
14344                    .arg(v0)
14345                    .arg(wq)
14346                    .arg(wk)
14347                    .arg(wv)
14348                    .arg(&mut *q)
14349                    .arg(&mut *k)
14350                    .arg(&mut *v)
14351                    .arg(&nc)
14352                    .arg(&rqi)
14353                    .arg(&rki)
14354                    .arg(pos)
14355                    .arg(&nhq)
14356                    .arg(&nhk)
14357                    .arg(&theta_scale)
14358                    .arg(&freq_scale)
14359                    .arg(&null)
14360                    .arg(&eps)
14361                    .arg(&mut *kc)
14362                    .arg(&mut *vc)
14363                    .arg(t_dev)
14364                    .arg(&ktb)
14365                    .arg(&vtb);
14366                unsafe {
14367                    b.launch(cfg)?;
14368                }
14369            }
14370        }
14371        Ok(())
14372    }
14373
14374    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
14375    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
14376    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
14377    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
14378    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
14379    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
14380    /// `head_dim` ([`Engine::full_width_rope_only`]).
14381    #[allow(clippy::too_many_arguments)]
14382    pub fn rms_norm_qkv_rope_append(
14383        &self,
14384        q0: &CudaSlice<f32>,
14385        k0: &CudaSlice<f32>,
14386        v0: &CudaSlice<f32>,
14387        wq: &CudaSlice<f32>,
14388        wk: &CudaSlice<f32>,
14389        wv: &CudaSlice<f32>,
14390        q: &mut CudaSlice<f32>,
14391        k: &mut CudaSlice<f32>,
14392        v: &mut CudaSlice<f32>,
14393        head_dim: usize,
14394        n_rot: usize,
14395        rq: usize,
14396        rk: usize,
14397        pos: &CudaSlice<i32>,
14398        nh_q: usize,
14399        nh_k: usize,
14400        base: f32,
14401        freq_scale: f32,
14402        ff: Option<&CudaSlice<f32>>,
14403        eps: f32,
14404        kc: &mut CudaSlice<u8>,
14405        vc: &mut CudaSlice<u8>,
14406        t: usize,
14407        k_tok_bytes: usize,
14408        v_tok_bytes: usize,
14409        g: bool,
14410    ) -> Result<(), Box<dyn std::error::Error>> {
14411        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
14412        let rows = rq + rk + rk;
14413        let theta_scale = base.powf(-2.0 / head_dim as f32);
14414        let (nc, rqi, rki, nhq, nhk) = (
14415            head_dim as i32,
14416            rq as i32,
14417            rk as i32,
14418            nh_q as i32,
14419            nh_k as i32,
14420        );
14421        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
14422        let ti = t as i32;
14423        if Self::pdl_on() && Self::pdl_wb_on() {
14424            use cudarc::driver::{DevicePtr, DevicePtrMut};
14425            let s = &self.gpu.stream();
14426            let (p0, _a0) = q0.device_ptr(s);
14427            let (p1, _a1) = k0.device_ptr(s);
14428            let (p2, _a2) = v0.device_ptr(s);
14429            let (pwq, _a3) = wq.device_ptr(s);
14430            let (pwk, _a4) = wk.device_ptr(s);
14431            let (pwv, _a5) = wv.device_ptr(s);
14432            let (pq, _a6) = q.device_ptr_mut(s);
14433            let (pk, _a7) = k.device_ptr_mut(s);
14434            let (pv, _a8) = v.device_ptr_mut(s);
14435            let (pp, _a9) = pos.device_ptr(s);
14436            let pff: u64 = match ff {
14437                Some(t) => {
14438                    let (p, _gg) = t.device_ptr(s);
14439                    p
14440                }
14441                None => 0,
14442            };
14443            let (pkc, _a10) = kc.device_ptr_mut(s);
14444            let (pvc, _a11) = vc.device_ptr_mut(s);
14445            let mut ps = [
14446                &p0 as *const _ as *mut std::ffi::c_void,
14447                &p1 as *const _ as *mut _,
14448                &p2 as *const _ as *mut _,
14449                &pwq as *const _ as *mut _,
14450                &pwk as *const _ as *mut _,
14451                &pwv as *const _ as *mut _,
14452                &pq as *const _ as *mut _,
14453                &pk as *const _ as *mut _,
14454                &pv as *const _ as *mut _,
14455                &nc as *const _ as *mut _,
14456                &rqi as *const _ as *mut _,
14457                &rki as *const _ as *mut _,
14458                &pp as *const _ as *mut _,
14459                &nhq as *const _ as *mut _,
14460                &nhk as *const _ as *mut _,
14461                &theta_scale as *const _ as *mut _,
14462                &freq_scale as *const _ as *mut _,
14463                &pff as *const _ as *mut _,
14464                &eps as *const _ as *mut _,
14465                &pkc as *const _ as *mut _,
14466                &pvc as *const _ as *mut _,
14467                &ti as *const _ as *mut _,
14468                &ktb as *const _ as *mut _,
14469                &vtb as *const _ as *mut _,
14470            ];
14471            unsafe {
14472                self.launch_pdl_flash(
14473                    g,
14474                    "rms_norm_qkv_rope_append_f32",
14475                    (rows as u32, 1, 1),
14476                    (rms_block(), 1, 1),
14477                    0,
14478                    &mut ps,
14479                )?;
14480            }
14481            return Ok(());
14482        }
14483        let f = if g {
14484            self.func_g("rms_norm_qkv_rope_append_f32")
14485        } else {
14486            self.func("rms_norm_qkv_rope_append_f32")
14487        };
14488        let cfg = LaunchConfig {
14489            grid_dim: (rows as u32, 1, 1),
14490            block_dim: (rms_block(), 1, 1),
14491            shared_mem_bytes: 0,
14492        };
14493        let __s_b = self.gpu.stream();
14494        let mut b = __s_b.launch_builder(&f);
14495        let null: u64 = 0;
14496        b.arg(q0)
14497            .arg(k0)
14498            .arg(v0)
14499            .arg(wq)
14500            .arg(wk)
14501            .arg(wv)
14502            .arg(&mut *q)
14503            .arg(&mut *k)
14504            .arg(&mut *v)
14505            .arg(&nc)
14506            .arg(&rqi)
14507            .arg(&rki)
14508            .arg(pos)
14509            .arg(&nhq)
14510            .arg(&nhk)
14511            .arg(&theta_scale)
14512            .arg(&freq_scale);
14513        match ff {
14514            Some(t) => {
14515                b.arg(t);
14516            }
14517            None => {
14518                b.arg(&null);
14519            }
14520        }
14521        b.arg(&eps)
14522            .arg(&mut *kc)
14523            .arg(&mut *vc)
14524            .arg(&ti)
14525            .arg(&ktb)
14526            .arg(&vtb);
14527        unsafe {
14528            b.launch(cfg)?;
14529        }
14530        Ok(())
14531    }
14532
14533    pub fn add_q8_1(
14534        &self,
14535        a: &CudaSlice<f32>,
14536        b: &CudaSlice<f32>,
14537        res: &mut CudaSlice<f32>,
14538        ncols: usize,
14539        nrows: usize,
14540    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14541        debug_assert!(ncols.is_multiple_of(128));
14542        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14543        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14544        let f = self.func("add_q8_1_f32");
14545        let cfg = LaunchConfig {
14546            grid_dim: (nrows as u32, 1, 1),
14547            block_dim: (rms_block(), 1, 1),
14548            shared_mem_bytes: 0,
14549        };
14550        let nc = ncols as i32;
14551        let __s_b2 = self.gpu.stream();
14552        let mut b2 = __s_b2.launch_builder(&f);
14553        b2.arg(a)
14554            .arg(b)
14555            .arg(&mut *res)
14556            .arg(&mut out_q)
14557            .arg(&mut out_d)
14558            .arg(&nc);
14559        unsafe {
14560            b2.launch(cfg)?;
14561        }
14562        Ok((out_q, out_d))
14563    }
14564
14565    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
14566    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
14567    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
14568    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14569    pub fn rms_pre_add_q8_1(
14570        &self,
14571        a: &CudaSlice<f32>,
14572        wa: &CudaSlice<f32>,
14573        b: &CudaSlice<f32>,
14574        res: &mut CudaSlice<f32>,
14575        ncols: usize,
14576        nrows: usize,
14577        eps: f32,
14578    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14579        debug_assert!(ncols.is_multiple_of(128));
14580        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14581        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14582        let f = self.func("rms_pre_add_q8_1_f32");
14583        let cfg = LaunchConfig {
14584            grid_dim: (nrows as u32, 1, 1),
14585            block_dim: (rms_block(), 1, 1),
14586            shared_mem_bytes: 0,
14587        };
14588        let (nc, ep) = (ncols as i32, eps);
14589        let __s_b2 = self.gpu.stream();
14590        let mut b2 = __s_b2.launch_builder(&f);
14591        b2.arg(a)
14592            .arg(wa)
14593            .arg(b)
14594            .arg(&mut *res)
14595            .arg(&mut out_q)
14596            .arg(&mut out_d)
14597            .arg(&nc)
14598            .arg(&ep);
14599        unsafe {
14600            b2.launch(cfg)?;
14601        }
14602        Ok((out_q, out_d))
14603    }
14604
14605    /// L2 norm per row (head_dim), no weight.
14606    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
14607    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
14608    pub fn l2_v2_on(ncols: usize) -> bool {
14609        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
14610    }
14611
14612    pub fn l2_norm_pp(
14613        &self,
14614        x: &CudaSlice<f32>,
14615        dst: &mut CudaSlice<f32>,
14616        dst16: Option<&mut CudaSlice<u8>>,
14617        ncols: usize,
14618        nrows: usize,
14619        eps: f32,
14620    ) -> Result<(), Box<dyn std::error::Error>> {
14621        if Self::l2_v2_on(ncols) {
14622            let f = self.func("l2_norm_pp_v2_f32");
14623            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
14624            let cfg = LaunchConfig {
14625                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
14626                block_dim: (256, 1, 1),
14627                shared_mem_bytes: 0,
14628            };
14629            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
14630            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
14631            let d16: u64 = match dst16 {
14632                Some(d) => self.addr_u8(d),
14633                None => 0,
14634            };
14635            let __s_b = self.gpu.stream();
14636            let mut b = __s_b.launch_builder(&f);
14637            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
14638            unsafe {
14639                b.launch(cfg)?;
14640            }
14641            return Ok(());
14642        }
14643        self.l2_norm(x, dst, ncols, nrows, eps)
14644    }
14645
14646    pub fn l2_norm(
14647        &self,
14648        x: &CudaSlice<f32>,
14649        dst: &mut CudaSlice<f32>,
14650        ncols: usize,
14651        nrows: usize,
14652        eps: f32,
14653    ) -> Result<(), Box<dyn std::error::Error>> {
14654        let f = self.func("l2_norm_f32");
14655        let cfg = LaunchConfig {
14656            grid_dim: (nrows as u32, 1, 1),
14657            block_dim: (256, 1, 1),
14658            shared_mem_bytes: 0,
14659        };
14660        let (nc, e) = (ncols as i32, eps);
14661        let __s_b = self.gpu.stream();
14662        let mut b = __s_b.launch_builder(&f);
14663        b.arg(x).arg(dst).arg(&nc).arg(&e);
14664        unsafe {
14665            b.launch(cfg)?;
14666        }
14667        Ok(())
14668    }
14669
14670    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
14671    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
14672    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
14673    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
14674    /// propagate through gdn_scan and flip argmax on marginal logits.
14675    pub fn l2_norm_decode(
14676        &self,
14677        x: &CudaSlice<f32>,
14678        dst: &mut CudaSlice<f32>,
14679        ncols: usize,
14680        nrows: usize,
14681        eps: f32,
14682    ) -> Result<(), Box<dyn std::error::Error>> {
14683        let f = self.func("l2_norm_f32");
14684        let cfg = LaunchConfig {
14685            grid_dim: (nrows as u32, 1, 1),
14686            block_dim: (32, 1, 1),
14687            shared_mem_bytes: 0,
14688        };
14689        let (nc, e) = (ncols as i32, eps);
14690        let __s_b = self.gpu.stream();
14691        let mut b = __s_b.launch_builder(&f);
14692        b.arg(x).arg(dst).arg(&nc).arg(&e);
14693        unsafe {
14694            b.launch(cfg)?;
14695        }
14696        Ok(())
14697    }
14698
14699    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
14700    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14701    pub fn rope_neox(
14702        &self,
14703        x: &mut CudaSlice<f32>,
14704        pos: &CudaSlice<i32>,
14705        head_dim: usize,
14706        n_dims: usize,
14707        n_heads: usize,
14708        n_tokens: usize,
14709        freq_base: f32,
14710        freq_scale: f32,
14711    ) -> Result<(), Box<dyn std::error::Error>> {
14712        let f = self.func("rope_neox_f32");
14713        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14714        let grid = (n_heads * n_tokens) as u32;
14715        let cfg = LaunchConfig {
14716            grid_dim: (grid, 1, 1),
14717            block_dim: ((head_dim / 2) as u32, 1, 1),
14718            shared_mem_bytes: 0,
14719        };
14720        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14721        let __s_b = self.gpu.stream();
14722        let mut b = __s_b.launch_builder(&f);
14723        b.arg(x)
14724            .arg(pos)
14725            .arg(&hd)
14726            .arg(&nd)
14727            .arg(&nh)
14728            .arg(&theta_scale)
14729            .arg(&freq_scale);
14730        unsafe {
14731            b.launch(cfg)?;
14732        }
14733        Ok(())
14734    }
14735
14736    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
14737    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14738    pub fn rope_neox_ff(
14739        &self,
14740        x: &mut CudaSlice<f32>,
14741        pos: &CudaSlice<i32>,
14742        head_dim: usize,
14743        n_dims: usize,
14744        n_heads: usize,
14745        n_tokens: usize,
14746        freq_base: f32,
14747        freq_scale: f32,
14748        ff: &CudaSlice<f32>,
14749    ) -> Result<(), Box<dyn std::error::Error>> {
14750        let f = self.func("rope_neox_ff_f32");
14751        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14752        let grid = (n_heads * n_tokens) as u32;
14753        let cfg = LaunchConfig {
14754            grid_dim: (grid, 1, 1),
14755            block_dim: ((head_dim / 2) as u32, 1, 1),
14756            shared_mem_bytes: 0,
14757        };
14758        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14759        let __s_b = self.gpu.stream();
14760        let mut b = __s_b.launch_builder(&f);
14761        b.arg(x)
14762            .arg(pos)
14763            .arg(&hd)
14764            .arg(&nd)
14765            .arg(&nh)
14766            .arg(&theta_scale)
14767            .arg(&freq_scale)
14768            .arg(ff);
14769        unsafe {
14770            b.launch(cfg)?;
14771        }
14772        Ok(())
14773    }
14774
14775    /// RoPE NEOX with per-dim freq factors AND the YaRN attention factor on cos/sin
14776    /// (qwen4_exp yarn lane — `rope_neox_ffm_f32`; ff = yarn_frequency_divisors, mscale =
14777    /// yarn_attention_factor). Identity inputs (ones, 1.0) reproduce `rope_neox` bit-for-bit.
14778    #[allow(clippy::too_many_arguments)]
14779    pub fn rope_neox_ffm(
14780        &self,
14781        x: &mut CudaSlice<f32>,
14782        pos: &CudaSlice<i32>,
14783        head_dim: usize,
14784        n_dims: usize,
14785        n_heads: usize,
14786        n_tokens: usize,
14787        freq_base: f32,
14788        freq_scale: f32,
14789        ff: &CudaSlice<f32>,
14790        mscale: f32,
14791    ) -> Result<(), Box<dyn std::error::Error>> {
14792        let f = self.func("rope_neox_ffm_f32");
14793        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14794        let grid = (n_heads * n_tokens) as u32;
14795        let cfg = LaunchConfig {
14796            grid_dim: (grid, 1, 1),
14797            block_dim: ((head_dim / 2) as u32, 1, 1),
14798            shared_mem_bytes: 0,
14799        };
14800        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14801        let __s_b = self.gpu.stream();
14802        let mut b = __s_b.launch_builder(&f);
14803        b.arg(x)
14804            .arg(pos)
14805            .arg(&hd)
14806            .arg(&nd)
14807            .arg(&nh)
14808            .arg(&theta_scale)
14809            .arg(&freq_scale)
14810            .arg(ff)
14811            .arg(&mscale);
14812        unsafe {
14813            b.launch(cfg)?;
14814        }
14815        Ok(())
14816    }
14817
14818    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
14819    #[allow(clippy::too_many_arguments)]
14820    pub fn rope_neox2(
14821        &self,
14822        q: &mut CudaSlice<f32>,
14823        k: &mut CudaSlice<f32>,
14824        pos: &CudaSlice<i32>,
14825        head_dim: usize,
14826        n_dims: usize,
14827        nh_q: usize,
14828        nh_k: usize,
14829        n_tokens: usize,
14830        freq_base: f32,
14831        freq_scale: f32,
14832        ff: Option<&CudaSlice<f32>>,
14833    ) -> Result<(), Box<dyn std::error::Error>> {
14834        let f = self.func("rope_neox2_f32");
14835        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14836        let grid = ((nh_q + nh_k) * n_tokens) as u32;
14837        let cfg = LaunchConfig {
14838            grid_dim: (grid, 1, 1),
14839            block_dim: ((head_dim / 2) as u32, 1, 1),
14840            shared_mem_bytes: 0,
14841        };
14842        let (hd, nd, nq, nk, nt) = (
14843            head_dim as i32,
14844            n_dims as i32,
14845            nh_q as i32,
14846            nh_k as i32,
14847            n_tokens as i32,
14848        );
14849        let __s_b = self.gpu.stream();
14850        let mut b = __s_b.launch_builder(&f);
14851        b.arg(q)
14852            .arg(k)
14853            .arg(pos)
14854            .arg(&hd)
14855            .arg(&nd)
14856            .arg(&nq)
14857            .arg(&nk)
14858            .arg(&nt)
14859            .arg(&theta_scale)
14860            .arg(&freq_scale);
14861        match ff {
14862            Some(ffv) => {
14863                b.arg(ffv);
14864                unsafe {
14865                    b.launch(cfg)?;
14866                }
14867            }
14868            None => {
14869                let null: u64 = 0;
14870                b.arg(&null);
14871                unsafe {
14872                    b.launch(cfg)?;
14873                }
14874            }
14875        }
14876        Ok(())
14877    }
14878
14879    /// gemma4 R1: dst = GELU_tanh(gate) * up.
14880    pub fn gelu_tanh_mul(
14881        &self,
14882        gate: &CudaSlice<f32>,
14883        up: &CudaSlice<f32>,
14884        dst: &mut CudaSlice<f32>,
14885        n: usize,
14886    ) -> Result<(), Box<dyn std::error::Error>> {
14887        let f = self.func("gelu_tanh_mul_f32");
14888        let cfg = LaunchConfig::for_num_elems(n as u32);
14889        let ni = n as i32;
14890        let __s_b = self.gpu.stream();
14891        let mut b = __s_b.launch_builder(&f);
14892        b.arg(gate).arg(up).arg(dst).arg(&ni);
14893        unsafe {
14894            b.launch(cfg)?;
14895        }
14896        Ok(())
14897    }
14898
14899    pub fn silu_mul(
14900        &self,
14901        gate: &CudaSlice<f32>,
14902        up: &CudaSlice<f32>,
14903        dst: &mut CudaSlice<f32>,
14904        n: usize,
14905    ) -> Result<(), Box<dyn std::error::Error>> {
14906        let f = self.func("silu_mul_f32");
14907        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
14908        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
14909        let ni = n as i32;
14910        let __s_b = self.gpu.stream();
14911        let mut b = __s_b.launch_builder(&f);
14912        b.arg(gate).arg(up).arg(dst).arg(&ni);
14913        unsafe {
14914            b.launch(cfg)?;
14915        }
14916        Ok(())
14917    }
14918
14919    /// SwiGLU twin using Memra's host-matching expf transcription.
14920    pub fn silu_mul_host_expf(
14921        &self,
14922        gate: &CudaSlice<f32>,
14923        up: &CudaSlice<f32>,
14924        dst: &mut CudaSlice<f32>,
14925        n: usize,
14926    ) -> Result<(), Box<dyn std::error::Error>> {
14927        let f = self.func("silu_mul_host_expf_f32");
14928        let cfg = LaunchConfig::for_num_elems(n as u32);
14929        let ni = n as i32;
14930        let __s_b = self.gpu.stream();
14931        let mut b = __s_b.launch_builder(&f);
14932        b.arg(gate).arg(up).arg(dst).arg(&ni);
14933        unsafe {
14934            b.launch(cfg)?;
14935        }
14936        Ok(())
14937    }
14938
14939    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
14940    pub fn silu_clamped_mul_host_expf(
14941        &self,
14942        gate: &CudaSlice<f32>,
14943        up: &CudaSlice<f32>,
14944        limit: f32,
14945        dst: &mut CudaSlice<f32>,
14946        n: usize,
14947    ) -> Result<(), Box<dyn std::error::Error>> {
14948        if !limit.is_finite() || limit <= 0.0 {
14949            return Err(
14950                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
14951            );
14952        }
14953        let f = self.func("silu_clamped_mul_host_expf_f32");
14954        let cfg = LaunchConfig::for_num_elems(n as u32);
14955        let ni = n as i32;
14956        let __s_b = self.gpu.stream();
14957        let mut b = __s_b.launch_builder(&f);
14958        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
14959        unsafe {
14960            b.launch(cfg)?;
14961        }
14962        Ok(())
14963    }
14964
14965    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
14966    /// for the down projection — kills the standalone convert pass. Bit-identical class.
14967    pub fn silu_mul_f16out(
14968        &self,
14969        gate: &CudaSlice<f32>,
14970        up: &CudaSlice<f32>,
14971        dst: &mut CudaSlice<f32>,
14972        dst16: &mut CudaSlice<u8>,
14973        n: usize,
14974    ) -> Result<(), Box<dyn std::error::Error>> {
14975        let f = self.func("silu_mul_f16out_f32");
14976        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
14977        let ni = n as i32;
14978        let __s_b = self.gpu.stream();
14979        let mut b = __s_b.launch_builder(&f);
14980        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
14981        unsafe {
14982            b.launch(cfg)?;
14983        }
14984        Ok(())
14985    }
14986
14987    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
14988    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
14989    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
14990    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
14991    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
14992    /// launches per dense FFN layer (the gate+up post-matmul scales).
14993    pub fn silu_mul_scaled(
14994        &self,
14995        gate: &CudaSlice<f32>,
14996        up: &CudaSlice<f32>,
14997        gs: f32,
14998        us: f32,
14999        dst: &mut CudaSlice<f32>,
15000        n: usize,
15001    ) -> Result<(), Box<dyn std::error::Error>> {
15002        let f = self.func("silu_mul_scaled_f32");
15003        let cfg = LaunchConfig::for_num_elems(n as u32);
15004        let ni = n as i32;
15005        let (gsf, usf) = (gs, us);
15006        let __s_b = self.gpu.stream();
15007        let mut b = __s_b.launch_builder(&f);
15008        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
15009        unsafe {
15010            b.launch(cfg)?;
15011        }
15012        Ok(())
15013    }
15014
15015    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
15016    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
15017    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
15018    #[allow(clippy::too_many_arguments)]
15019    pub fn swigluoai_mul_scaled(
15020        &self,
15021        gate: &CudaSlice<f32>,
15022        up: &CudaSlice<f32>,
15023        gs: f32,
15024        us: f32,
15025        alpha: f32,
15026        limit: f32,
15027        dst: &mut CudaSlice<f32>,
15028        n: usize,
15029    ) -> Result<(), Box<dyn std::error::Error>> {
15030        let f = self.func("swigluoai_mul_scaled_f32");
15031        let cfg = LaunchConfig::for_num_elems(n as u32);
15032        let ni = n as i32;
15033        let __s_b = self.gpu.stream();
15034        let mut b = __s_b.launch_builder(&f);
15035        b.arg(gate)
15036            .arg(up)
15037            .arg(&gs)
15038            .arg(&us)
15039            .arg(&alpha)
15040            .arg(&limit)
15041            .arg(dst)
15042            .arg(&ni);
15043        unsafe {
15044            b.launch(cfg)?;
15045        }
15046        Ok(())
15047    }
15048
15049    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
15050    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
15051    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
15052    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
15053    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
15054    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
15055    /// n must be a multiple of 32 (n_ff always is).
15056    pub fn silu_mul_scaled_q8_1(
15057        &self,
15058        gate: &CudaSlice<f32>,
15059        up: &CudaSlice<f32>,
15060        gs: f32,
15061        us: f32,
15062        n: usize,
15063    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15064        let f = self.func("silu_mul_scaled_q8_1");
15065        let nblk = n / 32;
15066        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
15067        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
15068        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
15069        let cfg = LaunchConfig::for_num_elems(n as u32);
15070        let (gsf, usf, ni) = (gs, us, n as i32);
15071        let __s_b = self.gpu.stream();
15072        let mut b = __s_b.launch_builder(&f);
15073        b.arg(gate)
15074            .arg(up)
15075            .arg(&gsf)
15076            .arg(&usf)
15077            .arg(&mut aq)
15078            .arg(&mut ad)
15079            .arg(&ni);
15080        unsafe {
15081            b.launch(cfg)?;
15082        }
15083        Ok((aq, ad))
15084    }
15085
15086    pub fn add(
15087        &self,
15088        a: &CudaSlice<f32>,
15089        b_in: &CudaSlice<f32>,
15090        dst: &mut CudaSlice<f32>,
15091        n: usize,
15092    ) -> Result<(), Box<dyn std::error::Error>> {
15093        let f = self.func("add_f32");
15094        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
15095        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15096        let ni = n as i32;
15097        let __s_bld = self.gpu.stream();
15098        let mut bld = __s_bld.launch_builder(&f);
15099        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
15100        unsafe {
15101            bld.launch(cfg)?;
15102        }
15103        Ok(())
15104    }
15105
15106    pub fn mul(
15107        &self,
15108        a: &CudaSlice<f32>,
15109        b_in: &CudaSlice<f32>,
15110        dst: &mut CudaSlice<f32>,
15111        n: usize,
15112    ) -> Result<(), Box<dyn std::error::Error>> {
15113        let f = self.func("mul_f32");
15114        let cfg = LaunchConfig::for_num_elems(n as u32);
15115        let ni = n as i32;
15116        let __s_bld = self.gpu.stream();
15117        let mut bld = __s_bld.launch_builder(&f);
15118        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
15119        unsafe {
15120            bld.launch(cfg)?;
15121        }
15122        Ok(())
15123    }
15124
15125    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
15126    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
15127    pub fn matmul(
15128        &self,
15129        w: &crate::model::GpuTensor,
15130        x: &CudaSlice<f32>,
15131        m: usize,
15132    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15133        use crate::model::GpuTensor;
15134        let in_f = w.in_features();
15135        let out_f = w.out_features();
15136        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
15137        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
15138        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
15139        // gives nothing). Quantize the activation once here then call the GEMM.
15140        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
15141        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
15142        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
15143        #[allow(non_snake_case)]
15144        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
15145        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
15146        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
15147            usize::MAX
15148        } else {
15149            16usize
15150        };
15151
15152        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
15153        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
15154        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
15155        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
15156        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
15157        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
15158        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
15159        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
15160        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
15161        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
15162        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
15163        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
15164        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
15165        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
15166        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
15167        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
15168        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
15169        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
15170        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
15171        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
15172        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
15173        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
15174        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
15175        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
15176        if m >= GEMM_M_THRESHOLD {
15177            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
15178                return Ok(y);
15179            }
15180            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
15181            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
15182            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
15183            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
15184            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
15185            // tile defaults differently by operand source.
15186            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15187                return Ok(y);
15188            }
15189            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
15190            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
15191            if let Some(y) = self.try_f16_gemm(w, x, m)? {
15192                return Ok(y);
15193            }
15194        }
15195        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
15196        // m threshold the rest of this method uses:
15197        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
15198        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
15199        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
15200        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
15201        //     across every tier by construction with no batched twin needed.
15202        //
15203        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
15204        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
15205        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
15206        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
15207        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
15208        // arms is what makes sure it never gets there.
15209        if let GpuTensor::Quant { qtype, .. } = w
15210            && *qtype == QT_F8_E4M3_BLK
15211        {
15212            if m >= GEMM_M_THRESHOLD
15213                && let Some(y) = self.try_e4m3_blk_prefill(w, x, m)?
15214            {
15215                return Ok(y);
15216            }
15217            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15218            if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
15219                return Ok(y);
15220            }
15221        }
15222        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
15223            return self.qmatvec_mmq(w, x, m);
15224        }
15225        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
15226            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15227            return self.qmatvec_gemm(w, &aq, &ad, m);
15228        }
15229        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
15230        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
15231        if m >= GEMM_M_THRESHOLD
15232            && let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)?
15233        {
15234            return Ok(y);
15235        }
15236        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
15237        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
15238        // to Stage-A f32-dequant (the correctness oracle path).
15239        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
15240        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
15241        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
15242        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
15243        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
15244        if m == 1
15245            && fast
15246            && let GpuTensor::Quant {
15247                bytes,
15248                qtype,
15249                row_bytes,
15250                rp,
15251                rp4,
15252                scale,
15253                ..
15254            } = w
15255            && self.mmvq_supports(*qtype)
15256        {
15257            // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
15258            // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
15259            // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
15260            let (bytes, rp) = match rp4 {
15261                Some(m4) => (m4, true),
15262                None => (bytes, *rp),
15263            };
15264            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15265            return self.qmatvec_mmvq(
15266                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
15267            );
15268        }
15269        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
15270        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
15271        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
15272        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
15273        // block below. MEMRA_NO_BATCHED -> per-m path.
15274        //
15275        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
15276        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
15277        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
15278        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
15279        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
15280        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
15281        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
15282        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
15283        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
15284        if (2..=16).contains(&m)
15285            && fast
15286            && std::env::var("MEMRA_NO_BATCHED").is_err()
15287            && (m <= 4 || Self::b8_enabled())
15288        {
15289            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
15290            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
15291            // is present (rp4) — the mirror pick below then routes to the _rp family.
15292            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
15293            // because the native e4m3 row layout is already aligned and needs no mirror.
15294            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
15295            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
15296            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
15297            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
15298            let m_ok = m <= 8
15299                || matches!(w, GpuTensor::Quant { qtype, .. }
15300                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
15301                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
15302            if m_ok
15303                && let GpuTensor::Quant {
15304                    bytes,
15305                    qtype,
15306                    row_bytes,
15307                    rp,
15308                    rp4,
15309                    ..
15310                } = w
15311                && self.batched_supports(*qtype)
15312                && self.mmvq_supports(*qtype)
15313            {
15314                let (bytes, rp) = match rp4 {
15315                    Some(m4) => (m4, true),
15316                    None => (bytes, *rp),
15317                };
15318                let mcols = Self::batched_mcols(m);
15319                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15320                let mut y = self.qmatvec_mmvq_batched(
15321                    bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
15322                )?;
15323                if let GpuTensor::Quant { scale, .. } = w
15324                    && *scale != 1.0
15325                {
15326                    self.scale_inplace(&mut y, *scale, m * out_f)?;
15327                }
15328                return Ok(y);
15329            }
15330        }
15331        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
15332        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
15333        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
15334        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
15335        // for this dtype, so the generic match below must never see it under `fast`.
15336        if fast
15337            && let GpuTensor::Quant {
15338                bytes,
15339                qtype,
15340                row_bytes,
15341                scale,
15342                ..
15343            } = w
15344            && *qtype == QT_F8_E4M3
15345        {
15346            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15347            return self.qmatvec_mmvq(
15348                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
15349            );
15350        }
15351        let mut y = match w {
15352            GpuTensor::Quant {
15353                bytes,
15354                qtype,
15355                row_bytes,
15356                ..
15357            } if fast && *qtype == QT_Q8_0 => {
15358                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15359            }
15360            GpuTensor::Quant {
15361                bytes,
15362                qtype,
15363                row_bytes,
15364                ..
15365            } if fast && *qtype == QT_Q4_K => {
15366                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15367            }
15368            GpuTensor::Quant {
15369                bytes,
15370                qtype,
15371                row_bytes,
15372                ..
15373            } if fast && *qtype == QT_Q6_K => {
15374                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15375            }
15376            GpuTensor::Quant {
15377                bytes,
15378                qtype,
15379                row_bytes,
15380                ..
15381            } if fast && *qtype == QT_Q5_K => {
15382                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15383            }
15384            GpuTensor::Quant {
15385                bytes,
15386                qtype,
15387                row_bytes,
15388                ..
15389            } if fast && *qtype == QT_Q3_K => {
15390                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15391            }
15392            GpuTensor::Quant {
15393                bytes,
15394                qtype,
15395                row_bytes,
15396                rp,
15397                ..
15398            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
15399                if *rp {
15400                    "qmatvec_nvfp4_dp4a_rp"
15401                } else {
15402                    "qmatvec_nvfp4_dp4a"
15403                },
15404                &bytes.slice(0..bytes.len()),
15405                x,
15406                m,
15407                in_f,
15408                out_f,
15409                *row_bytes,
15410            )?,
15411            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
15412            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
15413            // anomaly (research/kat-anomaly-20260802/).
15414            GpuTensor::Quant {
15415                bytes,
15416                qtype,
15417                row_bytes,
15418                ..
15419            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
15420                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15421            }
15422            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
15423            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
15424            // without first writing the matching kernel, or func() will panic
15425            // "kernel ... not in any fatbin".
15426            GpuTensor::Quant {
15427                bytes,
15428                qtype,
15429                row_bytes,
15430                rp,
15431                ..
15432            } =>
15433            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
15434            // deq(row,j) form cannot address the planes; same value/product order).
15435            {
15436                self.qmatvec(
15437                    bytes,
15438                    x,
15439                    m,
15440                    in_f,
15441                    out_f,
15442                    if *rp && *qtype == QT_NVFP4 {
15443                        QT_NVFP4_RP
15444                    } else {
15445                        *qtype
15446                    },
15447                    *row_bytes,
15448                )?
15449            }
15450            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
15451            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
15452            // cuBLASLt f32 GEMV as the Float arm.
15453            GpuTensor::FloatBf16 { data, .. } => {
15454                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
15455                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
15456                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
15457                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
15458                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
15459                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
15460                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f.is_multiple_of(8) {
15461                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15462                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
15463                    y
15464                } else {
15465                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
15466                }
15467            }
15468        };
15469        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
15470        if let GpuTensor::Quant { scale, .. } = w
15471            && *scale != 1.0
15472        {
15473            self.scale_inplace(&mut y, *scale, m * out_f)?;
15474        }
15475        Ok(y)
15476    }
15477
15478    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
15479    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
15480    ///
15481    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
15482    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
15483    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
15484    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
15485    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
15486    /// path must not pay an env lookup for a flag that is off.
15487    pub fn stage_a_raw_needed() -> bool {
15488        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15489        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
15490    }
15491
15492    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
15493    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
15494    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
15495        use crate::model::GpuTensor;
15496        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
15497            return false;
15498        }
15499        match w {
15500            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
15501            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
15502            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
15503            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
15504            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
15505            // block class has no fused twin yet, so each of its projections takes its own launch.
15506            GpuTensor::Quant { qtype, .. } => {
15507                matches!(
15508                    *qtype,
15509                    QT_Q8_0
15510                        | QT_Q4_K
15511                        | QT_Q6_K
15512                        | QT_Q5_K
15513                        | QT_Q3_K
15514                        | QT_NVFP4
15515                        | QT_F8_E4M3
15516                        | QT_F8_E4M3_BLK
15517                        | QT_Q4_0
15518                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
15519            }
15520            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15521        }
15522    }
15523
15524    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
15525    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
15526    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
15527    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
15528    pub fn matmul_pre(
15529        &self,
15530        w: &crate::model::GpuTensor,
15531        aq: &CudaSlice<i8>,
15532        ad: &CudaSlice<f32>,
15533        x_fallback: &CudaSlice<f32>,
15534        m: usize,
15535    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15536        use crate::model::GpuTensor;
15537        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
15538        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
15539        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
15540        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
15541        // rc=30013 dig, 2026-07-31).
15542        let x_raw_ok = x_fallback.len() >= m * w.in_features();
15543        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
15544        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
15545        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
15546            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
15547                return Ok(y);
15548            }
15549            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
15550            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
15551            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
15552                return Ok(y);
15553            }
15554            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
15555            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
15556                return Ok(y);
15557            }
15558        }
15559        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
15560        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
15561        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
15562        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
15563        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
15564        if m >= 16
15565            && x_raw_ok
15566            && !self.verify_exact_on()
15567            && let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)?
15568        {
15569            return Ok(y);
15570        }
15571        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15572            return Ok(y);
15573        }
15574        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
15575        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
15576        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
15577        // aq/ad.
15578        if m >= 16
15579            && w.out_features() >= 128
15580            && self.mmq_supports(w)
15581            && !self.verify_exact_on()
15582            && x_raw_ok
15583        {
15584            return self.qmatvec_mmq(w, x_fallback, m);
15585        }
15586        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
15587        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
15588        if m >= 16
15589            && x_raw_ok
15590            && !self.verify_exact_on()
15591            && let Some(y) =
15592                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
15593        {
15594            return Ok(y);
15595        }
15596        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
15597        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
15598        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
15599            return self.qmatvec_gemm(w, aq, ad, m);
15600        }
15601        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
15602        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
15603        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
15604        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
15605        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
15606        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
15607        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
15608        // which reads `m * in_f` floats out of a 0-byte allocation ->
15609        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
15610        // it poisons the context, so every LATER request in that process fails with an unrelated
15611        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
15612        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
15613        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
15614        // dense artifact and left the arm with no working truth instrument.
15615        //
15616        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
15617        // strictly better than an illegal address surfacing later at an unrelated sync point, and
15618        // an oracle that cannot run must say so rather than corrupt the context it runs in.
15619        if !self.uses_q8_1_fast(w) {
15620            if !x_raw_ok {
15621                return Err(format!(
15622                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
15623                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
15624                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
15625                     activation (see Engine::rms_norm_decode, which is bit-identical to \
15626                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
15627                    x_fallback.len(),
15628                    m,
15629                    w.in_features(),
15630                    m * w.in_features()
15631                )
15632                .into());
15633            }
15634            return self.matmul(w, x_fallback, m);
15635        }
15636        let in_f = w.in_features();
15637        let out_f = w.out_features();
15638        let (bytes, qtype, row_bytes, scale, rp) = match w {
15639            GpuTensor::Quant {
15640                bytes,
15641                qtype,
15642                row_bytes,
15643                scale,
15644                rp,
15645                ..
15646            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15647            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
15648        };
15649        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
15650        // the dp4a/oracle tails below keep the raw GGUF bytes.
15651        let (mbytes, mrp) = match w {
15652            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15653            _ => (bytes, rp),
15654        };
15655        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
15656        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
15657        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
15658        if m == 1 && self.mmvq_supports(qtype) {
15659            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
15660        }
15661        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
15662        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
15663        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
15664        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
15665        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
15666        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
15667        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
15668        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
15669        // m=5..8 on the old per-m path (b8-tier-only seam).
15670        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
15671        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
15672        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
15673        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
15674            && std::env::var("MEMRA_NO_BATCHED").is_err()
15675            && (m <= 4 || Self::b8_enabled())
15676            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
15677            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
15678            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
15679            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
15680                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
15681        {
15682            let mcols = Self::batched_mcols(m);
15683            return self.qmatvec_mmvq_batched(
15684                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
15685            );
15686        }
15687        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
15688        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
15689        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
15690        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
15691        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
15692        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
15693            let (b2, r2) = if qtype == QT_Q4_0 {
15694                (mbytes, mrp)
15695            } else {
15696                (bytes, rp)
15697            };
15698            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
15699        }
15700        let name = match qtype {
15701            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15702            QT_Q4_K => "qmatvec_q4_K_dp4a",
15703            QT_Q6_K => "qmatvec_q6_K_dp4a",
15704            QT_Q5_K => "qmatvec_q5_K_dp4a",
15705            QT_Q3_K => "qmatvec_q3_K_dp4a",
15706            QT_NVFP4 => {
15707                if rp {
15708                    "qmatvec_nvfp4_dp4a_rp"
15709                } else {
15710                    "qmatvec_nvfp4_dp4a"
15711                }
15712            }
15713            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15714            _ => unreachable!(),
15715        };
15716        let f = self.func(name);
15717        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15718        let cfg = LaunchConfig {
15719            grid_dim: (out_f as u32, m as u32, 1),
15720            block_dim: (128, 1, 1),
15721            shared_mem_bytes: 0,
15722        };
15723        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15724        let __s_b = self.gpu.stream();
15725        let mut b = __s_b.launch_builder(&f);
15726        b.arg(bytes)
15727            .arg(aq)
15728            .arg(ad)
15729            .arg(&mut y)
15730            .arg(&inf)
15731            .arg(&outf)
15732            .arg(&mi)
15733            .arg(&rb);
15734        unsafe {
15735            b.launch(cfg)?;
15736        }
15737        if scale != 1.0 {
15738            self.scale_inplace(&mut y, scale, m * out_f)?;
15739        }
15740        Ok(y)
15741    }
15742
15743    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
15744    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
15745    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
15746    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
15747    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
15748    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
15749    /// reduce as m=1); this method just forces that path unconditionally.
15750    pub fn matmul_decode_exact(
15751        &self,
15752        w: &crate::model::GpuTensor,
15753        x: &CudaSlice<f32>,
15754        m: usize,
15755    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15756        use crate::model::GpuTensor;
15757        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
15758        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
15759        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
15760        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
15761        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
15762        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
15763        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
15764        if let GpuTensor::Float { data, .. } = w {
15765            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
15766        }
15767        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
15768        // float linear (same n-independent reduction contract as the Float arm above).
15769        if let GpuTensor::FloatBf16 { data, .. } = w {
15770            let (in_f, out_f) = (w.in_features(), w.out_features());
15771            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
15772            // contract — the whole-weight f32 dequant disappears too).
15773            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
15774                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15775                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
15776                return Ok(y);
15777            }
15778            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
15779        }
15780        if !self.uses_q8_1_fast(w) {
15781            return self.matmul(w, x, m);
15782        }
15783        let in_f = w.in_features();
15784        let out_f = w.out_features();
15785        let (bytes, qtype, row_bytes, scale, rp) = match w {
15786            GpuTensor::Quant {
15787                bytes,
15788                qtype,
15789                row_bytes,
15790                scale,
15791                rp,
15792                ..
15793            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15794            _ => return self.matmul(w, x, m),
15795        };
15796        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
15797        // which does its own mirror pick).
15798        let (bytes, rp) = match w {
15799            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15800            _ => (bytes, rp),
15801        };
15802        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15803        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
15804        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
15805        // (token,row) by construction, which is exactly what this method exists to guarantee.
15806        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
15807            return Ok(y);
15808        }
15809        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
15810        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
15811        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
15812        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
15813        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
15814        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
15815        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
15816        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
15817        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
15818            && std::env::var("MEMRA_NO_BATCHED").is_err()
15819            && (m <= 4 || Self::b8_enabled())
15820            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
15821            // no mirror precondition, `rp` selects the layout only.
15822            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
15823                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
15824        {
15825            let mcols = Self::batched_mcols(m);
15826            return self.qmatvec_mmvq_batched(
15827                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
15828            );
15829        }
15830        if self.mmvq_supports(qtype) {
15831            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
15832            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
15833            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
15834        }
15835        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
15836        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
15837        self.matmul_pre(w, &aq, &ad, x, m)
15838    }
15839
15840    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
15841    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
15842    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
15843    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
15844    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
15845    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
15846    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
15847    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
15848    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
15849    pub fn matmul_decode_exact_pre(
15850        &self,
15851        w: &crate::model::GpuTensor,
15852        aq: &CudaSlice<i8>,
15853        ad: &CudaSlice<f32>,
15854        m: usize,
15855    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15856        use crate::model::GpuTensor;
15857        debug_assert!(
15858            self.uses_q8_1_fast(w),
15859            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
15860        );
15861        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
15862        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15863            return Ok(y);
15864        }
15865        let in_f = w.in_features();
15866        let out_f = w.out_features();
15867        let (bytes, qtype, row_bytes, scale, rp) = match w {
15868            GpuTensor::Quant {
15869                bytes,
15870                qtype,
15871                row_bytes,
15872                scale,
15873                rp,
15874                ..
15875            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15876            _ => {
15877                return Err(
15878                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
15879                );
15880            }
15881        };
15882        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
15883        let (bytes, rp) = match w {
15884            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15885            _ => (bytes, rp),
15886        };
15887        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
15888        if (2..=16).contains(&m)
15889            && self.batched_supports(qtype)
15890            && self.mmvq_supports(qtype)
15891            && std::env::var("MEMRA_NO_BATCHED").is_err()
15892            && (m <= 4 || Self::b8_enabled())
15893            && (m <= 8
15894                || qtype == QT_Q4_0
15895                || qtype == QT_Q6_K
15896                || qtype == QT_F8_E4M3
15897                || qtype == QT_NVFP4
15898                || qtype == QT_Q4_K
15899                || qtype == QT_Q5_K
15900                || qtype == QT_Q8_0)
15901        {
15902            let mcols = Self::batched_mcols(m);
15903            return self.qmatvec_mmvq_batched(
15904                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
15905            );
15906        }
15907        if self.mmvq_supports(qtype) {
15908            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
15909        }
15910        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
15911        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
15912        let x0 = self.zeros(0)?;
15913        self.matmul_pre(w, aq, ad, &x0, m)
15914    }
15915
15916    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
15917    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
15918    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
15919    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
15920    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
15921    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
15922    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
15923    /// per-tensor path.
15924    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
15925    pub fn matmul_decode_exact_dual_pre(
15926        &self,
15927        w0: &crate::model::GpuTensor,
15928        w1: &crate::model::GpuTensor,
15929        aq: &CudaSlice<i8>,
15930        ad: &CudaSlice<f32>,
15931        m: usize,
15932    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
15933    {
15934        use crate::model::GpuTensor;
15935        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15936        let on = *ON.get_or_init(|| {
15937            std::env::var("MEMRA_SPEC_DUAL_T")
15938                .map(|v| v != "0")
15939                .unwrap_or(true)
15940        });
15941        if !on
15942            || !(2..=7).contains(&m)
15943            || std::env::var("MEMRA_NO_BATCHED").is_ok()
15944            || !self.uses_q8_1_fast(w0)
15945            || !self.uses_q8_1_fast(w1)
15946        {
15947            return Ok(None);
15948        }
15949        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
15950        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
15951        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
15952        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
15953        if !self.mmvq_supports(QT_NVFP4) {
15954            return Ok(None);
15955        }
15956        let (in_f, out_f) = (w0.in_features(), w0.out_features());
15957        if w1.in_features() != in_f || w1.out_features() != out_f {
15958            return Ok(None);
15959        }
15960        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
15961            (
15962                GpuTensor::Quant {
15963                    bytes: b0,
15964                    qtype: q0,
15965                    row_bytes: rb0,
15966                    scale: s0,
15967                    rp: rp0,
15968                    rp4: None,
15969                    ..
15970                },
15971                GpuTensor::Quant {
15972                    bytes: b1,
15973                    qtype: q1,
15974                    row_bytes: rb1,
15975                    scale: s1,
15976                    rp: rp1,
15977                    rp4: None,
15978                    ..
15979                },
15980            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
15981                (b0, b1, *rb0, *s0, *s1, *rp0)
15982            }
15983            _ => return Ok(None),
15984        };
15985        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
15986        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
15987        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
15988        {
15989            return Ok(None);
15990        }
15991        let (y0, y1) =
15992            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
15993        Ok(Some(((y0, s0), (y1, s1))))
15994    }
15995
15996    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
15997    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
15998    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
15999    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
16000    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
16001    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
16002    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
16003    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
16004    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
16005    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
16006    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
16007    pub fn matmul_decode_exact_group4_pre(
16008        &self,
16009        ws: [&crate::model::GpuTensor; 4],
16010        aq: &CudaSlice<i8>,
16011        ad: &CudaSlice<f32>,
16012        m: usize,
16013    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16014        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16015        let on = *ON.get_or_init(|| {
16016            std::env::var("MEMRA_TK_GDN_GROUP")
16017                .map(|v| v != "0")
16018                .unwrap_or(true)
16019        });
16020        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
16021    }
16022
16023    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
16024    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
16025    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
16026    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
16027    pub fn matmul_decode_exact_group3_pre(
16028        &self,
16029        ws: [&crate::model::GpuTensor; 3],
16030        aq: &CudaSlice<i8>,
16031        ad: &CudaSlice<f32>,
16032        m: usize,
16033    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16034        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16035        let on = *ON.get_or_init(|| {
16036            std::env::var("MEMRA_TK_FA_GROUP")
16037                .map(|v| v != "0")
16038                .unwrap_or(true)
16039        });
16040        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
16041    }
16042
16043    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
16044    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
16045    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
16046    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
16047    fn matmul_decode_exact_group_pre(
16048        &self,
16049        ws: &[&crate::model::GpuTensor],
16050        aq: &CudaSlice<i8>,
16051        ad: &CudaSlice<f32>,
16052        m: usize,
16053        on: bool,
16054        tag: &'static str,
16055    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16056        use crate::model::GpuTensor;
16057        if !on
16058            || !(2..=16).contains(&m)
16059            || std::env::var("MEMRA_NO_BATCHED").is_ok()
16060            || (m > 4 && !Self::b8_enabled())
16061            || !self.mmvq_supports(QT_NVFP4)
16062            || !self.batched_supports(QT_NVFP4)
16063        {
16064            return Ok(None);
16065        }
16066        let in_f = ws[0].in_features();
16067        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
16068        for w in ws {
16069            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
16070                return Ok(None);
16071            }
16072            match w {
16073                GpuTensor::Quant {
16074                    bytes,
16075                    qtype,
16076                    scale,
16077                    rp: true,
16078                    rp4: None,
16079                    ..
16080                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
16081                    parts.push((bytes, w.out_features(), *scale));
16082                }
16083                _ => return Ok(None),
16084            }
16085        }
16086        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
16087        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16088        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16089        let mcols = if (5..=7).contains(&m) && b567 {
16090            m
16091        } else {
16092            Self::batched_mcols(m)
16093        };
16094        let kname: &'static str = match mcols {
16095            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
16096            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
16097            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
16098            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
16099            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
16100            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
16101            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
16102            _ => return Ok(None),
16103        };
16104        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
16105        // the second door's print on the slice-D battery — key the once-set by tag.
16106        if std::env::var("MEMRA_DEBUG").is_ok() {
16107            use std::sync::Mutex;
16108            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
16109            let mut seen = SEEN.lock().unwrap();
16110            if !seen.contains(&tag) {
16111                seen.push(tag);
16112                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
16113            }
16114        }
16115        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
16116        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
16117        let total: usize = parts.iter().map(|p| p.1).sum();
16118        let three = parts.len() == 3;
16119        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
16120        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
16121        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
16122        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
16123        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
16124        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
16125        let cfg = LaunchConfig {
16126            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16127            block_dim: (32, ROWS_PER_BLOCK, 1),
16128            shared_mem_bytes: 0,
16129        };
16130        let (inf, mi) = (in_f as i32, m as i32);
16131        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
16132        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
16133        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
16134        let s3 = if three { 1.0f32 } else { parts[3].2 };
16135        let w3 = if three { parts[0].0 } else { parts[3].0 };
16136        let f = self.func(kname);
16137        let __s_b = self.gpu.stream();
16138        let mut b = __s_b.launch_builder(&f);
16139        b.arg(parts[0].0)
16140            .arg(parts[1].0)
16141            .arg(parts[2].0)
16142            .arg(w3)
16143            .arg(aq)
16144            .arg(ad)
16145            .arg(&mut y0)
16146            .arg(&mut y1)
16147            .arg(&mut y2)
16148            .arg(&mut y3)
16149            .arg(&inf)
16150            .arg(&n0)
16151            .arg(&n1)
16152            .arg(&n2)
16153            .arg(&n3)
16154            .arg(&mi)
16155            .arg(&s0)
16156            .arg(&s1)
16157            .arg(&s2)
16158            .arg(&s3);
16159        unsafe {
16160            b.launch(cfg)?;
16161        }
16162        Ok(Some(if three {
16163            vec![y0, y1, y2]
16164        } else {
16165            vec![y0, y1, y2, y3]
16166        }))
16167    }
16168
16169    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
16170    /// launch computes both FFN projections of a verify batch — same activation, same shape,
16171    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
16172    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
16173    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
16174    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
16175    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
16176    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
16177    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
16178    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
16179    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
16180    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
16181    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
16182    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
16183    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
16184    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16185    pub fn matmul_decode_exact_dual(
16186        &self,
16187        w0: &crate::model::GpuTensor,
16188        w1: &crate::model::GpuTensor,
16189        x: &CudaSlice<f32>,
16190        m: usize,
16191    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
16192        use crate::model::GpuTensor;
16193        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16194        let on = *ON.get_or_init(|| {
16195            std::env::var("MEMRA_SPEC_DUAL_T")
16196                .map(|v| v != "0")
16197                .unwrap_or(true)
16198        });
16199        if !on
16200            || !(2..=4).contains(&m)
16201            || std::env::var("MEMRA_NO_BATCHED").is_ok()
16202            || !self.uses_q8_1_fast(w0)
16203            || !self.uses_q8_1_fast(w1)
16204        {
16205            return Ok(None);
16206        }
16207        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
16208        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
16209        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
16210        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
16211        if !self.mmvq_supports(QT_NVFP4) {
16212            return Ok(None);
16213        }
16214        let (in_f, out_f) = (w0.in_features(), w0.out_features());
16215        if w1.in_features() != in_f || w1.out_features() != out_f {
16216            return Ok(None);
16217        }
16218        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
16219            (
16220                GpuTensor::Quant {
16221                    bytes: b0,
16222                    qtype: q0,
16223                    row_bytes: rb0,
16224                    scale: s0,
16225                    rp: rp0,
16226                    rp4: None,
16227                    ..
16228                },
16229                GpuTensor::Quant {
16230                    bytes: b1,
16231                    qtype: q1,
16232                    row_bytes: rb1,
16233                    scale: s1,
16234                    rp: rp1,
16235                    rp4: None,
16236                    ..
16237                },
16238            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
16239                (b0, b1, *rb0, *s0, *s1, *rp0)
16240            }
16241            _ => return Ok(None),
16242        };
16243        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
16244        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
16245        if std::env::var("MEMRA_DEBUG").is_ok() {
16246            static ONCE: std::sync::Once = std::sync::Once::new();
16247            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
16248        }
16249        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16250        let (y0, y1) =
16251            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
16252        let mut y0 = y0;
16253        let mut y1 = y1;
16254        if s0 != 1.0 {
16255            self.scale_inplace(&mut y0, s0, m * out_f)?;
16256        }
16257        if s1 != 1.0 {
16258            self.scale_inplace(&mut y1, s1, m * out_f)?;
16259        }
16260        Ok(Some((y0, y1)))
16261    }
16262
16263    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
16264    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
16265    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
16266    /// twins (both buffers must be the repacked layout).
16267    #[allow(clippy::too_many_arguments)]
16268    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
16269    pub fn qmatvec_batched_dual_raw(
16270        &self,
16271        b0: &CudaSlice<u8>,
16272        b1: &CudaSlice<u8>,
16273        aq: &CudaSlice<i8>,
16274        ad: &CudaSlice<f32>,
16275        m: usize,
16276        in_f: usize,
16277        out_f: usize,
16278        row_bytes: usize,
16279        rp: bool,
16280    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16281        const ROWS_PER_BLOCK: u32 = 4;
16282        let mcols = Self::batched_mcols(m);
16283        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
16284        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
16285        let tiny_rp1 = rp
16286            && mcols == 4
16287            && out_f <= 128
16288            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
16289        let (name, rows_per_block) = if tiny_rp1 {
16290            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
16291        } else {
16292            match (mcols, rp, m) {
16293                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
16294                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
16295                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
16296                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
16297                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
16298                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
16299                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
16300                _ => {
16301                    return Err(
16302                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
16303                    );
16304                }
16305            }
16306        };
16307        let f = self.func(name);
16308        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
16309        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
16310        let cfg = LaunchConfig {
16311            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
16312            block_dim: (32, ROWS_PER_BLOCK, 1),
16313            shared_mem_bytes: 0,
16314        };
16315        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16316        let __s_b = self.gpu.stream();
16317        let mut b = __s_b.launch_builder(&f);
16318        b.arg(b0)
16319            .arg(b1)
16320            .arg(aq)
16321            .arg(ad)
16322            .arg(&mut y0)
16323            .arg(&mut y1)
16324            .arg(&inf)
16325            .arg(&outf)
16326            .arg(&mi)
16327            .arg(&rb);
16328        unsafe {
16329            b.launch(cfg)?;
16330        }
16331        Ok((y0, y1))
16332    }
16333
16334    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
16335    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
16336    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
16337    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
16338    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
16339    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
16340    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
16341    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
16342    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
16343    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
16344    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
16345    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16346    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
16347    pub fn matmul_pre_dual_noscale(
16348        &self,
16349        w0: &crate::model::GpuTensor,
16350        w1: &crate::model::GpuTensor,
16351        aq: &CudaSlice<i8>,
16352        ad: &CudaSlice<f32>,
16353        m: usize,
16354    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
16355    {
16356        use crate::model::GpuTensor;
16357        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
16358            return Ok(None);
16359        }
16360        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
16361        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
16362        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
16363        // would mix dispatch families across the pair — the exact class `q8_fused_params`
16364        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
16365        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
16366        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
16367        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
16368        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
16369        if !self.mmvq_supports(QT_NVFP4) {
16370            return Ok(None);
16371        }
16372        let (in_f, out_f) = (w0.in_features(), w0.out_features());
16373        if w1.in_features() != in_f || w1.out_features() != out_f {
16374            return Ok(None);
16375        }
16376        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
16377        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
16378        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
16379        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
16380        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
16381        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
16382        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
16383        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
16384        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
16385        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
16386        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
16387        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
16388        let no_mirror =
16389            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
16390        if self.q8_ffn_fuse2_on()
16391            && no_mirror(w0)
16392            && no_mirror(w1)
16393            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
16394        {
16395            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
16396            return Ok(Some(((y0, 1.0), (y1, 1.0))));
16397        }
16398        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
16399        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
16400        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
16401        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
16402        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
16403        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
16404        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
16405        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
16406        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
16407        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
16408            let (y0, y1) =
16409                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
16410            return Ok(Some(((y0, p0.3), (y1, p1.3))));
16411        }
16412        let (b0, q0, rb0, s0, rp0) = match w0 {
16413            GpuTensor::Quant {
16414                bytes,
16415                qtype,
16416                row_bytes,
16417                scale,
16418                rp,
16419                ..
16420            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16421            _ => return Ok(None),
16422        };
16423        let (b1, q1, rb1, s1, rp1) = match w1 {
16424            GpuTensor::Quant {
16425                bytes,
16426                qtype,
16427                row_bytes,
16428                scale,
16429                rp,
16430                ..
16431            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16432            _ => return Ok(None),
16433        };
16434        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
16435            return Ok(None);
16436        }
16437        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
16438        const RPW: u32 = 2;
16439        let rows_per_block = ROWS_PER_BLOCK * RPW;
16440        let f = self.func(if rp0 {
16441            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
16442        } else {
16443            "qmatvec_nvfp4_mmvq_dual_mr2"
16444        });
16445        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
16446        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
16447        let cfg = LaunchConfig {
16448            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
16449            block_dim: (32, ROWS_PER_BLOCK, 1),
16450            shared_mem_bytes: 0,
16451        };
16452        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
16453        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
16454        // yscale args stay 1.0 here (they exist for the single-tensor callers).
16455        let one = 1.0f32;
16456        let __s_b = self.gpu.stream();
16457        let mut b = __s_b.launch_builder(&f);
16458        b.arg(b0)
16459            .arg(b1)
16460            .arg(aq)
16461            .arg(ad)
16462            .arg(&mut y0)
16463            .arg(&mut y1)
16464            .arg(&inf)
16465            .arg(&outf)
16466            .arg(&mi)
16467            .arg(&rb)
16468            .arg(&one)
16469            .arg(&one);
16470        unsafe {
16471            b.launch(cfg)?;
16472        }
16473        Ok(Some(((y0, s0), (y1, s1))))
16474    }
16475
16476    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
16477    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
16478    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
16479    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
16480    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
16481    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
16482    /// back to the three singles.
16483    #[allow(clippy::too_many_arguments)]
16484    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16485    pub fn matmul_nvfp4_fused3(
16486        &self,
16487        w0: &crate::model::GpuTensor,
16488        w1: &crate::model::GpuTensor,
16489        w2: &crate::model::GpuTensor,
16490        aq: &CudaSlice<i8>,
16491        ad: &CudaSlice<f32>,
16492        m: usize,
16493    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
16494    {
16495        use crate::model::GpuTensor;
16496        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
16497        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
16498        // verbatim, weight rows read once for all m columns, bit-identical per
16499        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
16500        // segments would re-read the weight per row" note described the grid.y=m lift,
16501        // which this twin deliberately is NOT.
16502        if !self.mmvq_supports(QT_NVFP4)
16503            || !self.uses_q8_1_fast(w0)
16504            || !self.uses_q8_1_fast(w1)
16505            || !self.uses_q8_1_fast(w2)
16506        {
16507            return Ok(None);
16508        }
16509        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
16510        // door — same family and bit-identity law as the fused4 delegate above.
16511        if (9..=16).contains(&m) {
16512            return Ok(
16513                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
16514                    Some(mut ys) => {
16515                        let y2 = ys.pop().unwrap();
16516                        let y1 = ys.pop().unwrap();
16517                        let y0 = ys.pop().unwrap();
16518                        Some((y0, y1, y2))
16519                    }
16520                    None => None,
16521                },
16522            );
16523        }
16524        if !(1..=8).contains(&m) {
16525            return Ok(None);
16526        }
16527        if m > 1 {
16528            let in_f = w0.in_features();
16529            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
16530                || !self.batched_supports(QT_NVFP4)
16531                || std::env::var("MEMRA_NO_BATCHED").is_ok()
16532                || (m > 4 && !Self::b8_enabled())
16533                || !in_f.is_multiple_of(512)
16534                || in_f / 64 > 272
16535            {
16536                return Ok(None);
16537            }
16538        }
16539        let unpack = |w: &crate::model::GpuTensor| match w {
16540            GpuTensor::Quant {
16541                bytes,
16542                qtype,
16543                scale,
16544                rp,
16545                ..
16546            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16547            _ => None,
16548        };
16549        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
16550            return Ok(None);
16551        };
16552        let in_f = w0.in_features();
16553        if w1.in_features() != in_f || w2.in_features() != in_f {
16554            return Ok(None);
16555        }
16556        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
16557        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16558        const RPW: u32 = 2;
16559        let rows_pb = ROWS_PER_BLOCK * RPW;
16560        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16561        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16562        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16563        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
16564        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
16565        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16566        // only dereferenced for the launch-arg build inside this call.
16567        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
16568        if m > 1 {
16569            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
16570            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
16571                return Ok(None);
16572            }
16573            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
16574            let cfg = LaunchConfig {
16575                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
16576                block_dim: (32, ROWS_PER_BLOCK, 1),
16577                shared_mem_bytes: 0,
16578            };
16579            let __s_b = self.gpu.stream();
16580            let mut b = __s_b.launch_builder(&f);
16581            b.arg(b0)
16582                .arg(b1)
16583                .arg(b2)
16584                .arg(aq)
16585                .arg(ad)
16586                .arg(&mut y0)
16587                .arg(&mut y1)
16588                .arg(&mut y2)
16589                .arg(&inf)
16590                .arg(&oi0)
16591                .arg(&oi1)
16592                .arg(&oi2)
16593                .arg(&mi);
16594            unsafe {
16595                b.launch(cfg)?;
16596            }
16597            return Ok(Some((y0, y1, y2)));
16598        }
16599        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
16600        let cfg = LaunchConfig {
16601            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
16602            block_dim: (32, ROWS_PER_BLOCK, 1),
16603            shared_mem_bytes: 0,
16604        };
16605        let __s_b = self.gpu.stream();
16606        let mut b = __s_b.launch_builder(&f);
16607        b.arg(b0)
16608            .arg(b1)
16609            .arg(b2)
16610            .arg(aq)
16611            .arg(ad)
16612            .arg(&mut y0)
16613            .arg(&mut y1)
16614            .arg(&mut y2)
16615            .arg(&inf)
16616            .arg(&oi0)
16617            .arg(&oi1)
16618            .arg(&oi2)
16619            .arg(&mi)
16620            .arg(&p0.1)
16621            .arg(&p1.1)
16622            .arg(&p2.1);
16623        unsafe {
16624            b.launch(cfg)?;
16625        }
16626        Ok(Some((y0, y1, y2)))
16627    }
16628
16629    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
16630    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
16631    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
16632    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
16633    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
16634    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
16635    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
16636    /// same-binary interleaved A/B arm.
16637    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16638    pub fn matmul_nvfp4_fused2(
16639        &self,
16640        w0: &crate::model::GpuTensor,
16641        w1: &crate::model::GpuTensor,
16642        aq: &CudaSlice<i8>,
16643        ad: &CudaSlice<f32>,
16644        m: usize,
16645    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
16646        use crate::model::GpuTensor;
16647        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16648        let off =
16649            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
16650        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
16651        // read serves all m rows); the fused segments would re-read the weight per row.
16652        if off
16653            || m != 1
16654            || !self.mmvq_supports(QT_NVFP4)
16655            || !self.uses_q8_1_fast(w0)
16656            || !self.uses_q8_1_fast(w1)
16657        {
16658            return Ok(None);
16659        }
16660        let unpack = |w: &crate::model::GpuTensor| match w {
16661            GpuTensor::Quant {
16662                bytes,
16663                qtype,
16664                scale,
16665                rp,
16666                ..
16667            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16668            _ => None,
16669        };
16670        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
16671            return Ok(None);
16672        };
16673        let in_f = w0.in_features();
16674        if w1.in_features() != in_f {
16675            return Ok(None);
16676        }
16677        let (o0, o1) = (w0.out_features(), w1.out_features());
16678        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16679        const RPW: u32 = 2;
16680        let rows_pb = ROWS_PER_BLOCK * RPW;
16681        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16682        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
16683        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16684        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16685        let cfg = LaunchConfig {
16686            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
16687            block_dim: (32, ROWS_PER_BLOCK, 1),
16688            shared_mem_bytes: 0,
16689        };
16690        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
16691        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16692        // only dereferenced for the launch-arg build inside this call.
16693        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
16694        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
16695        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
16696        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
16697            {
16698                use cudarc::driver::{DevicePtr, DevicePtrMut};
16699                let s = &self.gpu.stream();
16700                let (pw0, _g0) = b0.device_ptr(s);
16701                let (pw1, _g1) = b1.device_ptr(s);
16702                let (paq, _g2) = aq.device_ptr(s);
16703                let (pad, _g3) = ad.device_ptr(s);
16704                let (py0, _g4) = y0.device_ptr_mut(s);
16705                let (py1, _g5) = y1.device_ptr_mut(s);
16706                let (s0, s1) = (p0.1, p1.1);
16707                let mut ps = [
16708                    &pw0 as *const _ as *mut std::ffi::c_void,
16709                    &pw1 as *const _ as *mut _,
16710                    &paq as *const _ as *mut _,
16711                    &pad as *const _ as *mut _,
16712                    &py0 as *const _ as *mut _,
16713                    &py1 as *const _ as *mut _,
16714                    &inf as *const _ as *mut _,
16715                    &oi0 as *const _ as *mut _,
16716                    &oi1 as *const _ as *mut _,
16717                    &mi as *const _ as *mut _,
16718                    &s0 as *const _ as *mut _,
16719                    &s1 as *const _ as *mut _,
16720                ];
16721                unsafe {
16722                    self.launch_pdl(
16723                        "qmatvec_nvfp4_mmvq_fused2_rp",
16724                        cfg.grid_dim,
16725                        cfg.block_dim,
16726                        &mut ps,
16727                    )?;
16728                }
16729            }
16730            return Ok(Some((y0, y1)));
16731        }
16732        let __s_b = self.gpu.stream();
16733        let mut b = __s_b.launch_builder(&f);
16734        b.arg(b0)
16735            .arg(b1)
16736            .arg(aq)
16737            .arg(ad)
16738            .arg(&mut y0)
16739            .arg(&mut y1)
16740            .arg(&inf)
16741            .arg(&oi0)
16742            .arg(&oi1)
16743            .arg(&mi)
16744            .arg(&p0.1)
16745            .arg(&p1.1);
16746        unsafe {
16747            b.launch(cfg)?;
16748        }
16749        Ok(Some((y0, y1)))
16750    }
16751
16752    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
16753    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
16754    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
16755    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
16756    pub fn matmul_nvfp4_fused2_into(
16757        &self,
16758        w0: &crate::model::GpuTensor,
16759        w1: &crate::model::GpuTensor,
16760        aq: &CudaSlice<i8>,
16761        ad: &CudaSlice<f32>,
16762        y0: &mut CudaSlice<f32>,
16763        y1: &mut CudaSlice<f32>,
16764    ) -> Result<bool, Box<dyn std::error::Error>> {
16765        use crate::model::GpuTensor;
16766        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16767        let off =
16768            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
16769        if off
16770            || !self.mmvq_supports(QT_NVFP4)
16771            || !self.uses_q8_1_fast(w0)
16772            || !self.uses_q8_1_fast(w1)
16773        {
16774            return Ok(false);
16775        }
16776        let unpack = |w: &crate::model::GpuTensor| match w {
16777            GpuTensor::Quant {
16778                bytes,
16779                qtype,
16780                scale,
16781                rp,
16782                ..
16783            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16784            _ => None,
16785        };
16786        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
16787            return Ok(false);
16788        };
16789        let in_f = w0.in_features();
16790        if w1.in_features() != in_f {
16791            return Ok(false);
16792        }
16793        let (o0, o1) = (w0.out_features(), w1.out_features());
16794        if y0.len() < o0 || y1.len() < o1 {
16795            return Ok(false);
16796        }
16797        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16798        const RPW: u32 = 2;
16799        let rows_pb = ROWS_PER_BLOCK * RPW;
16800        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16801        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
16802        let cfg = LaunchConfig {
16803            grid_dim: (nb(o0) + nb(o1), 1, 1),
16804            block_dim: (32, ROWS_PER_BLOCK, 1),
16805            shared_mem_bytes: 0,
16806        };
16807        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
16808        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16809        // only dereferenced for the launch-arg build inside this call.
16810        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
16811        let __s_b = self.gpu.stream();
16812        let mut b = __s_b.launch_builder(&f);
16813        b.arg(b0)
16814            .arg(b1)
16815            .arg(aq)
16816            .arg(ad)
16817            .arg(&mut *y0)
16818            .arg(&mut *y1)
16819            .arg(&inf)
16820            .arg(&oi0)
16821            .arg(&oi1)
16822            .arg(&mi)
16823            .arg(&p0.1)
16824            .arg(&p1.1);
16825        unsafe {
16826            b.launch(cfg)?;
16827        }
16828        Ok(true)
16829    }
16830
16831    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
16832    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
16833    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
16834    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
16835    #[allow(clippy::type_complexity)]
16836    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16837    pub fn matmul_nvfp4_fused4(
16838        &self,
16839        w0: &crate::model::GpuTensor,
16840        w1: &crate::model::GpuTensor,
16841        w2: &crate::model::GpuTensor,
16842        w3: &crate::model::GpuTensor,
16843        aq: &CudaSlice<i8>,
16844        ad: &CudaSlice<f32>,
16845        m: usize,
16846    ) -> Result<
16847        Option<(
16848            CudaSlice<f32>,
16849            CudaSlice<f32>,
16850            CudaSlice<f32>,
16851            CudaSlice<f32>,
16852        )>,
16853        Box<dyn std::error::Error>,
16854    > {
16855        use crate::model::GpuTensor;
16856        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
16857        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
16858        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
16859        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
16860        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
16861        // Admission mirrors the singles' batched gates below.
16862        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
16863            || !self.mmvq_supports(QT_NVFP4)
16864            || !self.uses_q8_1_fast(w0)
16865            || !self.uses_q8_1_fast(w1)
16866            || !self.uses_q8_1_fast(w2)
16867            || !self.uses_q8_1_fast(w3)
16868        {
16869            return Ok(None);
16870        }
16871        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
16872        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
16873        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
16874        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
16875        if (9..=16).contains(&m) {
16876            return Ok(
16877                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
16878                    Some(mut ys) => {
16879                        let y3 = ys.pop().unwrap();
16880                        let y2 = ys.pop().unwrap();
16881                        let y1 = ys.pop().unwrap();
16882                        let y0 = ys.pop().unwrap();
16883                        Some((y0, y1, y2, y3))
16884                    }
16885                    None => None,
16886                },
16887            );
16888        }
16889        if !(1..=8).contains(&m) {
16890            return Ok(None);
16891        }
16892        if m > 1 {
16893            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
16894            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
16895            let in_f = w0.in_features();
16896            if !self.batched_supports(QT_NVFP4)
16897                || std::env::var("MEMRA_NO_BATCHED").is_ok()
16898                || (m > 4 && !Self::b8_enabled())
16899                || !in_f.is_multiple_of(512)
16900                || in_f / 64 > 272
16901            {
16902                return Ok(None);
16903            }
16904        }
16905        let unpack = |w: &crate::model::GpuTensor| match w {
16906            GpuTensor::Quant {
16907                bytes,
16908                qtype,
16909                scale,
16910                rp,
16911                ..
16912            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16913            _ => None,
16914        };
16915        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
16916            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
16917        else {
16918            return Ok(None);
16919        };
16920        let in_f = w0.in_features();
16921        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
16922            return Ok(None);
16923        }
16924        let (o0, o1, o2, o3) = (
16925            w0.out_features(),
16926            w1.out_features(),
16927            w2.out_features(),
16928            w3.out_features(),
16929        );
16930        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16931        const RPW: u32 = 2;
16932        let rows_pb = ROWS_PER_BLOCK * RPW;
16933        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16934        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16935        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16936        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
16937        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
16938        let (inf, oi0, oi1, oi2, oi3, mi) = (
16939            in_f as i32,
16940            o0 as i32,
16941            o1 as i32,
16942            o2 as i32,
16943            o3 as i32,
16944            m as i32,
16945        );
16946        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16947        // only dereferenced for the launch-arg build inside this call.
16948        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
16949        if m > 1 {
16950            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
16951            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
16952            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
16953                return Ok(None);
16954            }
16955            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
16956            let cfg = LaunchConfig {
16957                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
16958                block_dim: (32, ROWS_PER_BLOCK, 1),
16959                shared_mem_bytes: 0,
16960            };
16961            let __s_b = self.gpu.stream();
16962            let mut b = __s_b.launch_builder(&f);
16963            b.arg(b0)
16964                .arg(b1)
16965                .arg(b2)
16966                .arg(b3)
16967                .arg(aq)
16968                .arg(ad)
16969                .arg(&mut y0)
16970                .arg(&mut y1)
16971                .arg(&mut y2)
16972                .arg(&mut y3)
16973                .arg(&inf)
16974                .arg(&oi0)
16975                .arg(&oi1)
16976                .arg(&oi2)
16977                .arg(&oi3)
16978                .arg(&mi);
16979            unsafe {
16980                b.launch(cfg)?;
16981            }
16982            return Ok(Some((y0, y1, y2, y3)));
16983        }
16984        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
16985        let cfg = LaunchConfig {
16986            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
16987            block_dim: (32, ROWS_PER_BLOCK, 1),
16988            shared_mem_bytes: 0,
16989        };
16990        let __s_b = self.gpu.stream();
16991        let mut b = __s_b.launch_builder(&f);
16992        b.arg(b0)
16993            .arg(b1)
16994            .arg(b2)
16995            .arg(b3)
16996            .arg(aq)
16997            .arg(ad)
16998            .arg(&mut y0)
16999            .arg(&mut y1)
17000            .arg(&mut y2)
17001            .arg(&mut y3)
17002            .arg(&inf)
17003            .arg(&oi0)
17004            .arg(&oi1)
17005            .arg(&oi2)
17006            .arg(&oi3)
17007            .arg(&mi)
17008            .arg(&p0.1)
17009            .arg(&p1.1)
17010            .arg(&p2.1)
17011            .arg(&p3.1);
17012        unsafe {
17013            b.launch(cfg)?;
17014        }
17015        Ok(Some((y0, y1, y2, y3)))
17016    }
17017
17018    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
17019    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
17020    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
17021    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
17022    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
17023    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
17024    /// back to the per-tensor path.
17025    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17026    pub fn matmul_q8_fused2(
17027        &self,
17028        w0: &crate::model::GpuTensor,
17029        w1: &crate::model::GpuTensor,
17030        aq: &CudaSlice<i8>,
17031        ad: &CudaSlice<f32>,
17032    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17033        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
17034        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
17035        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
17036        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
17037        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
17038        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17039            return Ok(Some(self.e4m3_fused2_core(
17040                p0.0,
17041                p1.0,
17042                aq,
17043                ad,
17044                w0.in_features(),
17045                p0.1,
17046                p1.1,
17047                p0.2,
17048                p0.3,
17049                p1.3,
17050            )?));
17051        }
17052        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
17053            return Ok(None);
17054        };
17055        Ok(Some(self.q8_fused2_core(
17056            p0.0,
17057            p1.0,
17058            aq,
17059            ad,
17060            w0.in_features(),
17061            p0.1,
17062            p1.1,
17063            p0.2,
17064        )?))
17065    }
17066
17067    #[allow(clippy::too_many_arguments)]
17068    fn q8_fused2_core(
17069        &self,
17070        b0: &CudaSlice<u8>,
17071        b1: &CudaSlice<u8>,
17072        aq: &CudaSlice<i8>,
17073        ad: &CudaSlice<f32>,
17074        in_f: usize,
17075        out0: usize,
17076        out1: usize,
17077        row_bytes: usize,
17078    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17079        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
17080        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
17081        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
17082        let f = self.func("qmatvec_q8_0_mmvq_fused2");
17083        let mut y0 = self.alloc_uninit::<f32>(out0)?;
17084        let mut y1 = self.alloc_uninit::<f32>(out1)?;
17085        let cfg = LaunchConfig {
17086            grid_dim: (nb0 + nb1, 1, 1),
17087            block_dim: (32, ROWS_PER_BLOCK, 1),
17088            shared_mem_bytes: 0,
17089        };
17090        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
17091        let __s_b = self.gpu.stream();
17092        let mut b = __s_b.launch_builder(&f);
17093        b.arg(b0)
17094            .arg(b1)
17095            .arg(aq)
17096            .arg(ad)
17097            .arg(&mut y0)
17098            .arg(&mut y1)
17099            .arg(&inf)
17100            .arg(&o0)
17101            .arg(&o1)
17102            .arg(&rbl);
17103        unsafe {
17104            b.launch(cfg)?;
17105        }
17106        Ok((y0, y1))
17107    }
17108
17109    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
17110    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
17111    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
17112    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
17113    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
17114    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17115    pub fn matmul_q8_fused2_x(
17116        &self,
17117        w0: &crate::model::GpuTensor,
17118        w1: &crate::model::GpuTensor,
17119        x: &CudaSlice<f32>,
17120    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17121        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
17122            return Ok(None);
17123        }
17124        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17125            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
17126            return Ok(Some(self.e4m3_fused2_core(
17127                p0.0,
17128                p1.0,
17129                &aq,
17130                &ad,
17131                w0.in_features(),
17132                p0.1,
17133                p1.1,
17134                p0.2,
17135                p0.3,
17136                p1.3,
17137            )?));
17138        }
17139        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
17140            return Ok(None);
17141        };
17142        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
17143        Ok(Some(self.q8_fused2_core(
17144            p0.0,
17145            p1.0,
17146            &aq,
17147            &ad,
17148            w0.in_features(),
17149            p0.1,
17150            p1.1,
17151            p0.2,
17152        )?))
17153    }
17154
17155    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
17156    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
17157    #[allow(clippy::too_many_arguments)]
17158    pub fn qmatvec_q8_fused2_raw(
17159        &self,
17160        b0: &CudaSlice<u8>,
17161        b1: &CudaSlice<u8>,
17162        x: &CudaSlice<f32>,
17163        in_f: usize,
17164        out0: usize,
17165        out1: usize,
17166        row_bytes: usize,
17167    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17168        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
17169        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
17170    }
17171
17172    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
17173    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
17174    /// (tensor,row) to three separate m=1 MMVQ launches.
17175    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
17176    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
17177    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17178    pub fn matmul_q4_fused3(
17179        &self,
17180        w0: &crate::model::GpuTensor,
17181        w1: &crate::model::GpuTensor,
17182        w2: &crate::model::GpuTensor,
17183        aq: &CudaSlice<i8>,
17184        ad: &CudaSlice<f32>,
17185    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17186    {
17187        use crate::model::GpuTensor;
17188        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17189            match w {
17190                GpuTensor::Quant {
17191                    qtype, row_bytes, ..
17192                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17193                _ => None,
17194            }
17195        };
17196        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
17197            return Ok(None);
17198        };
17199        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17200            return Ok(None);
17201        }
17202        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
17203        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
17204        // the separate matvecs (each routes its own rp).
17205        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17206            match w {
17207                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17208                    Some(m) => (m, true),
17209                    None => (bytes, *rp),
17210                },
17211                _ => unreachable!(),
17212            }
17213        }
17214        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17215        if rp0 != rp1 || rp1 != rp2 {
17216            return Ok(None);
17217        }
17218        let rp = rp0;
17219        let rpb: u32 = 4;
17220        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
17221        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
17222        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
17223        let mr1 = rp && Self::q40_mr1_on();
17224        let nb = |o: usize| {
17225            if mr1 {
17226                (o as u32).div_ceil(rpb)
17227            } else {
17228                (o as u32).div_ceil(2).div_ceil(rpb)
17229            }
17230        };
17231        let grid = nb(o0) + nb(o1) + nb(o2);
17232        let mut y0 = self.alloc_uninit::<f32>(o0)?;
17233        let mut y1 = self.alloc_uninit::<f32>(o1)?;
17234        let mut y2 = self.alloc_uninit::<f32>(o2)?;
17235        let f = self.func(if mr1 {
17236            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
17237        } else if rp {
17238            "qmatvec_q4_0_mmvq_fused3_rp"
17239        } else {
17240            "qmatvec_q4_0_mmvq_fused3"
17241        });
17242        let cfg = LaunchConfig {
17243            grid_dim: (grid, 1, 1),
17244            block_dim: (32, rpb, 1),
17245            shared_mem_bytes: 0,
17246        };
17247        let inf = w0.in_features() as i32;
17248        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
17249        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
17250        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
17251        // variant may take the programmatic-serialization launch.
17252        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17253            {
17254                use cudarc::driver::{DevicePtr, DevicePtrMut};
17255                let s = &self.gpu.stream();
17256                let (p0, _g0) = b0.device_ptr(s);
17257                let (p1, _g1) = b1.device_ptr(s);
17258                let (p2, _g2) = b2.device_ptr(s);
17259                let (paq, _g3) = aq.device_ptr(s);
17260                let (pad, _g4) = ad.device_ptr(s);
17261                let (py0, _g5) = y0.device_ptr_mut(s);
17262                let (py1, _g6) = y1.device_ptr_mut(s);
17263                let (py2, _g7) = y2.device_ptr_mut(s);
17264                let mut ps = [
17265                    &p0 as *const _ as *mut std::ffi::c_void,
17266                    &p1 as *const _ as *mut _,
17267                    &p2 as *const _ as *mut _,
17268                    &paq as *const _ as *mut _,
17269                    &pad as *const _ as *mut _,
17270                    &py0 as *const _ as *mut _,
17271                    &py1 as *const _ as *mut _,
17272                    &py2 as *const _ as *mut _,
17273                    &inf as *const _ as *mut _,
17274                    &oo0 as *const _ as *mut _,
17275                    &oo1 as *const _ as *mut _,
17276                    &oo2 as *const _ as *mut _,
17277                    &r0 as *const _ as *mut _,
17278                    &r1 as *const _ as *mut _,
17279                    &r2 as *const _ as *mut _,
17280                ];
17281                unsafe {
17282                    self.launch_pdl(
17283                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
17284                        (grid, 1, 1),
17285                        (32, rpb, 1),
17286                        &mut ps,
17287                    )?;
17288                }
17289            }
17290            return Ok(Some((y0, y1, y2)));
17291        }
17292        let __s_b = self.gpu.stream();
17293        let mut b = __s_b.launch_builder(&f);
17294        b.arg(b0)
17295            .arg(b1)
17296            .arg(b2)
17297            .arg(aq)
17298            .arg(ad)
17299            .arg(&mut y0)
17300            .arg(&mut y1)
17301            .arg(&mut y2)
17302            .arg(&inf)
17303            .arg(&oo0)
17304            .arg(&oo1)
17305            .arg(&oo2)
17306            .arg(&r0)
17307            .arg(&r1)
17308            .arg(&r2);
17309        unsafe {
17310            b.launch(cfg)?;
17311        }
17312        Ok(Some((y0, y1, y2)))
17313    }
17314
17315    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
17316    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
17317    #[allow(clippy::too_many_arguments)]
17318    pub fn matmul_q4_fused3_into(
17319        &self,
17320        w0: &crate::model::GpuTensor,
17321        w1: &crate::model::GpuTensor,
17322        w2: &crate::model::GpuTensor,
17323        aq: &CudaSlice<i8>,
17324        ad: &CudaSlice<f32>,
17325        y0: &mut CudaSlice<f32>,
17326        y1: &mut CudaSlice<f32>,
17327        y2: &mut CudaSlice<f32>,
17328    ) -> Result<bool, Box<dyn std::error::Error>> {
17329        use crate::model::GpuTensor;
17330        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17331            match w {
17332                GpuTensor::Quant {
17333                    qtype, row_bytes, ..
17334                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17335                _ => None,
17336            }
17337        };
17338        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
17339            return Ok(false);
17340        };
17341        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17342            return Ok(false);
17343        }
17344        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17345            match w {
17346                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17347                    Some(m) => (m, true),
17348                    None => (bytes, *rp),
17349                },
17350                _ => unreachable!(),
17351            }
17352        }
17353        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17354        if rp0 != rp1 || rp1 != rp2 {
17355            return Ok(false);
17356        }
17357        let rp = rp0;
17358        let rpb: u32 = 4;
17359        let mr1 = rp && Self::q40_mr1_on();
17360        let nb = |o: usize| {
17361            if mr1 {
17362                (o as u32).div_ceil(rpb)
17363            } else {
17364                (o as u32).div_ceil(2).div_ceil(rpb)
17365            }
17366        };
17367        let grid = nb(o0) + nb(o1) + nb(o2);
17368        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
17369        let f = self.func(if mr1 {
17370            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
17371        } else if rp {
17372            "qmatvec_q4_0_mmvq_fused3_rp"
17373        } else {
17374            "qmatvec_q4_0_mmvq_fused3"
17375        });
17376        let cfg = LaunchConfig {
17377            grid_dim: (grid, 1, 1),
17378            block_dim: (32, rpb, 1),
17379            shared_mem_bytes: 0,
17380        };
17381        let inf = w0.in_features() as i32;
17382        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
17383        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
17384        // PDL wave-A: identical to the owned twin (capture-lane parity).
17385        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17386            use cudarc::driver::{DevicePtr, DevicePtrMut};
17387            let s = &self.gpu.stream();
17388            let (p0, _g0) = b0.device_ptr(s);
17389            let (p1, _g1) = b1.device_ptr(s);
17390            let (p2, _g2) = b2.device_ptr(s);
17391            let (paq, _g3) = aq.device_ptr(s);
17392            let (pad, _g4) = ad.device_ptr(s);
17393            let (py0, _g5) = y0.device_ptr_mut(s);
17394            let (py1, _g6) = y1.device_ptr_mut(s);
17395            let (py2, _g7) = y2.device_ptr_mut(s);
17396            let mut ps = [
17397                &p0 as *const _ as *mut std::ffi::c_void,
17398                &p1 as *const _ as *mut _,
17399                &p2 as *const _ as *mut _,
17400                &paq as *const _ as *mut _,
17401                &pad as *const _ as *mut _,
17402                &py0 as *const _ as *mut _,
17403                &py1 as *const _ as *mut _,
17404                &py2 as *const _ as *mut _,
17405                &inf as *const _ as *mut _,
17406                &oo0 as *const _ as *mut _,
17407                &oo1 as *const _ as *mut _,
17408                &oo2 as *const _ as *mut _,
17409                &r0 as *const _ as *mut _,
17410                &r1 as *const _ as *mut _,
17411                &r2 as *const _ as *mut _,
17412            ];
17413            unsafe {
17414                self.launch_pdl(
17415                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
17416                    (grid, 1, 1),
17417                    (32, rpb, 1),
17418                    &mut ps,
17419                )?;
17420            }
17421            return Ok(true);
17422        }
17423        let __s_b = self.gpu.stream();
17424        let mut b = __s_b.launch_builder(&f);
17425        b.arg(b0)
17426            .arg(b1)
17427            .arg(b2)
17428            .arg(aq)
17429            .arg(ad)
17430            .arg(&mut *y0)
17431            .arg(&mut *y1)
17432            .arg(&mut *y2)
17433            .arg(&inf)
17434            .arg(&oo0)
17435            .arg(&oo1)
17436            .arg(&oo2)
17437            .arg(&r0)
17438            .arg(&r1)
17439            .arg(&r2);
17440        unsafe {
17441            b.launch(cfg)?;
17442        }
17443        Ok(true)
17444    }
17445
17446    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
17447    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17448    pub fn matmul_q4_fused2(
17449        &self,
17450        w0: &crate::model::GpuTensor,
17451        w1: &crate::model::GpuTensor,
17452        aq: &CudaSlice<i8>,
17453        ad: &CudaSlice<f32>,
17454    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17455        use crate::model::GpuTensor;
17456        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17457            match w {
17458                GpuTensor::Quant {
17459                    qtype, row_bytes, ..
17460                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17461                _ => None,
17462            }
17463        };
17464        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
17465            return Ok(None);
17466        };
17467        if w0.in_features() != w1.in_features() {
17468            return Ok(None);
17469        }
17470        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
17471        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17472            match w {
17473                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17474                    Some(m) => (m, true),
17475                    None => (bytes, *rp),
17476                },
17477                _ => unreachable!(),
17478            }
17479        }
17480        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17481        if rp0 != rp1 {
17482            return Ok(None);
17483        }
17484        let rp = rp0;
17485        let rpb: u32 = 4;
17486        // mr1 twin — see matmul_q4_fused3.
17487        let mr1 = rp && Self::q40_mr1_on();
17488        let nb = |o: usize| {
17489            if mr1 {
17490                (o as u32).div_ceil(rpb)
17491            } else {
17492                (o as u32).div_ceil(2).div_ceil(rpb)
17493            }
17494        };
17495        let grid = nb(o0) + nb(o1);
17496        let mut y0 = self.alloc_uninit::<f32>(o0)?;
17497        let mut y1 = self.alloc_uninit::<f32>(o1)?;
17498        let f = self.func(if mr1 {
17499            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
17500        } else if rp {
17501            "qmatvec_q4_0_mmvq_fused2_rp"
17502        } else {
17503            "qmatvec_q4_0_mmvq_fused2"
17504        });
17505        let cfg = LaunchConfig {
17506            grid_dim: (grid, 1, 1),
17507            block_dim: (32, rpb, 1),
17508            shared_mem_bytes: 0,
17509        };
17510        let inf = w0.in_features() as i32;
17511        let (oo0, oo1) = (o0 as i32, o1 as i32);
17512        let (r0, r1) = (rb0 as i64, rb1 as i64);
17513        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
17514        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17515            {
17516                use cudarc::driver::{DevicePtr, DevicePtrMut};
17517                let s = &self.gpu.stream();
17518                let (p0, _g0) = b0.device_ptr(s);
17519                let (p1, _g1) = b1.device_ptr(s);
17520                let (paq, _g2) = aq.device_ptr(s);
17521                let (pad, _g3) = ad.device_ptr(s);
17522                let (py0, _g4) = y0.device_ptr_mut(s);
17523                let (py1, _g5) = y1.device_ptr_mut(s);
17524                let mut ps = [
17525                    &p0 as *const _ as *mut std::ffi::c_void,
17526                    &p1 as *const _ as *mut _,
17527                    &paq as *const _ as *mut _,
17528                    &pad as *const _ as *mut _,
17529                    &py0 as *const _ as *mut _,
17530                    &py1 as *const _ as *mut _,
17531                    &inf as *const _ as *mut _,
17532                    &oo0 as *const _ as *mut _,
17533                    &oo1 as *const _ as *mut _,
17534                    &r0 as *const _ as *mut _,
17535                    &r1 as *const _ as *mut _,
17536                ];
17537                unsafe {
17538                    self.launch_pdl(
17539                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
17540                        (grid, 1, 1),
17541                        (32, rpb, 1),
17542                        &mut ps,
17543                    )?;
17544                }
17545            }
17546            return Ok(Some((y0, y1)));
17547        }
17548        let __s_b = self.gpu.stream();
17549        let mut b = __s_b.launch_builder(&f);
17550        b.arg(b0)
17551            .arg(b1)
17552            .arg(aq)
17553            .arg(ad)
17554            .arg(&mut y0)
17555            .arg(&mut y1)
17556            .arg(&inf)
17557            .arg(&oo0)
17558            .arg(&oo1)
17559            .arg(&r0)
17560            .arg(&r1);
17561        unsafe {
17562            b.launch(cfg)?;
17563        }
17564        Ok(Some((y0, y1)))
17565    }
17566
17567    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
17568    pub fn matmul_q4_fused2_into(
17569        &self,
17570        w0: &crate::model::GpuTensor,
17571        w1: &crate::model::GpuTensor,
17572        aq: &CudaSlice<i8>,
17573        ad: &CudaSlice<f32>,
17574        y0: &mut CudaSlice<f32>,
17575        y1: &mut CudaSlice<f32>,
17576    ) -> Result<bool, Box<dyn std::error::Error>> {
17577        use crate::model::GpuTensor;
17578        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17579            match w {
17580                GpuTensor::Quant {
17581                    qtype, row_bytes, ..
17582                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17583                _ => None,
17584            }
17585        };
17586        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
17587            return Ok(false);
17588        };
17589        if w0.in_features() != w1.in_features() {
17590            return Ok(false);
17591        }
17592        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17593            match w {
17594                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17595                    Some(m) => (m, true),
17596                    None => (bytes, *rp),
17597                },
17598                _ => unreachable!(),
17599            }
17600        }
17601        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17602        if rp0 != rp1 {
17603            return Ok(false);
17604        }
17605        let rp = rp0;
17606        let rpb: u32 = 4;
17607        let mr1 = rp && Self::q40_mr1_on();
17608        let nb = |o: usize| {
17609            if mr1 {
17610                (o as u32).div_ceil(rpb)
17611            } else {
17612                (o as u32).div_ceil(2).div_ceil(rpb)
17613            }
17614        };
17615        let grid = nb(o0) + nb(o1);
17616        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
17617        let f = self.func(if mr1 {
17618            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
17619        } else if rp {
17620            "qmatvec_q4_0_mmvq_fused2_rp"
17621        } else {
17622            "qmatvec_q4_0_mmvq_fused2"
17623        });
17624        let cfg = LaunchConfig {
17625            grid_dim: (grid, 1, 1),
17626            block_dim: (32, rpb, 1),
17627            shared_mem_bytes: 0,
17628        };
17629        let inf = w0.in_features() as i32;
17630        let (oo0, oo1) = (o0 as i32, o1 as i32);
17631        let (r0, r1) = (rb0 as i64, rb1 as i64);
17632        // PDL wave-A: identical to the owned twin (capture-lane parity).
17633        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17634            use cudarc::driver::{DevicePtr, DevicePtrMut};
17635            let s = &self.gpu.stream();
17636            let (p0, _g0) = b0.device_ptr(s);
17637            let (p1, _g1) = b1.device_ptr(s);
17638            let (paq, _g2) = aq.device_ptr(s);
17639            let (pad, _g3) = ad.device_ptr(s);
17640            let (py0, _g4) = y0.device_ptr_mut(s);
17641            let (py1, _g5) = y1.device_ptr_mut(s);
17642            let mut ps = [
17643                &p0 as *const _ as *mut std::ffi::c_void,
17644                &p1 as *const _ as *mut _,
17645                &paq as *const _ as *mut _,
17646                &pad as *const _ as *mut _,
17647                &py0 as *const _ as *mut _,
17648                &py1 as *const _ as *mut _,
17649                &inf as *const _ as *mut _,
17650                &oo0 as *const _ as *mut _,
17651                &oo1 as *const _ as *mut _,
17652                &r0 as *const _ as *mut _,
17653                &r1 as *const _ as *mut _,
17654            ];
17655            unsafe {
17656                self.launch_pdl(
17657                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
17658                    (grid, 1, 1),
17659                    (32, rpb, 1),
17660                    &mut ps,
17661                )?;
17662            }
17663            return Ok(true);
17664        }
17665        let __s_b = self.gpu.stream();
17666        let mut b = __s_b.launch_builder(&f);
17667        b.arg(b0)
17668            .arg(b1)
17669            .arg(aq)
17670            .arg(ad)
17671            .arg(&mut *y0)
17672            .arg(&mut *y1)
17673            .arg(&inf)
17674            .arg(&oo0)
17675            .arg(&oo1)
17676            .arg(&r0)
17677            .arg(&r1);
17678        unsafe {
17679            b.launch(cfg)?;
17680        }
17681        Ok(true)
17682    }
17683
17684    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
17685    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
17686    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
17687    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
17688    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17689    pub fn matmul_q4_fused2_batched(
17690        &self,
17691        w0: &crate::model::GpuTensor,
17692        w1: &crate::model::GpuTensor,
17693        aq: &CudaSlice<i8>,
17694        ad: &CudaSlice<f32>,
17695        m: usize,
17696    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17697        use crate::model::GpuTensor;
17698        if !(2..=8).contains(&m) {
17699            return Ok(None);
17700        }
17701        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17702            match w {
17703                GpuTensor::Quant {
17704                    qtype, row_bytes, ..
17705                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17706                _ => None,
17707            }
17708        };
17709        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
17710            return Ok(None);
17711        };
17712        if w0.in_features() != w1.in_features() {
17713            return Ok(None);
17714        }
17715        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17716            match w {
17717                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17718                    Some(mr) => (mr, true),
17719                    None => (bytes, *rp),
17720                },
17721                _ => unreachable!(),
17722            }
17723        }
17724        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17725        if !rp0 || !rp1 {
17726            return Ok(None);
17727        }
17728        let mcols = Self::batched_mcols(m);
17729        let rpb: u32 = 4;
17730        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
17731        let grid = nb(o0) + nb(o1);
17732        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17733        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17734        let f = self.func(match mcols {
17735            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
17736            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
17737            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
17738        });
17739        let cfg = LaunchConfig {
17740            grid_dim: (grid, 1, 1),
17741            block_dim: (32, rpb, 1),
17742            shared_mem_bytes: 0,
17743        };
17744        let inf = w0.in_features() as i32;
17745        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
17746        let rb = rb0 as i64;
17747        let __s_b = self.gpu.stream();
17748        let mut b = __s_b.launch_builder(&f);
17749        b.arg(b0)
17750            .arg(b1)
17751            .arg(aq)
17752            .arg(ad)
17753            .arg(&mut y0)
17754            .arg(&mut y1)
17755            .arg(&inf)
17756            .arg(&oo0)
17757            .arg(&oo1)
17758            .arg(&mi)
17759            .arg(&rb);
17760        unsafe {
17761            b.launch(cfg)?;
17762        }
17763        Ok(Some((y0, y1)))
17764    }
17765
17766    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
17767    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
17768    #[allow(clippy::too_many_arguments)]
17769    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17770    pub fn matmul_q4_fused3_batched(
17771        &self,
17772        w0: &crate::model::GpuTensor,
17773        w1: &crate::model::GpuTensor,
17774        w2: &crate::model::GpuTensor,
17775        aq: &CudaSlice<i8>,
17776        ad: &CudaSlice<f32>,
17777        m: usize,
17778    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17779    {
17780        use crate::model::GpuTensor;
17781        if !(2..=8).contains(&m) {
17782            return Ok(None);
17783        }
17784        let q4 = |w: &GpuTensor| -> Option<usize> {
17785            match w {
17786                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
17787                _ => None,
17788            }
17789        };
17790        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
17791            return Ok(None);
17792        };
17793        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17794            return Ok(None);
17795        }
17796        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17797            match w {
17798                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17799                    Some(mr) => (mr, true),
17800                    None => (bytes, *rp),
17801                },
17802                _ => unreachable!(),
17803            }
17804        }
17805        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17806        if !rp0 || !rp1 || !rp2 {
17807            return Ok(None);
17808        }
17809        let mcols = Self::batched_mcols(m);
17810        let rpb: u32 = 4;
17811        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
17812        let grid = nb(o0) + nb(o1) + nb(o2);
17813        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17814        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17815        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
17816        let f = self.func(match mcols {
17817            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
17818            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
17819            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
17820        });
17821        let cfg = LaunchConfig {
17822            grid_dim: (grid, 1, 1),
17823            block_dim: (32, rpb, 1),
17824            shared_mem_bytes: 0,
17825        };
17826        let inf = w0.in_features() as i32;
17827        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
17828        let rb = 0i64;
17829        let __s_b = self.gpu.stream();
17830        let mut b = __s_b.launch_builder(&f);
17831        b.arg(b0)
17832            .arg(b1)
17833            .arg(b2)
17834            .arg(aq)
17835            .arg(ad)
17836            .arg(&mut y0)
17837            .arg(&mut y1)
17838            .arg(&mut y2)
17839            .arg(&inf)
17840            .arg(&oo0)
17841            .arg(&oo1)
17842            .arg(&oo2)
17843            .arg(&mi)
17844            .arg(&rb);
17845        unsafe {
17846            b.launch(cfg)?;
17847        }
17848        Ok(Some((y0, y1, y2)))
17849    }
17850
17851    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17852    pub fn matmul_q8_fused3(
17853        &self,
17854        w0: &crate::model::GpuTensor,
17855        w1: &crate::model::GpuTensor,
17856        w2: &crate::model::GpuTensor,
17857        aq: &CudaSlice<i8>,
17858        ad: &CudaSlice<f32>,
17859    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17860    {
17861        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
17862        // are per-tensor FP8, so native residency without this arm meant three separate launches.
17863        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
17864            return Ok(Some(self.e4m3_fused3_core(
17865                p0.0,
17866                p1.0,
17867                p2.0,
17868                aq,
17869                ad,
17870                w0.in_features(),
17871                p0.1,
17872                p1.1,
17873                p2.1,
17874                p0.2,
17875                p0.3,
17876                p1.3,
17877                p2.3,
17878            )?));
17879        }
17880        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
17881            return Ok(None);
17882        };
17883        Ok(Some(self.q8_fused3_core(
17884            p0.0,
17885            p1.0,
17886            p2.0,
17887            aq,
17888            ad,
17889            w0.in_features(),
17890            p0.1,
17891            p1.1,
17892            p2.1,
17893            p0.2,
17894        )?))
17895    }
17896
17897    #[allow(clippy::too_many_arguments)]
17898    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17899    fn q8_fused3_core(
17900        &self,
17901        b0: &CudaSlice<u8>,
17902        b1: &CudaSlice<u8>,
17903        b2: &CudaSlice<u8>,
17904        aq: &CudaSlice<i8>,
17905        ad: &CudaSlice<f32>,
17906        in_f: usize,
17907        out0: usize,
17908        out1: usize,
17909        out2: usize,
17910        row_bytes: usize,
17911    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17912        const ROWS_PER_BLOCK: u32 = 4;
17913        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
17914        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
17915        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
17916        let f = self.func("qmatvec_q8_0_mmvq_fused3");
17917        let mut y0 = self.alloc_uninit::<f32>(out0)?;
17918        let mut y1 = self.alloc_uninit::<f32>(out1)?;
17919        let mut y2 = self.alloc_uninit::<f32>(out2)?;
17920        let cfg = LaunchConfig {
17921            grid_dim: (nb0 + nb1 + nb2, 1, 1),
17922            block_dim: (32, ROWS_PER_BLOCK, 1),
17923            shared_mem_bytes: 0,
17924        };
17925        let (inf, o0, o1, o2, rbl) = (
17926            in_f as i32,
17927            out0 as i32,
17928            out1 as i32,
17929            out2 as i32,
17930            row_bytes as i64,
17931        );
17932        let __s_b = self.gpu.stream();
17933        let mut b = __s_b.launch_builder(&f);
17934        b.arg(b0)
17935            .arg(b1)
17936            .arg(b2)
17937            .arg(aq)
17938            .arg(ad)
17939            .arg(&mut y0)
17940            .arg(&mut y1)
17941            .arg(&mut y2)
17942            .arg(&inf)
17943            .arg(&o0)
17944            .arg(&o1)
17945            .arg(&o2)
17946            .arg(&rbl);
17947        unsafe {
17948            b.launch(cfg)?;
17949        }
17950        Ok((y0, y1, y2))
17951    }
17952
17953    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
17954    #[allow(clippy::too_many_arguments)]
17955    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17956    pub fn qmatvec_q8_fused3_raw(
17957        &self,
17958        b0: &CudaSlice<u8>,
17959        b1: &CudaSlice<u8>,
17960        b2: &CudaSlice<u8>,
17961        x: &CudaSlice<f32>,
17962        in_f: usize,
17963        out0: usize,
17964        out1: usize,
17965        out2: usize,
17966        row_bytes: usize,
17967    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17968        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
17969        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
17970    }
17971
17972    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
17973    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
17974    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
17975    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
17976    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
17977    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
17978    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
17979    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
17980    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
17981    /// twin must not introduce a batched program the reference path would not run).
17982    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17983    pub fn matmul_q8_fused2_t(
17984        &self,
17985        w0: &crate::model::GpuTensor,
17986        w1: &crate::model::GpuTensor,
17987        aq: &CudaSlice<i8>,
17988        ad: &CudaSlice<f32>,
17989        m: usize,
17990    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17991        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
17992        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
17993        // fuses too — same template body, still bit-identical to the two _b8 launches.
17994        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
17995            return Ok(None);
17996        }
17997        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
17998        // so the fused b8 launch would introduce a batched program the reference path would not run.
17999        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18000            if m > 4 && !Self::b8_enabled() {
18001                return Ok(None);
18002            }
18003            return Ok(Some(self.e4m3_fused2_t_core(
18004                p0.0,
18005                p1.0,
18006                aq,
18007                ad,
18008                m,
18009                w0.in_features(),
18010                p0.1,
18011                p1.1,
18012                p0.2,
18013                p0.3,
18014                p1.3,
18015            )?));
18016        }
18017        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18018            return Ok(None);
18019        };
18020        Ok(Some(self.q8_fused2_t_core(
18021            p0.0,
18022            p1.0,
18023            aq,
18024            ad,
18025            m,
18026            w0.in_features(),
18027            p0.1,
18028            p1.1,
18029            p0.2,
18030        )?))
18031    }
18032
18033    #[allow(clippy::too_many_arguments)]
18034    fn q8_fused2_t_core(
18035        &self,
18036        b0: &CudaSlice<u8>,
18037        b1: &CudaSlice<u8>,
18038        aq: &CudaSlice<i8>,
18039        ad: &CudaSlice<f32>,
18040        m: usize,
18041        in_f: usize,
18042        out0: usize,
18043        out1: usize,
18044        row_bytes: usize,
18045    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18046        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18047        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18048        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18049        let f = self.func(match Self::batched_mcols(m) {
18050            2 => "qmatvec_q8_0_mmvq_fused2_b2",
18051            4 => "qmatvec_q8_0_mmvq_fused2_b4",
18052            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
18053            _ => "qmatvec_q8_0_mmvq_fused2_b8",
18054        });
18055        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18056        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18057        let cfg = LaunchConfig {
18058            grid_dim: (nb0 + nb1, 1, 1),
18059            block_dim: (32, ROWS_PER_BLOCK, 1),
18060            shared_mem_bytes: 0,
18061        };
18062        let (inf, o0, o1, mi, rbl) = (
18063            in_f as i32,
18064            out0 as i32,
18065            out1 as i32,
18066            m as i32,
18067            row_bytes as i64,
18068        );
18069        let __s_b = self.gpu.stream();
18070        let mut b = __s_b.launch_builder(&f);
18071        b.arg(b0)
18072            .arg(b1)
18073            .arg(aq)
18074            .arg(ad)
18075            .arg(&mut y0)
18076            .arg(&mut y1)
18077            .arg(&inf)
18078            .arg(&o0)
18079            .arg(&o1)
18080            .arg(&mi)
18081            .arg(&rbl);
18082        unsafe {
18083            b.launch(cfg)?;
18084        }
18085        Ok((y0, y1))
18086    }
18087
18088    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
18089    /// q8_1 quant of the [m, in_f] activation), no env gating.
18090    #[allow(clippy::too_many_arguments)]
18091    pub fn qmatvec_q8_fused2_t_raw(
18092        &self,
18093        b0: &CudaSlice<u8>,
18094        b1: &CudaSlice<u8>,
18095        x: &CudaSlice<f32>,
18096        m: usize,
18097        in_f: usize,
18098        out0: usize,
18099        out1: usize,
18100        row_bytes: usize,
18101    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18102        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18103        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
18104    }
18105
18106    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
18107    /// `matmul_q8_fused2_t` with three ranges.
18108    #[allow(clippy::too_many_arguments)]
18109    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18110    pub fn matmul_q8_fused3_t(
18111        &self,
18112        w0: &crate::model::GpuTensor,
18113        w1: &crate::model::GpuTensor,
18114        w2: &crate::model::GpuTensor,
18115        aq: &CudaSlice<i8>,
18116        ad: &CudaSlice<f32>,
18117        m: usize,
18118    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18119    {
18120        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
18121            return Ok(None);
18122        }
18123        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
18124            return Ok(Some(self.e4m3_fused3_t_core(
18125                p0.0,
18126                p1.0,
18127                p2.0,
18128                aq,
18129                ad,
18130                m,
18131                w0.in_features(),
18132                p0.1,
18133                p1.1,
18134                p2.1,
18135                p0.2,
18136                p0.3,
18137                p1.3,
18138                p2.3,
18139            )?));
18140        }
18141        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
18142            return Ok(None);
18143        };
18144        Ok(Some(self.q8_fused3_t_core(
18145            p0.0,
18146            p1.0,
18147            p2.0,
18148            aq,
18149            ad,
18150            m,
18151            w0.in_features(),
18152            p0.1,
18153            p1.1,
18154            p2.1,
18155            p0.2,
18156        )?))
18157    }
18158
18159    #[allow(clippy::too_many_arguments)]
18160    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18161    fn q8_fused3_t_core(
18162        &self,
18163        b0: &CudaSlice<u8>,
18164        b1: &CudaSlice<u8>,
18165        b2: &CudaSlice<u8>,
18166        aq: &CudaSlice<i8>,
18167        ad: &CudaSlice<f32>,
18168        m: usize,
18169        in_f: usize,
18170        out0: usize,
18171        out1: usize,
18172        out2: usize,
18173        row_bytes: usize,
18174    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18175        const ROWS_PER_BLOCK: u32 = 4;
18176        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18177        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18178        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18179        let f = self.func(if Self::batched_mcols(m) == 2 {
18180            "qmatvec_q8_0_mmvq_fused3_b2"
18181        } else {
18182            "qmatvec_q8_0_mmvq_fused3_b4"
18183        });
18184        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18185        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18186        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
18187        let cfg = LaunchConfig {
18188            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18189            block_dim: (32, ROWS_PER_BLOCK, 1),
18190            shared_mem_bytes: 0,
18191        };
18192        let (inf, o0, o1, o2, mi, rbl) = (
18193            in_f as i32,
18194            out0 as i32,
18195            out1 as i32,
18196            out2 as i32,
18197            m as i32,
18198            row_bytes as i64,
18199        );
18200        let __s_b = self.gpu.stream();
18201        let mut b = __s_b.launch_builder(&f);
18202        b.arg(b0)
18203            .arg(b1)
18204            .arg(b2)
18205            .arg(aq)
18206            .arg(ad)
18207            .arg(&mut y0)
18208            .arg(&mut y1)
18209            .arg(&mut y2)
18210            .arg(&inf)
18211            .arg(&o0)
18212            .arg(&o1)
18213            .arg(&o2)
18214            .arg(&mi)
18215            .arg(&rbl);
18216        unsafe {
18217            b.launch(cfg)?;
18218        }
18219        Ok((y0, y1, y2))
18220    }
18221
18222    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
18223    #[allow(clippy::too_many_arguments)]
18224    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18225    pub fn qmatvec_q8_fused3_t_raw(
18226        &self,
18227        b0: &CudaSlice<u8>,
18228        b1: &CudaSlice<u8>,
18229        b2: &CudaSlice<u8>,
18230        x: &CudaSlice<f32>,
18231        m: usize,
18232        in_f: usize,
18233        out0: usize,
18234        out1: usize,
18235        out2: usize,
18236        row_bytes: usize,
18237    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18238        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18239        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
18240    }
18241
18242    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
18243    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
18244    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
18245    pub fn q8_ffn_fuse2_on(&self) -> bool {
18246        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18247        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
18248    }
18249
18250    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
18251    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
18252    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
18253    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
18254    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
18255    #[allow(clippy::type_complexity)]
18256    fn q8_fused_params<'w, const N: usize>(
18257        &self,
18258        ws: &[&'w crate::model::GpuTensor; N],
18259    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
18260        use crate::model::GpuTensor;
18261        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
18262            return None;
18263        }
18264        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
18265            return None;
18266        }
18267        let in_f = ws[0].in_features();
18268        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
18269        for (i, w) in ws.iter().enumerate() {
18270            match w {
18271                GpuTensor::Quant {
18272                    bytes,
18273                    qtype,
18274                    row_bytes,
18275                    scale,
18276                    ..
18277                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
18278                    out[i] = Some((bytes, w.out_features(), *row_bytes))
18279                }
18280                _ => return None,
18281            }
18282        }
18283        Some(out.map(|o| o.unwrap()))
18284    }
18285
18286    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
18287    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
18288    pub fn e4m3_dual_on(&self) -> bool {
18289        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18290        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
18291    }
18292
18293    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
18294    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
18295    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
18296    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
18297    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
18298    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
18299    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
18300    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
18301    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
18302    ///     Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
18303    ///     there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
18304    #[allow(clippy::type_complexity)]
18305    fn e4m3_fused_params<'w, const N: usize>(
18306        &self,
18307        ws: &[&'w crate::model::GpuTensor; N],
18308    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
18309        use crate::model::GpuTensor;
18310        if !self.e4m3_dual_on() {
18311            return None;
18312        }
18313        let in_f = ws[0].in_features();
18314        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
18315        for (i, w) in ws.iter().enumerate() {
18316            match w {
18317                GpuTensor::Quant {
18318                    bytes,
18319                    qtype,
18320                    row_bytes,
18321                    scale,
18322                    rp,
18323                    rp4,
18324                    ..
18325                } if *qtype == QT_F8_E4M3
18326                    && w.in_features() == in_f
18327                    && *row_bytes == in_f
18328                    && !*rp
18329                    && rp4.is_none() =>
18330                {
18331                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
18332                }
18333                _ => return None,
18334            }
18335        }
18336        Some(out.map(|o| o.unwrap()))
18337    }
18338
18339    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
18340    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
18341    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
18342    #[allow(clippy::too_many_arguments)]
18343    fn e4m3_fused2_core(
18344        &self,
18345        b0: &CudaSlice<u8>,
18346        b1: &CudaSlice<u8>,
18347        aq: &CudaSlice<i8>,
18348        ad: &CudaSlice<f32>,
18349        in_f: usize,
18350        out0: usize,
18351        out1: usize,
18352        row_bytes: usize,
18353        ws0: f32,
18354        ws1: f32,
18355    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18356        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18357        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18358        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18359        let f = self.func("qmatvec_e4m3_mmvq_fused2");
18360        let mut y0 = self.alloc_uninit::<f32>(out0)?;
18361        let mut y1 = self.alloc_uninit::<f32>(out1)?;
18362        let cfg = LaunchConfig {
18363            grid_dim: (nb0 + nb1, 1, 1),
18364            block_dim: (32, ROWS_PER_BLOCK, 1),
18365            shared_mem_bytes: 0,
18366        };
18367        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
18368        let __s_b = self.gpu.stream();
18369        let mut b = __s_b.launch_builder(&f);
18370        b.arg(b0)
18371            .arg(b1)
18372            .arg(aq)
18373            .arg(ad)
18374            .arg(&mut y0)
18375            .arg(&mut y1)
18376            .arg(&inf)
18377            .arg(&o0)
18378            .arg(&o1)
18379            .arg(&rbl)
18380            .arg(&ws0)
18381            .arg(&ws1);
18382        unsafe {
18383            b.launch(cfg)?;
18384        }
18385        Ok((y0, y1))
18386    }
18387
18388    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
18389    #[allow(clippy::too_many_arguments)]
18390    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18391    fn e4m3_fused3_core(
18392        &self,
18393        b0: &CudaSlice<u8>,
18394        b1: &CudaSlice<u8>,
18395        b2: &CudaSlice<u8>,
18396        aq: &CudaSlice<i8>,
18397        ad: &CudaSlice<f32>,
18398        in_f: usize,
18399        out0: usize,
18400        out1: usize,
18401        out2: usize,
18402        row_bytes: usize,
18403        ws0: f32,
18404        ws1: f32,
18405        ws2: f32,
18406    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18407        const ROWS_PER_BLOCK: u32 = 4;
18408        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18409        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18410        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18411        let f = self.func("qmatvec_e4m3_mmvq_fused3");
18412        let mut y0 = self.alloc_uninit::<f32>(out0)?;
18413        let mut y1 = self.alloc_uninit::<f32>(out1)?;
18414        let mut y2 = self.alloc_uninit::<f32>(out2)?;
18415        let cfg = LaunchConfig {
18416            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18417            block_dim: (32, ROWS_PER_BLOCK, 1),
18418            shared_mem_bytes: 0,
18419        };
18420        let (inf, o0, o1, o2, rbl) = (
18421            in_f as i32,
18422            out0 as i32,
18423            out1 as i32,
18424            out2 as i32,
18425            row_bytes as i64,
18426        );
18427        let __s_b = self.gpu.stream();
18428        let mut b = __s_b.launch_builder(&f);
18429        b.arg(b0)
18430            .arg(b1)
18431            .arg(b2)
18432            .arg(aq)
18433            .arg(ad)
18434            .arg(&mut y0)
18435            .arg(&mut y1)
18436            .arg(&mut y2)
18437            .arg(&inf)
18438            .arg(&o0)
18439            .arg(&o1)
18440            .arg(&o2)
18441            .arg(&rbl)
18442            .arg(&ws0)
18443            .arg(&ws1)
18444            .arg(&ws2);
18445        unsafe {
18446            b.launch(cfg)?;
18447        }
18448        Ok((y0, y1, y2))
18449    }
18450
18451    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
18452    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
18453    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
18454    #[allow(clippy::too_many_arguments)]
18455    fn e4m3_fused2_t_core(
18456        &self,
18457        b0: &CudaSlice<u8>,
18458        b1: &CudaSlice<u8>,
18459        aq: &CudaSlice<i8>,
18460        ad: &CudaSlice<f32>,
18461        m: usize,
18462        in_f: usize,
18463        out0: usize,
18464        out1: usize,
18465        row_bytes: usize,
18466        ws0: f32,
18467        ws1: f32,
18468    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18469        const ROWS_PER_BLOCK: u32 = 4;
18470        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18471        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18472        let f = self.func(match Self::batched_mcols(m) {
18473            2 => "qmatvec_e4m3_mmvq_fused2_b2",
18474            4 => "qmatvec_e4m3_mmvq_fused2_b4",
18475            _ => "qmatvec_e4m3_mmvq_fused2_b8",
18476        });
18477        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18478        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18479        let cfg = LaunchConfig {
18480            grid_dim: (nb0 + nb1, 1, 1),
18481            block_dim: (32, ROWS_PER_BLOCK, 1),
18482            shared_mem_bytes: 0,
18483        };
18484        let (inf, o0, o1, mi, rbl) = (
18485            in_f as i32,
18486            out0 as i32,
18487            out1 as i32,
18488            m as i32,
18489            row_bytes as i64,
18490        );
18491        let __s_b = self.gpu.stream();
18492        let mut b = __s_b.launch_builder(&f);
18493        b.arg(b0)
18494            .arg(b1)
18495            .arg(aq)
18496            .arg(ad)
18497            .arg(&mut y0)
18498            .arg(&mut y1)
18499            .arg(&inf)
18500            .arg(&o0)
18501            .arg(&o1)
18502            .arg(&mi)
18503            .arg(&rbl);
18504        unsafe {
18505            b.launch(cfg)?;
18506        }
18507        if ws0 != 1.0 {
18508            self.scale_inplace(&mut y0, ws0, m * out0)?;
18509        }
18510        if ws1 != 1.0 {
18511            self.scale_inplace(&mut y1, ws1, m * out1)?;
18512        }
18513        Ok((y0, y1))
18514    }
18515
18516    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
18517    #[allow(clippy::too_many_arguments)]
18518    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18519    fn e4m3_fused3_t_core(
18520        &self,
18521        b0: &CudaSlice<u8>,
18522        b1: &CudaSlice<u8>,
18523        b2: &CudaSlice<u8>,
18524        aq: &CudaSlice<i8>,
18525        ad: &CudaSlice<f32>,
18526        m: usize,
18527        in_f: usize,
18528        out0: usize,
18529        out1: usize,
18530        out2: usize,
18531        row_bytes: usize,
18532        ws0: f32,
18533        ws1: f32,
18534        ws2: f32,
18535    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18536        const ROWS_PER_BLOCK: u32 = 4;
18537        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18538        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18539        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18540        let f = self.func(if Self::batched_mcols(m) == 2 {
18541            "qmatvec_e4m3_mmvq_fused3_b2"
18542        } else {
18543            "qmatvec_e4m3_mmvq_fused3_b4"
18544        });
18545        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18546        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18547        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
18548        let cfg = LaunchConfig {
18549            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18550            block_dim: (32, ROWS_PER_BLOCK, 1),
18551            shared_mem_bytes: 0,
18552        };
18553        let (inf, o0, o1, o2, mi, rbl) = (
18554            in_f as i32,
18555            out0 as i32,
18556            out1 as i32,
18557            out2 as i32,
18558            m as i32,
18559            row_bytes as i64,
18560        );
18561        let __s_b = self.gpu.stream();
18562        let mut b = __s_b.launch_builder(&f);
18563        b.arg(b0)
18564            .arg(b1)
18565            .arg(b2)
18566            .arg(aq)
18567            .arg(ad)
18568            .arg(&mut y0)
18569            .arg(&mut y1)
18570            .arg(&mut y2)
18571            .arg(&inf)
18572            .arg(&o0)
18573            .arg(&o1)
18574            .arg(&o2)
18575            .arg(&mi)
18576            .arg(&rbl);
18577        unsafe {
18578            b.launch(cfg)?;
18579        }
18580        if ws0 != 1.0 {
18581            self.scale_inplace(&mut y0, ws0, m * out0)?;
18582        }
18583        if ws1 != 1.0 {
18584            self.scale_inplace(&mut y1, ws1, m * out1)?;
18585        }
18586        if ws2 != 1.0 {
18587            self.scale_inplace(&mut y2, ws2, m * out2)?;
18588        }
18589        Ok((y0, y1, y2))
18590    }
18591
18592    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
18593    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
18594    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
18595    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
18596    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
18597    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
18598    ///
18599    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
18600    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
18601    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
18602    pub fn qmatvec_e4m3_blk_mmvq(
18603        &self,
18604        bytes: &CudaSlice<u8>,
18605        aq: &CudaSlice<i8>,
18606        ad: &CudaSlice<f32>,
18607        scales: &CudaSlice<f32>,
18608        m: usize,
18609        in_f: usize,
18610        out_f: usize,
18611        row_bytes: usize,
18612        scale_cols: usize,
18613    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18614        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
18615        self.qmatvec_e4m3_blk_mmvq_into(
18616            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
18617        )?;
18618        Ok(y)
18619    }
18620
18621    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
18622    #[allow(clippy::too_many_arguments)]
18623    pub fn qmatvec_e4m3_blk_mmvq_into(
18624        &self,
18625        bytes: &CudaSlice<u8>,
18626        aq: &CudaSlice<i8>,
18627        ad: &CudaSlice<f32>,
18628        scales: &CudaSlice<f32>,
18629        m: usize,
18630        in_f: usize,
18631        out_f: usize,
18632        row_bytes: usize,
18633        scale_cols: usize,
18634        y: &mut CudaSlice<f32>,
18635    ) -> Result<(), Box<dyn std::error::Error>> {
18636        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18637        let f = self.func("qmatvec_e4m3_blk_mmvq");
18638        let cfg = LaunchConfig {
18639            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
18640            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
18641            shared_mem_bytes: 0,                // warp-only reduce
18642        };
18643        let (inf, outf, mi, rb, sc) = (
18644            in_f as i32,
18645            out_f as i32,
18646            m as i32,
18647            row_bytes as i64,
18648            scale_cols as i32,
18649        );
18650        let __s_b = self.gpu.stream();
18651        let mut b = __s_b.launch_builder(&f);
18652        b.arg(bytes)
18653            .arg(aq)
18654            .arg(ad)
18655            .arg(scales)
18656            .arg(&mut *y)
18657            .arg(&inf)
18658            .arg(&outf)
18659            .arg(&mi)
18660            .arg(&rb)
18661            .arg(&sc);
18662        unsafe {
18663            b.launch(cfg)?;
18664        }
18665        Ok(())
18666    }
18667
18668    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
18669    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
18670    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
18671    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
18672    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
18673    #[allow(clippy::too_many_arguments)]
18674    pub fn qmatvec_e4m3_blk_mmvq_batched(
18675        &self,
18676        bytes: &CudaSlice<u8>,
18677        aq: &CudaSlice<i8>,
18678        ad: &CudaSlice<f32>,
18679        scales: &CudaSlice<f32>,
18680        m: usize,
18681        in_f: usize,
18682        out_f: usize,
18683        row_bytes: usize,
18684        scale_cols: usize,
18685        mcols: usize,
18686    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18687        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18688        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
18689        let name = match mcols {
18690            2 => "qmatvec_e4m3_blk_mmvq_b2",
18691            4 => "qmatvec_e4m3_blk_mmvq_b4",
18692            8 => "qmatvec_e4m3_blk_mmvq_b8",
18693            16 => "qmatvec_e4m3_blk_mmvq_b16",
18694            _ => {
18695                return Err(
18696                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
18697                );
18698            }
18699        };
18700        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18701        let f = self.func(name);
18702        let cfg = LaunchConfig {
18703            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18704            block_dim: (32, ROWS_PER_BLOCK, 1),
18705            shared_mem_bytes: 0,
18706        };
18707        let (inf, outf, mi, rb, sc) = (
18708            in_f as i32,
18709            out_f as i32,
18710            m as i32,
18711            row_bytes as i64,
18712            scale_cols as i32,
18713        );
18714        let __s_b = self.gpu.stream();
18715        let mut b = __s_b.launch_builder(&f);
18716        b.arg(bytes)
18717            .arg(aq)
18718            .arg(ad)
18719            .arg(scales)
18720            .arg(&mut y)
18721            .arg(&inf)
18722            .arg(&outf)
18723            .arg(&mi)
18724            .arg(&rb)
18725            .arg(&sc);
18726        unsafe {
18727            b.launch(cfg)?;
18728        }
18729        Ok(y)
18730    }
18731
18732    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
18733    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
18734    #[allow(clippy::too_many_arguments)]
18735    pub fn qmatvec_e4m3_blk_batched_raw(
18736        &self,
18737        bytes: &CudaSlice<u8>,
18738        x: &CudaSlice<f32>,
18739        scales: &CudaSlice<f32>,
18740        m: usize,
18741        in_f: usize,
18742        out_f: usize,
18743        row_bytes: usize,
18744        scale_cols: usize,
18745        mcols: usize,
18746    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18747        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18748        self.qmatvec_e4m3_blk_mmvq_batched(
18749            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
18750        )
18751    }
18752
18753    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
18754    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
18755    #[allow(clippy::too_many_arguments)]
18756    pub fn qmatvec_e4m3_blk_mmvq_raw(
18757        &self,
18758        bytes: &CudaSlice<u8>,
18759        x: &CudaSlice<f32>,
18760        scales: &CudaSlice<f32>,
18761        m: usize,
18762        in_f: usize,
18763        out_f: usize,
18764        row_bytes: usize,
18765        scale_cols: usize,
18766    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18767        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18768        self.qmatvec_e4m3_blk_mmvq(
18769            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
18770        )
18771    }
18772
18773    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
18774    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
18775    #[allow(clippy::too_many_arguments)]
18776    pub fn qmatvec_e4m3_fused2_raw(
18777        &self,
18778        b0: &CudaSlice<u8>,
18779        b1: &CudaSlice<u8>,
18780        x: &CudaSlice<f32>,
18781        in_f: usize,
18782        out0: usize,
18783        out1: usize,
18784        row_bytes: usize,
18785        ws0: f32,
18786        ws1: f32,
18787    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18788        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18789        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
18790    }
18791
18792    #[allow(clippy::too_many_arguments)]
18793    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18794    pub fn qmatvec_e4m3_fused3_raw(
18795        &self,
18796        b0: &CudaSlice<u8>,
18797        b1: &CudaSlice<u8>,
18798        b2: &CudaSlice<u8>,
18799        x: &CudaSlice<f32>,
18800        in_f: usize,
18801        out0: usize,
18802        out1: usize,
18803        out2: usize,
18804        row_bytes: usize,
18805        ws0: f32,
18806        ws1: f32,
18807        ws2: f32,
18808    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18809        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18810        self.e4m3_fused3_core(
18811            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
18812        )
18813    }
18814
18815    #[allow(clippy::too_many_arguments)]
18816    pub fn qmatvec_e4m3_fused2_t_raw(
18817        &self,
18818        b0: &CudaSlice<u8>,
18819        b1: &CudaSlice<u8>,
18820        x: &CudaSlice<f32>,
18821        m: usize,
18822        in_f: usize,
18823        out0: usize,
18824        out1: usize,
18825        row_bytes: usize,
18826        ws0: f32,
18827        ws1: f32,
18828    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18829        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18830        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
18831    }
18832
18833    #[allow(clippy::too_many_arguments)]
18834    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18835    pub fn qmatvec_e4m3_fused3_t_raw(
18836        &self,
18837        b0: &CudaSlice<u8>,
18838        b1: &CudaSlice<u8>,
18839        b2: &CudaSlice<u8>,
18840        x: &CudaSlice<f32>,
18841        m: usize,
18842        in_f: usize,
18843        out0: usize,
18844        out1: usize,
18845        out2: usize,
18846        row_bytes: usize,
18847        ws0: f32,
18848        ws1: f32,
18849        ws2: f32,
18850    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18851        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18852        self.e4m3_fused3_t_core(
18853            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
18854        )
18855    }
18856
18857    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
18858    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
18859    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
18860    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
18861    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
18862    ///
18863    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
18864    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
18865    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
18866    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
18867    fn try_e4m3_blk_pre(
18868        &self,
18869        w: &crate::model::GpuTensor,
18870        aq: &CudaSlice<i8>,
18871        ad: &CudaSlice<f32>,
18872        m: usize,
18873    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
18874        use crate::model::GpuTensor;
18875        if let GpuTensor::Quant {
18876            bytes,
18877            qtype,
18878            row_bytes,
18879            blk: Some(g),
18880            ..
18881        } = w
18882            && *qtype == QT_F8_E4M3_BLK
18883        {
18884            // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
18885            // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
18886            // below, so the decode-exactness contract is preserved at every width. Gated by
18887            // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
18888            // one rollback door covers every dtype's batched tier.
18889            if (2..=16).contains(&m)
18890                && std::env::var("MEMRA_NO_BATCHED").is_err()
18891                && (m <= 4 || Self::b8_enabled())
18892            {
18893                let mcols = Self::batched_mcols(m);
18894                return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
18895                    bytes,
18896                    aq,
18897                    ad,
18898                    &g.scales,
18899                    m,
18900                    w.in_features(),
18901                    w.out_features(),
18902                    *row_bytes,
18903                    g.cols,
18904                    mcols,
18905                )?));
18906            }
18907            return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
18908                bytes,
18909                aq,
18910                ad,
18911                &g.scales,
18912                m,
18913                w.in_features(),
18914                w.out_features(),
18915                *row_bytes,
18916                g.cols,
18917            )?));
18918        }
18919        Ok(None)
18920    }
18921
18922    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
18923    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
18924    ///
18925    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
18926    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
18927    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
18928    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
18929    /// prefill keeps the floor's arithmetic and the floor's kernels.
18930    ///
18931    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
18932    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
18933    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
18934    /// (projection, prefill call) and frees immediately.
18935    ///
18936    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
18937    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
18938    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
18939    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
18940    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
18941    /// single-variable comparison instead of a two-variable one.
18942    ///
18943    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
18944    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
18945    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
18946    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
18947    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
18948    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
18949    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
18950    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
18951    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
18952    ///
18953    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
18954    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
18955    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
18956    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
18957    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
18958    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
18959    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
18960    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
18961    /// because v2's denominator had its slab already resident while this class's floor must build it
18962    /// every call; same tile, opposite sign, because the question changed.
18963    ///
18964    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
18965    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
18966    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
18967    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
18968    fn try_e4m3_blk_prefill(
18969        &self,
18970        w: &crate::model::GpuTensor,
18971        x: &CudaSlice<f32>,
18972        m: usize,
18973    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
18974        use crate::model::GpuTensor;
18975        let GpuTensor::Quant {
18976            bytes,
18977            qtype,
18978            blk: Some(g),
18979            ..
18980        } = w
18981        else {
18982            return Ok(None);
18983        };
18984        if *qtype != QT_F8_E4M3_BLK {
18985            return Ok(None);
18986        }
18987        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
18988        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
18989        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
18990        // through to the dequant below when they do, never silently produce nothing.
18991        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
18992            return Ok(Some(y));
18993        }
18994        let (in_f, out_f) = (w.in_features(), w.out_features());
18995        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
18996        let tmp = GpuTensor::Quant {
18997            bytes: slab,
18998            qtype: QT_Q8_0,
18999            row_bytes: in_f / 32 * 34,
19000            ne: vec![in_f as u64, out_f as u64],
19001            scale: 1.0,
19002            rp: false,
19003            #[cfg(memra_cutlass)]
19004            cutlass: None,
19005            fp8: None,
19006            blk: None,
19007            f16: None,
19008            rp4: None,
19009        };
19010        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
19011        Ok(Some(self.matmul(&tmp, x, m)?))
19012    }
19013
19014    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19015    pub fn matmul_pre_noscale(
19016        &self,
19017        w: &crate::model::GpuTensor,
19018        aq: &CudaSlice<i8>,
19019        ad: &CudaSlice<f32>,
19020        m: usize,
19021    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
19022        use crate::model::GpuTensor;
19023        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
19024        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
19025        // rather than let the tail below refuse and cost the caller a re-dispatch.
19026        if m == 1
19027            && let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)?
19028        {
19029            return Ok(Some((y, 1.0)));
19030        }
19031        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
19032        if m != 1 || !self.uses_q8_1_fast(w) {
19033            return Ok(None);
19034        }
19035        let in_f = w.in_features();
19036        let out_f = w.out_features();
19037        let (bytes, qtype, row_bytes, scale, rp) = match w {
19038            GpuTensor::Quant {
19039                bytes,
19040                qtype,
19041                row_bytes,
19042                scale,
19043                rp,
19044                ..
19045            } => (bytes, *qtype, *row_bytes, *scale, *rp),
19046            _ => return Ok(None),
19047        };
19048        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
19049        if self.mmvq_supports(qtype) {
19050            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
19051            let (mbytes, mrp) = match w {
19052                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
19053                _ => (bytes, rp),
19054            };
19055            let y = self.qmatvec_mmvq(
19056                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
19057            )?;
19058            return Ok(Some((y, scale)));
19059        }
19060        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
19061        let name = match qtype {
19062            QT_Q8_0 => "qmatvec_q8_0_dp4a",
19063            QT_Q4_K => "qmatvec_q4_K_dp4a",
19064            QT_Q6_K => "qmatvec_q6_K_dp4a",
19065            QT_Q5_K => "qmatvec_q5_K_dp4a",
19066            QT_Q3_K => "qmatvec_q3_K_dp4a",
19067            QT_NVFP4 => {
19068                if rp {
19069                    "qmatvec_nvfp4_dp4a_rp"
19070                } else {
19071                    "qmatvec_nvfp4_dp4a"
19072                }
19073            }
19074            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
19075            _ => return Ok(None),
19076        };
19077        let f = self.func(name);
19078        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19079        let cfg = LaunchConfig {
19080            grid_dim: (out_f as u32, m as u32, 1),
19081            block_dim: (128, 1, 1),
19082            shared_mem_bytes: 0,
19083        };
19084        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19085        let __s_b = self.gpu.stream();
19086        let mut b = __s_b.launch_builder(&f);
19087        b.arg(bytes)
19088            .arg(aq)
19089            .arg(ad)
19090            .arg(&mut y)
19091            .arg(&inf)
19092            .arg(&outf)
19093            .arg(&mi)
19094            .arg(&rb);
19095        unsafe {
19096            b.launch(cfg)?;
19097        }
19098        Ok(Some((y, scale)))
19099    }
19100
19101    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
19102    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
19103    pub fn mmvq_supports(&self, qtype: i32) -> bool {
19104        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
19105        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
19106        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
19107        // is a pure function of the dtype — the decode-parity law holds under every env.
19108        if qtype == QT_F8_E4M3 {
19109            return true;
19110        }
19111        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
19112            return false;
19113        }
19114        matches!(
19115            qtype,
19116            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
19117        )
19118    }
19119
19120    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
19121    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
19122    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
19123    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
19124    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19125    pub fn qmatvec_mmvq(
19126        &self,
19127        bytes: &CudaSlice<u8>,
19128        aq: &CudaSlice<i8>,
19129        ad: &CudaSlice<f32>,
19130        m: usize,
19131        in_f: usize,
19132        out_f: usize,
19133        qtype: i32,
19134        row_bytes: usize,
19135        scale: f32,
19136        rp: bool,
19137    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19138        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
19139        self.qmatvec_mmvq_into(
19140            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
19141        )?;
19142        Ok(y)
19143    }
19144
19145    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
19146    #[allow(clippy::too_many_arguments)]
19147    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19148    pub fn qmatvec_mmvq_into(
19149        &self,
19150        bytes: &CudaSlice<u8>,
19151        aq: &CudaSlice<i8>,
19152        ad: &CudaSlice<f32>,
19153        m: usize,
19154        in_f: usize,
19155        out_f: usize,
19156        qtype: i32,
19157        row_bytes: usize,
19158        scale: f32,
19159        rp: bool,
19160        y: &mut CudaSlice<f32>,
19161    ) -> Result<(), Box<dyn std::error::Error>> {
19162        debug_assert!(y.len() >= m * out_f);
19163        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
19164        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
19165        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
19166        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
19167        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
19168        if qtype == QT_Q8_0
19169            && rp
19170            && m == 1
19171            && out_f >= 64
19172            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
19173            && {
19174                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19175                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
19176            }
19177        {
19178            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
19179            let cfg = LaunchConfig {
19180                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
19181                block_dim: (32, 2, 1),
19182                shared_mem_bytes: 0,
19183            };
19184            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
19185            let __s_b = self.gpu.stream();
19186            let mut b = __s_b.launch_builder(&f);
19187            b.arg(bytes)
19188                .arg(aq)
19189                .arg(ad)
19190                .arg(&mut *y)
19191                .arg(&inf)
19192                .arg(&outf)
19193                .arg(&mi)
19194                .arg(&rb);
19195            unsafe {
19196                b.launch(cfg)?;
19197            }
19198            if scale != 1.0 {
19199                self.scale_inplace(y, scale, out_f)?;
19200            }
19201            return Ok(());
19202        }
19203        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
19204        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
19205        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
19206        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
19207        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
19208        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
19209        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
19210        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
19211        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
19212            2
19213        } else {
19214            1
19215        };
19216        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
19217        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
19218        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
19219        // valid-window interleaved, bit-identical per row — same dot program).
19220        if m == 1 && qtype == QT_Q4_0 {
19221            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
19222            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
19223            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
19224            mr = *Q40MR.get_or_init(|| {
19225                std::env::var("MEMRA_Q40_MR")
19226                    .ok()
19227                    .and_then(|v| v.parse().ok())
19228                    .unwrap_or(1)
19229            });
19230        }
19231        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
19232        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
19233        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
19234        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
19235        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
19236        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
19237        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
19238        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
19239        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
19240        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
19241        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
19242        let q5_force = q5_mode.as_deref() == Some("2");
19243        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
19244        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
19245        let q5_il = qtype == QT_Q5_K
19246            && m == 1
19247            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
19248        if q5_il && !q5_force && out_f > 65536 {
19249            mr = 1;
19250        }
19251        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
19252        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
19253        if qtype == QT_Q4_0 && rp && mr != 1 {
19254            mr = 2;
19255        }
19256        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
19257        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
19258        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
19259        if qtype == QT_Q8_0 && rp {
19260            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
19261            mr = *Q80MR.get_or_init(|| {
19262                std::env::var("MEMRA_Q80_MR")
19263                    .ok()
19264                    .and_then(|v| v.parse().ok())
19265                    .unwrap_or(1)
19266            });
19267        }
19268        let name = match (qtype, mr, rp) {
19269            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
19270            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
19271            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
19272            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
19273            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
19274            (QT_Q5_K, 2, _) => {
19275                if q5_il {
19276                    "qmatvec_q5_K_mmvq_mr2_il"
19277                } else {
19278                    "qmatvec_q5_K_mmvq_mr2"
19279                }
19280            }
19281            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
19282            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
19283            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
19284            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
19285            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
19286            (QT_Q8_0, _, true)
19287                if in_f.is_multiple_of(1024) && {
19288                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19289                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
19290                } =>
19291            {
19292                "qmatvec_q8_0_mmvq_rpca"
19293            }
19294            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
19295            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
19296            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
19297            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
19298            // reach a GGUF-layout kernel or vice versa.
19299            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
19300            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
19301            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
19302            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
19303            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
19304            (QT_Q5_K, _, _) => {
19305                if q5_il {
19306                    "qmatvec_q5_K_mmvq_il"
19307                } else {
19308                    "qmatvec_q5_K_mmvq"
19309                }
19310            }
19311            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
19312            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
19313            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
19314            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
19315        };
19316        let f = self.func(name);
19317        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
19318        let rows_per_block = ROWS_PER_BLOCK * mr;
19319        let cfg = LaunchConfig {
19320            grid_dim: (
19321                (out_f as u32 + rows_per_block - 1) / rows_per_block,
19322                m as u32,
19323                1,
19324            ),
19325            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
19326            shared_mem_bytes: 0,                // warp-only reduce at m=1
19327        };
19328        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19329        let __s_b = self.gpu.stream();
19330        let mut b = __s_b.launch_builder(&f);
19331        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
19332        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
19333        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
19334        // weight_scale). Other mmvq kernels keep the 8-arg signature.
19335        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
19336            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
19337            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
19338            if Self::pdl_on()
19339                && Self::pdl_mmvq_on()
19340                && Self::pdl_nvfp4q8_on()
19341                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
19342            {
19343                use cudarc::driver::{DevicePtr, DevicePtrMut};
19344                let s = &self.gpu.stream();
19345                let (pw, _g0) = bytes.device_ptr(s);
19346                let (paq, _g1) = aq.device_ptr(s);
19347                let (pad, _g2) = ad.device_ptr(s);
19348                let (py, _g3) = y.device_ptr_mut(s);
19349                let mut ps = [
19350                    &pw as *const _ as *mut std::ffi::c_void,
19351                    &paq as *const _ as *mut _,
19352                    &pad as *const _ as *mut _,
19353                    &py as *const _ as *mut _,
19354                    &inf as *const _ as *mut _,
19355                    &outf as *const _ as *mut _,
19356                    &mi as *const _ as *mut _,
19357                    &rb as *const _ as *mut _,
19358                    &scale as *const _ as *mut _,
19359                ];
19360                unsafe {
19361                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
19362                }
19363                return Ok(());
19364            }
19365            b.arg(bytes)
19366                .arg(aq)
19367                .arg(ad)
19368                .arg(&mut *y)
19369                .arg(&inf)
19370                .arg(&outf)
19371                .arg(&mi)
19372                .arg(&rb)
19373                .arg(&scale);
19374            unsafe {
19375                b.launch(cfg)?;
19376            }
19377        } else if Self::pdl_on()
19378            && Self::pdl_mmvq_on()
19379            && (matches!(
19380                name,
19381                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
19382            ) || (Self::pdl_nvfp4q8_on()
19383                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
19384        {
19385            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
19386            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
19387            // names may take this launch (unmarked kernels would read unordered).
19388            {
19389                use cudarc::driver::{DevicePtr, DevicePtrMut};
19390                let s = &self.gpu.stream();
19391                let (pw, _g0) = bytes.device_ptr(s);
19392                let (paq, _g1) = aq.device_ptr(s);
19393                let (pad, _g2) = ad.device_ptr(s);
19394                let (py, _g3) = y.device_ptr_mut(s);
19395                let mut ps = [
19396                    &pw as *const _ as *mut std::ffi::c_void,
19397                    &paq as *const _ as *mut _,
19398                    &pad as *const _ as *mut _,
19399                    &py as *const _ as *mut _,
19400                    &inf as *const _ as *mut _,
19401                    &outf as *const _ as *mut _,
19402                    &mi as *const _ as *mut _,
19403                    &rb as *const _ as *mut _,
19404                ];
19405                unsafe {
19406                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
19407                }
19408            }
19409            if scale != 1.0 {
19410                self.scale_inplace(y, scale, m * out_f)?;
19411            }
19412        } else {
19413            b.arg(bytes)
19414                .arg(aq)
19415                .arg(ad)
19416                .arg(&mut *y)
19417                .arg(&inf)
19418                .arg(&outf)
19419                .arg(&mi)
19420                .arg(&rb);
19421            unsafe {
19422                b.launch(cfg)?;
19423            }
19424            if scale != 1.0 {
19425                self.scale_inplace(y, scale, m * out_f)?;
19426            }
19427        }
19428        Ok(())
19429    }
19430
19431    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
19432    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
19433    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
19434    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19435    pub fn qmatvec_mmvq_raw(
19436        &self,
19437        bytes: &CudaSlice<u8>,
19438        x: &CudaSlice<f32>,
19439        m: usize,
19440        in_f: usize,
19441        out_f: usize,
19442        qtype: i32,
19443        row_bytes: usize,
19444        rp: bool,
19445    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19446        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19447        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
19448    }
19449
19450    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
19451    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
19452    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
19453    pub fn batched_supports(&self, qtype: i32) -> bool {
19454        matches!(
19455            qtype,
19456            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
19457        )
19458    }
19459
19460    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
19461    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
19462    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
19463    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
19464    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
19465    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
19466    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
19467    pub fn iq_fast_enabled() -> bool {
19468        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19469        *ON.get_or_init(|| {
19470            std::env::var("MEMRA_IQ_FAST")
19471                .map(|v| v != "0")
19472                .unwrap_or(true)
19473        })
19474    }
19475
19476    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
19477    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
19478    pub fn b8_enabled() -> bool {
19479        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19480        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
19481    }
19482
19483    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
19484    pub fn batched_mcols(m: usize) -> usize {
19485        if m == 2 {
19486            2
19487        } else if m <= 4 {
19488            4
19489        } else if m <= 8 {
19490            8
19491        } else {
19492            16
19493        }
19494    }
19495
19496    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
19497    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
19498    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
19499    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
19500    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
19501        Some(match (qtype, mcols) {
19502            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
19503            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
19504            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
19505            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
19506            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
19507            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
19508            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
19509            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
19510            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
19511            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
19512            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
19513            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
19514            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
19515            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
19516            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
19517            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
19518            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
19519            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
19520            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
19521            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
19522            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
19523            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
19524            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
19525            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
19526            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
19527            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
19528            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
19529            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
19530            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
19531            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
19532            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
19533            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
19534            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
19535            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
19536            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
19537            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
19538            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
19539            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
19540            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
19541            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
19542            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
19543            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
19544            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
19545            _ => return None,
19546        })
19547    }
19548
19549    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
19550    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
19551    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
19552    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
19553    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
19554    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
19555    ///
19556    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
19557    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
19558    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
19559    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
19560    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
19561    /// msweep on all six 27B shapes (2026-07-03):
19562    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
19563    ///          it applies for b4 (-3..-14%), never loses;
19564    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
19565    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
19566    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
19567    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
19568    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
19569    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
19570    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
19571    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
19572    /// b2: in_f>=6144 -> r2, else base.
19573    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
19574    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
19575    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
19576    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
19577    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
19578    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
19579    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
19580    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
19581    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
19582    /// Device SM count (cached) — grid-fill policy input.
19583    pub fn sm_count(&self) -> i32 {
19584        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
19585        *SMS.get_or_init(|| {
19586            use cudarc::driver::sys::CUdevice_attribute_enum as A;
19587            self.gpu
19588                .ctx
19589                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
19590                .unwrap_or(82)
19591        })
19592    }
19593
19594    #[allow(clippy::too_many_arguments)]
19595    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19596    #[allow(clippy::if_same_then_else)] // allow: a fallback mapping table; distinct inputs deliberately share a target arm
19597    pub fn batched_variant(
19598        &self,
19599        _m: usize,
19600        in_f: usize,
19601        out_f: usize,
19602        qtype: i32,
19603        row_bytes: usize,
19604        mcols: usize,
19605        rp: bool,
19606    ) -> &'static str {
19607        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
19608        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
19609        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
19610        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
19611        if qtype == QT_Q8_0 {
19612            return if rp { "rp" } else { "base" };
19613        }
19614        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19615        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
19616            Ok("base") => "base",
19617            Ok("pf") => "pf",
19618            Ok("r2") => "r2",
19619            Ok("r2w8") => "r2w8",
19620            Ok("pfr2") => "pfr2",
19621            Ok("ca") => "ca",
19622            Ok("car2") => "car2",
19623            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
19624            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
19625            Ok("rp") => "rp",
19626            Ok("rpr2") => "rpr2",
19627            Ok("rpr2w8") => "rpr2w8",
19628            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
19629            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
19630            Ok("rpca") => "rpca",
19631            Ok("rpcar2") => "rpcar2",
19632            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
19633            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
19634            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
19635            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
19636            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
19637            // bit-identical to the decode path — measurement corpus ONLY, never auto).
19638            Ok("rpsc") => "rpsc",
19639            Ok("rpms") => "rpms",
19640            Ok("rpmsc") => "rpmsc",
19641            Ok("rpks") => "rpks",
19642            Ok("rpksc") => "rpksc",
19643            _ => "auto",
19644        });
19645        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
19646        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
19647        // shapes qualify; anything else falls back to the register variants.
19648        let ca_ok = qtype == QT_NVFP4 && row_bytes.is_multiple_of(16) && in_f.is_multiple_of(1024);
19649        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
19650        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
19651        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
19652        // forced MEMRA_MMVQ_BV values still work).
19653        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19654        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
19655        let sc_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(256) && (in_f / 64 <= 272);
19656        let ks_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(512) && (in_f / 64 <= 272);
19657        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
19658        let sms = *SMS.get_or_init(|| {
19659            use cudarc::driver::sys::CUdevice_attribute_enum as A;
19660            self.gpu
19661                .ctx
19662                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
19663                .unwrap_or(82)
19664        });
19665        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
19666        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
19667        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
19668        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
19669        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
19670        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
19671        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
19672        // AUTO RULE = the measured winners table (differs from NVFP4's!):
19673        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
19674        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
19675        //     r2 1258us) — kernels kept behind the force seam for the corpus;
19676        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
19677        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
19678        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
19679        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
19680        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
19681        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
19682        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
19683        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
19684        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
19685        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
19686        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
19687        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19688        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
19689            Ok("base") => "base",
19690            Ok("r2") => "r2",
19691            Ok("r2w8") => "r2w8",
19692            _ => "auto",
19693        });
19694        let variant: &'static str = if qtype == QT_Q4_0 {
19695            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
19696            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
19697            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
19698            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19699            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
19700                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
19701                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
19702                // + syncs cost more than the stalls, bank-pad made no difference);
19703                // register load-ahead flat (nvcc already reorders). The b-tier limiter
19704                // is still unidentified — see the jsonl row.
19705                Ok("base") => "base",
19706                Ok("r2") => "r2",
19707                Ok("ms") => "ms",
19708                Ok("sm") => "sm",
19709                Ok("la") => "la",
19710                _ => "auto",
19711            });
19712            let v = if q40 != "auto" {
19713                q40
19714            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
19715                "r2"
19716            } else {
19717                "base"
19718            };
19719            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
19720            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
19721            // and the limiter is the per-column activation load chain (long_scoreboard
19722            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
19723            if rp {
19724                match v {
19725                    "ms" => "r2ms_rp",
19726                    "sm" => "r2sm_rp",
19727                    "la" => "r2la_rp",
19728                    "r2" => "r2_rp",
19729                    _ => "rp",
19730                }
19731            } else if matches!(v, "ms" | "sm" | "la") {
19732                "r2"
19733            } else {
19734                v
19735            }
19736        } else if qtype != QT_NVFP4 && !kq_r2 {
19737            "base"
19738        } else if kq_r2 && rp {
19739            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
19740            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
19741            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
19742            "rp"
19743        } else if kq_r2 {
19744            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
19745            // mcols != 4 forced r2w8 falls to unbounded r2.
19746            if kq_bv != "auto" {
19747                if kq_bv == "r2w8" && mcols != 4 {
19748                    "r2"
19749                } else {
19750                    kq_bv
19751                }
19752            } else if bv != "auto" {
19753                match bv {
19754                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
19755                    "r2w8" | "rpr2w8" => {
19756                        if mcols != 4 {
19757                            "r2"
19758                        } else {
19759                            "r2w8"
19760                        }
19761                    }
19762                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
19763                }
19764            } else {
19765                #[allow(clippy::manual_div_ceil)]
19766                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19767                let blocks = (out_f + 7) / 8;
19768                let waves = blocks as f64 / (7 * sms as usize) as f64;
19769                let filled = blocks >= 4 * sms as usize;
19770                let use_r2 = if qtype == QT_Q4_K {
19771                    filled
19772                } else {
19773                    waves >= 2.0
19774                };
19775                if use_r2 { "r2" } else { "base" }
19776            }
19777        } else if bv != "auto" {
19778            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
19779            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
19780            // unsupported (shape, mcols) combos fall back to pf/r2.
19781            // On rp buffers, forced legacy names map to their rp twins (layout law).
19782            let v = if bv == "r2w8" && mcols == 2 {
19783                "r2"
19784            } else if bv == "ca" && (!ca_ok || mcols == 8) {
19785                "pf"
19786            } else if bv == "car2" && (!ca_ok || mcols == 8) {
19787                "r2"
19788            } else if bv == "pfr2" && mcols == 8 {
19789                "r2"
19790            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
19791                "rpr2"
19792            }
19793            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
19794            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
19795                if mcols == 8 { "rpr2w8" } else { "rpr2" }
19796            } else if bv == "rpcar2" && mcols == 2 {
19797                "rpca"
19798            }
19799            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
19800            // (rpms has no smem and no alignment need — always valid on rp buffers).
19801            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
19802                "rpr2"
19803            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
19804                "rpr2"
19805            } else {
19806                bv
19807            };
19808            if rp {
19809                match v {
19810                    "base" | "pf" | "ca" | "rp" => "rp",
19811                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
19812                    "r2w8" | "rpr2w8" => {
19813                        if mcols == 2 {
19814                            "rpr2"
19815                        } else {
19816                            "rpr2w8"
19817                        }
19818                    }
19819                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
19820                }
19821            } else {
19822                v
19823            }
19824        } else if mcols == 8 {
19825            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
19826            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
19827            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
19828            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
19829            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
19830            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
19831            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
19832            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
19833            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
19834            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
19835            if rp {
19836                if sc_ok { "rpsc" } else { "rpr2w8" }
19837            } else {
19838                "r2w8"
19839            }
19840        } else if mcols >= 4 {
19841            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
19842            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
19843            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
19844            #[allow(clippy::manual_div_ceil)]
19845            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19846            let blocks = (out_f + 7) / 8;
19847            let r7 = 7 * sms as usize;
19848            let r8 = 8 * sms as usize;
19849            let waves = blocks as f64 / r7 as f64;
19850            let filled = blocks >= 4 * sms as usize;
19851            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
19852            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
19853            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
19854            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
19855                // the extra residency drops the INTEGER wave count -> the straggler wave a
19856                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
19857                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
19858                if rp { "rpr2w8" } else { "r2w8" }
19859            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
19860                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
19861                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
19862                if rp { "rpr2" } else { "r2" }
19863            } else {
19864                // fractional straggler-wave window with no crossing, or grid too small to fill
19865                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
19866                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
19867                if rp { "rp" } else { "pf" }
19868            }
19869        } else if in_f >= 6144 {
19870            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
19871            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
19872            // stays.
19873            if rp { "rpr2" } else { "r2" }
19874        } else if rp {
19875            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
19876            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
19877            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
19878            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
19879            #[allow(clippy::manual_div_ceil)]
19880            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19881            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
19882            if sc_ok && (0.9..=1.1).contains(&waves) {
19883                "rpsc"
19884            } else {
19885                "rp"
19886            }
19887        } else {
19888            "base"
19889        };
19890        variant
19891    }
19892
19893    #[allow(clippy::too_many_arguments)]
19894    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19895    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19896    pub fn qmatvec_mmvq_batched(
19897        &self,
19898        bytes: &CudaSlice<u8>,
19899        aq: &CudaSlice<i8>,
19900        ad: &CudaSlice<f32>,
19901        m: usize,
19902        in_f: usize,
19903        out_f: usize,
19904        qtype: i32,
19905        row_bytes: usize,
19906        mcols: usize,
19907        scale: f32,
19908        rp: bool,
19909    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19910        const ROWS_PER_BLOCK: u32 = 4;
19911        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
19912        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
19913        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
19914        // weight keeps its rp-layout kernel family regardless of the override.
19915        let forced: Option<&'static str> = {
19916            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
19917            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
19918                .as_deref()
19919                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
19920        };
19921        let variant = match forced {
19922            Some(v) if !rp || v.contains("rp") => v,
19923            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
19924        };
19925        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
19926            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
19927        })?;
19928        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
19929        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
19930        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
19931        let variant = if mcols == 16 {
19932            if rp { "rp" } else { "base" }
19933        } else {
19934            variant
19935        };
19936        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
19937        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
19938        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
19939        // per-(token,row) chain (columns c >= m never execute in either form) ->
19940        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
19941        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
19942        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19943        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
19944        if b567
19945            && qtype == QT_NVFP4
19946            && rp
19947            && mcols == 8
19948            && (5..=7).contains(&m)
19949            && matches!(variant, "rpsc" | "rpr2w8")
19950        {
19951            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
19952            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
19953            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19954            let cfg = LaunchConfig {
19955                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
19956                block_dim: (32, ROWS_PER_BLOCK, 1),
19957                shared_mem_bytes: 0,
19958            };
19959            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19960            let __s_b = self.gpu.stream();
19961            let mut b = __s_b.launch_builder(&f);
19962            b.arg(bytes)
19963                .arg(aq)
19964                .arg(ad)
19965                .arg(&mut y)
19966                .arg(&inf)
19967                .arg(&outf)
19968                .arg(&mi)
19969                .arg(&rb);
19970            unsafe {
19971                b.launch(cfg)?;
19972            }
19973            if scale != 1.0 {
19974                self.scale_inplace(&mut y, scale, m * out_f)?;
19975            }
19976            return Ok(y);
19977        }
19978        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
19979            "base" => (base_name.into(), ROWS_PER_BLOCK),
19980            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
19981            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
19982            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
19983            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
19984            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
19985            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
19986            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
19987            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
19988            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
19989            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
19990            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
19991            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
19992            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
19993            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
19994        };
19995        debug_assert!(
19996            !rp || name.contains("_rp"),
19997            "rp weight dispatched to a GGUF-layout kernel"
19998        );
19999        let f = self.func(&name);
20000        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20001        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
20002        let smem = if name.contains("_r2sm_rp") {
20003            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
20004        } else {
20005            0
20006        };
20007        let cfg = LaunchConfig {
20008            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
20009            block_dim: (32, ROWS_PER_BLOCK, 1),
20010            shared_mem_bytes: smem,
20011        };
20012        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20013        let __s_b = self.gpu.stream();
20014        let mut b = __s_b.launch_builder(&f);
20015        b.arg(bytes)
20016            .arg(aq)
20017            .arg(ad)
20018            .arg(&mut y)
20019            .arg(&inf)
20020            .arg(&outf)
20021            .arg(&mi)
20022            .arg(&rb);
20023        unsafe {
20024            b.launch(cfg)?;
20025        }
20026        if scale != 1.0 {
20027            self.scale_inplace(&mut y, scale, m * out_f)?;
20028        }
20029        Ok(y)
20030    }
20031
20032    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
20033    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
20034    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
20035    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20036    pub fn qmatvec_batched_raw(
20037        &self,
20038        bytes: &CudaSlice<u8>,
20039        x: &CudaSlice<f32>,
20040        m: usize,
20041        in_f: usize,
20042        out_f: usize,
20043        qtype: i32,
20044        row_bytes: usize,
20045        mcols: usize,
20046        rp: bool,
20047    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20048        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20049        self.qmatvec_mmvq_batched(
20050            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
20051        )
20052    }
20053
20054    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
20055    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20056    pub fn qmatvec_nvfp4_batched_raw(
20057        &self,
20058        bytes: &CudaSlice<u8>,
20059        x: &CudaSlice<f32>,
20060        m: usize,
20061        in_f: usize,
20062        out_f: usize,
20063        row_bytes: usize,
20064        mcols: usize,
20065        rp: bool,
20066    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20067        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
20068    }
20069
20070    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
20071    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
20072    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
20073    fn try_fp4_gemm(
20074        &self,
20075        w: &crate::model::GpuTensor,
20076        x: &CudaSlice<f32>,
20077        m: usize,
20078        in_f: usize,
20079        out_f: usize,
20080    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20081        use crate::model::GpuTensor;
20082        if cfg!(memra_portable_cuda) {
20083            return Ok(None);
20084        }
20085        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which ONLY the sm_120a fatbin contains:
20086        // cu/qmatvec_gemm.cu omits it on portable builds (MEMRA_PORTABLE_CUDA) AND on sm_100a
20087        // (build.rs passes -DMEMRA_DISABLE_NATIVE_FP4=1 there — the mxf4 block-scale MMA is an
20088        // sm_120a instruction encoding). Refuse at the door on EVERY build that lacks it. The
20089        // portable refusal alone was an enumeration, not a property: a 100a build is not
20090        // portable, so `MEMRA_FP4=1` sailed past it into Engine::func's "kernel not in any
20091        // fatbin" panic — found by the 100a fatbin-lookup census, lane/glm5-b200-prep-20260901
20092        // (same enumeration-vs-property class as the 2026-08-23 stub-polarity fixes in build.rs).
20093        if std::env::var("MEMRA_FP4").is_ok() {
20094            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
20095            assert!(
20096                konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a"),
20097                "MEMRA_FP4 forces the native mxf4 block-scale GEMM (qmatvec_gemm_nvfp4_fp4), \
20098                 which only the sm_120a fatbin contains — this is an sm_{} build. Unset \
20099                 MEMRA_FP4; the W4A8 int8 path is the correct default for NVFP4 weights.",
20100                env!("MEMRA_BUILT_CUDA_ARCH")
20101            );
20102        }
20103        if std::env::var("MEMRA_FP4").is_err() {
20104            return Ok(None);
20105        }
20106        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
20107        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
20108        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
20109        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
20110        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
20111        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
20112        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
20113        // for the common no-macro-scale case.
20114        #[cfg(memra_cutlass)]
20115        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
20116            if let GpuTensor::Quant {
20117                bytes,
20118                qtype,
20119                scale,
20120                row_bytes,
20121                cutlass,
20122                ..
20123            } = w
20124            {
20125                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
20126                    if let Some(cw) = cutlass {
20127                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
20128                        let y = self.cutlass_fp4_gemm(
20129                            &cw.b_packed,
20130                            &cw.sfb_swizzled,
20131                            x,
20132                            *scale,
20133                            m,
20134                            out_f,
20135                            in_f,
20136                        )?;
20137                        return Ok(Some(y));
20138                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
20139                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
20140                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
20141                        // (the load-time repack ~doubles it) — needed for models that don't fit the
20142                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
20143                        let (b_packed, sfb_sw) =
20144                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
20145                        let y =
20146                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
20147                        return Ok(Some(y));
20148                    }
20149                }
20150            }
20151        }
20152        if let GpuTensor::Quant {
20153            bytes,
20154            qtype,
20155            row_bytes,
20156            scale,
20157            rp,
20158            ..
20159        } = w
20160        {
20161            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
20162            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
20163            if *qtype == QT_NVFP4 && in_f.is_multiple_of(64) && !*rp {
20164                let y =
20165                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
20166                return Ok(Some(y));
20167            }
20168        }
20169        Ok(None)
20170    }
20171
20172    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
20173    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
20174    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
20175    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20176    pub fn rms_norm_f16out(
20177        &self,
20178        x: &CudaSlice<f32>,
20179        w: &CudaSlice<f32>,
20180        dst: &mut CudaSlice<f32>,
20181        dst16: &mut CudaSlice<u8>,
20182        ncols: usize,
20183        nrows: usize,
20184        eps: f32,
20185    ) -> Result<(), Box<dyn std::error::Error>> {
20186        let f = self.func("rms_norm_f16out_f32");
20187        let cfg = LaunchConfig {
20188            grid_dim: (nrows as u32, 1, 1),
20189            block_dim: (rms_block(), 1, 1),
20190            shared_mem_bytes: 0,
20191        };
20192        let (nc, e) = (ncols as i32, eps);
20193        let __s_b = self.gpu.stream();
20194        let mut b = __s_b.launch_builder(&f);
20195        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
20196        unsafe {
20197            b.launch(cfg)?;
20198        }
20199        Ok(())
20200    }
20201
20202    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
20203    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
20204    #[allow(clippy::too_many_arguments)]
20205    pub fn add_rms_norm_f16out(
20206        &self,
20207        a: &CudaSlice<f32>,
20208        b: &CudaSlice<f32>,
20209        w: &CudaSlice<f32>,
20210        res: &mut CudaSlice<f32>,
20211        dst: &mut CudaSlice<f32>,
20212        dst16: &mut CudaSlice<u8>,
20213        ncols: usize,
20214        nrows: usize,
20215        eps: f32,
20216    ) -> Result<(), Box<dyn std::error::Error>> {
20217        let f = self.func("add_rms_norm_f16out_f32");
20218        let cfg = LaunchConfig {
20219            grid_dim: (nrows as u32, 1, 1),
20220            block_dim: (rms_block(), 1, 1),
20221            shared_mem_bytes: 0,
20222        };
20223        let (nc, e) = (ncols as i32, eps);
20224        let __s_lb = self.gpu.stream();
20225        let mut lb = __s_lb.launch_builder(&f);
20226        lb.arg(a)
20227            .arg(b)
20228            .arg(w)
20229            .arg(res)
20230            .arg(dst)
20231            .arg(dst16)
20232            .arg(&nc)
20233            .arg(&e);
20234        unsafe {
20235            lb.launch(cfg)?;
20236        }
20237        Ok(())
20238    }
20239
20240    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
20241    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
20242    pub fn matmul_group_xh(
20243        &self,
20244        ws: &[&crate::model::GpuTensor],
20245        x: &CudaSlice<f32>,
20246        xh: &CudaSlice<u8>,
20247        m: usize,
20248    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20249        let mut out = Vec::with_capacity(ws.len());
20250        let in_f = ws[0].in_features();
20251        for w in ws {
20252            if w.in_features() == in_f
20253                && m >= 16
20254                && !self.verify_exact_on()
20255                && let Some(y) = self.try_f16_gemm_pre(w, xh, m)?
20256            {
20257                out.push(y);
20258                continue;
20259            }
20260            out.push(self.matmul(w, x, m)?);
20261        }
20262        Ok(out)
20263    }
20264
20265    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
20266    /// GDN steps). Layouts [T, H].
20267    pub fn gdn_pad_mask(
20268        &self,
20269        beta: &mut CudaSlice<f32>,
20270        g_log: &mut CudaSlice<f32>,
20271        len_d: &CudaSlice<i32>,
20272        h: usize,
20273        t: usize,
20274    ) -> Result<(), Box<dyn std::error::Error>> {
20275        let f = self.func("gdn_pad_mask_f32");
20276        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
20277        let (hi, ti) = (h as i32, t as i32);
20278        let __s_b = self.gpu.stream();
20279        let mut b = __s_b.launch_builder(&f);
20280        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
20281        unsafe {
20282            b.launch(cfg)?;
20283        }
20284        Ok(())
20285    }
20286
20287    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
20288    /// gather for the padded prime graph's h_seed/hlast.
20289    pub fn row_gather_dev(
20290        &self,
20291        src: &CudaSlice<f32>,
20292        dst: &mut CudaSlice<f32>,
20293        len_d: &CudaSlice<i32>,
20294        ncols: usize,
20295    ) -> Result<(), Box<dyn std::error::Error>> {
20296        let f = self.func("row_gather_dev_f32");
20297        let cfg = LaunchConfig::for_num_elems(ncols as u32);
20298        let nc = ncols as i32;
20299        let __s_b = self.gpu.stream();
20300        let mut b = __s_b.launch_builder(&f);
20301        b.arg(src).arg(dst).arg(len_d).arg(&nc);
20302        unsafe {
20303            b.launch(cfg)?;
20304        }
20305        Ok(())
20306    }
20307
20308    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
20309    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
20310    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
20311    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
20312    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
20313    /// different in_f) falls back to its own `matmul` — behavior unchanged.
20314    pub fn matmul_group(
20315        &self,
20316        ws: &[&crate::model::GpuTensor],
20317        x: &CudaSlice<f32>,
20318        m: usize,
20319    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20320        use crate::model::GpuTensor;
20321        let mut out = Vec::with_capacity(ws.len());
20322        let any_mirror = ws
20323            .iter()
20324            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
20325        if m >= 16 && any_mirror && !self.verify_exact_on() {
20326            let in_f = ws[0].in_features();
20327            let xh = self.f16_act(x, m * in_f, in_f)?;
20328            for w in ws {
20329                if w.in_features() == in_f
20330                    && let Some(y) = self.try_f16_gemm_pre(w, &xh, m)?
20331                {
20332                    out.push(y);
20333                    continue;
20334                }
20335                out.push(self.matmul(w, x, m)?);
20336            }
20337            return Ok(out);
20338        }
20339        for w in ws {
20340            out.push(self.matmul(w, x, m)?);
20341        }
20342        Ok(out)
20343    }
20344
20345    /// Cross-request grouped matmul (task #13): run ONE projection group over the
20346    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
20347    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
20348    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
20349    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
20350    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
20351    pub fn matmul_group_multi(
20352        &self,
20353        ws: &[&crate::model::GpuTensor],
20354        xs: &[&CudaSlice<f32>],
20355        ms: &[usize],
20356    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
20357        assert_eq!(xs.len(), ms.len());
20358        let in_f = ws[0].in_features();
20359        let total: usize = ms.iter().sum();
20360        let mut xcat = self.uninit(total * in_f)?;
20361        let mut off = 0usize;
20362        for (x, &m) in xs.iter().zip(ms) {
20363            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
20364            off += m;
20365        }
20366        let ys = self.matmul_group(ws, &xcat, total)?;
20367        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
20368        for (w, y) in ws.iter().zip(ys) {
20369            let out_f = w.out_features();
20370            let mut off = 0usize;
20371            for (s, &m) in ms.iter().enumerate() {
20372                let mut ys_s = self.uninit(m * out_f)?;
20373                let src = y.slice(off * out_f..(off + m) * out_f);
20374                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
20375                out[s].push(ys_s);
20376                off += m;
20377            }
20378        }
20379        Ok(out)
20380    }
20381
20382    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
20383    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
20384    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
20385    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
20386    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
20387    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
20388    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
20389    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
20390    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
20391    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
20392        use crate::model::GpuTensor;
20393        if !legacy_quant_gemm_allowed(
20394            cfg!(memra_portable_cuda),
20395            cfg!(memra_hopper_mma),
20396            std::env::var_os("MEMRA_NO_GEMM").is_some(),
20397        ) {
20398            return false;
20399        }
20400        match w {
20401            GpuTensor::Quant { qtype, .. } => {
20402                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
20403                    || (*qtype == QT_NVFP4 && w.in_features().is_multiple_of(64))
20404            }
20405            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
20406        }
20407    }
20408
20409    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
20410    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
20411    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
20412    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
20413    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
20414    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
20415    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
20416    pub fn qmatvec_gemm(
20417        &self,
20418        w: &crate::model::GpuTensor,
20419        aq: &CudaSlice<i8>,
20420        ad: &CudaSlice<f32>,
20421        m: usize,
20422    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20423        use crate::model::GpuTensor;
20424        let in_f = w.in_features();
20425        let out_f = w.out_features();
20426        let (bytes, qtype, row_bytes, scale, rp) = match w {
20427            GpuTensor::Quant {
20428                bytes,
20429                qtype,
20430                row_bytes,
20431                scale,
20432                rp,
20433                ..
20434            } => (bytes, *qtype, *row_bytes, *scale, *rp),
20435            _ => unreachable!("gemm_supports guaranteed Quant"),
20436        };
20437        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
20438        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
20439        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
20440        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
20441        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
20442        if cfg!(memra_hopper_mma)
20443            && qtype == QT_Q8_0
20444            && out_f.is_multiple_of(64)
20445            && wgmma_gemm_enabled()
20446            && let GpuTensor::Quant { rp4: Some(m4), .. } = w
20447        {
20448            let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
20449            if scale != 1.0 {
20450                self.scale_inplace(&mut y, scale, m * out_f)?;
20451            }
20452            return Ok(y);
20453        }
20454        let name = match qtype {
20455            QT_Q8_0 => "qmatvec_gemm_q8_0",
20456            QT_Q4_K => "qmatvec_gemm_q4_K",
20457            QT_Q4_0 => {
20458                if rp {
20459                    "qmatvec_gemm_q4_0_rp"
20460                } else {
20461                    "qmatvec_gemm_q4_0"
20462                }
20463            }
20464            QT_Q5_K => "qmatvec_gemm_q5_K",
20465            QT_Q6_K => "qmatvec_gemm_q6_K",
20466            QT_NVFP4 => {
20467                if rp {
20468                    "qmatvec_gemm_nvfp4_rp"
20469                } else {
20470                    "qmatvec_gemm_nvfp4"
20471                }
20472            }
20473            _ => unreachable!(),
20474        };
20475        let f = self.func(name);
20476        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
20477        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
20478        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
20479        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
20480        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
20481        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
20482        let k1_tile = if is_k1 {
20483            k1_launch_override().unwrap_or((128, 128, 8))
20484        } else {
20485            (128, 128, 8)
20486        };
20487        let (bm, bn): (u32, u32) = if is_k1 {
20488            (k1_tile.0, k1_tile.1)
20489        } else {
20490            (64, 256)
20491        };
20492        let warps: u32 = if is_k1 {
20493            k1_tile.2
20494        } else {
20495            match qtype {
20496                QT_NVFP4 => 8,
20497                _ => 4,
20498            }
20499        };
20500        let cfg = LaunchConfig {
20501            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
20502            block_dim: (32, warps, 1),
20503            shared_mem_bytes: 0,
20504        };
20505        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20506        let __s_b = self.gpu.stream();
20507        let mut b = __s_b.launch_builder(&f);
20508        b.arg(bytes)
20509            .arg(aq)
20510            .arg(ad)
20511            .arg(&mut y)
20512            .arg(&inf)
20513            .arg(&outf)
20514            .arg(&mi)
20515            .arg(&rb);
20516        unsafe {
20517            b.launch(cfg)?;
20518        }
20519        if scale != 1.0 {
20520            self.scale_inplace(&mut y, scale, m * out_f)?;
20521        }
20522        Ok(y)
20523    }
20524
20525    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
20526    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
20527    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
20528    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
20529    #[allow(clippy::too_many_arguments)]
20530    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20531    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
20532    pub fn qmatvec_gemm_raw(
20533        &self,
20534        bytes: &CudaSlice<u8>,
20535        x: &CudaSlice<f32>,
20536        m: usize,
20537        in_f: usize,
20538        out_f: usize,
20539        qtype: i32,
20540        row_bytes: usize,
20541    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20542        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20543        let name = match qtype {
20544            QT_Q8_0 => "qmatvec_gemm_q8_0",
20545            QT_Q4_K => "qmatvec_gemm_q4_K",
20546            QT_Q4_0 => "qmatvec_gemm_q4_0",
20547            QT_Q5_K => "qmatvec_gemm_q5_K",
20548            QT_Q6_K => "qmatvec_gemm_q6_K",
20549            QT_NVFP4 => "qmatvec_gemm_nvfp4",
20550            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
20551            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
20552        };
20553        let f = self.func(name);
20554        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
20555        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
20556        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
20557        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
20558        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
20559        let k1_tile = if is_k1 {
20560            k1_launch_override().unwrap_or((128, 128, 8))
20561        } else {
20562            (128, 128, 8)
20563        };
20564        let (bm, bn): (u32, u32) = if is_k1 {
20565            (k1_tile.0, k1_tile.1)
20566        } else {
20567            (64, 256)
20568        };
20569        let warps: u32 = if is_k1 {
20570            k1_tile.2
20571        } else {
20572            match qtype {
20573                QT_NVFP4 | QT_NVFP4_RP => 8,
20574                _ => 4,
20575            }
20576        };
20577        let cfg = LaunchConfig {
20578            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
20579            block_dim: (32, warps, 1),
20580            shared_mem_bytes: 0,
20581        };
20582        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20583        let __s_b = self.gpu.stream();
20584        let mut b = __s_b.launch_builder(&f);
20585        b.arg(bytes)
20586            .arg(&aq)
20587            .arg(&ad)
20588            .arg(&mut y)
20589            .arg(&inf)
20590            .arg(&outf)
20591            .arg(&mi)
20592            .arg(&rb);
20593        unsafe {
20594            b.launch(cfg)?;
20595        }
20596        Ok(y)
20597    }
20598
20599    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
20600    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
20601    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
20602    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
20603    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
20604    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
20605    pub fn qmatvec_gemm_q8_0_wgmma_raw(
20606        &self,
20607        rp4: &CudaSlice<u8>,
20608        aq: &CudaSlice<i8>,
20609        ad: &CudaSlice<f32>,
20610        m: usize,
20611        in_f: usize,
20612        out_f: usize,
20613    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20614        assert!(
20615            out_f.is_multiple_of(64) && in_f.is_multiple_of(32),
20616            "wgmma GEMM needs out_f%64==0, in_f%32==0"
20617        );
20618        let f = self.func("qmatvec_gemm_q8_0_wgmma");
20619        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
20620        let cfg = LaunchConfig {
20621            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
20622            block_dim: (128, 1, 1),
20623            shared_mem_bytes: 0,
20624        };
20625        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
20626        let __s_b = self.gpu.stream();
20627        let mut b = __s_b.launch_builder(&f);
20628        b.arg(rp4)
20629            .arg(aq)
20630            .arg(ad)
20631            .arg(&mut y)
20632            .arg(&inf)
20633            .arg(&outf)
20634            .arg(&mi);
20635        unsafe {
20636            b.launch(cfg)?;
20637        }
20638        Ok(y)
20639    }
20640
20641    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
20642    pub fn scale_inplace(
20643        &self,
20644        y: &mut CudaSlice<f32>,
20645        s: f32,
20646        n: usize,
20647    ) -> Result<(), Box<dyn std::error::Error>> {
20648        let f = self.func("scale_f32");
20649        let cfg = LaunchConfig::for_num_elems(n as u32);
20650        let (sf, ni) = (s, n as i32);
20651        let __s_b = self.gpu.stream();
20652        let mut b = __s_b.launch_builder(&f);
20653        b.arg(y).arg(&sf).arg(&ni);
20654        unsafe {
20655            b.launch(cfg)?;
20656        }
20657        Ok(())
20658    }
20659
20660    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
20661    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
20662    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
20663    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
20664    pub fn bf16_to_f32(
20665        &self,
20666        data: &cudarc::driver::CudaView<'_, u8>,
20667        n: usize,
20668    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20669        let mut out = self.alloc_uninit::<f32>(n)?;
20670        let f = self.func("bf16_to_f32");
20671        let cfg = LaunchConfig::for_num_elems(n as u32);
20672        let ni = n as i32;
20673        let __s_b = self.gpu.stream();
20674        let mut b = __s_b.launch_builder(&f);
20675        b.arg(data).arg(&mut out).arg(&ni);
20676        unsafe {
20677            b.launch(cfg)?;
20678        }
20679        Ok(out)
20680    }
20681
20682    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
20683    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
20684    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
20685    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
20686    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
20687    /// calls, the spec-verify contract) vs plain linear.
20688    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20689    fn linear_bf16_chunked(
20690        &self,
20691        x: &CudaSlice<f32>,
20692        data: &CudaSlice<u8>,
20693        m: usize,
20694        in_f: usize,
20695        out_f: usize,
20696        exact: bool,
20697        canonical_chunk_rows: Option<usize>,
20698    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20699        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
20700        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
20701        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
20702        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20703        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20704        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20705        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
20706        let started = timing.then(std::time::Instant::now);
20707        let result =
20708            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
20709        if let Some(started) = started {
20710            use std::sync::atomic::Ordering;
20711            self.stream().synchronize()?;
20712            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
20713                + started.elapsed().as_nanos() as u64;
20714            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
20715                + (in_f * out_f * 2) as u64;
20716            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
20717            if calls.is_multiple_of(1024) {
20718                eprintln!(
20719                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
20720                     weight_gb={:.2}",
20721                    ns as f64 / 1.0e6,
20722                    ns as f64 / calls as f64 / 1.0e3,
20723                    wb as f64 / 1.0e9,
20724                );
20725            }
20726        }
20727        result
20728    }
20729
20730    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
20731    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
20732    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
20733    /// numeric-class doors (DEV_ROUTES precedent).
20734    pub(crate) fn bf16_mmv_on() -> bool {
20735        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20736        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
20737    }
20738
20739    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
20740    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
20741    fn matvec_bf16(
20742        &self,
20743        data: &CudaSlice<u8>,
20744        x: &CudaSlice<f32>,
20745        in_f: usize,
20746        out_f: usize,
20747    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20748        if data.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) {
20749            return Err(format!(
20750                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
20751                data.len(),
20752                x.len()
20753            )
20754            .into());
20755        }
20756        let mut y = self.alloc_uninit::<f32>(out_f)?;
20757        let f = self.func("matvec_bf16_f32acc");
20758        let cfg = LaunchConfig {
20759            grid_dim: (out_f as u32, 1, 1),
20760            block_dim: (mmv_block(), 1, 1),
20761            shared_mem_bytes: 0,
20762        };
20763        let ini = in_f as i32;
20764        let __s_bld = self.gpu.stream();
20765        let mut bld = __s_bld.launch_builder(&f);
20766        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
20767        unsafe {
20768            bld.launch(cfg)?;
20769        }
20770        Ok(y)
20771    }
20772
20773    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
20774    /// launches, a position upload, and the rope launch; the position is read directly from
20775    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
20776    #[allow(clippy::too_many_arguments)]
20777    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
20778    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
20779    /// Bit-identical to the split kernels; requires head_dim == 128 and
20780    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
20781    #[allow(clippy::too_many_arguments)]
20782    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
20783    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
20784    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
20785    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
20786    #[allow(clippy::too_many_arguments)]
20787    pub fn qk_norm_rope_append_inc_dcw_rows(
20788        &self,
20789        q_raw_t: &CudaSlice<f32>,
20790        k_raw_t: &CudaSlice<f32>,
20791        v_raw_t: &CudaSlice<f32>,
20792        qw: &CudaSlice<f32>,
20793        kw: &CudaSlice<f32>,
20794        q_out_t: &mut CudaSlice<f32>,
20795        k_out_t: &mut CudaSlice<f32>,
20796        tab: &CudaSlice<u64>,
20797        pos_t: &CudaSlice<i32>,
20798        same_session: bool,
20799        t: usize,
20800        kv_dim_k: usize,
20801        kv_dim_v: usize,
20802        k_tok_bytes: usize,
20803        v_tok_bytes: usize,
20804        head_dim: usize,
20805        n_dims: usize,
20806        nh_q: usize,
20807        nh_k: usize,
20808        eps: f32,
20809        freq_base: f32,
20810        freq_scale: f32,
20811        ff: Option<&CudaSlice<f32>>,
20812    ) -> Result<(), Box<dyn std::error::Error>> {
20813        if head_dim != 128
20814            || kv_dim_v != kv_dim_k
20815            || kv_dim_k != nh_k * head_dim
20816            || t == 0
20817            || t > 32
20818            || tab.len() < t * 6
20819            || pos_t.len() < t
20820            || q_raw_t.len() < t * nh_q * head_dim
20821            || k_raw_t.len() < t * nh_k * head_dim
20822            || v_raw_t.len() < t * kv_dim_v
20823            || q_out_t.len() < t * nh_q * head_dim
20824            || k_out_t.len() < t * nh_k * head_dim
20825        {
20826            return Err(format!(
20827                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
20828                 nh_q={nh_q} nh_k={nh_k}"
20829            )
20830            .into());
20831        }
20832        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
20833        let same_t: i32 = if same_session { t as i32 } else { 0 };
20834        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
20835        let cfg = LaunchConfig {
20836            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
20837            block_dim: (128, 1, 1),
20838            shared_mem_bytes: 0,
20839        };
20840        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
20841        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20842        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
20843        let null: u64 = 0;
20844        let __s_b = self.gpu.stream();
20845        let mut b = __s_b.launch_builder(&f);
20846        b.arg(q_raw_t)
20847            .arg(k_raw_t)
20848            .arg(v_raw_t)
20849            .arg(qw)
20850            .arg(kw)
20851            .arg(q_out_t)
20852            .arg(k_out_t)
20853            .arg(tab)
20854            .arg(pos_t)
20855            .arg(&same_t)
20856            .arg(&kvk)
20857            .arg(&kvv)
20858            .arg(&ktb)
20859            .arg(&vtb)
20860            .arg(&hd)
20861            .arg(&nd)
20862            .arg(&nq)
20863            .arg(&nk)
20864            .arg(&eps)
20865            .arg(&theta_scale)
20866            .arg(&freq_scale);
20867        match ff {
20868            Some(freqs) => {
20869                b.arg(freqs);
20870            }
20871            None => {
20872                b.arg(&null);
20873            }
20874        }
20875        unsafe {
20876            b.launch(cfg)?;
20877        }
20878        Ok(())
20879    }
20880
20881    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20882    pub fn qk_norm_rope_append_inc_dcw(
20883        &self,
20884        q_raw: &CudaSlice<f32>,
20885        k_raw: &CudaSlice<f32>,
20886        v_raw: &CudaSlice<f32>,
20887        qw: &CudaSlice<f32>,
20888        kw: &CudaSlice<f32>,
20889        q_out: &mut CudaSlice<f32>,
20890        k_out: &mut CudaSlice<f32>,
20891        pos: &CudaSlice<i32>,
20892        k_plane: &mut CudaSlice<u8>,
20893        v_plane: &mut CudaSlice<u8>,
20894        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
20895        // (single) writer, exactly like the split append+inc pair it replaces.
20896        len_dev: &CudaSlice<i32>,
20897        base_dev: Option<&CudaSlice<i32>>,
20898        done_ctr: &mut CudaSlice<u32>,
20899        kv_dim_k: usize,
20900        kv_dim_v: usize,
20901        k_tok_bytes: usize,
20902        v_tok_bytes: usize,
20903        head_dim: usize,
20904        n_dims: usize,
20905        nh_q: usize,
20906        nh_k: usize,
20907        eps: f32,
20908        freq_base: f32,
20909        freq_scale: f32,
20910        ff: Option<&CudaSlice<f32>>,
20911    ) -> Result<(), Box<dyn std::error::Error>> {
20912        if head_dim != 128
20913            || kv_dim_v != kv_dim_k
20914            || kv_dim_k != nh_k * head_dim
20915            || q_raw.len() < nh_q * head_dim
20916            || k_raw.len() < nh_k * head_dim
20917            || v_raw.len() < kv_dim_v
20918            || q_out.len() < nh_q * head_dim
20919            || k_out.len() < nh_k * head_dim
20920            || pos.is_empty()
20921            || done_ctr.is_empty()
20922        {
20923            return Err(format!(
20924                "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}"
20925            )
20926            .into());
20927        }
20928        let f = self.func("qk_norm_rope_append_inc_dcw");
20929        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
20930        let cfg = LaunchConfig {
20931            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
20932            block_dim: (128, 1, 1),
20933            shared_mem_bytes: 0,
20934        };
20935        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
20936        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20937        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
20938        let null: u64 = 0;
20939        let __s_b = self.gpu.stream();
20940        let mut b = __s_b.launch_builder(&f);
20941        b.arg(q_raw)
20942            .arg(k_raw)
20943            .arg(v_raw)
20944            .arg(qw)
20945            .arg(kw)
20946            .arg(q_out)
20947            .arg(k_out)
20948            .arg(pos)
20949            .arg(&mut *k_plane)
20950            .arg(&mut *v_plane)
20951            .arg(len_dev);
20952        match base_dev {
20953            Some(base) => {
20954                b.arg(base);
20955            }
20956            None => {
20957                b.arg(&null);
20958            }
20959        }
20960        b.arg(&mut *done_ctr)
20961            .arg(&kvk)
20962            .arg(&kvv)
20963            .arg(&ktb)
20964            .arg(&vtb)
20965            .arg(&hd)
20966            .arg(&nd)
20967            .arg(&nq)
20968            .arg(&eps)
20969            .arg(&theta_scale)
20970            .arg(&freq_scale);
20971        match ff {
20972            Some(freqs) => {
20973                b.arg(freqs);
20974            }
20975            None => {
20976                b.arg(&null);
20977            }
20978        }
20979        unsafe {
20980            b.launch(cfg)?;
20981        }
20982        Ok(())
20983    }
20984
20985    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20986    pub fn qk_norm_rope_into(
20987        &self,
20988        q_raw: &CudaSlice<f32>,
20989        k_raw: &CudaSlice<f32>,
20990        qw: &CudaSlice<f32>,
20991        kw: &CudaSlice<f32>,
20992        q_out: &mut CudaSlice<f32>,
20993        k_out: &mut CudaSlice<f32>,
20994        pos: &CudaSlice<i32>,
20995        head_dim: usize,
20996        n_dims: usize,
20997        nh_q: usize,
20998        nh_k: usize,
20999        eps: f32,
21000        freq_base: f32,
21001        freq_scale: f32,
21002        ff: Option<&CudaSlice<f32>>,
21003    ) -> Result<(), Box<dyn std::error::Error>> {
21004        if head_dim > 512
21005            || q_raw.len() < nh_q * head_dim
21006            || k_raw.len() < nh_k * head_dim
21007            || q_out.len() < nh_q * head_dim
21008            || k_out.len() < nh_k * head_dim
21009            || qw.len() < head_dim
21010            || kw.len() < head_dim
21011            || pos.is_empty()
21012        {
21013            return Err(format!(
21014                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
21015            )
21016            .into());
21017        }
21018        let f = self.func("qk_norm_rope_f32");
21019        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
21020        let cfg = LaunchConfig {
21021            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
21022            block_dim: (128, 1, 1),
21023            shared_mem_bytes: 0,
21024        };
21025        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
21026        let __s_b = self.gpu.stream();
21027        let mut b = __s_b.launch_builder(&f);
21028        b.arg(q_raw)
21029            .arg(k_raw)
21030            .arg(qw)
21031            .arg(kw)
21032            .arg(q_out)
21033            .arg(k_out)
21034            .arg(pos)
21035            .arg(&hd)
21036            .arg(&nd)
21037            .arg(&nq)
21038            .arg(&eps)
21039            .arg(&theta_scale)
21040            .arg(&freq_scale);
21041        match ff {
21042            Some(ffv) => {
21043                b.arg(ffv);
21044                unsafe {
21045                    b.launch(cfg)?;
21046                }
21047            }
21048            None => {
21049                let null: u64 = 0;
21050                b.arg(&null);
21051                unsafe {
21052                    b.launch(cfg)?;
21053                }
21054            }
21055        }
21056        Ok(())
21057    }
21058
21059    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
21060    /// launch computes a rank's whole O partial from its four canonical column blocks.
21061    #[allow(clippy::too_many_arguments)]
21062    pub fn matvec_f32_b4_into(
21063        &self,
21064        w: [&CudaSlice<f32>; 4],
21065        x: &CudaSlice<f32>,
21066        y: &mut CudaSlice<f32>,
21067        block_cols: usize,
21068        out_f: usize,
21069    ) -> Result<(), Box<dyn std::error::Error>> {
21070        if !block_cols.is_multiple_of(4)
21071            || x.len() < 4 * block_cols
21072            || y.len() < out_f
21073            || w.iter().any(|w| w.len() != out_f * block_cols)
21074        {
21075            return Err(format!(
21076                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
21077                x.len()
21078            )
21079            .into());
21080        }
21081        let f = self.func("matvec_f32_b4");
21082        let cfg = LaunchConfig {
21083            grid_dim: (out_f as u32, 1, 1),
21084            block_dim: (128, 1, 1),
21085            shared_mem_bytes: 0,
21086        };
21087        let (bc, of) = (block_cols as i32, out_f as i32);
21088        let __s_b = self.gpu.stream();
21089        let mut b = __s_b.launch_builder(&f);
21090        b.arg(w[0])
21091            .arg(w[1])
21092            .arg(w[2])
21093            .arg(w[3])
21094            .arg(x)
21095            .arg(y)
21096            .arg(&bc)
21097            .arg(&of);
21098        unsafe {
21099            b.launch(cfg)?;
21100        }
21101        Ok(())
21102    }
21103
21104    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
21105    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
21106    pub fn axpy_rows_seq_into(
21107        &self,
21108        x: &CudaSlice<f32>,
21109        w: &CudaSlice<f32>,
21110        y: &mut CudaSlice<f32>,
21111        width: usize,
21112        n_rows: usize,
21113    ) -> Result<(), Box<dyn std::error::Error>> {
21114        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
21115            return Err(format!(
21116                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
21117                x.len(),
21118                w.len(),
21119                y.len()
21120            )
21121            .into());
21122        }
21123        let f = self.func("axpy_rows_seq_f32");
21124        let cfg = LaunchConfig::for_num_elems(width as u32);
21125        let (wi, nr) = (width as i32, n_rows as i32);
21126        let __s_b = self.gpu.stream();
21127        let mut b = __s_b.launch_builder(&f);
21128        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
21129        unsafe {
21130            b.launch(cfg)?;
21131        }
21132        Ok(())
21133    }
21134
21135    /// Token-major sequential weighted row sums. Each token reduces exactly `slots` rows in
21136    /// canonical route order.
21137    pub fn axpy_rows_seq_tokens_into(
21138        &self,
21139        x: &CudaSlice<f32>,
21140        w: &CudaSlice<f32>,
21141        y: &mut CudaSlice<f32>,
21142        width: usize,
21143        slots: usize,
21144        tokens: usize,
21145    ) -> Result<(), Box<dyn std::error::Error>> {
21146        let rows = slots
21147            .checked_mul(tokens)
21148            .ok_or("axpy_rows_seq_tokens row count overflow")?;
21149        if x.len() < rows * width || w.len() < rows || y.len() < tokens * width {
21150            return Err(format!(
21151                "axpy_rows_seq_tokens geometry x={} w={} y={} width={width} \
21152                 slots={slots} tokens={tokens}",
21153                x.len(),
21154                w.len(),
21155                y.len()
21156            )
21157            .into());
21158        }
21159        let f = self.func("axpy_rows_seq_tokens_f32");
21160        let block = 256u32;
21161        let cfg = LaunchConfig {
21162            grid_dim: ((width as u32).div_ceil(block), tokens as u32, 1),
21163            block_dim: (block, 1, 1),
21164            shared_mem_bytes: 0,
21165        };
21166        let (wi, sl, tk) = (width as i32, slots as i32, tokens as i32);
21167        let __s_b = self.gpu.stream();
21168        let mut b = __s_b.launch_builder(&f);
21169        b.arg(x).arg(w).arg(y).arg(&wi).arg(&sl).arg(&tk);
21170        unsafe {
21171            b.launch(cfg)?;
21172        }
21173        Ok(())
21174    }
21175
21176    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
21177    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
21178    /// exact sequential FP chain of the base kernel over that window.
21179    #[allow(clippy::too_many_arguments)]
21180    pub fn axpy_rows_seq_md_off_into(
21181        &self,
21182        x: &CudaSlice<f32>,
21183        w_route: &CudaSlice<f32>,
21184        md: &CudaSlice<f32>,
21185        sel: &CudaSlice<i32>,
21186        y: &mut CudaSlice<f32>,
21187        width: usize,
21188        n_rows: usize,
21189        row0: usize,
21190    ) -> Result<(), Box<dyn std::error::Error>> {
21191        if x.len() < (row0 + n_rows) * width
21192            || w_route.len() < row0 + n_rows
21193            || sel.len() < row0 + n_rows
21194            || y.len() < width
21195        {
21196            return Err(format!(
21197                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
21198                 rows={n_rows} row0={row0}",
21199                x.len(),
21200                w_route.len(),
21201                sel.len(),
21202                y.len()
21203            )
21204            .into());
21205        }
21206        let f = self.func("axpy_rows_seq_md_off_f32");
21207        let cfg = LaunchConfig::for_num_elems(width as u32);
21208        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
21209        let __s_b = self.gpu.stream();
21210        let mut b = __s_b.launch_builder(&f);
21211        b.arg(x)
21212            .arg(w_route)
21213            .arg(md)
21214            .arg(sel)
21215            .arg(y)
21216            .arg(&wi)
21217            .arg(&nr)
21218            .arg(&r0);
21219        unsafe {
21220            b.launch(cfg)?;
21221        }
21222        Ok(())
21223    }
21224
21225    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
21226    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
21227    #[allow(clippy::too_many_arguments)]
21228    pub fn axpy_rows_seq_md_into(
21229        &self,
21230        x: &CudaSlice<f32>,
21231        w_route: &CudaSlice<f32>,
21232        md: &CudaSlice<f32>,
21233        sel: &CudaSlice<i32>,
21234        y: &mut CudaSlice<f32>,
21235        width: usize,
21236        n_rows: usize,
21237    ) -> Result<(), Box<dyn std::error::Error>> {
21238        if x.len() < n_rows * width
21239            || w_route.len() < n_rows
21240            || sel.len() < n_rows
21241            || y.len() < width
21242        {
21243            return Err(format!(
21244                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
21245                x.len(),
21246                w_route.len(),
21247                sel.len(),
21248                y.len()
21249            )
21250            .into());
21251        }
21252        let f = self.func("axpy_rows_seq_md_f32");
21253        let cfg = LaunchConfig::for_num_elems(width as u32);
21254        let (wi, nr) = (width as i32, n_rows as i32);
21255        let __s_b = self.gpu.stream();
21256        let mut b = __s_b.launch_builder(&f);
21257        b.arg(x)
21258            .arg(w_route)
21259            .arg(md)
21260            .arg(sel)
21261            .arg(y)
21262            .arg(&wi)
21263            .arg(&nr);
21264        unsafe {
21265            b.launch(cfg)?;
21266        }
21267        Ok(())
21268    }
21269
21270    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
21271    #[allow(clippy::too_many_arguments)]
21272    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
21273    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
21274    /// land column-major-of-rows: yq[c*out_q + row] etc.
21275    #[allow(clippy::too_many_arguments)]
21276    pub fn matvec_bf16_qkvg_tcol_into(
21277        &self,
21278        wq: &CudaSlice<u8>,
21279        wk: &CudaSlice<u8>,
21280        wv: &CudaSlice<u8>,
21281        wg: &CudaSlice<u8>,
21282        x_t: &CudaSlice<f32>,
21283        yq: &mut CudaSlice<f32>,
21284        yk: &mut CudaSlice<f32>,
21285        yv: &mut CudaSlice<f32>,
21286        yg: &mut CudaSlice<f32>,
21287        in_f: usize,
21288        out_q: usize,
21289        out_kv: usize,
21290        out_g: usize,
21291        t: usize,
21292    ) -> Result<(), Box<dyn std::error::Error>> {
21293        if t == 0
21294            || t > 8
21295            || !in_f.is_multiple_of(8)
21296            || x_t.len() < t * in_f
21297            || yq.len() < t * out_q
21298            || yk.len() < t * out_kv
21299            || yv.len() < t * out_kv
21300            || (out_g > 0 && yg.len() < t * out_g)
21301        {
21302            return Err("matvec_bf16_qkvg_tcol geometry".into());
21303        }
21304        let grid = out_q + 2 * out_kv + out_g;
21305        let cfg = LaunchConfig {
21306            grid_dim: (grid as u32, 1, 1),
21307            block_dim: (mmv_block(), 1, 1),
21308            shared_mem_bytes: 0,
21309        };
21310        let (ini, oq, okv, og, ti) = (
21311            in_f as i32,
21312            out_q as i32,
21313            out_kv as i32,
21314            out_g as i32,
21315            t as i32,
21316        );
21317        let __s_b = self.gpu.stream();
21318        // One runtime-T program for every live width. The compile-time 2/4/8 twins are
21319        // retained in the fatbin as research controls, but dispatching them by the current
21320        // batch width changes kernels inside a request when peers arrive or retire. That is
21321        // a load-history numeric-program switch, and their pre-twin TOKFP receipts did not
21322        // qualify it (Hermes `64fa2b55baf0d887`).
21323        let f = self.func("matvec_bf16_qkvg_tcol");
21324        let mut b = __s_b.launch_builder(&f);
21325        b.arg(wq)
21326            .arg(wk)
21327            .arg(wv)
21328            .arg(wg)
21329            .arg(x_t)
21330            .arg(yq)
21331            .arg(yk)
21332            .arg(yv)
21333            .arg(yg)
21334            .arg(&ini)
21335            .arg(&oq)
21336            .arg(&okv)
21337            .arg(&og)
21338            .arg(&ti);
21339        unsafe {
21340            b.launch(cfg)?;
21341        }
21342        Ok(())
21343    }
21344
21345    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
21346    pub fn matvec_bf16_qkvg_into(
21347        &self,
21348        wq: &CudaSlice<u8>,
21349        wk: &CudaSlice<u8>,
21350        wv: &CudaSlice<u8>,
21351        wg: &CudaSlice<u8>,
21352        x: &CudaSlice<f32>,
21353        yq: &mut CudaSlice<f32>,
21354        yk: &mut CudaSlice<f32>,
21355        yv: &mut CudaSlice<f32>,
21356        yg: &mut CudaSlice<f32>,
21357        in_f: usize,
21358        out_q: usize,
21359        out_kv: usize,
21360        out_g: usize,
21361    ) -> Result<(), Box<dyn std::error::Error>> {
21362        if !in_f.is_multiple_of(8)
21363            || wq.len() != out_q * in_f * 2
21364            || wk.len() != out_kv * in_f * 2
21365            || wv.len() != out_kv * in_f * 2
21366            || wg.len() < out_g * in_f * 2
21367            || x.len() < in_f
21368            || yq.len() < out_q
21369            || yk.len() < out_kv
21370            || yv.len() < out_kv
21371            || (out_g > 0 && yg.len() < out_g)
21372        {
21373            return Err(format!(
21374                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
21375            )
21376            .into());
21377        }
21378        let f = self.func("matvec_bf16_qkvg");
21379        let cfg = LaunchConfig {
21380            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
21381            block_dim: (mmv_block(), 1, 1),
21382            shared_mem_bytes: 0,
21383        };
21384        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
21385        let __s_b = self.gpu.stream();
21386        let mut b = __s_b.launch_builder(&f);
21387        b.arg(wq)
21388            .arg(wk)
21389            .arg(wv)
21390            .arg(wg)
21391            .arg(x)
21392            .arg(yq)
21393            .arg(yk)
21394            .arg(yv)
21395            .arg(yg)
21396            .arg(&inf)
21397            .arg(&oq)
21398            .arg(&okv)
21399            .arg(&og);
21400        unsafe {
21401            b.launch(cfg)?;
21402        }
21403        Ok(())
21404    }
21405
21406    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
21407    pub fn matvec_bf16_b4_into(
21408        &self,
21409        w: [&CudaSlice<u8>; 4],
21410        x: &CudaSlice<f32>,
21411        y: &mut CudaSlice<f32>,
21412        block_cols: usize,
21413        out_f: usize,
21414    ) -> Result<(), Box<dyn std::error::Error>> {
21415        if !block_cols.is_multiple_of(8)
21416            || x.len() < 4 * block_cols
21417            || y.len() < out_f
21418            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
21419        {
21420            return Err(format!(
21421                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
21422                x.len()
21423            )
21424            .into());
21425        }
21426        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
21427        // bit-identical per row (the second row's stream hides the first's reduce tail).
21428        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21429        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
21430        let f = self.func(if x2 {
21431            "matvec_bf16_b4_x2"
21432        } else {
21433            "matvec_bf16_b4"
21434        });
21435        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
21436        let cfg = LaunchConfig {
21437            grid_dim: (grid as u32, 1, 1),
21438            block_dim: (mmv_block(), 1, 1),
21439            shared_mem_bytes: 0,
21440        };
21441        let (bc, of) = (block_cols as i32, out_f as i32);
21442        let __s_b = self.gpu.stream();
21443        let mut b = __s_b.launch_builder(&f);
21444        b.arg(w[0])
21445            .arg(w[1])
21446            .arg(w[2])
21447            .arg(w[3])
21448            .arg(x)
21449            .arg(y)
21450            .arg(&bc)
21451            .arg(&of);
21452        unsafe {
21453            b.launch(cfg)?;
21454        }
21455        Ok(())
21456    }
21457
21458    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
21459    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
21460    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
21461    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
21462    /// t=1 program).
21463    pub fn matvec_bf16_b4_tcol_into(
21464        &self,
21465        w: [&CudaSlice<u8>; 4],
21466        x_t: &CudaSlice<f32>,
21467        y_t: &mut CudaSlice<f32>,
21468        block_cols: usize,
21469        out_f: usize,
21470        t: usize,
21471    ) -> Result<(), Box<dyn std::error::Error>> {
21472        if !block_cols.is_multiple_of(8)
21473            || t == 0
21474            || t > 8
21475            || x_t.len() < t * 4 * block_cols
21476            || y_t.len() < t * out_f
21477            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
21478        {
21479            return Err(format!(
21480                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
21481                x_t.len()
21482            )
21483            .into());
21484        }
21485        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
21486            return Err(
21487                "b4 tcol verify is qualified against the plain b4 kernel only \
21488                        (MEMRA_B4_X2=1 is a different t=1 program)"
21489                    .into(),
21490            );
21491        }
21492        // Keep one runtime-T program at every live width. Compile-time twins remain research
21493        // controls only; selecting them from the changing batch width switches programs
21494        // mid-request.
21495        let cfg = LaunchConfig {
21496            grid_dim: (out_f as u32, 1, 1),
21497            block_dim: (mmv_block(), 1, 1),
21498            shared_mem_bytes: 0,
21499        };
21500        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
21501        let __s_b = self.gpu.stream();
21502        let f = self.func("matvec_bf16_b4_tcol");
21503        let mut b = __s_b.launch_builder(&f);
21504        b.arg(w[0])
21505            .arg(w[1])
21506            .arg(w[2])
21507            .arg(w[3])
21508            .arg(x_t)
21509            .arg(y_t)
21510            .arg(&bc)
21511            .arg(&of)
21512            .arg(&ti);
21513        unsafe {
21514            b.launch(cfg)?;
21515        }
21516        Ok(())
21517    }
21518
21519    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
21520    /// q8_0 row bytes for an `in_f`-wide weight row: one 34-byte block per 32 weights.
21521    pub fn q8_0_row_bytes(in_f: usize) -> usize {
21522        in_f / 32 * 34
21523    }
21524
21525    /// Encode a resident bf16 weight slab into its q8_0 mirror (MEMRA_STEP_TP_W8). Runs once
21526    /// per matrix at load; the block program is the one `quant_K_block` writes for the KV
21527    /// cache, so the two formats cannot drift apart.
21528    pub fn encode_q8_0_from_bf16(
21529        &self,
21530        w_bf16: &CudaSlice<u8>,
21531        out: &mut CudaSlice<u8>,
21532        in_f: usize,
21533        out_f: usize,
21534    ) -> Result<(), Box<dyn std::error::Error>> {
21535        if !in_f.is_multiple_of(32)
21536            || w_bf16.len() < in_f * out_f * 2
21537            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
21538        {
21539            return Err(format!(
21540                "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
21541                w_bf16.len(),
21542                out.len()
21543            )
21544            .into());
21545        }
21546        let f = self.func("encode_q8_0_rows_from_bf16");
21547        // Flat 1D grid of (row, 32-block) pairs, 4 pairs per block: rows on grid.y would cap
21548        // at 65535 and the LM head has 128896 rows.
21549        const PAIRS_PER_BLOCK: u32 = 4;
21550        let pairs = (out_f * (in_f / 32)) as u64;
21551        let cfg = LaunchConfig {
21552            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
21553            block_dim: (32, PAIRS_PER_BLOCK, 1),
21554            shared_mem_bytes: 0,
21555        };
21556        let (ini, outi) = (in_f as i32, out_f as i32);
21557        let __s_b = self.gpu.stream();
21558        let mut b = __s_b.launch_builder(&f);
21559        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
21560        unsafe {
21561            b.launch(cfg)?;
21562        }
21563        Ok(())
21564    }
21565
21566    /// ROW-RANGE-VIEW twin of `encode_q8_0_from_bf16`. Identical kernel, identical launch
21567    /// geometry, identical per-row program: only the operand type differs, because the split
21568    /// decode paths hold their rows as a `CudaView` of the resident slab, not as an owned slab.
21569    pub fn encode_q8_0_from_bf16_view(
21570        &self,
21571        w_bf16: &cudarc::driver::CudaView<'_, u8>,
21572        out: &mut CudaSlice<u8>,
21573        in_f: usize,
21574        out_f: usize,
21575    ) -> Result<(), Box<dyn std::error::Error>> {
21576        if !in_f.is_multiple_of(32)
21577            || w_bf16.len() < in_f * out_f * 2
21578            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
21579        {
21580            return Err(format!(
21581                "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
21582                w_bf16.len(),
21583                out.len()
21584            )
21585            .into());
21586        }
21587        let f = self.func("encode_q8_0_rows_from_bf16");
21588        const PAIRS_PER_BLOCK: u32 = 4;
21589        let pairs = (out_f * (in_f / 32)) as u64;
21590        let cfg = LaunchConfig {
21591            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
21592            block_dim: (32, PAIRS_PER_BLOCK, 1),
21593            shared_mem_bytes: 0,
21594        };
21595        let (ini, outi) = (in_f as i32, out_f as i32);
21596        let __s_b = self.gpu.stream();
21597        let mut b = __s_b.launch_builder(&f);
21598        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
21599        unsafe {
21600            b.launch(cfg)?;
21601        }
21602        Ok(())
21603    }
21604
21605    /// Fused q8_0 QKV against a q8_1 activation (MEMRA_STEP_TP_W8): one launch over the
21606    /// stacked q/k/v rows, each row running the exact `qmatvec_q8_0_mmvq_rp` per-row program.
21607    /// Bit-identical to three per-matrix mmvq calls; it exists because those three launches
21608    /// plus the activation quantize measured SLOWER than the bf16 fused kernel.
21609    #[allow(clippy::too_many_arguments)]
21610    pub fn qmatvec_q8_0_qkv_rp_into(
21611        &self,
21612        wq: &CudaSlice<u8>,
21613        wk: &CudaSlice<u8>,
21614        wv: &CudaSlice<u8>,
21615        aq: &CudaSlice<i8>,
21616        ad: &CudaSlice<f32>,
21617        yq: &mut CudaSlice<f32>,
21618        yk: &mut CudaSlice<f32>,
21619        yv: &mut CudaSlice<f32>,
21620        in_f: usize,
21621        out_q: usize,
21622        out_kv: usize,
21623    ) -> Result<(), Box<dyn std::error::Error>> {
21624        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
21625        let rows = out_q + 2 * out_kv;
21626        let nblk = in_f / 32;
21627        if !in_f.is_multiple_of(32)
21628            || aq.len() < in_f
21629            || ad.len() < nblk
21630            || yq.len() < out_q
21631            || yk.len() < out_kv
21632            || yv.len() < out_kv
21633            || wq.len() < out_q * nblk * 34
21634            || wk.len() < out_kv * nblk * 34
21635            || wv.len() < out_kv * nblk * 34
21636        {
21637            return Err(
21638                format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
21639            );
21640        }
21641        let f = self.func("qmatvec_q8_0_qkv_rp");
21642        let cfg = LaunchConfig {
21643            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21644            block_dim: (32, ROWS_PER_BLOCK, 1),
21645            shared_mem_bytes: 0,
21646        };
21647        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
21648        let __s_b = self.gpu.stream();
21649        let mut b = __s_b.launch_builder(&f);
21650        b.arg(wq)
21651            .arg(wk)
21652            .arg(wv)
21653            .arg(aq)
21654            .arg(ad)
21655            .arg(yq)
21656            .arg(yk)
21657            .arg(yv)
21658            .arg(&ini)
21659            .arg(&oq)
21660            .arg(&okv);
21661        unsafe {
21662            b.launch(cfg)?;
21663        }
21664        Ok(())
21665    }
21666
21667    /// Fused q8_0 O projection over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8): one
21668    /// launch, one warp per output row, per-block reduce then add — the same shape
21669    /// `matvec_bf16_b4` uses, against a q8_1 activation.
21670    #[allow(clippy::too_many_arguments)]
21671    pub fn qmatvec_q8_0_b4_rp_into(
21672        &self,
21673        w: [&CudaSlice<u8>; 4],
21674        aq: &CudaSlice<i8>,
21675        ad: &CudaSlice<f32>,
21676        y: &mut CudaSlice<f32>,
21677        block_cols: usize,
21678        out_f: usize,
21679    ) -> Result<(), Box<dyn std::error::Error>> {
21680        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
21681        let nblk = block_cols / 32;
21682        if !block_cols.is_multiple_of(32)
21683            || aq.len() < 4 * block_cols
21684            || ad.len() < 4 * nblk
21685            || y.len() < out_f
21686            || w.iter().any(|p| p.len() < out_f * nblk * 34)
21687        {
21688            return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
21689        }
21690        let f = self.func("qmatvec_q8_0_b4_rp");
21691        let cfg = LaunchConfig {
21692            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21693            block_dim: (32, ROWS_PER_BLOCK, 1),
21694            shared_mem_bytes: 0,
21695        };
21696        let (bc, of) = (block_cols as i32, out_f as i32);
21697        let __s_b = self.gpu.stream();
21698        let mut b = __s_b.launch_builder(&f);
21699        b.arg(w[0])
21700            .arg(w[1])
21701            .arg(w[2])
21702            .arg(w[3])
21703            .arg(aq)
21704            .arg(ad)
21705            .arg(y)
21706            .arg(&bc)
21707            .arg(&of);
21708        unsafe {
21709            b.launch(cfg)?;
21710        }
21711        Ok(())
21712    }
21713
21714    /// T-column twin of `matvec_bf16_via_q8_mirror`: one q8 launch over all t rows, sharing the
21715    /// same pointer-keyed mirror cache and a t-wide q8_1 activation.
21716    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
21717    fn matvec_bf16_via_q8_mirror_t(
21718        &self,
21719        data: &CudaSlice<u8>,
21720        x: &CudaSlice<f32>,
21721        y: &mut CudaSlice<f32>,
21722        in_f: usize,
21723        out_f: usize,
21724        t: usize,
21725    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
21726        use cudarc::driver::DevicePtr;
21727        let key = {
21728            let s = self.gpu.stream();
21729            let (p, _g) = data.device_ptr(&s);
21730            (p, in_f as u32, out_f as u32)
21731        };
21732        {
21733            let mut mirrors = self
21734                .w8_mirrors
21735                .lock()
21736                .map_err(|_| "w8 mirror map is poisoned")?;
21737            if !mirrors.contains_key(&key) {
21738                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
21739                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
21740                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
21741                mirrors.insert(key, planar);
21742            }
21743        }
21744        let nblk = in_f / 32;
21745        // The t-wide activation scratch is keyed by (in_f, t-cap) so a wider walk regrows it.
21746        let akey = in_f * 64 + t.min(32);
21747        {
21748            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21749            if !act.contains_key(&akey) {
21750                let aq = self.alloc_i8_uninit(32 * in_f)?;
21751                let ad = self.alloc_uninit::<f32>(32 * nblk)?;
21752                act.insert(akey, (aq, ad));
21753            }
21754            let (aq, ad) = act.get_mut(&akey).expect("just inserted");
21755            self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
21756        }
21757        let mirrors = self
21758            .w8_mirrors
21759            .lock()
21760            .map_err(|_| "w8 mirror map is poisoned")?;
21761        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21762        let mirror = mirrors.get(&key).expect("built above");
21763        let (aq, ad) = act.get(&akey).expect("built above");
21764        const ROWS_PER_BLOCK: u32 = 4;
21765        let (ini, of) = (in_f as i32, out_f as i32);
21766        // MEMRA_Q8T_WONCE=1: the weight-once twin — one row grid, each weight int4 loaded once
21767        // and dotted against all t columns. The `_t` form re-streams the shared weights per
21768        // column through __ldcs (measured 1.43-1.67x a single-column call for 2 columns).
21769        if q8t_wonce_on() && t <= 32 {
21770            let f = self.func(if t <= 8 {
21771                "qmatvec_q8_0_rows_tw"
21772            } else {
21773                "qmatvec_q8_0_rows_tw32"
21774            });
21775            let cfg = LaunchConfig {
21776                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21777                block_dim: (32, ROWS_PER_BLOCK, 1),
21778                shared_mem_bytes: 0,
21779            };
21780            let ti = t as i32;
21781            let __s_b = self.gpu.stream();
21782            let mut b = __s_b.launch_builder(&f);
21783            b.arg(mirror)
21784                .arg(aq)
21785                .arg(ad)
21786                .arg(&mut *y)
21787                .arg(&ini)
21788                .arg(&of)
21789                .arg(&ti);
21790            unsafe {
21791                b.launch(cfg)?;
21792            }
21793            return Ok(Some(()));
21794        }
21795        let f = self.func("qmatvec_q8_0_rows_t");
21796        let cfg = LaunchConfig {
21797            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
21798            block_dim: (32, ROWS_PER_BLOCK, 1),
21799            shared_mem_bytes: 0,
21800        };
21801        let __s_b = self.gpu.stream();
21802        let mut b = __s_b.launch_builder(&f);
21803        b.arg(mirror)
21804            .arg(aq)
21805            .arg(ad)
21806            .arg(&mut *y)
21807            .arg(&ini)
21808            .arg(&of);
21809        unsafe {
21810            b.launch(cfg)?;
21811        }
21812        Ok(Some(()))
21813    }
21814
21815    /// Get-or-build this bf16 weight's q8_0 mirror and run the GEMV through it. Returns
21816    /// `None` when the shape has no mirror form, so the caller falls back to bf16.
21817    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
21818    fn matvec_bf16_via_q8_mirror(
21819        &self,
21820        data: &CudaSlice<u8>,
21821        x: &CudaSlice<f32>,
21822        y: &mut CudaSlice<f32>,
21823        in_f: usize,
21824        out_f: usize,
21825    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
21826        use cudarc::driver::DevicePtr;
21827        let key = {
21828            let s = self.gpu.stream();
21829            let (p, _g) = data.device_ptr(&s);
21830            (p, in_f as u32, out_f as u32)
21831        };
21832        {
21833            let mut mirrors = self
21834                .w8_mirrors
21835                .lock()
21836                .map_err(|_| "w8 mirror map is poisoned")?;
21837            if !mirrors.contains_key(&key) {
21838                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
21839                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
21840                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
21841                mirrors.insert(key, planar);
21842                // Which weights this half actually covers is not obvious from the call graph:
21843                // the head and the shared expert may reach the GPU through the rows fast path
21844                // or the fused dual-silu launcher instead of here. One line per mirror answers
21845                // that without a profiler (the hybrid half measured +0.1% and this is how we
21846                // find out whether it even fired).
21847                if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
21848                    eprintln!(
21849                        "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
21850                        mirrors.len()
21851                    );
21852                }
21853            }
21854        }
21855        let nblk = in_f / 32;
21856        {
21857            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21858            if !act.contains_key(&in_f) {
21859                let aq = self.alloc_uninit::<i8>(in_f)?;
21860                let ad = self.alloc_uninit::<f32>(nblk)?;
21861                act.insert(in_f, (aq, ad));
21862            }
21863            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
21864            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
21865        }
21866        let mirrors = self
21867            .w8_mirrors
21868            .lock()
21869            .map_err(|_| "w8 mirror map is poisoned")?;
21870        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21871        let mirror = mirrors.get(&key).expect("built above");
21872        let (aq, ad) = act.get(&in_f).expect("built above");
21873        self.qmatvec_mmvq_into(
21874            mirror,
21875            aq,
21876            ad,
21877            1,
21878            in_f,
21879            out_f,
21880            QT_Q8_0,
21881            Self::q8_0_row_bytes(in_f),
21882            1.0,
21883            true,
21884            y,
21885        )?;
21886        Ok(Some(()))
21887    }
21888
21889    /// T-column q8_0 QKV for the VERIFY walk (MEMRA_STEP_TP_W8). nsys put the bf16 twin
21890    /// `matvec_bf16_qkvg_tcol` at 12.3% of spec GPU time and `matvec_bf16_b4_tcol` at 24.8%:
21891    /// the W8 door had replaced only the decode kernels, so 37% of the verify still streamed
21892    /// bf16. Bit-identical to `t` separate `qmatvec_q8_0_qkv_rp` calls.
21893    #[allow(clippy::too_many_arguments)]
21894    pub fn qmatvec_q8_0_qkv_rp_t_into(
21895        &self,
21896        wq: &CudaSlice<u8>,
21897        wk: &CudaSlice<u8>,
21898        wv: &CudaSlice<u8>,
21899        aq: &CudaSlice<i8>,
21900        ad: &CudaSlice<f32>,
21901        yq: &mut CudaSlice<f32>,
21902        yk: &mut CudaSlice<f32>,
21903        yv: &mut CudaSlice<f32>,
21904        in_f: usize,
21905        out_q: usize,
21906        out_kv: usize,
21907        t: usize,
21908    ) -> Result<(), Box<dyn std::error::Error>> {
21909        const ROWS_PER_BLOCK: u32 = 4;
21910        let rows = out_q + 2 * out_kv;
21911        let nblk = in_f / 32;
21912        if !in_f.is_multiple_of(32)
21913            || t == 0
21914            || aq.len() < t * in_f
21915            || ad.len() < t * nblk
21916            || yq.len() < t * out_q
21917            || yk.len() < t * out_kv
21918            || yv.len() < t * out_kv
21919        {
21920            return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
21921        }
21922        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
21923        // MEMRA_Q8T_WONCE=1: weight-once twin — see qmatvec.cu's `_tw` block for why the `_t`
21924        // form re-streams the fully-shared QKV weights per column (__ldcs + column grid axis;
21925        // measured 1.67x a single-column call for 2 columns).
21926        if q8t_wonce_on() && t <= 32 {
21927            let f = self.func(if t <= 8 {
21928                "qmatvec_q8_0_qkv_rp_tw"
21929            } else {
21930                "qmatvec_q8_0_qkv_rp_tw32"
21931            });
21932            let cfg = LaunchConfig {
21933                grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21934                block_dim: (32, ROWS_PER_BLOCK, 1),
21935                shared_mem_bytes: 0,
21936            };
21937            let ti = t as i32;
21938            let __s_b = self.gpu.stream();
21939            let mut b = __s_b.launch_builder(&f);
21940            b.arg(wq)
21941                .arg(wk)
21942                .arg(wv)
21943                .arg(aq)
21944                .arg(ad)
21945                .arg(yq)
21946                .arg(yk)
21947                .arg(yv)
21948                .arg(&ini)
21949                .arg(&oq)
21950                .arg(&okv)
21951                .arg(&ti);
21952            unsafe {
21953                b.launch(cfg)?;
21954            }
21955            return Ok(());
21956        }
21957        let f = self.func("qmatvec_q8_0_qkv_rp_t");
21958        let cfg = LaunchConfig {
21959            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
21960            block_dim: (32, ROWS_PER_BLOCK, 1),
21961            shared_mem_bytes: 0,
21962        };
21963        let __s_b = self.gpu.stream();
21964        let mut b = __s_b.launch_builder(&f);
21965        b.arg(wq)
21966            .arg(wk)
21967            .arg(wv)
21968            .arg(aq)
21969            .arg(ad)
21970            .arg(yq)
21971            .arg(yk)
21972            .arg(yv)
21973            .arg(&ini)
21974            .arg(&oq)
21975            .arg(&okv);
21976        unsafe {
21977            b.launch(cfg)?;
21978        }
21979        Ok(())
21980    }
21981
21982    /// T-column q8_0 o_proj over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8, verify walk).
21983    /// Bit-identical to `t` separate `qmatvec_q8_0_b4_rp` calls.
21984    #[allow(clippy::too_many_arguments)]
21985    pub fn qmatvec_q8_0_b4_rp_t_into(
21986        &self,
21987        w: [&CudaSlice<u8>; 4],
21988        aq: &CudaSlice<i8>,
21989        ad: &CudaSlice<f32>,
21990        y: &mut CudaSlice<f32>,
21991        block_cols: usize,
21992        out_f: usize,
21993        t: usize,
21994    ) -> Result<(), Box<dyn std::error::Error>> {
21995        const ROWS_PER_BLOCK: u32 = 4;
21996        let nblk = block_cols / 32;
21997        if !block_cols.is_multiple_of(32)
21998            || t == 0
21999            || aq.len() < t * 4 * block_cols
22000            || ad.len() < t * 4 * nblk
22001            || y.len() < t * out_f
22002        {
22003            return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
22004        }
22005        let (bc, of) = (block_cols as i32, out_f as i32);
22006        // MEMRA_Q8T_WONCE=1: weight-once twin (see qmatvec.cu; `_t` measured 1.43x for 2 columns
22007        // on fully-shared o_proj weights).
22008        if q8t_wonce_on() && t <= 32 {
22009            let f = self.func(if t <= 8 {
22010                "qmatvec_q8_0_b4_rp_tw"
22011            } else {
22012                "qmatvec_q8_0_b4_rp_tw32"
22013            });
22014            let cfg = LaunchConfig {
22015                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
22016                block_dim: (32, ROWS_PER_BLOCK, 1),
22017                shared_mem_bytes: 0,
22018            };
22019            let ti = t as i32;
22020            let __s_b = self.gpu.stream();
22021            let mut b = __s_b.launch_builder(&f);
22022            b.arg(w[0])
22023                .arg(w[1])
22024                .arg(w[2])
22025                .arg(w[3])
22026                .arg(aq)
22027                .arg(ad)
22028                .arg(y)
22029                .arg(&bc)
22030                .arg(&of)
22031                .arg(&ti);
22032            unsafe {
22033                b.launch(cfg)?;
22034            }
22035            return Ok(());
22036        }
22037        let f = self.func("qmatvec_q8_0_b4_rp_t");
22038        let cfg = LaunchConfig {
22039            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
22040            block_dim: (32, ROWS_PER_BLOCK, 1),
22041            shared_mem_bytes: 0,
22042        };
22043        let __s_b = self.gpu.stream();
22044        let mut b = __s_b.launch_builder(&f);
22045        b.arg(w[0])
22046            .arg(w[1])
22047            .arg(w[2])
22048            .arg(w[3])
22049            .arg(aq)
22050            .arg(ad)
22051            .arg(y)
22052            .arg(&bc)
22053            .arg(&of);
22054        unsafe {
22055            b.launch(cfg)?;
22056        }
22057        Ok(())
22058    }
22059
22060    /// MEMRA_W8_VIEW: the q8_0 mirror for a bf16 GEMV whose weight is a ROW-RANGE VIEW.
22061    /// `MEMRA_W8_HYBRID` hangs off `matvec_bf16_into`, and the two split decode paths pinned in
22062    /// the step37 serving env send only their HI half there: HEAD_SPLIT runs
22063    /// `rank1.matvec_bf16_into(head_hi)` beside `e.matvec_bf16_view_into(head_lo)`, and
22064    /// SHEXP_OVERLAP does the same with the shared-expert down rows. The view launcher had no
22065    /// mirror, so the lo half kept streaming 2 B/w while its twin ran at 1.0625, and because the
22066    /// halves execute CONCURRENTLY on the two cards the critical path is the SLOW half.
22067    /// NUMERIC CLASS: identical to the rest of `MEMRA_STEP_TP_W8`, so it carries that argmax
22068    /// acceptance and that maxdiff class, not a new one. Default OFF until measured.
22069    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
22070    fn matvec_bf16_view_via_q8_mirror(
22071        &self,
22072        data: &cudarc::driver::CudaView<'_, u8>,
22073        x: &CudaSlice<f32>,
22074        y: &mut CudaSlice<f32>,
22075        in_f: usize,
22076        out_f: usize,
22077    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
22078        use cudarc::driver::DevicePtr;
22079        let key = {
22080            let s = self.gpu.stream();
22081            let (p, _g) = data.device_ptr(&s);
22082            (p, in_f as u32, out_f as u32)
22083        };
22084        {
22085            let mut mirrors = self
22086                .w8_mirrors
22087                .lock()
22088                .map_err(|_| "w8 mirror map is poisoned")?;
22089            if !mirrors.contains_key(&key) {
22090                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
22091                self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
22092                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
22093                mirrors.insert(key, planar);
22094                // Unconditional, once per distinct shape: a door with no announce cannot be read
22095                // in BOTH directions, and this lane was already burned once by a sweep that
22096                // inferred "never engages" from a log line that did not exist in the tree.
22097                eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
22098            }
22099        }
22100        let nblk = in_f / 32;
22101        {
22102            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
22103            if !act.contains_key(&in_f) {
22104                let aq = self.alloc_uninit::<i8>(in_f)?;
22105                let ad = self.alloc_uninit::<f32>(nblk)?;
22106                act.insert(in_f, (aq, ad));
22107            }
22108            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
22109            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
22110        }
22111        let mirrors = self
22112            .w8_mirrors
22113            .lock()
22114            .map_err(|_| "w8 mirror map is poisoned")?;
22115        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
22116        let mirror = mirrors.get(&key).expect("built above");
22117        let (aq, ad) = act.get(&in_f).expect("built above");
22118        self.qmatvec_mmvq_into(
22119            mirror,
22120            aq,
22121            ad,
22122            1,
22123            in_f,
22124            out_f,
22125            QT_Q8_0,
22126            Self::q8_0_row_bytes(in_f),
22127            1.0,
22128            true,
22129            y,
22130        )?;
22131        Ok(Some(()))
22132    }
22133
22134    pub fn matvec_bf16_into(
22135        &self,
22136        data: &CudaSlice<u8>,
22137        x: &CudaSlice<f32>,
22138        y: &mut CudaSlice<f32>,
22139        in_f: usize,
22140        out_f: usize,
22141    ) -> Result<(), Box<dyn std::error::Error>> {
22142        if data.len() != in_f * out_f * 2
22143            || x.len() < in_f
22144            || !in_f.is_multiple_of(8)
22145            || y.len() < out_f
22146        {
22147            return Err(format!(
22148                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22149                data.len(),
22150                x.len(),
22151                y.len()
22152            )
22153            .into());
22154        }
22155        // MEMRA_STEP_TP_W8, hybrid half: route this GEMV through a q8_0 mirror of the same
22156        // weight. Covers exactly the bf16 GEMVs that are NOT in a TP resident bank — the LM
22157        // head (324.4 -> 163.7 us measured), the shared-expert down rows (13.0 -> 5.6 us) and
22158        // the dense-FFN layers. Same numeric class as the QKV/o_proj arms (int8 dp4a with
22159        // per-32 scales), so it rides the same argmax acceptance; the bf16 slab stays resident
22160        // for prefill. The mirror builds on first use and is keyed by the slab's pointer.
22161        if step_tp_w8_on()
22162            && w8_hybrid_on()
22163            && in_f.is_multiple_of(32)
22164            && out_f >= 64
22165            && let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)?
22166        {
22167            return Ok(());
22168        }
22169        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
22170        // block, exact f32acc per-row program — cures the 1-iteration latency
22171        // starvation (shexp down measured 420GB/s at in_f=1280).
22172        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22173        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
22174            && in_f <= 2048;
22175        if x4 {
22176            let f = self.func("matvec_bf16_f32acc_x4");
22177            let cfg = LaunchConfig {
22178                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
22179                block_dim: (mmv_block(), 1, 1),
22180                shared_mem_bytes: 0,
22181            };
22182            let (ini, outi) = (in_f as i32, out_f as i32);
22183            let __s_b = self.gpu.stream();
22184            let mut b = __s_b.launch_builder(&f);
22185            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
22186            unsafe {
22187                b.launch(cfg)?;
22188            }
22189            return Ok(());
22190        }
22191        let f = self.func("matvec_bf16_f32acc");
22192        let cfg = LaunchConfig {
22193            grid_dim: (out_f as u32, 1, 1),
22194            block_dim: (mmv_block(), 1, 1),
22195            shared_mem_bytes: 0,
22196        };
22197        let ini = in_f as i32;
22198        let __s_b = self.gpu.stream();
22199        let mut b = __s_b.launch_builder(&f);
22200        b.arg(data).arg(x).arg(y).arg(&ini);
22201        unsafe {
22202            b.launch(cfg)?;
22203        }
22204        Ok(())
22205    }
22206
22207    /// BF16 matvec over activation/output views. Automatic TP4 attention keeps each rank's
22208    /// O-projection input and canonical partial inside persistent slabs, so copying either view
22209    /// into a temporary allocation would give back the bandwidth and allocator win that TP is
22210    /// meant to provide.
22211    pub fn matvec_bf16_views_into(
22212        &self,
22213        data: &CudaSlice<u8>,
22214        x: &cudarc::driver::CudaView<'_, f32>,
22215        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
22216        in_f: usize,
22217        out_f: usize,
22218    ) -> Result<(), Box<dyn std::error::Error>> {
22219        if data.len() != in_f * out_f * 2
22220            || x.len() < in_f
22221            || !in_f.is_multiple_of(8)
22222            || y.len() < out_f
22223        {
22224            return Err(format!(
22225                "matvec_bf16_views_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22226                data.len(),
22227                x.len(),
22228                y.len()
22229            )
22230            .into());
22231        }
22232        let f = self.func("matvec_bf16_f32acc");
22233        let cfg = LaunchConfig {
22234            grid_dim: (out_f as u32, 1, 1),
22235            block_dim: (mmv_block(), 1, 1),
22236            shared_mem_bytes: 0,
22237        };
22238        let ini = in_f as i32;
22239        let __s_b = self.gpu.stream();
22240        let mut b = __s_b.launch_builder(&f);
22241        b.arg(data).arg(x).arg(y).arg(&ini);
22242        unsafe {
22243            b.launch(cfg)?;
22244        }
22245        Ok(())
22246    }
22247
22248    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
22249    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
22250    pub fn matvec_bf16_view_into(
22251        &self,
22252        data: &cudarc::driver::CudaView<'_, u8>,
22253        x: &CudaSlice<f32>,
22254        y: &mut CudaSlice<f32>,
22255        in_f: usize,
22256        out_f: usize,
22257    ) -> Result<(), Box<dyn std::error::Error>> {
22258        if data.len() != in_f * out_f * 2
22259            || x.len() < in_f
22260            || !in_f.is_multiple_of(8)
22261            || y.len() < out_f
22262        {
22263            return Err(format!(
22264                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22265                data.len(),
22266                x.len(),
22267                y.len()
22268            )
22269            .into());
22270        }
22271        if w8_view_on()
22272            && step_tp_w8_on()
22273            && w8_hybrid_on()
22274            && in_f.is_multiple_of(32)
22275            && out_f >= 64
22276            && let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)?
22277        {
22278            return Ok(());
22279        }
22280        let f = self.func("matvec_bf16_f32acc");
22281        let cfg = LaunchConfig {
22282            grid_dim: (out_f as u32, 1, 1),
22283            block_dim: (mmv_block(), 1, 1),
22284            shared_mem_bytes: 0,
22285        };
22286        let ini = in_f as i32;
22287        let __s_b = self.gpu.stream();
22288        let mut b = __s_b.launch_builder(&f);
22289        b.arg(data).arg(x).arg(y).arg(&ini);
22290        unsafe {
22291            b.launch(cfg)?;
22292        }
22293        Ok(())
22294    }
22295
22296    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
22297    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
22298    pub fn matvec_bf16_raw_out(
22299        &self,
22300        w: &CudaSlice<u8>,
22301        x: &CudaSlice<f32>,
22302        y_raw: u64,
22303        in_f: usize,
22304        out_f: usize,
22305    ) -> Result<(), Box<dyn std::error::Error>> {
22306        if w.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) || y_raw == 0 {
22307            return Err("matvec_bf16_raw_out geometry".into());
22308        }
22309        let f = self.func("matvec_bf16_f32acc");
22310        let cfg = LaunchConfig {
22311            grid_dim: (out_f as u32, 1, 1),
22312            block_dim: (mmv_block(), 1, 1),
22313            shared_mem_bytes: 0,
22314        };
22315        let ini = in_f as i32;
22316        let __s_b = self.gpu.stream();
22317        let mut b = __s_b.launch_builder(&f);
22318        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
22319        unsafe {
22320            b.launch(cfg)?;
22321        }
22322        Ok(())
22323    }
22324
22325    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
22326    /// UVA pointers so the caller passes persistent-static rows without holding locks).
22327    /// Exact per-element sequence of the split add + add_scaled_rows pair.
22328    pub fn add3_raw(
22329        &self,
22330        a: &CudaSlice<f32>,
22331        b: &CudaSlice<f32>,
22332        sh_raw: u64,
22333        scale_raw: u64,
22334        dst: &mut CudaSlice<f32>,
22335        n: usize,
22336    ) -> Result<(), Box<dyn std::error::Error>> {
22337        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
22338            return Err("add3_raw geometry".into());
22339        }
22340        let f = self.func("add3_f32");
22341        let cfg = LaunchConfig {
22342            grid_dim: ((n as u32).div_ceil(256), 1, 1),
22343            block_dim: (256, 1, 1),
22344            shared_mem_bytes: 0,
22345        };
22346        let ni = n as i32;
22347        let __s_b = self.gpu.stream();
22348        let mut bld = __s_b.launch_builder(&f);
22349        bld.arg(a)
22350            .arg(b)
22351            .arg(&sh_raw)
22352            .arg(&scale_raw)
22353            .arg(dst)
22354            .arg(&ni);
22355        unsafe {
22356            bld.launch(cfg)?;
22357        }
22358        Ok(())
22359    }
22360
22361    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
22362    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
22363    pub fn matvec_bf16_down_addscale_into(
22364        &self,
22365        w: &CudaSlice<u8>,
22366        x: &CudaSlice<f32>,
22367        scale: &CudaSlice<f32>,
22368        dst: &mut CudaSlice<f32>,
22369        in_f: usize,
22370        out_f: usize,
22371    ) -> Result<(), Box<dyn std::error::Error>> {
22372        if w.len() != in_f * out_f * 2
22373            || x.len() < in_f
22374            || !in_f.is_multiple_of(8)
22375            || dst.len() < out_f
22376            || scale.is_empty()
22377        {
22378            return Err("matvec_bf16_down_addscale geometry".into());
22379        }
22380        let f = self.func("matvec_bf16_down_addscale");
22381        let cfg = LaunchConfig {
22382            grid_dim: (out_f as u32, 1, 1),
22383            block_dim: (mmv_block(), 1, 1),
22384            shared_mem_bytes: 0,
22385        };
22386        let ini = in_f as i32;
22387        let __s_b = self.gpu.stream();
22388        let mut b = __s_b.launch_builder(&f);
22389        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
22390        unsafe {
22391            b.launch(cfg)?;
22392        }
22393        Ok(())
22394    }
22395
22396    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
22397    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
22398    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
22399    #[allow(clippy::too_many_arguments)]
22400    pub fn matvec_bf16_dual_silu_rows_into(
22401        &self,
22402        wg: &CudaSlice<u8>,
22403        wu: &CudaSlice<u8>,
22404        x: &CudaSlice<f32>,
22405        act: &mut CudaSlice<f32>,
22406        in_f: usize,
22407        out_f: usize,
22408        limit: Option<f32>,
22409        t: usize,
22410    ) -> Result<(), Box<dyn std::error::Error>> {
22411        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
22412            return Err("matvec_bf16_dual_silu_rows geometry".into());
22413        }
22414        let f = self.func("matvec_bf16_dual_silu_rows");
22415        let cfg = LaunchConfig {
22416            grid_dim: (out_f as u32, t as u32, 1),
22417            block_dim: (mmv_block(), 1, 1),
22418            shared_mem_bytes: 0,
22419        };
22420        let (ini, outi) = (in_f as i32, out_f as i32);
22421        let lim = limit.unwrap_or(0.0);
22422        let __s_b = self.gpu.stream();
22423        let mut b = __s_b.launch_builder(&f);
22424        b.arg(wg)
22425            .arg(wu)
22426            .arg(x)
22427            .arg(&mut *act)
22428            .arg(&ini)
22429            .arg(&outi)
22430            .arg(&lim);
22431        unsafe {
22432            b.launch(cfg)?;
22433        }
22434        Ok(())
22435    }
22436
22437    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
22438    pub fn matvec_bf16_rows_into(
22439        &self,
22440        w: &CudaSlice<u8>,
22441        x: &CudaSlice<f32>,
22442        y: &mut CudaSlice<f32>,
22443        in_f: usize,
22444        out_f: usize,
22445        t: usize,
22446    ) -> Result<(), Box<dyn std::error::Error>> {
22447        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || !in_f.is_multiple_of(8)
22448        {
22449            return Err("matvec_bf16_rows geometry".into());
22450        }
22451        // MEMRA_STEP_TP_W8 + MEMRA_W8_HYBRID, t > 1: the VERIFY walk's shexp/dense rows land
22452        // here too (`matvec_bf16_f32acc_x4_rows` was 78 launches/round at 56.5 us in a spec
22453        // capture, ~162 ms of GPU over 37 rounds), and the t==1 gate below skipped them. The
22454        // t-column q8 kernel is bit-identical to t single-row calls.
22455        if (2..=32).contains(&t)
22456            && step_tp_w8_on()
22457            && w8_hybrid_on()
22458            && in_f.is_multiple_of(32)
22459            && out_f >= 64
22460            && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
22461        {
22462            return Ok(());
22463        }
22464        // MEMRA_STEP_TP_W8: the LM head reaches the device HERE, not through
22465        // matvec_bf16_into — the W8 trace showed the hybrid half building mirrors only for
22466        // in_f=1280 out_f=4096 (the shared-expert down rows, which SHEXP_OVERLAP already
22467        // hides, hence its +0.1%). Route the t=1 decode row through the q8 mirror; wider t
22468        // (the verify walk) keeps bf16 so the prefill class is untouched.
22469        if t == 1
22470            && step_tp_w8_on()
22471            && w8_hybrid_on()
22472            && in_f.is_multiple_of(32)
22473            && out_f >= 64
22474            && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
22475        {
22476            return Ok(());
22477        }
22478        // MEMRA_BF16_TCOLS_WIDE (lane/glm5-matvec door T, default ON since 2026-08-31): t=2..=16 rides the
22479        // weight-once t-column class instead of the grid.y=t per-token weight re-read below.
22480        // Placed AFTER the W8-mirror intercepts (their precedence unchanged). Bit-identical
22481        // per (row, token) to the _rows kernel by the tcols class's standing construction
22482        // (order-pinned per-token chains + the identical red[256] tree); the motivating call
22483        // is the DFlash2 drafter's t=15 block-head matmul, which re-read the 1.269 GB lm
22484        // head 15x per spec round. Rollback seam: unset or =0 falls through unchanged.
22485        if (2..=16).contains(&t) && bf16_tcols_wide_on() {
22486            if BF16_TCOLS_WIDE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
22487                eprintln!(
22488                    "[bf16-tcols-wide] engaged: t={t} in_f={in_f} out_f={out_f} rides the \
22489                     weight-once tcols class (MEMRA_BF16_TCOLS_WIDE=1)"
22490                );
22491            }
22492            if t <= 8 {
22493                return self.matvec_bf16_tcols_into(w, x, y, in_f, out_f, t);
22494            }
22495            return self.matvec_bf16_tcols16_into(w, x, y, in_f, out_f, t);
22496        }
22497        let f = self.func("matvec_bf16_f32acc_x4_rows");
22498        let cfg = LaunchConfig {
22499            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
22500            block_dim: (mmv_block(), 1, 1),
22501            shared_mem_bytes: 0,
22502        };
22503        let (ini, outi) = (in_f as i32, out_f as i32);
22504        let __s_b = self.gpu.stream();
22505        let mut b = __s_b.launch_builder(&f);
22506        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
22507        unsafe {
22508            b.launch(cfg)?;
22509        }
22510        Ok(())
22511    }
22512
22513    /// One aligned K-range partial of the t=1 BF16 row matvec. This is the row-parallel TP
22514    /// building block: every rank computes all output rows over a disjoint K range from its
22515    /// compact local activation shard, then the persistent replicated-row collective sums the
22516    /// partials. The kernel preserves the unsplit program inside each range; the cross-rank
22517    /// association is separately gated.
22518    #[allow(clippy::too_many_arguments)]
22519    pub fn matvec_bf16_col_range_into(
22520        &self,
22521        w: &CudaSlice<u8>,
22522        x: &CudaSlice<f32>,
22523        y: &mut CudaSlice<f32>,
22524        in_f: usize,
22525        out_f: usize,
22526        k_start: usize,
22527        k_len: usize,
22528    ) -> Result<(), Box<dyn std::error::Error>> {
22529        let weight_bytes = in_f
22530            .checked_mul(out_f)
22531            .and_then(|elements| elements.checked_mul(2))
22532            .ok_or("BF16 column-range matvec geometry")?;
22533        let k_end = k_start
22534            .checked_add(k_len)
22535            .ok_or("BF16 column-range matvec geometry")?;
22536        if w.len() < weight_bytes
22537            || x.len() < k_len
22538            || y.len() < out_f
22539            || in_f > i32::MAX as usize
22540            || out_f > i32::MAX as usize
22541            || k_start > i32::MAX as usize
22542            || k_len > i32::MAX as usize
22543            || in_f == 0
22544            || out_f == 0
22545            || k_len == 0
22546            || !in_f.is_multiple_of(8)
22547            || !k_start.is_multiple_of(8)
22548            || !k_len.is_multiple_of(8)
22549            || k_end > in_f
22550        {
22551            return Err("BF16 column-range matvec geometry".into());
22552        }
22553        let function = self.func("matvec_bf16_f32acc_x4_range");
22554        let config = LaunchConfig {
22555            grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
22556            block_dim: (mmv_block(), 1, 1),
22557            shared_mem_bytes: 0,
22558        };
22559        let (in_f, out_f, k_start, k_len) =
22560            (in_f as i32, out_f as i32, k_start as i32, k_len as i32);
22561        let stream = self.gpu.stream();
22562        let mut launch = stream.launch_builder(&function);
22563        launch
22564            .arg(w)
22565            .arg(x)
22566            .arg(y)
22567            .arg(&in_f)
22568            .arg(&out_f)
22569            .arg(&k_start)
22570            .arg(&k_len);
22571        unsafe {
22572            launch.launch(config)?;
22573        }
22574        Ok(())
22575    }
22576
22577    /// T-COLUMN twin of the bf16 rows matvec (lane/glm5-verify-batch, the
22578    /// varlen-batched-cores pattern): one block owns 4 output rows for ALL t tokens, so the
22579    /// weight pack is read ONCE and reused across tokens — vs the `_rows` twin's grid.y=t
22580    /// per-token weight re-read. Per-(row,token) BIT-IDENTICAL to the t=1 program by
22581    /// construction (order-pinned single-chain accumulators, identical shared-tree reduce
22582    /// per token — LAW:vl-bit-identity-order-pinning); the `glm5_verify_batch_gpu` tcols
22583    /// bit-gate holds it. t is bounded by the kernel's MEMRA_BF16_TCOLS_MAX = 8.
22584    pub fn matvec_bf16_tcols_into(
22585        &self,
22586        w: &CudaSlice<u8>,
22587        x: &CudaSlice<f32>,
22588        y: &mut CudaSlice<f32>,
22589        in_f: usize,
22590        out_f: usize,
22591        t: usize,
22592    ) -> Result<(), Box<dyn std::error::Error>> {
22593        if x.len() < t * in_f
22594            || y.len() < t * out_f
22595            || !(2..=8).contains(&t)
22596            || !in_f.is_multiple_of(8)
22597        {
22598            return Err("matvec_bf16_tcols geometry".into());
22599        }
22600        // MEMRA_BF16_TCOLS_X1 (lane/glm5-matvec door X, default ON since 2026-08-31): one row per block
22601        // (grid.x = out_f) — 4x the wave count on the ~one-wave trunk grids (census: same
22602        // kernel runs 59% of peak at 512..2048 blocks, 80% at 38720). Per-row body and
22603        // reduce tree verbatim — bit-identical per (row, token). Rollback: unset or =0.
22604        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the chosen grid
22605        // form takes its `_rf` fused-reduce-tail twin — one barrier sequence shared by the t
22606        // columns plus intra-warp shuffles at the identical pairing (9t -> 3 barriers per
22607        // block). Composes with door X (grid choice first, tail twin second). Requires a
22608        // power-of-two block (the fused tail must pass exactly through s=32); any other
22609        // MEMRA_MMV_BLOCK falls through to the standing tree. Rollback: unset or =0.
22610        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
22611        if rf
22612            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
22613                == 0
22614        {
22615            eprintln!(
22616                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
22617                 shared across the t token columns + intra-warp shuffles at the identical \
22618                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
22619            );
22620        }
22621        let x1 = bf16_tcols_x1_on();
22622        let (fname, grid_x) = match (x1, rf) {
22623            (true, true) => ("matvec_bf16_f32acc_x1_tcols_rf", out_f as u32),
22624            (true, false) => ("matvec_bf16_f32acc_x1_tcols", out_f as u32),
22625            (false, true) => ("matvec_bf16_f32acc_x4_tcols_rf", out_f.div_ceil(4) as u32),
22626            (false, false) => ("matvec_bf16_f32acc_x4_tcols", out_f.div_ceil(4) as u32),
22627        };
22628        if x1 && BF16_TCOLS_X1_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
22629            eprintln!(
22630                "[bf16-tcols-x1] engaged: one-row-per-block tcols grid \
22631                 (MEMRA_BF16_TCOLS_X1=1)"
22632            );
22633        }
22634        let f = self.func(fname);
22635        let cfg = LaunchConfig {
22636            grid_dim: (grid_x, 1, 1),
22637            block_dim: (mmv_block(), 1, 1),
22638            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
22639        };
22640        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22641        let __s_b = self.gpu.stream();
22642        let mut b = __s_b.launch_builder(&f);
22643        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22644        unsafe {
22645            b.launch(cfg)?;
22646        }
22647        Ok(())
22648    }
22649
22650    /// WIDE-T twin of [`Self::matvec_bf16_tcols_into`] (lane/glm5-matvec door T,
22651    /// `MEMRA_BF16_TCOLS_WIDE`): t = 9..=16 through the SEPARATE `..._tcols16` kernel — its
22652    /// acc[16] register footprint never touches the priced t<=8 class (the qmatvec `_tw32`
22653    /// acc-sizing lesson). Bit-identical per (row, token) to the t=1 program by the same
22654    /// order-pinned construction; gated by `glm5_matvec_doors_gpu`.
22655    pub fn matvec_bf16_tcols16_into(
22656        &self,
22657        w: &CudaSlice<u8>,
22658        x: &CudaSlice<f32>,
22659        y: &mut CudaSlice<f32>,
22660        in_f: usize,
22661        out_f: usize,
22662        t: usize,
22663    ) -> Result<(), Box<dyn std::error::Error>> {
22664        if x.len() < t * in_f
22665            || y.len() < t * out_f
22666            || !(9..=16).contains(&t)
22667            || !in_f.is_multiple_of(8)
22668        {
22669            return Err("matvec_bf16_tcols16 geometry".into());
22670        }
22671        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the wide-t twin
22672        // takes its `_rf` fused tail too — the drafter head's t=15 is the extreme case (135
22673        // barriers -> 6 per block). Same power-of-two block guard as the t<=8 dispatch.
22674        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
22675        if rf
22676            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
22677                == 0
22678        {
22679            eprintln!(
22680                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
22681                 shared across the t token columns + intra-warp shuffles at the identical \
22682                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
22683            );
22684        }
22685        let f = self.func(if rf {
22686            "matvec_bf16_f32acc_x4_tcols16_rf"
22687        } else {
22688            "matvec_bf16_f32acc_x4_tcols16"
22689        });
22690        let cfg = LaunchConfig {
22691            grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
22692            block_dim: (mmv_block(), 1, 1),
22693            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
22694        };
22695        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22696        let __s_b = self.gpu.stream();
22697        let mut b = __s_b.launch_builder(&f);
22698        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22699        unsafe {
22700            b.launch(cfg)?;
22701        }
22702        Ok(())
22703    }
22704
22705    /// GATE-ONLY launcher for door R arms no route dispatches (`glm5_matvec_doors_gpu`):
22706    /// the shifted-pairing RED twin (`matvec_bf16_f32acc_x1_tcols_rf_redshift`, the arm that
22707    /// proves the bit bar can see an association change) and the `_rf` twins at t=1 (the
22708    /// routed launchers refuse t<2; the door-R bar covers t=1..=16, so the degenerate
22709    /// column-loop bounds are gated here). `kernel` is an ALLOWLIST, not a name proxy.
22710    #[allow(clippy::too_many_arguments)]
22711    pub fn matvec_bf16_tcols_gate_kernel_into(
22712        &self,
22713        kernel: &str,
22714        w: &CudaSlice<u8>,
22715        x: &CudaSlice<f32>,
22716        y: &mut CudaSlice<f32>,
22717        in_f: usize,
22718        out_f: usize,
22719        t: usize,
22720    ) -> Result<(), Box<dyn std::error::Error>> {
22721        let (grid_x, t_max) = match kernel {
22722            "matvec_bf16_f32acc_x1_tcols_rf" | "matvec_bf16_f32acc_x1_tcols_rf_redshift" => {
22723                (out_f as u32, 8usize)
22724            }
22725            "matvec_bf16_f32acc_x4_tcols_rf" => (out_f.div_ceil(4) as u32, 8usize),
22726            "matvec_bf16_f32acc_x4_tcols16_rf" => (out_f.div_ceil(4) as u32, 16usize),
22727            _ => return Err("matvec_bf16_tcols_gate_kernel_into: unknown kernel".into()),
22728        };
22729        if x.len() < t * in_f
22730            || y.len() < t * out_f
22731            || !(1..=t_max).contains(&t)
22732            || !in_f.is_multiple_of(8)
22733            || !mmv_block().is_power_of_two()
22734        {
22735            return Err("matvec_bf16_tcols_gate_kernel geometry".into());
22736        }
22737        let f = self.func(kernel);
22738        let cfg = LaunchConfig {
22739            grid_dim: (grid_x, 1, 1),
22740            block_dim: (mmv_block(), 1, 1),
22741            shared_mem_bytes: (t as u32) * mmv_block() * 4,
22742        };
22743        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22744        let __s_b = self.gpu.stream();
22745        let mut b = __s_b.launch_builder(&f);
22746        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22747        unsafe {
22748            b.launch(cfg)?;
22749        }
22750        Ok(())
22751    }
22752
22753    /// DECODE-EXACT matmul for the glm5 verify-batch walk (lane/glm5-verify-batch): the
22754    /// exact `matmul_decode_exact` dispatch with ONE addition — FloatBf16 weights at
22755    /// t=2..=8 under `MEMRA_BF16_MMV` ride the tcols twin above (weight read once for all
22756    /// t rows). Refused back to `matmul_decode_exact` whenever the t=1 decode chain would
22757    /// ride the W8 q8-mirror class instead of the bf16 rows kernel (the decode-parity law:
22758    /// the m>1 class must equal the m=1 class). Every class stays per-row bit-exact vs
22759    /// the t=1 chain; only the verify-batch walk calls this.
22760    pub fn matmul_rows_exact(
22761        &self,
22762        w: &crate::model::GpuTensor,
22763        x: &CudaSlice<f32>,
22764        m: usize,
22765    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22766        use crate::model::GpuTensor;
22767        if let GpuTensor::FloatBf16 { data, .. } = w
22768            && (2..=8).contains(&m)
22769            && Self::bf16_mmv_on()
22770            && w.in_features().is_multiple_of(8)
22771            && !(step_tp_w8_on() && w8_hybrid_on())
22772        {
22773            let (in_f, out_f) = (w.in_features(), w.out_features());
22774            // Door W: rows-exact is verify-walk-only by contract, so its y is a pooled
22775            // draw (vws_uninit == alloc_uninit with the door off).
22776            let mut y = self.vws_uninit(m * out_f)?;
22777            self.matvec_bf16_tcols_into(data, x, &mut y, in_f, out_f, m)?;
22778            return Ok(y);
22779        }
22780        self.matmul_decode_exact(w, x, m)
22781    }
22782
22783    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22784    pub fn matvec_bf16_dual_silu_into(
22785        &self,
22786        wg: &CudaSlice<u8>,
22787        wu: &CudaSlice<u8>,
22788        x: &CudaSlice<f32>,
22789        act: &mut CudaSlice<f32>,
22790        in_f: usize,
22791        out_f: usize,
22792        limit: Option<f32>,
22793    ) -> Result<(), Box<dyn std::error::Error>> {
22794        if wg.len() != in_f * out_f * 2
22795            || wu.len() != in_f * out_f * 2
22796            || x.len() < in_f
22797            || !in_f.is_multiple_of(8)
22798            || act.len() < out_f
22799        {
22800            return Err("matvec_bf16_dual_silu geometry".into());
22801        }
22802        let f = self.func("matvec_bf16_dual_silu");
22803        let cfg = LaunchConfig {
22804            grid_dim: (out_f as u32, 1, 1),
22805            block_dim: (mmv_block(), 1, 1),
22806            shared_mem_bytes: 0,
22807        };
22808        let (ini, outi) = (in_f as i32, out_f as i32);
22809        let lim = limit.unwrap_or(0.0);
22810        let __s_b = self.gpu.stream();
22811        let mut b = __s_b.launch_builder(&f);
22812        b.arg(wg)
22813            .arg(wu)
22814            .arg(x)
22815            .arg(act)
22816            .arg(&ini)
22817            .arg(&outi)
22818            .arg(&lim);
22819        unsafe {
22820            b.launch(cfg)?;
22821        }
22822        Ok(())
22823    }
22824
22825    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
22826    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
22827    #[allow(clippy::too_many_arguments)]
22828    pub fn matvec_bf16_dual_view_into(
22829        &self,
22830        wg: &cudarc::driver::CudaView<'_, u8>,
22831        wu: &cudarc::driver::CudaView<'_, u8>,
22832        x: &CudaSlice<f32>,
22833        yg: &mut CudaSlice<f32>,
22834        yu: &mut CudaSlice<f32>,
22835        in_f: usize,
22836        out_f: usize,
22837    ) -> Result<(), Box<dyn std::error::Error>> {
22838        if wg.len() != in_f * out_f * 2
22839            || wu.len() != in_f * out_f * 2
22840            || x.len() < in_f
22841            || !in_f.is_multiple_of(8)
22842            || yg.len() < out_f
22843            || yu.len() < out_f
22844        {
22845            return Err(format!(
22846                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
22847                wg.len(),
22848                wu.len(),
22849                x.len()
22850            )
22851            .into());
22852        }
22853        let f = self.func("matvec_bf16_dual");
22854        let cfg = LaunchConfig {
22855            grid_dim: ((2 * out_f) as u32, 1, 1),
22856            block_dim: (mmv_block(), 1, 1),
22857            shared_mem_bytes: 0,
22858        };
22859        let (ini, outi) = (in_f as i32, out_f as i32);
22860        let __s_b = self.gpu.stream();
22861        let mut b = __s_b.launch_builder(&f);
22862        b.arg(wg)
22863            .arg(wu)
22864            .arg(x)
22865            .arg(yg)
22866            .arg(yu)
22867            .arg(&ini)
22868            .arg(&outi);
22869        unsafe {
22870            b.launch(cfg)?;
22871        }
22872        Ok(())
22873    }
22874
22875    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
22876    #[allow(clippy::too_many_arguments)]
22877    pub fn matvec_bf16_dual_into(
22878        &self,
22879        wg: &CudaSlice<u8>,
22880        wu: &CudaSlice<u8>,
22881        x: &CudaSlice<f32>,
22882        yg: &mut CudaSlice<f32>,
22883        yu: &mut CudaSlice<f32>,
22884        in_f: usize,
22885        out_f: usize,
22886    ) -> Result<(), Box<dyn std::error::Error>> {
22887        if wg.len() != in_f * out_f * 2
22888            || wu.len() != in_f * out_f * 2
22889            || x.len() < in_f
22890            || !in_f.is_multiple_of(8)
22891            || yg.len() < out_f
22892            || yu.len() < out_f
22893        {
22894            return Err(format!(
22895                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
22896                wg.len(),
22897                wu.len(),
22898                x.len()
22899            )
22900            .into());
22901        }
22902        let f = self.func("matvec_bf16_dual");
22903        let cfg = LaunchConfig {
22904            grid_dim: ((2 * out_f) as u32, 1, 1),
22905            block_dim: (mmv_block(), 1, 1),
22906            shared_mem_bytes: 0,
22907        };
22908        let (ini, outi) = (in_f as i32, out_f as i32);
22909        let __s_b = self.gpu.stream();
22910        let mut b = __s_b.launch_builder(&f);
22911        b.arg(wg)
22912            .arg(wu)
22913            .arg(x)
22914            .arg(yg)
22915            .arg(yu)
22916            .arg(&ini)
22917            .arg(&outi);
22918        unsafe {
22919            b.launch(cfg)?;
22920        }
22921        Ok(())
22922    }
22923
22924    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
22925    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
22926    #[allow(dead_code)] // allow: base form of the matvec_bf16_dual_* family; kept as the reference entry point
22927    pub(crate) fn matvec_bf16_dual(
22928        &self,
22929        wg: &CudaSlice<u8>,
22930        wu: &CudaSlice<u8>,
22931        x: &CudaSlice<f32>,
22932        in_f: usize,
22933        out_f: usize,
22934    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22935        if wg.len() != in_f * out_f * 2
22936            || wu.len() != in_f * out_f * 2
22937            || x.len() < in_f
22938            || !in_f.is_multiple_of(8)
22939        {
22940            return Err(format!(
22941                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
22942                wg.len(),
22943                wu.len(),
22944                x.len()
22945            )
22946            .into());
22947        }
22948        let mut yg = self.alloc_uninit::<f32>(out_f)?;
22949        let mut yu = self.alloc_uninit::<f32>(out_f)?;
22950        let f = self.func("matvec_bf16_dual");
22951        let cfg = LaunchConfig {
22952            grid_dim: ((2 * out_f) as u32, 1, 1),
22953            block_dim: (mmv_block(), 1, 1),
22954            shared_mem_bytes: 0,
22955        };
22956        let (ini, outi) = (in_f as i32, out_f as i32);
22957        let __s_b = self.gpu.stream();
22958        let mut b = __s_b.launch_builder(&f);
22959        b.arg(wg)
22960            .arg(wu)
22961            .arg(x)
22962            .arg(&mut yg)
22963            .arg(&mut yu)
22964            .arg(&ini)
22965            .arg(&outi);
22966        unsafe {
22967            b.launch(cfg)?;
22968        }
22969        Ok((yg, yu))
22970    }
22971
22972    #[allow(clippy::too_many_arguments)]
22973    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
22974    fn linear_bf16_chunked_inner(
22975        &self,
22976        x: &CudaSlice<f32>,
22977        data: &CudaSlice<u8>,
22978        m: usize,
22979        in_f: usize,
22980        out_f: usize,
22981        exact: bool,
22982        canonical_chunk_rows: Option<usize>,
22983    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22984        const CHUNK_BYTES: usize = 256 << 20;
22985        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
22986        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
22987        if m == 1
22988            && !exact
22989            && canonical_chunk_rows.is_none()
22990            && in_f.is_multiple_of(8)
22991            && Self::bf16_mmv_on()
22992        {
22993            return self.matvec_bf16(data, x, in_f, out_f);
22994        }
22995        // MEMRA_PP_BF16: prefill on the RESIDENT bf16 bytes through cuBLASLt tensor cores.
22996        // Below this door the whole weight is dequanted to f32 and multiplied without tensor
22997        // cores — the step37 prime's 14x gap to vLLM. `exact` and canonical-chunk callers are
22998        // numerical programs with their own equality gates and are left alone.
22999        if m >= 16
23000            && !exact
23001            && canonical_chunk_rows.is_none()
23002            && data.len() == in_f * out_f * 2
23003            && crate::f16_ffi::pp_bf16_enabled()
23004        {
23005            // None = cuBLASLt declined this shape (it announced which one); fall through to the
23006            // f32 dequant GEMM below, which is always correct.
23007            if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
23008                return Ok(y);
23009            }
23010        }
23011        let row_bytes = in_f
23012            .checked_mul(std::mem::size_of::<f32>())
23013            .ok_or("BF16 chunk row byte count overflow")?;
23014        if row_bytes == 0 || out_f == 0 {
23015            return Err("BF16 chunk dimensions must be nonzero".into());
23016        }
23017        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
23018        let chunk_rows = match canonical_chunk_rows {
23019            Some(0) => {
23020                return Err("canonical BF16 chunk rows must be nonzero".into());
23021            }
23022            Some(rows) if rows > max_chunk_rows => {
23023                return Err(format!(
23024                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
23025                )
23026                .into());
23027            }
23028            Some(rows) if out_f % rows != 0 => {
23029                return Err(format!(
23030                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
23031                )
23032                .into());
23033            }
23034            Some(rows) => rows,
23035            None => max_chunk_rows,
23036        };
23037        if chunk_rows >= out_f {
23038            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
23039            return if exact {
23040                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
23041            } else {
23042                self.linear(x, &wf32, m, in_f, out_f)
23043            };
23044        }
23045        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
23046        let mut r0 = 0usize;
23047        while r0 < out_f {
23048            let rows = chunk_rows.min(out_f - r0);
23049            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
23050            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
23051            let yc = if exact {
23052                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
23053            } else {
23054                self.linear(x, &wf32, m, in_f, rows)?
23055            };
23056            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
23057            for mi in 0..m {
23058                let src = yc.slice(mi * rows..(mi + 1) * rows);
23059                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
23060                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
23061            }
23062            r0 += rows;
23063        }
23064        Ok(y)
23065    }
23066
23067    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
23068    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
23069    /// chunked BF16 numerical program instead of re-encoding the weight.
23070    pub fn linear_bf16_resident(
23071        &self,
23072        x: &CudaSlice<f32>,
23073        data: &CudaSlice<u8>,
23074        m: usize,
23075        in_f: usize,
23076        out_f: usize,
23077    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23078        if data.len() != in_f * out_f * 2 {
23079            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
23080        }
23081        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
23082    }
23083
23084    /// Execute a resident BF16 projection as fixed-width output-row chunks.
23085    ///
23086    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
23087    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
23088    /// model topology rather than the active rank count.
23089    pub fn linear_bf16_resident_canonical_rows(
23090        &self,
23091        x: &CudaSlice<f32>,
23092        data: &CudaSlice<u8>,
23093        m: usize,
23094        in_f: usize,
23095        out_f: usize,
23096        canonical_chunk_rows: usize,
23097    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23098        if data.len() != in_f * out_f * 2 {
23099            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
23100        }
23101        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
23102    }
23103
23104    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
23105    ///
23106    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
23107    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
23108    pub fn linear_f32_resident_canonical_rows(
23109        &self,
23110        x: &CudaSlice<f32>,
23111        data: &CudaSlice<f32>,
23112        m: usize,
23113        in_f: usize,
23114        out_f: usize,
23115        canonical_chunk_rows: usize,
23116    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23117        self.linear_f32_resident_canonical_rows_inner(
23118            x,
23119            data,
23120            m,
23121            in_f,
23122            out_f,
23123            canonical_chunk_rows,
23124            false,
23125        )
23126    }
23127
23128    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
23129    ///
23130    /// The projection shapes and values are identical to
23131    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
23132    /// changes, replacing one device copy per token with one placement kernel per output chunk.
23133    pub fn linear_f32_resident_canonical_rows_strided(
23134        &self,
23135        x: &CudaSlice<f32>,
23136        data: &CudaSlice<f32>,
23137        m: usize,
23138        in_f: usize,
23139        out_f: usize,
23140        canonical_chunk_rows: usize,
23141    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23142        self.linear_f32_resident_canonical_rows_inner(
23143            x,
23144            data,
23145            m,
23146            in_f,
23147            out_f,
23148            canonical_chunk_rows,
23149            true,
23150        )
23151    }
23152
23153    #[allow(clippy::too_many_arguments)]
23154    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23155    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
23156    fn linear_f32_resident_canonical_rows_inner(
23157        &self,
23158        x: &CudaSlice<f32>,
23159        data: &CudaSlice<f32>,
23160        m: usize,
23161        in_f: usize,
23162        out_f: usize,
23163        canonical_chunk_rows: usize,
23164        strided_output: bool,
23165    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23166        if data.len() != in_f * out_f {
23167            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
23168        }
23169        if canonical_chunk_rows == 0
23170            || canonical_chunk_rows > out_f
23171            || out_f % canonical_chunk_rows != 0
23172        {
23173            return Err(format!(
23174                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
23175            )
23176            .into());
23177        }
23178        if canonical_chunk_rows == out_f {
23179            return self.linear(x, data, m, in_f, out_f);
23180        }
23181
23182        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
23183        let input = x.slice(0..x.len());
23184        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
23185            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
23186            if m == 1 {
23187                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
23188                self.linear_device_into(
23189                    &input,
23190                    &weights,
23191                    &mut destination,
23192                    1,
23193                    in_f,
23194                    canonical_chunk_rows,
23195                )?;
23196                continue;
23197            }
23198            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
23199            if strided_output {
23200                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
23201            } else {
23202                for token in 0..m {
23203                    let source = chunk
23204                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
23205                    let mut destination =
23206                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
23207                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
23208                }
23209            }
23210        }
23211        Ok(y)
23212    }
23213
23214    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
23215    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
23216    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
23217    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
23218    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
23219    pub fn linear_f32_resident_canonical_rows_t1_into(
23220        &self,
23221        x: &CudaSlice<f32>,
23222        data: &CudaSlice<f32>,
23223        y: &mut CudaSlice<f32>,
23224        in_f: usize,
23225        out_f: usize,
23226        canonical_chunk_rows: usize,
23227    ) -> Result<(), Box<dyn std::error::Error>> {
23228        if data.len() != in_f * out_f {
23229            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
23230        }
23231        if y.len() != out_f || x.len() != in_f {
23232            return Err(format!(
23233                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
23234                x.len(),
23235                y.len()
23236            )
23237            .into());
23238        }
23239        if canonical_chunk_rows == 0
23240            || canonical_chunk_rows > out_f
23241            || out_f % canonical_chunk_rows != 0
23242        {
23243            return Err(format!(
23244                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
23245            )
23246            .into());
23247        }
23248        let input = x.slice(0..x.len());
23249        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
23250            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
23251            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
23252            self.linear_device_into(
23253                &input,
23254                &weights,
23255                &mut destination,
23256                1,
23257                in_f,
23258                canonical_chunk_rows,
23259            )?;
23260        }
23261        Ok(())
23262    }
23263
23264    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
23265    /// without the allocation, for workspace-resident operands.
23266    pub fn linear_t1_into(
23267        &self,
23268        x: &cudarc::driver::CudaView<'_, f32>,
23269        w: &cudarc::driver::CudaView<'_, f32>,
23270        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
23271        in_f: usize,
23272        out_f: usize,
23273    ) -> Result<(), Box<dyn std::error::Error>> {
23274        self.linear_device_into(x, w, y, 1, in_f, out_f)
23275    }
23276
23277    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
23278    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
23279    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
23280    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
23281    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
23282    /// router/shexp sites and matmul_decode_exact's Float arm.
23283    pub fn linear_decode_exact(
23284        &self,
23285        x: &CudaSlice<f32>,
23286        w: &CudaSlice<f32>,
23287        m_tokens: usize,
23288        in_f: usize,
23289        out_f: usize,
23290    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23291        if m_tokens == 1 {
23292            return self.linear(x, w, 1, in_f, out_f);
23293        }
23294        let xv = self.view(x, m_tokens * in_f);
23295        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
23296        for t in 0..m_tokens {
23297            let row = xv.slice(t * in_f..(t + 1) * in_f);
23298            let mut xr = self.alloc_uninit::<f32>(in_f)?;
23299            self.copy_view_into(&mut xr, 0, &row, in_f)?;
23300            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
23301            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
23302        }
23303        Ok(y)
23304    }
23305
23306    pub fn linear(
23307        &self,
23308        x: &CudaSlice<f32>,
23309        w: &CudaSlice<f32>,
23310        m_tokens: usize,
23311        in_f: usize,
23312        out_f: usize,
23313    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23314        self.linear_device(x, w, m_tokens, in_f, out_f)
23315    }
23316
23317    fn linear_device<I>(
23318        &self,
23319        x: &I,
23320        w: &I,
23321        m_tokens: usize,
23322        in_f: usize,
23323        out_f: usize,
23324    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
23325    where
23326        I: cudarc::driver::DevicePtr<f32>,
23327    {
23328        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
23329        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
23330        Ok(c)
23331    }
23332
23333    fn linear_device_into<I, O>(
23334        &self,
23335        x: &I,
23336        w: &I,
23337        c: &mut O,
23338        m_tokens: usize,
23339        in_f: usize,
23340        out_f: usize,
23341    ) -> Result<(), Box<dyn std::error::Error>>
23342    where
23343        I: cudarc::driver::DevicePtr<f32>,
23344        O: cudarc::driver::DevicePtrMut<f32>,
23345    {
23346        use cudarc::cublaslt::{Matmul, MatmulConfig};
23347        let cfg = MatmulConfig {
23348            transa: true,
23349            transb: false,
23350            transc: false,
23351            m: out_f as u64,
23352            n: m_tokens as u64,
23353            k: in_f as u64,
23354            alpha: 1.0,
23355            lda: in_f as i64,
23356            ldb: in_f as i64,
23357            beta: 0.0,
23358            ldc: out_f as i64,
23359            stride_a: None,
23360            stride_b: None,
23361            stride_c: None,
23362            stride_bias: None,
23363            batch_size: None,
23364        };
23365        let blas = self.gpu.blas();
23366        unsafe {
23367            blas.matmul(cfg, w, x, c, None, None)?;
23368        }
23369        Ok(())
23370    }
23371
23372    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
23373    ///
23374    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
23375    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
23376    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
23377    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
23378    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
23379    /// launch error mid-request.
23380    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23381    pub fn sdpa_naive(
23382        &self,
23383        q: &CudaSlice<f32>,
23384        k: &CudaSlice<f32>,
23385        v: &CudaSlice<f32>,
23386        o: &mut CudaSlice<f32>,
23387        head_dim: usize,
23388        n_head: usize,
23389        n_head_kv: usize,
23390        t: usize,
23391        t_kv: usize,
23392        scale: f32,
23393        causal: bool,
23394    ) -> Result<(), Box<dyn std::error::Error>> {
23395        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
23396            return self.sdpa_naive_gmem(
23397                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23398            );
23399        }
23400        let f = self.func("sdpa_naive_f32");
23401        let cfg = LaunchConfig {
23402            grid_dim: (n_head as u32, t as u32, 1),
23403            block_dim: (128, 1, 1),
23404            shared_mem_bytes: (t_kv * 4) as u32,
23405        };
23406        let (hd, nh, nhkv, ti, tkvi, cz) = (
23407            head_dim as i32,
23408            n_head as i32,
23409            n_head_kv as i32,
23410            t as i32,
23411            t_kv as i32,
23412            causal as i32,
23413        );
23414        let __s_b = self.gpu.stream();
23415        let mut b = __s_b.launch_builder(&f);
23416        b.arg(q)
23417            .arg(k)
23418            .arg(v)
23419            .arg(o)
23420            .arg(&hd)
23421            .arg(&nh)
23422            .arg(&nhkv)
23423            .arg(&ti)
23424            .arg(&tkvi)
23425            .arg(&scale)
23426            .arg(&cz);
23427        unsafe {
23428            b.launch(cfg)?;
23429        }
23430        Ok(())
23431    }
23432
23433    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
23434    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
23435    /// of dynamic shared memory: identical loop structure and reduction order, so the output
23436    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
23437    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
23438    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
23439    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
23440    /// T==T_kv caller cannot silently allocate tens of GB.
23441    #[allow(clippy::too_many_arguments)]
23442    pub fn sdpa_naive_gmem(
23443        &self,
23444        q: &CudaSlice<f32>,
23445        k: &CudaSlice<f32>,
23446        v: &CudaSlice<f32>,
23447        o: &mut CudaSlice<f32>,
23448        head_dim: usize,
23449        n_head: usize,
23450        n_head_kv: usize,
23451        t: usize,
23452        t_kv: usize,
23453        scale: f32,
23454        causal: bool,
23455    ) -> Result<(), Box<dyn std::error::Error>> {
23456        let ws_len = n_head
23457            .checked_mul(t)
23458            .and_then(|x| x.checked_mul(t_kv))
23459            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
23460        let ws_bytes = ws_len
23461            .checked_mul(std::mem::size_of::<f32>())
23462            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
23463        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
23464            return Err(format!(
23465                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
23466                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
23467                 needs a tiled/flash kernel, not the naive oracle"
23468            )
23469            .into());
23470        }
23471        let mut scores = self.uninit(ws_len)?;
23472        let f = self.func("sdpa_naive_gmem_f32");
23473        let cfg = LaunchConfig {
23474            grid_dim: (n_head as u32, t as u32, 1),
23475            block_dim: (128, 1, 1),
23476            shared_mem_bytes: 0,
23477        };
23478        let (hd, nh, nhkv, ti, tkvi, cz) = (
23479            head_dim as i32,
23480            n_head as i32,
23481            n_head_kv as i32,
23482            t as i32,
23483            t_kv as i32,
23484            causal as i32,
23485        );
23486        let __s_b = self.gpu.stream();
23487        let mut b = __s_b.launch_builder(&f);
23488        b.arg(q)
23489            .arg(k)
23490            .arg(v)
23491            .arg(o)
23492            .arg(&mut scores)
23493            .arg(&hd)
23494            .arg(&nh)
23495            .arg(&nhkv)
23496            .arg(&ti)
23497            .arg(&tkvi)
23498            .arg(&scale)
23499            .arg(&cz);
23500        unsafe {
23501            b.launch(cfg)?;
23502        }
23503        Ok(())
23504    }
23505
23506    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
23507    /// bidirectional image islands. `span_id` labels each absolute kv position
23508    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
23509    /// reproducing the reference's non-causal image batch. window 0 = no window.
23510    #[allow(clippy::too_many_arguments)]
23511    pub fn sdpa_naive_island(
23512        &self,
23513        q: &CudaSlice<f32>,
23514        k: &CudaSlice<f32>,
23515        v: &CudaSlice<f32>,
23516        o: &mut CudaSlice<f32>,
23517        span_id: &CudaSlice<i32>,
23518        head_dim: usize,
23519        n_head: usize,
23520        n_head_kv: usize,
23521        t: usize,
23522        t_kv: usize,
23523        scale: f32,
23524        window: usize,
23525    ) -> Result<(), Box<dyn std::error::Error>> {
23526        let f = self.func("sdpa_naive_island_f32");
23527        let cfg = LaunchConfig {
23528            grid_dim: (n_head as u32, t as u32, 1),
23529            block_dim: (128, 1, 1),
23530            shared_mem_bytes: (t_kv * 4) as u32,
23531        };
23532        let (hd, nh, nhkv, ti, tkvi, wi) = (
23533            head_dim as i32,
23534            n_head as i32,
23535            n_head_kv as i32,
23536            t as i32,
23537            t_kv as i32,
23538            window as i32,
23539        );
23540        let __s_b = self.gpu.stream();
23541        let mut b = __s_b.launch_builder(&f);
23542        b.arg(q)
23543            .arg(k)
23544            .arg(v)
23545            .arg(o)
23546            .arg(span_id)
23547            .arg(&hd)
23548            .arg(&nh)
23549            .arg(&nhkv)
23550            .arg(&ti)
23551            .arg(&tkvi)
23552            .arg(&scale)
23553            .arg(&wi);
23554        unsafe {
23555            b.launch(cfg)?;
23556        }
23557        Ok(())
23558    }
23559
23560    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
23561    #[allow(clippy::too_many_arguments)]
23562    pub fn sdpa_naive_w(
23563        &self,
23564        q: &CudaSlice<f32>,
23565        k: &CudaSlice<f32>,
23566        v: &CudaSlice<f32>,
23567        o: &mut CudaSlice<f32>,
23568        head_dim: usize,
23569        n_head: usize,
23570        n_head_kv: usize,
23571        t: usize,
23572        t_kv: usize,
23573        scale: f32,
23574        causal: bool,
23575        window: usize,
23576    ) -> Result<(), Box<dyn std::error::Error>> {
23577        let f = self.func("sdpa_naive_w_f32");
23578        let cfg = LaunchConfig {
23579            grid_dim: (n_head as u32, t as u32, 1),
23580            block_dim: (128, 1, 1),
23581            shared_mem_bytes: (t_kv * 4) as u32,
23582        };
23583        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
23584            head_dim as i32,
23585            n_head as i32,
23586            n_head_kv as i32,
23587            t as i32,
23588            t_kv as i32,
23589            causal as i32,
23590            window as i32,
23591        );
23592        let __s_b = self.gpu.stream();
23593        let mut b = __s_b.launch_builder(&f);
23594        b.arg(q)
23595            .arg(k)
23596            .arg(v)
23597            .arg(o)
23598            .arg(&hd)
23599            .arg(&nh)
23600            .arg(&nhkv)
23601            .arg(&ti)
23602            .arg(&tkvi)
23603            .arg(&scale)
23604            .arg(&cz)
23605            .arg(&wi);
23606        unsafe {
23607            b.launch(cfg)?;
23608        }
23609        Ok(())
23610    }
23611
23612    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
23613    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
23614    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
23615    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
23616    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
23617    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
23618    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
23619    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
23620    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
23621    #[allow(clippy::too_many_arguments)]
23622    pub fn sdpa_naive_w_lo(
23623        &self,
23624        q: &CudaSlice<f32>,
23625        k: &CudaSlice<f32>,
23626        v: &CudaSlice<f32>,
23627        o: &mut CudaSlice<f32>,
23628        head_dim: usize,
23629        n_head: usize,
23630        n_head_kv: usize,
23631        t: usize,
23632        t_kv: usize,
23633        scale: f32,
23634        causal: bool,
23635        window: usize,
23636    ) -> Result<(), Box<dyn std::error::Error>> {
23637        let kv_lo = if window > 0 {
23638            (t_kv - t + 1).saturating_sub(window)
23639        } else {
23640            0
23641        };
23642        let smem = (t_kv - kv_lo) * 4;
23643        if smem > 48 * 1024 {
23644            return Err(format!(
23645                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
23646                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
23647                 a window this wide needs the multi-pass long-ctx kernel"
23648            )
23649            .into());
23650        }
23651        let f = self.func("sdpa_naive_w_lo_f32");
23652        let cfg = LaunchConfig {
23653            grid_dim: (n_head as u32, t as u32, 1),
23654            block_dim: (128, 1, 1),
23655            shared_mem_bytes: smem as u32,
23656        };
23657        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
23658            head_dim as i32,
23659            n_head as i32,
23660            n_head_kv as i32,
23661            t as i32,
23662            t_kv as i32,
23663            causal as i32,
23664            window as i32,
23665            kv_lo as i32,
23666        );
23667        let __s_b = self.gpu.stream();
23668        let mut b = __s_b.launch_builder(&f);
23669        b.arg(q)
23670            .arg(k)
23671            .arg(v)
23672            .arg(o)
23673            .arg(&hd)
23674            .arg(&nh)
23675            .arg(&nhkv)
23676            .arg(&ti)
23677            .arg(&tkvi)
23678            .arg(&scale)
23679            .arg(&cz)
23680            .arg(&wi)
23681            .arg(&lo);
23682        unsafe {
23683            b.launch(cfg)?;
23684        }
23685        Ok(())
23686    }
23687
23688    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
23689    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23690    pub fn sdpa_naive_view(
23691        &self,
23692        q: &CudaSlice<f32>,
23693        k: &cudarc::driver::CudaView<f32>,
23694        v: &cudarc::driver::CudaView<f32>,
23695        o: &mut CudaSlice<f32>,
23696        head_dim: usize,
23697        n_head: usize,
23698        n_head_kv: usize,
23699        t: usize,
23700        t_kv: usize,
23701        scale: f32,
23702        causal: bool,
23703    ) -> Result<(), Box<dyn std::error::Error>> {
23704        let f = self.func("sdpa_naive_f32");
23705        let cfg = LaunchConfig {
23706            grid_dim: (n_head as u32, t as u32, 1),
23707            block_dim: (128, 1, 1),
23708            shared_mem_bytes: (t_kv * 4) as u32,
23709        };
23710        let (hd, nh, nhkv, ti, tkvi, cz) = (
23711            head_dim as i32,
23712            n_head as i32,
23713            n_head_kv as i32,
23714            t as i32,
23715            t_kv as i32,
23716            causal as i32,
23717        );
23718        let __s_b = self.gpu.stream();
23719        let mut b = __s_b.launch_builder(&f);
23720        b.arg(q)
23721            .arg(k)
23722            .arg(v)
23723            .arg(o)
23724            .arg(&hd)
23725            .arg(&nh)
23726            .arg(&nhkv)
23727            .arg(&ti)
23728            .arg(&tkvi)
23729            .arg(&scale)
23730            .arg(&cz);
23731        unsafe {
23732            b.launch(cfg)?;
23733        }
23734        Ok(())
23735    }
23736
23737    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
23738    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
23739    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
23740    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
23741    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
23742    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
23743    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
23744    #[allow(clippy::too_many_arguments)]
23745    pub fn fa_dequant_kv_view_f32(
23746        &self,
23747        k: &cudarc::driver::CudaView<u8>,
23748        v: &cudarc::driver::CudaView<u8>,
23749        kf: &mut CudaSlice<f32>,
23750        vf: &mut CudaSlice<f32>,
23751        kv_dim_k: usize,
23752        kv_dim_v: usize,
23753        t_kv: usize,
23754        k_tok_bytes: usize,
23755        v_tok_bytes: usize,
23756        g: bool,
23757    ) -> Result<(), Box<dyn std::error::Error>> {
23758        let f = if g {
23759            self.func_g("fa_dequant_kv_ws_f32")
23760        } else {
23761            self.func("fa_dequant_kv_ws_f32")
23762        };
23763        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
23764        #[allow(clippy::manual_div_ceil)]
23765        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23766        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23767        let cfg = LaunchConfig {
23768            grid_dim: (nblk.max(1), 1, 1),
23769            block_dim: (256, 1, 1),
23770            shared_mem_bytes: 0,
23771        };
23772        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
23773        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23774        let __s_b = self.gpu.stream();
23775        let mut b = __s_b.launch_builder(&f);
23776        b.arg(k)
23777            .arg(v)
23778            .arg(&mut *kf)
23779            .arg(&mut *vf)
23780            .arg(&kdk)
23781            .arg(&kdv)
23782            .arg(&tkvi)
23783            .arg(&ktb)
23784            .arg(&vtb);
23785        unsafe {
23786            b.launch(cfg)?;
23787        }
23788        Ok(())
23789    }
23790
23791    #[allow(clippy::too_many_arguments)]
23792    pub fn sdpa_naive_quantized_view(
23793        &self,
23794        q: &CudaSlice<f32>,
23795        k: &cudarc::driver::CudaView<u8>,
23796        v: &cudarc::driver::CudaView<u8>,
23797        o: &mut CudaSlice<f32>,
23798        head_dim: usize,
23799        n_head: usize,
23800        n_head_kv: usize,
23801        t: usize,
23802        t_kv: usize,
23803        scale: f32,
23804        causal: bool,
23805        k_tok_bytes: usize,
23806        v_tok_bytes: usize,
23807    ) -> Result<(), Box<dyn std::error::Error>> {
23808        let kv_dim = n_head_kv * head_dim;
23809        let mut kf = self.uninit(t_kv * kv_dim)?;
23810        let mut vf = self.uninit(t_kv * kv_dim)?;
23811        let f = self.func("fa_dequant_kv_ws_f32");
23812        let total = (2 * t_kv * kv_dim) as u64;
23813        #[allow(clippy::manual_div_ceil)]
23814        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23815        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23816        let cfg = LaunchConfig {
23817            grid_dim: (nblk.max(1), 1, 1),
23818            block_dim: (256, 1, 1),
23819            shared_mem_bytes: 0,
23820        };
23821        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
23822        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
23823        let __s_b = self.gpu.stream();
23824        let mut b = __s_b.launch_builder(&f);
23825        b.arg(k)
23826            .arg(v)
23827            .arg(&mut kf)
23828            .arg(&mut vf)
23829            .arg(&kv_dim_i)
23830            .arg(&kv_dim_i)
23831            .arg(&t_kv_i)
23832            .arg(&k_tok_bytes_i)
23833            .arg(&v_tok_bytes_i);
23834        unsafe { b.launch(cfg)? };
23835        self.sdpa_naive(
23836            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23837        )
23838    }
23839
23840    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
23841    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
23842    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
23843    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
23844    /// unwindowed function above and produces bit-identical output at window == 0.
23845    ///
23846    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
23847    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
23848    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
23849    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
23850    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
23851    #[allow(clippy::too_many_arguments)]
23852    pub fn sdpa_naive_w_quantized_view(
23853        &self,
23854        q: &CudaSlice<f32>,
23855        k: &cudarc::driver::CudaView<u8>,
23856        v: &cudarc::driver::CudaView<u8>,
23857        o: &mut CudaSlice<f32>,
23858        head_dim: usize,
23859        n_head: usize,
23860        n_head_kv: usize,
23861        t: usize,
23862        t_kv: usize,
23863        scale: f32,
23864        causal: bool,
23865        window: usize,
23866        k_tok_bytes: usize,
23867        v_tok_bytes: usize,
23868    ) -> Result<(), Box<dyn std::error::Error>> {
23869        let kv_dim = n_head_kv * head_dim;
23870        let mut kf = self.uninit(t_kv * kv_dim)?;
23871        let mut vf = self.uninit(t_kv * kv_dim)?;
23872        let f = self.func("fa_dequant_kv_ws_f32");
23873        let total = (2 * t_kv * kv_dim) as u64;
23874        #[allow(clippy::manual_div_ceil)]
23875        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23876        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23877        let cfg = LaunchConfig {
23878            grid_dim: (nblk.max(1), 1, 1),
23879            block_dim: (256, 1, 1),
23880            shared_mem_bytes: 0,
23881        };
23882        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
23883        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
23884        let __s_b = self.gpu.stream();
23885        let mut b = __s_b.launch_builder(&f);
23886        b.arg(k)
23887            .arg(v)
23888            .arg(&mut kf)
23889            .arg(&mut vf)
23890            .arg(&kv_dim_i)
23891            .arg(&kv_dim_i)
23892            .arg(&t_kv_i)
23893            .arg(&k_tok_bytes_i)
23894            .arg(&v_tok_bytes_i);
23895        unsafe { b.launch(cfg)? };
23896        self.sdpa_naive_w(
23897            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
23898        )
23899    }
23900
23901    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
23902    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
23903    /// Q/K/V/O [head_dim, n_head(_kv), T].
23904    #[allow(clippy::too_many_arguments)]
23905    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23906    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23907    pub fn fa_prefill(
23908        &self,
23909        q: &CudaSlice<f32>,
23910        k: &CudaSlice<f32>,
23911        v: &CudaSlice<f32>,
23912        o: &mut CudaSlice<f32>,
23913        head_dim: usize,
23914        n_head: usize,
23915        n_head_kv: usize,
23916        t: usize,
23917        t_kv: usize,
23918        scale: f32,
23919        causal: bool,
23920    ) -> Result<(), Box<dyn std::error::Error>> {
23921        if portable_mma_gated() {
23922            return self.sdpa_naive(
23923                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23924            );
23925        }
23926        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
23927        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
23928        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
23929        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
23930        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
23931        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
23932        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
23933        let fa3_on = head_dim == 256
23934            && causal
23935            && t == t_kv
23936            && match std::env::var("MEMRA_FA3").as_deref() {
23937                Ok("0") => false,
23938                // The force arm consults the arch now: the bf16 stage below calls
23939                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
23940                // a portable build. Refuse at the switch, not at the lookup.
23941                Ok("1") => {
23942                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
23943                    true
23944                }
23945                _ => cfg!(memra_hopper_mma),
23946            };
23947        if fa3_on {
23948            let n = t * n_head * head_dim;
23949            let nkv = t * n_head_kv * head_dim;
23950            let mut q16 = self.alloc_u8_uninit(n * 2)?;
23951            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
23952            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
23953            self.f32_to_bf16_into(q, &mut q16, n)?;
23954            self.f32_to_bf16_into(k, &mut k16, nkv)?;
23955            self.f32_to_bf16_into(v, &mut v16, nkv)?;
23956            let rc = {
23957                use cudarc::driver::{DevicePtr, DevicePtrMut};
23958                let stream = self.gpu.stream();
23959                let (qp, _g1) = q16.device_ptr(&stream);
23960                let (kp, _g2) = k16.device_ptr(&stream);
23961                let (vp, _g3) = v16.device_ptr(&stream);
23962                let (op, _g4) = o.device_ptr_mut(&stream);
23963                unsafe {
23964                    memra_fa3_prefill(
23965                        qp as *const core::ffi::c_void,
23966                        kp as *const core::ffi::c_void,
23967                        vp as *const core::ffi::c_void,
23968                        op as *mut f32,
23969                        t as i32,
23970                        n_head as i32,
23971                        n_head_kv as i32,
23972                        head_dim as i32,
23973                        scale,
23974                        stream.cu_stream() as *mut core::ffi::c_void,
23975                    )
23976                }
23977            };
23978            if rc != 0 {
23979                return Err(format!("memra_fa3_prefill rc={rc}").into());
23980            }
23981            return Ok(());
23982        }
23983        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
23984        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
23985        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
23986        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
23987        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23988        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
23989        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
23990            const BLOCK_Q: usize = 64;
23991            const BKX: usize = 32;
23992            let f = self.func("fa_prefill_bf16_p1");
23993            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
23994                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
23995            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23996            f.set_attribute(
23997                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23998                shmem as i32,
23999            )?;
24000            let cfg = LaunchConfig {
24001                grid_dim: (
24002                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24003                    n_head as u32,
24004                    1,
24005                ),
24006                block_dim: (32, 4, 1),
24007                shared_mem_bytes: shmem,
24008            };
24009            let (hd, nh, nhkv, ti, tkvi, cz) = (
24010                head_dim as i32,
24011                n_head as i32,
24012                n_head_kv as i32,
24013                t as i32,
24014                t_kv as i32,
24015                causal as i32,
24016            );
24017            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24018            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24019            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24020            let __s_b = self.gpu.stream();
24021            let mut b = __s_b.launch_builder(&f);
24022            b.arg(&qb)
24023                .arg(&kb)
24024                .arg(&vb)
24025                .arg(o)
24026                .arg(&hd)
24027                .arg(&nh)
24028                .arg(&nhkv)
24029                .arg(&ti)
24030                .arg(&tkvi)
24031                .arg(&scale)
24032                .arg(&cz);
24033            unsafe {
24034                b.launch(cfg)?;
24035            }
24036            return Ok(());
24037        }
24038        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
24039        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
24040        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
24041        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
24042        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
24043        const BK: usize = 32;
24044        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
24045        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
24046        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
24047        let (block_q, warps, w2_sfx): (usize, u32, &str) =
24048            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
24049        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
24050        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
24051        // other head_dims to sdpa_naive before reaching here.
24052        let hd_sfx = fa_hd_suffix(head_dim)?;
24053        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
24054        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
24055        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
24056        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
24057        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
24058        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
24059        let (kb16, vb16) = if bf16kv {
24060            let n = t_kv * n_head_kv * head_dim;
24061            let mut kb = self.alloc_u8_uninit(n * 2)?;
24062            let mut vb = self.alloc_u8_uninit(n * 2)?;
24063            let fcv = self.func("f32_to_bf16_bulk");
24064            let ni = n as i64;
24065            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24066            let __s_b = self.gpu.stream();
24067            let mut b = __s_b.launch_builder(&fcv);
24068            b.arg(k).arg(&mut kb).arg(&ni);
24069            unsafe {
24070                b.launch(cfgc)?;
24071            }
24072            let __s_b = self.gpu.stream();
24073            let mut b = __s_b.launch_builder(&fcv);
24074            b.arg(v).arg(&mut vb).arg(&ni);
24075            unsafe {
24076                b.launch(cfgc)?;
24077            }
24078            (Some(kb), Some(vb))
24079        } else {
24080            (None, None)
24081        };
24082        let f = self.func(&if bf16kv {
24083            format!("fa_prefill_bf16kv_pp{hd_sfx}")
24084        } else {
24085            format!(
24086                "fa_prefill_f32{}{}{hd_sfx}",
24087                if floor { "" } else { "_pp" },
24088                if floor { "" } else { w2_sfx }
24089            )
24090        });
24091        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
24092        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
24093        let kv_stages = if bf16kv { 2 } else { 1 };
24094        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
24095            + 4 * (block_q * BK + 2 * block_q)) as u32;
24096        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24097        f.set_attribute(
24098            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24099            shmem as i32,
24100        )?;
24101        let cfg = LaunchConfig {
24102            grid_dim: (
24103                (t as u32 + block_q as u32 - 1) / block_q as u32,
24104                n_head as u32,
24105                1,
24106            ),
24107            block_dim: (32, warps, 1),
24108            shared_mem_bytes: shmem,
24109        };
24110        let (hd, nh, nhkv, ti, tkvi, cz) = (
24111            head_dim as i32,
24112            n_head as i32,
24113            n_head_kv as i32,
24114            t as i32,
24115            t_kv as i32,
24116            causal as i32,
24117        );
24118        let __s_b = self.gpu.stream();
24119        let mut b = __s_b.launch_builder(&f);
24120        b.arg(q);
24121        match (&kb16, &vb16) {
24122            (Some(kb), Some(vb)) => {
24123                b.arg(kb).arg(vb);
24124            }
24125            _ => {
24126                b.arg(k).arg(v);
24127            }
24128        }
24129        b.arg(o)
24130            .arg(&hd)
24131            .arg(&nh)
24132            .arg(&nhkv)
24133            .arg(&ti)
24134            .arg(&tkvi)
24135            .arg(&scale)
24136            .arg(&cz);
24137        unsafe {
24138            b.launch(cfg)?;
24139        }
24140        Ok(())
24141    }
24142
24143    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
24144    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
24145    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
24146    #[allow(clippy::too_many_arguments)]
24147    pub fn fa_prefill_w(
24148        &self,
24149        q: &CudaSlice<f32>,
24150        k: &CudaSlice<f32>,
24151        v: &CudaSlice<f32>,
24152        o: &mut CudaSlice<f32>,
24153        head_dim: usize,
24154        n_head: usize,
24155        n_head_kv: usize,
24156        t: usize,
24157        t_kv: usize,
24158        scale: f32,
24159        causal: bool,
24160        window: usize,
24161    ) -> Result<(), Box<dyn std::error::Error>> {
24162        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
24163        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
24164        if portable_mma_gated() {
24165            return self.sdpa_naive_w(
24166                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
24167            );
24168        }
24169        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
24170        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
24171        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
24172        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24173        let faw_f32 =
24174            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
24175        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
24176        self.fa_prefill_w_arm(
24177            q,
24178            k,
24179            v,
24180            o,
24181            head_dim,
24182            n_head,
24183            n_head_kv,
24184            t,
24185            t_kv,
24186            scale,
24187            causal,
24188            window,
24189            floor || faw_f32,
24190            floor,
24191        )
24192    }
24193
24194    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
24195    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
24196    #[allow(clippy::too_many_arguments)]
24197    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
24198    pub fn fa_prefill_w_pre(
24199        &self,
24200        qb: &CudaSlice<u8>,
24201        kb: &CudaSlice<u8>,
24202        vb: &CudaSlice<u8>,
24203        o: &mut CudaSlice<f32>,
24204        head_dim: usize,
24205        n_head: usize,
24206        n_head_kv: usize,
24207        t: usize,
24208        t_kv: usize,
24209        scale: f32,
24210        causal: bool,
24211        window: usize,
24212        v_f16: bool,
24213    ) -> Result<(), Box<dyn std::error::Error>> {
24214        const BLOCK_Q: usize = 64;
24215        const BK: usize = 32;
24216        debug_assert_eq!(head_dim, 256);
24217        let hp = fa_f16pv_on()
24218            && faw_hp_on()
24219            && n_head.is_multiple_of(2)
24220            && (n_head / n_head_kv).is_multiple_of(2);
24221        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
24222        if hp {
24223            const BLOCK_QH: usize = 32;
24224            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
24225            // else re-encode through the pooled scratch (stream-ordered reuse).
24226            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
24227            let vh: &CudaSlice<u8> = if v_f16 {
24228                vb
24229            } else {
24230                let n = t_kv * n_head_kv * head_dim;
24231                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
24232                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
24233                }
24234                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
24235                vguard.as_ref().unwrap()
24236            };
24237            let f = self.func("fa_prefill_w_bf16_p1h2");
24238            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
24239            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24240            f.set_attribute(
24241                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24242                shmem as i32,
24243            )?;
24244            let cfg = LaunchConfig {
24245                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
24246                block_dim: (32, 4, 1),
24247                shared_mem_bytes: shmem,
24248            };
24249            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24250                head_dim as i32,
24251                n_head as i32,
24252                n_head_kv as i32,
24253                t as i32,
24254                t_kv as i32,
24255                causal as i32,
24256                window as i32,
24257            );
24258            let __s_b = self.gpu.stream();
24259            let mut b = __s_b.launch_builder(&f);
24260            b.arg(qb)
24261                .arg(kb)
24262                .arg(vh)
24263                .arg(o)
24264                .arg(&hd)
24265                .arg(&nh)
24266                .arg(&nhkv)
24267                .arg(&ti)
24268                .arg(&tkvi)
24269                .arg(&scale)
24270                .arg(&cz)
24271                .arg(&wi);
24272            unsafe {
24273                b.launch(cfg)?;
24274            }
24275            return Ok(());
24276        }
24277        let f = self.func("fa_prefill_w_bf16_p1");
24278        let shmem =
24279            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24280        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24281        f.set_attribute(
24282            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24283            shmem as i32,
24284        )?;
24285        let cfg = LaunchConfig {
24286            grid_dim: (
24287                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24288                n_head as u32,
24289                1,
24290            ),
24291            block_dim: (32, 4, 1),
24292            shared_mem_bytes: shmem,
24293        };
24294        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24295            head_dim as i32,
24296            n_head as i32,
24297            n_head_kv as i32,
24298            t as i32,
24299            t_kv as i32,
24300            causal as i32,
24301            window as i32,
24302        );
24303        let __s_b = self.gpu.stream();
24304        let mut b = __s_b.launch_builder(&f);
24305        b.arg(qb)
24306            .arg(kb)
24307            .arg(vb)
24308            .arg(o)
24309            .arg(&hd)
24310            .arg(&nh)
24311            .arg(&nhkv)
24312            .arg(&ti)
24313            .arg(&tkvi)
24314            .arg(&scale)
24315            .arg(&cz)
24316            .arg(&wi);
24317        unsafe {
24318            b.launch(cfg)?;
24319        }
24320        Ok(())
24321    }
24322
24323    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
24324    #[allow(clippy::too_many_arguments)]
24325    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
24326    pub fn fa_prefill_w_arm(
24327        &self,
24328        q: &CudaSlice<f32>,
24329        k: &CudaSlice<f32>,
24330        v: &CudaSlice<f32>,
24331        o: &mut CudaSlice<f32>,
24332        head_dim: usize,
24333        n_head: usize,
24334        n_head_kv: usize,
24335        t: usize,
24336        t_kv: usize,
24337        scale: f32,
24338        causal: bool,
24339        window: usize,
24340        f32_stage: bool,
24341        floor: bool,
24342    ) -> Result<(), Box<dyn std::error::Error>> {
24343        const BLOCK_Q: usize = 64;
24344        const BK: usize = 32;
24345        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
24346        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
24347        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
24348        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
24349        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24350        let p1 = !floor
24351            && !f32_stage
24352            && *P1_ON.get_or_init(|| {
24353                std::env::var("MEMRA_FAW_P1")
24354                    .map(|v| v != "0")
24355                    .unwrap_or(true)
24356            });
24357        let hp = p1
24358            && fa_f16pv_on()
24359            && faw_hp_on()
24360            && n_head.is_multiple_of(2)
24361            && (n_head / n_head_kv).is_multiple_of(2);
24362        if hp {
24363            const BLOCK_QH: usize = 32;
24364            let f = self.func("fa_prefill_w_bf16_p1h2");
24365            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
24366            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24367            f.set_attribute(
24368                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24369                shmem as i32,
24370            )?;
24371            let cfg = LaunchConfig {
24372                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
24373                block_dim: (32, 4, 1),
24374                shared_mem_bytes: shmem,
24375            };
24376            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24377                head_dim as i32,
24378                n_head as i32,
24379                n_head_kv as i32,
24380                t as i32,
24381                t_kv as i32,
24382                causal as i32,
24383                window as i32,
24384            );
24385            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24386            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24387            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
24388            let __s_b = self.gpu.stream();
24389            let mut b = __s_b.launch_builder(&f);
24390            b.arg(&qb)
24391                .arg(&kb)
24392                .arg(&vh)
24393                .arg(o)
24394                .arg(&hd)
24395                .arg(&nh)
24396                .arg(&nhkv)
24397                .arg(&ti)
24398                .arg(&tkvi)
24399                .arg(&scale)
24400                .arg(&cz)
24401                .arg(&wi);
24402            unsafe {
24403                b.launch(cfg)?;
24404            }
24405            return Ok(());
24406        }
24407        if p1 {
24408            let f = self.func("fa_prefill_w_bf16_p1");
24409            let shmem =
24410                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24411            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24412            f.set_attribute(
24413                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24414                shmem as i32,
24415            )?;
24416            let cfg = LaunchConfig {
24417                grid_dim: (
24418                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24419                    n_head as u32,
24420                    1,
24421                ),
24422                block_dim: (32, 4, 1),
24423                shared_mem_bytes: shmem,
24424            };
24425            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24426                head_dim as i32,
24427                n_head as i32,
24428                n_head_kv as i32,
24429                t as i32,
24430                t_kv as i32,
24431                causal as i32,
24432                window as i32,
24433            );
24434            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24435            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24436            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24437            let __s_b = self.gpu.stream();
24438            let mut b = __s_b.launch_builder(&f);
24439            b.arg(&qb)
24440                .arg(&kb)
24441                .arg(&vb)
24442                .arg(o)
24443                .arg(&hd)
24444                .arg(&nh)
24445                .arg(&nhkv)
24446                .arg(&ti)
24447                .arg(&tkvi)
24448                .arg(&scale)
24449                .arg(&cz)
24450                .arg(&wi);
24451            unsafe {
24452                b.launch(cfg)?;
24453            }
24454            return Ok(());
24455        }
24456        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
24457        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
24458        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24459        let g4 = !floor
24460            && !f32_stage
24461            && n_head_kv == 1
24462            && n_head.is_multiple_of(4)
24463            && *G4_ON.get_or_init(|| {
24464                std::env::var("MEMRA_FAW_G4")
24465                    .map(|v| v != "0")
24466                    .unwrap_or(true)
24467            });
24468        if g4 {
24469            const SP_M: usize = 16;
24470            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
24471            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
24472            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24473            let o2 = *O2_ON.get_or_init(|| {
24474                std::env::var("MEMRA_FAW_O2")
24475                    .map(|v| v != "0")
24476                    .unwrap_or(true)
24477            });
24478            let f = self.func(if o2 {
24479                "fa_prefill_w_bf16_g4o2"
24480            } else {
24481                "fa_prefill_w_bf16_g4"
24482            });
24483            let shmem = if o2 {
24484                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
24485            } else {
24486                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
24487                    as u32
24488            };
24489            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24490            f.set_attribute(
24491                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24492                shmem as i32,
24493            )?;
24494            let cfg = LaunchConfig {
24495                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
24496                block_dim: (32, 4, 1),
24497                shared_mem_bytes: shmem,
24498            };
24499            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24500                head_dim as i32,
24501                n_head as i32,
24502                n_head_kv as i32,
24503                t as i32,
24504                t_kv as i32,
24505                causal as i32,
24506                window as i32,
24507            );
24508            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24509            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24510            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24511            let __s_b = self.gpu.stream();
24512            let mut b = __s_b.launch_builder(&f);
24513            b.arg(&qb)
24514                .arg(&kb)
24515                .arg(&vb)
24516                .arg(o)
24517                .arg(&hd)
24518                .arg(&nh)
24519                .arg(&nhkv)
24520                .arg(&ti)
24521                .arg(&tkvi)
24522                .arg(&scale)
24523                .arg(&cz)
24524                .arg(&wi);
24525            unsafe {
24526                b.launch(cfg)?;
24527            }
24528            return Ok(());
24529        }
24530        let f = self.func(if floor {
24531            "fa_prefill_w_f32"
24532        } else if f32_stage {
24533            "fa_prefill_w_f32_pp"
24534        } else {
24535            "fa_prefill_w_bf16_pp"
24536        });
24537        let shmem =
24538            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24539        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24540        f.set_attribute(
24541            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24542            shmem as i32,
24543        )?;
24544        let cfg = LaunchConfig {
24545            grid_dim: (
24546                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24547                n_head as u32,
24548                1,
24549            ),
24550            block_dim: (32, 4, 1),
24551            shared_mem_bytes: shmem,
24552        };
24553        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24554            head_dim as i32,
24555            n_head as i32,
24556            n_head_kv as i32,
24557            t as i32,
24558            t_kv as i32,
24559            causal as i32,
24560            window as i32,
24561        );
24562        if f32_stage {
24563            let __s_b = self.gpu.stream();
24564            let mut b = __s_b.launch_builder(&f);
24565            b.arg(q)
24566                .arg(k)
24567                .arg(v)
24568                .arg(o)
24569                .arg(&hd)
24570                .arg(&nh)
24571                .arg(&nhkv)
24572                .arg(&ti)
24573                .arg(&tkvi)
24574                .arg(&scale)
24575                .arg(&cz)
24576                .arg(&wi);
24577            unsafe {
24578                b.launch(cfg)?;
24579            }
24580        } else {
24581            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24582            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24583            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24584            let __s_b = self.gpu.stream();
24585            let mut b = __s_b.launch_builder(&f);
24586            b.arg(&qb)
24587                .arg(&kb)
24588                .arg(&vb)
24589                .arg(o)
24590                .arg(&hd)
24591                .arg(&nh)
24592                .arg(&nhkv)
24593                .arg(&ti)
24594                .arg(&tkvi)
24595                .arg(&scale)
24596                .arg(&cz)
24597                .arg(&wi);
24598            unsafe {
24599                b.launch(cfg)?;
24600            }
24601        }
24602        Ok(())
24603    }
24604
24605    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
24606    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
24607    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
24608    #[allow(clippy::too_many_arguments)]
24609    pub fn fa_prefill_hd512(
24610        &self,
24611        q: &CudaSlice<f32>,
24612        k: &CudaSlice<f32>,
24613        v: &CudaSlice<f32>,
24614        o: &mut CudaSlice<f32>,
24615        head_dim: usize,
24616        n_head: usize,
24617        n_head_kv: usize,
24618        t: usize,
24619        t_kv: usize,
24620        scale: f32,
24621        causal: bool,
24622    ) -> Result<(), Box<dyn std::error::Error>> {
24623        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
24624        if portable_mma_gated() {
24625            return self.sdpa_naive(
24626                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
24627            );
24628        }
24629        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
24630        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
24631        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
24632        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
24633        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
24634        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24635        let f32_stage =
24636            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
24637        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
24638        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
24639        // Own numeric config (partial-sum order) — battery-gated.
24640        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24641        let sp = !f32_stage
24642            && *SP_ON.get_or_init(|| {
24643                std::env::var("MEMRA_FA512_SP")
24644                    .map(|v| v != "0")
24645                    .unwrap_or(true)
24646            });
24647        self.fa_prefill_hd512_arm(
24648            q,
24649            k,
24650            v,
24651            o,
24652            head_dim,
24653            n_head,
24654            n_head_kv,
24655            t,
24656            t_kv,
24657            scale,
24658            causal,
24659            f32_stage,
24660            sp,
24661            sp && fa_f16pv_on(),
24662        )
24663    }
24664
24665    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
24666    #[allow(clippy::too_many_arguments)]
24667    pub fn fa_prefill_hd512_pre(
24668        &self,
24669        qb: &CudaSlice<u8>,
24670        kb: &CudaSlice<u8>,
24671        vb: &CudaSlice<u8>,
24672        o: &mut CudaSlice<f32>,
24673        head_dim: usize,
24674        n_head: usize,
24675        n_head_kv: usize,
24676        t: usize,
24677        t_kv: usize,
24678        scale: f32,
24679        causal: bool,
24680        v_f16: bool,
24681    ) -> Result<(), Box<dyn std::error::Error>> {
24682        debug_assert_eq!(head_dim, 512);
24683        const SP_M: usize = 16;
24684        const BKS: usize = 32;
24685        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
24686        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
24687        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
24688        let f16pv = fa_f16pv_on();
24689        let nw = if f16pv { fa512_wide_warps() } else { 2 };
24690        let hp = f16pv
24691            && fa512_hp_on()
24692            && n_head.is_multiple_of(2)
24693            && (n_head / n_head_kv).is_multiple_of(2);
24694        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
24695        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
24696        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
24697            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
24698            let n = t_kv * n_head_kv * head_dim;
24699            let need = n * 2;
24700            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
24701                *vguard = Some(self.alloc_uninit::<u8>(need)?);
24702            }
24703            let dst = vguard.as_mut().unwrap();
24704            self.bf16_to_f16_into(vb, n, dst)?;
24705            vguard.as_ref().unwrap()
24706        } else {
24707            vb
24708        };
24709        let f = self.func(if hp {
24710            "fa_prefill_bf16_hd512_sp16h2"
24711        } else {
24712            match (f16pv, nw) {
24713                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
24714                (true, _) => "fa_prefill_bf16_hd512_sp16",
24715                _ => "fa_prefill_bf16_hd512_sp",
24716            }
24717        });
24718        let (nwarp, npart) = if hp {
24719            (4usize, 4usize)
24720        } else if nw > 2 {
24721            (nw, nw)
24722        } else {
24723            (2, 1)
24724        };
24725        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
24726        let shmem = if hp {
24727            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
24728                as u32
24729        } else {
24730            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
24731                + 4 * (npart * SP_M * BKS + SP_M)) as u32
24732        };
24733        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24734        f.set_attribute(
24735            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24736            shmem as i32,
24737        )?;
24738        let grid_y = if hp {
24739            (n_head / 2) as u32
24740        } else {
24741            n_head as u32
24742        };
24743        let cfg = LaunchConfig {
24744            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
24745            block_dim: (32, nwarp as u32, 1),
24746            shared_mem_bytes: shmem,
24747        };
24748        let (hd, nh, nhkv, ti, tkvi, cz) = (
24749            head_dim as i32,
24750            n_head as i32,
24751            n_head_kv as i32,
24752            t as i32,
24753            t_kv as i32,
24754            causal as i32,
24755        );
24756        let __s_b = self.gpu.stream();
24757        let mut b = __s_b.launch_builder(&f);
24758        b.arg(qb)
24759            .arg(kb)
24760            .arg(vref)
24761            .arg(o)
24762            .arg(&hd)
24763            .arg(&nh)
24764            .arg(&nhkv)
24765            .arg(&ti)
24766            .arg(&tkvi)
24767            .arg(&scale)
24768            .arg(&cz);
24769        unsafe {
24770            b.launch(cfg)?;
24771        }
24772        Ok(())
24773    }
24774
24775    /// Absorbed-form MLA prefill attention over a DSA-GATHERED index list, on tensor cores —
24776    /// the MEMRA_MLA_TC_PREFILL kernel (`fa_mla_gathered_bf16`, cu/flash_attn.cu). One CTA per
24777    /// (query, 16-head band); the query's index list is shared across heads (the DSA indexer
24778    /// mixes heads BEFORE top-k), which is exactly what gives the MMA its m axis. V is K
24779    /// (NoPE latent rows), so the kernel is `kv_rank == 512, d_rope == 0` ONLY and this
24780    /// launcher refuses anything else rather than approximate.
24781    #[allow(clippy::too_many_arguments)]
24782    pub fn mla_attn_gathered_tc(
24783        &self,
24784        q_lat_bf: &CudaSlice<u8>,   // [t_q, n_head, 512] bf16
24785        cache_bf: &CudaSlice<u8>,   // [t_kv, 512] bf16 latent rows
24786        idx: &CudaSlice<i32>,       // [t_q, width], ascending, -1 trailing
24787        o_lat: &mut CudaSlice<f32>, // [t_q, n_head, 512] f32
24788        n_head: usize,
24789        kv_rank: usize,
24790        t_q: usize,
24791        width: usize,
24792        scale: f32,
24793    ) -> Result<(), Box<dyn std::error::Error>> {
24794        if kv_rank != 512 {
24795            return Err(format!(
24796                "mla_attn_gathered_tc is stamped at kv_rank 512 (the glm5_next latent width); \
24797                 got {kv_rank} — the caller's door must fall back to the f32 gathered kernel"
24798            )
24799            .into());
24800        }
24801        if t_q == 0 || n_head == 0 {
24802            return Ok(());
24803        }
24804        const SP_M: usize = 16;
24805        const BKS: usize = 32;
24806        const HD: usize = 512;
24807        let f = self.func("fa_mla_gathered_bf16");
24808        // sQ + sK (V aliases K) + sP bf16, sS + sL f32, sIdx i32.
24809        let shmem =
24810            (2 * (SP_M * HD + BKS * HD + SP_M * BKS) + 4 * (SP_M * BKS + SP_M) + 4 * BKS) as u32;
24811        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24812        f.set_attribute(
24813            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24814            shmem as i32,
24815        )?;
24816        let cfg = LaunchConfig {
24817            grid_dim: (t_q as u32, (n_head as u32).div_ceil(SP_M as u32), 1),
24818            block_dim: (32, 2, 1),
24819            shared_mem_bytes: shmem,
24820        };
24821        let (nh, tq, w) = (n_head as i32, t_q as i32, width as i32);
24822        let __s_b = self.gpu.stream();
24823        let mut b = __s_b.launch_builder(&f);
24824        b.arg(q_lat_bf)
24825            .arg(cache_bf)
24826            .arg(idx)
24827            .arg(o_lat)
24828            .arg(&nh)
24829            .arg(&tq)
24830            .arg(&w)
24831            .arg(&scale);
24832        unsafe {
24833            b.launch(cfg)?;
24834        }
24835        Ok(())
24836    }
24837
24838    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
24839    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
24840    #[allow(clippy::too_many_arguments)]
24841    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
24842    pub fn fa_prefill_hd512_arm(
24843        &self,
24844        q: &CudaSlice<f32>,
24845        k: &CudaSlice<f32>,
24846        v: &CudaSlice<f32>,
24847        o: &mut CudaSlice<f32>,
24848        head_dim: usize,
24849        n_head: usize,
24850        n_head_kv: usize,
24851        t: usize,
24852        t_kv: usize,
24853        scale: f32,
24854        causal: bool,
24855        f32_stage: bool,
24856        sp: bool,
24857        f16pv: bool,
24858    ) -> Result<(), Box<dyn std::error::Error>> {
24859        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
24860        if sp && !f32_stage {
24861            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
24862            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
24863            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
24864            const SP_M: usize = 16;
24865            const BKS: usize = 32;
24866            let nw = if f16pv { fa512_wide_warps() } else { 2 };
24867            let hp = f16pv
24868                && fa512_hp_on()
24869                && n_head.is_multiple_of(2)
24870                && (n_head / n_head_kv).is_multiple_of(2);
24871            let f = self.func(if hp {
24872                "fa_prefill_bf16_hd512_sp16h2"
24873            } else {
24874                match (f16pv, nw) {
24875                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
24876                    (true, _) => "fa_prefill_bf16_hd512_sp16",
24877                    _ => "fa_prefill_bf16_hd512_sp",
24878                }
24879            });
24880            let (nwarp, npart) = if hp {
24881                (4usize, 4usize)
24882            } else if nw > 2 {
24883                (nw, nw)
24884            } else {
24885                (2, 1)
24886            };
24887            let shmem = if hp {
24888                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
24889                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
24890            } else {
24891                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
24892                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
24893            };
24894            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24895            f.set_attribute(
24896                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24897                shmem as i32,
24898            )?;
24899            let grid_y = if hp {
24900                (n_head / 2) as u32
24901            } else {
24902                n_head as u32
24903            };
24904            let cfg = LaunchConfig {
24905                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
24906                block_dim: (32, nwarp as u32, 1),
24907                shared_mem_bytes: shmem,
24908            };
24909            let (hd, nh, nhkv, ti, tkvi, cz) = (
24910                head_dim as i32,
24911                n_head as i32,
24912                n_head_kv as i32,
24913                t as i32,
24914                t_kv as i32,
24915                causal as i32,
24916            );
24917            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24918            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24919            let vb = if f16pv {
24920                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
24921            } else {
24922                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
24923            };
24924            let __s_b = self.gpu.stream();
24925            let mut b = __s_b.launch_builder(&f);
24926            b.arg(&qb)
24927                .arg(&kb)
24928                .arg(&vb)
24929                .arg(o)
24930                .arg(&hd)
24931                .arg(&nh)
24932                .arg(&nhkv)
24933                .arg(&ti)
24934                .arg(&tkvi)
24935                .arg(&scale)
24936                .arg(&cz);
24937            unsafe {
24938                b.launch(cfg)?;
24939            }
24940            return Ok(());
24941        }
24942        const BLOCK_Q: usize = 32;
24943        const BK: usize = 32;
24944        const HALF: usize = 256;
24945        let f = self.func(if f32_stage {
24946            "fa_prefill_f32_hd512"
24947        } else {
24948            "fa_prefill_bf16_hd512"
24949        });
24950        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
24951        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
24952            + 4 * BLOCK_Q) as u32;
24953        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24954        f.set_attribute(
24955            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24956            shmem as i32,
24957        )?;
24958        let cfg = LaunchConfig {
24959            grid_dim: (
24960                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24961                n_head as u32,
24962                2,
24963            ),
24964            block_dim: (32, 2, 1),
24965            shared_mem_bytes: shmem,
24966        };
24967        let (hd, nh, nhkv, ti, tkvi, cz) = (
24968            head_dim as i32,
24969            n_head as i32,
24970            n_head_kv as i32,
24971            t as i32,
24972            t_kv as i32,
24973            causal as i32,
24974        );
24975        if f32_stage {
24976            let __s_b = self.gpu.stream();
24977            let mut b = __s_b.launch_builder(&f);
24978            b.arg(q)
24979                .arg(k)
24980                .arg(v)
24981                .arg(o)
24982                .arg(&hd)
24983                .arg(&nh)
24984                .arg(&nhkv)
24985                .arg(&ti)
24986                .arg(&tkvi)
24987                .arg(&scale)
24988                .arg(&cz);
24989            unsafe {
24990                b.launch(cfg)?;
24991            }
24992        } else {
24993            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24994            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24995            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24996            let __s_b = self.gpu.stream();
24997            let mut b = __s_b.launch_builder(&f);
24998            b.arg(&qb)
24999                .arg(&kb)
25000                .arg(&vb)
25001                .arg(o)
25002                .arg(&hd)
25003                .arg(&nh)
25004                .arg(&nhkv)
25005                .arg(&ti)
25006                .arg(&tkvi)
25007                .arg(&scale)
25008                .arg(&cz);
25009            unsafe {
25010                b.launch(cfg)?;
25011            }
25012        }
25013        Ok(())
25014    }
25015
25016    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
25017    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
25018    /// separate f32_to_bf16 the FA entries would run).
25019    #[allow(clippy::too_many_arguments)]
25020    pub fn rope_neox2_bf16e(
25021        &self,
25022        q: &mut CudaSlice<f32>,
25023        k: &mut CudaSlice<f32>,
25024        qb: &mut CudaSlice<u8>,
25025        kb: &mut CudaSlice<u8>,
25026        pos: &CudaSlice<i32>,
25027        head_dim: usize,
25028        n_dims: usize,
25029        nh_q: usize,
25030        nh_k: usize,
25031        n_tokens: usize,
25032        base: f32,
25033        freq_scale: f32,
25034        ff: Option<&CudaSlice<f32>>,
25035    ) -> Result<(), Box<dyn std::error::Error>> {
25036        let f = self.func("rope_neox2_bf16e_f32");
25037        let rows = ((nh_q + nh_k) * n_tokens) as u32;
25038        let cfg = LaunchConfig {
25039            grid_dim: (rows, 1, 1),
25040            block_dim: ((head_dim / 2) as u32, 1, 1),
25041            shared_mem_bytes: 0,
25042        };
25043        let theta_scale = base.powf(-2.0 / n_dims as f32);
25044        let (hd, nd, nhq, nhk, nt) = (
25045            head_dim as i32,
25046            n_dims as i32,
25047            nh_q as i32,
25048            nh_k as i32,
25049            n_tokens as i32,
25050        );
25051        let __s_b = self.gpu.stream();
25052        let mut b = __s_b.launch_builder(&f);
25053        match ff {
25054            Some(t) => {
25055                b.arg(&mut *q)
25056                    .arg(&mut *k)
25057                    .arg(&mut *qb)
25058                    .arg(&mut *kb)
25059                    .arg(pos)
25060                    .arg(&hd)
25061                    .arg(&nd)
25062                    .arg(&nhq)
25063                    .arg(&nhk)
25064                    .arg(&nt)
25065                    .arg(&theta_scale)
25066                    .arg(&freq_scale)
25067                    .arg(t);
25068                unsafe {
25069                    b.launch(cfg)?;
25070                }
25071            }
25072            None => {
25073                let null: u64 = 0;
25074                b.arg(&mut *q)
25075                    .arg(&mut *k)
25076                    .arg(&mut *qb)
25077                    .arg(&mut *kb)
25078                    .arg(pos)
25079                    .arg(&hd)
25080                    .arg(&nd)
25081                    .arg(&nhq)
25082                    .arg(&nhk)
25083                    .arg(&nt)
25084                    .arg(&theta_scale)
25085                    .arg(&freq_scale)
25086                    .arg(&null);
25087                unsafe {
25088                    b.launch(cfg)?;
25089                }
25090            }
25091        }
25092        Ok(())
25093    }
25094
25095    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
25096    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
25097    pub fn f32_to_bf16(
25098        &self,
25099        x: &CudaSlice<f32>,
25100        n: usize,
25101    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25102        assert!(
25103            n.is_multiple_of(4),
25104            "f32_to_bf16 requires n % 4 == 0, got {n}"
25105        );
25106        let mut y = self.alloc_uninit::<u8>(n * 2)?;
25107        let f = self.func("f32_to_bf16_flat");
25108        let n_i = n as i64;
25109        let cfg = LaunchConfig {
25110            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
25111            block_dim: (256, 1, 1),
25112            shared_mem_bytes: 0,
25113        };
25114        let __s_b = self.gpu.stream();
25115        let mut b = __s_b.launch_builder(&f);
25116        b.arg(x).arg(&mut y).arg(&n_i);
25117        unsafe {
25118            b.launch(cfg)?;
25119        }
25120        Ok(y)
25121    }
25122
25123    pub fn f32_to_f16(
25124        &self,
25125        x: &CudaSlice<f32>,
25126        n: usize,
25127    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25128        assert!(
25129            n.is_multiple_of(4),
25130            "f32_to_f16 requires n % 4 == 0, got {n}"
25131        );
25132        let mut y = self.alloc_uninit::<u8>(n * 2)?;
25133        let f = self.func("f32_to_f16_flat");
25134        let n_i = n as i64;
25135        let cfg = LaunchConfig {
25136            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
25137            block_dim: (256, 1, 1),
25138            shared_mem_bytes: 0,
25139        };
25140        let __s_b = self.gpu.stream();
25141        let mut b = __s_b.launch_builder(&f);
25142        b.arg(x).arg(&mut y).arg(&n_i);
25143        unsafe {
25144            b.launch(cfg)?;
25145        }
25146        Ok(y)
25147    }
25148
25149    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
25150    pub fn bf16_to_f16(
25151        &self,
25152        xb: &CudaSlice<u8>,
25153        n: usize,
25154    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25155        let mut y = self.alloc_uninit::<u8>(n * 2)?;
25156        self.bf16_to_f16_into(xb, n, &mut y)?;
25157        Ok(y)
25158    }
25159
25160    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
25161    pub fn bf16_to_f16_into(
25162        &self,
25163        xb: &CudaSlice<u8>,
25164        n: usize,
25165        y: &mut CudaSlice<u8>,
25166    ) -> Result<(), Box<dyn std::error::Error>> {
25167        assert!(
25168            n.is_multiple_of(2),
25169            "bf16_to_f16 requires n % 2 == 0, got {n}"
25170        );
25171        assert!(y.len() >= n * 2);
25172        let f = self.func("bf16_to_f16_flat");
25173        let n2 = (n / 2) as i64;
25174        let cfg = LaunchConfig {
25175            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
25176            block_dim: (256, 1, 1),
25177            shared_mem_bytes: 0,
25178        };
25179        let __s_b = self.gpu.stream();
25180        let mut b = __s_b.launch_builder(&f);
25181        b.arg(xb).arg(y).arg(&n2);
25182        unsafe {
25183            b.launch(cfg)?;
25184        }
25185        Ok(())
25186    }
25187
25188    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
25189    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
25190    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
25191    /// head_dim in {256, 128}, bf16kv lane on.
25192    #[allow(clippy::too_many_arguments)]
25193    pub fn fa_prefill_vl8(
25194        &self,
25195        seqs: &[FaSeqVl],
25196        head_dim: usize,
25197        n_head: usize,
25198        n_head_kv: usize,
25199        scale: f32,
25200    ) -> Result<(), Box<dyn std::error::Error>> {
25201        const BK: usize = 32;
25202        let b = seqs.len();
25203        assert!((1..=8).contains(&b));
25204        let mut packed = [FaSeqVl::default(); 8];
25205        packed[..b].copy_from_slice(seqs);
25206        let v = FaVl8(packed);
25207        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25208        let ept = (n_head_kv * head_dim) as i32;
25209        {
25210            let f = self.func("fa_mirror_vl");
25211            let max_n = (max_t as i64) * ept as i64;
25212            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25213            for which in 0..2i32 {
25214                let cfg = LaunchConfig {
25215                    grid_dim: (blocks, 1, b as u32),
25216                    block_dim: (256, 1, 1),
25217                    shared_mem_bytes: 0,
25218                };
25219                let __s_lb = self.gpu.stream();
25220                let mut lb = __s_lb.launch_builder(&f);
25221                lb.arg(&v).arg(&ept).arg(&which);
25222                unsafe {
25223                    lb.launch(cfg)?;
25224                }
25225            }
25226        }
25227        let hd_sfx = fa_hd_suffix(head_dim)?;
25228        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
25229        let block_q = 64usize;
25230        let kv_stages = 2usize;
25231        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
25232            + 4 * (block_q * BK + 2 * block_q)) as u32;
25233        use cudarc::driver::sys::CUfunction_attribute_enum as A;
25234        f.set_attribute(
25235            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25236            shmem as i32,
25237        )?;
25238        let cfg = LaunchConfig {
25239            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
25240            block_dim: (32, 4, 1),
25241            shared_mem_bytes: shmem,
25242        };
25243        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
25244        let __s_lb = self.gpu.stream();
25245        let mut lb = __s_lb.launch_builder(&f);
25246        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
25247        unsafe {
25248            lb.launch(cfg)?;
25249        }
25250        Ok(())
25251    }
25252
25253    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
25254    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
25255    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
25256    #[allow(clippy::too_many_arguments)]
25257    pub fn attn_pre_vl8(
25258        &self,
25259        seqs: &[AttnPreVl],
25260        wq: &CudaSlice<f32>,
25261        wk: &CudaSlice<f32>,
25262        head_dim: usize,
25263        rope_dims: usize,
25264        n_head: usize,
25265        n_head_kv: usize,
25266        eps: f32,
25267        freq_base: f32,
25268        freq_scale: f32,
25269        kv_dim_k: usize,
25270        kv_dim_v: usize,
25271        k_tok_bytes: usize,
25272        v_tok_bytes: usize,
25273    ) -> Result<(), Box<dyn std::error::Error>> {
25274        let b = seqs.len();
25275        assert!((1..=8).contains(&b));
25276        let mut packed = [AttnPreVl::default(); 8];
25277        packed[..b].copy_from_slice(seqs);
25278        let v = AttnPreVl8(packed);
25279        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25280        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
25281        {
25282            let f = self.func("q_gate_split_vl");
25283            let n = max_t * (n_head * head_dim) as u32;
25284            let cfg = LaunchConfig {
25285                grid_dim: (n.div_ceil(256), 1, b as u32),
25286                block_dim: (256, 1, 1),
25287                shared_mem_bytes: 0,
25288            };
25289            let __s_lb = self.gpu.stream();
25290            let mut lb = __s_lb.launch_builder(&f);
25291            lb.arg(&v).arg(&hd).arg(&nh);
25292            unsafe {
25293                lb.launch(cfg)?;
25294            }
25295        }
25296        {
25297            let f = self.func("attn_rms_vl");
25298            let cfg = LaunchConfig {
25299                grid_dim: (max_t * n_head as u32, 2, b as u32),
25300                block_dim: (rms_block(), 1, 1),
25301                shared_mem_bytes: 0,
25302            };
25303            let __s_lb = self.gpu.stream();
25304            let mut lb = __s_lb.launch_builder(&f);
25305            lb.arg(&v)
25306                .arg(wq)
25307                .arg(wk)
25308                .arg(&hd)
25309                .arg(&nh)
25310                .arg(&nhkv)
25311                .arg(&eps);
25312            unsafe {
25313                lb.launch(cfg)?;
25314            }
25315        }
25316        {
25317            let f = self.func("attn_rope_vl");
25318            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
25319            let nd = rope_dims as i32;
25320            let cfg = LaunchConfig {
25321                grid_dim: (max_t * n_head as u32, 2, b as u32),
25322                block_dim: ((head_dim / 2) as u32, 1, 1),
25323                shared_mem_bytes: 0,
25324            };
25325            let __s_lb = self.gpu.stream();
25326            let mut lb = __s_lb.launch_builder(&f);
25327            lb.arg(&v)
25328                .arg(&hd)
25329                .arg(&nd)
25330                .arg(&nh)
25331                .arg(&nhkv)
25332                .arg(&theta_scale)
25333                .arg(&freq_scale);
25334            unsafe {
25335                lb.launch(cfg)?;
25336            }
25337        }
25338        {
25339            let f = self.func("append_kv_vl");
25340            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
25341            let cfg = LaunchConfig {
25342                grid_dim: (nblk, max_t, b as u32),
25343                block_dim: (32, 1, 1),
25344                shared_mem_bytes: 0,
25345            };
25346            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
25347            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25348            let __s_lb = self.gpu.stream();
25349            let mut lb = __s_lb.launch_builder(&f);
25350            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
25351            unsafe {
25352                lb.launch(cfg)?;
25353            }
25354        }
25355        Ok(())
25356    }
25357
25358    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
25359    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
25360    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
25361    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
25362    #[allow(clippy::too_many_arguments)]
25363    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25364    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25365    pub fn fa_prefill_view(
25366        &self,
25367        q: &CudaSlice<f32>,
25368        k: &cudarc::driver::CudaView<u8>,
25369        v: &cudarc::driver::CudaView<u8>,
25370        o: &mut CudaSlice<f32>,
25371        head_dim: usize,
25372        n_head: usize,
25373        n_head_kv: usize,
25374        t: usize,
25375        t_kv: usize,
25376        scale: f32,
25377        causal: bool,
25378        k_tok_bytes: usize,
25379        v_tok_bytes: usize,
25380        g: bool,
25381    ) -> Result<(), Box<dyn std::error::Error>> {
25382        if portable_mma_gated() {
25383            return self.sdpa_naive_quantized_view(
25384                q,
25385                k,
25386                v,
25387                o,
25388                head_dim,
25389                n_head,
25390                n_head_kv,
25391                t,
25392                t_kv,
25393                scale,
25394                causal,
25395                k_tok_bytes,
25396                v_tok_bytes,
25397            );
25398        }
25399        const BLOCK_Q: usize = 64;
25400        const BK: usize = 32;
25401        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
25402        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
25403        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
25404        let f = if g {
25405            self.func_g(&name)
25406        } else {
25407            self.func(&name)
25408        };
25409        let shmem =
25410            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
25411        use cudarc::driver::sys::CUfunction_attribute_enum as A;
25412        f.set_attribute(
25413            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25414            shmem as i32,
25415        )?;
25416        let cfg = LaunchConfig {
25417            grid_dim: (
25418                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25419                n_head as u32,
25420                1,
25421            ),
25422            block_dim: (32, 4, 1),
25423            shared_mem_bytes: shmem,
25424        };
25425        let (hd, nh, nhkv, ti, tkvi, cz) = (
25426            head_dim as i32,
25427            n_head as i32,
25428            n_head_kv as i32,
25429            t as i32,
25430            t_kv as i32,
25431            causal as i32,
25432        );
25433        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25434        let __s_b = self.gpu.stream();
25435        let mut b = __s_b.launch_builder(&f);
25436        b.arg(q)
25437            .arg(k)
25438            .arg(v)
25439            .arg(o)
25440            .arg(&hd)
25441            .arg(&nh)
25442            .arg(&nhkv)
25443            .arg(&ti)
25444            .arg(&tkvi)
25445            .arg(&scale)
25446            .arg(&cz)
25447            .arg(&ktb)
25448            .arg(&vtb);
25449        unsafe {
25450            b.launch(cfg)?;
25451        }
25452        Ok(())
25453    }
25454
25455    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
25456    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
25457    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
25458    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
25459    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
25460    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
25461    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
25462    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
25463    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
25464    #[allow(clippy::too_many_arguments)]
25465    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25466    pub fn fa_prefill_view_ws(
25467        &self,
25468        q: &CudaSlice<f32>,
25469        k: &cudarc::driver::CudaView<u8>,
25470        v: &cudarc::driver::CudaView<u8>,
25471        o: &mut CudaSlice<f32>,
25472        head_dim: usize,
25473        n_head: usize,
25474        n_head_kv: usize,
25475        t: usize,
25476        t_kv: usize,
25477        scale: f32,
25478        causal: bool,
25479        k_tok_bytes: usize,
25480        v_tok_bytes: usize,
25481        g: bool,
25482    ) -> Result<(), Box<dyn std::error::Error>> {
25483        if portable_mma_gated() {
25484            return self.sdpa_naive_quantized_view(
25485                q,
25486                k,
25487                v,
25488                o,
25489                head_dim,
25490                n_head,
25491                n_head_kv,
25492                t,
25493                t_kv,
25494                scale,
25495                causal,
25496                k_tok_bytes,
25497                v_tok_bytes,
25498            );
25499        }
25500        const BLOCK_Q: usize = 64;
25501        const BK: usize = 32;
25502        let kv_dim_k = n_head_kv * head_dim;
25503        let kv_dim_v = n_head_kv * head_dim;
25504        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
25505        let v_ws_bytes = t_kv * kv_dim_v * 2;
25506        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
25507        let mut guard = self.prime_deqw_ws.lock().unwrap();
25508        let need_grow = match guard.as_ref() {
25509            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
25510            None => true,
25511        };
25512        if need_grow {
25513            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
25514            let (ck, cv) = guard
25515                .as_ref()
25516                .map(|(a, b)| (a.len(), b.len()))
25517                .unwrap_or((0, 0));
25518            *guard = Some((
25519                self.alloc_u8(grow(ck, k_ws_bytes))?,
25520                self.alloc_u8(grow(cv, v_ws_bytes))?,
25521            ));
25522        }
25523        let (kw, vw) = guard.as_mut().unwrap();
25524        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
25525        {
25526            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
25527            let f = if g {
25528                self.func_g("fa_dequant_kv_ws_bf16")
25529            } else {
25530                self.func("fa_dequant_kv_ws_bf16")
25531            };
25532            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
25533            #[allow(clippy::manual_div_ceil)]
25534            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25535            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
25536            let cfg = LaunchConfig {
25537                grid_dim: (nblk.max(1), 1, 1),
25538                block_dim: (256, 1, 1),
25539                shared_mem_bytes: 0,
25540            };
25541            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
25542            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25543            let __s_b = self.gpu.stream();
25544            let mut b = __s_b.launch_builder(&f);
25545            b.arg(k)
25546                .arg(v)
25547                .arg(&mut *kw)
25548                .arg(&mut *vw)
25549                .arg(&kdk)
25550                .arg(&kdv)
25551                .arg(&tkvi)
25552                .arg(&ktb)
25553                .arg(&vtb);
25554            unsafe {
25555                b.launch(cfg)?;
25556            }
25557        }
25558        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
25559        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
25560        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
25561        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
25562        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
25563        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
25564        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
25565        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
25566            .map(|v| v != "0")
25567            .unwrap_or(true);
25568        {
25569            let hd_sfx = fa_hd_suffix(head_dim)?;
25570            let f = self.func(&format!(
25571                "fa_prefill_qw{}{hd_sfx}",
25572                if db { "_db" } else { "" }
25573            ));
25574            let shmem = if db {
25575                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
25576                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
25577            } else {
25578                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
25579            };
25580            use cudarc::driver::sys::CUfunction_attribute_enum as A;
25581            f.set_attribute(
25582                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25583                shmem as i32,
25584            )?;
25585            let cfg = LaunchConfig {
25586                grid_dim: (
25587                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25588                    n_head as u32,
25589                    1,
25590                ),
25591                block_dim: (32, 4, 1),
25592                shared_mem_bytes: shmem,
25593            };
25594            let (hd, nh, nhkv, ti, tkvi, cz) = (
25595                head_dim as i32,
25596                n_head as i32,
25597                n_head_kv as i32,
25598                t as i32,
25599                t_kv as i32,
25600                causal as i32,
25601            );
25602            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
25603            let __s_b = self.gpu.stream();
25604            let mut b = __s_b.launch_builder(&f);
25605            b.arg(q)
25606                .arg(&*kw)
25607                .arg(&*vw)
25608                .arg(o)
25609                .arg(&hd)
25610                .arg(&nh)
25611                .arg(&nhkv)
25612                .arg(&ti)
25613                .arg(&tkvi)
25614                .arg(&scale)
25615                .arg(&cz)
25616                .arg(&kdk)
25617                .arg(&kdv);
25618            unsafe {
25619                b.launch(cfg)?;
25620            }
25621        }
25622        Ok(())
25623    }
25624
25625    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
25626    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
25627    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
25628    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
25629    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
25630    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
25631    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
25632    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
25633    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
25634    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
25635    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
25636    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
25637    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
25638    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
25639    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
25640    #[allow(clippy::too_many_arguments)]
25641    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25642    pub fn fa_prefill_view_ws_w_hd128(
25643        &self,
25644        q: &CudaSlice<f32>,
25645        k: &cudarc::driver::CudaView<u8>,
25646        v: &cudarc::driver::CudaView<u8>,
25647        o: &mut CudaSlice<f32>,
25648        head_dim: usize,
25649        n_head: usize,
25650        n_head_kv: usize,
25651        t: usize,
25652        t_kv: usize,
25653        scale: f32,
25654        causal: bool,
25655        window: usize,
25656        k_tok_bytes: usize,
25657        v_tok_bytes: usize,
25658    ) -> Result<(), Box<dyn std::error::Error>> {
25659        assert_eq!(
25660            head_dim, 128,
25661            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
25662        );
25663        if portable_mma_gated() {
25664            return self.sdpa_naive_w_quantized_view(
25665                q,
25666                k,
25667                v,
25668                o,
25669                head_dim,
25670                n_head,
25671                n_head_kv,
25672                t,
25673                t_kv,
25674                scale,
25675                causal,
25676                window,
25677                k_tok_bytes,
25678                v_tok_bytes,
25679            );
25680        }
25681        const BLOCK_Q: usize = 64;
25682        const BK: usize = 32;
25683        let kv_dim_k = n_head_kv * head_dim;
25684        let kv_dim_v = n_head_kv * head_dim;
25685        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
25686        let v_ws_bytes = t_kv * kv_dim_v * 2;
25687        let mut guard = self.prime_deqw_ws.lock().unwrap();
25688        let need_grow = match guard.as_ref() {
25689            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
25690            None => true,
25691        };
25692        if need_grow {
25693            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
25694            let (ck, cv) = guard
25695                .as_ref()
25696                .map(|(a, b)| (a.len(), b.len()))
25697                .unwrap_or((0, 0));
25698            *guard = Some((
25699                self.alloc_u8(grow(ck, k_ws_bytes))?,
25700                self.alloc_u8(grow(cv, v_ws_bytes))?,
25701            ));
25702        }
25703        let (kw, vw) = guard.as_mut().unwrap();
25704        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
25705        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
25706        {
25707            let f = self.func("fa_dequant_kv_ws_bf16");
25708            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
25709            #[allow(clippy::manual_div_ceil)]
25710            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25711            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
25712            let cfg = LaunchConfig {
25713                grid_dim: (nblk.max(1), 1, 1),
25714                block_dim: (256, 1, 1),
25715                shared_mem_bytes: 0,
25716            };
25717            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
25718            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25719            let __s_b = self.gpu.stream();
25720            let mut b = __s_b.launch_builder(&f);
25721            b.arg(k)
25722                .arg(v)
25723                .arg(&mut *kw)
25724                .arg(&mut *vw)
25725                .arg(&kdk)
25726                .arg(&kdv)
25727                .arg(&tkvi)
25728                .arg(&ktb)
25729                .arg(&vtb);
25730            unsafe {
25731                b.launch(cfg)?;
25732            }
25733        }
25734        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
25735        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
25736            .map(|v| v != "0")
25737            .unwrap_or(true);
25738        {
25739            let f = self.func(if db {
25740                "fa_prefill_qw_db_w_hd128"
25741            } else {
25742                "fa_prefill_qw_w_hd128"
25743            });
25744            let shmem = if db {
25745                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
25746            } else {
25747                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
25748            };
25749            use cudarc::driver::sys::CUfunction_attribute_enum as A;
25750            f.set_attribute(
25751                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25752                shmem as i32,
25753            )?;
25754            let cfg = LaunchConfig {
25755                grid_dim: (
25756                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25757                    n_head as u32,
25758                    1,
25759                ),
25760                block_dim: (32, 4, 1),
25761                shared_mem_bytes: shmem,
25762            };
25763            let (hd, nh, nhkv, ti, tkvi, cz) = (
25764                head_dim as i32,
25765                n_head as i32,
25766                n_head_kv as i32,
25767                t as i32,
25768                t_kv as i32,
25769                causal as i32,
25770            );
25771            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
25772            let __s_b = self.gpu.stream();
25773            let mut b = __s_b.launch_builder(&f);
25774            b.arg(q)
25775                .arg(&*kw)
25776                .arg(&*vw)
25777                .arg(o)
25778                .arg(&hd)
25779                .arg(&nh)
25780                .arg(&nhkv)
25781                .arg(&ti)
25782                .arg(&tkvi)
25783                .arg(&scale)
25784                .arg(&cz)
25785                .arg(&kdk)
25786                .arg(&kdv)
25787                .arg(&wnd);
25788            unsafe {
25789                b.launch(cfg)?;
25790            }
25791        }
25792        Ok(())
25793    }
25794
25795    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
25796    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
25797    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
25798    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25799    pub fn fa_decode(
25800        &self,
25801        q: &CudaSlice<f32>,
25802        k: &cudarc::driver::CudaView<u8>,
25803        v: &cudarc::driver::CudaView<u8>,
25804        o: &mut CudaSlice<f32>,
25805        head_dim: usize,
25806        n_head: usize,
25807        n_head_kv: usize,
25808        t_kv: usize,
25809        scale: f32,
25810        k_tok_bytes: usize,
25811        v_tok_bytes: usize,
25812    ) -> Result<(), Box<dyn std::error::Error>> {
25813        self.fa_decode_kvmod(
25814            q,
25815            k,
25816            v,
25817            o,
25818            head_dim,
25819            n_head,
25820            n_head_kv,
25821            t_kv,
25822            scale,
25823            k_tok_bytes,
25824            v_tok_bytes,
25825            false,
25826        )
25827    }
25828
25829    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
25830    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
25831    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
25832    #[allow(clippy::too_many_arguments)]
25833    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
25834    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
25835    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
25836    #[allow(clippy::too_many_arguments)]
25837    #[allow(clippy::too_many_arguments)]
25838    fn fa_decode_scalar_unified(
25839        &self,
25840        q: &cudarc::driver::CudaView<f32>,
25841        k: &cudarc::driver::CudaView<u8>,
25842        v: &cudarc::driver::CudaView<u8>,
25843        o: &mut cudarc::driver::CudaViewMut<f32>,
25844        head_dim: usize,
25845        n_head: usize,
25846        n_head_kv: usize,
25847        t_kv_host: usize,
25848        t_kv_dev: Option<&CudaSlice<i32>>,
25849        scale: f32,
25850        n_splits: usize,
25851        split_keys: usize,
25852        k_tok_bytes: usize,
25853        v_tok_bytes: usize,
25854        g: bool,
25855        part_o: &mut CudaSlice<f32>,
25856        part_m: &mut CudaSlice<f32>,
25857        part_l: &mut CudaSlice<f32>,
25858        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
25859    ) -> Result<(), Box<dyn std::error::Error>> {
25860        let f = if g {
25861            self.func_g("fa_decode_f32")
25862        } else {
25863            self.fa_func("fa_decode_f32", head_dim)
25864        };
25865        let cfg = LaunchConfig {
25866            grid_dim: (n_head as u32, n_splits as u32, 1),
25867            block_dim: (head_dim as u32, 1, 1),
25868            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
25869        };
25870        let (hd, nh, nhkv, nsp) = (
25871            head_dim as i32,
25872            n_head as i32,
25873            n_head_kv as i32,
25874            n_splits as i32,
25875        );
25876        let (ktb, vtb, tkvi, ski) = (
25877            k_tok_bytes as i64,
25878            v_tok_bytes as i64,
25879            t_kv_host as i32,
25880            split_keys as i32,
25881        );
25882        let __s_b = self.gpu.stream();
25883        let mut b = __s_b.launch_builder(&f);
25884        match t_kv_dev {
25885            Some(d) => {
25886                b.arg(q)
25887                    .arg(k)
25888                    .arg(v)
25889                    .arg(&mut *part_o)
25890                    .arg(&mut *part_m)
25891                    .arg(&mut *part_l)
25892                    .arg(&hd)
25893                    .arg(&nh)
25894                    .arg(&nhkv)
25895                    .arg(&tkvi)
25896                    .arg(d)
25897                    .arg(&scale)
25898                    .arg(&nsp)
25899                    .arg(&ski)
25900                    .arg(&ktb)
25901                    .arg(&vtb);
25902                unsafe {
25903                    b.launch(cfg)?;
25904                }
25905            }
25906            None => {
25907                let null: u64 = 0;
25908                b.arg(q)
25909                    .arg(k)
25910                    .arg(v)
25911                    .arg(&mut *part_o)
25912                    .arg(&mut *part_m)
25913                    .arg(&mut *part_l)
25914                    .arg(&hd)
25915                    .arg(&nh)
25916                    .arg(&nhkv)
25917                    .arg(&tkvi)
25918                    .arg(&null)
25919                    .arg(&scale)
25920                    .arg(&nsp)
25921                    .arg(&ski)
25922                    .arg(&ktb)
25923                    .arg(&vtb);
25924                unsafe {
25925                    b.launch(cfg)?;
25926                }
25927            }
25928        }
25929        let cfg2 = LaunchConfig {
25930            grid_dim: (n_head as u32, 1, 1),
25931            block_dim: (head_dim as u32, 1, 1),
25932            shared_mem_bytes: 0,
25933        };
25934        if let Some((oq, od)) = q8_out {
25935            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
25936            let fc = if g {
25937                self.func_g("fa_decode_combine_q8_1")
25938            } else {
25939                self.fa_func("fa_decode_combine_q8_1", head_dim)
25940            };
25941            let __s_b2 = self.gpu.stream();
25942            let mut b2 = __s_b2.launch_builder(&fc);
25943            b2.arg(&*part_o)
25944                .arg(&*part_m)
25945                .arg(&*part_l)
25946                .arg(oq)
25947                .arg(od)
25948                .arg(&hd)
25949                .arg(&nh)
25950                .arg(&nsp);
25951            unsafe {
25952                b2.launch(cfg2)?;
25953            }
25954            return Ok(());
25955        }
25956        let fc = if g {
25957            self.func_g("fa_decode_combine_f32")
25958        } else {
25959            self.fa_func("fa_decode_combine_f32", head_dim)
25960        };
25961        let __s_b2 = self.gpu.stream();
25962        let mut b2 = __s_b2.launch_builder(&fc);
25963        b2.arg(&*part_o)
25964            .arg(&*part_m)
25965            .arg(&*part_l)
25966            .arg(o)
25967            .arg(&hd)
25968            .arg(&nh)
25969            .arg(&nsp);
25970        unsafe {
25971            b2.launch(cfg2)?;
25972        }
25973        Ok(())
25974    }
25975
25976    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25977    pub fn fa_decode_kvmod(
25978        &self,
25979        q: &CudaSlice<f32>,
25980        k: &cudarc::driver::CudaView<u8>,
25981        v: &cudarc::driver::CudaView<u8>,
25982        o: &mut CudaSlice<f32>,
25983        head_dim: usize,
25984        n_head: usize,
25985        n_head_kv: usize,
25986        t_kv: usize,
25987        scale: f32,
25988        k_tok_bytes: usize,
25989        v_tok_bytes: usize,
25990        g: bool,
25991    ) -> Result<(), Box<dyn std::error::Error>> {
25992        let q_view = q.as_view();
25993        let mut o_view = o.as_view_mut();
25994        self.fa_decode_kvmod_view(
25995            &q_view,
25996            k,
25997            v,
25998            &mut o_view,
25999            head_dim,
26000            n_head,
26001            n_head_kv,
26002            t_kv,
26003            scale,
26004            k_tok_bytes,
26005            v_tok_bytes,
26006            g,
26007        )
26008    }
26009
26010    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
26011    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
26012    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
26013    /// per-session KV view and FA launch.
26014    #[allow(clippy::too_many_arguments)]
26015    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26016    pub fn fa_decode_kvmod_view(
26017        &self,
26018        q: &cudarc::driver::CudaView<f32>,
26019        k: &cudarc::driver::CudaView<u8>,
26020        v: &cudarc::driver::CudaView<u8>,
26021        o: &mut cudarc::driver::CudaViewMut<f32>,
26022        head_dim: usize,
26023        n_head: usize,
26024        n_head_kv: usize,
26025        t_kv: usize,
26026        scale: f32,
26027        k_tok_bytes: usize,
26028        v_tok_bytes: usize,
26029        g: bool,
26030    ) -> Result<(), Box<dyn std::error::Error>> {
26031        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
26032        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
26033        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
26034        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
26035        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
26036        //
26037        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
26038        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
26039        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
26040        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
26041        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
26042        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
26043        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
26044        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
26045        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
26046        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
26047        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
26048        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
26049        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
26050        // fall to the exact scalar there instead of the broken register arm.
26051        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
26052        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
26053        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
26054        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
26055        if g && head_dim == 256 && !fa_v4_at(t_kv) {
26056            fa_vec = false;
26057        }
26058        let sp = fa_split_keys(t_kv, n_head_kv);
26059        let n_splits = if fa_vec {
26060            ((t_kv + sp - 1) / sp).max(1)
26061        } else {
26062            ((t_kv + 255) / 256).max(1)
26063        };
26064        let o_len = n_head * n_splits * head_dim;
26065        let ml_len = n_head * n_splits;
26066        let mut part_guard = self.fa_part_pool.lock().unwrap();
26067        if part_guard
26068            .as_ref()
26069            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26070            .unwrap_or(true)
26071        {
26072            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26073            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26074            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26075            // later live allocations land at those addresses, and the next graph REPLAY writes
26076            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26077            // output corruption began the burst after the trunk's t_kv growth first realloc'd
26078            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26079            // the baked addresses alive (single-stream: eager writes the new buffers, replays
26080            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26081            // (total retired < final size).
26082            let old = part_guard.take();
26083            let (co, cm) = old
26084                .as_ref()
26085                .map(|pp| (pp.0.len(), pp.1.len()))
26086                .unwrap_or((0, 0));
26087            if let Some(old) = old {
26088                self.fa_part_retired.lock().unwrap().push(old);
26089            }
26090            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26091                eprintln!(
26092                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26093                    co, o_len, cm, ml_len
26094                );
26095            }
26096            *part_guard =
26097                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
26098        }
26099        let pg = part_guard.as_mut().unwrap();
26100        self.gpu
26101            .stream()
26102            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
26103        self.gpu
26104            .stream()
26105            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
26106        self.gpu
26107            .stream()
26108            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
26109        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
26110        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
26111        let (hd, nh, nhkv, tkvi, nsp) = (
26112            head_dim as i32,
26113            n_head as i32,
26114            n_head_kv as i32,
26115            t_kv as i32,
26116            n_splits as i32,
26117        );
26118        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26119        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
26120        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
26121        // silently truncating the accumulator.
26122        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
26123        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
26124        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
26125        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
26126        // 178.4 -> 173.7 when 512 rode vec unconditionally).
26127        let fa512_min = fa512_min_tkv();
26128        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
26129        // g-module keeps the v4 pick (its class is not the depth-decay class).
26130        let deep = fa_vec
26131            && head_dim == 256
26132            && fa_v4_at(t_kv)
26133            && !g
26134            && fa_deep_at(t_kv)
26135            && !matches!(fa_v4_mode(), "noB3" | "stage");
26136        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
26137            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
26138            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
26139            let gqa = (n_head / n_head_kv).max(1) as u32;
26140            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
26141            (
26142                fv,
26143                LaunchConfig {
26144                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26145                    block_dim: (32, gqa, 1),
26146                    shared_mem_bytes: 0,
26147                },
26148            )
26149        } else if fa_vec && head_dim <= 256 {
26150            let gqa = (n_head / n_head_kv).max(1) as u32;
26151            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
26152            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
26153            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
26154            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
26155            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
26156            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
26157            // dequant each tile ONCE per block.
26158            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
26159            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
26160            // there by 12x — latency, not bandwidth, rules small KV).
26161            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26162            let smem_tkv = *SMEM_TKV.get_or_init(|| {
26163                std::env::var("MEMRA_FA_SMEM_TKV")
26164                    .ok()
26165                    .and_then(|v| v.parse().ok())
26166                    .unwrap_or_else(|| {
26167                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
26168                    })
26169            });
26170            if fa_v4_at(t_kv) && head_dim == 256 {
26171                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
26172                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
26173                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
26174                let v4name = match fa_v4_mode() {
26175                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
26176                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
26177                    _ if deep => "fa_decode_vec_q_v4_deep",
26178                    _ => "fa_decode_vec_q_v4",
26179                };
26180                let fv = if g {
26181                    self.func_g(v4name)
26182                } else {
26183                    self.func(v4name)
26184                };
26185                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
26186                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
26187                let shmem = (if deep { 12160 } else { 11520 }
26188                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
26189                use cudarc::driver::sys::CUfunction_attribute_enum as A;
26190                fv.set_attribute(
26191                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26192                    shmem as i32,
26193                )?;
26194                (
26195                    fv,
26196                    LaunchConfig {
26197                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26198                        block_dim: (32, gqa, 1),
26199                        shared_mem_bytes: shmem,
26200                    },
26201                )
26202            } else if fa_v3_active(head_dim) {
26203                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
26204                // smem = sV only (half of v2's).
26205                let fv = if g {
26206                    self.func_g("fa_decode_vec_q_v3")
26207                } else {
26208                    self.func("fa_decode_vec_q_v3")
26209                };
26210                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
26211                (
26212                    fv,
26213                    LaunchConfig {
26214                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26215                        block_dim: (32, gqa, 1),
26216                        shared_mem_bytes: shmem,
26217                    },
26218                )
26219            } else if fa_v2_on() {
26220                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
26221                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
26222                // partials; same 32KB sK+sV tile as the smem twin.
26223                let fv = if g {
26224                    self.func_g("fa_decode_vec_q_v2")
26225                } else {
26226                    self.func("fa_decode_vec_q_v2")
26227                };
26228                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
26229                (
26230                    fv,
26231                    LaunchConfig {
26232                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26233                        block_dim: (32, gqa, 1),
26234                        shared_mem_bytes: shmem,
26235                    },
26236                )
26237            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
26238            {
26239                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
26240                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
26241                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
26242                let fv = if g {
26243                    self.func_g("fa_decode_vec_q_smem")
26244                } else {
26245                    self.func("fa_decode_vec_q_smem")
26246                };
26247                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
26248                use cudarc::driver::sys::CUfunction_attribute_enum as A;
26249                fv.set_attribute(
26250                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26251                    shmem as i32,
26252                )?;
26253                (
26254                    fv,
26255                    LaunchConfig {
26256                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26257                        block_dim: (32, gqa, 1),
26258                        shared_mem_bytes: shmem,
26259                    },
26260                )
26261            } else {
26262                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
26263                // dequant, zero dynamic shared memory.
26264                let fv = if g {
26265                    self.func_g("fa_decode_vec_q")
26266                } else {
26267                    self.func("fa_decode_vec_q")
26268                };
26269                (
26270                    fv,
26271                    LaunchConfig {
26272                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26273                        block_dim: (32, gqa, 1),
26274                        shared_mem_bytes: 0,
26275                    },
26276                )
26277            }
26278        } else {
26279            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
26280            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
26281            return self.fa_decode_scalar_unified(
26282                q,
26283                k,
26284                v,
26285                o,
26286                head_dim,
26287                n_head,
26288                n_head_kv,
26289                t_kv,
26290                None,
26291                scale,
26292                n_splits,
26293                if fa_vec { sp } else { 256 },
26294                k_tok_bytes,
26295                v_tok_bytes,
26296                g,
26297                part_o,
26298                part_m,
26299                part_l,
26300                None,
26301            );
26302        };
26303        let __s_b = self.gpu.stream();
26304        let mut b = __s_b.launch_builder(&f);
26305        b.arg(q)
26306            .arg(k)
26307            .arg(v)
26308            .arg(&mut *part_o)
26309            .arg(&mut *part_m)
26310            .arg(&mut *part_l)
26311            .arg(&hd)
26312            .arg(&nh)
26313            .arg(&nhkv)
26314            .arg(&tkvi)
26315            .arg(&scale)
26316            .arg(&nsp)
26317            .arg(&ktb)
26318            .arg(&vtb);
26319        unsafe {
26320            b.launch(cfg)?;
26321        }
26322        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
26323        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
26324        let (fc, cfg2) = (
26325            if g {
26326                self.func_g("fa_decode_combine_f32")
26327            } else {
26328                self.fa_func("fa_decode_combine_f32", head_dim)
26329            },
26330            LaunchConfig {
26331                grid_dim: (n_head as u32, 1, 1),
26332                block_dim: (head_dim as u32, 1, 1),
26333                shared_mem_bytes: 0,
26334            },
26335        );
26336        let __s_b2 = self.gpu.stream();
26337        let mut b2 = __s_b2.launch_builder(&fc);
26338        b2.arg(&*part_o)
26339            .arg(&*part_m)
26340            .arg(&*part_l)
26341            .arg(o)
26342            .arg(&hd)
26343            .arg(&nh)
26344            .arg(&nsp);
26345        unsafe {
26346            b2.launch(cfg2)?;
26347        }
26348        Ok(())
26349    }
26350
26351    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
26352    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
26353    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
26354    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
26355    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
26356    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
26357    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
26358    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
26359    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
26360    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
26361    #[allow(clippy::too_many_arguments)]
26362    pub fn fa_decode_batch_seqs_v4(
26363        &self,
26364        q: &CudaSlice<f32>,
26365        kv_ptrs: &cudarc::driver::CudaView<u64>,
26366        pos_seq: &CudaSlice<i32>,
26367        o: &mut CudaSlice<f32>,
26368        head_dim: usize,
26369        n_head: usize,
26370        n_head_kv: usize,
26371        b_n: usize,
26372        t_kv_max: usize,
26373        scale: f32,
26374        split_keys: usize,
26375        k_tok_bytes: usize,
26376        v_tok_bytes: usize,
26377    ) -> Result<(), Box<dyn std::error::Error>> {
26378        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
26379        #[allow(clippy::manual_div_ceil)]
26380        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26381        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
26382        let o_len = b_n * n_head * n_splits_max * head_dim;
26383        let ml_len = b_n * n_head * n_splits_max;
26384        let mut part_guard = self.fa_part_pool.lock().unwrap();
26385        if part_guard
26386            .as_ref()
26387            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26388            .unwrap_or(true)
26389        {
26390            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26391            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26392            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26393            // later live allocations land at those addresses, and the next graph REPLAY writes
26394            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26395            // output corruption began the burst after the trunk's t_kv growth first realloc'd
26396            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26397            // the baked addresses alive (single-stream: eager writes the new buffers, replays
26398            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26399            // (total retired < final size).
26400            let old = part_guard.take();
26401            let (co, cm) = old
26402                .as_ref()
26403                .map(|pp| (pp.0.len(), pp.1.len()))
26404                .unwrap_or((0, 0));
26405            if let Some(old) = old {
26406                self.fa_part_retired.lock().unwrap().push(old);
26407            }
26408            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26409                eprintln!(
26410                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26411                    co, o_len, cm, ml_len
26412                );
26413            }
26414            *part_guard =
26415                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
26416        }
26417        let pg = part_guard.as_mut().unwrap();
26418        self.gpu
26419            .stream()
26420            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
26421        self.gpu
26422            .stream()
26423            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
26424        self.gpu
26425            .stream()
26426            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
26427        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
26428        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
26429        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
26430        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26431        let gqa = (n_head / n_head_kv).max(1) as u32;
26432        let f = self.func("fa_decode_vec_q_seqs_v4");
26433        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
26434        let shmem = (11520 + 32 * head_dim * 2) as u32;
26435        use cudarc::driver::sys::CUfunction_attribute_enum as A;
26436        f.set_attribute(
26437            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26438            shmem as i32,
26439        )?;
26440        let cfg = LaunchConfig {
26441            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
26442            block_dim: (32, gqa, 1),
26443            shared_mem_bytes: shmem,
26444        };
26445        {
26446            let __s_b = self.gpu.stream();
26447            let mut b = __s_b.launch_builder(&f);
26448            b.arg(q)
26449                .arg(kv_ptrs)
26450                .arg(pos_seq)
26451                .arg(&mut *part_o)
26452                .arg(&mut *part_m)
26453                .arg(&mut *part_l)
26454                .arg(&hd)
26455                .arg(&nh)
26456                .arg(&nhkv)
26457                .arg(&scale)
26458                .arg(&nspm)
26459                .arg(&spk)
26460                .arg(&ktb)
26461                .arg(&vtb);
26462            unsafe {
26463                b.launch(cfg)?;
26464            }
26465        }
26466        let fc = self.func("fa_decode_combine_seqs");
26467        let cfg2 = LaunchConfig {
26468            grid_dim: (n_head as u32, b_n as u32, 1),
26469            block_dim: (head_dim as u32, 1, 1),
26470            shared_mem_bytes: 0,
26471        };
26472        let __s_b2 = self.gpu.stream();
26473        let mut b2 = __s_b2.launch_builder(&fc);
26474        b2.arg(&*part_o)
26475            .arg(&*part_m)
26476            .arg(&*part_l)
26477            .arg(o)
26478            .arg(&hd)
26479            .arg(&nh)
26480            .arg(pos_seq)
26481            .arg(&nspm)
26482            .arg(&spk);
26483        unsafe {
26484            b2.launch(cfg2)?;
26485        }
26486        Ok(())
26487    }
26488
26489    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
26490    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
26491    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
26492    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
26493    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
26494    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
26495    #[allow(clippy::too_many_arguments)]
26496    pub fn append_kv_quantized_seqs(
26497        &self,
26498        k_rows: &CudaSlice<f32>,
26499        v_rows: &CudaSlice<f32>,
26500        kv_ptrs: &cudarc::driver::CudaView<u64>,
26501        pos_seq: &CudaSlice<i32>,
26502        b_n: usize,
26503        kv_dim_k: usize,
26504        kv_dim_v: usize,
26505        k_tok_bytes: usize,
26506        v_tok_bytes: usize,
26507    ) -> Result<(), Box<dyn std::error::Error>> {
26508        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
26509        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
26510        let cfg = LaunchConfig {
26511            grid_dim: (nblk, b_n as u32, 1),
26512            block_dim: (32, 1, 1),
26513            shared_mem_bytes: 0,
26514        };
26515        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
26516        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26517        let __s_b = self.gpu.stream();
26518        let mut b = __s_b.launch_builder(&f);
26519        b.arg(k_rows)
26520            .arg(v_rows)
26521            .arg(kv_ptrs)
26522            .arg(pos_seq)
26523            .arg(&kdk)
26524            .arg(&kdv)
26525            .arg(&ktb)
26526            .arg(&vtb);
26527        unsafe {
26528            b.launch(cfg)?;
26529        }
26530        Ok(())
26531    }
26532
26533    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
26534    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
26535    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
26536    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
26537    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
26538    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
26539        std::env::var("MEMRA_NO_FA_VEC").is_err()
26540            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
26541            && base_len + 1 >= fa_vec_min_tkv()
26542            && head_dim <= 256
26543            && head_dim.is_multiple_of(32)
26544    }
26545
26546    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
26547    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
26548    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
26549    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
26550    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
26551    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
26552    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
26553    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
26554    #[allow(clippy::too_many_arguments)]
26555    pub fn fa_decode_rows(
26556        &self,
26557        q: &CudaSlice<f32>,
26558        k: &cudarc::driver::CudaView<u8>,
26559        v: &cudarc::driver::CudaView<u8>,
26560        o: &mut CudaSlice<f32>,
26561        head_dim: usize,
26562        n_head: usize,
26563        n_head_kv: usize,
26564        base_len: usize,
26565        t: usize,
26566        scale: f32,
26567        k_tok_bytes: usize,
26568        v_tok_bytes: usize,
26569        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
26570        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
26571        // keep the host arg. None is a bug for hd512 (asserted below).
26572        base_dev: Option<(&CudaSlice<i32>, i32)>,
26573        // K and V planes hold the same values (gemma globals, wv:=wk): pick
26574        // the _kv twin — V plane never read, value rides the q8_0 key dq.
26575        kv_shared: bool,
26576        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
26577        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
26578        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
26579        g: bool,
26580        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
26581        // (hd512 path) — the standalone quantize launch folds away.
26582        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
26583    ) -> Result<(), Box<dyn std::error::Error>> {
26584        debug_assert!(
26585            base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim.is_multiple_of(32)
26586        );
26587        let t_kv_max = base_len + t; // LAST row's key bound
26588        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
26589        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
26590        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
26591        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
26592        // (parity law), so the partition is freely tunable — verify and decode move together.
26593        if head_dim == 512 {
26594            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26595            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
26596            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
26597            let v = *SP512.get_or_init(|| {
26598                std::env::var("MEMRA_FA_SP512")
26599                    .ok()
26600                    .and_then(|x| x.parse().ok())
26601                    .unwrap_or(0)
26602            });
26603            sp = if v >= 8 {
26604                v
26605            } else {
26606                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
26607            };
26608        }
26609        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
26610        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26611        let gqa = (n_head / n_head_kv).max(1) as u32;
26612        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
26613        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
26614        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
26615        // the different partition changes the combine's FP order (greedy tie flips at depth;
26616        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
26617        // consecutive rows by their OWN ladder value and launch once per group — each row then
26618        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
26619        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
26620        // sp override is t_kv-independent by construction).
26621        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
26622        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
26623            groups.push((0, t, sp));
26624        } else {
26625            let mut r0 = 0usize;
26626            while r0 < t {
26627                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
26628                let mut r1 = r0 + 1;
26629                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
26630                    r1 += 1;
26631                }
26632                groups.push((r0, r1 - r0, sp_g));
26633                r0 = r1;
26634            }
26635        }
26636        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
26637        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
26638        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
26639        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26640        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
26641            std::env::var("MEMRA_FA_SMEM_TKV")
26642                .ok()
26643                .and_then(|v| v.parse().ok())
26644                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
26645        });
26646        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
26647        let v3 = fa_v3_active(head_dim);
26648        let smem_rows =
26649            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
26650        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
26651        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
26652        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
26653        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
26654        let _ = kv_shared;
26655        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
26656        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
26657        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
26658        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
26659        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
26660        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
26661        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
26662        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
26663        // (kv_head, split) stages its tile once and loops the rows over it — kills the
26664        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
26665        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
26666        // shared by every hd512 caller through this wrapper (decode+verify flip together;
26667        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
26668        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
26669        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
26670        // not unpack-bound; jsonl 2026-07-14.
26671        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26672        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
26673        let tb512 = head_dim == 512
26674            && sp <= 32
26675            && n_head / n_head_kv.max(1) <= 16
26676            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
26677        let fname = if tb512 {
26678            "fa_decode_vec_q_rows_v4_512_tb"
26679        } else if i2 {
26680            "fa_decode_vec_q_rows_dpl16_i2"
26681        } else if head_dim == 512 {
26682            "fa_decode_vec_q_rows_dpl16"
26683        }
26684        // gemma globals (parity law)
26685        else if v4 {
26686            "fa_decode_vec_q_rows_v4"
26687        } else if v3 {
26688            "fa_decode_vec_q_rows_v3"
26689        } else if fa_v2_on() {
26690            "fa_decode_vec_q_rows_v2"
26691        } else if smem_rows {
26692            "fa_decode_vec_q_rows_smem"
26693        } else {
26694            "fa_decode_vec_q_rows"
26695        };
26696        let f = if head_dim == 512 {
26697            self.fa_func(fname, head_dim)
26698        } else if g {
26699            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
26700            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
26701            // g-module rows against decode's g-module v4 — different programs, short-VG
26702            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
26703            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
26704            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
26705            // dq macros are format-aware.
26706            self.func_g(if smem_rows {
26707                "fa_decode_vec_q_rows"
26708            } else {
26709                fname
26710            })
26711        } else {
26712            self.func(fname)
26713        };
26714        let shmem = if tb512 {
26715            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
26716            let gk = Self::gkv_on();
26717            let sh =
26718                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
26719            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26720            f.set_attribute(
26721                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26722                sh as i32,
26723            )?;
26724            sh
26725        } else if v4 || v3 || smem_rows || fa_v2_on() {
26726            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
26727            let sh = (if v4 {
26728                11520 + 32 * head_dim * if g { 1 } else { 2 }
26729            } else if v3 {
26730                32 * head_dim * 2
26731            } else {
26732                2 * 32 * head_dim * 2
26733            }) as u32;
26734            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26735            f.set_attribute(
26736                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26737                sh as i32,
26738            )?;
26739            sh
26740        } else {
26741            0
26742        };
26743        // Per-GROUP launches (single group in the common case — identical to the pre-fix
26744        // single launch there): each group gets its own partials (the rows kernel indexes
26745        // partials by its LOCAL grid.z row) and q/o row-offset views.
26746        for &(r0, t_g, sp_g) in &groups {
26747            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
26748            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
26749            let base_i = (base_len + r0) as i32;
26750            let o_len = t_g * n_head * n_splits_g * head_dim;
26751            let ml_len = t_g * n_head * n_splits_g;
26752            let mut part_guard = self.fa_part_pool.lock().unwrap();
26753            if part_guard
26754                .as_ref()
26755                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26756                .unwrap_or(true)
26757            {
26758                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26759                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26760                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26761                // later live allocations land at those addresses, and the next graph REPLAY writes
26762                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26763                // output corruption began the burst after the trunk's t_kv growth first realloc'd
26764                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26765                // the baked addresses alive (single-stream: eager writes the new buffers, replays
26766                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26767                // (total retired < final size).
26768                let old = part_guard.take();
26769                let (co, cm) = old
26770                    .as_ref()
26771                    .map(|pp| (pp.0.len(), pp.1.len()))
26772                    .unwrap_or((0, 0));
26773                if let Some(old) = old {
26774                    self.fa_part_retired.lock().unwrap().push(old);
26775                }
26776                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26777                    eprintln!(
26778                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26779                        co, o_len, cm, ml_len
26780                    );
26781                }
26782                *part_guard =
26783                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
26784            }
26785            let pg = part_guard.as_mut().unwrap();
26786            self.gpu
26787                .stream()
26788                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
26789            self.gpu
26790                .stream()
26791                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
26792            self.gpu
26793                .stream()
26794                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
26795            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
26796            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
26797            let qv = self.view(q, t * n_head * head_dim);
26798            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
26799            let cfg = LaunchConfig {
26800                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
26801                block_dim: (32, gqa, 1),
26802                shared_mem_bytes: shmem,
26803            };
26804            {
26805                let __s_b = self.gpu.stream();
26806                let mut b = __s_b.launch_builder(&f);
26807                if tb512 {
26808                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
26809                    let (bd, plus) =
26810                        base_dev.expect("hd512 rows twin requires a device base counter");
26811                    let plus_g = plus + r0 as i32;
26812                    let nr = t_g as i32;
26813                    if Self::pdl_on() && Self::pdl_wb_on() {
26814                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
26815                        use cudarc::driver::{DevicePtr, DevicePtrMut};
26816                        let s = &self.gpu.stream();
26817                        let (pq, _b0) = q_g.device_ptr(s);
26818                        let (pk, _b1) = k.device_ptr(s);
26819                        let (pv, _b2) = v.device_ptr(s);
26820                        let (po, _b3) = part_o.device_ptr_mut(s);
26821                        let (pm, _b4) = part_m.device_ptr_mut(s);
26822                        let (pl, _b5) = part_l.device_ptr_mut(s);
26823                        let (pb, _b6) = bd.device_ptr(s);
26824                        let mut ps = [
26825                            &pq as *const _ as *mut std::ffi::c_void,
26826                            &pk as *const _ as *mut _,
26827                            &pv as *const _ as *mut _,
26828                            &po as *const _ as *mut _,
26829                            &pm as *const _ as *mut _,
26830                            &pl as *const _ as *mut _,
26831                            &hd as *const _ as *mut _,
26832                            &nh as *const _ as *mut _,
26833                            &nhkv as *const _ as *mut _,
26834                            &pb as *const _ as *mut _,
26835                            &plus_g as *const _ as *mut _,
26836                            &scale as *const _ as *mut _,
26837                            &nspm as *const _ as *mut _,
26838                            &spk as *const _ as *mut _,
26839                            &ktb as *const _ as *mut _,
26840                            &vtb as *const _ as *mut _,
26841                            &nr as *const _ as *mut _,
26842                        ];
26843                        unsafe {
26844                            self.launch_pdl_flash(
26845                                Self::gkv_on(),
26846                                "fa_decode_vec_q_rows_v4_512_tb",
26847                                (n_head_kv as u32, n_splits_g as u32, 1),
26848                                (32, gqa, 1),
26849                                shmem,
26850                                &mut ps,
26851                            )?;
26852                        }
26853                    } else {
26854                        let cfg_tb = LaunchConfig {
26855                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
26856                            block_dim: (32, gqa, 1),
26857                            shared_mem_bytes: shmem,
26858                        };
26859                        b.arg(&q_g)
26860                            .arg(k)
26861                            .arg(v)
26862                            .arg(&mut *part_o)
26863                            .arg(&mut *part_m)
26864                            .arg(&mut *part_l)
26865                            .arg(&hd)
26866                            .arg(&nh)
26867                            .arg(&nhkv)
26868                            .arg(bd)
26869                            .arg(&plus_g)
26870                            .arg(&scale)
26871                            .arg(&nspm)
26872                            .arg(&spk)
26873                            .arg(&ktb)
26874                            .arg(&vtb)
26875                            .arg(&nr);
26876                        unsafe {
26877                            b.launch(cfg_tb)?;
26878                        }
26879                    }
26880                } else if head_dim == 512 {
26881                    let (bd, plus) =
26882                        base_dev.expect("hd512 rows twin requires a device base counter");
26883                    let plus_g = plus + r0 as i32;
26884                    b.arg(&q_g)
26885                        .arg(k)
26886                        .arg(v)
26887                        .arg(&mut *part_o)
26888                        .arg(&mut *part_m)
26889                        .arg(&mut *part_l)
26890                        .arg(&hd)
26891                        .arg(&nh)
26892                        .arg(&nhkv)
26893                        .arg(bd)
26894                        .arg(&plus_g)
26895                        .arg(&scale)
26896                        .arg(&nspm)
26897                        .arg(&spk)
26898                        .arg(&ktb)
26899                        .arg(&vtb);
26900                    unsafe {
26901                        b.launch(cfg)?;
26902                    }
26903                } else {
26904                    b.arg(&q_g)
26905                        .arg(k)
26906                        .arg(v)
26907                        .arg(&mut *part_o)
26908                        .arg(&mut *part_m)
26909                        .arg(&mut *part_l)
26910                        .arg(&hd)
26911                        .arg(&nh)
26912                        .arg(&nhkv)
26913                        .arg(&base_i)
26914                        .arg(&scale)
26915                        .arg(&nspm)
26916                        .arg(&spk)
26917                        .arg(&ktb)
26918                        .arg(&vtb);
26919                    unsafe {
26920                        b.launch(cfg)?;
26921                    }
26922                }
26923            }
26924            let cfg2 = LaunchConfig {
26925                grid_dim: (n_head as u32, t_g as u32, 1),
26926                block_dim: (head_dim as u32, 1, 1),
26927                shared_mem_bytes: 0,
26928            };
26929            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
26930            if head_dim == 512 {
26931                // device-len combine (shared by verify/eager/graph — parity by symbol): the
26932                // per-row n_splits derives from the SAME counter the rows kernel read.
26933                let (bd, plus) = base_dev.unwrap();
26934                let plus_g = plus + r0 as i32;
26935                if let Some((oq, od)) = q8_out.as_mut() {
26936                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
26937                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
26938                    if Self::pdl_on() && Self::pdl_wb_on() {
26939                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
26940                        use cudarc::driver::{DevicePtr, DevicePtrMut};
26941                        let s = &self.gpu.stream();
26942                        let (po, _g0) = part_o.device_ptr(s);
26943                        let (pm, _g1) = part_m.device_ptr(s);
26944                        let (pl, _g2) = part_l.device_ptr(s);
26945                        let (pq, _g3) = oq.device_ptr_mut(s);
26946                        let (pd, _g4) = od.device_ptr_mut(s);
26947                        let (pb, _g5) = bd.device_ptr(s);
26948                        let mut ps = [
26949                            &po as *const _ as *mut std::ffi::c_void,
26950                            &pm as *const _ as *mut _,
26951                            &pl as *const _ as *mut _,
26952                            &pq as *const _ as *mut _,
26953                            &pd as *const _ as *mut _,
26954                            &hd as *const _ as *mut _,
26955                            &nh as *const _ as *mut _,
26956                            &pb as *const _ as *mut _,
26957                            &plus_g as *const _ as *mut _,
26958                            &nspm as *const _ as *mut _,
26959                            &spk as *const _ as *mut _,
26960                        ];
26961                        unsafe {
26962                            self.launch_pdl_flash(
26963                                Self::gkv_on(),
26964                                "fa_decode_combine_rows_dc_q8_1",
26965                                cfg2.grid_dim,
26966                                cfg2.block_dim,
26967                                0,
26968                                &mut ps,
26969                            )?;
26970                        }
26971                        continue;
26972                    }
26973                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
26974                    let __s_b2 = self.gpu.stream();
26975                    let mut b2 = __s_b2.launch_builder(&fc);
26976                    b2.arg(&*part_o)
26977                        .arg(&*part_m)
26978                        .arg(&*part_l)
26979                        .arg(&mut **oq)
26980                        .arg(&mut **od)
26981                        .arg(&hd)
26982                        .arg(&nh)
26983                        .arg(bd)
26984                        .arg(&plus_g)
26985                        .arg(&nspm)
26986                        .arg(&spk);
26987                    unsafe {
26988                        b2.launch(cfg2)?;
26989                    }
26990                    continue;
26991                }
26992                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
26993                let __s_b2 = self.gpu.stream();
26994                let mut b2 = __s_b2.launch_builder(&fc);
26995                b2.arg(&*part_o)
26996                    .arg(&*part_m)
26997                    .arg(&*part_l)
26998                    .arg(&mut o_g)
26999                    .arg(&hd)
27000                    .arg(&nh)
27001                    .arg(bd)
27002                    .arg(&plus_g)
27003                    .arg(&nspm)
27004                    .arg(&spk);
27005                unsafe {
27006                    b2.launch(cfg2)?;
27007                }
27008            } else {
27009                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
27010                // leave the caller's pair unwritten (consumer would read garbage).
27011                assert!(
27012                    q8_out.is_none(),
27013                    "rows q8 emit requires the hd512 dc combine"
27014                );
27015                let fc = self.func("fa_decode_combine_rows");
27016                let __s_b2 = self.gpu.stream();
27017                let mut b2 = __s_b2.launch_builder(&fc);
27018                b2.arg(&*part_o)
27019                    .arg(&*part_m)
27020                    .arg(&*part_l)
27021                    .arg(&mut o_g)
27022                    .arg(&hd)
27023                    .arg(&nh)
27024                    .arg(&base_i)
27025                    .arg(&nspm)
27026                    .arg(&spk);
27027                unsafe {
27028                    b2.launch(cfg2)?;
27029                }
27030            }
27031        }
27032        Ok(())
27033    }
27034
27035    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
27036    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
27037    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
27038    #[allow(clippy::too_many_arguments)]
27039    pub fn fa_decode_rows_w(
27040        &self,
27041        q: &CudaSlice<f32>,
27042        k: &cudarc::driver::CudaView<u8>,
27043        v: &cudarc::driver::CudaView<u8>,
27044        o: &mut CudaSlice<f32>,
27045        head_dim: usize,
27046        n_head: usize,
27047        n_head_kv: usize,
27048        base_dev: &CudaSlice<i32>,
27049        base_plus: i32,
27050        t: usize,
27051        scale: f32,
27052        window: usize,
27053        k_tok_bytes: usize,
27054        v_tok_bytes: usize,
27055        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
27056    ) -> Result<(), Box<dyn std::error::Error>> {
27057        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
27058        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
27059        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
27060        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
27061        debug_assert!(head_dim == 256);
27062        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
27063        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
27064        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
27065        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
27066        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
27067        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
27068        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
27069        let sp = {
27070            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
27071            let v = *SPW.get_or_init(|| {
27072                std::env::var("MEMRA_FA_SPW")
27073                    .ok()
27074                    .and_then(|x| x.parse().ok())
27075                    .unwrap_or(0)
27076            });
27077            if v >= 8 {
27078                v
27079            } else {
27080                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
27081            }
27082        };
27083        #[allow(clippy::manual_div_ceil)]
27084        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27085        let n_splits_max = (window + sp - 1) / sp;
27086        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27087        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
27088        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27089        let gqa = (n_head / n_head_kv).max(1) as u32;
27090        let o_len = t * n_head * n_splits_max * head_dim;
27091        let ml_len = t * n_head * n_splits_max;
27092        let mut part_guard = self.fa_part_pool.lock().unwrap();
27093        if part_guard
27094            .as_ref()
27095            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27096            .unwrap_or(true)
27097        {
27098            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27099            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27100            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27101            // later live allocations land at those addresses, and the next graph REPLAY writes
27102            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27103            // output corruption began the burst after the trunk's t_kv growth first realloc'd
27104            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27105            // the baked addresses alive (single-stream: eager writes the new buffers, replays
27106            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27107            // (total retired < final size).
27108            let old = part_guard.take();
27109            let (co, cm) = old
27110                .as_ref()
27111                .map(|pp| (pp.0.len(), pp.1.len()))
27112                .unwrap_or((0, 0));
27113            if let Some(old) = old {
27114                self.fa_part_retired.lock().unwrap().push(old);
27115            }
27116            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27117                eprintln!(
27118                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27119                    co, o_len, cm, ml_len
27120                );
27121            }
27122            *part_guard =
27123                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27124        }
27125        let pg = part_guard.as_mut().unwrap();
27126        self.gpu
27127            .stream()
27128            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27129        self.gpu
27130            .stream()
27131            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27132        self.gpu
27133            .stream()
27134            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27135        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27136        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
27137        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
27138        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
27139        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
27140        // floor (deep-ctx broadcast win); register twin between.
27141        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
27142        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
27143            std::env::var("MEMRA_FA_SMEM_TKV")
27144                .ok()
27145                .and_then(|v| v.parse().ok())
27146                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
27147        });
27148        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
27149        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
27150        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
27151        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
27152        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
27153        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27154        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
27155        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
27156        // per (lane, format-module) keeps parity structural; the old register-i2 detour
27157        // (-33%) is retired.
27158        let wg = Self::wkv_on();
27159        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
27160        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
27161        let sp2 =
27162            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
27163        if sp2 {
27164            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
27165            if Self::pdl_on() && Self::pdl_wb_on() {
27166                // wave-B2b: flavor mirrors wg.
27167                use cudarc::driver::{DevicePtr, DevicePtrMut};
27168                let s = &self.gpu.stream();
27169                let (pq, _b0) = q.device_ptr(s);
27170                let (pk, _b1) = k.device_ptr(s);
27171                let (pv, _b2) = v.device_ptr(s);
27172                let (po, _b3) = part_o.device_ptr_mut(s);
27173                let (pm, _b4) = part_m.device_ptr_mut(s);
27174                let (pl, _b5) = part_l.device_ptr_mut(s);
27175                let (pb, _b6) = base_dev.device_ptr(s);
27176                let mut ps = [
27177                    &pq as *const _ as *mut std::ffi::c_void,
27178                    &pk as *const _ as *mut _,
27179                    &pv as *const _ as *mut _,
27180                    &po as *const _ as *mut _,
27181                    &pm as *const _ as *mut _,
27182                    &pl as *const _ as *mut _,
27183                    &hd as *const _ as *mut _,
27184                    &nh as *const _ as *mut _,
27185                    &nhkv as *const _ as *mut _,
27186                    &pb as *const _ as *mut _,
27187                    &base_plus as *const _ as *mut _,
27188                    &scale as *const _ as *mut _,
27189                    &nspm as *const _ as *mut _,
27190                    &spk as *const _ as *mut _,
27191                    &ktb as *const _ as *mut _,
27192                    &vtb as *const _ as *mut _,
27193                    &wini as *const _ as *mut _,
27194                ];
27195                unsafe {
27196                    self.launch_pdl_flash(
27197                        wg,
27198                        "fa_decode_vec_q_rows_v4_w_sp",
27199                        (n_head_kv as u32, n_splits_max as u32, t as u32),
27200                        (32, gqa + 1, 1),
27201                        sh,
27202                        &mut ps,
27203                    )?;
27204                }
27205            } else {
27206                let f = if wg {
27207                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
27208                } else {
27209                    self.func("fa_decode_vec_q_rows_v4_w_sp")
27210                };
27211                f.set_attribute(
27212                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27213                    sh as i32,
27214                )?;
27215                let cfg = LaunchConfig {
27216                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27217                    block_dim: (32, gqa + 1, 1),
27218                    shared_mem_bytes: sh,
27219                };
27220                let __s_b = self.gpu.stream();
27221                let mut b = __s_b.launch_builder(&f);
27222                b.arg(q)
27223                    .arg(k)
27224                    .arg(v)
27225                    .arg(&mut *part_o)
27226                    .arg(&mut *part_m)
27227                    .arg(&mut *part_l)
27228                    .arg(&hd)
27229                    .arg(&nh)
27230                    .arg(&nhkv)
27231                    .arg(base_dev)
27232                    .arg(&base_plus)
27233                    .arg(&scale)
27234                    .arg(&nspm)
27235                    .arg(&spk)
27236                    .arg(&ktb)
27237                    .arg(&vtb)
27238                    .arg(&wini);
27239                unsafe {
27240                    b.launch(cfg)?;
27241                }
27242            }
27243        } else {
27244            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
27245                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
27246                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
27247                use cudarc::driver::{DevicePtr, DevicePtrMut};
27248                let s = &self.gpu.stream();
27249                let (pq, _b0) = q.device_ptr(s);
27250                let (pk, _b1) = k.device_ptr(s);
27251                let (pv, _b2) = v.device_ptr(s);
27252                let (po, _b3) = part_o.device_ptr_mut(s);
27253                let (pm, _b4) = part_m.device_ptr_mut(s);
27254                let (pl, _b5) = part_l.device_ptr_mut(s);
27255                let (pb, _b6) = base_dev.device_ptr(s);
27256                let mut ps = [
27257                    &pq as *const _ as *mut std::ffi::c_void,
27258                    &pk as *const _ as *mut _,
27259                    &pv as *const _ as *mut _,
27260                    &po as *const _ as *mut _,
27261                    &pm as *const _ as *mut _,
27262                    &pl as *const _ as *mut _,
27263                    &hd as *const _ as *mut _,
27264                    &nh as *const _ as *mut _,
27265                    &nhkv as *const _ as *mut _,
27266                    &pb as *const _ as *mut _,
27267                    &base_plus as *const _ as *mut _,
27268                    &scale as *const _ as *mut _,
27269                    &nspm as *const _ as *mut _,
27270                    &spk as *const _ as *mut _,
27271                    &ktb as *const _ as *mut _,
27272                    &vtb as *const _ as *mut _,
27273                    &wini as *const _ as *mut _,
27274                ];
27275                unsafe {
27276                    self.launch_pdl_flash(
27277                        wg,
27278                        "fa_decode_vec_q_rows_v4_w",
27279                        (n_head_kv as u32, n_splits_max as u32, t as u32),
27280                        (32, gqa, 1),
27281                        sh,
27282                        &mut ps,
27283                    )?;
27284                }
27285            } else {
27286                let pick = |name: &str| {
27287                    if wg {
27288                        self.func_g(name)
27289                    } else {
27290                        self.func(name)
27291                    }
27292                };
27293                let (f, sh) = if fa_v4_at(window) {
27294                    let f = pick("fa_decode_vec_q_rows_v4_w");
27295                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
27296                } else if smem_tkv > 0 && window >= smem_tkv {
27297                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
27298                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
27299                    (
27300                        pick("fa_decode_vec_q_rows_smem_w"),
27301                        (2 * 32 * head_dim * 2) as u32,
27302                    )
27303                } else {
27304                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
27305                };
27306                f.set_attribute(
27307                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27308                    sh as i32,
27309                )?;
27310                let cfg = LaunchConfig {
27311                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27312                    block_dim: (32, gqa, 1),
27313                    shared_mem_bytes: sh,
27314                };
27315                let __s_b = self.gpu.stream();
27316                let mut b = __s_b.launch_builder(&f);
27317                b.arg(q)
27318                    .arg(k)
27319                    .arg(v)
27320                    .arg(&mut *part_o)
27321                    .arg(&mut *part_m)
27322                    .arg(&mut *part_l)
27323                    .arg(&hd)
27324                    .arg(&nh)
27325                    .arg(&nhkv)
27326                    .arg(base_dev)
27327                    .arg(&base_plus)
27328                    .arg(&scale)
27329                    .arg(&nspm)
27330                    .arg(&spk)
27331                    .arg(&ktb)
27332                    .arg(&vtb)
27333                    .arg(&wini);
27334                unsafe {
27335                    b.launch(cfg)?;
27336                }
27337            }
27338        }
27339        let cfg2 = LaunchConfig {
27340            grid_dim: (n_head as u32, t as u32, 1),
27341            block_dim: (head_dim as u32, 1, 1),
27342            shared_mem_bytes: 0,
27343        };
27344        if let Some((oq, od)) = q8_out {
27345            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
27346            // consumes the pair directly; the standalone quantize launch folds away.
27347            if Self::pdl_on() && Self::pdl_wb_on() {
27348                // wave-B2: flavor mirrors the builder's wg choice.
27349                use cudarc::driver::{DevicePtr, DevicePtrMut};
27350                let s = &self.gpu.stream();
27351                let (po, _g0) = part_o.device_ptr(s);
27352                let (pm, _g1) = part_m.device_ptr(s);
27353                let (pl, _g2) = part_l.device_ptr(s);
27354                let (pq, _g3) = oq.device_ptr_mut(s);
27355                let (pd, _g4) = od.device_ptr_mut(s);
27356                let mut ps = [
27357                    &po as *const _ as *mut std::ffi::c_void,
27358                    &pm as *const _ as *mut _,
27359                    &pl as *const _ as *mut _,
27360                    &pq as *const _ as *mut _,
27361                    &pd as *const _ as *mut _,
27362                    &hd as *const _ as *mut _,
27363                    &nh as *const _ as *mut _,
27364                    &nspm as *const _ as *mut _,
27365                    &spk as *const _ as *mut _,
27366                    &wini as *const _ as *mut _,
27367                ];
27368                unsafe {
27369                    self.launch_pdl_flash(
27370                        wg,
27371                        "fa_decode_combine_rows_w_q8_1",
27372                        cfg2.grid_dim,
27373                        cfg2.block_dim,
27374                        0,
27375                        &mut ps,
27376                    )?;
27377                }
27378                return Ok(());
27379            }
27380            let fc = if wg {
27381                self.func_g("fa_decode_combine_rows_w_q8_1")
27382            } else {
27383                self.func("fa_decode_combine_rows_w_q8_1")
27384            };
27385            let __s_b2 = self.gpu.stream();
27386            let mut b2 = __s_b2.launch_builder(&fc);
27387            b2.arg(&*part_o)
27388                .arg(&*part_m)
27389                .arg(&*part_l)
27390                .arg(oq)
27391                .arg(od)
27392                .arg(&hd)
27393                .arg(&nh)
27394                .arg(&nspm)
27395                .arg(&spk)
27396                .arg(&wini);
27397            unsafe {
27398                b2.launch(cfg2)?;
27399            }
27400            return Ok(());
27401        }
27402        let fc = if wg {
27403            self.func_g("fa_decode_combine_rows_w")
27404        } else {
27405            self.func("fa_decode_combine_rows_w")
27406        };
27407        let __s_b2 = self.gpu.stream();
27408        let mut b2 = __s_b2.launch_builder(&fc);
27409        b2.arg(&*part_o)
27410            .arg(&*part_m)
27411            .arg(&*part_l)
27412            .arg(o)
27413            .arg(&hd)
27414            .arg(&nh)
27415            .arg(&nspm)
27416            .arg(&spk)
27417            .arg(&wini);
27418        unsafe {
27419            b2.launch(cfg2)?;
27420        }
27421        Ok(())
27422    }
27423
27424    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
27425    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
27426    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
27427    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
27428    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
27429    #[allow(clippy::too_many_arguments)]
27430    pub fn fa_decode_rows_dc(
27431        &self,
27432        q: &CudaSlice<f32>,
27433        k: &cudarc::driver::CudaView<u8>,
27434        v: &cudarc::driver::CudaView<u8>,
27435        o: &mut CudaSlice<f32>,
27436        head_dim: usize,
27437        n_head: usize,
27438        n_head_kv: usize,
27439        base_dev: &CudaSlice<i32>,
27440        t_kv_upper: usize,
27441        t: usize,
27442        scale: f32,
27443        k_tok_bytes: usize,
27444        v_tok_bytes: usize,
27445        base_plus: i32,
27446        g: bool,
27447    ) -> Result<(), Box<dyn std::error::Error>> {
27448        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
27449        assert!(
27450            v4 || fa_v3_active(head_dim),
27451            "stream fa rows requires the v3 or v4 lane"
27452        );
27453        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
27454        if v4 {
27455            let sp = fa_split_keys(t_kv_upper, n_head_kv);
27456            #[allow(clippy::manual_div_ceil)]
27457            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27458            let n_splits_max = (t_kv_upper + sp - 1) / sp;
27459            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27460            let (nspm, spk) = (n_splits_max as i32, sp as i32);
27461            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27462            let gqa = (n_head / n_head_kv).max(1) as u32;
27463            let o_len = t * n_head * n_splits_max * head_dim;
27464            let ml_len = t * n_head * n_splits_max;
27465            let mut part_guard = self.fa_part_pool.lock().unwrap();
27466            if part_guard
27467                .as_ref()
27468                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27469                .unwrap_or(true)
27470            {
27471                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27472                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27473                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27474                // later live allocations land at those addresses, and the next graph REPLAY writes
27475                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27476                // output corruption began the burst after the trunk's t_kv growth first realloc'd
27477                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27478                // the baked addresses alive (single-stream: eager writes the new buffers, replays
27479                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27480                // (total retired < final size).
27481                let old = part_guard.take();
27482                let (co, cm) = old
27483                    .as_ref()
27484                    .map(|pp| (pp.0.len(), pp.1.len()))
27485                    .unwrap_or((0, 0));
27486                if let Some(old) = old {
27487                    self.fa_part_retired.lock().unwrap().push(old);
27488                }
27489                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27490                    eprintln!(
27491                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27492                        co, o_len, cm, ml_len
27493                    );
27494                }
27495                *part_guard =
27496                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27497            }
27498            let pg = part_guard.as_mut().unwrap();
27499            self.gpu
27500                .stream()
27501                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27502            self.gpu
27503                .stream()
27504                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27505            self.gpu
27506                .stream()
27507                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27508            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27509            let f = if g {
27510                self.func_g("fa_decode_vec_q_rows_v4_dc")
27511            } else {
27512                self.func("fa_decode_vec_q_rows_v4_dc")
27513            };
27514            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
27515            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27516            f.set_attribute(
27517                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27518                sh as i32,
27519            )?;
27520            let cfg = LaunchConfig {
27521                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27522                block_dim: (32, gqa, 1),
27523                shared_mem_bytes: sh,
27524            };
27525            let __s_b = self.gpu.stream();
27526            let mut b = __s_b.launch_builder(&f);
27527            b.arg(q)
27528                .arg(k)
27529                .arg(v)
27530                .arg(&mut *part_o)
27531                .arg(&mut *part_m)
27532                .arg(&mut *part_l)
27533                .arg(&hd)
27534                .arg(&nh)
27535                .arg(&nhkv)
27536                .arg(base_dev)
27537                .arg(&base_plus)
27538                .arg(&scale)
27539                .arg(&nspm)
27540                .arg(&spk)
27541                .arg(&ktb)
27542                .arg(&vtb);
27543            unsafe {
27544                b.launch(cfg)?;
27545            }
27546            let fc = self.func("fa_decode_combine_rows_dc");
27547            let cfg2 = LaunchConfig {
27548                grid_dim: (n_head as u32, t as u32, 1),
27549                block_dim: (head_dim as u32, 1, 1),
27550                shared_mem_bytes: 0,
27551            };
27552            let __s_b2 = self.gpu.stream();
27553            let mut b2 = __s_b2.launch_builder(&fc);
27554            b2.arg(&*part_o)
27555                .arg(&*part_m)
27556                .arg(&*part_l)
27557                .arg(o)
27558                .arg(&hd)
27559                .arg(&nh)
27560                .arg(base_dev)
27561                .arg(&base_plus)
27562                .arg(&nspm)
27563                .arg(&spk);
27564            unsafe {
27565                b2.launch(cfg2)?;
27566            }
27567            return Ok(());
27568        }
27569        let sp = fa_split_keys(t_kv_upper, n_head_kv);
27570        #[allow(clippy::manual_div_ceil)]
27571        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27572        let n_splits_max = (t_kv_upper + sp - 1) / sp;
27573        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27574        let (nspm, spk) = (n_splits_max as i32, sp as i32);
27575        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27576        let gqa = (n_head / n_head_kv).max(1) as u32;
27577        let o_len = t * n_head * n_splits_max * head_dim;
27578        let ml_len = t * n_head * n_splits_max;
27579        let mut part_guard = self.fa_part_pool.lock().unwrap();
27580        if part_guard
27581            .as_ref()
27582            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27583            .unwrap_or(true)
27584        {
27585            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27586            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27587            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27588            // later live allocations land at those addresses, and the next graph REPLAY writes
27589            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27590            // output corruption began the burst after the trunk's t_kv growth first realloc'd
27591            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27592            // the baked addresses alive (single-stream: eager writes the new buffers, replays
27593            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27594            // (total retired < final size).
27595            let old = part_guard.take();
27596            let (co, cm) = old
27597                .as_ref()
27598                .map(|pp| (pp.0.len(), pp.1.len()))
27599                .unwrap_or((0, 0));
27600            if let Some(old) = old {
27601                self.fa_part_retired.lock().unwrap().push(old);
27602            }
27603            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27604                eprintln!(
27605                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27606                    co, o_len, cm, ml_len
27607                );
27608            }
27609            *part_guard =
27610                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27611        }
27612        let pg = part_guard.as_mut().unwrap();
27613        self.gpu
27614            .stream()
27615            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27616        self.gpu
27617            .stream()
27618            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27619        self.gpu
27620            .stream()
27621            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27622        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27623        let f = self.func("fa_decode_vec_q_rows_v3_dc");
27624        let sh = (32 * head_dim * 2) as u32;
27625        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27626        f.set_attribute(
27627            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27628            sh as i32,
27629        )?;
27630        let cfg = LaunchConfig {
27631            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27632            block_dim: (32, gqa, 1),
27633            shared_mem_bytes: sh,
27634        };
27635        let __s_b = self.gpu.stream();
27636        let mut b = __s_b.launch_builder(&f);
27637        b.arg(q)
27638            .arg(k)
27639            .arg(v)
27640            .arg(&mut *part_o)
27641            .arg(&mut *part_m)
27642            .arg(&mut *part_l)
27643            .arg(&hd)
27644            .arg(&nh)
27645            .arg(&nhkv)
27646            .arg(base_dev)
27647            .arg(&scale)
27648            .arg(&nspm)
27649            .arg(&spk)
27650            .arg(&ktb)
27651            .arg(&vtb);
27652        unsafe {
27653            b.launch(cfg)?;
27654        }
27655        let fc = self.func("fa_decode_combine_rows_dc");
27656        let cfg2 = LaunchConfig {
27657            grid_dim: (n_head as u32, t as u32, 1),
27658            block_dim: (head_dim as u32, 1, 1),
27659            shared_mem_bytes: 0,
27660        };
27661        let plus0 = 0i32;
27662        let __s_b2 = self.gpu.stream();
27663        let mut b2 = __s_b2.launch_builder(&fc);
27664        b2.arg(&*part_o)
27665            .arg(&*part_m)
27666            .arg(&*part_l)
27667            .arg(o)
27668            .arg(&hd)
27669            .arg(&nh)
27670            .arg(base_dev)
27671            .arg(&plus0)
27672            .arg(&nspm)
27673            .arg(&spk);
27674        unsafe {
27675            b2.launch(cfg2)?;
27676        }
27677        Ok(())
27678    }
27679
27680    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
27681    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
27682    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
27683    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
27684    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
27685    ///
27686    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
27687    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
27688    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
27689    /// grouping (different but mathematically-equal log-sum-exp merge).
27690    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
27691    pub fn fa_decode_dc(
27692        &self,
27693        q: &CudaSlice<f32>,
27694        k: &cudarc::driver::CudaView<u8>,
27695        v: &cudarc::driver::CudaView<u8>,
27696        o: &mut CudaSlice<f32>,
27697        head_dim: usize,
27698        n_head: usize,
27699        n_head_kv: usize,
27700        t_kv_dev: &CudaSlice<i32>,
27701        bucket_max: usize,
27702        scale: f32,
27703        k_tok_bytes: usize,
27704        v_tok_bytes: usize,
27705        g: bool,
27706    ) -> Result<(), Box<dyn std::error::Error>> {
27707        self.fa_decode_dc_q8(
27708            q,
27709            k,
27710            v,
27711            o,
27712            head_dim,
27713            n_head,
27714            n_head_kv,
27715            t_kv_dev,
27716            bucket_max,
27717            scale,
27718            k_tok_bytes,
27719            v_tok_bytes,
27720            g,
27721            None,
27722        )
27723    }
27724
27725    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
27726    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
27727    #[allow(clippy::too_many_arguments)]
27728    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27729    pub fn fa_decode_dc_q8(
27730        &self,
27731        q: &CudaSlice<f32>,
27732        k: &cudarc::driver::CudaView<u8>,
27733        v: &cudarc::driver::CudaView<u8>,
27734        o: &mut CudaSlice<f32>,
27735        head_dim: usize,
27736        n_head: usize,
27737        n_head_kv: usize,
27738        t_kv_dev: &CudaSlice<i32>,
27739        bucket_max: usize,
27740        scale: f32,
27741        k_tok_bytes: usize,
27742        v_tok_bytes: usize,
27743        g: bool,
27744        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
27745    ) -> Result<(), Box<dyn std::error::Error>> {
27746        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
27747        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
27748        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
27749        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
27750        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
27751        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
27752        // 2026-07-12).
27753        let mut fa_vec =
27754            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
27755        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
27756            fa_vec = false;
27757        } // mirror kvmod/geom
27758        let sp = fa_split_keys(bucket_max, n_head_kv);
27759        let n_splits = if fa_vec {
27760            ((bucket_max + sp - 1) / sp).max(1)
27761        } else {
27762            ((bucket_max + 255) / 256).max(1)
27763        };
27764        let o_len = n_head * n_splits * head_dim;
27765        let ml_len = n_head * n_splits;
27766        let mut part_guard = self.fa_part_pool.lock().unwrap();
27767        if part_guard
27768            .as_ref()
27769            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27770            .unwrap_or(true)
27771        {
27772            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27773            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27774            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27775            // later live allocations land at those addresses, and the next graph REPLAY writes
27776            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27777            // output corruption began the burst after the trunk's t_kv growth first realloc'd
27778            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27779            // the baked addresses alive (single-stream: eager writes the new buffers, replays
27780            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27781            // (total retired < final size).
27782            let old = part_guard.take();
27783            let (co, cm) = old
27784                .as_ref()
27785                .map(|pp| (pp.0.len(), pp.1.len()))
27786                .unwrap_or((0, 0));
27787            if let Some(old) = old {
27788                self.fa_part_retired.lock().unwrap().push(old);
27789            }
27790            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27791                eprintln!(
27792                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27793                    co, o_len, cm, ml_len
27794                );
27795            }
27796            *part_guard =
27797                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27798        }
27799        let pg = part_guard.as_mut().unwrap();
27800        self.gpu
27801            .stream()
27802            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27803        self.gpu
27804            .stream()
27805            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27806        self.gpu
27807            .stream()
27808            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27809        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27810        let (hd, nh, nhkv, nsp) = (
27811            head_dim as i32,
27812            n_head as i32,
27813            n_head_kv as i32,
27814            n_splits as i32,
27815        );
27816        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27817        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
27818        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
27819        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
27820        let deep = fa_vec
27821            && head_dim == 256
27822            && fa_v4_at(bucket_max)
27823            && !g
27824            && fa_deep_at(bucket_max)
27825            && !matches!(fa_v4_mode(), "noB3" | "stage");
27826        let (f, cfg) = if fa_vec
27827            && head_dim == 512
27828            && bucket_max >= {
27829                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
27830                *FA512_MIN_DC.get_or_init(|| {
27831                    std::env::var("MEMRA_FA512_MIN")
27832                        .ok()
27833                        .and_then(|v| v.parse().ok())
27834                        .unwrap_or(512)
27835                })
27836            } {
27837            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
27838            let gqa = (n_head / n_head_kv).max(1) as u32;
27839            (
27840                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
27841                LaunchConfig {
27842                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27843                    block_dim: (32, gqa, 1),
27844                    shared_mem_bytes: 0,
27845                },
27846            )
27847        } else if fa_vec && head_dim == 512 {
27848            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
27849            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
27850            let q_view = q.as_view();
27851            let mut o_view = o.as_view_mut();
27852            return self.fa_decode_scalar_unified(
27853                &q_view,
27854                k,
27855                v,
27856                &mut o_view,
27857                head_dim,
27858                n_head,
27859                n_head_kv,
27860                0,
27861                Some(t_kv_dev),
27862                scale,
27863                n_splits,
27864                sp,
27865                k_tok_bytes,
27866                v_tok_bytes,
27867                g,
27868                &mut *part_o,
27869                &mut *part_m,
27870                &mut *part_l,
27871                q8_out,
27872            );
27873        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
27874            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
27875            // incl the g-module route + raw-e4m3 sV sizing.
27876            let gqa = (n_head / n_head_kv).max(1) as u32;
27877            let fv = if g {
27878                self.func_g("fa_decode_vec_q_v4_dc")
27879            } else if deep {
27880                self.func("fa_decode_vec_q_v4_deep_dc")
27881            } else {
27882                self.func("fa_decode_vec_q_v4_dc")
27883            };
27884            let shmem =
27885                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
27886            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27887            fv.set_attribute(
27888                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27889                shmem as i32,
27890            )?;
27891            (
27892                fv,
27893                LaunchConfig {
27894                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27895                    block_dim: (32, gqa, 1),
27896                    shared_mem_bytes: shmem,
27897                },
27898            )
27899        } else if fa_vec && fa_v3_active(head_dim) {
27900            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
27901            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
27902            let gqa = (n_head / n_head_kv).max(1) as u32;
27903            let fv = if g {
27904                self.func_g("fa_decode_vec_q_v3_dc")
27905            } else {
27906                self.func("fa_decode_vec_q_v3_dc")
27907            };
27908            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
27909            (
27910                fv,
27911                LaunchConfig {
27912                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27913                    block_dim: (32, gqa, 1),
27914                    shared_mem_bytes: shmem,
27915                },
27916            )
27917        } else if fa_vec && fa_v2_on() {
27918            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
27919            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
27920            // a numeric config; eager, rows-verify and graph all switch together).
27921            let gqa = (n_head / n_head_kv).max(1) as u32;
27922            let fv = if g {
27923                self.func_g("fa_decode_vec_q_v2_dc")
27924            } else {
27925                self.func("fa_decode_vec_q_v2_dc")
27926            };
27927            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
27928            (
27929                fv,
27930                LaunchConfig {
27931                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27932                    block_dim: (32, gqa, 1),
27933                    shared_mem_bytes: shmem,
27934                },
27935            )
27936        } else if fa_vec {
27937            let gqa = (n_head / n_head_kv).max(1) as u32;
27938            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
27939            let fv = if g {
27940                self.func_g("fa_decode_vec_q_dc")
27941            } else {
27942                self.func("fa_decode_vec_q_dc")
27943            };
27944            (
27945                fv,
27946                LaunchConfig {
27947                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27948                    block_dim: (32, gqa, 1),
27949                    shared_mem_bytes: 0,
27950                },
27951            )
27952        } else {
27953            let q_view = q.as_view();
27954            let mut o_view = o.as_view_mut();
27955            return self.fa_decode_scalar_unified(
27956                &q_view,
27957                k,
27958                v,
27959                &mut o_view,
27960                head_dim,
27961                n_head,
27962                n_head_kv,
27963                0,
27964                Some(t_kv_dev),
27965                scale,
27966                n_splits,
27967                if fa_vec { sp } else { 256 },
27968                k_tok_bytes,
27969                v_tok_bytes,
27970                g,
27971                &mut *part_o,
27972                &mut *part_m,
27973                &mut *part_l,
27974                q8_out,
27975            );
27976        };
27977        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
27978        let __s_b = self.gpu.stream();
27979        let mut b = __s_b.launch_builder(&f);
27980        b.arg(q)
27981            .arg(k)
27982            .arg(v)
27983            .arg(&mut *part_o)
27984            .arg(&mut *part_m)
27985            .arg(&mut *part_l)
27986            .arg(&hd)
27987            .arg(&nh)
27988            .arg(&nhkv)
27989            .arg(t_kv_dev)
27990            .arg(&scale)
27991            .arg(&nsp)
27992            .arg(&ski)
27993            .arg(&ktb)
27994            .arg(&vtb);
27995        unsafe {
27996            b.launch(cfg)?;
27997        }
27998        let cfg2 = LaunchConfig {
27999            grid_dim: (n_head as u32, 1, 1),
28000            block_dim: (head_dim as u32, 1, 1),
28001            shared_mem_bytes: 0,
28002        };
28003        if let Some((oq, od)) = q8_out {
28004            let fc = if g {
28005                self.func_g("fa_decode_combine_q8_1")
28006            } else {
28007                self.fa_func("fa_decode_combine_q8_1", head_dim)
28008            };
28009            let __s_b2 = self.gpu.stream();
28010            let mut b2 = __s_b2.launch_builder(&fc);
28011            b2.arg(&*part_o)
28012                .arg(&*part_m)
28013                .arg(&*part_l)
28014                .arg(oq)
28015                .arg(od)
28016                .arg(&hd)
28017                .arg(&nh)
28018                .arg(&nsp);
28019            unsafe {
28020                b2.launch(cfg2)?;
28021            }
28022            return Ok(());
28023        }
28024        let fc = if g {
28025            self.func_g("fa_decode_combine_f32")
28026        } else {
28027            self.fa_func("fa_decode_combine_f32", head_dim)
28028        };
28029        let __s_b2 = self.gpu.stream();
28030        let mut b2 = __s_b2.launch_builder(&fc);
28031        b2.arg(&*part_o)
28032            .arg(&*part_m)
28033            .arg(&*part_l)
28034            .arg(o)
28035            .arg(&hd)
28036            .arg(&nh)
28037            .arg(&nsp);
28038        unsafe {
28039            b2.launch(cfg2)?;
28040        }
28041        Ok(())
28042    }
28043
28044    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
28045    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
28046    /// at equal rows.
28047    #[allow(clippy::too_many_arguments)]
28048    pub fn append_kv_quantized_dcw(
28049        &self,
28050        k_row: &CudaSlice<f32>,
28051        v_row: &CudaSlice<f32>,
28052        kc: &mut CudaSlice<u8>,
28053        vc: &mut CudaSlice<u8>,
28054        len_dev: &CudaSlice<i32>,
28055        base_dev: Option<&CudaSlice<i32>>,
28056        kv_dim_k: usize,
28057        kv_dim_v: usize,
28058        k_tok_bytes: usize,
28059        v_tok_bytes: usize,
28060    ) -> Result<(), Box<dyn std::error::Error>> {
28061        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
28062        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
28063        let cfg = LaunchConfig {
28064            grid_dim: (nblk, 1, 1),
28065            block_dim: (32, 1, 1),
28066            shared_mem_bytes: 0,
28067        };
28068        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28069        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28070        let null: u64 = 0;
28071        let __s_b = self.gpu.stream();
28072        let mut b = __s_b.launch_builder(&f);
28073        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
28074        match base_dev {
28075            Some(base) => {
28076                b.arg(base);
28077            }
28078            None => {
28079                b.arg(&null);
28080            }
28081        }
28082        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
28083        unsafe {
28084            b.launch(cfg)?;
28085        }
28086        Ok(())
28087    }
28088
28089    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
28090    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
28091        let f = self.func("inc_i32");
28092        let cfg = LaunchConfig {
28093            grid_dim: (1, 1, 1),
28094            block_dim: (1, 1, 1),
28095            shared_mem_bytes: 0,
28096        };
28097        let __s_b = self.gpu.stream();
28098        let mut b = __s_b.launch_builder(&f);
28099        b.arg(counter);
28100        unsafe {
28101            b.launch(cfg)?;
28102        }
28103        Ok(())
28104    }
28105
28106    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
28107    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
28108    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
28109    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
28110    /// kernel class on this lane); callers keep eager below the vec floor and for any other
28111    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
28112    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
28113    /// alive across bucket growth.
28114    #[allow(clippy::too_many_arguments)]
28115    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
28116    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
28117    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
28118    /// THE ONE PLACE THE FA PARTIAL POOL IS ALLOCATED.
28119    ///
28120    /// Eight call sites grow this pool and all eight retire-on-grow correctly, but only ONE
28121    /// of them carried the `[fa-pool] grow` receipt, so that receipt under-reported grows by
28122    /// seven eighths and no grow could honestly be dated against a request. Routing every
28123    /// grower through here makes the count real. The receipt names the site so a ladder can
28124    /// be attributed, and stays bounded so a pathological ladder cannot flood a serving log.
28125    ///
28126    /// `MEMRA_FA_PART_ZERO=1` (DEFAULT OFF, diagnostic only) zeroes the fresh buffers. A grow
28127    /// hands every subsequent launch three UNINITIALIZED banks; if the poison is a combine
28128    /// reading a partial bank its producer never wrote, that makes every row and every head
28129    /// non-finite at once, which is the shape the level-2 bad-row bitmap reports at the
28130    /// global-attention join.
28131    ///
28132    /// READ IT IN ONE DIRECTION ONLY. Zeroed banks carry m = 0.0, not NEG_INF, so the
28133    /// empty-split no-op guard never engages: a bank that is entirely unwritten still
28134    /// combines to L = 0 and O/L = 0/0 = NaN. So **silence under this arm convicts the pool;
28135    /// continued trapping acquits nothing**, because only the PARTIALLY unwritten class (real
28136    /// splits beside stale zeroed ones) goes quiet. Discriminator, never a fix, and never a
28137    /// serving arm: where it does go quiet the output is still wrong, it just looks plausible.
28138    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28139    fn fa_part_alloc(
28140        &self,
28141        o_len: usize,
28142        ml_len: usize,
28143        co: usize,
28144        cm: usize,
28145    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
28146        static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
28147        let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
28148        if n < 64 {
28149            eprintln!(
28150                "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
28151                self.ctx().ordinal(),
28152                fa_part_zero_on()
28153            );
28154        }
28155        let mut po = self.alloc_uninit::<f32>(o_len)?;
28156        let mut pm = self.alloc_uninit::<f32>(ml_len)?;
28157        let mut pl = self.alloc_uninit::<f32>(ml_len)?;
28158        if fa_part_zero_on() {
28159            self.gpu.stream().memset_zeros(&mut po)?;
28160            self.gpu.stream().memset_zeros(&mut pm)?;
28161            self.gpu.stream().memset_zeros(&mut pl)?;
28162        }
28163        Ok((po, pm, pl))
28164    }
28165
28166    fn fa_part_pool_grow(
28167        &self,
28168        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
28169        o_len: usize,
28170        ml_len: usize,
28171    ) -> Result<(), Box<dyn std::error::Error>> {
28172        if part_guard
28173            .as_ref()
28174            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28175            .unwrap_or(true)
28176        {
28177            let old = part_guard.take();
28178            let (co, cm) = old
28179                .as_ref()
28180                .map(|pp| (pp.0.len(), pp.1.len()))
28181                .unwrap_or((0, 0));
28182            if let Some(old) = old {
28183                self.fa_part_retired.lock().unwrap().push(old);
28184            }
28185            // GROW RECEIPT. This pool is grow-only, retires-on-grow and never frees, and every
28186            // FA decode/verify launch in the process reads and writes it. A grow is therefore a
28187            // process-lifetime EVENT — new addresses, a retired buffer kept alive forever, and
28188            // a different partial layout — and it is invisible in every log we have. The step37
28189            // spec fault is clean for the first two or three requests of a process and then
28190            // poisons trunk layer 20 (research: MEMRA_SPEC_NAN_SCAN), which is exactly the
28191            // shape a mid-life pool grow would produce, so the grows have to be datable
28192            // against the requests. Cap raised from 8 after the first run measured FOUR
28193            // grows per device (380928 -> 761856 -> 1523712 -> 3047424): with two devices the
28194            // 8 slots were spent before any grow could be dated against a request, which was
28195            // the entire point of the receipt. Still bounded so a pathological ladder cannot
28196            // flood a serving log.
28197            *part_guard =
28198                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28199        }
28200        Ok(())
28201    }
28202
28203    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
28204    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
28205    pub fn fa_dcw_pool_ensure(
28206        &self,
28207        head_dim: usize,
28208        n_head: usize,
28209        n_head_kv: usize,
28210        bucket_max: usize,
28211    ) -> Result<(), Box<dyn std::error::Error>> {
28212        let sp = fa_split_keys(bucket_max, n_head_kv);
28213        #[allow(clippy::manual_div_ceil)]
28214        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28215        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28216        let o_len = n_head * n_splits * head_dim;
28217        let ml_len = n_head * n_splits;
28218        let mut part_guard = self.fa_part_pool.lock().unwrap();
28219        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
28220    }
28221
28222    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
28223    /// appended; one launch walks the KV stream once with two query rows (per-row causal
28224    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
28225    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
28226    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
28227    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
28228    /// outputs (the head gate fuses into the combine as in the t=1 path).
28229    #[allow(clippy::too_many_arguments)]
28230    pub fn fa_decode_dcw2(
28231        &self,
28232        q2: &CudaSlice<f32>,
28233        k_ring: &cudarc::driver::CudaView<u8>,
28234        v_ring: &cudarc::driver::CudaView<u8>,
28235        o2: &mut CudaSlice<f32>,
28236        head_dim: usize,
28237        n_head: usize,
28238        n_head_kv: usize,
28239        len_dev: &CudaSlice<i32>,
28240        base_dev: Option<&CudaSlice<i32>>,
28241        window: usize,
28242        bucket_max: usize,
28243        scale: f32,
28244        k_tok_bytes: usize,
28245        v_tok_bytes: usize,
28246        gate2: &CudaSlice<f32>,
28247    ) -> Result<(), Box<dyn std::error::Error>> {
28248        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
28249        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
28250            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
28251        }
28252        let sp = fa_split_keys(bucket_max, n_head_kv);
28253        #[allow(clippy::manual_div_ceil)]
28254        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28255        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28256        // Partials for BOTH rows: row-major halves.
28257        let o_len = 2 * n_head * n_splits * head_dim;
28258        let ml_len = 2 * n_head * n_splits;
28259        let mut part_guard = self.fa_part_pool.lock().unwrap();
28260        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28261        let pg = part_guard.as_mut().unwrap();
28262        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28263        let (hd, nh, nhkv, nsp) = (
28264            head_dim as i32,
28265            n_head as i32,
28266            n_head_kv as i32,
28267            n_splits as i32,
28268        );
28269        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28270        let (ski, win) = (sp as i32, window as i32);
28271        let gqa = (n_head / n_head_kv).max(1) as u32;
28272        let smem = (32 * head_dim * 2) as u32;
28273        let f = self.func("fa_decode_vec_q_v3_dcw2");
28274        let cfg = LaunchConfig {
28275            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28276            block_dim: (32, gqa, 1),
28277            shared_mem_bytes: smem,
28278        };
28279        let null: u64 = 0;
28280        {
28281            let __s_b = self.gpu.stream();
28282            let mut b = __s_b.launch_builder(&f);
28283            b.arg(q2)
28284                .arg(k_ring)
28285                .arg(v_ring)
28286                .arg(&mut *part_o)
28287                .arg(&mut *part_m)
28288                .arg(&mut *part_l)
28289                .arg(&hd)
28290                .arg(&nh)
28291                .arg(&nhkv)
28292                .arg(len_dev);
28293            match base_dev {
28294                Some(base) => {
28295                    b.arg(base);
28296                }
28297                None => {
28298                    b.arg(&null);
28299                }
28300            }
28301            b.arg(&win)
28302                .arg(&scale)
28303                .arg(&nsp)
28304                .arg(&ski)
28305                .arg(&ktb)
28306                .arg(&vtb);
28307            unsafe {
28308                b.launch(cfg)?;
28309            }
28310        }
28311        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
28312        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
28313        // one launch covers both rows with the exact t=1 program per (row, head).
28314        let fc = {
28315            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28316            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28317                self.func("fa_decode_combine_gate_f32_s")
28318            } else {
28319                self.func("fa_decode_combine_gate_f32")
28320            }
28321        };
28322        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
28323        let nh2 = (2 * n_head) as i32;
28324        let cfg2 = LaunchConfig {
28325            grid_dim: ((2 * n_head) as u32, 1, 1),
28326            block_dim: (head_dim as u32, 1, 1),
28327            shared_mem_bytes: if combine_shared {
28328                (2 * n_splits * 4) as u32
28329            } else {
28330                0
28331            },
28332        };
28333        let __s_b2 = self.gpu.stream();
28334        let mut b2 = __s_b2.launch_builder(&fc);
28335        b2.arg(&*part_o)
28336            .arg(&*part_m)
28337            .arg(&*part_l)
28338            .arg(gate2)
28339            .arg(o2)
28340            .arg(&hd)
28341            .arg(&nh2)
28342            .arg(&nsp);
28343        unsafe {
28344            b2.launch(cfg2)?;
28345        }
28346        Ok(())
28347    }
28348
28349    /// T-ROW dcw decode attention over a per-row session table (the per-session
28350    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
28351    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
28352    /// program verbatim with that row's ring/len/base and its own split geometry, so each
28353    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
28354    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
28355    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
28356    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
28357    #[allow(clippy::too_many_arguments)]
28358    pub fn fa_decode_dcw_rows(
28359        &self,
28360        q_rows: &CudaSlice<f32>,
28361        tab: &CudaSlice<u64>,
28362        o_rows: &mut CudaSlice<f32>,
28363        t: usize,
28364        head_dim: usize,
28365        n_head: usize,
28366        n_head_kv: usize,
28367        window: usize,
28368        max_ns: usize,
28369        scale: f32,
28370        k_tok_bytes: usize,
28371        v_tok_bytes: usize,
28372        gate_rows: &CudaSlice<f32>,
28373    ) -> Result<(), Box<dyn std::error::Error>> {
28374        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
28375            || head_dim > 256
28376            || !head_dim.is_multiple_of(32)
28377            || !fa_v3_on()
28378        {
28379            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
28380        }
28381        if fa_sm_count() < 128
28382            || std::env::var("MEMRA_FA_SPLIT").is_ok()
28383            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
28384            || std::env::var("MEMRA_FA_SP16").is_ok()
28385        {
28386            return Err(
28387                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
28388                 (or a <128-SM rig) keep the per-row path"
28389                    .into(),
28390            );
28391        }
28392        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
28393            return Err("fa_decode_dcw_rows geometry".into());
28394        }
28395        let o_len = t * n_head * max_ns * head_dim;
28396        let ml_len = t * n_head * max_ns;
28397        let mut part_guard = self.fa_part_pool.lock().unwrap();
28398        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28399        let pg = part_guard.as_mut().unwrap();
28400        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28401        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
28402        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28403        let (win, mns) = (window as i32, max_ns as i32);
28404        let gqa = (n_head / n_head_kv).max(1) as u32;
28405        let smem = (32 * head_dim * 2) as u32;
28406        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
28407        let cfg = LaunchConfig {
28408            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
28409            block_dim: (32, gqa, 1),
28410            shared_mem_bytes: smem,
28411        };
28412        {
28413            let __s_b = self.gpu.stream();
28414            let mut b = __s_b.launch_builder(&f);
28415            b.arg(q_rows)
28416                .arg(tab)
28417                .arg(&mut *part_o)
28418                .arg(&mut *part_m)
28419                .arg(&mut *part_l)
28420                .arg(&hd)
28421                .arg(&nh)
28422                .arg(&nhkv)
28423                .arg(&win)
28424                .arg(&scale)
28425                .arg(&mns)
28426                .arg(&ktb)
28427                .arg(&vtb);
28428            unsafe {
28429                b.launch(cfg)?;
28430            }
28431        }
28432        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
28433        // row r head h reads its own partial bank; splits past a row's ns_eff carry
28434        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
28435        let fc = {
28436            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28437            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28438                self.func("fa_decode_combine_gate_f32_s")
28439            } else {
28440                self.func("fa_decode_combine_gate_f32")
28441            }
28442        };
28443        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
28444        let nht = (t * n_head) as i32;
28445        let cfg2 = LaunchConfig {
28446            grid_dim: ((t * n_head) as u32, 1, 1),
28447            block_dim: (head_dim as u32, 1, 1),
28448            shared_mem_bytes: if combine_shared {
28449                (2 * max_ns * 4) as u32
28450            } else {
28451                0
28452            },
28453        };
28454        let __s_b2 = self.gpu.stream();
28455        let mut b2 = __s_b2.launch_builder(&fc);
28456        b2.arg(&*part_o)
28457            .arg(&*part_m)
28458            .arg(&*part_l)
28459            .arg(gate_rows)
28460            .arg(o_rows)
28461            .arg(&hd)
28462            .arg(&nht)
28463            .arg(&mns);
28464        unsafe {
28465            b2.launch(cfg2)?;
28466        }
28467        Ok(())
28468    }
28469
28470    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
28471    pub fn fa_decode_dcw(
28472        &self,
28473        q: &CudaSlice<f32>,
28474        k_ring: &cudarc::driver::CudaView<u8>,
28475        v_ring: &cudarc::driver::CudaView<u8>,
28476        o: &mut CudaSlice<f32>,
28477        head_dim: usize,
28478        n_head: usize,
28479        n_head_kv: usize,
28480        len_dev: &CudaSlice<i32>,
28481        base_dev: Option<&CudaSlice<i32>>,
28482        window: usize,
28483        bucket_max: usize,
28484        scale: f32,
28485        k_tok_bytes: usize,
28486        v_tok_bytes: usize,
28487        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
28488        // one launch saved); `o` then receives the GATED output and the caller skips its
28489        // attn_head_gate call.
28490        fused_gate: Option<&CudaSlice<f32>>,
28491    ) -> Result<(), Box<dyn std::error::Error>> {
28492        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
28493        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
28494            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"
28495                .into());
28496        }
28497        let sp = fa_split_keys(bucket_max, n_head_kv);
28498        #[allow(clippy::manual_div_ceil)]
28499        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28500        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28501        let o_len = n_head * n_splits * head_dim;
28502        let ml_len = n_head * n_splits;
28503        let mut part_guard = self.fa_part_pool.lock().unwrap();
28504        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28505        let pg = part_guard.as_mut().unwrap();
28506        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
28507        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
28508        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
28509        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
28510        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28511        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
28512        // finds the attention children BY their three-memset signature and updates the
28513        // memset widths per bucket — capturing without them silently kills retargeting
28514        // (battery-v8 token drift, 2026-08-21).
28515        let memset_on = *MEMSET_ON
28516            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
28517            || crate::tp::token_graph_building();
28518        if memset_on {
28519            self.gpu
28520                .stream()
28521                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28522            self.gpu
28523                .stream()
28524                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28525            self.gpu
28526                .stream()
28527                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28528        }
28529        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28530        let (hd, nh, nhkv, nsp) = (
28531            head_dim as i32,
28532            n_head as i32,
28533            n_head_kv as i32,
28534            n_splits as i32,
28535        );
28536        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28537        let (ski, win) = (sp as i32, window as i32);
28538        let gqa = (n_head / n_head_kv).max(1) as u32;
28539        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
28540        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
28541        // see fa_dec_v3_walk_u). Same launch geometry.
28542        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28543        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
28544        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
28545            Ok("2") => 2,
28546            Ok("1") => 1,
28547            _ => 0,
28548        });
28549        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
28550        // permission-blocked in this container and the module params are not exposed, so this
28551        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
28552        // prints cumulative cycle shares every 430 launches.
28553        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28554        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
28555        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
28556            std::sync::Mutex::new(None);
28557        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
28558        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
28559        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
28560        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28561        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
28562            && (n_head / n_head_kv).is_multiple_of(2)
28563            && (n_head / n_head_kv) >= 2;
28564        let f = if fprof {
28565            self.func("fa_decode_vec_q_v3_dcw_prof")
28566        } else if hs2 {
28567            self.func("fa_decode_vec_q_v3_dcw_hs2")
28568        } else if hoist == 2 {
28569            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
28570            self.func("fa_decode_vec_q_v3_dcw_hc")
28571        } else if hoist == 1 {
28572            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
28573            self.func("fa_decode_vec_q_v3_dcw_h")
28574        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
28575            self.func("fa_decode_vec_q_v3_dcw_u8")
28576        } else {
28577            self.func("fa_decode_vec_q_v3_dcw")
28578        };
28579        let cfg = LaunchConfig {
28580            grid_dim: if hs2 {
28581                ((2 * n_head_kv) as u32, n_splits as u32, 1)
28582            } else {
28583                (n_head_kv as u32, n_splits as u32, 1)
28584            },
28585            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
28586            shared_mem_bytes: smem,
28587        };
28588        let null: u64 = 0;
28589        let __s_b = self.gpu.stream();
28590        let mut b = __s_b.launch_builder(&f);
28591        b.arg(q)
28592            .arg(k_ring)
28593            .arg(v_ring)
28594            .arg(&mut *part_o)
28595            .arg(&mut *part_m)
28596            .arg(&mut *part_l)
28597            .arg(&hd)
28598            .arg(&nh)
28599            .arg(&nhkv)
28600            .arg(len_dev);
28601        match base_dev {
28602            Some(base) => {
28603                b.arg(base);
28604            }
28605            None => {
28606                b.arg(&null);
28607            }
28608        }
28609        b.arg(&win)
28610            .arg(&scale)
28611            .arg(&nsp)
28612            .arg(&ski)
28613            .arg(&ktb)
28614            .arg(&vtb);
28615        if fprof {
28616            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
28617            if guard
28618                .as_ref()
28619                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
28620            {
28621                *guard = Some((self.ctx().ordinal(), self.htod_u64(&[0u64; 8])?));
28622            }
28623            let (_, buf) = guard.as_mut().expect("armed above");
28624            b.arg(&*buf);
28625            unsafe {
28626                b.launch(cfg)?;
28627            }
28628            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
28629            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
28630            if n.is_multiple_of(430) {
28631                self.stream().synchronize()?;
28632                let h = self.dtoh_u64(buf)?;
28633                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
28634                let tot: u64 = h[..6].iter().sum();
28635                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
28636                for (i, name) in phases.iter().enumerate() {
28637                    let pct = if tot > 0 {
28638                        h[i] as f64 / tot as f64 * 100.0
28639                    } else {
28640                        0.0
28641                    };
28642                    line.push_str(&format!(" {name}={pct:.1}%"));
28643                }
28644                if h[6] > 0 {
28645                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
28646                }
28647                eprintln!("{line}");
28648            }
28649        } else {
28650            unsafe {
28651                b.launch(cfg)?;
28652            }
28653        }
28654        let mut combine_shared = false;
28655        let fc = if fused_gate.is_some() {
28656            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
28657            // n_splits-deep dependent global load chain every thread used to walk twice).
28658            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28659            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28660                combine_shared = true;
28661                self.func("fa_decode_combine_gate_f32_s")
28662            } else {
28663                self.func("fa_decode_combine_gate_f32")
28664            }
28665        } else {
28666            self.fa_func("fa_decode_combine_f32", head_dim)
28667        };
28668        let cfg2 = LaunchConfig {
28669            grid_dim: (n_head as u32, 1, 1),
28670            block_dim: (head_dim as u32, 1, 1),
28671            shared_mem_bytes: if combine_shared {
28672                (2 * n_splits * 4) as u32
28673            } else {
28674                0
28675            },
28676        };
28677        let __s_b2 = self.gpu.stream();
28678        let mut b2 = __s_b2.launch_builder(&fc);
28679        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
28680        if let Some(gate_row) = fused_gate {
28681            b2.arg(gate_row);
28682        }
28683        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
28684        unsafe {
28685            b2.launch(cfg2)?;
28686        }
28687        Ok(())
28688    }
28689
28690    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
28691    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
28692    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
28693    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
28694    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
28695    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28696    pub fn fa_geom_eager(
28697        &self,
28698        t_kv: usize,
28699        head_dim: usize,
28700        n_head_kv: usize,
28701        g: bool,
28702    ) -> (bool, usize) {
28703        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
28704        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
28705        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
28706        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
28707        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
28708        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
28709        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
28710        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
28711        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
28712        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
28713        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim.is_multiple_of(32));
28714        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
28715        // family; everything else falls to the g-module scalar.
28716        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
28717        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
28718        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
28719        if g && head_dim == 256 && !fa_v4_at(t_kv) {
28720            fa_vec = false;
28721        }
28722        let sp = fa_split_keys(t_kv, n_head_kv);
28723        let n_splits = if fa_vec {
28724            ((t_kv + sp - 1) / sp).max(1)
28725        } else {
28726            ((t_kv + 255) / 256).max(1)
28727        };
28728        (fa_vec, n_splits)
28729    }
28730
28731    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
28732    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
28733    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
28734    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
28735    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
28736    pub fn fa_bucket_key(
28737        &self,
28738        t_kv: usize,
28739        head_dim: usize,
28740        n_head_kv: usize,
28741        g: bool,
28742    ) -> (bool, usize) {
28743        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
28744    }
28745
28746    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
28747    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
28748    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
28749    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
28750    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
28751    /// device data) — every per-step varying scalar must come from a device counter. Returns the
28752    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
28753    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
28754    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
28755    /// replays (transients returning to the pool get reused by unrelated work and corrupt
28756    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
28757    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28758    pub fn capture_graph_retained<F>(
28759        &self,
28760        step: F,
28761    ) -> Result<
28762        (
28763            cudarc::driver::CudaGraph,
28764            Vec<Box<dyn std::any::Any + Send>>,
28765        ),
28766        Box<dyn std::error::Error>,
28767    >
28768    where
28769        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28770    {
28771        use cudarc::driver::sys::CUgraphInstantiate_flags;
28772        self.capture_graph_retained_flags(
28773            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28774            step,
28775        )
28776    }
28777
28778    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
28779    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
28780    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
28781    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
28782    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28783    pub fn capture_graph_retained_flags<F>(
28784        &self,
28785        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
28786        mut step: F,
28787    ) -> Result<
28788        (
28789            cudarc::driver::CudaGraph,
28790            Vec<Box<dyn std::any::Any + Send>>,
28791        ),
28792        Box<dyn std::error::Error>,
28793    >
28794    where
28795        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28796    {
28797        use cudarc::driver::sys::CUstreamCaptureMode;
28798        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
28799        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
28800        // while the capture region is open become dead copy NODES replayed every launch
28801        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
28802        // warmup runs allocate the same transient sequence at the same pool addresses, so
28803        // retaining the warmup clones preserves the draft-graph fix without polluting the
28804        // captured graph.
28805        self.capture_keep.lock().unwrap().clear();
28806        let was_tracking = self.gpu.ctx.is_event_tracking();
28807        if was_tracking {
28808            unsafe {
28809                self.gpu.ctx.disable_event_tracking();
28810            }
28811        }
28812        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28813            self.capture_keep_on
28814                .store(true, std::sync::atomic::Ordering::Relaxed);
28815            let w = (|| {
28816                step(self)?;
28817                step(self)
28818            })();
28819            self.capture_keep_on
28820                .store(false, std::sync::atomic::Ordering::Relaxed);
28821            w?;
28822            self.gpu.stream().synchronize()?;
28823            self.gpu
28824                .stream()
28825                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28826            let r = step(self);
28827            let g = self.gpu.stream().end_capture(flags);
28828            r?;
28829            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28830            graph.upload()?;
28831            Ok(graph)
28832        };
28833        let result = run();
28834        self.capture_keep_on
28835            .store(false, std::sync::atomic::Ordering::Relaxed);
28836        if was_tracking {
28837            unsafe {
28838                self.gpu.ctx.enable_event_tracking();
28839            }
28840        }
28841        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
28842        Ok((result?, keeper))
28843    }
28844
28845    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
28846    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
28847    /// alloc-free with persistent operands, and their bodies carry device side effects
28848    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
28849    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
28850    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28851    pub fn capture_graph_retained_nowarm<F>(
28852        &self,
28853        mut step: F,
28854    ) -> Result<
28855        (
28856            cudarc::driver::CudaGraph,
28857            Vec<Box<dyn std::any::Any + Send>>,
28858        ),
28859        Box<dyn std::error::Error>,
28860    >
28861    where
28862        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28863    {
28864        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
28865        let was_tracking = self.gpu.ctx.is_event_tracking();
28866        if was_tracking {
28867            unsafe {
28868                self.gpu.ctx.disable_event_tracking();
28869            }
28870        }
28871        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28872            self.gpu.stream().synchronize()?;
28873            self.gpu
28874                .stream()
28875                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28876            let r = step(self);
28877            let g = self.gpu.stream().end_capture(
28878                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28879            );
28880            r?;
28881            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28882            graph.upload()?;
28883            Ok(graph)
28884        };
28885        let result = run();
28886        if was_tracking {
28887            unsafe {
28888                self.gpu.ctx.enable_event_tracking();
28889            }
28890        }
28891        Ok((result?, Vec::new()))
28892    }
28893
28894    pub fn capture_graph<F>(
28895        &self,
28896        mut step: F,
28897    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
28898    where
28899        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28900    {
28901        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
28902        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
28903        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
28904        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
28905        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
28906        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
28907        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
28908        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
28909        let was_tracking = self.gpu.ctx.is_event_tracking();
28910        if was_tracking {
28911            unsafe {
28912                self.gpu.ctx.disable_event_tracking();
28913            }
28914        }
28915        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
28916        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
28917        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
28918        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
28919        // measure that scan's real cost on the generic path. Diagnostic door only; the
28920        // default stays AUTO_FREE until a measured A/B justifies moving it.
28921        let iflag = {
28922            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
28923            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
28924                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
28925                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
28926                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
28927                Ok("priority") => {
28928                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
28929                }
28930                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28931            })
28932        };
28933        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
28934        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
28935        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
28936        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
28937        // eager step executions and are node-count-invariant. Printing the split bounds the
28938        // refactor's ceiling instead of assuming it.
28939        let ct = {
28940            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28941            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
28942        };
28943        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
28944        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
28945        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
28946        // chased, and node-count-invariant, so no capture-body refactor could touch it.
28947        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
28948        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
28949        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
28950        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
28951        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
28952        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
28953        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
28954        // grow and never frees, resident counters/scratch, cache set in place), and the
28955        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
28956        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
28957        // settling and pool mapping. Arbitrated adversarially, not by taste:
28958        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
28959        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
28960        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
28961        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
28962        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
28963        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
28964        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
28965        let warmups = {
28966            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
28967            *W.get_or_init(|| {
28968                std::env::var("MEMRA_GRAPH_WARMUPS")
28969                    .ok()
28970                    .and_then(|v| v.parse().ok())
28971                    .filter(|n| *n >= 1)
28972                    .unwrap_or(1)
28973            })
28974        };
28975        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28976            let t_w = std::time::Instant::now();
28977            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
28978            for _ in 0..warmups {
28979                step(self)?;
28980            }
28981            self.gpu.stream().synchronize()?;
28982            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
28983            // capture the third run.
28984            let t_c = std::time::Instant::now();
28985            self.gpu
28986                .stream()
28987                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28988            // If the body errors mid-capture, end the capture before propagating so the stream isn't
28989            // left in a capturing state.
28990            let r = step(self);
28991            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
28992            let t_i = std::time::Instant::now();
28993            let g = self.gpu.stream().end_capture(iflag);
28994            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
28995            r?;
28996            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28997            let t_u = std::time::Instant::now();
28998            graph.upload()?;
28999            if ct {
29000                println!(
29001                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
29002                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
29003                    t_u.elapsed().as_secs_f64() * 1e3
29004                );
29005            }
29006            Ok(graph)
29007        };
29008        let result = run();
29009        if was_tracking {
29010            unsafe {
29011                self.gpu.ctx.enable_event_tracking();
29012            }
29013        }
29014        result
29015    }
29016
29017    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
29018    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29019    pub fn gdn_scan_s128_view(
29020        &self,
29021        q: &CudaSlice<f32>,
29022        k: &CudaSlice<f32>,
29023        v: &CudaSlice<f32>,
29024        g: &CudaSlice<f32>,
29025        beta: &CudaSlice<f32>,
29026        state_in: &cudarc::driver::CudaView<f32>,
29027        state_out: &mut cudarc::driver::CudaViewMut<f32>,
29028        o: &mut CudaSlice<f32>,
29029        n_head: usize,
29030        t: usize,
29031        scale: f32,
29032    ) -> Result<(), Box<dyn std::error::Error>> {
29033        let f = self.func("gdn_scan_s128");
29034        const S_V: u32 = 128;
29035        const WARP: u32 = 32;
29036        const COLS: u32 = 4;
29037        let cfg = LaunchConfig {
29038            grid_dim: (n_head as u32, 1, S_V / COLS),
29039            block_dim: (WARP, COLS, 1),
29040            shared_mem_bytes: 0,
29041        };
29042        let (h, ti) = (n_head as i32, t as i32);
29043        let __s_b = self.gpu.stream();
29044        let mut b = __s_b.launch_builder(&f);
29045        b.arg(q)
29046            .arg(k)
29047            .arg(v)
29048            .arg(g)
29049            .arg(beta)
29050            .arg(state_in)
29051            .arg(state_out)
29052            .arg(o)
29053            .arg(&h)
29054            .arg(&ti)
29055            .arg(&scale);
29056        unsafe {
29057            b.launch(cfg)?;
29058        }
29059        Ok(())
29060    }
29061
29062    /// conv1d where the input is a CudaView (resident conv state assembled in place).
29063    #[allow(clippy::too_many_arguments)]
29064    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29065    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29066    pub fn ssm_conv1d_view(
29067        &self,
29068        x: &cudarc::driver::CudaView<f32>,
29069        w: &CudaSlice<f32>,
29070        y: &mut CudaSlice<f32>,
29071        conv_dim: usize,
29072        t: usize,
29073        d_conv: usize,
29074        silu: bool,
29075    ) -> Result<(), Box<dyn std::error::Error>> {
29076        let f = self.func("ssm_conv1d_silu_f32");
29077        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
29078        let cfg = LaunchConfig {
29079            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
29080            block_dim: (256, 1, 1),
29081            shared_mem_bytes: 0,
29082        };
29083        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
29084        let __s_b = self.gpu.stream();
29085        let mut b = __s_b.launch_builder(&f);
29086        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
29087        unsafe {
29088            b.launch(cfg)?;
29089        }
29090        Ok(())
29091    }
29092
29093    /// Depthwise causal conv1d + optional SiLU.
29094    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
29095    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
29096    /// FUSED prefill conv (token-major input, zero left-state): replaces
29097    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
29098    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
29099    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29100    pub fn ssm_conv1d_tm(
29101        &self,
29102        qkv_tm: &CudaSlice<f32>,
29103        w: &CudaSlice<f32>,
29104        y: &mut CudaSlice<f32>,
29105        conv_dim: usize,
29106        t: usize,
29107        d_conv: usize,
29108    ) -> Result<(), Box<dyn std::error::Error>> {
29109        let f = self.func("ssm_conv1d_tm_f32");
29110        let cfg = LaunchConfig {
29111            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29112            block_dim: (256, 1, 1),
29113            shared_mem_bytes: 0,
29114        };
29115        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29116        let __s_b = self.gpu.stream();
29117        let mut b = __s_b.launch_builder(&f);
29118        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
29119        unsafe {
29120            b.launch(cfg)?;
29121        }
29122        Ok(())
29123    }
29124
29125    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
29126    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
29127    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
29128    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
29129    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
29130    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
29131    /// columns; the final ring == what T sequential decode ring rolls leave).
29132    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29133    pub fn ssm_conv1d_tm_state(
29134        &self,
29135        qkv_tm: &CudaSlice<f32>,
29136        conv_state: &mut CudaSlice<f32>,
29137        w: &CudaSlice<f32>,
29138        y: &mut CudaSlice<f32>,
29139        conv_dim: usize,
29140        t: usize,
29141        d_conv: usize,
29142    ) -> Result<(), Box<dyn std::error::Error>> {
29143        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
29144    }
29145
29146    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
29147    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
29148    #[allow(clippy::too_many_arguments)]
29149    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29150    pub fn ssm_conv1d_tm_state_pad(
29151        &self,
29152        qkv_tm: &CudaSlice<f32>,
29153        conv_state: &mut CudaSlice<f32>,
29154        w: &CudaSlice<f32>,
29155        y: &mut CudaSlice<f32>,
29156        conv_dim: usize,
29157        t: usize,
29158        d_conv: usize,
29159        pad_len: Option<&CudaSlice<i32>>,
29160    ) -> Result<(), Box<dyn std::error::Error>> {
29161        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
29162        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
29163        // the window kernel both read the pre-roll ring; the roll launches after both) — but
29164        // cloning first keeps the ordering trivially correct under any future stream split.
29165        let ring_old = if t < d_conv - 1 {
29166            Some(self.clone_dtod(conv_state)?)
29167        } else {
29168            None
29169        };
29170        {
29171            let f = self.func("ssm_conv1d_tm_state_f32");
29172            let cfg = LaunchConfig {
29173                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29174                block_dim: (256, 1, 1),
29175                shared_mem_bytes: 0,
29176            };
29177            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29178            let __s_b = self.gpu.stream();
29179            let mut b = __s_b.launch_builder(&f);
29180            b.arg(qkv_tm)
29181                .arg(&*conv_state)
29182                .arg(w)
29183                .arg(y)
29184                .arg(&cd)
29185                .arg(&ti)
29186                .arg(&dc);
29187            unsafe {
29188                b.launch(cfg)?;
29189            }
29190        }
29191        match (ring_old, pad_len) {
29192            (None, Some(len_d)) => {
29193                let f = self.func("ssm_conv_ring_update_dev_f32");
29194                let n = conv_dim * (d_conv - 1);
29195                let cfg = LaunchConfig::for_num_elems(n as u32);
29196                let (cd, dc) = (conv_dim as i32, d_conv as i32);
29197                let __s_b = self.gpu.stream();
29198                let mut b = __s_b.launch_builder(&f);
29199                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
29200                unsafe {
29201                    b.launch(cfg)?;
29202                }
29203            }
29204            (None, None) => {
29205                let f = self.func("ssm_conv_ring_update_f32");
29206                let n = conv_dim * (d_conv - 1);
29207                let cfg = LaunchConfig::for_num_elems(n as u32);
29208                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29209                let __s_b = self.gpu.stream();
29210                let mut b = __s_b.launch_builder(&f);
29211                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
29212                unsafe {
29213                    b.launch(cfg)?;
29214                }
29215            }
29216            (Some(old), _) => {
29217                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
29218            }
29219        }
29220        Ok(())
29221    }
29222
29223    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
29224    #[allow(clippy::too_many_arguments)]
29225    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29226    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29227    pub fn ssm_conv1d_tm_state_pad_v(
29228        &self,
29229        qkv_tm: &cudarc::driver::CudaView<f32>,
29230        conv_state: &mut CudaSlice<f32>,
29231        w: &CudaSlice<f32>,
29232        y: &mut CudaSlice<f32>,
29233        conv_dim: usize,
29234        t: usize,
29235        d_conv: usize,
29236        pad_len: Option<&CudaSlice<i32>>,
29237    ) -> Result<(), Box<dyn std::error::Error>> {
29238        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
29239        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
29240        // the window kernel both read the pre-roll ring; the roll launches after both) — but
29241        // cloning first keeps the ordering trivially correct under any future stream split.
29242        let ring_old = if t < d_conv - 1 {
29243            Some(self.clone_dtod(conv_state)?)
29244        } else {
29245            None
29246        };
29247        {
29248            let f = self.func("ssm_conv1d_tm_state_f32");
29249            let cfg = LaunchConfig {
29250                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29251                block_dim: (256, 1, 1),
29252                shared_mem_bytes: 0,
29253            };
29254            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29255            let __s_b = self.gpu.stream();
29256            let mut b = __s_b.launch_builder(&f);
29257            b.arg(qkv_tm)
29258                .arg(&*conv_state)
29259                .arg(w)
29260                .arg(y)
29261                .arg(&cd)
29262                .arg(&ti)
29263                .arg(&dc);
29264            unsafe {
29265                b.launch(cfg)?;
29266            }
29267        }
29268        match (ring_old, pad_len) {
29269            (None, Some(len_d)) => {
29270                let f = self.func("ssm_conv_ring_update_dev_f32");
29271                let n = conv_dim * (d_conv - 1);
29272                let cfg = LaunchConfig::for_num_elems(n as u32);
29273                let (cd, dc) = (conv_dim as i32, d_conv as i32);
29274                let __s_b = self.gpu.stream();
29275                let mut b = __s_b.launch_builder(&f);
29276                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
29277                unsafe {
29278                    b.launch(cfg)?;
29279                }
29280            }
29281            (None, None) => {
29282                let f = self.func("ssm_conv_ring_update_f32");
29283                let n = conv_dim * (d_conv - 1);
29284                let cfg = LaunchConfig::for_num_elems(n as u32);
29285                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29286                let __s_b = self.gpu.stream();
29287                let mut b = __s_b.launch_builder(&f);
29288                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
29289                unsafe {
29290                    b.launch(cfg)?;
29291                }
29292            }
29293            (Some(_), _) => unreachable!(
29294                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
29295            ),
29296        }
29297        Ok(())
29298    }
29299
29300    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
29301    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
29302    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
29303    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
29304    pub fn ssm_conv_ring_rebuild(
29305        &self,
29306        qkv_tm: &CudaSlice<f32>,
29307        ring_old: &CudaSlice<f32>,
29308        conv_state: &mut CudaSlice<f32>,
29309        conv_dim: usize,
29310        tc: usize,
29311        d_conv: usize,
29312    ) -> Result<(), Box<dyn std::error::Error>> {
29313        let f = self.func("ssm_conv_ring_rebuild_f32");
29314        let n = conv_dim * (d_conv - 1);
29315        let cfg = LaunchConfig::for_num_elems(n as u32);
29316        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
29317        let __s_b = self.gpu.stream();
29318        let mut b = __s_b.launch_builder(&f);
29319        b.arg(qkv_tm)
29320            .arg(ring_old)
29321            .arg(conv_state)
29322            .arg(&cd)
29323            .arg(&ti)
29324            .arg(&dc);
29325        unsafe {
29326            b.launch(cfg)?;
29327        }
29328        Ok(())
29329    }
29330
29331    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
29332    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
29333    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
29334    /// the argmax + run-spec gates are the authority.
29335    #[allow(clippy::too_many_arguments)]
29336    pub fn gdn_prep_decode(
29337        &self,
29338        conv_out: &CudaSlice<f32>,
29339        beta_raw: &CudaSlice<f32>,
29340        alpha: &CudaSlice<f32>,
29341        dt_bias: &CudaSlice<f32>,
29342        a: &CudaSlice<f32>,
29343        q_l2: &mut CudaSlice<f32>,
29344        k_l2: &mut CudaSlice<f32>,
29345        v_g: &mut CudaSlice<f32>,
29346        beta: &mut CudaSlice<f32>,
29347        g_log: &mut CudaSlice<f32>,
29348        d_state: usize,
29349        num_v: usize,
29350        num_k: usize,
29351        key_dim: usize,
29352        eps: f32,
29353    ) -> Result<(), Box<dyn std::error::Error>> {
29354        let f = self.func("gdn_prep_decode_f32");
29355        let cfg = LaunchConfig {
29356            grid_dim: (num_v as u32, 1, 1),
29357            block_dim: (32, 4, 1),
29358            shared_mem_bytes: 0,
29359        };
29360        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
29361        let __s_b = self.gpu.stream();
29362        let mut b = __s_b.launch_builder(&f);
29363        b.arg(conv_out)
29364            .arg(beta_raw)
29365            .arg(alpha)
29366            .arg(dt_bias)
29367            .arg(a)
29368            .arg(q_l2)
29369            .arg(k_l2)
29370            .arg(v_g)
29371            .arg(beta)
29372            .arg(g_log)
29373            .arg(&ds)
29374            .arg(&nv)
29375            .arg(&nk)
29376            .arg(&kd)
29377            .arg(&eps);
29378        unsafe {
29379            b.launch(cfg)?;
29380        }
29381        Ok(())
29382    }
29383
29384    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
29385    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
29386    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
29387    #[allow(clippy::too_many_arguments)]
29388    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29389    pub fn ssm_conv1d_gdn(
29390        &self,
29391        qkv_tm: &CudaSlice<f32>,
29392        w: &CudaSlice<f32>,
29393        q_g: &mut CudaSlice<f32>,
29394        k_g: &mut CudaSlice<f32>,
29395        v_g: &mut CudaSlice<f32>,
29396        conv_dim: usize,
29397        t: usize,
29398        d_conv: usize,
29399        d_state: usize,
29400        num_v: usize,
29401        num_k: usize,
29402        key_dim: usize,
29403    ) -> Result<(), Box<dyn std::error::Error>> {
29404        let f = self.func("ssm_conv1d_gdn_f32");
29405        let cfg = LaunchConfig {
29406            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29407            block_dim: (256, 1, 1),
29408            shared_mem_bytes: 0,
29409        };
29410        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29411        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
29412        let __s_b = self.gpu.stream();
29413        let mut b = __s_b.launch_builder(&f);
29414        b.arg(qkv_tm)
29415            .arg(w)
29416            .arg(q_g)
29417            .arg(k_g)
29418            .arg(v_g)
29419            .arg(&cd)
29420            .arg(&ti)
29421            .arg(&dc)
29422            .arg(&ds)
29423            .arg(&nv)
29424            .arg(&nk)
29425            .arg(&kd);
29426        unsafe {
29427            b.launch(cfg)?;
29428        }
29429        Ok(())
29430    }
29431
29432    #[allow(clippy::too_many_arguments)]
29433    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29434    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29435    pub fn ssm_conv1d(
29436        &self,
29437        x: &CudaSlice<f32>,
29438        w: &CudaSlice<f32>,
29439        y: &mut CudaSlice<f32>,
29440        conv_dim: usize,
29441        t: usize,
29442        d_conv: usize,
29443        silu: bool,
29444    ) -> Result<(), Box<dyn std::error::Error>> {
29445        let f = self.func("ssm_conv1d_silu_f32");
29446        let cfg = LaunchConfig {
29447            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
29448            block_dim: (256, 1, 1),
29449            shared_mem_bytes: 0,
29450        };
29451        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
29452        let __s_b = self.gpu.stream();
29453        let mut b = __s_b.launch_builder(&f);
29454        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
29455        unsafe {
29456            b.launch(cfg)?;
29457        }
29458        Ok(())
29459    }
29460
29461    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
29462    /// o:[128,H,T]. Single sequence.
29463    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29464    pub fn gdn_scan_s128(
29465        &self,
29466        q: &CudaSlice<f32>,
29467        k: &CudaSlice<f32>,
29468        v: &CudaSlice<f32>,
29469        g: &CudaSlice<f32>,
29470        beta: &CudaSlice<f32>,
29471        state_in: &CudaSlice<f32>,
29472        state_out: &mut CudaSlice<f32>,
29473        o: &mut CudaSlice<f32>,
29474        n_head: usize,
29475        t: usize,
29476        scale: f32,
29477    ) -> Result<(), Box<dyn std::error::Error>> {
29478        let f = self.func("gdn_scan_s128");
29479        const S_V: u32 = 128;
29480        const WARP: u32 = 32;
29481        const COLS_PER_BLOCK: u32 = 4;
29482        let cfg = LaunchConfig {
29483            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
29484            block_dim: (WARP, COLS_PER_BLOCK, 1),
29485            shared_mem_bytes: 0,
29486        };
29487        let (h, ti) = (n_head as i32, t as i32);
29488        let __s_b = self.gpu.stream();
29489        let mut b = __s_b.launch_builder(&f);
29490        b.arg(q)
29491            .arg(k)
29492            .arg(v)
29493            .arg(g)
29494            .arg(beta)
29495            .arg(state_in)
29496            .arg(state_out)
29497            .arg(o)
29498            .arg(&h)
29499            .arg(&ti)
29500            .arg(&scale);
29501        unsafe {
29502            b.launch(cfg)?;
29503        }
29504        Ok(())
29505    }
29506
29507    // ==== B2' batched decode state ops (decode_batch.rs) ====
29508    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
29509    // Bodies are the single-seq kernels per sequence — bit-identical per row.
29510
29511    #[allow(clippy::too_many_arguments)]
29512    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29513    pub fn ssm_conv1d_fused_decode_b(
29514        &self,
29515        qkv_cols: &CudaSlice<f32>,
29516        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
29517        w: &CudaSlice<f32>,
29518        conv_outs: &mut CudaSlice<f32>,
29519        conv_dim: usize,
29520        d_conv: usize,
29521        b_n: usize,
29522    ) -> Result<(), Box<dyn std::error::Error>> {
29523        let f = self.func("ssm_conv1d_fused_decode_b_f32");
29524        let cfg = LaunchConfig {
29525            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
29526            block_dim: (256, 1, 1),
29527            shared_mem_bytes: 0,
29528        };
29529        let (cd, dc) = (conv_dim as i32, d_conv as i32);
29530        let __s_b = self.gpu.stream();
29531        let mut b = __s_b.launch_builder(&f);
29532        b.arg(qkv_cols)
29533            .arg(conv_state_ptrs)
29534            .arg(w)
29535            .arg(conv_outs)
29536            .arg(&cd)
29537            .arg(&dc);
29538        unsafe {
29539            b.launch(cfg)?;
29540        }
29541        Ok(())
29542    }
29543
29544    #[allow(clippy::too_many_arguments)]
29545    pub fn gdn_prep_decode_b(
29546        &self,
29547        conv_outs: &CudaSlice<f32>,
29548        beta_raws: &CudaSlice<f32>,
29549        alphas: &CudaSlice<f32>,
29550        dt_bias: &CudaSlice<f32>,
29551        a: &CudaSlice<f32>,
29552        q_l2: &mut CudaSlice<f32>,
29553        k_l2: &mut CudaSlice<f32>,
29554        v_g: &mut CudaSlice<f32>,
29555        beta: &mut CudaSlice<f32>,
29556        g_log: &mut CudaSlice<f32>,
29557        d_state: usize,
29558        num_v: usize,
29559        num_k: usize,
29560        key_dim: usize,
29561        eps: f32,
29562        conv_dim: usize,
29563        b_n: usize,
29564    ) -> Result<(), Box<dyn std::error::Error>> {
29565        let f = self.func("gdn_prep_decode_b_f32");
29566        let cfg = LaunchConfig {
29567            grid_dim: (num_v as u32, 1, b_n as u32),
29568            block_dim: (32, 4, 1),
29569            shared_mem_bytes: 0,
29570        };
29571        let (ds, nv, nk, kd, cd) = (
29572            d_state as i32,
29573            num_v as i32,
29574            num_k as i32,
29575            key_dim as i32,
29576            conv_dim as i32,
29577        );
29578        let __s_b = self.gpu.stream();
29579        let mut b = __s_b.launch_builder(&f);
29580        b.arg(conv_outs)
29581            .arg(beta_raws)
29582            .arg(alphas)
29583            .arg(dt_bias)
29584            .arg(a)
29585            .arg(q_l2)
29586            .arg(k_l2)
29587            .arg(v_g)
29588            .arg(beta)
29589            .arg(g_log)
29590            .arg(&ds)
29591            .arg(&nv)
29592            .arg(&nk)
29593            .arg(&kd)
29594            .arg(&eps)
29595            .arg(&cd);
29596        unsafe {
29597            b.launch(cfg)?;
29598        }
29599        Ok(())
29600    }
29601
29602    #[allow(clippy::too_many_arguments)]
29603    pub fn gdn_scan_s128_batched(
29604        &self,
29605        q: &CudaSlice<f32>,
29606        k: &CudaSlice<f32>,
29607        v: &CudaSlice<f32>,
29608        g: &CudaSlice<f32>,
29609        beta: &CudaSlice<f32>,
29610        state_in_ptrs: &cudarc::driver::CudaView<u64>,
29611        state_out_ptrs: &cudarc::driver::CudaView<u64>,
29612        o: &mut CudaSlice<f32>,
29613        n_head: usize,
29614        b_n: usize,
29615        scale: f32,
29616    ) -> Result<(), Box<dyn std::error::Error>> {
29617        let f = self.func("gdn_scan_s128_b");
29618        const S_V: u32 = 128;
29619        const WARP: u32 = 32;
29620        const COLS_PER_BLOCK: u32 = 4;
29621        let cfg = LaunchConfig {
29622            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
29623            block_dim: (WARP, COLS_PER_BLOCK, 1),
29624            shared_mem_bytes: 0,
29625        };
29626        let h = n_head as i32;
29627        let __s_b = self.gpu.stream();
29628        let mut b = __s_b.launch_builder(&f);
29629        b.arg(q)
29630            .arg(k)
29631            .arg(v)
29632            .arg(g)
29633            .arg(beta)
29634            .arg(state_in_ptrs)
29635            .arg(state_out_ptrs)
29636            .arg(o)
29637            .arg(&h)
29638            .arg(&scale);
29639        unsafe {
29640            b.launch(cfg)?;
29641        }
29642        Ok(())
29643    }
29644
29645    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
29646    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
29647    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
29648    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
29649    /// numeric class; only the pointer arithmetic moved host-side.
29650    #[allow(clippy::too_many_arguments)]
29651    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29652    pub fn ssm_conv1d_fused_decode_b_view(
29653        &self,
29654        qkv_cols: &cudarc::driver::CudaView<f32>,
29655        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
29656        w: &CudaSlice<f32>,
29657        conv_outs: &mut CudaSlice<f32>,
29658        conv_dim: usize,
29659        d_conv: usize,
29660        b_n: usize,
29661    ) -> Result<(), Box<dyn std::error::Error>> {
29662        let f = self.func("ssm_conv1d_fused_decode_b_f32");
29663        let cfg = LaunchConfig {
29664            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
29665            block_dim: (256, 1, 1),
29666            shared_mem_bytes: 0,
29667        };
29668        let (cd, dc) = (conv_dim as i32, d_conv as i32);
29669        let __s_b = self.gpu.stream();
29670        let mut b = __s_b.launch_builder(&f);
29671        b.arg(qkv_cols)
29672            .arg(conv_state_ptrs)
29673            .arg(w)
29674            .arg(conv_outs)
29675            .arg(&cd)
29676            .arg(&dc);
29677        unsafe {
29678            b.launch(cfg)?;
29679        }
29680        Ok(())
29681    }
29682
29683    #[allow(clippy::too_many_arguments)]
29684    pub fn gdn_prep_decode_b_view(
29685        &self,
29686        conv_outs: &CudaSlice<f32>,
29687        beta_raws: &cudarc::driver::CudaView<f32>,
29688        alphas: &cudarc::driver::CudaView<f32>,
29689        dt_bias: &CudaSlice<f32>,
29690        a: &CudaSlice<f32>,
29691        q_l2: &mut CudaSlice<f32>,
29692        k_l2: &mut CudaSlice<f32>,
29693        v_g: &mut CudaSlice<f32>,
29694        beta: &mut CudaSlice<f32>,
29695        g_log: &mut CudaSlice<f32>,
29696        d_state: usize,
29697        num_v: usize,
29698        num_k: usize,
29699        key_dim: usize,
29700        eps: f32,
29701        conv_dim: usize,
29702        b_n: usize,
29703    ) -> Result<(), Box<dyn std::error::Error>> {
29704        let f = self.func("gdn_prep_decode_b_f32");
29705        let cfg = LaunchConfig {
29706            grid_dim: (num_v as u32, 1, b_n as u32),
29707            block_dim: (32, 4, 1),
29708            shared_mem_bytes: 0,
29709        };
29710        let (ds, nv, nk, kd, cd) = (
29711            d_state as i32,
29712            num_v as i32,
29713            num_k as i32,
29714            key_dim as i32,
29715            conv_dim as i32,
29716        );
29717        let __s_b = self.gpu.stream();
29718        let mut b = __s_b.launch_builder(&f);
29719        b.arg(conv_outs)
29720            .arg(beta_raws)
29721            .arg(alphas)
29722            .arg(dt_bias)
29723            .arg(a)
29724            .arg(q_l2)
29725            .arg(k_l2)
29726            .arg(v_g)
29727            .arg(beta)
29728            .arg(g_log)
29729            .arg(&ds)
29730            .arg(&nv)
29731            .arg(&nk)
29732            .arg(&kd)
29733            .arg(&eps)
29734            .arg(&cd);
29735        unsafe {
29736            b.launch(cfg)?;
29737        }
29738        Ok(())
29739    }
29740
29741    #[allow(clippy::too_many_arguments)]
29742    pub fn gdn_scan_s128_batched_view(
29743        &self,
29744        q: &CudaSlice<f32>,
29745        k: &CudaSlice<f32>,
29746        v: &CudaSlice<f32>,
29747        g: &CudaSlice<f32>,
29748        beta: &CudaSlice<f32>,
29749        state_in_ptrs: &cudarc::driver::CudaView<u64>,
29750        state_out_ptrs: &cudarc::driver::CudaView<u64>,
29751        o: &mut cudarc::driver::CudaViewMut<f32>,
29752        n_head: usize,
29753        b_n: usize,
29754        scale: f32,
29755    ) -> Result<(), Box<dyn std::error::Error>> {
29756        let f = self.func("gdn_scan_s128_b");
29757        const S_V: u32 = 128;
29758        const WARP: u32 = 32;
29759        const COLS_PER_BLOCK: u32 = 4;
29760        let cfg = LaunchConfig {
29761            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
29762            block_dim: (WARP, COLS_PER_BLOCK, 1),
29763            shared_mem_bytes: 0,
29764        };
29765        let h = n_head as i32;
29766        let __s_b = self.gpu.stream();
29767        let mut b = __s_b.launch_builder(&f);
29768        b.arg(q)
29769            .arg(k)
29770            .arg(v)
29771            .arg(g)
29772            .arg(beta)
29773            .arg(state_in_ptrs)
29774            .arg(state_out_ptrs)
29775            .arg(o)
29776            .arg(&h)
29777            .arg(&scale);
29778        unsafe {
29779            b.launch(cfg)?;
29780        }
29781        Ok(())
29782    }
29783
29784    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
29785    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
29786    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
29787    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
29788    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
29789    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
29790    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
29791    /// identity law); prime_cache/forward/forward_last are the only callers.
29792    pub fn gdn_chunked_enabled() -> bool {
29793        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
29794        *E.get_or_init(|| {
29795            std::env::var("MEMRA_GDN_CHUNKED")
29796                .map(|v| v != "0")
29797                .unwrap_or(true)
29798        })
29799    }
29800
29801    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
29802    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
29803    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
29804    /// of 32 in [32, 128] (kernel row mappings require it).
29805    pub fn gdn_chunk_size() -> usize {
29806        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29807        *C.get_or_init(|| {
29808            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
29809                .ok()
29810                .and_then(|v| v.parse().ok())
29811                .unwrap_or(32);
29812            c.clamp(32, 128) / 32 * 32
29813        })
29814    }
29815
29816    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
29817    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
29818    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
29819    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
29820    #[allow(clippy::too_many_arguments)]
29821    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
29822    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
29823    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
29824    #[allow(clippy::too_many_arguments)]
29825    pub fn gdn_chunk_k123(
29826        &self,
29827        q: &CudaSlice<f32>,
29828        k: &CudaSlice<f32>,
29829        v: &CudaSlice<f32>,
29830        g: &CudaSlice<f32>,
29831        beta: &CudaSlice<f32>,
29832        wb16: Option<&mut CudaSlice<u8>>,
29833        n_head: usize,
29834        t: usize,
29835        c: usize,
29836        hk: usize,
29837        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
29838    ) -> Result<
29839        (
29840            CudaSlice<f32>,
29841            CudaSlice<f32>,
29842            CudaSlice<f32>,
29843            CudaSlice<f32>,
29844        ),
29845        Box<dyn std::error::Error>,
29846    > {
29847        const D: usize = 128;
29848        let h = n_head;
29849        #[allow(clippy::manual_div_ceil)]
29850        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29851        let nc = (t + c - 1) / c;
29852        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
29853        let mut gcum = self.uninit(t * h)?;
29854        let mut a = self.uninit(nc * h * c * c)?;
29855        let mut p = self.uninit(nc * h * c * c)?;
29856        let mut u = self.uninit(nc * h * c * D)?;
29857        let mut w = self.uninit(nc * h * c * D)?;
29858        {
29859            // K1
29860            let f = self.func("gdn_chunk_cumgate_f32");
29861            let cfg = LaunchConfig {
29862                grid_dim: (nc as u32, h as u32, 1),
29863                block_dim: (32, 1, 1),
29864                shared_mem_bytes: 0,
29865            };
29866            let __s_b = self.gpu.stream();
29867            let mut b = __s_b.launch_builder(&f);
29868            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
29869            unsafe {
29870                b.launch(cfg)?;
29871            }
29872        }
29873        if let Some((qb, kb, pb)) = k2w {
29874            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
29875            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
29876            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
29877            let f = self.func("gdn_k2_wgmma");
29878            let cfg = LaunchConfig {
29879                grid_dim: (nc as u32, h as u32, 1),
29880                block_dim: (128, 1, 1),
29881                shared_mem_bytes: 0,
29882            };
29883            let hki = hk as i32;
29884            let __s_b = self.gpu.stream();
29885            let mut b = __s_b.launch_builder(&f);
29886            b.arg(qb)
29887                .arg(kb)
29888                .arg(&gcum)
29889                .arg(beta)
29890                .arg(&mut a)
29891                .arg(&mut *pb)
29892                .arg(&hi)
29893                .arg(&ti)
29894                .arg(&ci)
29895                .arg(&hki);
29896            unsafe {
29897                b.launch(cfg)?;
29898            }
29899        } else if c <= 64 && !portable_mma_gated() {
29900            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
29901            let f = self.func("gdn_chunk_attn_f32");
29902            f.set_attribute(
29903                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29904                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
29905            )?;
29906            #[allow(clippy::manual_div_ceil)]
29907            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29908            let jt = ((c + 31) / 32) as u32;
29909            let cfg = LaunchConfig {
29910                grid_dim: (nc as u32, h as u32, jt),
29911                block_dim: (256, 1, 1),
29912                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
29913            };
29914            let hki = hk as i32;
29915            let __s_b = self.gpu.stream();
29916            let mut b = __s_b.launch_builder(&f);
29917            b.arg(q)
29918                .arg(k)
29919                .arg(&gcum)
29920                .arg(beta)
29921                .arg(&mut a)
29922                .arg(&mut p)
29923                .arg(&hi)
29924                .arg(&ti)
29925                .arg(&ci)
29926                .arg(&hki);
29927            unsafe {
29928                b.launch(cfg)?;
29929            }
29930        } else {
29931            // K2 generic (C = 128, or the portable target's low-smem fallback)
29932            assert!(
29933                hk == h,
29934                "generic K2 is broadcast-only (de-broadcast rides C==32)"
29935            );
29936            let f = self.func("gdn_chunk_attn_g_f32");
29937            let cfg = LaunchConfig {
29938                grid_dim: (nc as u32, h as u32, 1),
29939                block_dim: (32, 8, 1),
29940                shared_mem_bytes: 0,
29941            };
29942            let __s_b = self.gpu.stream();
29943            let mut b = __s_b.launch_builder(&f);
29944            b.arg(q)
29945                .arg(k)
29946                .arg(&gcum)
29947                .arg(beta)
29948                .arg(&mut a)
29949                .arg(&mut p)
29950                .arg(&hi)
29951                .arg(&ti)
29952                .arg(&ci);
29953            unsafe {
29954                b.launch(cfg)?;
29955            }
29956        }
29957        {
29958            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
29959            let cfg = LaunchConfig {
29960                grid_dim: (nc as u32, h as u32, 1),
29961                block_dim: (256, 1, 1),
29962                shared_mem_bytes: 0,
29963            };
29964            match c {
29965                32 | 64 => {
29966                    let f = self.func(if c == 32 {
29967                        "gdn_chunk_solve32_f32"
29968                    } else {
29969                        "gdn_chunk_solve64_f32"
29970                    });
29971                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
29972                    let wb: u64 = match wb16 {
29973                        Some(d) => self.addr_u8(d),
29974                        None => 0,
29975                    };
29976                    let hki = hk as i32;
29977                    let __s_b = self.gpu.stream();
29978                    let mut b = __s_b.launch_builder(&f);
29979                    b.arg(v)
29980                        .arg(k)
29981                        .arg(&a)
29982                        .arg(&gcum)
29983                        .arg(&mut u)
29984                        .arg(&mut w)
29985                        .arg(&wb)
29986                        .arg(&hi)
29987                        .arg(&ti)
29988                        .arg(&hki);
29989                    unsafe {
29990                        b.launch(cfg)?;
29991                    }
29992                }
29993                _ => {
29994                    assert!(hk == h, "generic K3 is broadcast-only");
29995                    let f = self.func("gdn_chunk_solve_f32");
29996                    let __s_b = self.gpu.stream();
29997                    let mut b = __s_b.launch_builder(&f);
29998                    b.arg(v)
29999                        .arg(k)
30000                        .arg(&a)
30001                        .arg(&gcum)
30002                        .arg(&mut u)
30003                        .arg(&mut w)
30004                        .arg(&hi)
30005                        .arg(&ti)
30006                        .arg(&ci);
30007                    unsafe {
30008                        b.launch(cfg)?;
30009                    }
30010                }
30011            }
30012        }
30013        Ok((gcum, p, u, w))
30014    }
30015
30016    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
30017    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
30018    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
30019    pub fn gdn_db_on() -> bool {
30020        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
30021    }
30022
30023    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
30024    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
30025    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
30026    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
30027    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
30028    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
30029    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
30030    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
30031    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
30032        !portable_mma_gated()
30033            && c == 32
30034            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
30035                Ok("1") => true,
30036                Ok("0") => false,
30037                _ => gdn_mma_default_on(),
30038            }
30039    }
30040
30041    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
30042    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
30043    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
30044    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
30045    /// force would silently produce garbage. Required since the sm_120a mma default
30046    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
30047    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
30048        cfg!(memra_hopper_mma)
30049            && self.gdn_mma_enabled(c)
30050            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
30051    }
30052
30053    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
30054    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
30055    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
30056    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
30057    #[allow(clippy::too_many_arguments)]
30058    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30059    pub fn ssm_conv1d_gdn_state_pad(
30060        &self,
30061        qkv_tm: &cudarc::driver::CudaView<f32>,
30062        conv_state: &mut CudaSlice<f32>,
30063        w: &CudaSlice<f32>,
30064        q_g: &mut CudaSlice<f32>,
30065        k_g: &mut CudaSlice<f32>,
30066        v_g: &mut CudaSlice<f32>,
30067        conv_dim: usize,
30068        t: usize,
30069        d_conv: usize,
30070        d_state: usize,
30071        num_v: usize,
30072        num_k: usize,
30073        key_dim: usize,
30074        hk: usize,
30075        pad_len: Option<&CudaSlice<i32>>,
30076    ) -> Result<(), Box<dyn std::error::Error>> {
30077        assert!(
30078            t >= d_conv - 1,
30079            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
30080        );
30081        {
30082            let f = self.func("ssm_conv1d_gdn_state_f32");
30083            let cfg = LaunchConfig {
30084                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
30085                block_dim: (256, 1, 1),
30086                shared_mem_bytes: 0,
30087            };
30088            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
30089            let (ds, nv, nk, kd, hki) = (
30090                d_state as i32,
30091                num_v as i32,
30092                num_k as i32,
30093                key_dim as i32,
30094                hk as i32,
30095            );
30096            let __s_b = self.gpu.stream();
30097            let mut b = __s_b.launch_builder(&f);
30098            b.arg(qkv_tm)
30099                .arg(&*conv_state)
30100                .arg(w)
30101                .arg(q_g)
30102                .arg(k_g)
30103                .arg(v_g)
30104                .arg(&cd)
30105                .arg(&ti)
30106                .arg(&dc)
30107                .arg(&ds)
30108                .arg(&nv)
30109                .arg(&nk)
30110                .arg(&kd)
30111                .arg(&hki);
30112            unsafe {
30113                b.launch(cfg)?;
30114            }
30115        }
30116        match pad_len {
30117            Some(len_d) => {
30118                let f = self.func("ssm_conv_ring_update_dev_f32");
30119                let n = conv_dim * (d_conv - 1);
30120                let cfg = LaunchConfig::for_num_elems(n as u32);
30121                let (cd, dc) = (conv_dim as i32, d_conv as i32);
30122                let __s_b = self.gpu.stream();
30123                let mut b = __s_b.launch_builder(&f);
30124                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
30125                unsafe {
30126                    b.launch(cfg)?;
30127                }
30128            }
30129            None => {
30130                let f = self.func("ssm_conv_ring_update_f32");
30131                let n = conv_dim * (d_conv - 1);
30132                let cfg = LaunchConfig::for_num_elems(n as u32);
30133                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
30134                let __s_b = self.gpu.stream();
30135                let mut b = __s_b.launch_builder(&f);
30136                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
30137                unsafe {
30138                    b.launch(cfg)?;
30139                }
30140            }
30141        }
30142        Ok(())
30143    }
30144
30145    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
30146    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
30147    /// K2/K3 can write them.
30148    pub fn gdn_chunk_alloc(
30149        &self,
30150        n_head: usize,
30151        t: usize,
30152        c: usize,
30153        hk: usize,
30154    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
30155        const D: usize = 128;
30156        assert!(
30157            c == 32,
30158            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
30159        );
30160        let h = n_head;
30161        #[allow(clippy::manual_div_ceil)]
30162        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30163        let nc = (t + c - 1) / c;
30164        Ok(GdnChunkBufs {
30165            gcum: self.uninit(t * h)?,
30166            a: self.uninit(nc * h * c * c)?,
30167            p: self.uninit(nc * h * c * c)?,
30168            u: self.uninit(nc * h * c * D)?,
30169            w: self.uninit(nc * h * c * D)?,
30170            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
30171            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
30172            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
30173            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
30174            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
30175            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
30176            o: self.uninit(D * h * t)?,
30177            t,
30178            nc,
30179        })
30180    }
30181
30182    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
30183    pub fn f32_to_bf16_v(
30184        &self,
30185        x: &cudarc::driver::CudaView<f32>,
30186        dst: &mut CudaSlice<u8>,
30187        n: usize,
30188    ) -> Result<(), Box<dyn std::error::Error>> {
30189        let f = self.func("f32_to_bf16_bulk");
30190        let ni = n as i64;
30191        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
30192        let __s_b = self.gpu.stream();
30193        let mut b = __s_b.launch_builder(&f);
30194        b.arg(x).arg(dst).arg(&ni);
30195        unsafe {
30196            b.launch(cfg)?;
30197        }
30198        Ok(())
30199    }
30200
30201    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
30202    pub fn f32_to_bf16_into(
30203        &self,
30204        x: &CudaSlice<f32>,
30205        dst: &mut CudaSlice<u8>,
30206        n: usize,
30207    ) -> Result<(), Box<dyn std::error::Error>> {
30208        let f = self.func("f32_to_bf16_bulk");
30209        let ni = n as i64;
30210        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
30211        let __s_b = self.gpu.stream();
30212        let mut b = __s_b.launch_builder(&f);
30213        b.arg(x).arg(dst).arg(&ni);
30214        unsafe {
30215            b.launch(cfg)?;
30216        }
30217        Ok(())
30218    }
30219
30220    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
30221    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
30222    pub fn gdn_chunk_k123_vl8(
30223        &self,
30224        seqs: &[GdnSeqVl],
30225        n_head: usize,
30226        hk: usize,
30227        wq: Option<&GdnWVl8>,
30228    ) -> Result<(), Box<dyn std::error::Error>> {
30229        let b = seqs.len();
30230        assert!((1..=8).contains(&b), "gdn_chunk_k123_vl8: 1..=8 sequences");
30231        let mut packed = [GdnSeqVl::default(); 8];
30232        packed[..b].copy_from_slice(seqs);
30233        let v = GdnVl8(packed);
30234        let (hi, ci) = (n_head as i32, 32i32);
30235        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
30236        {
30237            let f = self.func("gdn_chunk_cumgate_vl");
30238            let cfg = LaunchConfig {
30239                grid_dim: (max_nc, n_head as u32, b as u32),
30240                block_dim: (32, 1, 1),
30241                shared_mem_bytes: 0,
30242            };
30243            let __s_lb = self.gpu.stream();
30244            let mut lb = __s_lb.launch_builder(&f);
30245            lb.arg(&v).arg(&hi).arg(&ci);
30246            unsafe {
30247                lb.launch(cfg)?;
30248            }
30249        }
30250        let hki = hk as i32;
30251        if let Some(w) = wq {
30252            // K2-wgmma vl twin (writes A + pre-masked Pb16)
30253            let f = self.func("gdn_k2_wgmma_vl");
30254            let cfg = LaunchConfig {
30255                grid_dim: (max_nc, n_head as u32, b as u32),
30256                block_dim: (128, 1, 1),
30257                shared_mem_bytes: 0,
30258            };
30259            let __s_lb = self.gpu.stream();
30260            let mut lb = __s_lb.launch_builder(&f);
30261            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
30262            unsafe {
30263                lb.launch(cfg)?;
30264            }
30265        } else {
30266            let f = self.func("gdn_chunk_attn_vl");
30267            f.set_attribute(
30268                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30269                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
30270            )?;
30271            let cfg = LaunchConfig {
30272                grid_dim: (max_nc, n_head as u32, b as u32),
30273                block_dim: (256, 1, 1),
30274                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
30275            };
30276            let __s_lb = self.gpu.stream();
30277            let mut lb = __s_lb.launch_builder(&f);
30278            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30279            unsafe {
30280                lb.launch(cfg)?;
30281            }
30282        }
30283        {
30284            let f = self.func("gdn_chunk_solve32_vl");
30285            let cfg = LaunchConfig {
30286                grid_dim: (max_nc, n_head as u32, b as u32),
30287                block_dim: (256, 1, 1),
30288                shared_mem_bytes: 0,
30289            };
30290            let __s_lb = self.gpu.stream();
30291            let mut lb = __s_lb.launch_builder(&f);
30292            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30293            unsafe {
30294                lb.launch(cfg)?;
30295            }
30296        }
30297        Ok(())
30298    }
30299
30300    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
30301    /// fused gate-prep, 5 launches for every sequence (per-element math identical
30302    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
30303    #[allow(clippy::too_many_arguments)]
30304    pub fn gdn_prep_vl8(
30305        &self,
30306        seqs: &[GdnPrepVl],
30307        conv_w: &CudaSlice<f32>,
30308        dt_bias: &CudaSlice<f32>,
30309        a: &CudaSlice<f32>,
30310        conv_dim: usize,
30311        d_conv: usize,
30312        d_state: usize,
30313        num_v: usize,
30314        num_k: usize,
30315        key_dim: usize,
30316        hk: usize,
30317        eps: f32,
30318    ) -> Result<(), Box<dyn std::error::Error>> {
30319        let b = seqs.len();
30320        assert!((1..=8).contains(&b));
30321        let mut packed = [GdnPrepVl::default(); 8];
30322        packed[..b].copy_from_slice(seqs);
30323        let v = GdnPrepVl8(packed);
30324        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
30325        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
30326        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
30327        assert!(
30328            conv_fuse || hk == num_v,
30329            "de-broadcast requires the fused conv"
30330        );
30331        if conv_fuse {
30332            let f = self.func("ssm_conv1d_gdn_state_vl");
30333            let cfg = LaunchConfig {
30334                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
30335                block_dim: (256, 1, 1),
30336                shared_mem_bytes: 0,
30337            };
30338            let (dsi, nvi, nki, kdi, hki) = (
30339                d_state as i32,
30340                num_v as i32,
30341                num_k as i32,
30342                key_dim as i32,
30343                hk as i32,
30344            );
30345            let __s_lb = self.gpu.stream();
30346            let mut lb = __s_lb.launch_builder(&f);
30347            lb.arg(&v)
30348                .arg(conv_w)
30349                .arg(&cdi)
30350                .arg(&dci)
30351                .arg(&dsi)
30352                .arg(&nvi)
30353                .arg(&nki)
30354                .arg(&kdi)
30355                .arg(&hki);
30356            unsafe {
30357                lb.launch(cfg)?;
30358            }
30359        } else {
30360            let f = self.func("ssm_conv1d_tm_state_vl");
30361            let cfg = LaunchConfig {
30362                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
30363                block_dim: (256, 1, 1),
30364                shared_mem_bytes: 0,
30365            };
30366            let __s_lb = self.gpu.stream();
30367            let mut lb = __s_lb.launch_builder(&f);
30368            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
30369            unsafe {
30370                lb.launch(cfg)?;
30371            }
30372        }
30373        {
30374            let f = self.func("ssm_conv_ring_update_vl");
30375            let n = (conv_dim * (d_conv - 1)) as u32;
30376            let cfg = LaunchConfig {
30377                grid_dim: (n.div_ceil(256), 1, b as u32),
30378                block_dim: (256, 1, 1),
30379                shared_mem_bytes: 0,
30380            };
30381            let __s_lb = self.gpu.stream();
30382            let mut lb = __s_lb.launch_builder(&f);
30383            lb.arg(&v).arg(&cdi).arg(&dci);
30384            unsafe {
30385                lb.launch(cfg)?;
30386            }
30387        }
30388        if !conv_fuse {
30389            let f = self.func("qkv_to_gdn_repack_vl");
30390            let n = max_t * (num_v * d_state) as u32;
30391            let cfg = LaunchConfig {
30392                grid_dim: (n.div_ceil(256), 1, b as u32),
30393                block_dim: (256, 1, 1),
30394                shared_mem_bytes: 0,
30395            };
30396            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
30397            let __s_lb = self.gpu.stream();
30398            let mut lb = __s_lb.launch_builder(&f);
30399            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
30400            unsafe {
30401                lb.launch(cfg)?;
30402            }
30403        }
30404        if Self::l2_v2_on(d_state) {
30405            let f = self.func("gdn_l2_v2_vl");
30406            let cfg = LaunchConfig {
30407                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
30408                block_dim: (256, 1, 1),
30409                shared_mem_bytes: 0,
30410            };
30411            let (dsi, nvi) = (d_state as i32, hk as i32);
30412            let __s_lb = self.gpu.stream();
30413            let mut lb = __s_lb.launch_builder(&f);
30414            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
30415            unsafe {
30416                lb.launch(cfg)?;
30417            }
30418        } else {
30419            let f = self.func("gdn_l2_vl");
30420            let cfg = LaunchConfig {
30421                grid_dim: (max_t * hk as u32, 2, b as u32),
30422                block_dim: (256, 1, 1),
30423                shared_mem_bytes: 0,
30424            };
30425            let (dsi, nvi) = (d_state as i32, hk as i32);
30426            let __s_lb = self.gpu.stream();
30427            let mut lb = __s_lb.launch_builder(&f);
30428            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
30429            unsafe {
30430                lb.launch(cfg)?;
30431            }
30432        }
30433        {
30434            let f = self.func("gdn_gate_prep_vl");
30435            let n = max_t * num_v as u32;
30436            let cfg = LaunchConfig {
30437                grid_dim: (n.div_ceil(256), 1, b as u32),
30438                block_dim: (256, 1, 1),
30439                shared_mem_bytes: 0,
30440            };
30441            let nvi = num_v as i32;
30442            let __s_lb = self.gpu.stream();
30443            let mut lb = __s_lb.launch_builder(&f);
30444            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
30445            unsafe {
30446                lb.launch(cfg)?;
30447            }
30448        }
30449        Ok(())
30450    }
30451
30452    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
30453    pub fn gdn_mirror_vl8(
30454        &self,
30455        seqs: &[GdnSeqVl],
30456        n_head: usize,
30457        which: i32,
30458        hk: usize,
30459    ) -> Result<(), Box<dyn std::error::Error>> {
30460        let b = seqs.len();
30461        assert!((1..=8).contains(&b));
30462        let mut packed = [GdnSeqVl::default(); 8];
30463        packed[..b].copy_from_slice(seqs);
30464        let v = GdnVl8(packed);
30465        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
30466        let max_n = seqs
30467            .iter()
30468            .map(|s| {
30469                if which == 0 {
30470                    s.t as i64 * ept as i64
30471                } else {
30472                    s.nc as i64 * ept as i64 * 32
30473                }
30474            })
30475            .max()
30476            .unwrap();
30477        let f = self.func("gdn_mirror_vl");
30478        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
30479        let cfg = LaunchConfig {
30480            grid_dim: (blocks, 1, b as u32),
30481            block_dim: (256, 1, 1),
30482            shared_mem_bytes: 0,
30483        };
30484        let __s_lb = self.gpu.stream();
30485        let mut lb = __s_lb.launch_builder(&f);
30486        lb.arg(&v).arg(&ept).arg(&which);
30487        unsafe {
30488            lb.launch(cfg)?;
30489        }
30490        Ok(())
30491    }
30492
30493    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
30494    pub fn gdn_tail_vl8(
30495        &self,
30496        seqs: &[GdnPrepVl],
30497        norm_w: &CudaSlice<f32>,
30498        d_state: usize,
30499        num_v: usize,
30500        eps: f32,
30501    ) -> Result<(), Box<dyn std::error::Error>> {
30502        let b = seqs.len();
30503        assert!((1..=8).contains(&b));
30504        let mut packed = [GdnPrepVl::default(); 8];
30505        packed[..b].copy_from_slice(seqs);
30506        let v = GdnPrepVl8(packed);
30507        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
30508        let f = self.func("gated_rmsnorm_f16out_vl");
30509        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
30510        let cfg = LaunchConfig {
30511            grid_dim: (max_t * num_v as u32, 1, b as u32),
30512            block_dim: (128, 1, 1),
30513            shared_mem_bytes: 0,
30514        };
30515        let (dsi, nvi) = (d_state as i32, num_v as i32);
30516        let __s_lb = self.gpu.stream();
30517        let mut lb = __s_lb.launch_builder(&f);
30518        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
30519        unsafe {
30520            lb.launch(cfg)?;
30521        }
30522        Ok(())
30523    }
30524
30525    /// Raw device address helpers for the varlen by-value arg struct (single-stream
30526    /// launches; every buffer outlives the call — the f16 FFI discipline).
30527    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
30528        use cudarc::driver::DevicePtr;
30529        let s = self.gpu.stream();
30530        let (p, _g) = x.device_ptr(&s);
30531        p
30532    }
30533    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
30534        use cudarc::driver::DevicePtrMut;
30535        let s = self.gpu.stream();
30536        let (p, _g) = x.device_ptr_mut(&s);
30537        p
30538    }
30539    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
30540        use cudarc::driver::DevicePtr;
30541        let s = self.gpu.stream();
30542        let (p, _g) = x.device_ptr(&s);
30543        p
30544    }
30545    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
30546        use cudarc::driver::DevicePtr;
30547        let s = self.gpu.stream();
30548        let (p, _g) = x.device_ptr(&s);
30549        p
30550    }
30551
30552    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
30553    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
30554    /// launches, so this is strictly bit-gateable against them).
30555    pub fn gdn_chunk_vl8(
30556        &self,
30557        seqs: &[GdnSeqVl],
30558        n_head: usize,
30559        scale: f32,
30560        hk: usize,
30561        wq: Option<&GdnWVl8>,
30562    ) -> Result<(), Box<dyn std::error::Error>> {
30563        const NSPLIT: u32 = 4;
30564        let b = seqs.len();
30565        assert!((1..=8).contains(&b), "gdn_chunk_vl8: 1..=8 sequences");
30566        let mut packed = [GdnSeqVl::default(); 8];
30567        packed[..b].copy_from_slice(seqs);
30568        let v = GdnVl8(packed);
30569        let (hi, ci) = (n_head as i32, 32i32);
30570        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
30571        let hki = hk as i32;
30572        if let Some(w) = wq {
30573            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
30574            let f = self.func("gdn_k45_wgmma_vl");
30575            let cfg = LaunchConfig {
30576                grid_dim: (n_head as u32, NSPLIT, b as u32),
30577                block_dim: (256, 1, 1),
30578                shared_mem_bytes: 0,
30579            };
30580            let __s_lb = self.gpu.stream();
30581            let mut lb = __s_lb.launch_builder(&f);
30582            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
30583            unsafe {
30584                lb.launch(cfg)?;
30585            }
30586            let _ = max_nc;
30587            return Ok(());
30588        }
30589        {
30590            let f = self.func("gdn_chunk_state_mma_vl");
30591            let cfg = LaunchConfig {
30592                grid_dim: (n_head as u32, NSPLIT, b as u32),
30593                block_dim: (256, 1, 1),
30594                shared_mem_bytes: 0,
30595            };
30596            let __s_lb = self.gpu.stream();
30597            let mut lb = __s_lb.launch_builder(&f);
30598            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30599            unsafe {
30600                lb.launch(cfg)?;
30601            }
30602        }
30603        {
30604            let f = self.func("gdn_chunk_output_mma_vl");
30605            let cfg = LaunchConfig {
30606                grid_dim: (max_nc, n_head as u32, b as u32),
30607                block_dim: (256, 1, 1),
30608                shared_mem_bytes: 0,
30609            };
30610            let __s_lb = self.gpu.stream();
30611            let mut lb = __s_lb.launch_builder(&f);
30612            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
30613            unsafe {
30614                lb.launch(cfg)?;
30615            }
30616        }
30617        Ok(())
30618    }
30619    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
30620    pub fn gdn_scan_chunked(
30621        &self,
30622        q: &CudaSlice<f32>,
30623        k: &CudaSlice<f32>,
30624        v: &CudaSlice<f32>,
30625        g: &CudaSlice<f32>,
30626        beta: &CudaSlice<f32>,
30627        kb16_pre: Option<&CudaSlice<u8>>,
30628        qb16_pre: Option<&CudaSlice<u8>>,
30629        state_in: &CudaSlice<f32>,
30630        state_out: &mut CudaSlice<f32>,
30631        o: &mut CudaSlice<f32>,
30632        n_head: usize,
30633        t: usize,
30634        scale: f32,
30635        c: usize,
30636        hk: usize,
30637    ) -> Result<(), Box<dyn std::error::Error>> {
30638        const D: usize = 128;
30639        const NSPLIT: u32 = 4;
30640        assert!(
30641            (1..=128).contains(&c),
30642            "gdn_scan_chunked: C must be in 1..=128"
30643        );
30644        let h = n_head;
30645        #[allow(clippy::manual_div_ceil)]
30646        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30647        let nc = (t + c - 1) / c;
30648        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
30649        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
30650        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
30651        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
30652        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
30653        let gdn_mma_pre = !portable_mma_gated()
30654            && c == 32
30655            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
30656                Ok("1") => true,
30657                Ok("0") => false,
30658                _ => gdn_mma_default_on(),
30659            };
30660        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
30661            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
30662        } else {
30663            None
30664        };
30665        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
30666        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
30667        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
30668        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
30669        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
30670            && gdn_mma_pre
30671            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
30672        let nk = t * hk * D;
30673        let mut kb16_local: Option<CudaSlice<u8>> = None;
30674        if gdn_mma_pre && kb16_pre.is_none() {
30675            let mut kb = self.alloc_u8_uninit(nk * 2)?;
30676            let f = self.func("f32_to_bf16_bulk");
30677            let n2 = nk as i64;
30678            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
30679            let __s_b = self.gpu.stream();
30680            let mut b = __s_b.launch_builder(&f);
30681            b.arg(k).arg(&mut kb).arg(&n2);
30682            unsafe {
30683                b.launch(cfg2)?;
30684            }
30685            kb16_local = Some(kb);
30686        }
30687        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
30688        if let Some(kb) = kb16_pre {
30689            assert!(kb.len() >= nk * 2, "kb16_pre too small");
30690        }
30691        let mut qb16: Option<CudaSlice<u8>> = None;
30692        let mut pb16: Option<CudaSlice<u8>> = None;
30693        if gdn_wgmma_pre {
30694            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
30695            // the standalone bulk cvt only serves callers without the prep mirror.
30696            if qb16_pre.is_none() {
30697                let mut qb = self.alloc_u8_uninit(nk * 2)?;
30698                let f = self.func("f32_to_bf16_bulk");
30699                let n2 = nk as i64;
30700                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
30701                let __s_b = self.gpu.stream();
30702                let mut b = __s_b.launch_builder(&f);
30703                b.arg(q).arg(&mut qb).arg(&n2);
30704                unsafe {
30705                    b.launch(cfg2)?;
30706                }
30707                qb16 = Some(qb);
30708            } else if let Some(qb) = qb16_pre {
30709                assert!(qb.len() >= nk * 2, "qb16_pre too small");
30710            }
30711            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
30712        }
30713        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
30714        let k2w = if gdn_wgmma_pre {
30715            Some((
30716                *qb16_ref0.as_ref().unwrap(),
30717                *kb16_ref0.as_ref().unwrap(),
30718                pb16.as_mut().unwrap(),
30719            ))
30720        } else {
30721            None
30722        };
30723        let (gcum, p, u, w) =
30724            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
30725        let _ = &w;
30726        let mut y = self.uninit(nc * h * c * D)?;
30727        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
30728        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
30729        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
30730        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
30731        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
30732        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
30733        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
30734        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
30735        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
30736        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
30737        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
30738        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
30739        // sites must agree or the pre-work arms while the scan takes the scalar route.
30740        let gdn_mma = !portable_mma_gated()
30741            && c == 32
30742            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
30743                Ok("1") => true,
30744                Ok("0") => false,
30745                _ => gdn_mma_default_on(),
30746            };
30747        if gdn_mma {
30748            let wb16 = wb16_pre
30749                .take()
30750                .expect("mma path pre-allocates wb16 (K3 store fold)");
30751            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
30752            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
30753            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
30754            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
30755            // pass runs inside the persistent-M kernel; Y and Ssnap are never
30756            // materialized. New numeric class (gk folds into k^T instead of ys) —
30757            // explicit opt-in until the state-carry battery promotes it. Env read per
30758            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
30759            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
30760            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
30761            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
30762            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
30763            if gdn_wgmma_pre {
30764                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
30765                let qb16 = qb16_ref0.unwrap();
30766                let pb16 = pb16.as_ref().unwrap();
30767                {
30768                    let f = self.func("gdn_k45_wgmma");
30769                    let cfg = LaunchConfig {
30770                        grid_dim: (h as u32, 4, 1),
30771                        block_dim: (256, 1, 1),
30772                        shared_mem_bytes: 0,
30773                    };
30774                    let hki = hk as i32;
30775                    let __s_b = self.gpu.stream();
30776                    let mut b = __s_b.launch_builder(&f);
30777                    b.arg(kb16_ref)
30778                        .arg(&gcum)
30779                        .arg(beta)
30780                        .arg(&u)
30781                        .arg(&wb16)
30782                        .arg(qb16)
30783                        .arg(pb16)
30784                        .arg(o)
30785                        .arg(&scale)
30786                        .arg(state_in)
30787                        .arg(&mut *state_out)
30788                        .arg(&hi)
30789                        .arg(&ti)
30790                        .arg(&ci)
30791                        .arg(&hki);
30792                    unsafe {
30793                        b.launch(cfg)?;
30794                    }
30795                }
30796                return Ok(());
30797            }
30798            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
30799            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
30800            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
30801            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
30802            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
30803            {
30804                let f = self.func("gdn_chunk_state_mma");
30805                let cfg = LaunchConfig {
30806                    grid_dim: (h as u32, NSPLIT, 1),
30807                    block_dim: (256, 1, 1),
30808                    shared_mem_bytes: 0,
30809                };
30810                let hki = hk as i32;
30811                let __s_b = self.gpu.stream();
30812                let mut b = __s_b.launch_builder(&f);
30813                b.arg(kb16_ref)
30814                    .arg(&gcum)
30815                    .arg(beta)
30816                    .arg(&u)
30817                    .arg(&wb16)
30818                    .arg(&mut y16)
30819                    .arg(&mut ssnap16)
30820                    .arg(state_in)
30821                    .arg(&mut *state_out)
30822                    .arg(&hi)
30823                    .arg(&ti)
30824                    .arg(&ci)
30825                    .arg(&hki);
30826                unsafe {
30827                    b.launch(cfg)?;
30828                }
30829            }
30830            {
30831                // K5-mma (bf16 St/Y consumers)
30832                let f = self.func("gdn_chunk_output_mma");
30833                #[allow(clippy::manual_div_ceil)]
30834                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30835                let jt = ((c + 31) / 32) as u32;
30836                let cfg = LaunchConfig {
30837                    grid_dim: (nc as u32, h as u32, jt),
30838                    block_dim: (256, 1, 1),
30839                    shared_mem_bytes: 0,
30840                };
30841                let hki = hk as i32;
30842                let __s_b = self.gpu.stream();
30843                let mut b = __s_b.launch_builder(&f);
30844                b.arg(q)
30845                    .arg(&gcum)
30846                    .arg(&p)
30847                    .arg(&y16)
30848                    .arg(&ssnap16)
30849                    .arg(o)
30850                    .arg(&hi)
30851                    .arg(&ti)
30852                    .arg(&ci)
30853                    .arg(&scale)
30854                    .arg(&hki);
30855                unsafe {
30856                    b.launch(cfg)?;
30857                }
30858            }
30859            return Ok(());
30860        }
30861        {
30862            // K4 (sequential over chunks inside; blocks col-partition the state)
30863            let f = self.func("gdn_chunk_state_f32");
30864            let cfg = LaunchConfig {
30865                grid_dim: (h as u32, NSPLIT, 1),
30866                block_dim: (256, 1, 1),
30867                shared_mem_bytes: 0,
30868            };
30869            let __s_b = self.gpu.stream();
30870            let mut b = __s_b.launch_builder(&f);
30871            b.arg(k)
30872                .arg(&gcum)
30873                .arg(beta)
30874                .arg(&u)
30875                .arg(&w)
30876                .arg(&mut y)
30877                .arg(&mut ssnap)
30878                .arg(state_in)
30879                .arg(&mut *state_out)
30880                .arg(&hi)
30881                .arg(&ti)
30882                .arg(&ci);
30883            unsafe {
30884                b.launch(cfg)?;
30885            }
30886        }
30887        {
30888            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
30889            let f = self.func("gdn_chunk_output_f32");
30890            #[allow(clippy::manual_div_ceil)]
30891            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30892            let jt = ((c + 31) / 32) as u32;
30893            let cfg = LaunchConfig {
30894                grid_dim: (nc as u32, h as u32, jt),
30895                block_dim: (256, 1, 1),
30896                shared_mem_bytes: 0,
30897            };
30898            let __s_b = self.gpu.stream();
30899            let mut b = __s_b.launch_builder(&f);
30900            b.arg(q)
30901                .arg(&gcum)
30902                .arg(&p)
30903                .arg(&y)
30904                .arg(&ssnap)
30905                .arg(o)
30906                .arg(&hi)
30907                .arg(&ti)
30908                .arg(&ci)
30909                .arg(&scale);
30910            unsafe {
30911                b.launch(cfg)?;
30912            }
30913        }
30914        Ok(())
30915    }
30916
30917    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
30918    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
30919    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
30920    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
30921    ///
30922    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
30923    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
30924    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
30925    #[allow(clippy::too_many_arguments)]
30926    #[allow(clippy::too_many_arguments)]
30927    pub fn gdn_scan_prefill(
30928        &self,
30929        q: &CudaSlice<f32>,
30930        k: &CudaSlice<f32>,
30931        v: &CudaSlice<f32>,
30932        g: &CudaSlice<f32>,
30933        beta: &CudaSlice<f32>,
30934        kb16_pre: Option<&CudaSlice<u8>>,
30935        qb16_pre: Option<&CudaSlice<u8>>,
30936        state_in: &CudaSlice<f32>,
30937        state_out: &mut CudaSlice<f32>,
30938        o: &mut CudaSlice<f32>,
30939        n_head: usize,
30940        t: usize,
30941        scale: f32,
30942        hk: usize,
30943    ) -> Result<(), Box<dyn std::error::Error>> {
30944        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
30945            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
30946            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
30947        }
30948        if Self::gdn_chunked_enabled() && t >= 16 {
30949            self.gdn_scan_chunked(
30950                q,
30951                k,
30952                v,
30953                g,
30954                beta,
30955                kb16_pre,
30956                qb16_pre,
30957                state_in,
30958                state_out,
30959                o,
30960                n_head,
30961                t,
30962                scale,
30963                Self::gdn_chunk_size(),
30964                hk,
30965            )
30966        } else {
30967            assert!(
30968                hk == n_head,
30969                "s128 scan is broadcast-only (prep guarantees by predicate)"
30970            );
30971            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
30972        }
30973    }
30974
30975    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
30976    #[allow(clippy::too_many_arguments)]
30977    fn gdn_scan_diff(
30978        &self,
30979        q: &CudaSlice<f32>,
30980        k: &CudaSlice<f32>,
30981        v: &CudaSlice<f32>,
30982        g: &CudaSlice<f32>,
30983        beta: &CudaSlice<f32>,
30984        state_in: &CudaSlice<f32>,
30985        state_out: &mut CudaSlice<f32>,
30986        o: &mut CudaSlice<f32>,
30987        n_head: usize,
30988        t: usize,
30989        scale: f32,
30990    ) -> Result<(), Box<dyn std::error::Error>> {
30991        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
30992        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
30993        let mut o_c = self.uninit(o.len())?;
30994        let mut st_c = self.uninit(state_out.len())?;
30995        self.gdn_scan_chunked(
30996            q,
30997            k,
30998            v,
30999            g,
31000            beta,
31001            None,
31002            None,
31003            state_in,
31004            &mut st_c,
31005            &mut o_c,
31006            n_head,
31007            t,
31008            scale,
31009            Self::gdn_chunk_size(),
31010            n_head,
31011        )?;
31012        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
31013        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
31014        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
31015        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
31016            let mut max_abs = 0f32;
31017            let mut max_rel = 0f32;
31018            let mut sum_rel = 0f64;
31019            for (x, y) in a.iter().zip(b) {
31020                let ad = (x - y).abs();
31021                let rel = ad / x.abs().max(y.abs()).max(1e-3);
31022                if ad > max_abs {
31023                    max_abs = ad;
31024                }
31025                if rel > max_rel {
31026                    max_rel = rel;
31027                }
31028                sum_rel += rel as f64;
31029            }
31030            (max_abs, max_rel, sum_rel / a.len() as f64)
31031        };
31032        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
31033        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
31034        println!(
31035            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
31036                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
31037            Self::gdn_chunk_size()
31038        );
31039        Ok(())
31040    }
31041
31042    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
31043    pub fn gdn_glog(
31044        &self,
31045        alpha: &CudaSlice<f32>,
31046        dt_bias: &CudaSlice<f32>,
31047        a: &CudaSlice<f32>,
31048        g_log: &mut CudaSlice<f32>,
31049        n_head: usize,
31050        t: usize,
31051    ) -> Result<(), Box<dyn std::error::Error>> {
31052        let f = self.func("gdn_glog_f32");
31053        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
31054        let (h, ti) = (n_head as i32, t as i32);
31055        let __s_b = self.gpu.stream();
31056        let mut b = __s_b.launch_builder(&f);
31057        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
31058        unsafe {
31059            b.launch(cfg)?;
31060        }
31061        Ok(())
31062    }
31063
31064    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
31065    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
31066    pub fn sigmoid_v(
31067        &self,
31068        x: &cudarc::driver::CudaView<f32>,
31069        y: &mut CudaSlice<f32>,
31070        n: usize,
31071    ) -> Result<(), Box<dyn std::error::Error>> {
31072        let f = self.func("sigmoid_f32");
31073        let cfg = LaunchConfig::for_num_elems(n as u32);
31074        let ni = n as i32;
31075        let __s_b = self.gpu.stream();
31076        let mut b = __s_b.launch_builder(&f);
31077        b.arg(x).arg(y).arg(&ni);
31078        unsafe {
31079            b.launch(cfg)?;
31080        }
31081        Ok(())
31082    }
31083
31084    pub fn gdn_glog_v(
31085        &self,
31086        alpha: &cudarc::driver::CudaView<f32>,
31087        dt_bias: &CudaSlice<f32>,
31088        a: &CudaSlice<f32>,
31089        g_log: &mut CudaSlice<f32>,
31090        n_head: usize,
31091        t: usize,
31092    ) -> Result<(), Box<dyn std::error::Error>> {
31093        let f = self.func("gdn_glog_f32");
31094        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
31095        let (h, ti) = (n_head as i32, t as i32);
31096        let __s_b = self.gpu.stream();
31097        let mut b = __s_b.launch_builder(&f);
31098        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
31099        unsafe {
31100            b.launch(cfg)?;
31101        }
31102        Ok(())
31103    }
31104
31105    pub fn sigmoid(
31106        &self,
31107        x: &CudaSlice<f32>,
31108        y: &mut CudaSlice<f32>,
31109        n: usize,
31110    ) -> Result<(), Box<dyn std::error::Error>> {
31111        let f = self.func("sigmoid_f32");
31112        let cfg = LaunchConfig::for_num_elems(n as u32);
31113        let ni = n as i32;
31114        let __s_b = self.gpu.stream();
31115        let mut b = __s_b.launch_builder(&f);
31116        b.arg(x).arg(y).arg(&ni);
31117        unsafe {
31118            b.launch(cfg)?;
31119        }
31120        Ok(())
31121    }
31122
31123    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
31124    /// (replaces sigmoid + mul + convert). Bit-identical class.
31125    pub fn sig_mul_f16out(
31126        &self,
31127        a: &CudaSlice<f32>,
31128        g: &CudaSlice<f32>,
31129        dst: &mut CudaSlice<f32>,
31130        dst16: &mut CudaSlice<u8>,
31131        n: usize,
31132    ) -> Result<(), Box<dyn std::error::Error>> {
31133        let f = self.func("sig_mul_f16out_f32");
31134        let cfg = LaunchConfig::for_num_elems(n as u32);
31135        let ni = n as i32;
31136        let __s_b = self.gpu.stream();
31137        let mut b = __s_b.launch_builder(&f);
31138        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
31139        unsafe {
31140            b.launch(cfg)?;
31141        }
31142        Ok(())
31143    }
31144
31145    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
31146    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
31147    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
31148    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
31149    ///
31150    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
31151    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
31152    /// applies the wrong number of distinct gate values.
31153    #[allow(clippy::too_many_arguments)]
31154    pub fn attn_head_gate(
31155        &self,
31156        a: &CudaSlice<f32>,
31157        g: &CudaSlice<f32>,
31158        dst: &mut CudaSlice<f32>,
31159        dst16: Option<&mut CudaSlice<u8>>,
31160        head_dim: usize,
31161        n_head: usize,
31162        t: usize,
31163    ) -> Result<(), Box<dyn std::error::Error>> {
31164        let f = self.func("attn_head_gate_f32");
31165        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
31166        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
31167        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
31168        let d16: u64 = match dst16 {
31169            Some(d) => self.addr_u8(d),
31170            None => 0,
31171        };
31172        let __s_b = self.gpu.stream();
31173        let mut b = __s_b.launch_builder(&f);
31174        b.arg(a)
31175            .arg(g)
31176            .arg(dst)
31177            .arg(&d16)
31178            .arg(&hd)
31179            .arg(&nh)
31180            .arg(&ti);
31181        unsafe {
31182            b.launch(cfg)?;
31183        }
31184        Ok(())
31185    }
31186
31187    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
31188    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
31189    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
31190    ///
31191    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
31192    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
31193    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
31194    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
31195    #[allow(clippy::too_many_arguments)]
31196    pub fn swiglu_clamped_mul_scaled(
31197        &self,
31198        gate: &CudaSlice<f32>,
31199        up: &CudaSlice<f32>,
31200        gs: f32,
31201        us: f32,
31202        limit: f32,
31203        dst: &mut CudaSlice<f32>,
31204        n: usize,
31205    ) -> Result<(), Box<dyn std::error::Error>> {
31206        debug_assert!(
31207            limit > 1e-6,
31208            "swiglu_clamped needs a live limit; use silu_mul_scaled"
31209        );
31210        let f = self.func("swiglu_clamped_mul_scaled_f32");
31211        let cfg = LaunchConfig::for_num_elems(n as u32);
31212        let ni = n as i32;
31213        let __s_b = self.gpu.stream();
31214        let mut b = __s_b.launch_builder(&f);
31215        b.arg(gate)
31216            .arg(up)
31217            .arg(&gs)
31218            .arg(&us)
31219            .arg(&limit)
31220            .arg(dst)
31221            .arg(&ni);
31222        unsafe {
31223            b.launch(cfg)?;
31224        }
31225        Ok(())
31226    }
31227
31228    /// glm5_next PRE-clamped SwiGLU: `dst = silu(min(gate*gs, limit)) * clamp(up*us, +-limit)`.
31229    /// The gate clamp is BEFORE silu and one-sided — vendor `Glm5NextTextMLP.forward` /
31230    /// `Glm5NextTextExperts._apply_gate`, one `swiglu_limit` shared by the dense MLP, the routed
31231    /// experts and the shared expert on every layer.
31232    ///
31233    /// This is NOT `swiglu_clamped_mul_scaled` (step35 clamps the silu OUTPUT) and NOT
31234    /// `swigluoai_mul_scaled` (alpha-swish plus a `1 +` linear term). Same caller contract as the
31235    /// post-clamp sibling: `limit > 1e-6`, else the plain `silu_mul_scaled` path.
31236    #[allow(clippy::too_many_arguments)]
31237    pub fn swiglu_preclamped_mul_scaled(
31238        &self,
31239        gate: &CudaSlice<f32>,
31240        up: &CudaSlice<f32>,
31241        gs: f32,
31242        us: f32,
31243        limit: f32,
31244        dst: &mut CudaSlice<f32>,
31245        n: usize,
31246    ) -> Result<(), Box<dyn std::error::Error>> {
31247        debug_assert!(
31248            limit > 1e-6,
31249            "swiglu_preclamped needs a live limit; use silu_mul_scaled"
31250        );
31251        let f = self.func("swiglu_preclamped_mul_scaled_f32");
31252        let cfg = LaunchConfig::for_num_elems(n as u32);
31253        let ni = n as i32;
31254        let __s_b = self.gpu.stream();
31255        let mut b = __s_b.launch_builder(&f);
31256        b.arg(gate)
31257            .arg(up)
31258            .arg(&gs)
31259            .arg(&us)
31260            .arg(&limit)
31261            .arg(dst)
31262            .arg(&ni);
31263        unsafe {
31264            b.launch(cfg)?;
31265        }
31266        Ok(())
31267    }
31268
31269    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
31270    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31271    pub fn gated_rmsnorm(
31272        &self,
31273        o: &CudaSlice<f32>,
31274        w: &CudaSlice<f32>,
31275        z: &CudaSlice<f32>,
31276        dst: &mut CudaSlice<f32>,
31277        ncols: usize,
31278        nrows: usize,
31279        eps: f32,
31280    ) -> Result<(), Box<dyn std::error::Error>> {
31281        let f = self.func("gated_rmsnorm_f32");
31282        let cfg = LaunchConfig {
31283            grid_dim: (nrows as u32, 1, 1),
31284            block_dim: (128, 1, 1),
31285            shared_mem_bytes: 0,
31286        };
31287        let (nc, e) = (ncols as i32, eps);
31288        let __s_b = self.gpu.stream();
31289        let mut b = __s_b.launch_builder(&f);
31290        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
31291        unsafe {
31292            b.launch(cfg)?;
31293        }
31294        Ok(())
31295    }
31296
31297    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
31298    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
31299    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31300    pub fn gated_rmsnorm_f16out(
31301        &self,
31302        o: &CudaSlice<f32>,
31303        w: &CudaSlice<f32>,
31304        z: &CudaSlice<f32>,
31305        dst: &mut CudaSlice<f32>,
31306        dst16: &mut CudaSlice<u8>,
31307        ncols: usize,
31308        nrows: usize,
31309        eps: f32,
31310    ) -> Result<(), Box<dyn std::error::Error>> {
31311        let f = self.func("gated_rmsnorm_f16out_f32");
31312        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
31313        let cfg = LaunchConfig {
31314            grid_dim: (nrows as u32, 1, 1),
31315            block_dim: (128, 1, 1),
31316            shared_mem_bytes: 0,
31317        };
31318        let (nc, e) = (ncols as i32, eps);
31319        let __s_b = self.gpu.stream();
31320        let mut b = __s_b.launch_builder(&f);
31321        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
31322        unsafe {
31323            b.launch(cfg)?;
31324        }
31325        Ok(())
31326    }
31327
31328    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
31329    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
31330    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
31331    #[allow(clippy::too_many_arguments)]
31332    pub fn add_rms_norm_zq8(
31333        &self,
31334        a: &CudaSlice<f32>,
31335        b_in: &CudaSlice<f32>,
31336        w: &CudaSlice<f32>,
31337        res: &mut CudaSlice<f32>,
31338        z: &mut CudaSlice<f32>,
31339        ncols: usize,
31340        nrows: usize,
31341        eps: f32,
31342    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
31343        assert!(ncols.is_multiple_of(32));
31344        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
31345        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
31346        let f = self.func("add_rms_norm_zq8");
31347        let cfg = LaunchConfig {
31348            grid_dim: (nrows as u32, 1, 1),
31349            block_dim: (1024, 1, 1),
31350            shared_mem_bytes: 0,
31351        };
31352        let (nc, ep) = (ncols as i32, eps);
31353        let __s_b = self.gpu.stream();
31354        let mut b = __s_b.launch_builder(&f);
31355        b.arg(a)
31356            .arg(b_in)
31357            .arg(w)
31358            .arg(res)
31359            .arg(z)
31360            .arg(&mut q)
31361            .arg(&mut d)
31362            .arg(&nc)
31363            .arg(&ep);
31364        unsafe {
31365            b.launch(cfg)?;
31366        }
31367        Ok((q, d))
31368    }
31369
31370    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
31371    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
31372    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
31373    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
31374    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31375    pub fn gated_rmsnorm_zv(
31376        &self,
31377        o: &CudaSlice<f32>,
31378        w: &CudaSlice<f32>,
31379        z: &cudarc::driver::CudaView<f32>,
31380        dst: &mut CudaSlice<f32>,
31381        ncols: usize,
31382        nrows: usize,
31383        eps: f32,
31384    ) -> Result<(), Box<dyn std::error::Error>> {
31385        let f = self.func("gated_rmsnorm_f32");
31386        let cfg = LaunchConfig {
31387            grid_dim: (nrows as u32, 1, 1),
31388            block_dim: (128, 1, 1),
31389            shared_mem_bytes: 0,
31390        };
31391        let (nc, e) = (ncols as i32, eps);
31392        let __s_b = self.gpu.stream();
31393        let mut b = __s_b.launch_builder(&f);
31394        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
31395        unsafe {
31396            b.launch(cfg)?;
31397        }
31398        Ok(())
31399    }
31400
31401    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31402    pub fn gated_rmsnorm_f16out_zv(
31403        &self,
31404        o: &CudaSlice<f32>,
31405        w: &CudaSlice<f32>,
31406        z: &cudarc::driver::CudaView<f32>,
31407        dst: &mut CudaSlice<f32>,
31408        dst16: &mut CudaSlice<u8>,
31409        ncols: usize,
31410        nrows: usize,
31411        eps: f32,
31412    ) -> Result<(), Box<dyn std::error::Error>> {
31413        let f = self.func("gated_rmsnorm_f16out_f32");
31414        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
31415        let cfg = LaunchConfig {
31416            grid_dim: (nrows as u32, 1, 1),
31417            block_dim: (128, 1, 1),
31418            shared_mem_bytes: 0,
31419        };
31420        let (nc, e) = (ncols as i32, eps);
31421        let __s_b = self.gpu.stream();
31422        let mut b = __s_b.launch_builder(&f);
31423        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
31424        unsafe {
31425            b.launch(cfg)?;
31426        }
31427        Ok(())
31428    }
31429
31430    pub fn gated_rmsnorm_q8_1(
31431        &self,
31432        o: &CudaSlice<f32>,
31433        w: &CudaSlice<f32>,
31434        z: &CudaSlice<f32>,
31435        ncols: usize,
31436        nrows: usize,
31437        eps: f32,
31438    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
31439        assert!(ncols.is_multiple_of(32));
31440        let f = self.func("gated_rmsnorm_q8_1");
31441        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
31442        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
31443        let cfg = LaunchConfig {
31444            grid_dim: (nrows as u32, 1, 1),
31445            block_dim: (128, 1, 1),
31446            shared_mem_bytes: 0,
31447        };
31448        let (nc, ep) = (ncols as i32, eps);
31449        let __s_b = self.gpu.stream();
31450        let mut b = __s_b.launch_builder(&f);
31451        b.arg(o)
31452            .arg(w)
31453            .arg(z)
31454            .arg(&mut out_q)
31455            .arg(&mut out_d)
31456            .arg(&nc)
31457            .arg(&ep);
31458        unsafe {
31459            b.launch(cfg)?;
31460        }
31461        Ok((out_q, out_d))
31462    }
31463
31464    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
31465    pub fn transpose(
31466        &self,
31467        inp: &CudaSlice<f32>,
31468        rows: usize,
31469        cols: usize,
31470    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31471        let f = self.func("transpose_f32");
31472        let mut out = self.zeros(rows * cols)?;
31473        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
31474        let (r, c) = (rows as i32, cols as i32);
31475        let __s_b = self.gpu.stream();
31476        let mut b = __s_b.launch_builder(&f);
31477        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
31478        unsafe {
31479            b.launch(cfg)?;
31480        }
31481        Ok(out)
31482    }
31483
31484    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
31485    pub fn repeat_heads(
31486        &self,
31487        inp: &CudaSlice<f32>,
31488        out: &mut CudaSlice<f32>,
31489        head_dim: usize,
31490        n_in: usize,
31491        n_out: usize,
31492        t: usize,
31493    ) -> Result<(), Box<dyn std::error::Error>> {
31494        let f = self.func("repeat_heads_f32");
31495        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
31496        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
31497        let __s_b = self.gpu.stream();
31498        let mut b = __s_b.launch_builder(&f);
31499        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
31500        unsafe {
31501            b.launch(cfg)?;
31502        }
31503        Ok(())
31504    }
31505
31506    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
31507    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
31508    ///
31509    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
31510    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
31511    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
31512    pub fn q_gate_split(
31513        &self,
31514        qf: &CudaSlice<f32>,
31515        q_out: &mut CudaSlice<f32>,
31516        gate_out: &mut CudaSlice<f32>,
31517        head_dim: usize,
31518        n_head: usize,
31519        t: usize,
31520    ) -> Result<(), Box<dyn std::error::Error>> {
31521        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
31522        let out_need = head_dim * n_head * t;
31523        if q_out.len() < out_need || gate_out.len() < out_need {
31524            return Err(format!(
31525                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
31526                q_out.len(),
31527                gate_out.len()
31528            )
31529            .into());
31530        }
31531        let f = self.func("q_gate_split_f32");
31532        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
31533        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
31534        let __s_b = self.gpu.stream();
31535        let mut b = __s_b.launch_builder(&f);
31536        b.arg(qf)
31537            .arg(q_out)
31538            .arg(gate_out)
31539            .arg(&hd)
31540            .arg(&nh)
31541            .arg(&ti);
31542        unsafe {
31543            b.launch(cfg)?;
31544        }
31545        Ok(())
31546    }
31547
31548    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
31549    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
31550    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
31551    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31552    pub fn qkv_to_gdn_repack(
31553        &self,
31554        conv_out: &CudaSlice<f32>,
31555        q_g: &mut CudaSlice<f32>,
31556        k_g: &mut CudaSlice<f32>,
31557        v_g: &mut CudaSlice<f32>,
31558        d_state: usize,
31559        num_v: usize,
31560        num_k: usize,
31561        key_dim: usize,
31562        t: usize,
31563    ) -> Result<(), Box<dyn std::error::Error>> {
31564        let f = self.func("qkv_to_gdn_repack_f32");
31565        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
31566        let (ds, nv, nk, kd, ti) = (
31567            d_state as i32,
31568            num_v as i32,
31569            num_k as i32,
31570            key_dim as i32,
31571            t as i32,
31572        );
31573        let __s_b = self.gpu.stream();
31574        let mut b = __s_b.launch_builder(&f);
31575        b.arg(conv_out)
31576            .arg(q_g)
31577            .arg(k_g)
31578            .arg(v_g)
31579            .arg(&ds)
31580            .arg(&nv)
31581            .arg(&nk)
31582            .arg(&kd)
31583            .arg(&ti);
31584        unsafe {
31585            b.launch(cfg)?;
31586        }
31587        Ok(())
31588    }
31589
31590    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
31591    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
31592    pub fn conv_left_pad(
31593        &self,
31594        src: &CudaSlice<f32>,
31595        dst: &mut CudaSlice<f32>,
31596        conv_dim: usize,
31597        t: usize,
31598        pad: usize,
31599    ) -> Result<(), Box<dyn std::error::Error>> {
31600        let f = self.func("conv_left_pad_f32");
31601        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
31602        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
31603        let __s_b = self.gpu.stream();
31604        let mut b = __s_b.launch_builder(&f);
31605        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
31606        unsafe {
31607            b.launch(cfg)?;
31608        }
31609        Ok(())
31610    }
31611
31612    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
31613    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
31614    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
31615    pub fn conv_assemble_and_roll(
31616        &self,
31617        qkv_col: &CudaSlice<f32>,
31618        conv_state: &mut CudaSlice<f32>,
31619        conv_in: &mut CudaSlice<f32>,
31620        conv_dim: usize,
31621        pad: usize,
31622    ) -> Result<(), Box<dyn std::error::Error>> {
31623        let f = self.func("conv_assemble_and_roll_f32");
31624        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
31625        let (cd, p) = (conv_dim as i32, pad as i32);
31626        let __s_b = self.gpu.stream();
31627        let mut b = __s_b.launch_builder(&f);
31628        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
31629        unsafe {
31630            b.launch(cfg)?;
31631        }
31632        Ok(())
31633    }
31634
31635    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
31636    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
31637    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
31638    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
31639    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
31640    pub fn ssm_conv1d_fused_decode(
31641        &self,
31642        qkv_col: &CudaSlice<f32>,
31643        conv_state: &mut CudaSlice<f32>,
31644        w: &CudaSlice<f32>,
31645        conv_out: &mut CudaSlice<f32>,
31646        conv_dim: usize,
31647        d_conv: usize,
31648    ) -> Result<(), Box<dyn std::error::Error>> {
31649        let f = self.func("ssm_conv1d_fused_decode_f32");
31650        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
31651        let (cd, dc) = (conv_dim as i32, d_conv as i32);
31652        let __s_b = self.gpu.stream();
31653        let mut b = __s_b.launch_builder(&f);
31654        b.arg(qkv_col)
31655            .arg(conv_state)
31656            .arg(w)
31657            .arg(conv_out)
31658            .arg(&cd)
31659            .arg(&dc);
31660        unsafe {
31661            b.launch(cfg)?;
31662        }
31663        Ok(())
31664    }
31665
31666    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
31667    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
31668    pub fn slice_range(
31669        &self,
31670        src: &CudaSlice<f32>,
31671        start: usize,
31672        len: usize,
31673    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31674        let host = self.gpu.stream().clone_dtoh(src)?;
31675        self.gpu.stream().synchronize()?;
31676        self.htod(&host[start..start + len])
31677    }
31678}
31679
31680#[cfg(test)]
31681mod target_dispatch_tests {
31682    use super::legacy_quant_gemm_allowed;
31683
31684    #[test]
31685    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
31686        // sm_120a native lane
31687        assert!(legacy_quant_gemm_allowed(false, false, false));
31688        assert!(!legacy_quant_gemm_allowed(false, false, true));
31689        // pure portable lane (sm_89): gated
31690        assert!(!legacy_quant_gemm_allowed(true, false, false));
31691        assert!(!legacy_quant_gemm_allowed(true, false, true));
31692        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
31693        assert!(legacy_quant_gemm_allowed(true, true, false));
31694        assert!(!legacy_quant_gemm_allowed(true, true, true));
31695    }
31696
31697    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
31698    #[test]
31699    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
31700        assert!(!legacy_quant_gemm_allowed(
31701            cfg!(memra_portable_cuda),
31702            cfg!(memra_hopper_mma),
31703            false
31704        ));
31705    }
31706
31707    #[cfg(memra_hopper_mma)]
31708    #[test]
31709    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
31710        assert!(legacy_quant_gemm_allowed(
31711            cfg!(memra_portable_cuda),
31712            cfg!(memra_hopper_mma),
31713            false
31714        ));
31715        assert!(super::portable_mma_gated() == false);
31716    }
31717}
31718
31719/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
31720/// inherent methods (inherent methods win name resolution, so no recursion).
31721impl memra_kv::KvDev for Engine {
31722    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31723        Engine::zeros(self, n)
31724    }
31725    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31726        Engine::uninit(self, n)
31727    }
31728    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
31729        Engine::alloc_u8(self, n)
31730    }
31731    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
31732        Engine::htod_i32(self, v)
31733    }
31734    fn clone_dtod(
31735        &self,
31736        src: &CudaSlice<f32>,
31737    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31738        Engine::clone_dtod(self, src)
31739    }
31740    fn copy_into(
31741        &self,
31742        dst: &mut CudaSlice<f32>,
31743        off: usize,
31744        src: &CudaSlice<f32>,
31745        len: usize,
31746    ) -> Result<(), Box<dyn std::error::Error>> {
31747        Engine::copy_into(self, dst, off, src, len)
31748    }
31749    fn copy_range_into(
31750        &self,
31751        dst: &mut CudaSlice<f32>,
31752        dst_off: usize,
31753        src: &CudaSlice<f32>,
31754        src_off: usize,
31755        len: usize,
31756    ) -> Result<(), Box<dyn std::error::Error>> {
31757        Engine::copy_range_into(self, dst, dst_off, src, src_off, len)
31758    }
31759    fn set_i32_one(
31760        &self,
31761        d: &mut CudaSlice<i32>,
31762        v: i32,
31763    ) -> Result<(), Box<dyn std::error::Error>> {
31764        Engine::set_i32_one(self, d, v)
31765    }
31766}
31767
31768#[cfg(test)]
31769mod fused_gate_bounds_tests {
31770    use super::*;
31771
31772    /// The fused `[q|gate]` split's read-site guard, on the device.
31773    ///
31774    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
31775    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
31776    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
31777    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
31778    /// `FusedQGateExtent` before the launch.
31779    ///
31780    /// Catch demonstration for this test (guard temporarily removed, then restored):
31781    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
31782    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
31783    /// the call returns `Err`. Receipt in the lane report.
31784    #[test]
31785    #[ignore = "requires a CUDA GPU"]
31786    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
31787        let e = Engine::new(0).unwrap();
31788        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
31789        let fused = 2 * head_dim * n_head * t;
31790        let out_n = head_dim * n_head * t;
31791
31792        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
31793        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
31794        let mut q = e.uninit(out_n).unwrap();
31795        let mut gate = e.uninit(out_n).unwrap();
31796        let err = e
31797            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
31798            .expect_err("half-width wq must be refused, not read past")
31799            .to_string();
31800        assert!(err.contains("NO fused gate"), "{err}");
31801        assert!(err.contains(&format!("{fused}")), "{err}");
31802
31803        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
31804        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
31805        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
31806        let wide = e.htod(&host).unwrap();
31807        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
31808            .expect("full-width wq splits");
31809        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
31810        for tok in 0..t {
31811            for hh in 0..n_head {
31812                for d in 0..head_dim {
31813                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
31814                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
31815                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
31816                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
31817                }
31818            }
31819        }
31820
31821        // undersized destinations are refused too (the other half of the extent contract)
31822        let mut small = e.uninit(out_n - 1).unwrap();
31823        assert!(
31824            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
31825                .is_err()
31826        );
31827    }
31828}
31829
31830/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
31831/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
31832/// any launch, so the refusal is testable without a device.
31833#[cfg(test)]
31834mod fused_rope_width_tests {
31835    use super::Engine;
31836
31837    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
31838    /// safetensors route derives the same), which is why the fusion is legal there today.
31839    #[test]
31840    fn full_width_is_accepted() {
31841        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
31842        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
31843        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
31844    }
31845
31846    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
31847    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
31848    ///
31849    /// ```text
31850    /// attention.key_length     512   rope.dimension_count     512   (global class)
31851    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
31852    /// ```
31853    ///
31854    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
31855    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
31856    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
31857    /// instead of a silently over-rotated head.
31858    #[test]
31859    fn gemma4_official_artifact_widths_pass() {
31860        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
31861        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
31862    }
31863
31864    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
31865    /// with no `n_dims`, silently rotating the pass-through band.
31866    #[test]
31867    fn partial_rotary_is_refused_with_the_geometry_named() {
31868        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
31869        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
31870            .expect_err("partial rotary must refuse");
31871        let msg = err.to_string();
31872        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
31873        assert!(msg.contains("n_rot 64"), "{msg}");
31874        assert!(msg.contains("head_dim 256"), "{msg}");
31875        assert!(
31876            msg.contains("64..256"),
31877            "names the band it would corrupt: {msg}"
31878        );
31879        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
31880        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
31881        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
31882        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
31883    }
31884}