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    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
5854    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
5855    pub fn htod_u64_into(
5856        &self,
5857        v: &[u64],
5858        dst: &mut CudaSlice<u64>,
5859    ) -> Result<(), Box<dyn std::error::Error>> {
5860        let mut view = dst.slice_mut(0..v.len());
5861        self.gpu.stream().memcpy_htod(v, &mut view)?;
5862        Ok(())
5863    }
5864
5865    /// f32 twin of [`Self::htod_u64_into`] (the MoE vrows scale tables through the
5866    /// verify-walk workspace, door W).
5867    pub fn htod_f32_into(
5868        &self,
5869        v: &[f32],
5870        dst: &mut CudaSlice<f32>,
5871    ) -> Result<(), Box<dyn std::error::Error>> {
5872        let mut view = dst.slice_mut(0..v.len());
5873        self.gpu.stream().memcpy_htod(v, &mut view)?;
5874        Ok(())
5875    }
5876
5877    /// `htod_f32_into` landing at an element offset: `dst[off..off+v.len()] = v`. The EP
5878    /// dispatch-diet's bulk peer-row return lands the peer's compact block directly into the
5879    /// pair-slab tail with ONE upload instead of a per-row scatter.
5880    pub fn htod_f32_into_at(
5881        &self,
5882        v: &[f32],
5883        dst: &mut CudaSlice<f32>,
5884        off: usize,
5885    ) -> Result<(), Box<dyn std::error::Error>> {
5886        if off + v.len() > dst.len() {
5887            return Err(format!(
5888                "htod_f32_into_at range {}..{} exceeds dst {}",
5889                off,
5890                off + v.len(),
5891                dst.len()
5892            )
5893            .into());
5894        }
5895        let mut view = dst.slice_mut(off..off + v.len());
5896        self.gpu.stream().memcpy_htod(v, &mut view)?;
5897        Ok(())
5898    }
5899
5900    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
5901    /// device pointer-table entry at run time, so a captured graph follows the gdn
5902    /// ping-pong through the same table its scan kernels read — a baked memcpy node
5903    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
5904    pub fn copy_indirect_src_f32(
5905        &self,
5906        src_entry: &cudarc::driver::CudaView<u64>,
5907        dst: &mut CudaSlice<f32>,
5908        dst_off: usize,
5909        words: usize,
5910    ) -> Result<(), Box<dyn std::error::Error>> {
5911        let f = self.func("copy_indirect_src_f32");
5912        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
5913        let wi = words as i32;
5914        let cfg = LaunchConfig {
5915            grid_dim: (chunks, 1, 1),
5916            block_dim: (256, 1, 1),
5917            shared_mem_bytes: 0,
5918        };
5919        let mut dv = dst.slice_mut(dst_off..dst_off + words);
5920        let __s = self.gpu.stream();
5921        let mut b = __s.launch_builder(&f);
5922        b.arg(src_entry).arg(&mut dv).arg(&wi);
5923        unsafe {
5924            b.launch(cfg)?;
5925        }
5926        Ok(())
5927    }
5928
5929    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
5930    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
5931        self.alloc_uninit::<i8>(n)
5932    }
5933
5934    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
5935    #[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
5936    pub fn qmatvec(
5937        &self,
5938        w: &CudaSlice<u8>,
5939        x: &CudaSlice<f32>,
5940        m: usize,
5941        in_f: usize,
5942        out_f: usize,
5943        qtype: i32,
5944        row_bytes: usize,
5945    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5946        let f = self.func("qmatvec_f32");
5947        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
5948        let cfg = LaunchConfig {
5949            grid_dim: (out_f as u32, m as u32, 1),
5950            block_dim: (256, 1, 1),
5951            shared_mem_bytes: 0,
5952        };
5953        let (inf, outf, mi, qt, rb) =
5954            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5955        let __s_b = self.gpu.stream();
5956        let mut b = __s_b.launch_builder(&f);
5957        b.arg(w)
5958            .arg(x)
5959            .arg(&mut y)
5960            .arg(&inf)
5961            .arg(&outf)
5962            .arg(&mi)
5963            .arg(&qt)
5964            .arg(&rb);
5965        unsafe {
5966            b.launch(cfg)?;
5967        }
5968        Ok(y)
5969    }
5970
5971    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
5972    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5973        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
5974        self.keep_if_capturing(&s);
5975        Ok(s)
5976    }
5977
5978    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
5979    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
5980    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
5981    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5982        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
5983        self.keep_if_capturing(&s);
5984        Ok(s)
5985    }
5986
5987    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
5988    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
5989    pub fn memset_zeros_view(
5990        &self,
5991        dst: &mut cudarc::driver::CudaViewMut<f32>,
5992    ) -> Result<(), Box<dyn std::error::Error>> {
5993        self.gpu.stream().memset_zeros(dst)?;
5994        Ok(())
5995    }
5996
5997    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
5998    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
5999    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
6000    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
6001    /// stream would require an event).
6002    pub fn stage_expert(
6003        &self,
6004        host_bytes: &[u8],
6005        scratch: &mut CudaSlice<u8>,
6006        off: usize,
6007    ) -> Result<(), Box<dyn std::error::Error>> {
6008        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
6009        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
6010        Ok(())
6011    }
6012
6013    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
6014    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
6015    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
6016    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
6017    /// One CTA per token row, 256 threads (one per expert).
6018    pub fn moe_router_topk(
6019        &self,
6020        logits: &CudaSlice<f32>,
6021        t: usize,
6022        n_expert: usize,
6023        n_used: usize,
6024    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6025        let f = self.func("moe_router_topk_f32");
6026        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
6027        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
6028        let cfg = LaunchConfig {
6029            grid_dim: (t as u32, 1, 1),
6030            block_dim: (n_expert as u32, 1, 1),
6031            shared_mem_bytes: 0,
6032        };
6033        let (ne, nu) = (n_expert as i32, n_used as i32);
6034        let __s_b = self.gpu.stream();
6035        let mut b = __s_b.launch_builder(&f);
6036        b.arg(logits)
6037            .arg(&mut sel_idx)
6038            .arg(&mut sel_w)
6039            .arg(&ne)
6040            .arg(&nu);
6041        unsafe {
6042            b.launch(cfg)?;
6043        }
6044        Ok((sel_idx, sel_w))
6045    }
6046
6047    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
6048    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
6049    pub fn moe_router_topk_scaled(
6050        &self,
6051        logits: &CudaSlice<f32>,
6052        t: usize,
6053        n_expert: usize,
6054        n_used: usize,
6055        ex_scale: &CudaSlice<f32>,
6056    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6057        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
6058        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
6059        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
6060        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
6061        let f = self.func("moe_router_topk_scaled_f32");
6062        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
6063        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
6064        let cfg = LaunchConfig {
6065            grid_dim: (t as u32, 1, 1),
6066            block_dim: (n_expert as u32, 1, 1),
6067            shared_mem_bytes: 0,
6068        };
6069        let (ne, nu) = (n_expert as i32, n_used as i32);
6070        let __s_b = self.gpu.stream();
6071        let mut b = __s_b.launch_builder(&f);
6072        b.arg(logits)
6073            .arg(&mut sel_idx)
6074            .arg(&mut sel_w)
6075            .arg(&ne)
6076            .arg(&nu)
6077            .arg(ex_scale);
6078        unsafe {
6079            b.launch(cfg)?;
6080        }
6081        Ok((sel_idx, sel_w))
6082    }
6083
6084    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
6085    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
6086    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
6087    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
6088    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
6089    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
6090    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
6091    pub fn moe_router_topk_host(
6092        &self,
6093        logits: &CudaSlice<f32>,
6094        t: usize,
6095        n_expert: usize,
6096        n_used: usize,
6097    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6098        let f = self.func("moe_router_topk_f32");
6099        let n = t * n_used;
6100        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
6101        let mut sel_w = self.alloc_uninit::<f32>(n)?;
6102        let cfg = LaunchConfig {
6103            grid_dim: (t as u32, 1, 1),
6104            block_dim: (n_expert as u32, 1, 1),
6105            shared_mem_bytes: 0,
6106        };
6107        let (ne, nu) = (n_expert as i32, n_used as i32);
6108        let __s_b = self.gpu.stream();
6109        let mut b = __s_b.launch_builder(&f);
6110        b.arg(logits)
6111            .arg(&mut sel_idx)
6112            .arg(&mut sel_w)
6113            .arg(&ne)
6114            .arg(&nu);
6115        unsafe {
6116            b.launch(cfg)?;
6117        }
6118        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
6119        let bytes = n * 8;
6120        let mut guard = self.router_stage.lock().unwrap();
6121        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
6122            *guard = Some(PinnedStage::new(bytes.max(4096))?);
6123        }
6124        let stage = guard.as_mut().unwrap();
6125        let (si, sw) = unsafe {
6126            (
6127                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
6128                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
6129            )
6130        };
6131        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
6132        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
6133        self.gpu.stream().synchronize()?; // ONE sync for both
6134        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
6135    }
6136
6137    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
6138    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
6139    /// original expert ids before top-k. Exact key ties choose the smaller original id.
6140    #[allow(clippy::too_many_arguments)]
6141    pub fn moe_router_sigmoid_topk(
6142        &self,
6143        logits: &CudaSlice<f32>,
6144        t: usize,
6145        n_expert: usize,
6146        n_used: usize,
6147        active_count: usize,
6148        correction_bias: &CudaSlice<f32>,
6149        active: &CudaSlice<u8>,
6150        scaling_factor: f32,
6151        route_norm: bool,
6152    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6153        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6154        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
6155            return Err(format!(
6156                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
6157            )
6158            .into());
6159        }
6160        if logits.len() < t * n_expert
6161            || correction_bias.len() != n_expert
6162            || active.len() != n_expert
6163        {
6164            return Err(format!(
6165                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
6166                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
6167            ).into());
6168        }
6169        let f = self.func(crate::sigmoid_topk_kernel(
6170            crate::sig_expf_dev_on(),
6171            crate::topk_fast_on(),
6172            n_used,
6173        ));
6174        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
6175        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
6176        let threads = n_expert.div_ceil(32) * 32;
6177        let cfg = LaunchConfig {
6178            grid_dim: (t as u32, 1, 1),
6179            block_dim: (threads as u32, 1, 1),
6180            shared_mem_bytes: 0,
6181        };
6182        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
6183        let __s_b = self.gpu.stream();
6184        let mut b = __s_b.launch_builder(&f);
6185        b.arg(logits)
6186            .arg(correction_bias)
6187            .arg(active)
6188            .arg(&mut sel_idx)
6189            .arg(&mut sel_w)
6190            .arg(&ne)
6191            .arg(&nu)
6192            .arg(&scaling_factor)
6193            .arg(&rn);
6194        unsafe {
6195            b.launch(cfg)?;
6196        }
6197        Ok((sel_idx, sel_w))
6198    }
6199
6200    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
6201    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
6202    #[allow(clippy::too_many_arguments)]
6203    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
6204    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
6205    /// the model engine can wait on it with a same-device stream memop.
6206    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
6207        if ptr == 0 {
6208            return Err("ring_flag_raw: unarmed flag".into());
6209        }
6210        let f = self.func("memra_ring_flag");
6211        let cfg = LaunchConfig {
6212            grid_dim: (1, 1, 1),
6213            block_dim: (32, 1, 1),
6214            shared_mem_bytes: 0,
6215        };
6216        let __s_b = self.gpu.stream();
6217        let mut b = __s_b.launch_builder(&f);
6218        b.arg(&ptr).arg(&value);
6219        unsafe {
6220            b.launch(cfg)?;
6221        }
6222        Ok(())
6223    }
6224
6225    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
6226    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
6227    pub fn moe_sel_w_mirror(
6228        &self,
6229        sel_src: &CudaSlice<i32>,
6230        w_src: &CudaSlice<f32>,
6231        sel_dst: &mut CudaSlice<i32>,
6232        w_dst: &mut CudaSlice<f32>,
6233        n: usize,
6234    ) -> Result<(), Box<dyn std::error::Error>> {
6235        if n == 0
6236            || n > i32::MAX as usize
6237            || sel_src.len() < n
6238            || w_src.len() < n
6239            || sel_dst.len() < n
6240            || w_dst.len() < n
6241        {
6242            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
6243        }
6244        let f = self.func("moe_sel_w_mirror");
6245        let threads = if n <= 32 { 32 } else { 128 };
6246        let cfg = LaunchConfig {
6247            grid_dim: ((n as u32).div_ceil(threads), 1, 1),
6248            block_dim: (threads, 1, 1),
6249            shared_mem_bytes: 0,
6250        };
6251        let ni = n as i32;
6252        let __s_b = self.gpu.stream();
6253        let mut b = __s_b.launch_builder(&f);
6254        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
6255        unsafe {
6256            b.launch(cfg)?;
6257        }
6258        Ok(())
6259    }
6260
6261    /// One-launch W4A16 EP staging: peer-read the active f32 input plus routed ids/weights from
6262    /// the root device, round the input directly into the rank-local BF16 buffer, and mirror the
6263    /// fixed route metadata. The caller orders root production with an entry event.
6264    #[allow(clippy::too_many_arguments)]
6265    pub fn nvfp4_ep_stage_inputs(
6266        &self,
6267        input_src: &CudaSlice<f32>,
6268        sel_src: &CudaSlice<i32>,
6269        w_src: &CudaSlice<f32>,
6270        input_bf16_dst: &mut CudaSlice<u8>,
6271        sel_dst: &mut CudaSlice<i32>,
6272        w_dst: &mut CudaSlice<f32>,
6273        input_values: usize,
6274        pairs: usize,
6275        copy_weights: bool,
6276    ) -> Result<(), Box<dyn std::error::Error>> {
6277        if input_values == 0
6278            || pairs == 0
6279            || input_src.len() < input_values
6280            || sel_src.len() < pairs
6281            || w_src.len() < pairs
6282            || input_bf16_dst.len() < 2 * input_values
6283            || sel_dst.len() < pairs
6284            || w_dst.len() < pairs
6285        {
6286            return Err(format!(
6287                "W4A16 EP stage geometry input={} sel={} weights={} input_bf16={} \
6288                 sel_dst={} weights_dst={} active={input_values} pairs={pairs}",
6289                input_src.len(),
6290                sel_src.len(),
6291                w_src.len(),
6292                input_bf16_dst.len(),
6293                sel_dst.len(),
6294                w_dst.len(),
6295            )
6296            .into());
6297        }
6298        let f = self.func("nvfp4_ep_stage_inputs");
6299        let n = input_values.max(pairs);
6300        let cfg = LaunchConfig::for_num_elems(n as u32);
6301        let (input_values, pairs, copy_weights) =
6302            (input_values as i32, pairs as i32, i32::from(copy_weights));
6303        let __s_b = self.gpu.stream();
6304        let mut b = __s_b.launch_builder(&f);
6305        b.arg(input_src)
6306            .arg(sel_src)
6307            .arg(w_src)
6308            .arg(input_bf16_dst)
6309            .arg(sel_dst)
6310            .arg(w_dst)
6311            .arg(&input_values)
6312            .arg(&pairs)
6313            .arg(&copy_weights);
6314        unsafe {
6315            b.launch(cfg)?;
6316        }
6317        Ok(())
6318    }
6319
6320    /// Capture-safe twin of `nvfp4_ep_stage_inputs`: the three sources are persistent raw
6321    /// device addresses owned by the root engine. Destinations remain rank-local typed slices.
6322    #[allow(clippy::too_many_arguments)]
6323    pub fn nvfp4_ep_stage_inputs_raw(
6324        &self,
6325        input_src: u64,
6326        sel_src: u64,
6327        w_src: u64,
6328        input_bf16_dst: &mut CudaSlice<u8>,
6329        sel_dst: &mut CudaSlice<i32>,
6330        w_dst: &mut CudaSlice<f32>,
6331        input_values: usize,
6332        pairs: usize,
6333        copy_weights: bool,
6334    ) -> Result<(), Box<dyn std::error::Error>> {
6335        if input_src == 0
6336            || sel_src == 0
6337            || w_src == 0
6338            || input_values == 0
6339            || pairs == 0
6340            || input_bf16_dst.len() < 2 * input_values
6341            || sel_dst.len() < pairs
6342            || w_dst.len() < pairs
6343        {
6344            return Err(format!(
6345                "W4A16 EP raw stage geometry input={input_src:#x} sel={sel_src:#x} \
6346                 weights={w_src:#x} input_bf16={} sel_dst={} weights_dst={} \
6347                 active={input_values} pairs={pairs}",
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    #[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
6377    pub fn moe_router_sigmoid_topk_into(
6378        &self,
6379        logits: &CudaSlice<f32>,
6380        t: usize,
6381        n_expert: usize,
6382        n_used: usize,
6383        active_count: usize,
6384        correction_bias: &CudaSlice<f32>,
6385        active: &CudaSlice<u8>,
6386        scaling_factor: f32,
6387        route_norm: bool,
6388        sel_idx: &mut CudaSlice<i32>,
6389        sel_w: &mut CudaSlice<f32>,
6390    ) -> Result<(), Box<dyn std::error::Error>> {
6391        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6392        if n_expert == 0
6393            || n_expert > 1024
6394            || n_used == 0
6395            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
6396            || n_used > n_expert
6397            || logits.len() < t * n_expert
6398            || correction_bias.len() != n_expert
6399            || active.len() != n_expert
6400            || sel_idx.len() < t * n_used
6401            || sel_w.len() < t * n_used
6402        {
6403            return Err("sigmoid router _into geometry mismatch".into());
6404        }
6405        let f = self.func(crate::sigmoid_topk_kernel(
6406            crate::sig_expf_dev_on(),
6407            crate::topk_fast_on(),
6408            n_used,
6409        ));
6410        let threads = n_expert.div_ceil(32) * 32;
6411        let cfg = LaunchConfig {
6412            grid_dim: (t as u32, 1, 1),
6413            block_dim: (threads as u32, 1, 1),
6414            shared_mem_bytes: 0,
6415        };
6416        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
6417        let __s_b = self.gpu.stream();
6418        let mut b = __s_b.launch_builder(&f);
6419        b.arg(logits)
6420            .arg(correction_bias)
6421            .arg(active)
6422            .arg(&mut *sel_idx)
6423            .arg(&mut *sel_w)
6424            .arg(&ne)
6425            .arg(&nu)
6426            .arg(&scaling_factor)
6427            .arg(&rn);
6428        unsafe {
6429            b.launch(cfg)?;
6430        }
6431        Ok(())
6432    }
6433
6434    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
6435    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
6436    #[allow(clippy::too_many_arguments)]
6437    pub fn moe_router_sigmoid_topk_host(
6438        &self,
6439        logits: &CudaSlice<f32>,
6440        t: usize,
6441        n_expert: usize,
6442        n_used: usize,
6443        active_count: usize,
6444        correction_bias: &CudaSlice<f32>,
6445        active: &CudaSlice<u8>,
6446        scaling_factor: f32,
6447        route_norm: bool,
6448    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6449        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
6450            logits,
6451            t,
6452            n_expert,
6453            n_used,
6454            active_count,
6455            correction_bias,
6456            active,
6457            scaling_factor,
6458            route_norm,
6459        )?;
6460        let n = t * n_used;
6461        let bytes = n * 8;
6462        let mut guard = self.router_stage.lock().unwrap();
6463        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
6464            *guard = Some(PinnedStage::new(bytes.max(4096))?);
6465        }
6466        let stage = guard.as_mut().unwrap();
6467        let (si, sw) = unsafe {
6468            (
6469                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
6470                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
6471            )
6472        };
6473        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
6474        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
6475        self.gpu.stream().synchronize()?;
6476        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
6477    }
6478
6479    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
6480    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
6481    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
6482    pub fn stage_expert_async(
6483        &self,
6484        host_bytes: &[u8],
6485        scratch: &mut CudaSlice<u8>,
6486        off: usize,
6487    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
6488        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
6489        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
6490        Ok(self.copy_stream.record_event(None)?)
6491    }
6492
6493    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
6494    pub fn compute_wait(
6495        &self,
6496        ev: &cudarc::driver::CudaEvent,
6497    ) -> Result<(), Box<dyn std::error::Error>> {
6498        self.gpu.stream().wait(ev)?;
6499        Ok(())
6500    }
6501
6502    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
6503    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
6504    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
6505    /// CudaView base+offset pointer is honored by the launch arg.
6506    #[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
6507    pub fn qmatvec_view(
6508        &self,
6509        w: &CudaSlice<u8>,
6510        range: std::ops::Range<usize>,
6511        x: &cudarc::driver::CudaView<f32>,
6512        m: usize,
6513        in_f: usize,
6514        out_f: usize,
6515        qtype: i32,
6516        row_bytes: usize,
6517    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6518        self.qmatvec_view_inner(w, range, x, m, in_f, out_f, qtype, row_bytes)
6519    }
6520
6521    /// W4A16 expert matvec: round the floating activation to checkpoint BF16 before the
6522    /// existing f32-dequant weight dot. The output remains f32. This is selected per model by
6523    /// `MoeWeights`; it is not a process-global NVFP4 policy.
6524    #[allow(clippy::too_many_arguments)]
6525    pub fn qmatvec_view_bf16_activation(
6526        &self,
6527        w: &CudaSlice<u8>,
6528        range: std::ops::Range<usize>,
6529        x: &cudarc::driver::CudaView<f32>,
6530        m: usize,
6531        in_f: usize,
6532        out_f: usize,
6533        qtype: i32,
6534        row_bytes: usize,
6535    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6536        let n = m * in_f;
6537        if x.len() != n {
6538            return Err(format!(
6539                "W4A16 BF16 activation input length {} != {m}x{in_f}",
6540                x.len()
6541            )
6542            .into());
6543        }
6544        let mut x_bf16 = self.alloc_u8_uninit(n * 2)?;
6545        self.f32_to_bf16_v(x, &mut x_bf16, n)?;
6546        let x_f32 = self.bf16_to_f32(&x_bf16.slice(0..n * 2), n)?;
6547        self.qmatvec_view_inner(
6548            w,
6549            range,
6550            &x_f32.slice(0..n),
6551            m,
6552            in_f,
6553            out_f,
6554            qtype,
6555            row_bytes,
6556        )
6557    }
6558
6559    #[allow(clippy::too_many_arguments)]
6560    fn qmatvec_view_inner(
6561        &self,
6562        w: &CudaSlice<u8>,
6563        range: std::ops::Range<usize>,
6564        x: &cudarc::driver::CudaView<f32>,
6565        m: usize,
6566        in_f: usize,
6567        out_f: usize,
6568        qtype: i32,
6569        row_bytes: usize,
6570    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6571        let f = self.func("qmatvec_f32");
6572        let wv = w.slice(range); // CudaView<u8>, offset honored
6573        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6574        let cfg = LaunchConfig {
6575            grid_dim: (out_f as u32, m as u32, 1),
6576            block_dim: (256, 1, 1),
6577            shared_mem_bytes: 0,
6578        };
6579        let (inf, outf, mi, qt, rb) =
6580            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6581        let __s_b = self.gpu.stream();
6582        let mut b = __s_b.launch_builder(&f);
6583        b.arg(&wv)
6584            .arg(x)
6585            .arg(&mut y)
6586            .arg(&inf)
6587            .arg(&outf)
6588            .arg(&mi)
6589            .arg(&qt)
6590            .arg(&rb);
6591        unsafe {
6592            b.launch(cfg)?;
6593        }
6594        Ok(y)
6595    }
6596
6597    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
6598    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
6599    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
6600    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
6601    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
6602    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
6603    #[allow(clippy::too_many_arguments)]
6604    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
6605    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
6606    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
6607    pub fn moe_gate_up_silu8_q8(
6608        &self,
6609        gp: WPtr8,
6610        up: WPtr8,
6611        aq: &CudaSlice<i8>,
6612        ad: &CudaSlice<f32>,
6613        in_f: usize,
6614        n_ff: usize,
6615        n_used: usize,
6616        qt_g: i32,
6617        qt_u: i32,
6618        rb_g: usize,
6619        rb_u: usize,
6620    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6621        let f = self.func("moe_gate_up_silu8_q8");
6622        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6623        let cfg = LaunchConfig {
6624            grid_dim: (n_ff as u32, n_used as u32, 1),
6625            block_dim: (32, 1, 1),
6626            shared_mem_bytes: 0,
6627        };
6628        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
6629        let __s_b = self.gpu.stream();
6630        let mut b = __s_b.launch_builder(&f);
6631        b.arg(&gp)
6632            .arg(&up)
6633            .arg(aq)
6634            .arg(ad)
6635            .arg(&mut act)
6636            .arg(&inf)
6637            .arg(&nff)
6638            .arg(&qt_g)
6639            .arg(&qt_u)
6640            .arg(&rbg)
6641            .arg(&rbu);
6642        unsafe {
6643            b.launch(cfg)?;
6644        }
6645        Ok(act)
6646    }
6647
6648    /// The PRE-clamped, macro-folding twin of [`Engine::moe_gate_up_silu8_q8`] — the kernel
6649    /// class for any MoE family whose activation clamps the gate BEFORE the silu (glm5_next is
6650    /// the first such family; the door names the arithmetic, not the family).
6651    ///
6652    /// Same grid/block/dots/warp reduction; the epilogue is
6653    /// `silu(min(gate*gs, limit)) * clamp(up*us, ±limit)` — `swiglu_preclamped_mul_scaled_f32`'s
6654    /// expression verbatim — and `gs`/`us` carry the SELECTED experts' NVFP4 `weight_scale_2`
6655    /// macro scales in router slot order (1.0 for a macro-free bank).
6656    ///
6657    /// `limit` must be live: at `limit == 0` every gate collapses to `silu(0) == 0`, so a caller
6658    /// with no clamp belongs on the plain-SiLU sibling, not here. Same contract as
6659    /// [`Engine::swiglu_preclamped_mul_scaled`].
6660    #[allow(clippy::too_many_arguments)]
6661    pub fn moe_gate_up_preclamp8_q8(
6662        &self,
6663        gp: WPtr8,
6664        up: WPtr8,
6665        aq: &CudaSlice<i8>,
6666        ad: &CudaSlice<f32>,
6667        gs: F32x8,
6668        us: F32x8,
6669        limit: f32,
6670        in_f: usize,
6671        n_ff: usize,
6672        n_used: usize,
6673        qt_g: i32,
6674        qt_u: i32,
6675        rb_g: usize,
6676        rb_u: usize,
6677    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6678        debug_assert!(
6679            limit > 1e-6,
6680            "moe_gate_up_preclamp8_q8 needs a live limit; use moe_gate_up_silu8_q8"
6681        );
6682        let f = self.func("moe_gate_up_preclamp8_q8");
6683        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6684        let cfg = LaunchConfig {
6685            grid_dim: (n_ff as u32, n_used as u32, 1),
6686            block_dim: (32, 1, 1),
6687            shared_mem_bytes: 0,
6688        };
6689        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
6690        let __s_b = self.gpu.stream();
6691        let mut b = __s_b.launch_builder(&f);
6692        b.arg(&gp)
6693            .arg(&up)
6694            .arg(aq)
6695            .arg(ad)
6696            .arg(&gs)
6697            .arg(&us)
6698            .arg(&limit)
6699            .arg(&mut act)
6700            .arg(&inf)
6701            .arg(&nff)
6702            .arg(&qt_g)
6703            .arg(&qt_u)
6704            .arg(&rbg)
6705            .arg(&rbu);
6706        unsafe {
6707            b.launch(cfg)?;
6708        }
6709        Ok(act)
6710    }
6711
6712    #[allow(clippy::too_many_arguments)]
6713    pub fn moe_down8_fma_q8(
6714        &self,
6715        dp: WPtr8,
6716        w: F32x8,
6717        aq2: &CudaSlice<i8>,
6718        ad2: &CudaSlice<f32>,
6719        dst: &mut cudarc::driver::CudaViewMut<f32>,
6720        in_f: usize,
6721        out_f: usize,
6722        n_used: usize,
6723        qt: i32,
6724        rb: usize,
6725    ) -> Result<(), Box<dyn std::error::Error>> {
6726        let f = self.func("moe_down8_fma_q8");
6727        let cfg = LaunchConfig {
6728            grid_dim: (out_f as u32, 1, 1),
6729            block_dim: (32, 1, 1),
6730            shared_mem_bytes: 0,
6731        };
6732        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
6733        let __s_b = self.gpu.stream();
6734        let mut b = __s_b.launch_builder(&f);
6735        b.arg(&dp)
6736            .arg(&w)
6737            .arg(aq2)
6738            .arg(ad2)
6739            .arg(dst)
6740            .arg(&inf)
6741            .arg(&outf)
6742            .arg(&nu)
6743            .arg(&qt)
6744            .arg(&rbi);
6745        unsafe {
6746            b.launch(cfg)?;
6747        }
6748        Ok(())
6749    }
6750
6751    /// DEVICE-SIDE build of the verify-rows pair's pointer/scale tables (door D,
6752    /// `MEMRA_MOE_VROWS_DEV_TABLES`) from the router's own device selection. Replaces the host
6753    /// loop plus its two pageable HtoD, and lets the caller skip the router's pinned readback
6754    /// and its full `cuStreamSynchronize` entirely. Arithmetic is term-for-term the host loop's
6755    /// (see the kernel comment in `qmatvec.cu`), so the tables — and therefore every downstream
6756    /// byte — are identical.
6757    ///
6758    /// `macros` is the model's immutable `(gate, up, down)` `weight_scale_2` host planes, or
6759    /// `None` for a non-macro bank (the kernel then takes 1.0f, `macro_scale`'s own answer).
6760    /// The planes get a resident device mirror keyed by `(il, plane)` on first use — uploading
6761    /// them per call would ADD three HtoD to a door whose purpose is removing two.
6762    #[allow(clippy::too_many_arguments)]
6763    // allow: the parameter list mirrors the kernel/FFI/call contract
6764    pub fn moe_vrows_tables_from_sel(
6765        &self,
6766        sel: &CudaSlice<i32>,
6767        selw: &CudaSlice<f32>,
6768        il: u16,
6769        macros: Option<(&[f32], &[f32], &[f32])>,
6770        (pg, pu, pd): (u64, u64, u64),
6771        (sg, su, sd): (usize, usize, usize),
6772        n_pairs: usize,
6773        ptrs: &mut CudaSlice<u64>,
6774        scl: &mut CudaSlice<f32>,
6775    ) -> Result<(), Box<dyn std::error::Error>> {
6776        debug_assert!(sel.len() >= n_pairs && selw.len() >= n_pairs);
6777        // `>=` not `==`: door E appends a fourth (expert-major order) plane to the same table.
6778        debug_assert!(ptrs.len() >= 3 * n_pairs);
6779        debug_assert_eq!(scl.len(), 3 * n_pairs);
6780        // Resident macro mirrors, uploaded once per (layer, plane). The guard is held across the
6781        // launch because `CudaSlice` is not clonable — the same shape as the w8-mirror sites.
6782        let mut mac = self
6783            .vrows_macro_dev
6784            .lock()
6785            .map_err(|_| "vrows macro mirror map is poisoned")?;
6786        if let Some((hg, hu, hd)) = macros {
6787            for (plane, host) in [(0u8, hg), (1u8, hu), (2u8, hd)] {
6788                // `entry` rather than contains_key+insert: the upload is fallible, so it lands in
6789                // the Vacant arm instead of an `or_insert_with` closure.
6790                if let std::collections::hash_map::Entry::Vacant(slot) = mac.entry((il, plane)) {
6791                    slot.insert(self.htod(host)?);
6792                }
6793            }
6794        }
6795        // Absent macro planes: the three kernel pointers must still be legal device addresses,
6796        // so the call aliases the selection weights and never dereferences them (have_macros=0).
6797        let (mg, mu, md, have) = match macros {
6798            Some(_) => (
6799                mac.get(&(il, 0)).expect("gate macro mirror built above"),
6800                mac.get(&(il, 1)).expect("up macro mirror built above"),
6801                mac.get(&(il, 2)).expect("down macro mirror built above"),
6802                1i32,
6803            ),
6804            None => (selw, selw, selw, 0i32),
6805        };
6806        let f = self.func("moe_vrows_tables_from_sel");
6807        let threads = 128u32;
6808        let cfg = LaunchConfig {
6809            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
6810            block_dim: (threads, 1, 1),
6811            shared_mem_bytes: 0,
6812        };
6813        let (sgi, sui, sdi) = (sg as i64, su as i64, sd as i64);
6814        let (npi, havei) = (n_pairs as i32, have);
6815        let __s_b = self.gpu.stream();
6816        let mut b = __s_b.launch_builder(&f);
6817        b.arg(sel)
6818            .arg(selw)
6819            .arg(mg)
6820            .arg(mu)
6821            .arg(md)
6822            .arg(&mut *ptrs)
6823            .arg(&mut *scl)
6824            .arg(&pg)
6825            .arg(&pu)
6826            .arg(&pd)
6827            .arg(&sgi)
6828            .arg(&sui)
6829            .arg(&sdi)
6830            .arg(&npi)
6831            .arg(&havei);
6832        unsafe {
6833            b.launch(cfg)?;
6834        }
6835        Ok(())
6836    }
6837
6838    /// DEVICE-SIDE build of the verify-rows pair's EXPERT-MAJOR order plane (door E,
6839    /// `MEMRA_MOE_VROWS_DEDUP_ORDER`) from the router's own device selection, written into the
6840    /// pointer table's fourth plane `ptrs[3*n_pairs ..)`. Bit-identical to
6841    /// [`crate::vrows_expert_major_order`]: both are a stable order on `(expert id, pair index)`,
6842    /// the kernel by counting rank (see its comment in `qmatvec.cu`), the host by a stable sort.
6843    ///
6844    /// This launch exists ONLY in the door-D (device tables) arm — the host arm appends the plane
6845    /// to the vector it already uploads, so it costs zero extra transfers there. Cost in the
6846    /// device arm: 42 launches/round = ~0.093 ms at the box's 2.216 us eager-launch constant,
6847    /// against a predicted -2.17 ms/round; folding it into `moe_vrows_tables_from_sel` (same
6848    /// inputs, same one-thread-per-pair grid) is the named follow-up that recovers it.
6849    pub fn moe_vrows_order_from_sel(
6850        &self,
6851        sel: &CudaSlice<i32>,
6852        n_pairs: usize,
6853        ptrs: &mut CudaSlice<u64>,
6854    ) -> Result<(), Box<dyn std::error::Error>> {
6855        debug_assert!(sel.len() >= n_pairs);
6856        debug_assert!(
6857            ptrs.len() >= 4 * n_pairs,
6858            "the order plane lives at ptrs[3*n_pairs .. 4*n_pairs)"
6859        );
6860        let f = self.func("moe_vrows_order_from_sel");
6861        let threads = 128u32;
6862        let cfg = LaunchConfig {
6863            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
6864            block_dim: (threads, 1, 1),
6865            shared_mem_bytes: 0,
6866        };
6867        let np = n_pairs as i32;
6868        let __s_b = self.gpu.stream();
6869        let mut b = __s_b.launch_builder(&f);
6870        b.arg(sel).arg(&mut *ptrs).arg(&np);
6871        unsafe {
6872            b.launch(cfg)?;
6873        }
6874        Ok(())
6875    }
6876
6877    /// Verify-rows twin of [`Self::moe_gate_up_preclamp8_q8`] (lane/glm5-vrest): one launch
6878    /// covers ALL `n_pairs = t * n_used` routed pairs of a spec-verify batch. `ptrs` /
6879    /// `scl` are the `[3 * n_pairs]` plane-major (gate | up | down) expert-pointer and
6880    /// scale tables (gs | us | w*macro_down); per pair the kernel body is the t=1 fused
6881    /// epilogue's verbatim, bit-gated per row vs the sequential chain.
6882    #[allow(clippy::too_many_arguments)]
6883    // allow: the parameter list mirrors the kernel/FFI/call contract
6884    pub fn moe_gate_up_preclamp8_q8_rows(
6885        &self,
6886        ptrs: &CudaSlice<u64>,
6887        scl: &CudaSlice<f32>,
6888        aq: &CudaSlice<i8>,
6889        ad: &CudaSlice<f32>,
6890        limit: f32,
6891        in_f: usize,
6892        n_ff: usize,
6893        n_used: usize,
6894        n_pairs: usize,
6895        qt_g: i32,
6896        qt_u: i32,
6897        rb_g: usize,
6898        rb_u: usize,
6899    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6900        debug_assert!(
6901            limit > 1e-6,
6902            "moe_gate_up_preclamp8_q8_rows needs a live limit; the kernel collapses every gate \
6903             to silu(0) at limit 0"
6904        );
6905        debug_assert!(ptrs.len() >= 3 * n_pairs);
6906        debug_assert_eq!(scl.len(), 3 * n_pairs);
6907        // MEMRA_MOE_VROWS_DEDUP_ORDER (lane/glm5-dedup door E, default OFF): the `_ord` twin —
6908        // pair index the FASTEST grid dimension, walked in expert-major order from the table's
6909        // fourth plane, so two verify rows sharing an expert read the identical gate/up rows in
6910        // adjacent blocks. `ptrs.len() >= 4*n_pairs` is a REQUIREMENT not a hint: the door engages
6911        // only when the caller actually built the order plane, so a direct launcher call with the
6912        // shipped 3-plane table (every standing gate) keeps the shipped program. Door M wins the
6913        // tie by being tested first — the two are refused together rather than crossed.
6914        let packed = moe_vrows_pack_on();
6915        let ordered =
6916            !packed && moe_vrows_dedup_order_on() && ptrs.len() >= 4 * n_pairs && n_ff <= 65535;
6917        let (f, cfg) = if packed {
6918            if MOE_VROWS_PACK_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
6919                eprintln!(
6920                    "[moe-vrows-pack] engaged: 4-warp blocks on the verify-rows MoE pair \
6921                     (MEMRA_MOE_VROWS_PACK=1)"
6922                );
6923            }
6924            (
6925                self.func("moe_gate_up_preclamp8_q8_rows_w4"),
6926                LaunchConfig {
6927                    grid_dim: ((n_ff as u32).div_ceil(4), n_pairs as u32, 1),
6928                    block_dim: (32, 4, 1),
6929                    shared_mem_bytes: 0,
6930                },
6931            )
6932        } else if ordered {
6933            if MOE_VROWS_DEDUP_ORDER_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
6934                == 0
6935            {
6936                eprintln!(
6937                    "[moe-vrows-dedup-order] engaged: verify-rows gate/up walks the pair union \
6938                     EXPERT-MAJOR with the pair index as the fastest grid dimension, so the \
6939                     21.96%-measured repeat visits read a shared expert slab's rows in adjacent \
6940                     blocks (MEMRA_MOE_VROWS_DEDUP_ORDER=1)"
6941                );
6942            }
6943            (
6944                self.func("moe_gate_up_preclamp8_q8_rows_ord"),
6945                LaunchConfig {
6946                    grid_dim: (n_pairs as u32, n_ff as u32, 1),
6947                    block_dim: (32, 1, 1),
6948                    shared_mem_bytes: 0,
6949                },
6950            )
6951        } else {
6952            (
6953                self.func("moe_gate_up_preclamp8_q8_rows"),
6954                LaunchConfig {
6955                    grid_dim: (n_ff as u32, n_pairs as u32, 1),
6956                    block_dim: (32, 1, 1),
6957                    shared_mem_bytes: 0,
6958                },
6959            )
6960        };
6961        // Door W: the vrows launcher is verify-walk-only; act is a pooled draw.
6962        let mut act = self.vws_uninit(n_pairs * n_ff)?;
6963        let (inf, nff, nu, np) = (in_f as i32, n_ff as i32, n_used as i32, n_pairs as i32);
6964        let (rbg, rbu) = (rb_g as i64, rb_u as i64);
6965        let __s_b = self.gpu.stream();
6966        let mut b = __s_b.launch_builder(&f);
6967        b.arg(ptrs)
6968            .arg(scl)
6969            .arg(aq)
6970            .arg(ad)
6971            .arg(&limit)
6972            .arg(&mut act)
6973            .arg(&inf)
6974            .arg(&nff)
6975            .arg(&nu)
6976            .arg(&np)
6977            .arg(&qt_g)
6978            .arg(&qt_u)
6979            .arg(&rbg)
6980            .arg(&rbu);
6981        unsafe {
6982            b.launch(cfg)?;
6983        }
6984        Ok(act)
6985    }
6986
6987    /// Verify-rows twin of [`Self::moe_down8_fma_q8`] (lane/glm5-vrest): every verify row's
6988    /// slot-ordered down+FMA chain in one launch. `dst` is `[t, out_f]`, fully overwritten;
6989    /// `ptrs`/`scl` are the same tables the gate/up rows launch consumed (down plane).
6990    #[allow(clippy::too_many_arguments)]
6991    // allow: the parameter list mirrors the kernel/FFI/call contract
6992    pub fn moe_down8_fma_q8_rows(
6993        &self,
6994        ptrs: &CudaSlice<u64>,
6995        scl: &CudaSlice<f32>,
6996        aq2: &CudaSlice<i8>,
6997        ad2: &CudaSlice<f32>,
6998        dst: &mut CudaSlice<f32>,
6999        in_f: usize,
7000        out_f: usize,
7001        n_used: usize,
7002        n_pairs: usize,
7003        qt: i32,
7004        rb: usize,
7005    ) -> Result<(), Box<dyn std::error::Error>> {
7006        debug_assert!(ptrs.len() >= 3 * n_pairs);
7007        debug_assert_eq!(scl.len(), 3 * n_pairs);
7008        debug_assert_eq!(n_pairs % n_used, 0, "pairs are dense slot-major");
7009        let t = n_pairs / n_used;
7010        debug_assert!(dst.len() >= t * out_f);
7011        // MEMRA_MOE_VROWS_PACK (door M): the _w4 twin, same packing as the gate/up launch.
7012        let packed = moe_vrows_pack_on();
7013        // MEMRA_MOE_VROWS_DOWN_TMAJ (door E-down): grid transposed to (t, out_f) — token fastest —
7014        // so the t verify rows at one output row are adjacent blocks and a repeated expert's down
7015        // row is read once for every token that shares it. The slot-ordered __fmaf_rn chain is
7016        // inside the block and keeps its ORIGINAL slot order; only the grid moves. Needs no table
7017        // plane (the down chain cannot be permuted), so it composes with either table provenance.
7018        let tmaj = !packed && moe_vrows_down_tmaj_on() && out_f <= 65535;
7019        let (f, cfg) = if packed {
7020            (
7021                self.func("moe_down8_fma_q8_rows_w4"),
7022                LaunchConfig {
7023                    grid_dim: ((out_f as u32).div_ceil(4), t as u32, 1),
7024                    block_dim: (32, 4, 1),
7025                    shared_mem_bytes: 0,
7026                },
7027            )
7028        } else if tmaj {
7029            if MOE_VROWS_DOWN_TMAJ_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
7030                == 0
7031            {
7032                eprintln!(
7033                    "[moe-vrows-down-tmaj] engaged: verify-rows down/FMA grid transposed to \
7034                     (t, out_f) so the verify rows at one output row are adjacent blocks; the \
7035                     slot-ordered FMA chain is unchanged (MEMRA_MOE_VROWS_DOWN_TMAJ=1)"
7036                );
7037            }
7038            (
7039                self.func("moe_down8_fma_q8_rows_tmaj"),
7040                LaunchConfig {
7041                    grid_dim: (t as u32, out_f as u32, 1),
7042                    block_dim: (32, 1, 1),
7043                    shared_mem_bytes: 0,
7044                },
7045            )
7046        } else {
7047            (
7048                self.func("moe_down8_fma_q8_rows"),
7049                LaunchConfig {
7050                    grid_dim: (out_f as u32, t as u32, 1),
7051                    block_dim: (32, 1, 1),
7052                    shared_mem_bytes: 0,
7053                },
7054            )
7055        };
7056        let (inf, outf, nu, np, rbi) = (
7057            in_f as i32,
7058            out_f as i32,
7059            n_used as i32,
7060            n_pairs as i32,
7061            rb as i64,
7062        );
7063        let __s_b = self.gpu.stream();
7064        let mut b = __s_b.launch_builder(&f);
7065        b.arg(ptrs)
7066            .arg(scl)
7067            .arg(aq2)
7068            .arg(ad2)
7069            .arg(dst)
7070            .arg(&inf)
7071            .arg(&outf)
7072            .arg(&nu)
7073            .arg(&np)
7074            .arg(&qt)
7075            .arg(&rbi);
7076        unsafe {
7077            b.launch(cfg)?;
7078        }
7079        Ok(())
7080    }
7081
7082    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
7083    #[allow(clippy::too_many_arguments)]
7084    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7085    #[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
7086    pub fn qmatvec_expert_q8(
7087        &self,
7088        w: &CudaSlice<u8>,
7089        range: std::ops::Range<usize>,
7090        aq: &CudaSlice<i8>,
7091        ad: &CudaSlice<f32>,
7092        m: usize,
7093        in_f: usize,
7094        out_f: usize,
7095        qtype: i32,
7096        row_bytes: usize,
7097    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7098        let f = self.func("qmatvec_expert_q8");
7099        let wv = w.slice(range);
7100        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7101        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
7102        let cfg = LaunchConfig {
7103            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
7104            block_dim: (32, ROWS, 1),
7105            shared_mem_bytes: 0,
7106        };
7107        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7108        let __s_b = self.gpu.stream();
7109        let mut b = __s_b.launch_builder(&f);
7110        b.arg(&wv)
7111            .arg(aq)
7112            .arg(ad)
7113            .arg(&mut y)
7114            .arg(&inf)
7115            .arg(&outf)
7116            .arg(&mi)
7117            .arg(&qtype)
7118            .arg(&rbi);
7119        unsafe {
7120            b.launch(cfg)?;
7121        }
7122        Ok(y)
7123    }
7124
7125    #[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
7126    pub fn moe_gate_up_silu8(
7127        &self,
7128        gp: WPtr8,
7129        up: WPtr8,
7130        x: &cudarc::driver::CudaView<f32>,
7131        in_f: usize,
7132        n_ff: usize,
7133        n_used: usize,
7134        qt_g: i32,
7135        qt_u: i32,
7136        rb_g: usize,
7137        rb_u: usize,
7138    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7139        let f = self.func("moe_gate_up_silu8_f32");
7140        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
7141        let cfg = LaunchConfig {
7142            grid_dim: (n_ff as u32, n_used as u32, 1),
7143            block_dim: (256, 1, 1),
7144            shared_mem_bytes: 0,
7145        };
7146        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7147        let __s_b = self.gpu.stream();
7148        let mut b = __s_b.launch_builder(&f);
7149        b.arg(&gp)
7150            .arg(&up)
7151            .arg(x)
7152            .arg(&mut act)
7153            .arg(&inf)
7154            .arg(&nff)
7155            .arg(&qt_g)
7156            .arg(&qt_u)
7157            .arg(&rbg)
7158            .arg(&rbu);
7159        unsafe {
7160            b.launch(cfg)?;
7161        }
7162        Ok(act)
7163    }
7164
7165    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
7166    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
7167    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
7168    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
7169    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
7170    #[allow(clippy::too_many_arguments)]
7171    pub fn moe_down8_fma_into(
7172        &self,
7173        dp: WPtr8,
7174        w: F32x8,
7175        act: &CudaSlice<f32>,
7176        dst: &mut cudarc::driver::CudaViewMut<f32>,
7177        in_f: usize,
7178        out_f: usize,
7179        n_used: usize,
7180        qt: i32,
7181        rb: usize,
7182    ) -> Result<(), Box<dyn std::error::Error>> {
7183        let f = self.func("moe_down8_fma_f32");
7184        let cfg = LaunchConfig {
7185            grid_dim: (out_f as u32, 1, 1),
7186            block_dim: (256, 1, 1),
7187            shared_mem_bytes: 0,
7188        };
7189        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7190        let __s_b = self.gpu.stream();
7191        let mut b = __s_b.launch_builder(&f);
7192        b.arg(&dp)
7193            .arg(&w)
7194            .arg(act)
7195            .arg(dst)
7196            .arg(&inf)
7197            .arg(&outf)
7198            .arg(&nu)
7199            .arg(&qt)
7200            .arg(&rbv);
7201        unsafe {
7202            b.launch(cfg)?;
7203        }
7204        Ok(())
7205    }
7206
7207    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
7208    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
7209    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
7210    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
7211    #[allow(clippy::too_many_arguments)]
7212    /// dp4a q8 twin of the _dev pair (resident-experts arc).
7213    ///
7214    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
7215    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
7216    /// down's FMA chain stays slot-ordered serial). Seams:
7217    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
7218    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
7219    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
7220    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
7221    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
7222    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
7223    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
7224    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
7225    ///                       only) | w8h2 (h2 x slot-parallel)
7226    #[allow(clippy::too_many_arguments)]
7227    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
7228    #[allow(clippy::too_many_arguments)]
7229    #[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
7230    pub fn moe_pairs_matvec_q8(
7231        &self,
7232        table: &CudaSlice<u64>,
7233        proj: i32,
7234        pair_tok: &CudaSlice<i32>,
7235        pair_ex: &CudaSlice<i32>,
7236        aq: &CudaSlice<i8>,
7237        ad: &CudaSlice<f32>,
7238        in_f: usize,
7239        out_f: usize,
7240        n_expert: usize,
7241        n_pairs: usize,
7242        qtype: i32,
7243        row_bytes: usize,
7244    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7245        let f = self.func("moe_pairs_matvec_q8");
7246        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7247        const ROWS: u32 = 4;
7248        let cfg = LaunchConfig {
7249            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
7250            block_dim: (32, ROWS, 1),
7251            shared_mem_bytes: 0,
7252        };
7253        let (inf, outf, ne, np, rbi) = (
7254            in_f as i32,
7255            out_f as i32,
7256            n_expert as i32,
7257            n_pairs as i32,
7258            row_bytes as i64,
7259        );
7260        let __s_b = self.gpu.stream();
7261        let mut b = __s_b.launch_builder(&f);
7262        b.arg(table)
7263            .arg(&proj)
7264            .arg(pair_tok)
7265            .arg(pair_ex)
7266            .arg(aq)
7267            .arg(ad)
7268            .arg(&mut y)
7269            .arg(&inf)
7270            .arg(&outf)
7271            .arg(&ne)
7272            .arg(&np)
7273            .arg(&qtype)
7274            .arg(&rbi);
7275        unsafe {
7276            b.launch(cfg)?;
7277        }
7278        Ok(y)
7279    }
7280
7281    /// Expert-major pair matvec (weight-reuse across each expert's token group).
7282    #[allow(clippy::too_many_arguments)]
7283    #[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
7284    pub fn moe_pairs_matvec_q8_em(
7285        &self,
7286        table: &CudaSlice<u64>,
7287        proj: i32,
7288        ex_ids: &CudaSlice<i32>,
7289        ex_off: &CudaSlice<i32>,
7290        ex_pairs: &CudaSlice<i32>,
7291        pair_tok: &CudaSlice<i32>,
7292        aq: &CudaSlice<i8>,
7293        ad: &CudaSlice<f32>,
7294        in_f: usize,
7295        out_f: usize,
7296        n_expert: usize,
7297        n_active: usize,
7298        n_pairs: usize,
7299        qtype: i32,
7300        row_bytes: usize,
7301    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7302        let f = self.func("moe_pairs_matvec_q8_em");
7303        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7304        const ROWS: u32 = 4;
7305        let cfg = LaunchConfig {
7306            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
7307            block_dim: (32, ROWS, 1),
7308            shared_mem_bytes: 0,
7309        };
7310        let (inf, outf, ne, na, rbi) = (
7311            in_f as i32,
7312            out_f as i32,
7313            n_expert as i32,
7314            n_active as i32,
7315            row_bytes as i64,
7316        );
7317        let __s_b = self.gpu.stream();
7318        let mut b = __s_b.launch_builder(&f);
7319        b.arg(table)
7320            .arg(&proj)
7321            .arg(ex_ids)
7322            .arg(ex_off)
7323            .arg(ex_pairs)
7324            .arg(pair_tok)
7325            .arg(aq)
7326            .arg(ad)
7327            .arg(&mut y)
7328            .arg(&inf)
7329            .arg(&outf)
7330            .arg(&ne)
7331            .arg(&na)
7332            .arg(&qtype)
7333            .arg(&rbi);
7334        unsafe {
7335            b.launch(cfg)?;
7336        }
7337        Ok(y)
7338    }
7339
7340    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
7341    // weight group once per (row,group) then dp4a's across the expert's token group.
7342    #[allow(clippy::too_many_arguments)]
7343    #[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
7344    pub fn moe_pairs_matvec_q8_dec(
7345        &self,
7346        table: &CudaSlice<u64>,
7347        proj: i32,
7348        ex_ids: &CudaSlice<i32>,
7349        ex_off: &CudaSlice<i32>,
7350        ex_pairs: &CudaSlice<i32>,
7351        pair_tok: &CudaSlice<i32>,
7352        aq: &CudaSlice<i8>,
7353        ad: &CudaSlice<f32>,
7354        in_f: usize,
7355        out_f: usize,
7356        n_expert: usize,
7357        n_active: usize,
7358        n_pairs: usize,
7359        qtype: i32,
7360        row_bytes: usize,
7361    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7362        let f = self.func("moe_pairs_matvec_q8_dec");
7363        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
7364        const ROWS: u32 = 4;
7365        let cfg = LaunchConfig {
7366            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
7367            block_dim: (32, ROWS, 1),
7368            shared_mem_bytes: 0,
7369        };
7370        let (inf, outf, ne, na, rbi) = (
7371            in_f as i32,
7372            out_f as i32,
7373            n_expert as i32,
7374            n_active as i32,
7375            row_bytes as i64,
7376        );
7377        let __s_b = self.gpu.stream();
7378        let mut b = __s_b.launch_builder(&f);
7379        b.arg(table)
7380            .arg(&proj)
7381            .arg(ex_ids)
7382            .arg(ex_off)
7383            .arg(ex_pairs)
7384            .arg(pair_tok)
7385            .arg(aq)
7386            .arg(ad)
7387            .arg(&mut y)
7388            .arg(&inf)
7389            .arg(&outf)
7390            .arg(&ne)
7391            .arg(&na)
7392            .arg(&qtype)
7393            .arg(&rbi);
7394        unsafe {
7395            b.launch(cfg)?;
7396        }
7397        Ok(y)
7398    }
7399
7400    pub fn moe_pairs_gelu_mul(
7401        &self,
7402        gate: &CudaSlice<f32>,
7403        up: &CudaSlice<f32>,
7404        n: usize,
7405    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7406        let f = self.func("moe_pairs_gelu_mul");
7407        let mut act = self.alloc_uninit::<f32>(n)?;
7408        let cfg = LaunchConfig::for_num_elems(n as u32);
7409        let nl = n as i64;
7410        let __s_b = self.gpu.stream();
7411        let mut b = __s_b.launch_builder(&f);
7412        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
7413        unsafe {
7414            b.launch(cfg)?;
7415        }
7416        Ok(act)
7417    }
7418
7419    pub fn moe_pairs_silu_mul(
7420        &self,
7421        gate: &CudaSlice<f32>,
7422        up: &CudaSlice<f32>,
7423        n: usize,
7424    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7425        let f = self.func("moe_pairs_silu_mul");
7426        let mut act = self.alloc_uninit::<f32>(n)?;
7427        let cfg = LaunchConfig::for_num_elems(n as u32);
7428        let nl = n as i64;
7429        let __s_b = self.gpu.stream();
7430        let mut b = __s_b.launch_builder(&f);
7431        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
7432        unsafe {
7433            b.launch(cfg)?;
7434        }
7435        Ok(act)
7436    }
7437
7438    #[allow(clippy::too_many_arguments)]
7439    #[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
7440    pub fn moe_pairs_scatter(
7441        &self,
7442        y_down: &CudaSlice<f32>,
7443        pair_w: &CudaSlice<f32>,
7444        tok_pair_off: &CudaSlice<i32>,
7445        tok_pair_ids: &CudaSlice<i32>,
7446        moe_out: &mut CudaSlice<f32>,
7447        t: usize,
7448        n_embd: usize,
7449    ) -> Result<(), Box<dyn std::error::Error>> {
7450        let f = self.func("moe_pairs_scatter");
7451        let cfg = LaunchConfig {
7452            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
7453            block_dim: (256, 1, 1),
7454            shared_mem_bytes: 0,
7455        };
7456        let ne = n_embd as i32;
7457        let __s_b = self.gpu.stream();
7458        let mut b = __s_b.launch_builder(&f);
7459        b.arg(y_down)
7460            .arg(pair_w)
7461            .arg(tok_pair_off)
7462            .arg(tok_pair_ids)
7463            .arg(moe_out)
7464            .arg(&ne);
7465        unsafe {
7466            b.launch(cfg)?;
7467        }
7468        Ok(())
7469    }
7470
7471    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
7472    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
7473    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
7474    #[allow(clippy::too_many_arguments)]
7475    pub fn moe_gate_up_gelu8_dev_q8(
7476        &self,
7477        table: &CudaSlice<u64>,
7478        sel: &cudarc::driver::CudaView<i32>,
7479        aq: &CudaSlice<i8>,
7480        ad: &CudaSlice<f32>,
7481        in_f: usize,
7482        n_ff: usize,
7483        n_used: usize,
7484        n_expert: usize,
7485        qt_g: i32,
7486        qt_u: i32,
7487        rb_g: usize,
7488        rb_u: usize,
7489    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7490        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7491        let (inf, nff, ne, rbg, rbu) = (
7492            in_f as i32,
7493            n_ff as i32,
7494            n_expert as i32,
7495            rb_g as i64,
7496            rb_u as i64,
7497        );
7498        let f = self.func("moe_gate_up_gelu8_dev_q8");
7499        let cfg = LaunchConfig {
7500            grid_dim: (n_ff as u32, n_used as u32, 1),
7501            block_dim: (32, 1, 1),
7502            shared_mem_bytes: 0,
7503        };
7504        let __s_b = self.gpu.stream();
7505        let mut b = __s_b.launch_builder(&f);
7506        b.arg(table)
7507            .arg(sel)
7508            .arg(aq)
7509            .arg(ad)
7510            .arg(&mut act)
7511            .arg(&inf)
7512            .arg(&nff)
7513            .arg(&ne)
7514            .arg(&qt_g)
7515            .arg(&qt_u)
7516            .arg(&rbg)
7517            .arg(&rbu);
7518        unsafe {
7519            b.launch(cfg)?;
7520        }
7521        Ok(act)
7522    }
7523
7524    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
7525    #[allow(clippy::too_many_arguments)]
7526    pub fn moe_gate_up_gelu8_dev_q8_rows(
7527        &self,
7528        table: &CudaSlice<u64>,
7529        sel: &CudaSlice<i32>,
7530        aq: &CudaSlice<i8>,
7531        ad: &CudaSlice<f32>,
7532        t: usize,
7533        in_f: usize,
7534        n_ff: usize,
7535        n_used: usize,
7536        n_expert: usize,
7537        qt_g: i32,
7538        qt_u: i32,
7539        rb_g: usize,
7540        rb_u: usize,
7541    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7542        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
7543        let (inf, nff, ne, rbg, rbu, nu) = (
7544            in_f as i32,
7545            n_ff as i32,
7546            n_expert as i32,
7547            rb_g as i64,
7548            rb_u as i64,
7549            n_used as i32,
7550        );
7551        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
7552        let cfg = LaunchConfig {
7553            grid_dim: (n_ff as u32, n_used as u32, t as u32),
7554            block_dim: (32, 1, 1),
7555            shared_mem_bytes: 0,
7556        };
7557        let __s_b = self.gpu.stream();
7558        let mut b = __s_b.launch_builder(&f);
7559        b.arg(table)
7560            .arg(sel)
7561            .arg(aq)
7562            .arg(ad)
7563            .arg(&mut act)
7564            .arg(&inf)
7565            .arg(&nff)
7566            .arg(&ne)
7567            .arg(&qt_g)
7568            .arg(&qt_u)
7569            .arg(&rbg)
7570            .arg(&rbu)
7571            .arg(&nu);
7572        unsafe {
7573            b.launch(cfg)?;
7574        }
7575        Ok(act)
7576    }
7577
7578    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
7579    #[allow(clippy::too_many_arguments)]
7580    pub fn moe_gate_up_gelu8_dev_q8_csr(
7581        &self,
7582        table: &CudaSlice<u64>,
7583        sel: &CudaSlice<i32>,
7584        aq: &CudaSlice<i8>,
7585        ad: &CudaSlice<f32>,
7586        n_pairs: usize,
7587        in_f: usize,
7588        n_ff: usize,
7589        n_used: usize,
7590        n_expert: usize,
7591        qt_g: i32,
7592        qt_u: i32,
7593        rb_g: usize,
7594        rb_u: usize,
7595    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7596        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7597        let (inf, nff, ne, rbg, rbu, nu, npi) = (
7598            in_f as i32,
7599            n_ff as i32,
7600            n_expert as i32,
7601            rb_g as i64,
7602            rb_u as i64,
7603            n_used as i32,
7604            n_pairs as i32,
7605        );
7606        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
7607        let cfg = LaunchConfig {
7608            grid_dim: (n_ff as u32, n_pairs as u32, 1),
7609            block_dim: (32, 1, 1),
7610            shared_mem_bytes: 0,
7611        };
7612        let __s_b = self.gpu.stream();
7613        let mut b = __s_b.launch_builder(&f);
7614        b.arg(table)
7615            .arg(sel)
7616            .arg(aq)
7617            .arg(ad)
7618            .arg(&mut act)
7619            .arg(&inf)
7620            .arg(&nff)
7621            .arg(&ne)
7622            .arg(&qt_g)
7623            .arg(&qt_u)
7624            .arg(&rbg)
7625            .arg(&rbu)
7626            .arg(&nu)
7627            .arg(&npi);
7628        unsafe {
7629            b.launch(cfg)?;
7630        }
7631        Ok(act)
7632    }
7633
7634    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
7635    #[allow(clippy::too_many_arguments)]
7636    pub fn moe_down8_fma_dev_q8_rows_g(
7637        &self,
7638        table: &CudaSlice<u64>,
7639        sel: &CudaSlice<i32>,
7640        w: &CudaSlice<f32>,
7641        aq2: &CudaSlice<i8>,
7642        ad2: &CudaSlice<f32>,
7643        dst: &mut CudaSlice<f32>,
7644        t: usize,
7645        in_f: usize,
7646        out_f: usize,
7647        n_used: usize,
7648        n_expert: usize,
7649        qt: i32,
7650        rb: usize,
7651    ) -> Result<(), Box<dyn std::error::Error>> {
7652        let (inf, outf, nu, ne, rbi) = (
7653            in_f as i32,
7654            out_f as i32,
7655            n_used as i32,
7656            n_expert as i32,
7657            rb as i64,
7658        );
7659        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
7660        // eight warps, then replay the original slot-ordered FMA chain. Every
7661        // other shape retains the generic one-warp rows kernel.
7662        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
7663        let f = self.func(if step_b1_w8 {
7664            "moe_down8_fma_dev_q8_rows_w8"
7665        } else {
7666            "moe_down8_fma_dev_q8_rows_g"
7667        });
7668        let cfg = LaunchConfig {
7669            grid_dim: (out_f as u32, 1, t as u32),
7670            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
7671            shared_mem_bytes: 0,
7672        };
7673        let __s_b = self.gpu.stream();
7674        let mut b = __s_b.launch_builder(&f);
7675        b.arg(table)
7676            .arg(sel)
7677            .arg(w)
7678            .arg(aq2)
7679            .arg(ad2)
7680            .arg(dst)
7681            .arg(&inf)
7682            .arg(&outf)
7683            .arg(&nu)
7684            .arg(&ne)
7685            .arg(&qt)
7686            .arg(&rbi);
7687        unsafe {
7688            b.launch(cfg)?;
7689        }
7690        Ok(())
7691    }
7692
7693    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
7694    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
7695    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
7696    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
7697        let (out_f, in_f) = (2048usize, 2816usize);
7698        let nblk = in_f / 32;
7699        let mut seed = 0x9E3779B97F4A7C15u64;
7700        let mut rng = move || {
7701            seed = seed
7702                .wrapping_mul(6364136223846793005)
7703                .wrapping_add(1442695040888963407);
7704            (seed >> 33) as u8
7705        };
7706        let mut w = vec![0u8; out_f * nblk * 18];
7707        for b in w.iter_mut() {
7708            *b = rng();
7709        }
7710        for r in 0..out_f {
7711            for g in 0..nblk {
7712                let off = (r * nblk + g) * 18;
7713                w[off] = 0x00;
7714                w[off + 1] = 0x2C; // sane half d
7715            }
7716        }
7717        let qplane = out_f * nblk * 16;
7718        let mut wrp = vec![0u8; w.len()];
7719        for r in 0..out_f {
7720            for g in 0..nblk {
7721                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
7722                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
7723                    .copy_from_slice(&src[0..2]);
7724                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
7725            }
7726        }
7727        let w_d = self.htod_bytes(&w)?;
7728        let wrp_d = self.htod_bytes(&wrp)?;
7729        let mut aq = vec![0i8; m * in_f];
7730        for v in aq.iter_mut() {
7731            *v = rng() as i8;
7732        }
7733        let aq_d = self.htod_i8(&aq)?;
7734        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
7735        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
7736        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
7737        const RPB: u32 = 4;
7738        let cfg = LaunchConfig {
7739            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
7740            block_dim: (32, RPB, 1),
7741            shared_mem_bytes: 0,
7742        };
7743        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
7744        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
7745        let fb = self.func("qmatvec_q4_0_mmvq_b4");
7746        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
7747        {
7748            let __s_b = self.gpu.stream();
7749            let mut b = __s_b.launch_builder(&fb);
7750            b.arg(&w_d)
7751                .arg(&aq_d)
7752                .arg(&ad_d)
7753                .arg(&mut y0)
7754                .arg(&inf)
7755                .arg(&outf)
7756                .arg(&mi)
7757                .arg(&rb);
7758            unsafe {
7759                b.launch(cfg)?;
7760            }
7761            let __s_b = self.gpu.stream();
7762            let mut b = __s_b.launch_builder(&fr);
7763            b.arg(&wrp_d)
7764                .arg(&aq_d)
7765                .arg(&ad_d)
7766                .arg(&mut y1)
7767                .arg(&inf)
7768                .arg(&outf)
7769                .arg(&mi)
7770                .arg(&qp);
7771            unsafe {
7772                b.launch(cfg)?;
7773            }
7774        }
7775        self.gpu.stream().synchronize()?;
7776        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
7777        let nd = h0
7778            .iter()
7779            .zip(&h1)
7780            .filter(|(a, b)| a.to_bits() != b.to_bits())
7781            .count();
7782        if nd != 0 {
7783            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
7784        }
7785        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
7786            self.gpu.stream().synchronize()?;
7787            let t0 = std::time::Instant::now();
7788            for _ in 0..500 {
7789                if rp {
7790                    let __s_b = self.gpu.stream();
7791                    let mut b = __s_b.launch_builder(&fr);
7792                    b.arg(&wrp_d)
7793                        .arg(&aq_d)
7794                        .arg(&ad_d)
7795                        .arg(&mut y1)
7796                        .arg(&inf)
7797                        .arg(&outf)
7798                        .arg(&mi)
7799                        .arg(&qp);
7800                    unsafe {
7801                        b.launch(cfg)?;
7802                    }
7803                } else {
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                }
7818            }
7819            self.gpu.stream().synchronize()?;
7820            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
7821        };
7822        let _ = time(false)?;
7823        let _ = time(true)?; // warm
7824        Ok((time(false)?, time(true)?))
7825    }
7826
7827    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
7828    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
7829    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
7830    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
7831    pub fn build_q4_rp4(
7832        &self,
7833        t: &mut crate::model::GpuTensor,
7834    ) -> Result<(), Box<dyn std::error::Error>> {
7835        use crate::model::GpuTensor;
7836        let GpuTensor::Quant {
7837            bytes,
7838            qtype,
7839            row_bytes,
7840            ne,
7841            rp4,
7842            ..
7843        } = t
7844        else {
7845            return Ok(());
7846        };
7847        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
7848            return Ok(());
7849        }
7850        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7851        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
7852            return Ok(());
7853        }
7854        let nblk = in_f / 32;
7855        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
7856        let f = self.func("q4_0_split_rp_build");
7857        let n = (out_f * nblk) as i32;
7858        let cfg = LaunchConfig {
7859            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
7860            block_dim: (256, 1, 1),
7861            shared_mem_bytes: 0,
7862        };
7863        let (of, nb) = (out_f as i32, nblk as i32);
7864        let _ = n;
7865        let __s_b = self.gpu.stream();
7866        let mut b = __s_b.launch_builder(&f);
7867        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
7868        unsafe {
7869            b.launch(cfg)?;
7870        }
7871        *rp4 = Some(dst);
7872        Ok(())
7873    }
7874
7875    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
7876    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
7877    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
7878    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
7879    pub fn build_q8_rp4(
7880        &self,
7881        t: &mut crate::model::GpuTensor,
7882    ) -> Result<(), Box<dyn std::error::Error>> {
7883        use crate::model::GpuTensor;
7884        let GpuTensor::Quant {
7885            bytes,
7886            qtype,
7887            row_bytes,
7888            ne,
7889            rp4,
7890            ..
7891        } = t
7892        else {
7893            return Ok(());
7894        };
7895        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
7896            return Ok(());
7897        }
7898        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7899        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
7900            return Ok(());
7901        }
7902        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
7903        Ok(())
7904    }
7905
7906    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
7907    /// mirror without a GpuTensor (same kernel the loader path above uses).
7908    pub fn build_q8_rp4_raw(
7909        &self,
7910        bytes: &CudaSlice<u8>,
7911        in_f: usize,
7912        out_f: usize,
7913    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7914        assert!(in_f.is_multiple_of(32));
7915        let nblk = in_f / 32;
7916        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
7917        let f = self.func("q8_0_split_rp_build");
7918        let cfg = LaunchConfig {
7919            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
7920            block_dim: (256, 1, 1),
7921            shared_mem_bytes: 0,
7922        };
7923        let (of, nb) = (out_f as i32, nblk as i32);
7924        let __s_b = self.gpu.stream();
7925        let mut b = __s_b.launch_builder(&f);
7926        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
7927        unsafe {
7928            b.launch(cfg)?;
7929        }
7930        Ok(dst)
7931    }
7932
7933    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
7934    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
7935    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
7936    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
7937    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
7938    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
7939    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
7940    pub fn build_q4k_rp4(
7941        &self,
7942        t: &mut crate::model::GpuTensor,
7943    ) -> Result<(), Box<dyn std::error::Error>> {
7944        use crate::model::GpuTensor;
7945        let GpuTensor::Quant {
7946            bytes,
7947            qtype,
7948            row_bytes,
7949            ne,
7950            rp4,
7951            ..
7952        } = t
7953        else {
7954            return Ok(());
7955        };
7956        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
7957            return Ok(());
7958        }
7959        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7960        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
7961            return Ok(());
7962        }
7963        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
7964        Ok(())
7965    }
7966
7967    pub fn build_q6k_rp4(
7968        &self,
7969        t: &mut crate::model::GpuTensor,
7970    ) -> Result<(), Box<dyn std::error::Error>> {
7971        use crate::model::GpuTensor;
7972        let GpuTensor::Quant {
7973            bytes,
7974            qtype,
7975            row_bytes,
7976            ne,
7977            rp4,
7978            ..
7979        } = t
7980        else {
7981            return Ok(());
7982        };
7983        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
7984            return Ok(());
7985        }
7986        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
7987        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
7988            return Ok(());
7989        }
7990        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
7991        Ok(())
7992    }
7993
7994    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
7995    pub fn build_kq_rp4_raw(
7996        &self,
7997        bytes: &CudaSlice<u8>,
7998        in_f: usize,
7999        out_f: usize,
8000        qtype: i32,
8001    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8002        assert!(in_f.is_multiple_of(256));
8003        let nsbk = in_f / 256;
8004        let (sb_bytes, kname) = match qtype {
8005            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
8006            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
8007            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
8008        };
8009        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
8010        let f = self.func(kname);
8011        let cfg = LaunchConfig {
8012            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
8013            block_dim: (256, 1, 1),
8014            shared_mem_bytes: 0,
8015        };
8016        let (of, nb) = (out_f as i32, nsbk as i32);
8017        let __s_b = self.gpu.stream();
8018        let mut b = __s_b.launch_builder(&f);
8019        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
8020        unsafe {
8021            b.launch(cfg)?;
8022        }
8023        Ok(dst)
8024    }
8025
8026    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
8027    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
8028    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
8029    pub fn kqrp_enabled() -> bool {
8030        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8031        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
8032            Ok("0") => false,
8033            Ok(_) => true,
8034            Err(_) => cfg!(memra_hopper_mma),
8035        })
8036    }
8037
8038    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
8039    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
8040    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
8041    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
8042    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
8043    pub fn build_q4_rp_swap(
8044        &self,
8045        t: &mut crate::model::GpuTensor,
8046    ) -> Result<bool, Box<dyn std::error::Error>> {
8047        use crate::model::GpuTensor;
8048        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
8049        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
8050        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
8051        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
8052        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
8053        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
8054        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
8055        // this fn's OWN builder serves may ever be swapped; everything else refuses
8056        // here, regardless of walk ordering.
8057        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
8058            return Ok(false);
8059        }
8060        self.build_q4_rp4(t)?;
8061        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
8062        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
8063            return Ok(false);
8064        };
8065        match rp4.take() {
8066            Some(split) => {
8067                *bytes = split; // the GGUF-layout buffer drops here
8068                *rp = true;
8069                Ok(true)
8070            }
8071            None => Ok(false),
8072        }
8073    }
8074
8075    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
8076    pub fn q4rp_enabled() -> bool {
8077        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8078        *ON.get_or_init(|| {
8079            std::env::var("MEMRA_Q4RP")
8080                .map(|v| v != "0")
8081                .unwrap_or(true)
8082        })
8083    }
8084
8085    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
8086    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
8087    #[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
8088    pub fn copy_rows_strided(
8089        &self,
8090        src: &CudaSlice<f32>,
8091        dst: &mut CudaSlice<f32>,
8092        row_elems: usize,
8093        n_rows: usize,
8094        src_stride: usize,
8095        src_off: usize,
8096    ) -> Result<(), Box<dyn std::error::Error>> {
8097        let f = self.func("copy_rows_strided_f32");
8098        let cfg = LaunchConfig {
8099            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
8100            block_dim: (256, 1, 1),
8101            shared_mem_bytes: 0,
8102        };
8103        let (re, nr) = (row_elems as i32, n_rows as i32);
8104        let (st, off) = (src_stride as i64, src_off as i64);
8105        let __s_b = self.gpu.stream();
8106        let mut b = __s_b.launch_builder(&f);
8107        b.arg(src)
8108            .arg(&mut *dst)
8109            .arg(&re)
8110            .arg(&nr)
8111            .arg(&st)
8112            .arg(&off);
8113        unsafe {
8114            b.launch(cfg)?;
8115        }
8116        Ok(())
8117    }
8118
8119    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
8120    ///
8121    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
8122    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
8123    /// one peer copy per token.
8124    #[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
8125    pub fn place_rows_strided(
8126        &self,
8127        src: &CudaSlice<f32>,
8128        dst: &mut CudaSlice<f32>,
8129        row_elems: usize,
8130        n_rows: usize,
8131        dst_stride: usize,
8132        dst_off: usize,
8133    ) -> Result<(), Box<dyn std::error::Error>> {
8134        if row_elems == 0 || n_rows == 0 {
8135            return Err("strided row placement requires nonzero rows and row width".into());
8136        }
8137        let src_len = n_rows
8138            .checked_mul(row_elems)
8139            .ok_or("strided row placement source size overflow")?;
8140        let dst_len = n_rows
8141            .checked_sub(1)
8142            .and_then(|rows| rows.checked_mul(dst_stride))
8143            .and_then(|base| base.checked_add(dst_off))
8144            .and_then(|base| base.checked_add(row_elems))
8145            .ok_or("strided row placement destination size overflow")?;
8146        let row_end = dst_off
8147            .checked_add(row_elems)
8148            .ok_or("strided row placement row size overflow")?;
8149        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
8150            return Err(format!(
8151                "strided row placement geometry mismatch: src={} need_src={src_len} \
8152                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
8153                 dst_stride={dst_stride} dst_off={dst_off}",
8154                src.len(),
8155                dst.len(),
8156            )
8157            .into());
8158        }
8159        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
8160            return Err("strided row placement exceeds CUDA kernel geometry".into());
8161        }
8162        let f = self.func("place_rows_strided_f32");
8163        let cfg = LaunchConfig {
8164            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
8165            block_dim: (256, 1, 1),
8166            shared_mem_bytes: 0,
8167        };
8168        let (re, nr) = (row_elems as i32, n_rows as i32);
8169        let (st, off) = (dst_stride as i64, dst_off as i64);
8170        let __s_b = self.gpu.stream();
8171        let mut b = __s_b.launch_builder(&f);
8172        b.arg(src)
8173            .arg(&mut *dst)
8174            .arg(&re)
8175            .arg(&nr)
8176            .arg(&st)
8177            .arg(&off);
8178        unsafe {
8179            b.launch(cfg)?;
8180        }
8181        Ok(())
8182    }
8183
8184    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
8185    pub fn u32_set_k(
8186        &self,
8187        dst: &mut CudaSlice<u32>,
8188        v: u32,
8189        idx: usize,
8190    ) -> Result<(), Box<dyn std::error::Error>> {
8191        let f = self.func("u32_set_k");
8192        let cfg = LaunchConfig {
8193            grid_dim: (1, 1, 1),
8194            block_dim: (1, 1, 1),
8195            shared_mem_bytes: 0,
8196        };
8197        let ii = idx as i32;
8198        let __s_b = self.gpu.stream();
8199        let mut b = __s_b.launch_builder(&f);
8200        b.arg(dst).arg(&v).arg(&ii);
8201        unsafe {
8202            b.launch(cfg)?;
8203        }
8204        Ok(())
8205    }
8206
8207    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
8208    pub fn i32_add_k(
8209        &self,
8210        d: &mut CudaSlice<i32>,
8211        v: i32,
8212    ) -> Result<(), Box<dyn std::error::Error>> {
8213        let f = self.func("i32_add_k");
8214        let cfg = LaunchConfig {
8215            grid_dim: (1, 1, 1),
8216            block_dim: (32, 1, 1),
8217            shared_mem_bytes: 0,
8218        };
8219        let __s_b = self.gpu.stream();
8220        let mut b = __s_b.launch_builder(&f);
8221        b.arg(d).arg(&v);
8222        unsafe {
8223            b.launch(cfg)?;
8224        }
8225        Ok(())
8226    }
8227
8228    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
8229    pub fn i32_iota_from(
8230        &self,
8231        ctr: &CudaSlice<i32>,
8232        dst: &mut CudaSlice<i32>,
8233        n: usize,
8234    ) -> Result<(), Box<dyn std::error::Error>> {
8235        let f = self.func("i32_iota_from");
8236        let cfg = LaunchConfig::for_num_elems(n as u32);
8237        let ni = n as i32;
8238        let __s_b = self.gpu.stream();
8239        let mut b = __s_b.launch_builder(&f);
8240        b.arg(ctr).arg(dst).arg(&ni);
8241        unsafe {
8242            b.launch(cfg)?;
8243        }
8244        Ok(())
8245    }
8246
8247    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
8248    pub fn u32_map_k(
8249        &self,
8250        buf: &mut CudaSlice<u32>,
8251        map: &CudaSlice<u32>,
8252        idx: usize,
8253    ) -> Result<(), Box<dyn std::error::Error>> {
8254        let f = self.func("u32_map_k");
8255        let cfg = LaunchConfig {
8256            grid_dim: (1, 1, 1),
8257            block_dim: (1, 1, 1),
8258            shared_mem_bytes: 0,
8259        };
8260        let ii = idx as i32;
8261        let __s_b = self.gpu.stream();
8262        let mut b = __s_b.launch_builder(&f);
8263        b.arg(buf).arg(map).arg(&ii);
8264        unsafe {
8265            b.launch(cfg)?;
8266        }
8267        Ok(())
8268    }
8269
8270    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
8271    #[allow(clippy::too_many_arguments)]
8272    pub fn u32_pack2(
8273        &self,
8274        a: &CudaSlice<u32>,
8275        off_a: usize,
8276        n1: usize,
8277        b_in: &CudaSlice<u32>,
8278        n2: usize,
8279        out: &mut CudaSlice<u32>,
8280    ) -> Result<(), Box<dyn std::error::Error>> {
8281        let f = self.func("u32_pack2");
8282        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
8283        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
8284        let __s_b = self.gpu.stream();
8285        let mut b = __s_b.launch_builder(&f);
8286        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
8287        unsafe {
8288            b.launch(cfg)?;
8289        }
8290        Ok(())
8291    }
8292
8293    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
8294    pub fn moe_w_exscale(
8295        &self,
8296        w: &mut CudaSlice<f32>,
8297        sel: &CudaSlice<i32>,
8298        s: &CudaSlice<f32>,
8299        n: usize,
8300    ) -> Result<(), Box<dyn std::error::Error>> {
8301        let f = self.func("moe_w_exscale");
8302        let cfg = LaunchConfig::for_num_elems(n as u32);
8303        let ni = n as i32;
8304        let __s_b = self.gpu.stream();
8305        let mut b = __s_b.launch_builder(&f);
8306        b.arg(w).arg(sel).arg(s).arg(&ni);
8307        unsafe {
8308            b.launch(cfg)?;
8309        }
8310        Ok(())
8311    }
8312
8313    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
8314    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
8315    pub fn moe_w_scale_by_expert(
8316        &self,
8317        w: &mut CudaSlice<f32>,
8318        sel: &CudaSlice<i32>,
8319        macros: &CudaSlice<f32>,
8320        n_expert: usize,
8321        n: usize,
8322    ) -> Result<(), Box<dyn std::error::Error>> {
8323        let f = self.func("moe_w_scale_by_expert");
8324        let cfg = LaunchConfig {
8325            grid_dim: (n.div_ceil(64) as u32, 1, 1),
8326            block_dim: (64, 1, 1),
8327            shared_mem_bytes: 0,
8328        };
8329        let (ne, nn) = (n_expert as i32, n as i32);
8330        let __s_b = self.gpu.stream();
8331        let mut b = __s_b.launch_builder(&f);
8332        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
8333        unsafe {
8334            b.launch(cfg)?;
8335        }
8336        Ok(())
8337    }
8338
8339    #[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
8340    pub fn moe_gate_up_silu8_dev_q8(
8341        &self,
8342        table: &CudaSlice<u64>,
8343        sel: &cudarc::driver::CudaView<i32>,
8344        aq: &CudaSlice<i8>,
8345        ad: &CudaSlice<f32>,
8346        in_f: usize,
8347        n_ff: usize,
8348        n_used: usize,
8349        n_expert: usize,
8350        qt_g: i32,
8351        qt_u: i32,
8352        rb_g: usize,
8353        rb_u: usize,
8354        macros: &CudaSlice<f32>,
8355    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8356        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
8357        let (mode, wpb) = GU.get_or_init(|| {
8358            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
8359            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
8360                .ok()
8361                .and_then(|v| v.parse().ok())
8362                .unwrap_or(4u32)
8363                .clamp(1, 16);
8364            (mode, wpb)
8365        });
8366        let (mode, wpb) = (mode.as_str(), *wpb);
8367        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8368        let (inf, nff, ne, rbg, rbu) = (
8369            in_f as i32,
8370            n_ff as i32,
8371            n_expert as i32,
8372            rb_g as i64,
8373            rb_u as i64,
8374        );
8375        let (f, cfg) = match mode {
8376            "1" | "2" | "4" => {
8377                let rpw: u32 = mode.parse().unwrap();
8378                let f = self.func(match rpw {
8379                    1 => "moe_gate_up_silu8_dev_q8_r1",
8380                    2 => "moe_gate_up_silu8_dev_q8_r2",
8381                    _ => "moe_gate_up_silu8_dev_q8_r4",
8382                });
8383                let rows_per_block = (rpw * wpb) as usize;
8384                let gx = n_ff.div_ceil(rows_per_block) as u32;
8385                (
8386                    f,
8387                    LaunchConfig {
8388                        grid_dim: (gx, n_used as u32, 1),
8389                        block_dim: (32, wpb, 1),
8390                        shared_mem_bytes: 0,
8391                    },
8392                )
8393            }
8394            "j8" if n_used <= 32 => (
8395                self.func("moe_gate_up_silu8_dev_q8_j8"),
8396                LaunchConfig {
8397                    grid_dim: (n_ff as u32, 1, 1),
8398                    block_dim: (32, n_used as u32, 1),
8399                    shared_mem_bytes: 0,
8400                },
8401            ),
8402            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
8403            "vsm2" => {
8404                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
8405                let sh = (rb_g + rb_u) as u32;
8406                use cudarc::driver::sys::CUfunction_attribute_enum as A;
8407                f.set_attribute(
8408                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
8409                    sh as i32,
8410                )?;
8411                (
8412                    f,
8413                    LaunchConfig {
8414                        grid_dim: (n_ff as u32, n_used as u32, 1),
8415                        block_dim: (32, 1, 1),
8416                        shared_mem_bytes: sh,
8417                    },
8418                )
8419            }
8420            "vsm" => {
8421                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
8422                let sh = (rb_g + rb_u) as u32;
8423                use cudarc::driver::sys::CUfunction_attribute_enum as A;
8424                f.set_attribute(
8425                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
8426                    sh as i32,
8427                )?;
8428                (
8429                    f,
8430                    LaunchConfig {
8431                        grid_dim: (n_ff as u32, n_used as u32, 1),
8432                        block_dim: (32, 1, 1),
8433                        shared_mem_bytes: sh,
8434                    },
8435                )
8436            }
8437            "sg" => (
8438                self.func("moe_gate_up_silu8_dev_q8_sg"),
8439                LaunchConfig {
8440                    grid_dim: (n_ff as u32, n_used as u32, 1),
8441                    block_dim: (32, 1, 1),
8442                    shared_mem_bytes: 0,
8443                },
8444            ),
8445            "j8sg" if n_used <= 32 => (
8446                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
8447                LaunchConfig {
8448                    grid_dim: (n_ff as u32, 1, 1),
8449                    block_dim: (32, n_used as u32, 1),
8450                    shared_mem_bytes: 0,
8451                },
8452            ),
8453            "u64" if in_f == 2048 => (
8454                self.func("moe_gate_up_silu8_dev_q8_u64"),
8455                LaunchConfig {
8456                    grid_dim: (n_ff as u32, n_used as u32, 1),
8457                    block_dim: (32, 1, 1),
8458                    shared_mem_bytes: 0,
8459                },
8460            ),
8461            "gs4" if in_f == 2048 => (
8462                self.func("moe_gate_up_silu8_dev_q8_gs4"),
8463                LaunchConfig {
8464                    grid_dim: (n_ff as u32, n_used as u32, 1),
8465                    block_dim: (32, 4, 1),
8466                    shared_mem_bytes: 0,
8467                },
8468            ),
8469            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
8470            "v" | "" => (
8471                self.func("moe_gate_up_silu8_dev_q8_v"),
8472                LaunchConfig {
8473                    grid_dim: (n_ff as u32, n_used as u32, 1),
8474                    block_dim: (32, 1, 1),
8475                    shared_mem_bytes: 0,
8476                },
8477            ),
8478            "s2" => (
8479                self.func("moe_gate_up_silu8_dev_q8_s2"),
8480                LaunchConfig {
8481                    grid_dim: (n_ff as u32, n_used as u32, 1),
8482                    block_dim: (32, 2, 1),
8483                    shared_mem_bytes: 0,
8484                },
8485            ),
8486            "s2z" => {
8487                let rz = wpb.min(16); // s2z smem tile is [16][2]
8488                (
8489                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
8490                    LaunchConfig {
8491                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
8492                        block_dim: (32, 2, rz),
8493                        shared_mem_bytes: 0,
8494                    },
8495                )
8496            }
8497            _ => (
8498                self.func("moe_gate_up_silu8_dev_q8"),
8499                LaunchConfig {
8500                    grid_dim: (n_ff as u32, n_used as u32, 1),
8501                    block_dim: (32, 1, 1),
8502                    shared_mem_bytes: 0,
8503                },
8504            ),
8505        };
8506        let __s_b = self.gpu.stream();
8507        let mut b = __s_b.launch_builder(&f);
8508        b.arg(table)
8509            .arg(sel)
8510            .arg(aq)
8511            .arg(ad)
8512            .arg(&mut act)
8513            .arg(&inf)
8514            .arg(&nff)
8515            .arg(&ne)
8516            .arg(&qt_g)
8517            .arg(&qt_u)
8518            .arg(&rbg)
8519            .arg(&rbu)
8520            .arg(macros);
8521        unsafe {
8522            b.launch(cfg)?;
8523        }
8524        Ok(act)
8525    }
8526
8527    #[allow(clippy::too_many_arguments)]
8528    pub fn moe_down8_fma_dev_q8(
8529        &self,
8530        table: &CudaSlice<u64>,
8531        sel: &cudarc::driver::CudaView<i32>,
8532        w: &cudarc::driver::CudaView<f32>,
8533        aq2: &CudaSlice<i8>,
8534        ad2: &CudaSlice<f32>,
8535        dst: &mut cudarc::driver::CudaViewMut<f32>,
8536        in_f: usize,
8537        out_f: usize,
8538        n_used: usize,
8539        n_expert: usize,
8540        qt: i32,
8541        rb: usize,
8542    ) -> Result<(), Box<dyn std::error::Error>> {
8543        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
8544        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
8545        let (inf, outf, nu, ne, rbi) = (
8546            in_f as i32,
8547            out_f as i32,
8548            n_used as i32,
8549            n_expert as i32,
8550            rb as i64,
8551        );
8552        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
8553        // the h2 twins are nsb==16 (in_f==512) shape-gated.
8554        let (f, cfg) = match mode.as_str() {
8555            m @ ("1" | "2" | "4") if n_used <= 8 => {
8556                let rpw: usize = m.parse().unwrap();
8557                let f = self.func(match rpw {
8558                    1 => "moe_down8_fma_dev_q8_w8r1",
8559                    2 => "moe_down8_fma_dev_q8_w8r2",
8560                    _ => "moe_down8_fma_dev_q8_w8r4",
8561                });
8562                (
8563                    f,
8564                    LaunchConfig {
8565                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
8566                        block_dim: (32, n_used as u32, 1),
8567                        shared_mem_bytes: 0,
8568                    },
8569                )
8570            }
8571            "h2" if in_f == 512 => (
8572                self.func("moe_down8_fma_dev_q8_h2"),
8573                LaunchConfig {
8574                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8575                    block_dim: (32, 1, 1),
8576                    shared_mem_bytes: 0,
8577                },
8578            ),
8579            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
8580            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
8581            "" if in_f == 704 && n_used <= 8 => (
8582                self.func("moe_down8_fma_dev_q8_w8r2"),
8583                LaunchConfig {
8584                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8585                    block_dim: (32, n_used as u32, 1),
8586                    shared_mem_bytes: 0,
8587                },
8588            ),
8589            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
8590            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
8591            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
8592            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
8593                self.func("moe_down8_fma_dev_q8_w8h2v"),
8594                LaunchConfig {
8595                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8596                    block_dim: (32, n_used as u32, 1),
8597                    shared_mem_bytes: 0,
8598                },
8599            ),
8600            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
8601                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
8602                LaunchConfig {
8603                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8604                    block_dim: (32, n_used as u32, 1),
8605                    shared_mem_bytes: 0,
8606                },
8607            ),
8608            "w8h2r2" if in_f == 512 && n_used <= 8 => (
8609                self.func("moe_down8_fma_dev_q8_w8h2r2"),
8610                LaunchConfig {
8611                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8612                    block_dim: (32, n_used as u32, 1),
8613                    shared_mem_bytes: 0,
8614                },
8615            ),
8616            "w8h2" if in_f == 512 && n_used <= 8 => (
8617                self.func("moe_down8_fma_dev_q8_w8h2"),
8618                LaunchConfig {
8619                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8620                    block_dim: (32, n_used as u32, 1),
8621                    shared_mem_bytes: 0,
8622                },
8623            ),
8624            _ => (
8625                self.func("moe_down8_fma_dev_q8"),
8626                LaunchConfig {
8627                    grid_dim: (out_f as u32, 1, 1),
8628                    block_dim: (32, 1, 1),
8629                    shared_mem_bytes: 0,
8630                },
8631            ),
8632        };
8633        let __s_b = self.gpu.stream();
8634        let mut b = __s_b.launch_builder(&f);
8635        b.arg(table)
8636            .arg(sel)
8637            .arg(w)
8638            .arg(aq2)
8639            .arg(ad2)
8640            .arg(dst)
8641            .arg(&inf)
8642            .arg(&outf)
8643            .arg(&nu)
8644            .arg(&ne)
8645            .arg(&qt)
8646            .arg(&rbi);
8647        unsafe {
8648            b.launch(cfg)?;
8649        }
8650        Ok(())
8651    }
8652
8653    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
8654    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
8655    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
8656    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
8657    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
8658    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
8659    #[allow(clippy::too_many_arguments)]
8660    pub fn moe_gate_up_silu8_dev_q8_rows(
8661        &self,
8662        table: &CudaSlice<u64>,
8663        sel: &CudaSlice<i32>,
8664        aq: &CudaSlice<i8>,
8665        ad: &CudaSlice<f32>,
8666        t: usize,
8667        in_f: usize,
8668        n_ff: usize,
8669        n_used: usize,
8670        n_expert: usize,
8671        qt_g: i32,
8672        qt_u: i32,
8673        rb_g: usize,
8674        rb_u: usize,
8675        macros: &CudaSlice<f32>,
8676    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8677        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
8678        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
8679        let cfg = LaunchConfig {
8680            grid_dim: (n_ff as u32, n_used as u32, t as u32),
8681            block_dim: (32, 1, 1),
8682            shared_mem_bytes: 0,
8683        };
8684        let (inf, nff, ne, nu, rbg, rbu) = (
8685            in_f as i32,
8686            n_ff as i32,
8687            n_expert as i32,
8688            n_used as i32,
8689            rb_g as i64,
8690            rb_u as i64,
8691        );
8692        let __s_b = self.gpu.stream();
8693        let mut b = __s_b.launch_builder(&f);
8694        b.arg(table)
8695            .arg(sel)
8696            .arg(aq)
8697            .arg(ad)
8698            .arg(&mut act)
8699            .arg(&inf)
8700            .arg(&nff)
8701            .arg(&ne)
8702            .arg(&qt_g)
8703            .arg(&qt_u)
8704            .arg(&rbg)
8705            .arg(&rbu)
8706            .arg(&nu)
8707            .arg(macros);
8708        unsafe {
8709            b.launch(cfg)?;
8710        }
8711        Ok(act)
8712    }
8713
8714    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
8715    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
8716    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
8717    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
8718    #[allow(clippy::too_many_arguments)]
8719    pub fn moe_down8_fma_dev_q8_rows(
8720        &self,
8721        table: &CudaSlice<u64>,
8722        sel: &CudaSlice<i32>,
8723        w: &CudaSlice<f32>,
8724        aq2: &CudaSlice<i8>,
8725        ad2: &CudaSlice<f32>,
8726        dst: &mut CudaSlice<f32>,
8727        t: usize,
8728        in_f: usize,
8729        out_f: usize,
8730        n_used: usize,
8731        n_expert: usize,
8732        qt: i32,
8733        rb: usize,
8734    ) -> Result<(), Box<dyn std::error::Error>> {
8735        assert!(
8736            in_f == 512 && n_used <= 8,
8737            "down rows twin is w8h2v shape-gated"
8738        );
8739        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
8740        let cfg = LaunchConfig {
8741            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
8742            block_dim: (32, n_used as u32, 1),
8743            shared_mem_bytes: 0,
8744        };
8745        let (inf, outf, nu, ne, rbi) = (
8746            in_f as i32,
8747            out_f as i32,
8748            n_used as i32,
8749            n_expert as i32,
8750            rb as i64,
8751        );
8752        let __s_b = self.gpu.stream();
8753        let mut b = __s_b.launch_builder(&f);
8754        b.arg(table)
8755            .arg(sel)
8756            .arg(w)
8757            .arg(aq2)
8758            .arg(ad2)
8759            .arg(dst)
8760            .arg(&inf)
8761            .arg(&outf)
8762            .arg(&nu)
8763            .arg(&ne)
8764            .arg(&qt)
8765            .arg(&rbi);
8766        unsafe {
8767            b.launch(cfg)?;
8768        }
8769        Ok(())
8770    }
8771
8772    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
8773    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
8774    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
8775    #[allow(clippy::too_many_arguments)]
8776    pub fn moe_gate_up_silu8_dev_q8_csr(
8777        &self,
8778        table: &CudaSlice<u64>,
8779        sel: &CudaSlice<i32>,
8780        aq: &CudaSlice<i8>,
8781        ad: &CudaSlice<f32>,
8782        n_pairs: usize,
8783        in_f: usize,
8784        n_ff: usize,
8785        n_used: usize,
8786        n_expert: usize,
8787        qt_g: i32,
8788        qt_u: i32,
8789        rb_g: usize,
8790        rb_u: usize,
8791    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8792        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
8793        // host gate guarantees qt_g == qt_u within a supported class.
8794        let f = if qt_g == crate::QT_NVFP4 {
8795            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
8796        } else {
8797            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
8798        };
8799        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
8800        let cfg = LaunchConfig {
8801            grid_dim: (n_ff as u32, n_pairs as u32, 1),
8802            block_dim: (32, 1, 1),
8803            shared_mem_bytes: 0,
8804        };
8805        let (inf, nff, ne, nu, npi, rbg, rbu) = (
8806            in_f as i32,
8807            n_ff as i32,
8808            n_expert as i32,
8809            n_used as i32,
8810            n_pairs as i32,
8811            rb_g as i64,
8812            rb_u as i64,
8813        );
8814        let __s_b = self.gpu.stream();
8815        let mut b = __s_b.launch_builder(&f);
8816        b.arg(table)
8817            .arg(sel)
8818            .arg(aq)
8819            .arg(ad)
8820            .arg(&mut act)
8821            .arg(&inf)
8822            .arg(&nff)
8823            .arg(&ne)
8824            .arg(&qt_g)
8825            .arg(&qt_u)
8826            .arg(&rbg)
8827            .arg(&rbu)
8828            .arg(&nu)
8829            .arg(&npi);
8830        unsafe {
8831            b.launch(cfg)?;
8832        }
8833        Ok(act)
8834    }
8835
8836    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
8837    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
8838    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
8839    #[allow(clippy::too_many_arguments)]
8840    pub fn moe_down8_fma_dev_q8_variant(
8841        &self,
8842        variant: &str,
8843        table: &CudaSlice<u64>,
8844        sel: &cudarc::driver::CudaView<i32>,
8845        w: &cudarc::driver::CudaView<f32>,
8846        aq2: &CudaSlice<i8>,
8847        ad2: &CudaSlice<f32>,
8848        dst: &mut cudarc::driver::CudaViewMut<f32>,
8849        in_f: usize,
8850        out_f: usize,
8851        n_used: usize,
8852        n_expert: usize,
8853        qt: i32,
8854        rb: usize,
8855    ) -> Result<(), Box<dyn std::error::Error>> {
8856        let (inf, outf, nu, ne, rbi) = (
8857            in_f as i32,
8858            out_f as i32,
8859            n_used as i32,
8860            n_expert as i32,
8861            rb as i64,
8862        );
8863        let (f, cfg) = match variant {
8864            "w8h2" | "w8h2v" => (
8865                self.func(if variant == "w8h2" {
8866                    "moe_down8_fma_dev_q8_w8h2"
8867                } else {
8868                    "moe_down8_fma_dev_q8_w8h2v"
8869                }),
8870                LaunchConfig {
8871                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
8872                    block_dim: (32, n_used as u32, 1),
8873                    shared_mem_bytes: 0,
8874                },
8875            ),
8876            "w8h2r2" | "w8h2r2v" => (
8877                self.func(if variant == "w8h2r2" {
8878                    "moe_down8_fma_dev_q8_w8h2r2"
8879                } else {
8880                    "moe_down8_fma_dev_q8_w8h2r2v"
8881                }),
8882                LaunchConfig {
8883                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
8884                    block_dim: (32, n_used as u32, 1),
8885                    shared_mem_bytes: 0,
8886                },
8887            ),
8888            _ => (
8889                self.func("moe_down8_fma_dev_q8"),
8890                LaunchConfig {
8891                    grid_dim: (out_f as u32, 1, 1),
8892                    block_dim: (32, 1, 1),
8893                    shared_mem_bytes: 0,
8894                },
8895            ),
8896        };
8897        let __s_b = self.gpu.stream();
8898        let mut b = __s_b.launch_builder(&f);
8899        b.arg(table)
8900            .arg(sel)
8901            .arg(w)
8902            .arg(aq2)
8903            .arg(ad2)
8904            .arg(dst)
8905            .arg(&inf)
8906            .arg(&outf)
8907            .arg(&nu)
8908            .arg(&ne)
8909            .arg(&qt)
8910            .arg(&rbi);
8911        unsafe {
8912            b.launch(cfg)?;
8913        }
8914        Ok(())
8915    }
8916
8917    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
8918    #[allow(clippy::too_many_arguments)]
8919    pub fn moe_gate_up_silu8_dev_q8_variant(
8920        &self,
8921        variant: &str,
8922        table: &CudaSlice<u64>,
8923        sel: &cudarc::driver::CudaView<i32>,
8924        aq: &CudaSlice<i8>,
8925        ad: &CudaSlice<f32>,
8926        in_f: usize,
8927        n_ff: usize,
8928        n_used: usize,
8929        n_expert: usize,
8930        qt_g: i32,
8931        qt_u: i32,
8932        rb_g: usize,
8933        rb_u: usize,
8934    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8935        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8936        let (inf, nff, ne, rbg, rbu) = (
8937            in_f as i32,
8938            n_ff as i32,
8939            n_expert as i32,
8940            rb_g as i64,
8941            rb_u as i64,
8942        );
8943        let f = self.func(if variant == "v" {
8944            "moe_gate_up_silu8_dev_q8_v"
8945        } else {
8946            "moe_gate_up_silu8_dev_q8"
8947        });
8948        let cfg = LaunchConfig {
8949            grid_dim: (n_ff as u32, n_used as u32, 1),
8950            block_dim: (32, 1, 1),
8951            shared_mem_bytes: 0,
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(aq)
8958            .arg(ad)
8959            .arg(&mut act)
8960            .arg(&inf)
8961            .arg(&nff)
8962            .arg(&ne)
8963            .arg(&qt_g)
8964            .arg(&qt_u)
8965            .arg(&rbg)
8966            .arg(&rbu);
8967        unsafe {
8968            b.launch(cfg)?;
8969        }
8970        Ok(act)
8971    }
8972
8973    #[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
8974    pub fn moe_gate_up_silu8_dev(
8975        &self,
8976        table: &CudaSlice<u64>,
8977        sel: &cudarc::driver::CudaView<i32>,
8978        x: &cudarc::driver::CudaView<f32>,
8979        in_f: usize,
8980        n_ff: usize,
8981        n_used: usize,
8982        n_expert: usize,
8983        qt_g: i32,
8984        qt_u: i32,
8985        rb_g: usize,
8986        rb_u: usize,
8987        macros: &CudaSlice<f32>,
8988    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8989        let f = self.func("moe_gate_up_silu8_dev");
8990        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
8991        let cfg = LaunchConfig {
8992            grid_dim: (n_ff as u32, n_used as u32, 1),
8993            block_dim: (256, 1, 1),
8994            shared_mem_bytes: 0,
8995        };
8996        let (inf, nff, ne, rbg, rbu) = (
8997            in_f as i32,
8998            n_ff as i32,
8999            n_expert as i32,
9000            rb_g as i64,
9001            rb_u as i64,
9002        );
9003        let __s_b = self.gpu.stream();
9004        let mut b = __s_b.launch_builder(&f);
9005        b.arg(table)
9006            .arg(sel)
9007            .arg(x)
9008            .arg(&mut act)
9009            .arg(&inf)
9010            .arg(&nff)
9011            .arg(&ne)
9012            .arg(&qt_g)
9013            .arg(&qt_u)
9014            .arg(&rbg)
9015            .arg(&rbu)
9016            .arg(macros);
9017        unsafe {
9018            b.launch(cfg)?;
9019        }
9020        Ok(act)
9021    }
9022
9023    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
9024    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
9025    #[allow(clippy::too_many_arguments)]
9026    pub fn moe_down8_fma_dev(
9027        &self,
9028        table: &CudaSlice<u64>,
9029        sel: &cudarc::driver::CudaView<i32>,
9030        w: &cudarc::driver::CudaView<f32>,
9031        act: &CudaSlice<f32>,
9032        dst: &mut cudarc::driver::CudaViewMut<f32>,
9033        in_f: usize,
9034        out_f: usize,
9035        n_used: usize,
9036        n_expert: usize,
9037        qt: i32,
9038        rb: usize,
9039    ) -> Result<(), Box<dyn std::error::Error>> {
9040        let f = self.func("moe_down8_fma_dev");
9041        let cfg = LaunchConfig {
9042            grid_dim: (out_f as u32, 1, 1),
9043            block_dim: (256, 1, 1),
9044            shared_mem_bytes: 0,
9045        };
9046        let (inf, outf, nu, ne, rbv) = (
9047            in_f as i32,
9048            out_f as i32,
9049            n_used as i32,
9050            n_expert as i32,
9051            rb as i64,
9052        );
9053        let __s_b = self.gpu.stream();
9054        let mut b = __s_b.launch_builder(&f);
9055        b.arg(table)
9056            .arg(sel)
9057            .arg(w)
9058            .arg(act)
9059            .arg(dst)
9060            .arg(&inf)
9061            .arg(&outf)
9062            .arg(&nu)
9063            .arg(&ne)
9064            .arg(&qt)
9065            .arg(&rbv);
9066        unsafe {
9067            b.launch(cfg)?;
9068        }
9069        Ok(())
9070    }
9071
9072    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
9073    pub fn axpy_into(
9074        &self,
9075        src: &CudaSlice<f32>,
9076        alpha: f32,
9077        dst: &mut cudarc::driver::CudaViewMut<f32>,
9078        n: usize,
9079    ) -> Result<(), Box<dyn std::error::Error>> {
9080        let f = self.func("axpy_f32");
9081        let cfg = LaunchConfig::for_num_elems(n as u32);
9082        let (a, ni) = (alpha, n as i32);
9083        let __s_b = self.gpu.stream();
9084        let mut b = __s_b.launch_builder(&f);
9085        b.arg(src).arg(dst).arg(&a).arg(&ni);
9086        unsafe {
9087            b.launch(cfg)?;
9088        }
9089        Ok(())
9090    }
9091
9092    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
9093    pub fn axpy_host_into(
9094        &self,
9095        src: &cudarc::driver::CudaView<'_, f32>,
9096        alpha: f32,
9097        dst: &mut cudarc::driver::CudaViewMut<f32>,
9098        n: usize,
9099    ) -> Result<(), Box<dyn std::error::Error>> {
9100        let f = self.func("axpy_host_f32");
9101        let cfg = LaunchConfig::for_num_elems(n as u32);
9102        let (a, ni) = (alpha, n as i32);
9103        let __s_b = self.gpu.stream();
9104        let mut b = __s_b.launch_builder(&f);
9105        b.arg(src).arg(dst).arg(&a).arg(&ni);
9106        unsafe {
9107            b.launch(cfg)?;
9108        }
9109        Ok(())
9110    }
9111
9112    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
9113    pub fn add_scaled_rows(
9114        &self,
9115        src: &CudaSlice<f32>,
9116        scale: &CudaSlice<f32>,
9117        dst: &mut CudaSlice<f32>,
9118        ncols: usize,
9119        nrows: usize,
9120    ) -> Result<(), Box<dyn std::error::Error>> {
9121        let f = self.func("add_scaled_rows_f32");
9122        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9123        let (nc, nr) = (ncols as i32, nrows as i32);
9124        let __s_b = self.gpu.stream();
9125        let mut b = __s_b.launch_builder(&f);
9126        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
9127        unsafe {
9128            b.launch(cfg)?;
9129        }
9130        Ok(())
9131    }
9132
9133    /// `add_scaled_rows` with an all-ones scale drawn from the resident ones buffer (door H,
9134    /// `MEMRA_HTOD_DIET`) — the UNGATED shared-expert add, without re-uploading the
9135    /// constant every MoE layer-call. Same kernel, same values: the buffer may be longer than
9136    /// `nrows` because `add_scaled_rows_f32` reads only `scale[0..nrows]`.
9137    pub fn add_scaled_rows_ones(
9138        &self,
9139        src: &CudaSlice<f32>,
9140        dst: &mut CudaSlice<f32>,
9141        ncols: usize,
9142        nrows: usize,
9143    ) -> Result<(), Box<dyn std::error::Error>> {
9144        let mut guard = self
9145            .shexp_ones
9146            .lock()
9147            .map_err(|_| "shexp ones buffer is poisoned")?;
9148        if guard.as_ref().map(|b| b.len() < nrows).unwrap_or(true) {
9149            // One upload per process (or per growth step): the serving shapes are t <= 8 for the
9150            // verify walk and the prime's chunk width otherwise.
9151            *guard = Some(self.htod(&vec![1.0f32; nrows.max(64)])?);
9152        }
9153        let ones = guard.as_ref().expect("just ensured");
9154        let f = self.func("add_scaled_rows_f32");
9155        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9156        let (nc, nr) = (ncols as i32, nrows as i32);
9157        let __s_b = self.gpu.stream();
9158        let mut b = __s_b.launch_builder(&f);
9159        b.arg(src).arg(ones).arg(&mut *dst).arg(&nc).arg(&nr);
9160        unsafe {
9161            b.launch(cfg)?;
9162        }
9163        Ok(())
9164    }
9165
9166    /// The `len_d` i32 mirror store, door H aware (`MEMRA_HTOD_DIET`): the async
9167    /// [`Self::i32_set_k`] launch when the door is on, else the shipped synchronizing pageable
9168    /// `memcpy_htod`. Identical value into the identical slot, both stream-ordered.
9169    pub fn i32_mirror_store(
9170        &self,
9171        dst: &mut CudaSlice<i32>,
9172        v: i32,
9173    ) -> Result<(), Box<dyn std::error::Error>> {
9174        if crate::htod_diet_on() {
9175            HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9176            return self.i32_set_k(dst, v);
9177        }
9178        self.gpu.stream().memcpy_htod(&[v], dst)?;
9179        Ok(())
9180    }
9181
9182    /// y[r, :] *= s[r] in place (per-CSR-row macro scale for the grouped prime's gate/up —
9183    /// silu is nonlinear, so per-expert NVFP4 macros must land before it).
9184    pub fn scale_rows(
9185        &self,
9186        y: &mut CudaSlice<f32>,
9187        s: &CudaSlice<f32>,
9188        ncols: usize,
9189        nrows: usize,
9190    ) -> Result<(), Box<dyn std::error::Error>> {
9191        let f = self.func("scale_rows_f32");
9192        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
9193        let (nc, nr) = (ncols as i32, nrows as i32);
9194        let __s_b = self.gpu.stream();
9195        let mut b = __s_b.launch_builder(&f);
9196        b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
9197        unsafe {
9198            b.launch(cfg)?;
9199        }
9200        Ok(())
9201    }
9202
9203    /// Fused grouped-prime tail: join both rank partials (canonical shard order), permute
9204    /// CSR->pair via `inv`, weight, and scatter to tokens in one pass — replaces
9205    /// rows_permute + add + scatter and the three large temporaries they needed.
9206    #[allow(clippy::too_many_arguments)]
9207    pub fn moe_prime_join_scatter(
9208        &self,
9209        y0: &CudaSlice<f32>,
9210        y1: &CudaSlice<f32>,
9211        inv: &CudaSlice<i32>,
9212        w: &CudaSlice<f32>,
9213        out: &mut CudaSlice<f32>,
9214        ncols: usize,
9215        n_used: usize,
9216        t: usize,
9217    ) -> Result<(), Box<dyn std::error::Error>> {
9218        let f = self.func("moe_prime_join_scatter_f32");
9219        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9220        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9221        let __s_b = self.gpu.stream();
9222        let mut b = __s_b.launch_builder(&f);
9223        b.arg(y0)
9224            .arg(y1)
9225            .arg(inv)
9226            .arg(w)
9227            .arg(&mut *out)
9228            .arg(&nc)
9229            .arg(&nu)
9230            .arg(&ti);
9231        unsafe {
9232            b.launch(cfg)?;
9233        }
9234        Ok(())
9235    }
9236
9237    /// out[t, :] += sum_j w[t*n_used+j] * y[t*n_used+j, :], the j-sum sequential per thread —
9238    /// a pinned per-token reduction order, never atomics (the grouped prime's scatter).
9239    pub fn moe_pairs_weighted_scatter(
9240        &self,
9241        y: &CudaSlice<f32>,
9242        w: &CudaSlice<f32>,
9243        out: &mut CudaSlice<f32>,
9244        ncols: usize,
9245        n_used: usize,
9246        t: usize,
9247    ) -> Result<(), Box<dyn std::error::Error>> {
9248        let f = self.func("moe_pairs_weighted_scatter_f32");
9249        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9250        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9251        let __s_b = self.gpu.stream();
9252        let mut b = __s_b.launch_builder(&f);
9253        b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
9254        unsafe {
9255            b.launch(cfg)?;
9256        }
9257        Ok(())
9258    }
9259
9260    // ======== A2 GROUPED MoE PREFILL KERNELS ========
9261
9262    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
9263    pub fn gather_rows(
9264        &self,
9265        src: &CudaSlice<f32>,
9266        idx: &CudaSlice<i32>,
9267        dst: &mut CudaSlice<f32>,
9268        ncols: usize,
9269        m_e: usize,
9270    ) -> Result<(), Box<dyn std::error::Error>> {
9271        let f = self.func("gather_rows_f32");
9272        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
9273        let (nc, me) = (ncols as i32, m_e as i32);
9274        let __s_b = self.gpu.stream();
9275        let mut b = __s_b.launch_builder(&f);
9276        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
9277        unsafe {
9278            b.launch(cfg)?;
9279        }
9280        Ok(())
9281    }
9282
9283    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
9284    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
9285    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
9286    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
9287    #[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
9288    pub fn scatter_slot(
9289        &self,
9290        src: &CudaSlice<f32>,
9291        tok_idx: &CudaSlice<i32>,
9292        slot_idx: &CudaSlice<i32>,
9293        weight: &CudaSlice<f32>,
9294        dst: &mut CudaSlice<f32>,
9295        wbuf: &mut CudaSlice<f32>,
9296        ncols: usize,
9297        n_used: usize,
9298        m_e: usize,
9299    ) -> Result<(), Box<dyn std::error::Error>> {
9300        let f = self.func("scatter_add_slot_f32");
9301        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
9302        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
9303        let __s_b = self.gpu.stream();
9304        let mut b = __s_b.launch_builder(&f);
9305        b.arg(src)
9306            .arg(tok_idx)
9307            .arg(slot_idx)
9308            .arg(weight)
9309            .arg(dst)
9310            .arg(wbuf)
9311            .arg(&nc)
9312            .arg(&nu)
9313            .arg(&me);
9314        unsafe {
9315            b.launch(cfg)?;
9316        }
9317        Ok(())
9318    }
9319
9320    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
9321    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
9322    /// Uses FMA for bit-identity with the sequential axpy path.
9323    pub fn reduce_slots(
9324        &self,
9325        slots: &CudaSlice<f32>,
9326        wbuf: &CudaSlice<f32>,
9327        dst: &mut CudaSlice<f32>,
9328        ncols: usize,
9329        n_used: usize,
9330        t: usize,
9331    ) -> Result<(), Box<dyn std::error::Error>> {
9332        let f = self.func("reduce_slots_f32");
9333        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9334        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9335        let __s_b = self.gpu.stream();
9336        let mut b = __s_b.launch_builder(&f);
9337        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
9338        unsafe {
9339            b.launch(cfg)?;
9340        }
9341        Ok(())
9342    }
9343
9344    /// Canonical slot-order reduction with separately rounded multiply and add.
9345    ///
9346    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
9347    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
9348    pub fn reduce_slots_host(
9349        &self,
9350        slots: &CudaSlice<f32>,
9351        wbuf: &CudaSlice<f32>,
9352        dst: &mut CudaSlice<f32>,
9353        ncols: usize,
9354        n_used: usize,
9355        t: usize,
9356    ) -> Result<(), Box<dyn std::error::Error>> {
9357        let f = self.func("reduce_slots_host_f32");
9358        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
9359        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
9360        let __s_b = self.gpu.stream();
9361        let mut b = __s_b.launch_builder(&f);
9362        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
9363        unsafe {
9364            b.launch(cfg)?;
9365        }
9366        Ok(())
9367    }
9368
9369    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
9370    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
9371    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
9372    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
9373    /// GPU time, ~half of it redundant re-quantization of the same row.
9374    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
9375    pub fn quantize_q8_1_view(
9376        &self,
9377        x: &cudarc::driver::CudaView<f32>,
9378        m: usize,
9379        in_f: usize,
9380    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9381        let f = self.func("quantize_q8_1");
9382        let nblk = in_f / 32;
9383        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
9384        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
9385        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9386        let (inf, mi) = (in_f as i32, m as i32);
9387        let __s_b = self.gpu.stream();
9388        let mut b = __s_b.launch_builder(&f);
9389        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
9390        unsafe {
9391            b.launch(cfg)?;
9392        }
9393        Ok((q, d))
9394    }
9395
9396    pub fn quantize_q8_1(
9397        &self,
9398        x: &CudaSlice<f32>,
9399        m: usize,
9400        in_f: usize,
9401    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9402        let nblk = in_f / 32;
9403        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
9404        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
9405        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
9406        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9407        let (inf, mi) = (in_f as i32, m as i32);
9408        if Self::pdl_on() && Self::pdl_wb_on() {
9409            {
9410                use cudarc::driver::{DevicePtr, DevicePtrMut};
9411                let s = &self.gpu.stream();
9412                let (px, _g0) = x.device_ptr(s);
9413                let (pq, _g1) = q.device_ptr_mut(s);
9414                let (pd, _g2) = d.device_ptr_mut(s);
9415                let mut ps = [
9416                    &px as *const _ as *mut std::ffi::c_void,
9417                    &pq as *const _ as *mut _,
9418                    &pd as *const _ as *mut _,
9419                    &inf as *const _ as *mut _,
9420                    &mi as *const _ as *mut _,
9421                ];
9422                unsafe {
9423                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9424                }
9425            }
9426            return Ok((q, d));
9427        }
9428        let f = self.func("quantize_q8_1");
9429        let __s_b = self.gpu.stream();
9430        let mut b = __s_b.launch_builder(&f);
9431        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
9432        unsafe {
9433            b.launch(cfg)?;
9434        }
9435        Ok((q, d))
9436    }
9437
9438    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
9439    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
9440    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
9441    pub fn quantize_fp4_act(
9442        &self,
9443        x: &CudaSlice<f32>,
9444        m: usize,
9445        in_f: usize,
9446    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
9447        let f = self.func("quantize_fp4_act");
9448        let nb16 = in_f / 16;
9449        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
9450        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
9451        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
9452        let (inf, mi) = (in_f as i32, m as i32);
9453        let __s_b = self.gpu.stream();
9454        let mut b = __s_b.launch_builder(&f);
9455        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
9456        unsafe {
9457            b.launch(cfg)?;
9458        }
9459        Ok((aq4, ad4))
9460    }
9461
9462    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
9463    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
9464    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
9465    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
9466    #[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
9467    pub fn qmatvec_gemm_nvfp4_fp4(
9468        &self,
9469        bytes: &CudaSlice<u8>,
9470        x: &CudaSlice<f32>,
9471        m: usize,
9472        in_f: usize,
9473        out_f: usize,
9474        row_bytes: usize,
9475        scale: f32,
9476    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9477        assert!(
9478            in_f.is_multiple_of(64),
9479            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
9480        );
9481        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
9482        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
9483        if scale != 1.0 {
9484            self.scale_inplace(&mut y, scale, m * out_f)?;
9485        }
9486        Ok(y)
9487    }
9488
9489    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
9490    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
9491    #[allow(clippy::too_many_arguments)]
9492    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9493    #[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
9494    fn fp4_gemm_launch(
9495        &self,
9496        bytes: &CudaSlice<u8>,
9497        aq4: &CudaSlice<u32>,
9498        ad4: &CudaSlice<u8>,
9499        m: usize,
9500        in_f: usize,
9501        out_f: usize,
9502        row_bytes: usize,
9503    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9504        let f = self.func("qmatvec_gemm_nvfp4_fp4");
9505        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9506        const BM: u32 = 64;
9507        const BN: u32 = 256;
9508        let cfg = LaunchConfig {
9509            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
9510            block_dim: (32, 4, 1),
9511            shared_mem_bytes: 0,
9512        };
9513        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9514        let __s_b = self.gpu.stream();
9515        let mut b = __s_b.launch_builder(&f);
9516        b.arg(bytes)
9517            .arg(aq4)
9518            .arg(ad4)
9519            .arg(&mut y)
9520            .arg(&inf)
9521            .arg(&outf)
9522            .arg(&mi)
9523            .arg(&rb);
9524        unsafe {
9525            b.launch(cfg)?;
9526        }
9527        Ok(y)
9528    }
9529
9530    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
9531    pub fn qmatvec_gemm_nvfp4_fp4_raw(
9532        &self,
9533        bytes: &CudaSlice<u8>,
9534        x: &CudaSlice<f32>,
9535        m: usize,
9536        in_f: usize,
9537        out_f: usize,
9538        row_bytes: usize,
9539    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9540        assert!(
9541            in_f.is_multiple_of(64),
9542            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
9543        );
9544        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
9545        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
9546    }
9547
9548    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
9549    pub fn qmatvec_q8_0_fast(
9550        &self,
9551        w: &CudaSlice<u8>,
9552        x: &CudaSlice<f32>,
9553        m: usize,
9554        in_f: usize,
9555        out_f: usize,
9556        row_bytes: usize,
9557    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9558        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9559        let f = self.func("qmatvec_q8_0_dp4a");
9560        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9561        let cfg = LaunchConfig {
9562            grid_dim: (out_f as u32, m as u32, 1),
9563            block_dim: (128, 1, 1),
9564            shared_mem_bytes: 0,
9565        };
9566        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9567        let __s_b = self.gpu.stream();
9568        let mut b = __s_b.launch_builder(&f);
9569        b.arg(w)
9570            .arg(&aq)
9571            .arg(&ad)
9572            .arg(&mut y)
9573            .arg(&inf)
9574            .arg(&outf)
9575            .arg(&mi)
9576            .arg(&rb);
9577        unsafe {
9578            b.launch(cfg)?;
9579        }
9580        Ok(y)
9581    }
9582
9583    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
9584    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9585    pub fn qmatvec_q4_K_fast(
9586        &self,
9587        w: &CudaSlice<u8>,
9588        x: &CudaSlice<f32>,
9589        m: usize,
9590        in_f: usize,
9591        out_f: usize,
9592        row_bytes: usize,
9593    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9594        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9595        let f = self.func("qmatvec_q4_K_dp4a");
9596        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9597        let cfg = LaunchConfig {
9598            grid_dim: (out_f as u32, m as u32, 1),
9599            block_dim: (128, 1, 1),
9600            shared_mem_bytes: 0,
9601        };
9602        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9603        let __s_b = self.gpu.stream();
9604        let mut b = __s_b.launch_builder(&f);
9605        b.arg(w)
9606            .arg(&aq)
9607            .arg(&ad)
9608            .arg(&mut y)
9609            .arg(&inf)
9610            .arg(&outf)
9611            .arg(&mi)
9612            .arg(&rb);
9613        unsafe {
9614            b.launch(cfg)?;
9615        }
9616        Ok(y)
9617    }
9618
9619    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
9620    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9621    pub fn qmatvec_q6_K_fast(
9622        &self,
9623        w: &CudaSlice<u8>,
9624        x: &CudaSlice<f32>,
9625        m: usize,
9626        in_f: usize,
9627        out_f: usize,
9628        row_bytes: usize,
9629    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9630        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9631        let f = self.func("qmatvec_q6_K_dp4a");
9632        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9633        let cfg = LaunchConfig {
9634            grid_dim: (out_f as u32, m as u32, 1),
9635            block_dim: (128, 1, 1),
9636            shared_mem_bytes: 0,
9637        };
9638        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9639        let __s_b = self.gpu.stream();
9640        let mut b = __s_b.launch_builder(&f);
9641        b.arg(w)
9642            .arg(&aq)
9643            .arg(&ad)
9644            .arg(&mut y)
9645            .arg(&inf)
9646            .arg(&outf)
9647            .arg(&mi)
9648            .arg(&rb);
9649        unsafe {
9650            b.launch(cfg)?;
9651        }
9652        Ok(y)
9653    }
9654
9655    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
9656    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9657    pub fn qmatvec_q5_K_fast(
9658        &self,
9659        w: &CudaSlice<u8>,
9660        x: &CudaSlice<f32>,
9661        m: usize,
9662        in_f: usize,
9663        out_f: usize,
9664        row_bytes: usize,
9665    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9666        self.qmatvec_dp4a_named(
9667            "qmatvec_q5_K_dp4a",
9668            &w.slice(0..w.len()),
9669            x,
9670            m,
9671            in_f,
9672            out_f,
9673            row_bytes,
9674        )
9675    }
9676    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
9677    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9678    pub fn qmatvec_q3_K_fast(
9679        &self,
9680        w: &CudaSlice<u8>,
9681        x: &CudaSlice<f32>,
9682        m: usize,
9683        in_f: usize,
9684        out_f: usize,
9685        row_bytes: usize,
9686    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9687        self.qmatvec_dp4a_named(
9688            "qmatvec_q3_K_dp4a",
9689            &w.slice(0..w.len()),
9690            x,
9691            m,
9692            in_f,
9693            out_f,
9694            row_bytes,
9695        )
9696    }
9697    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
9698    pub fn qmatvec_nvfp4_fast_rp(
9699        &self,
9700        w: &CudaSlice<u8>,
9701        x: &CudaSlice<f32>,
9702        m: usize,
9703        in_f: usize,
9704        out_f: usize,
9705        row_bytes: usize,
9706    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9707        assert!(
9708            in_f.is_multiple_of(64),
9709            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9710        );
9711        self.qmatvec_dp4a_named(
9712            "qmatvec_nvfp4_dp4a_rp",
9713            &w.slice(0..w.len()),
9714            x,
9715            m,
9716            in_f,
9717            out_f,
9718            row_bytes,
9719        )
9720    }
9721    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
9722    pub fn qmatvec_nvfp4_fast(
9723        &self,
9724        w: &cudarc::driver::CudaView<'_, u8>,
9725        x: &CudaSlice<f32>,
9726        m: usize,
9727        in_f: usize,
9728        out_f: usize,
9729        row_bytes: usize,
9730    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9731        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
9732        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
9733        assert!(
9734            in_f.is_multiple_of(64),
9735            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9736        );
9737        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
9738    }
9739    /// Slot-major-layout twin of `qmatvec_nvfp4_fast`: bit-identical per row, coalesced
9740    /// reads. Since the 2026-08-29 `MEMRA_NVFP4_BANK_V2` door removal its only in-tree
9741    /// producer of slot-major banks is the EP2 whole-expert bank build; this is EP2's
9742    /// host-canonical oracle reader (plus offline harnesses like moe_tp2_repro).
9743    pub fn qmatvec_nvfp4_fast_v2(
9744        &self,
9745        w: &cudarc::driver::CudaView<'_, u8>,
9746        x: &CudaSlice<f32>,
9747        m: usize,
9748        in_f: usize,
9749        out_f: usize,
9750        row_bytes: usize,
9751    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9752        assert!(
9753            in_f.is_multiple_of(64),
9754            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9755        );
9756        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
9757    }
9758    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
9759    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
9760    pub fn qmatvec_iq4_XS_fast(
9761        &self,
9762        w: &CudaSlice<u8>,
9763        x: &CudaSlice<f32>,
9764        m: usize,
9765        in_f: usize,
9766        out_f: usize,
9767        row_bytes: usize,
9768    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9769        self.qmatvec_dp4a_named(
9770            "qmatvec_iq4_XS_dp4a",
9771            &w.slice(0..w.len()),
9772            x,
9773            m,
9774            in_f,
9775            out_f,
9776            row_bytes,
9777        )
9778    }
9779
9780    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
9781    #[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
9782    fn qmatvec_dp4a_named(
9783        &self,
9784        name: &str,
9785        w: &cudarc::driver::CudaView<'_, u8>,
9786        x: &CudaSlice<f32>,
9787        m: usize,
9788        in_f: usize,
9789        out_f: usize,
9790        row_bytes: usize,
9791    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9792        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9793        let f = self.func(name);
9794        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
9795        let cfg = LaunchConfig {
9796            grid_dim: (out_f as u32, m as u32, 1),
9797            block_dim: (128, 1, 1),
9798            shared_mem_bytes: 0,
9799        };
9800        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9801        let __s_b = self.gpu.stream();
9802        let mut b = __s_b.launch_builder(&f);
9803        b.arg(w)
9804            .arg(&aq)
9805            .arg(&ad)
9806            .arg(&mut y)
9807            .arg(&inf)
9808            .arg(&outf)
9809            .arg(&mi)
9810            .arg(&rb);
9811        unsafe {
9812            b.launch(cfg)?;
9813        }
9814        Ok(y)
9815    }
9816
9817    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
9818    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
9819    /// its output); this entry exists so a routed-expert program can quantize one activation
9820    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
9821    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
9822    #[allow(clippy::too_many_arguments)]
9823    pub fn qmatvec_nvfp4_fast_prequant_into(
9824        &self,
9825        w: &CudaSlice<u8>,
9826        aq: &CudaSlice<i8>,
9827        ad: &CudaSlice<f32>,
9828        y: &mut CudaSlice<f32>,
9829        m: usize,
9830        in_f: usize,
9831        out_f: usize,
9832        row_bytes: usize,
9833    ) -> Result<(), Box<dyn std::error::Error>> {
9834        assert!(
9835            in_f.is_multiple_of(64),
9836            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
9837        );
9838        if y.len() < m * out_f {
9839            return Err(format!(
9840                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
9841                y.len()
9842            )
9843            .into());
9844        }
9845        let f = self.func("qmatvec_nvfp4_dp4a");
9846        let cfg = LaunchConfig {
9847            grid_dim: (out_f as u32, m as u32, 1),
9848            block_dim: (128, 1, 1),
9849            shared_mem_bytes: 0,
9850        };
9851        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
9852        let __s_b = self.gpu.stream();
9853        let mut b = __s_b.launch_builder(&f);
9854        b.arg(w)
9855            .arg(aq)
9856            .arg(ad)
9857            .arg(y)
9858            .arg(&inf)
9859            .arg(&outf)
9860            .arg(&mi)
9861            .arg(&rb);
9862        unsafe {
9863            b.launch(cfg)?;
9864        }
9865        Ok(())
9866    }
9867
9868    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
9869    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
9870    #[allow(clippy::too_many_arguments)]
9871    pub fn matvec_f32_qkv_into(
9872        &self,
9873        wq: &CudaSlice<f32>,
9874        wk: &CudaSlice<f32>,
9875        wv: &CudaSlice<f32>,
9876        wg: &CudaSlice<f32>,
9877        x: &CudaSlice<f32>,
9878        yq: &mut CudaSlice<f32>,
9879        yk: &mut CudaSlice<f32>,
9880        yv: &mut CudaSlice<f32>,
9881        yg: &mut CudaSlice<f32>,
9882        in_f: usize,
9883        out_q: usize,
9884        out_kv: usize,
9885        out_g: usize,
9886    ) -> Result<(), Box<dyn std::error::Error>> {
9887        if !in_f.is_multiple_of(4)
9888            || wq.len() != out_q * in_f
9889            || wk.len() != out_kv * in_f
9890            || wv.len() != out_kv * in_f
9891            || wg.len() < out_g * in_f
9892            || x.len() < in_f
9893            || yq.len() < out_q
9894            || yk.len() < out_kv
9895            || yv.len() < out_kv
9896            || (out_g > 0 && yg.len() < out_g)
9897        {
9898            return Err(format!(
9899                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
9900                 wq={} wk={} wv={} wg={}",
9901                wq.len(),
9902                wk.len(),
9903                wv.len(),
9904                wg.len()
9905            )
9906            .into());
9907        }
9908        let f = self.func("matvec_f32_qkv");
9909        let cfg = LaunchConfig {
9910            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
9911            block_dim: (128, 1, 1),
9912            shared_mem_bytes: 0,
9913        };
9914        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
9915        let __s_b = self.gpu.stream();
9916        let mut b = __s_b.launch_builder(&f);
9917        b.arg(wq)
9918            .arg(wk)
9919            .arg(wv)
9920            .arg(wg)
9921            .arg(x)
9922            .arg(yq)
9923            .arg(yk)
9924            .arg(yv)
9925            .arg(yg)
9926            .arg(&inf)
9927            .arg(&oq)
9928            .arg(&okv)
9929            .arg(&og);
9930        unsafe {
9931            b.launch(cfg)?;
9932        }
9933        Ok(())
9934    }
9935
9936    /// PROGRAM 2 (`MEMRA_NVFP4_SEL_GU`, default OFF): the routed gate and up sweeps in ONE
9937    /// launch. The two sweeps share `sel`/`aq`/`ad` and have identical geometry, so blocks
9938    /// `[0,out_f)` run the exact `_sel_v2` body on the GATE bank and `[out_f,2*out_f)` on the UP
9939    /// bank — per-row BIT-IDENTICAL to two `qmatvec_nvfp4_sel_into` calls, with half the sweep
9940    /// launches and double the grid fill.
9941    ///
9942    /// SLOT-MAJOR ONLY, and the caller proves it: the kernel reads the slot-major byte map, so
9943    /// this refuses banks that do not carry it rather than trusting an env door. In the removed
9944    /// implementation this fusion auto-armed on `nvfp4_bank_v2_on()` with NO door of its own,
9945    /// which is one of the three programs that moved together behind one env var and made the
9946    /// 2026-08-29 bisect unable to name a mechanism (DIAGNOSIS.md).
9947    #[allow(clippy::too_many_arguments)]
9948    pub fn qmatvec_nvfp4_sel_gu_into(
9949        &self,
9950        gate_bank: &CudaSlice<u8>,
9951        up_bank: &CudaSlice<u8>,
9952        sel: &CudaSlice<i32>,
9953        aq: &CudaSlice<i8>,
9954        ad: &CudaSlice<f32>,
9955        yg: &mut CudaSlice<f32>,
9956        yu: &mut CudaSlice<f32>,
9957        n_sel: usize,
9958        in_f: usize,
9959        out_f: usize,
9960        row_bytes: usize,
9961        expert_stride: usize,
9962        slot_major: bool,
9963    ) -> Result<(), Box<dyn std::error::Error>> {
9964        assert!(
9965            in_f.is_multiple_of(64),
9966            "NVFP4 dp4a requires in_f % 64 == 0"
9967        );
9968        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
9969            return Err("NVFP4 gu sel geometry".into());
9970        }
9971        if !slot_major {
9972            return Err(
9973                "NVFP4 gu sel fusion reads slot-major rows: these banks are block_nvfp4 \
9974                        v1 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
9975                    .into(),
9976            );
9977        }
9978        // MEMRA_NVFP4_SEL_GU_RPW=2|4 (sub-door, default OFF, UNPRICED): multirow twin — the
9979        // activation group is read once and reused across RPW rows' gate+up dots. Per-row
9980        // accumulation order and reduce tree are the base kernel's -> bit-identical.
9981        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9982        let rpw = *RPW.get_or_init(|| {
9983            std::env::var("MEMRA_NVFP4_SEL_GU_RPW")
9984                .ok()
9985                .and_then(|v| v.parse().ok())
9986                .filter(|r| *r == 2 || *r == 4)
9987                .unwrap_or(1)
9988        });
9989        let rpw = if out_f.is_multiple_of(rpw) { rpw } else { 1 };
9990        // MEMRA_NVFP4_SEL_GU_WPR=1 (sub-door, default OFF, UNPRICED): warp-per-row.
9991        // NUMERIC-CLASS — the per-row REDUCTION ORDER changes, so a bit tape cannot apply and
9992        // acceptance is the argmax gate plus the boot battery (the QKV_FUSED/BF16_MMV class).
9993        // It is deliberately NOT part of this lane's priced arms, which are all bit-gateable.
9994        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9995        let wpr =
9996            *WPR.get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_GU_WPR").as_deref() == Ok("1"));
9997        let f = self.func(match (wpr, rpw) {
9998            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
9999            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
10000            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
10001            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
10002        });
10003        let cfg = LaunchConfig {
10004            grid_dim: if wpr {
10005                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
10006            } else if rpw == 1 {
10007                ((2 * out_f) as u32, n_sel as u32, 1)
10008            } else {
10009                ((out_f / rpw) as u32, n_sel as u32, 1)
10010            },
10011            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
10012            shared_mem_bytes: 0,
10013        };
10014        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10015        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10016        let (ars, adrs) = (0i64, 0i64);
10017        let __s_b = self.gpu.stream();
10018        let mut b = __s_b.launch_builder(&f);
10019        b.arg(gate_bank)
10020            .arg(up_bank)
10021            .arg(sel)
10022            .arg(aq)
10023            .arg(ad)
10024            .arg(yg)
10025            .arg(yu)
10026            .arg(&inf)
10027            .arg(&outf)
10028            .arg(&ns)
10029            .arg(&rb)
10030            .arg(&es)
10031            .arg(&ars)
10032            .arg(&adrs);
10033        unsafe {
10034            b.launch(cfg)?;
10035        }
10036        Ok(())
10037    }
10038
10039    /// PROGRAM 3 (`MEMRA_NVFP4_SEL_DOWN8`, **default ON since 2026-09-01**): the DOWN sweep and
10040    /// the route-weight
10041    /// combine in ONE launch (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm
10042    /// ported to the NVFP4 banks). Block = `(32, n_sel)`: one warp per slot instead of one warp
10043    /// per (row, slot), and the `n_sel x out_f` partial buffer disappears. BIT-IDENTICAL to
10044    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce tree,
10045    /// same slot-ordered combine chain.
10046    ///
10047    /// Requires slot-major rows and `nsb <= 32` (the fit-block class the reduce identity is
10048    /// argued at). The removed implementation refused on `!nvfp4_bank_v2_on()`; it now refuses on
10049    /// the LAYOUT THE CALLER READ OFF THE BANK, so the guard cannot disagree with the bytes.
10050    #[allow(clippy::too_many_arguments)]
10051    pub fn qmatvec_nvfp4_sel_down8_into(
10052        &self,
10053        bank: &CudaSlice<u8>,
10054        sel: &CudaSlice<i32>,
10055        aq: &CudaSlice<i8>,
10056        ad: &CudaSlice<f32>,
10057        route_w: &CudaSlice<f32>,
10058        md: &CudaSlice<f32>,
10059        dst: &mut CudaSlice<f32>,
10060        n_sel: usize,
10061        in_f: usize,
10062        out_f: usize,
10063        row_bytes: usize,
10064        expert_stride: usize,
10065        act_row_stride: usize,
10066        ad_row_stride: usize,
10067        slot_major: bool,
10068    ) -> Result<(), Box<dyn std::error::Error>> {
10069        if !in_f.is_multiple_of(64)
10070            || n_sel == 0
10071            || n_sel > 8
10072            || (in_f >> 5) > 32
10073            || dst.len() < out_f
10074            || sel.len() < n_sel
10075            || route_w.len() < n_sel
10076        {
10077            return Err(format!(
10078                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
10079                dst.len()
10080            )
10081            .into());
10082        }
10083        if !slot_major {
10084            return Err(
10085                "NVFP4 sel down8 reads slot-major rows: this shard is block_nvfp4 v1 \
10086                        (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
10087                    .into(),
10088            );
10089        }
10090        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
10091        let cfg = LaunchConfig {
10092            grid_dim: (out_f as u32, 1, 1),
10093            block_dim: (32, n_sel as u32, 1),
10094            shared_mem_bytes: 0,
10095        };
10096        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10097        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10098        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
10099        let __s_b = self.gpu.stream();
10100        let mut b = __s_b.launch_builder(&f);
10101        b.arg(bank)
10102            .arg(sel)
10103            .arg(aq)
10104            .arg(ad)
10105            .arg(route_w)
10106            .arg(md)
10107            .arg(dst)
10108            .arg(&inf)
10109            .arg(&outf)
10110            .arg(&ns)
10111            .arg(&rb)
10112            .arg(&es)
10113            .arg(&ars)
10114            .arg(&adrs);
10115        unsafe {
10116            b.launch(cfg)?;
10117        }
10118        Ok(())
10119    }
10120
10121    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
10122    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
10123    #[allow(clippy::too_many_arguments)]
10124    pub fn qmatvec_nvfp4_sel_gu_ep_into(
10125        &self,
10126        gate_bank: &CudaSlice<u8>,
10127        up_bank: &CudaSlice<u8>,
10128        sel: &CudaSlice<i32>,
10129        aq: &CudaSlice<i8>,
10130        ad: &CudaSlice<f32>,
10131        yg: &mut CudaSlice<f32>,
10132        yu: &mut CudaSlice<f32>,
10133        n_sel: usize,
10134        in_f: usize,
10135        out_f: usize,
10136        row_bytes: usize,
10137        expert_stride: usize,
10138        owner: usize,
10139    ) -> Result<(), Box<dyn std::error::Error>> {
10140        assert!(
10141            in_f.is_multiple_of(64),
10142            "NVFP4 dp4a requires in_f % 64 == 0"
10143        );
10144        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
10145            return Err("NVFP4 gu ep geometry".into());
10146        }
10147        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
10148        let cfg = LaunchConfig {
10149            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
10150            block_dim: (128, 1, 1),
10151            shared_mem_bytes: 0,
10152        };
10153        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
10154        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10155        let (ars, adrs) = (0i64, 0i64);
10156        let __s_b = self.gpu.stream();
10157        let mut b = __s_b.launch_builder(&f);
10158        b.arg(gate_bank)
10159            .arg(up_bank)
10160            .arg(sel)
10161            .arg(aq)
10162            .arg(ad)
10163            .arg(yg)
10164            .arg(yu)
10165            .arg(&inf)
10166            .arg(&outf)
10167            .arg(&ns)
10168            .arg(&rb)
10169            .arg(&es)
10170            .arg(&ars)
10171            .arg(&adrs)
10172            .arg(&own);
10173        unsafe {
10174            b.launch(cfg)?;
10175        }
10176        Ok(())
10177    }
10178
10179    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
10180    #[allow(clippy::too_many_arguments)]
10181    pub fn silu_mul_scaled_q8_1_sel_ep_into(
10182        &self,
10183        gate: &CudaSlice<f32>,
10184        up: &CudaSlice<f32>,
10185        gmac: &CudaSlice<f32>,
10186        umac: &CudaSlice<f32>,
10187        sel: &CudaSlice<i32>,
10188        limit: Option<f32>,
10189        out_q: &mut CudaSlice<i8>,
10190        out_d: &mut CudaSlice<f32>,
10191        n_per: usize,
10192        n_sel: usize,
10193        owner: usize,
10194    ) -> Result<(), Box<dyn std::error::Error>> {
10195        if !n_per.is_multiple_of(32)
10196            || out_q.len() < n_sel * n_per
10197            || out_d.len() < n_sel * n_per / 32
10198        {
10199            return Err("NVFP4 silu ep geometry".into());
10200        }
10201        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
10202        let warps = n_sel * n_per / 32;
10203        let cfg = LaunchConfig {
10204            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
10205            block_dim: (128, 1, 1),
10206            shared_mem_bytes: 0,
10207        };
10208        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
10209        let (lim, has) = match limit {
10210            Some(l) => (l, 1i32),
10211            None => (0.0f32, 0i32),
10212        };
10213        let __s_b = self.gpu.stream();
10214        let mut b = __s_b.launch_builder(&f);
10215        b.arg(gate)
10216            .arg(up)
10217            .arg(gmac)
10218            .arg(umac)
10219            .arg(sel)
10220            .arg(&lim)
10221            .arg(&has)
10222            .arg(out_q)
10223            .arg(out_d)
10224            .arg(&np)
10225            .arg(&ns)
10226            .arg(&own);
10227        unsafe {
10228            b.launch(cfg)?;
10229        }
10230        Ok(())
10231    }
10232
10233    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
10234    #[allow(clippy::too_many_arguments)]
10235    pub fn qmatvec_nvfp4_sel_down8_ep_into(
10236        &self,
10237        bank: &CudaSlice<u8>,
10238        sel: &CudaSlice<i32>,
10239        aq: &CudaSlice<i8>,
10240        ad: &CudaSlice<f32>,
10241        route_w: &CudaSlice<f32>,
10242        md: &CudaSlice<f32>,
10243        dst: &mut CudaSlice<f32>,
10244        n_sel: usize,
10245        in_f: usize,
10246        out_f: usize,
10247        row_bytes: usize,
10248        expert_stride: usize,
10249        act_row_stride: usize,
10250        ad_row_stride: usize,
10251        owner: usize,
10252    ) -> Result<(), Box<dyn std::error::Error>> {
10253        if !in_f.is_multiple_of(64)
10254            || n_sel == 0
10255            || n_sel > 8
10256            || (in_f >> 5) > 64
10257            || dst.len() < out_f
10258        {
10259            return Err("NVFP4 down8 ep geometry".into());
10260        }
10261        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
10262        let cfg = LaunchConfig {
10263            grid_dim: (out_f as u32, 1, 1),
10264            block_dim: (32, n_sel as u32, 1),
10265            shared_mem_bytes: 0,
10266        };
10267        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
10268        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10269        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
10270        let __s_b = self.gpu.stream();
10271        let mut b = __s_b.launch_builder(&f);
10272        b.arg(bank)
10273            .arg(sel)
10274            .arg(aq)
10275            .arg(ad)
10276            .arg(route_w)
10277            .arg(md)
10278            .arg(dst)
10279            .arg(&inf)
10280            .arg(&outf)
10281            .arg(&ns)
10282            .arg(&rb)
10283            .arg(&es)
10284            .arg(&ars)
10285            .arg(&adrs)
10286            .arg(&own);
10287        unsafe {
10288            b.launch(cfg)?;
10289        }
10290        Ok(())
10291    }
10292
10293    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
10294    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
10295    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
10296    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
10297    /// kernel — the batching only removes host launch latency.
10298    ///
10299    /// `slot_major` names the LAYOUT OF THE BYTES AT `bank` and is REQUIRED, never defaulted:
10300    /// true routes the `_sel_v2` reader (slot g's 16 qs bytes at `g*16`, scale tail at
10301    /// `nslots*16`), false the block_nvfp4 v1 reader. The caller reads it off the resident bank
10302    /// (`ResidentNvfp4{Column,Row}BankRank::slot_major`) — never off an env door, and never with
10303    /// a default. A defaulted layout scalar in exactly this position is what produced the
10304    /// 2026-08-29 step37 corruption (`kq_fetch(..., int in_f = 0)`,
10305    /// research/step37-bankv3-20260901/DIAGNOSIS.md).
10306    #[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
10307    pub fn qmatvec_nvfp4_sel_into(
10308        &self,
10309        bank: &CudaSlice<u8>,
10310        sel: &CudaSlice<i32>,
10311        aq: &CudaSlice<i8>,
10312        ad: &CudaSlice<f32>,
10313        y: &mut CudaSlice<f32>,
10314        n_sel: usize,
10315        in_f: usize,
10316        out_f: usize,
10317        row_bytes: usize,
10318        expert_stride: usize,
10319        act_row_stride: usize,
10320        ad_row_stride: usize,
10321        slot_major: bool,
10322    ) -> Result<(), Box<dyn std::error::Error>> {
10323        assert!(
10324            in_f.is_multiple_of(64),
10325            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10326        );
10327        if y.len() < n_sel * out_f || sel.len() < n_sel {
10328            return Err(format!(
10329                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
10330                y.len(),
10331                sel.len()
10332            )
10333            .into());
10334        }
10335        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
10336        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
10337        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
10338        // sequential-rows variant was flat). Default stays the single-row form.
10339        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
10340        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
10341        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
10342        let mode = *MR.get_or_init(|| {
10343            if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
10344                2
10345            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
10346                1
10347            } else {
10348                0
10349            }
10350        });
10351        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
10352        // Mode 3 is the SLOT-MAJOR reader, and it is chosen by the BANK's layout, not by an env
10353        // door: `MEMRA_SEL_MR`/`MEMRA_SEL_STREAM` are v1-layout probes and cannot read these
10354        // bytes at all, so the layout overrides them rather than racing them.
10355        let mode = if slot_major { 3 } else { mode };
10356        // MEMRA_NVFP4_SEL_SM_STREAM=1 (sub-door of MEMRA_NVFP4_BANK_SM, default OFF, UNPRICED):
10357        // 8 contiguous rows per block with next-row int4 prefetch. Needs 16B-aligned rows
10358        // (step37 gate/up 2304B yes, down 360B no -> single-row) and one slot per thread.
10359        static SM_STREAM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10360        let sm_stream = mode == 3
10361            && *SM_STREAM
10362                .get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_SM_STREAM").as_deref() == Ok("1"))
10363            && row_bytes.is_multiple_of(16)
10364            && in_f <= 4096;
10365        let kname = match (mode, sm_stream) {
10366            (3, true) => "qmatvec_nvfp4_dp4a_sel_v2s",
10367            (3, false) => "qmatvec_nvfp4_dp4a_sel_v2",
10368            (2, _) => "qmatvec_nvfp4_dp4a_sel_stream",
10369            (1, _) => "qmatvec_nvfp4_dp4a_sel_mr4",
10370            _ => "qmatvec_nvfp4_dp4a_sel",
10371        };
10372        // ENGAGEMENT RECEIPT for PROGRAM 1, one line per distinct (kernel, geometry) pair. The
10373        // door being SET in the environment does not prove the slot-major READER ran; only the
10374        // selected kernel name does. Without this, a pricing cell that reports a flat delta
10375        // cannot distinguish "the program is worth nothing" from "the program never ran" — the
10376        // defect the MEMRA_BF16_MMV lane hit when its engagement grep returned 0 in both arms.
10377        {
10378            static SEEN_SEL: std::sync::Mutex<Vec<(&'static str, usize, usize)>> =
10379                std::sync::Mutex::new(Vec::new());
10380            let combo = (kname, in_f, out_f);
10381            let mut seen = SEEN_SEL.lock().unwrap();
10382            if !seen.contains(&combo) {
10383                seen.push(combo);
10384                eprintln!(
10385                    "[nvfp4-sel] kernel={kname} slot_major={slot_major} in_f={in_f} \
10386                     out_f={out_f} nsb={} row_bytes={row_bytes}",
10387                    in_f >> 5
10388                );
10389            }
10390        }
10391        let f = self.func(kname);
10392        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
10393        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
10394        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
10395        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
10396        let nsb = in_f >> 5;
10397        let fit_block: u32 = if (mode == 0 || mode == 3) && !sm_stream && nsb <= 32 {
10398            32
10399        } else if mode == 1 {
10400            512
10401        } else {
10402            128
10403        };
10404        let cfg = LaunchConfig {
10405            grid_dim: (
10406                if sm_stream {
10407                    (out_f as u32).div_ceil(8)
10408                } else {
10409                    match mode {
10410                        2 => (out_f as u32).div_ceil(16),
10411                        1 => (out_f as u32).div_ceil(4),
10412                        _ => out_f as u32,
10413                    }
10414                },
10415                n_sel as u32,
10416                1,
10417            ),
10418            block_dim: (fit_block, 1, 1),
10419            shared_mem_bytes: 0,
10420        };
10421        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10422        let (rb, es, ars, adrs) = (
10423            row_bytes as i64,
10424            expert_stride as i64,
10425            act_row_stride as i64,
10426            ad_row_stride as i64,
10427        );
10428        let __s_b = self.gpu.stream();
10429        let mut b = __s_b.launch_builder(&f);
10430        b.arg(bank)
10431            .arg(sel)
10432            .arg(aq)
10433            .arg(ad)
10434            .arg(y)
10435            .arg(&inf)
10436            .arg(&outf)
10437            .arg(&ns)
10438            .arg(&rb)
10439            .arg(&es)
10440            .arg(&ars)
10441            .arg(&adrs);
10442        unsafe {
10443            b.launch(cfg)?;
10444        }
10445        Ok(())
10446    }
10447
10448    /// W4A16 selected-expert gate+up pair. `x_bf16` contains checkpoint-rounded BF16
10449    /// activations; selected ids are local to the rank's contiguous expert bank.
10450    #[allow(clippy::too_many_arguments)]
10451    pub fn qmatvec_nvfp4_bf16_sel_dual_rows_into(
10452        &self,
10453        gate_bank: &CudaSlice<u8>,
10454        up_bank: &CudaSlice<u8>,
10455        sel: &CudaSlice<i32>,
10456        token_rows: &CudaSlice<i32>,
10457        x_bf16: &CudaSlice<u8>,
10458        gate_out: &mut CudaSlice<f32>,
10459        up_out: &mut CudaSlice<f32>,
10460        n_sel: usize,
10461        in_f: usize,
10462        out_f: usize,
10463        row_bytes: usize,
10464        expert_stride: usize,
10465        tokens: usize,
10466    ) -> Result<(), Box<dyn std::error::Error>> {
10467        if !in_f.is_multiple_of(64)
10468            || sel.len() < n_sel
10469            || token_rows.len() < n_sel
10470            || gate_out.len() < n_sel * out_f
10471            || up_out.len() < n_sel * out_f
10472            || x_bf16.len() < 2 * in_f * tokens
10473        {
10474            return Err(format!(
10475                "W4A16 NVFP4 dual selected rows geometry sel={} token_rows={} x={} gate={} up={} \
10476                 n_sel={n_sel} tokens={tokens} in={in_f} out={out_f}",
10477                sel.len(),
10478                token_rows.len(),
10479                x_bf16.len(),
10480                gate_out.len(),
10481                up_out.len(),
10482            )
10483            .into());
10484        }
10485        let adjacent_rows = tokens > 1;
10486        let f = if adjacent_rows {
10487            self.func("qmatvec_nvfp4_bf16_sel_quad_rows")
10488        } else {
10489            self.func("qmatvec_nvfp4_bf16_sel_dual_rows")
10490        };
10491        let cfg = LaunchConfig {
10492            grid_dim: (
10493                if adjacent_rows {
10494                    out_f.div_ceil(2) as u32
10495                } else {
10496                    (2 * out_f) as u32
10497                },
10498                n_sel as u32,
10499                1,
10500            ),
10501            block_dim: (256, 1, 1),
10502            shared_mem_bytes: 0,
10503        };
10504        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10505        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10506        let __s_b = self.gpu.stream();
10507        let mut b = __s_b.launch_builder(&f);
10508        b.arg(gate_bank)
10509            .arg(up_bank)
10510            .arg(sel)
10511            .arg(token_rows)
10512            .arg(x_bf16)
10513            .arg(gate_out)
10514            .arg(up_out)
10515            .arg(&inf)
10516            .arg(&outf)
10517            .arg(&ns)
10518            .arg(&rb)
10519            .arg(&es);
10520        unsafe {
10521            b.launch(cfg)?;
10522        }
10523        Ok(())
10524    }
10525
10526    /// Device-routed W4A16 gate+up over fixed token/slot rows. Selection ids remain global;
10527    /// each rank rejects non-owned slots and translates owned ids into its local expert bank.
10528    #[allow(clippy::too_many_arguments)]
10529    pub fn qmatvec_nvfp4_bf16_ep_dual_slots_into(
10530        &self,
10531        gate_bank: &CudaSlice<u8>,
10532        up_bank: &CudaSlice<u8>,
10533        sel: &CudaSlice<i32>,
10534        x_bf16: &CudaSlice<u8>,
10535        gate_out: &mut CudaSlice<f32>,
10536        up_out: &mut CudaSlice<f32>,
10537        n_pairs: usize,
10538        top_k: usize,
10539        in_f: usize,
10540        out_f: usize,
10541        owner_start: usize,
10542        owner_end: usize,
10543        row_bytes: usize,
10544        expert_stride: usize,
10545    ) -> Result<(), Box<dyn std::error::Error>> {
10546        let tokens = n_pairs.div_ceil(top_k);
10547        if top_k == 0
10548            || owner_start >= owner_end
10549            || !in_f.is_multiple_of(64)
10550            || sel.len() < n_pairs
10551            || gate_out.len() < n_pairs * out_f
10552            || up_out.len() < n_pairs * out_f
10553            || x_bf16.len() < 2 * in_f * tokens
10554        {
10555            return Err(format!(
10556                "W4A16 NVFP4 device EP dual-slot geometry sel={} x={} gate={} up={} \
10557                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10558                 owner={owner_start}..{owner_end}",
10559                sel.len(),
10560                x_bf16.len(),
10561                gate_out.len(),
10562                up_out.len(),
10563            )
10564            .into());
10565        }
10566        let pair_parallel = tokens > 1;
10567        let f = if pair_parallel {
10568            self.func("qmatvec_nvfp4_bf16_ep_quad_pairs")
10569        } else {
10570            self.func("qmatvec_nvfp4_bf16_ep_dual_slots")
10571        };
10572        let cfg = LaunchConfig {
10573            grid_dim: (
10574                if pair_parallel {
10575                    out_f.div_ceil(2) as u32
10576                } else {
10577                    (2 * out_f) as u32
10578                },
10579                if pair_parallel { n_pairs as u32 } else { 1 },
10580                1,
10581            ),
10582            block_dim: (256, 1, 1),
10583            shared_mem_bytes: 0,
10584        };
10585        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10586        let (os, oe) = (owner_start as i32, owner_end as i32);
10587        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10588        let __s_b = self.gpu.stream();
10589        let mut b = __s_b.launch_builder(&f);
10590        b.arg(gate_bank)
10591            .arg(up_bank)
10592            .arg(sel)
10593            .arg(x_bf16)
10594            .arg(gate_out)
10595            .arg(up_out)
10596            .arg(&inf)
10597            .arg(&outf)
10598            .arg(&np)
10599            .arg(&tk)
10600            .arg(&os)
10601            .arg(&oe)
10602            .arg(&rb)
10603            .arg(&es);
10604        unsafe {
10605            b.launch(cfg)?;
10606        }
10607        Ok(())
10608    }
10609
10610    /// Optional A8 t=1 gate+up program over global fixed slots.
10611    #[allow(clippy::too_many_arguments)]
10612    pub fn qmatvec_nvfp4_q8_ep_dual_slots_into(
10613        &self,
10614        gate_bank: &CudaSlice<u8>,
10615        up_bank: &CudaSlice<u8>,
10616        sel: &CudaSlice<i32>,
10617        aq: &CudaSlice<i8>,
10618        ad: &CudaSlice<f32>,
10619        gate_out: &mut CudaSlice<f32>,
10620        up_out: &mut CudaSlice<f32>,
10621        n_pairs: usize,
10622        top_k: usize,
10623        in_f: usize,
10624        out_f: usize,
10625        owner_start: usize,
10626        owner_end: usize,
10627        row_bytes: usize,
10628        expert_stride: usize,
10629    ) -> Result<(), Box<dyn std::error::Error>> {
10630        let tokens = n_pairs.div_ceil(top_k);
10631        if top_k == 0
10632            || owner_start >= owner_end
10633            || !in_f.is_multiple_of(64)
10634            || sel.len() < n_pairs
10635            || aq.len() < tokens * in_f
10636            || ad.len() < tokens * (in_f / 32)
10637            || gate_out.len() < n_pairs * out_f
10638            || up_out.len() < n_pairs * out_f
10639        {
10640            return Err(format!(
10641                "W4A8 NVFP4 device EP gate/up geometry sel={} aq={} ad={} gate={} up={} \
10642                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10643                 owner={owner_start}..{owner_end}",
10644                sel.len(),
10645                aq.len(),
10646                ad.len(),
10647                gate_out.len(),
10648                up_out.len(),
10649            )
10650            .into());
10651        }
10652        let f = self.func("qmatvec_nvfp4_q8_ep_dual_slots");
10653        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
10654        let cfg = LaunchConfig {
10655            grid_dim: (out_f as u32, 1, 1),
10656            block_dim: (threads, 1, 1),
10657            shared_mem_bytes: 0,
10658        };
10659        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10660        let (os, oe) = (owner_start as i32, owner_end as i32);
10661        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10662        let __s_b = self.gpu.stream();
10663        let mut b = __s_b.launch_builder(&f);
10664        b.arg(gate_bank)
10665            .arg(up_bank)
10666            .arg(sel)
10667            .arg(aq)
10668            .arg(ad)
10669            .arg(gate_out)
10670            .arg(up_out)
10671            .arg(&inf)
10672            .arg(&outf)
10673            .arg(&np)
10674            .arg(&tk)
10675            .arg(&os)
10676            .arg(&oe)
10677            .arg(&rb)
10678            .arg(&es);
10679        unsafe {
10680            b.launch(cfg)?;
10681        }
10682        Ok(())
10683    }
10684
10685    /// Known-good paired gate+up Q8 schedule: one CTA owns the same output row in both banks,
10686    /// shares the activation bytes, and retains one independent accumulator/reduction per bank.
10687    #[allow(clippy::too_many_arguments)]
10688    pub fn qmatvec_nvfp4_q8_ep_paired_slots_into(
10689        &self,
10690        gate_bank: &CudaSlice<u8>,
10691        up_bank: &CudaSlice<u8>,
10692        sel: &CudaSlice<i32>,
10693        aq: &CudaSlice<i8>,
10694        ad: &CudaSlice<f32>,
10695        gate_out: &mut CudaSlice<f32>,
10696        up_out: &mut CudaSlice<f32>,
10697        n_pairs: usize,
10698        top_k: usize,
10699        in_f: usize,
10700        out_f: usize,
10701        owner_start: usize,
10702        owner_end: usize,
10703        row_bytes: usize,
10704        expert_stride: usize,
10705    ) -> Result<(), Box<dyn std::error::Error>> {
10706        let tokens = n_pairs.div_ceil(top_k);
10707        if top_k == 0
10708            || owner_start >= owner_end
10709            || !in_f.is_multiple_of(64)
10710            || sel.len() < n_pairs
10711            || aq.len() < tokens * in_f
10712            || ad.len() < tokens * (in_f / 32)
10713            || gate_out.len() < n_pairs * out_f
10714            || up_out.len() < n_pairs * out_f
10715        {
10716            return Err(format!(
10717                "W4A8 NVFP4 paired gate/up geometry sel={} aq={} ad={} gate={} up={} \
10718                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
10719                 owner={owner_start}..{owner_end}",
10720                sel.len(),
10721                aq.len(),
10722                ad.len(),
10723                gate_out.len(),
10724                up_out.len(),
10725            )
10726            .into());
10727        }
10728        let f = self.func("qmatvec_nvfp4_q8_ep_paired_slots");
10729        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
10730        let cfg = LaunchConfig {
10731            grid_dim: (out_f as u32, 1, 1),
10732            block_dim: (threads, 1, 1),
10733            shared_mem_bytes: 0,
10734        };
10735        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
10736        let (os, oe) = (owner_start as i32, owner_end as i32);
10737        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10738        let __s_b = self.gpu.stream();
10739        let mut b = __s_b.launch_builder(&f);
10740        b.arg(gate_bank)
10741            .arg(up_bank)
10742            .arg(sel)
10743            .arg(aq)
10744            .arg(ad)
10745            .arg(gate_out)
10746            .arg(up_out)
10747            .arg(&inf)
10748            .arg(&outf)
10749            .arg(&np)
10750            .arg(&tk)
10751            .arg(&os)
10752            .arg(&oe)
10753            .arg(&rb)
10754            .arg(&es);
10755        unsafe {
10756            b.launch(cfg)?;
10757        }
10758        Ok(())
10759    }
10760
10761    /// W4A16 selected down rows scattered into canonical global pair positions on the root.
10762    #[allow(clippy::too_many_arguments)]
10763    pub fn qmatvec_nvfp4_bf16_sel_down_rows_raw(
10764        &self,
10765        bank: &CudaSlice<u8>,
10766        sel: &CudaSlice<i32>,
10767        global_pairs: &CudaSlice<i32>,
10768        activation_bf16: &CudaSlice<u8>,
10769        macros_down: &CudaSlice<f32>,
10770        dst_raw: u64,
10771        n_sel: usize,
10772        in_f: usize,
10773        out_f: usize,
10774        row_bytes: usize,
10775        expert_stride: usize,
10776        total_pairs: usize,
10777    ) -> Result<(), Box<dyn std::error::Error>> {
10778        if !in_f.is_multiple_of(64)
10779            || sel.len() < n_sel
10780            || global_pairs.len() < n_sel
10781            || activation_bf16.len() < 2 * n_sel * in_f
10782            || dst_raw == 0
10783        {
10784            return Err(format!(
10785                "W4A16 NVFP4 down rows geometry sel={} pairs={} act={} dst_raw={dst_raw:#x} \
10786                 n_sel={n_sel} total_pairs={total_pairs} in={in_f} out={out_f}",
10787                sel.len(),
10788                global_pairs.len(),
10789                activation_bf16.len(),
10790            )
10791            .into());
10792        }
10793        let f = self.func("qmatvec_nvfp4_bf16_sel_down_rows");
10794        let cfg = LaunchConfig {
10795            grid_dim: (out_f.div_ceil(2) as u32, n_sel as u32, 1),
10796            block_dim: (256, 1, 1),
10797            shared_mem_bytes: 0,
10798        };
10799        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
10800        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10801        let __s_b = self.gpu.stream();
10802        let mut b = __s_b.launch_builder(&f);
10803        b.arg(bank)
10804            .arg(sel)
10805            .arg(global_pairs)
10806            .arg(activation_bf16)
10807            .arg(macros_down)
10808            .arg(&dst_raw)
10809            .arg(&inf)
10810            .arg(&outf)
10811            .arg(&ns)
10812            .arg(&rb)
10813            .arg(&es);
10814        unsafe {
10815            b.launch(cfg)?;
10816        }
10817        Ok(())
10818    }
10819
10820    /// Device-routed W4A16 down rows. Exactly one owner rank writes each global token/slot row
10821    /// into the root device's peer-accessible slab.
10822    #[allow(clippy::too_many_arguments)]
10823    pub fn qmatvec_nvfp4_bf16_ep_down_slots_raw(
10824        &self,
10825        bank: &CudaSlice<u8>,
10826        sel: &CudaSlice<i32>,
10827        activation_bf16: &CudaSlice<u8>,
10828        macros_down: &CudaSlice<f32>,
10829        dst_raw: u64,
10830        n_pairs: usize,
10831        in_f: usize,
10832        out_f: usize,
10833        owner_start: usize,
10834        owner_end: usize,
10835        row_bytes: usize,
10836        expert_stride: usize,
10837    ) -> Result<(), Box<dyn std::error::Error>> {
10838        if owner_start >= owner_end
10839            || !in_f.is_multiple_of(64)
10840            || sel.len() < n_pairs
10841            || activation_bf16.len() < 2 * n_pairs * in_f
10842            || dst_raw == 0
10843        {
10844            return Err(format!(
10845                "W4A16 NVFP4 device EP down-slot geometry sel={} act={} dst_raw={dst_raw:#x} \
10846                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
10847                sel.len(),
10848                activation_bf16.len(),
10849            )
10850            .into());
10851        }
10852        let f = self.func("qmatvec_nvfp4_bf16_ep_down_slots");
10853        let cfg = LaunchConfig {
10854            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
10855            block_dim: (256, 1, 1),
10856            shared_mem_bytes: 0,
10857        };
10858        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
10859        let (os, oe) = (owner_start as i32, owner_end as i32);
10860        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10861        let __s_b = self.gpu.stream();
10862        let mut b = __s_b.launch_builder(&f);
10863        b.arg(bank)
10864            .arg(sel)
10865            .arg(activation_bf16)
10866            .arg(macros_down)
10867            .arg(&dst_raw)
10868            .arg(&inf)
10869            .arg(&outf)
10870            .arg(&np)
10871            .arg(&os)
10872            .arg(&oe)
10873            .arg(&rb)
10874            .arg(&es);
10875        unsafe {
10876            b.launch(cfg)?;
10877        }
10878        Ok(())
10879    }
10880
10881    /// Pair-parallel multi-token twin of `qmatvec_nvfp4_bf16_ep_down_slots_raw`.
10882    #[allow(clippy::too_many_arguments)]
10883    pub fn qmatvec_nvfp4_bf16_ep_down_pairs_raw(
10884        &self,
10885        bank: &CudaSlice<u8>,
10886        sel: &CudaSlice<i32>,
10887        activation_bf16: &CudaSlice<u8>,
10888        macros_down: &CudaSlice<f32>,
10889        dst_raw: u64,
10890        n_pairs: usize,
10891        in_f: usize,
10892        out_f: usize,
10893        owner_start: usize,
10894        owner_end: usize,
10895        row_bytes: usize,
10896        expert_stride: usize,
10897    ) -> Result<(), Box<dyn std::error::Error>> {
10898        if owner_start >= owner_end
10899            || !in_f.is_multiple_of(64)
10900            || sel.len() < n_pairs
10901            || activation_bf16.len() < 2 * n_pairs * in_f
10902            || dst_raw == 0
10903        {
10904            return Err(format!(
10905                "W4A16 NVFP4 device EP down-pair geometry sel={} act={} dst_raw={dst_raw:#x} \
10906                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
10907                sel.len(),
10908                activation_bf16.len(),
10909            )
10910            .into());
10911        }
10912        let f = self.func("qmatvec_nvfp4_bf16_ep_down_pairs");
10913        let cfg = LaunchConfig {
10914            grid_dim: (out_f.div_ceil(2) as u32, n_pairs as u32, 1),
10915            block_dim: (256, 1, 1),
10916            shared_mem_bytes: 0,
10917        };
10918        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
10919        let (os, oe) = (owner_start as i32, owner_end as i32);
10920        let (rb, es) = (row_bytes as i64, expert_stride as i64);
10921        let __s_b = self.gpu.stream();
10922        let mut b = __s_b.launch_builder(&f);
10923        b.arg(bank)
10924            .arg(sel)
10925            .arg(activation_bf16)
10926            .arg(macros_down)
10927            .arg(&dst_raw)
10928            .arg(&inf)
10929            .arg(&outf)
10930            .arg(&np)
10931            .arg(&os)
10932            .arg(&oe)
10933            .arg(&rb)
10934            .arg(&es);
10935        unsafe {
10936            b.launch(cfg)?;
10937        }
10938        Ok(())
10939    }
10940
10941    /// Host-expf W4A16 SwiGLU selected rows, rounded directly to BF16 for the down projection.
10942    #[allow(clippy::too_many_arguments)]
10943    pub fn silu_mul_scaled_host_expf_bf16_sel_into(
10944        &self,
10945        gate: &CudaSlice<f32>,
10946        up: &CudaSlice<f32>,
10947        gate_macros: &CudaSlice<f32>,
10948        up_macros: &CudaSlice<f32>,
10949        sel: &CudaSlice<i32>,
10950        limit: Option<f32>,
10951        output_bf16: &mut CudaSlice<u8>,
10952        n_per: usize,
10953        n_sel: usize,
10954    ) -> Result<(), Box<dyn std::error::Error>> {
10955        let n = n_per * n_sel;
10956        if sel.len() < n_sel || gate.len() < n || up.len() < n || output_bf16.len() < 2 * n {
10957            return Err(format!(
10958                "W4A16 selected activation geometry sel={} gate={} up={} out={} \
10959                 n_per={n_per} n_sel={n_sel}",
10960                sel.len(),
10961                gate.len(),
10962                up.len(),
10963                output_bf16.len(),
10964            )
10965            .into());
10966        }
10967        let (limit, has_limit) = match limit {
10968            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
10969            Some(limit) => {
10970                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
10971            }
10972            None => (0.0f32, 0i32),
10973        };
10974        let f = self.func("silu_mul_scaled_host_expf_bf16_sel");
10975        let cfg = LaunchConfig::for_num_elems(n as u32);
10976        let (np, ns) = (n_per as i32, n_sel as i32);
10977        let __s_b = self.gpu.stream();
10978        let mut b = __s_b.launch_builder(&f);
10979        b.arg(gate)
10980            .arg(up)
10981            .arg(gate_macros)
10982            .arg(up_macros)
10983            .arg(sel)
10984            .arg(&limit)
10985            .arg(&has_limit)
10986            .arg(output_bf16)
10987            .arg(&np)
10988            .arg(&ns);
10989        unsafe {
10990            b.launch(cfg)?;
10991        }
10992        Ok(())
10993    }
10994
10995    /// Device-routed fixed token/slot W4A16 activation. Global expert ids are translated into
10996    /// rank-local macro rows only on the owning rank.
10997    #[allow(clippy::too_many_arguments)]
10998    pub fn silu_mul_scaled_host_expf_bf16_ep_slots_into(
10999        &self,
11000        gate: &CudaSlice<f32>,
11001        up: &CudaSlice<f32>,
11002        gate_macros: &CudaSlice<f32>,
11003        up_macros: &CudaSlice<f32>,
11004        sel: &CudaSlice<i32>,
11005        owner_start: usize,
11006        owner_end: usize,
11007        limit: Option<f32>,
11008        output_bf16: &mut CudaSlice<u8>,
11009        n_per: usize,
11010        n_pairs: usize,
11011    ) -> Result<(), Box<dyn std::error::Error>> {
11012        let n = n_per * n_pairs;
11013        if owner_start >= owner_end
11014            || sel.len() < n_pairs
11015            || gate.len() < n
11016            || up.len() < n
11017            || output_bf16.len() < 2 * n
11018        {
11019            return Err(format!(
11020                "W4A16 device EP activation geometry sel={} gate={} up={} out={} \
11021                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11022                sel.len(),
11023                gate.len(),
11024                up.len(),
11025                output_bf16.len(),
11026            )
11027            .into());
11028        }
11029        let (limit, has_limit) = match limit {
11030            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11031            Some(limit) => {
11032                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
11033            }
11034            None => (0.0f32, 0i32),
11035        };
11036        let f = self.func("silu_mul_scaled_host_expf_bf16_ep_slots");
11037        let cfg = LaunchConfig::for_num_elems(n as u32);
11038        let (np, pairs) = (n_per as i32, n_pairs as i32);
11039        let (os, oe) = (owner_start as i32, owner_end as i32);
11040        let __s_b = self.gpu.stream();
11041        let mut b = __s_b.launch_builder(&f);
11042        b.arg(gate)
11043            .arg(up)
11044            .arg(gate_macros)
11045            .arg(up_macros)
11046            .arg(sel)
11047            .arg(&limit)
11048            .arg(&has_limit)
11049            .arg(output_bf16)
11050            .arg(&np)
11051            .arg(&pairs)
11052            .arg(&os)
11053            .arg(&oe);
11054        unsafe {
11055            b.launch(cfg)?;
11056        }
11057        Ok(())
11058    }
11059
11060    /// Optional A8 host-expf SwiGLU over global fixed slots.
11061    #[allow(clippy::too_many_arguments)]
11062    pub fn silu_mul_scaled_host_expf_q8_ep_slots_into(
11063        &self,
11064        gate: &CudaSlice<f32>,
11065        up: &CudaSlice<f32>,
11066        gate_macros: &CudaSlice<f32>,
11067        up_macros: &CudaSlice<f32>,
11068        sel: &CudaSlice<i32>,
11069        owner_start: usize,
11070        owner_end: usize,
11071        limit: Option<f32>,
11072        output_q8: &mut CudaSlice<i8>,
11073        output_scales: &mut CudaSlice<f32>,
11074        n_per: usize,
11075        n_pairs: usize,
11076    ) -> Result<(), Box<dyn std::error::Error>> {
11077        let n = n_per * n_pairs;
11078        if owner_start >= owner_end
11079            || !n_per.is_multiple_of(32)
11080            || sel.len() < n_pairs
11081            || gate.len() < n
11082            || up.len() < n
11083            || output_q8.len() < n
11084            || output_scales.len() < n / 32
11085        {
11086            return Err(format!(
11087                "W4A8 device EP activation geometry sel={} gate={} up={} q8={} scales={} \
11088                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11089                sel.len(),
11090                gate.len(),
11091                up.len(),
11092                output_q8.len(),
11093                output_scales.len(),
11094            )
11095            .into());
11096        }
11097        let (limit, has_limit) = match limit {
11098            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11099            Some(limit) => {
11100                return Err(format!("W4A8 selected activation limit {limit} is invalid").into());
11101            }
11102            None => (0.0f32, 0i32),
11103        };
11104        let f = self.func("silu_mul_scaled_host_expf_q8_ep_slots");
11105        let warps = n / 32;
11106        let cfg = LaunchConfig {
11107            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
11108            block_dim: (128, 1, 1),
11109            shared_mem_bytes: 0,
11110        };
11111        let (np, pairs) = (n_per as i32, n_pairs as i32);
11112        let (os, oe) = (owner_start as i32, owner_end as i32);
11113        let __s_b = self.gpu.stream();
11114        let mut b = __s_b.launch_builder(&f);
11115        b.arg(gate)
11116            .arg(up)
11117            .arg(gate_macros)
11118            .arg(up_macros)
11119            .arg(sel)
11120            .arg(&limit)
11121            .arg(&has_limit)
11122            .arg(output_q8)
11123            .arg(output_scales)
11124            .arg(&np)
11125            .arg(&pairs)
11126            .arg(&os)
11127            .arg(&oe);
11128        unsafe {
11129            b.launch(cfg)?;
11130        }
11131        Ok(())
11132    }
11133
11134    /// W4A16 selected-expert down projection plus owner-local route combine. The destination may
11135    /// reside in the model engine's peer-accessible root pool.
11136    #[allow(clippy::too_many_arguments)]
11137    pub fn qmatvec_nvfp4_bf16_sel_down_fma_into(
11138        &self,
11139        bank: &CudaSlice<u8>,
11140        sel: &CudaSlice<i32>,
11141        activation_bf16: &CudaSlice<u8>,
11142        route_weights: &CudaSlice<f32>,
11143        macros_down: &CudaSlice<f32>,
11144        dst: &mut cudarc::driver::CudaViewMut<f32>,
11145        n_sel: usize,
11146        in_f: usize,
11147        out_f: usize,
11148        row_bytes: usize,
11149        expert_stride: usize,
11150    ) -> Result<(), Box<dyn std::error::Error>> {
11151        if !in_f.is_multiple_of(64)
11152            || sel.len() < n_sel
11153            || route_weights.len() < n_sel
11154            || activation_bf16.len() < 2 * n_sel * in_f
11155            || dst.len() < out_f
11156        {
11157            return Err(format!(
11158                "W4A16 NVFP4 down selected geometry sel={} act={} weights={} dst={} \
11159                 n_sel={n_sel} in={in_f} out={out_f}",
11160                sel.len(),
11161                activation_bf16.len(),
11162                route_weights.len(),
11163                dst.len(),
11164            )
11165            .into());
11166        }
11167        let f = self.func("qmatvec_nvfp4_bf16_sel_down_fma");
11168        let cfg = LaunchConfig {
11169            grid_dim: (out_f as u32, 1, 1),
11170            block_dim: (256, 1, 1),
11171            shared_mem_bytes: 0,
11172        };
11173        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11174        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11175        let __s_b = self.gpu.stream();
11176        let mut b = __s_b.launch_builder(&f);
11177        b.arg(bank)
11178            .arg(sel)
11179            .arg(activation_bf16)
11180            .arg(route_weights)
11181            .arg(macros_down)
11182            .arg(dst)
11183            .arg(&inf)
11184            .arg(&outf)
11185            .arg(&ns)
11186            .arg(&rb)
11187            .arg(&es);
11188        unsafe {
11189            b.launch(cfg)?;
11190        }
11191        Ok(())
11192    }
11193
11194    /// Device-routed t=1 W4A16 down projection plus owner-local weighted combine.
11195    #[allow(clippy::too_many_arguments)]
11196    pub fn qmatvec_nvfp4_bf16_ep_down_fma_into(
11197        &self,
11198        bank: &CudaSlice<u8>,
11199        sel: &CudaSlice<i32>,
11200        activation_bf16: &CudaSlice<u8>,
11201        route_weights: &CudaSlice<f32>,
11202        macros_down: &CudaSlice<f32>,
11203        dst: &mut cudarc::driver::CudaViewMut<f32>,
11204        n_pairs: usize,
11205        in_f: usize,
11206        out_f: usize,
11207        owner_start: usize,
11208        owner_end: usize,
11209        row_bytes: usize,
11210        expert_stride: usize,
11211    ) -> Result<(), Box<dyn std::error::Error>> {
11212        if owner_start >= owner_end
11213            || !in_f.is_multiple_of(64)
11214            || sel.len() < n_pairs
11215            || route_weights.len() < n_pairs
11216            || activation_bf16.len() < 2 * n_pairs * in_f
11217            || dst.len() < out_f
11218        {
11219            return Err(format!(
11220                "W4A16 device EP down-FMA geometry sel={} act={} weights={} dst={} \
11221                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
11222                sel.len(),
11223                activation_bf16.len(),
11224                route_weights.len(),
11225                dst.len(),
11226            )
11227            .into());
11228        }
11229        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
11230        let cfg = LaunchConfig {
11231            grid_dim: (out_f as u32, 1, 1),
11232            block_dim: (256, 1, 1),
11233            shared_mem_bytes: 0,
11234        };
11235        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11236        let (os, oe) = (owner_start as i32, owner_end as i32);
11237        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11238        let __s_b = self.gpu.stream();
11239        let mut b = __s_b.launch_builder(&f);
11240        b.arg(bank)
11241            .arg(sel)
11242            .arg(activation_bf16)
11243            .arg(route_weights)
11244            .arg(macros_down)
11245            .arg(dst)
11246            .arg(&inf)
11247            .arg(&outf)
11248            .arg(&np)
11249            .arg(&os)
11250            .arg(&oe)
11251            .arg(&rb)
11252            .arg(&es);
11253        unsafe {
11254            b.launch(cfg)?;
11255        }
11256        Ok(())
11257    }
11258
11259    /// Capture-safe twin of `qmatvec_nvfp4_bf16_ep_down_fma_into`. The destination is one
11260    /// rank-owned row inside a persistent root-device slab.
11261    #[allow(clippy::too_many_arguments)]
11262    pub fn qmatvec_nvfp4_bf16_ep_down_fma_raw(
11263        &self,
11264        bank: &CudaSlice<u8>,
11265        sel: &CudaSlice<i32>,
11266        activation_bf16: &CudaSlice<u8>,
11267        route_weights: &CudaSlice<f32>,
11268        macros_down: &CudaSlice<f32>,
11269        dst_raw: u64,
11270        n_pairs: usize,
11271        in_f: usize,
11272        out_f: usize,
11273        owner_start: usize,
11274        owner_end: usize,
11275        row_bytes: usize,
11276        expert_stride: usize,
11277    ) -> Result<(), Box<dyn std::error::Error>> {
11278        if dst_raw == 0
11279            || owner_start >= owner_end
11280            || !in_f.is_multiple_of(64)
11281            || sel.len() < n_pairs
11282            || route_weights.len() < n_pairs
11283            || activation_bf16.len() < 2 * n_pairs * in_f
11284        {
11285            return Err(format!(
11286                "W4A16 device EP raw down-FMA geometry sel={} act={} weights={} \
11287                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11288                 owner={owner_start}..{owner_end}",
11289                sel.len(),
11290                activation_bf16.len(),
11291                route_weights.len(),
11292            )
11293            .into());
11294        }
11295        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
11296        let cfg = LaunchConfig {
11297            grid_dim: (out_f as u32, 1, 1),
11298            block_dim: (256, 1, 1),
11299            shared_mem_bytes: 0,
11300        };
11301        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11302        let (os, oe) = (owner_start as i32, owner_end as i32);
11303        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11304        let __s_b = self.gpu.stream();
11305        let mut b = __s_b.launch_builder(&f);
11306        b.arg(bank)
11307            .arg(sel)
11308            .arg(activation_bf16)
11309            .arg(route_weights)
11310            .arg(macros_down)
11311            .arg(&dst_raw)
11312            .arg(&inf)
11313            .arg(&outf)
11314            .arg(&np)
11315            .arg(&os)
11316            .arg(&oe)
11317            .arg(&rb)
11318            .arg(&es);
11319        unsafe {
11320            b.launch(cfg)?;
11321        }
11322        Ok(())
11323    }
11324
11325    /// Optional A8 fixed-slot down rows. Each owner rank writes its selected pair rows directly
11326    /// into the root slot slab; the root applies route weights in canonical token/slot order.
11327    #[allow(clippy::too_many_arguments)]
11328    pub fn qmatvec_nvfp4_q8_ep_down_slots_raw(
11329        &self,
11330        bank: &CudaSlice<u8>,
11331        sel: &CudaSlice<i32>,
11332        aq: &CudaSlice<i8>,
11333        ad: &CudaSlice<f32>,
11334        macros_down: &CudaSlice<f32>,
11335        dst_raw: u64,
11336        n_pairs: usize,
11337        in_f: usize,
11338        out_f: usize,
11339        owner_start: usize,
11340        owner_end: usize,
11341        row_bytes: usize,
11342        expert_stride: usize,
11343    ) -> Result<(), Box<dyn std::error::Error>> {
11344        if dst_raw == 0
11345            || owner_start >= owner_end
11346            || !in_f.is_multiple_of(64)
11347            || sel.len() < n_pairs
11348            || aq.len() < n_pairs * in_f
11349            || ad.len() < n_pairs * (in_f / 32)
11350        {
11351            return Err(format!(
11352                "W4A8 device EP raw down-slot geometry sel={} aq={} ad={} \
11353                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11354                 owner={owner_start}..{owner_end}",
11355                sel.len(),
11356                aq.len(),
11357                ad.len(),
11358            )
11359            .into());
11360        }
11361        let f = self.func("qmatvec_nvfp4_q8_ep_down_slots");
11362        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11363        let cfg = LaunchConfig {
11364            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
11365            block_dim: (threads, 1, 1),
11366            shared_mem_bytes: 0,
11367        };
11368        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11369        let (os, oe) = (owner_start as i32, owner_end as i32);
11370        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11371        let __s_b = self.gpu.stream();
11372        let mut b = __s_b.launch_builder(&f);
11373        b.arg(bank)
11374            .arg(sel)
11375            .arg(aq)
11376            .arg(ad)
11377            .arg(macros_down)
11378            .arg(&dst_raw)
11379            .arg(&inf)
11380            .arg(&outf)
11381            .arg(&np)
11382            .arg(&os)
11383            .arg(&oe)
11384            .arg(&rb)
11385            .arg(&es);
11386        unsafe {
11387            b.launch(cfg)?;
11388        }
11389        Ok(())
11390    }
11391
11392    /// Historical A8 t=1 down + owner-local route combine into a persistent root row.
11393    #[allow(clippy::too_many_arguments)]
11394    pub fn qmatvec_nvfp4_q8_ep_down_fma_raw(
11395        &self,
11396        bank: &CudaSlice<u8>,
11397        sel: &CudaSlice<i32>,
11398        aq: &CudaSlice<i8>,
11399        ad: &CudaSlice<f32>,
11400        route_weights: &CudaSlice<f32>,
11401        macros_down: &CudaSlice<f32>,
11402        dst_raw: u64,
11403        n_pairs: usize,
11404        in_f: usize,
11405        out_f: usize,
11406        owner_start: usize,
11407        owner_end: usize,
11408        row_bytes: usize,
11409        expert_stride: usize,
11410    ) -> Result<(), Box<dyn std::error::Error>> {
11411        if dst_raw == 0
11412            || owner_start >= owner_end
11413            || !in_f.is_multiple_of(64)
11414            || sel.len() < n_pairs
11415            || aq.len() < n_pairs * in_f
11416            || ad.len() < n_pairs * (in_f / 32)
11417            || route_weights.len() < n_pairs
11418        {
11419            return Err(format!(
11420                "W4A8 device EP raw down-FMA geometry sel={} aq={} ad={} weights={} \
11421                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
11422                 owner={owner_start}..{owner_end}",
11423                sel.len(),
11424                aq.len(),
11425                ad.len(),
11426                route_weights.len(),
11427            )
11428            .into());
11429        }
11430        let f = self.func("qmatvec_nvfp4_q8_ep_down_fma");
11431        let cfg = LaunchConfig {
11432            grid_dim: (out_f as u32, 1, 1),
11433            block_dim: (256, 1, 1),
11434            shared_mem_bytes: 0,
11435        };
11436        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11437        let (os, oe) = (owner_start as i32, owner_end as i32);
11438        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11439        let __s_b = self.gpu.stream();
11440        let mut b = __s_b.launch_builder(&f);
11441        b.arg(bank)
11442            .arg(sel)
11443            .arg(aq)
11444            .arg(ad)
11445            .arg(route_weights)
11446            .arg(macros_down)
11447            .arg(&dst_raw)
11448            .arg(&inf)
11449            .arg(&outf)
11450            .arg(&np)
11451            .arg(&os)
11452            .arg(&oe)
11453            .arg(&rb)
11454            .arg(&es);
11455        unsafe {
11456            b.launch(cfg)?;
11457        }
11458        Ok(())
11459    }
11460
11461    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
11462    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
11463    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
11464    /// takes the plain SiLU kernel.
11465    #[allow(clippy::too_many_arguments)]
11466    pub fn silu_mul_scaled_q8_1_sel_into(
11467        &self,
11468        gate: &CudaSlice<f32>,
11469        up: &CudaSlice<f32>,
11470        gmac: &CudaSlice<f32>,
11471        umac: &CudaSlice<f32>,
11472        sel: &CudaSlice<i32>,
11473        limit: Option<f32>,
11474        out_q: &mut CudaSlice<i8>,
11475        out_d: &mut CudaSlice<f32>,
11476        n_per: usize,
11477        n_sel: usize,
11478    ) -> Result<(), Box<dyn std::error::Error>> {
11479        let n = n_per * n_sel;
11480        if !n_per.is_multiple_of(32) || out_q.len() < n || out_d.len() < n / 32 {
11481            return Err(format!(
11482                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
11483                out_q.len(),
11484                out_d.len()
11485            )
11486            .into());
11487        }
11488        if let Some(limit) = limit {
11489            if limit <= 1e-6 {
11490                return Err(format!(
11491                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
11492                )
11493                .into());
11494            }
11495            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
11496            let cfg = LaunchConfig::for_num_elems(n as u32);
11497            let (np, ns) = (n_per as i32, n_sel as i32);
11498            let __s_b = self.gpu.stream();
11499            let mut b = __s_b.launch_builder(&f);
11500            b.arg(gate)
11501                .arg(up)
11502                .arg(gmac)
11503                .arg(umac)
11504                .arg(sel)
11505                .arg(&limit)
11506                .arg(out_q)
11507                .arg(out_d)
11508                .arg(&np)
11509                .arg(&ns);
11510            unsafe {
11511                b.launch(cfg)?;
11512            }
11513            return Ok(());
11514        }
11515        let f = self.func("silu_mul_scaled_q8_1_sel");
11516        let cfg = LaunchConfig::for_num_elems(n as u32);
11517        let (np, ns) = (n_per as i32, n_sel as i32);
11518        let __s_b = self.gpu.stream();
11519        let mut b = __s_b.launch_builder(&f);
11520        b.arg(gate)
11521            .arg(up)
11522            .arg(gmac)
11523            .arg(umac)
11524            .arg(sel)
11525            .arg(out_q)
11526            .arg(out_d)
11527            .arg(&np)
11528            .arg(&ns);
11529        unsafe {
11530            b.launch(cfg)?;
11531        }
11532        Ok(())
11533    }
11534
11535    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11536        Ok(self.gpu.stream().clone_htod(v)?)
11537    }
11538    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
11539        Ok(self.gpu.stream().clone_htod(v)?)
11540    }
11541    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
11542    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
11543        Ok(self.gpu.stream().clone_htod(v)?)
11544    }
11545    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
11546        Ok(self.gpu.stream().clone_htod(v)?)
11547    }
11548    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
11549    pub fn dtoh_view(
11550        &self,
11551        d: &cudarc::driver::CudaView<f32>,
11552    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11553        let v = self.gpu.stream().clone_dtoh(d)?;
11554        self.gpu.stream().synchronize()?;
11555        Ok(v)
11556    }
11557    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11558        let v = self.gpu.stream().clone_dtoh(d)?;
11559        self.gpu.stream().synchronize()?;
11560        Ok(v)
11561    }
11562    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
11563    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
11564    /// issuing them together avoids a second stream synchronization in every trunk layer.
11565    pub fn dtoh_pair(
11566        &self,
11567        a: &CudaSlice<f32>,
11568        b: &CudaSlice<f32>,
11569    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
11570        let av = self.gpu.stream().clone_dtoh(a)?;
11571        let bv = self.gpu.stream().clone_dtoh(b)?;
11572        self.gpu.stream().synchronize()?;
11573        Ok((av, bv))
11574    }
11575    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
11576    /// cross a shape-sensitive host boundary.
11577    pub fn dtoh_pair_views(
11578        &self,
11579        a: &cudarc::driver::CudaView<f32>,
11580        b: &cudarc::driver::CudaView<f32>,
11581    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
11582        let av = self.gpu.stream().clone_dtoh(a)?;
11583        let bv = self.gpu.stream().clone_dtoh(b)?;
11584        self.gpu.stream().synchronize()?;
11585        Ok((av, bv))
11586    }
11587    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
11588    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
11589        let v = self.gpu.stream().clone_dtoh(d)?;
11590        self.gpu.stream().synchronize()?;
11591        Ok(v)
11592    }
11593    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
11594    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
11595        let v = self.gpu.stream().clone_dtoh(d)?;
11596        self.gpu.stream().synchronize()?;
11597        Ok(v)
11598    }
11599    pub fn dtoh_u8_view(
11600        &self,
11601        d: &cudarc::driver::CudaView<u8>,
11602    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
11603        let v = self.gpu.stream().clone_dtoh(d)?;
11604        self.gpu.stream().synchronize()?;
11605        Ok(v)
11606    }
11607    /// D2H copy of the first `n` bytes of `d` into a pinned CACHEABLE host buffer: the
11608    /// prefix-cache host-tier demote primitive (lane/kv-host-spill-20260830). Queued on the
11609    /// worker stream and synchronized before returning, exactly like `dtoh_u8`: v1 keeps every
11610    /// host-tier copy on the CUDA owner thread (the HY3 spill law). SEAM (named, not built): an
11611    /// overlapped copy-stream variant would queue this on a dedicated D2H stream with an event
11612    /// handshake against the compute stream; build it only with a tick-stall receipt that says
11613    /// the sync copy is the bottleneck.
11614    pub fn dtoh_u8_into_pinned(
11615        &self,
11616        d: &CudaSlice<u8>,
11617        dst: &mut PinnedHostBuf,
11618        n: usize,
11619    ) -> Result<(), Box<dyn std::error::Error>> {
11620        if n > d.len() || n > dst.len() {
11621            return Err(format!(
11622                "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
11623                d.len(),
11624                dst.len(),
11625            )
11626            .into());
11627        }
11628        if n == 0 {
11629            return Ok(());
11630        }
11631        let host = &mut dst.as_mut_slice()[..n];
11632        self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
11633        self.gpu.stream().synchronize()?;
11634        Ok(())
11635    }
11636    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11637        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11638        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
11639        self.keep_if_capturing(&s);
11640        Ok(s)
11641    }
11642
11643    /// Take the pooled hc-glue decode workspace (MEMRA_HC_DECODE_WS) for one step's walk; put
11644    /// it back with [`Self::hyper_ws_put`]. A `None` here means another walk holds it (or it
11645    /// was never built) — the caller allocates fresh, which is always correct.
11646    pub(crate) fn hyper_ws_take(&self) -> Option<crate::hyper::HyperDecodeWs> {
11647        self.hyper_decode_ws.lock().unwrap().take()
11648    }
11649
11650    pub(crate) fn hyper_ws_put(&self, ws: crate::hyper::HyperDecodeWs) {
11651        *self.hyper_decode_ws.lock().unwrap() = Some(ws);
11652    }
11653
11654    // ---- Verify-walk workspace (MEMRA_VERIFY_WS, door W — see VerifyWs). ----
11655    // take/recycle are no-ops with the door off, so every OFF-arm call site is byte-for-byte
11656    // the shipped program (fresh alloc, ordinary async free). All pooled sites are
11657    // verify-walk-only by construction (rows-exact matmuls, the KDA Rows stash arm, the MoE
11658    // vrows staging), and the pool is per-engine = per-stream: recycle-then-reuse carries the
11659    // same stream-ordering guarantee the async allocator's free-then-alloc does.
11660
11661    /// Pool-or-alloc f32 scratch for a verify-walk site (uninit contract unchanged).
11662    pub(crate) fn vws_uninit(
11663        &self,
11664        n: usize,
11665    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11666        if verify_ws_on() {
11667            let mut ws = self.verify_ws.lock().unwrap();
11668            let ws = &mut *ws;
11669            if let Some(s) = VerifyWs::take(&mut ws.f32_pool, &mut ws.held_bytes, n) {
11670                if VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
11671                    eprintln!(
11672                        "[glm5-verify-ws] engaged: verify-walk buffers recycling through \
11673                         the size-keyed pool (MEMRA_GLM5_VERIFY_WS=1)"
11674                    );
11675                }
11676                return Ok(s);
11677            }
11678        }
11679        self.alloc_uninit::<f32>(n)
11680    }
11681
11682    /// Pool-or-alloc i8 scratch (q8_1 activation planes).
11683    pub(crate) fn vws_uninit_i8(
11684        &self,
11685        n: usize,
11686    ) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
11687        if verify_ws_on() {
11688            let mut ws = self.verify_ws.lock().unwrap();
11689            let ws = &mut *ws;
11690            if let Some(s) = VerifyWs::take(&mut ws.i8_pool, &mut ws.held_bytes, n) {
11691                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11692                return Ok(s);
11693            }
11694        }
11695        self.alloc_uninit::<i8>(n)
11696    }
11697
11698    /// Pool-or-alloc u64 scratch (the MoE vrows pointer tables).
11699    pub(crate) fn vws_uninit_u64(
11700        &self,
11701        n: usize,
11702    ) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
11703        if verify_ws_on() {
11704            let mut ws = self.verify_ws.lock().unwrap();
11705            let ws = &mut *ws;
11706            if let Some(s) = VerifyWs::take(&mut ws.u64_pool, &mut ws.held_bytes, n) {
11707                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11708                return Ok(s);
11709            }
11710        }
11711        self.alloc_uninit::<u64>(n)
11712    }
11713
11714    /// Return a dead verify-walk buffer to the pool (no-op with the door off: the buffer
11715    /// drops to the ordinary async free, the shipped program).
11716    pub(crate) fn vws_recycle(&self, s: CudaSlice<f32>) {
11717        if verify_ws_on() {
11718            let mut ws = self.verify_ws.lock().unwrap();
11719            let ws = &mut *ws;
11720            VerifyWs::put(&mut ws.f32_pool, &mut ws.held_bytes, s);
11721        }
11722    }
11723
11724    /// i8 twin of [`Self::vws_recycle`].
11725    pub(crate) fn vws_recycle_i8(&self, s: CudaSlice<i8>) {
11726        if verify_ws_on() {
11727            let mut ws = self.verify_ws.lock().unwrap();
11728            let ws = &mut *ws;
11729            VerifyWs::put(&mut ws.i8_pool, &mut ws.held_bytes, s);
11730        }
11731    }
11732
11733    /// u64 twin of [`Self::vws_recycle`].
11734    pub(crate) fn vws_recycle_u64(&self, s: CudaSlice<u64>) {
11735        if verify_ws_on() {
11736            let mut ws = self.verify_ws.lock().unwrap();
11737            let ws = &mut *ws;
11738            VerifyWs::put(&mut ws.u64_pool, &mut ws.held_bytes, s);
11739        }
11740    }
11741
11742    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
11743    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
11744    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
11745    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
11746    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
11747    /// back (or kept resident for graph replay). Returns the device token buffer.
11748    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
11749    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
11750    pub fn prob_of_token_device(
11751        &self,
11752        logits: &CudaSlice<f32>,
11753        tok: &CudaSlice<u32>,
11754        n_vocab: usize,
11755    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11756        let nb = ARGMAX_NB;
11757        let mut part = self.alloc_uninit::<f32>(nb)?;
11758        let mut p = self.alloc_uninit::<f32>(1)?;
11759        let f1 = self.func("prob_of_token_partial_f32");
11760        let cfg1 = LaunchConfig {
11761            grid_dim: (nb as u32, 1, 1),
11762            block_dim: (256, 1, 1),
11763            shared_mem_bytes: 0,
11764        };
11765        let nv = n_vocab as i32;
11766        let __s_b1 = self.gpu.stream();
11767        let mut b1 = __s_b1.launch_builder(&f1);
11768        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
11769        unsafe {
11770            b1.launch(cfg1)?;
11771        }
11772        let f2 = self.func("prob_of_token_final_f32");
11773        let cfg2 = LaunchConfig {
11774            grid_dim: (1, 1, 1),
11775            block_dim: (256, 1, 1),
11776            shared_mem_bytes: 0,
11777        };
11778        let nbi = nb as i32;
11779        let __s_b2 = self.gpu.stream();
11780        let mut b2 = __s_b2.launch_builder(&f2);
11781        b2.arg(&part).arg(&mut p).arg(&nbi);
11782        unsafe {
11783            b2.launch(cfg2)?;
11784        }
11785        Ok(p)
11786    }
11787
11788    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
11789    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
11790    /// where the host reads the p-min confidence between replays. Same kernels, same math.
11791    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
11792    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
11793    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
11794    pub fn prob_of_token_device_col(
11795        &self,
11796        logits: &CudaSlice<f32>,
11797        tok_all: &CudaSlice<u32>,
11798        tok_idx: usize,
11799        p_out: &mut CudaSlice<f32>,
11800        p_idx: usize,
11801        n_vocab: usize,
11802    ) -> Result<(), Box<dyn std::error::Error>> {
11803        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
11804        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
11805        let nb = ARGMAX_NB;
11806        let mut part = self.alloc_uninit::<f32>(nb)?;
11807        let f1 = self.func("prob_of_token_partial_f32");
11808        let cfg1 = LaunchConfig {
11809            grid_dim: (nb as u32, 1, 1),
11810            block_dim: (256, 1, 1),
11811            shared_mem_bytes: 0,
11812        };
11813        let nv = n_vocab as i32;
11814        let __s_b1 = self.gpu.stream();
11815        let mut b1 = __s_b1.launch_builder(&f1);
11816        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
11817        unsafe {
11818            b1.launch(cfg1)?;
11819        }
11820        let f2 = self.func("prob_of_token_final_f32");
11821        let cfg2 = LaunchConfig {
11822            grid_dim: (1, 1, 1),
11823            block_dim: (256, 1, 1),
11824            shared_mem_bytes: 0,
11825        };
11826        let nbi = nb as i32;
11827        let __s_b2 = self.gpu.stream();
11828        let mut b2 = __s_b2.launch_builder(&f2);
11829        b2.arg(&part).arg(&mut p_v).arg(&nbi);
11830        unsafe {
11831            b2.launch(cfg2)?;
11832        }
11833        Ok(())
11834    }
11835
11836    pub fn prob_of_token_device_into(
11837        &self,
11838        logits: &CudaSlice<f32>,
11839        tok: &CudaSlice<u32>,
11840        p_out: &mut CudaSlice<f32>,
11841        n_vocab: usize,
11842    ) -> Result<(), Box<dyn std::error::Error>> {
11843        let nb = ARGMAX_NB;
11844        let mut part = self.alloc_uninit::<f32>(nb)?;
11845        let f1 = self.func("prob_of_token_partial_f32");
11846        let cfg1 = LaunchConfig {
11847            grid_dim: (nb as u32, 1, 1),
11848            block_dim: (256, 1, 1),
11849            shared_mem_bytes: 0,
11850        };
11851        let nv = n_vocab as i32;
11852        let __s_b1 = self.gpu.stream();
11853        let mut b1 = __s_b1.launch_builder(&f1);
11854        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
11855        unsafe {
11856            b1.launch(cfg1)?;
11857        }
11858        let f2 = self.func("prob_of_token_final_f32");
11859        let cfg2 = LaunchConfig {
11860            grid_dim: (1, 1, 1),
11861            block_dim: (256, 1, 1),
11862            shared_mem_bytes: 0,
11863        };
11864        let nbi = nb as i32;
11865        let __s_b2 = self.gpu.stream();
11866        let mut b2 = __s_b2.launch_builder(&f2);
11867        b2.arg(&part).arg(p_out).arg(&nbi);
11868        unsafe {
11869            b2.launch(cfg2)?;
11870        }
11871        Ok(())
11872    }
11873
11874    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
11875    /// (graph-constant params, device-varying index). Capture-safe.
11876    pub fn u32_hist_append(
11877        &self,
11878        tok: &CudaSlice<u32>,
11879        hist: &mut CudaSlice<u32>,
11880        idx: &mut CudaSlice<i32>,
11881    ) -> Result<(), Box<dyn std::error::Error>> {
11882        let f = self.func("u32_hist_append");
11883        let cfg = LaunchConfig {
11884            grid_dim: (1, 1, 1),
11885            block_dim: (32, 1, 1),
11886            shared_mem_bytes: 0,
11887        };
11888        let __s_b = self.gpu.stream();
11889        let mut b = __s_b.launch_builder(&f);
11890        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
11891        unsafe {
11892            b.launch(cfg)?;
11893        }
11894        Ok(())
11895    }
11896
11897    pub fn argmax_token_device(
11898        &self,
11899        logits: &CudaSlice<f32>,
11900        n_vocab: usize,
11901    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
11902        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
11903        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
11904        Ok(tok)
11905    }
11906    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
11907    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
11908    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
11909    /// pointer is baked once and the token id never round-trips to host inside steady state. The
11910    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
11911    /// captured passes bake fixed addresses.
11912    pub fn argmax_token_device_into(
11913        &self,
11914        logits: &CudaSlice<f32>,
11915        tok: &mut CudaSlice<u32>,
11916        n_vocab: usize,
11917    ) -> Result<(), Box<dyn std::error::Error>> {
11918        let nb = ARGMAX_NB;
11919        let f1 = self.func("argmax_partial_f32");
11920        let f2 = self.func("argmax_final_f32");
11921        let mut guard = self.argmax_partials.lock().unwrap();
11922        if guard.is_none() {
11923            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
11924            // buffers carry no cudarc events (illegal inside capture).
11925            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
11926            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
11927            *guard = Some((pv, pi));
11928        }
11929        let (part_v, part_i) = guard.as_mut().unwrap();
11930        let nv = n_vocab as i32;
11931        let nbi = nb as i32;
11932        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
11933        let cfg1 = LaunchConfig {
11934            grid_dim: (nb as u32, 1, 1),
11935            block_dim: (256, 1, 1),
11936            shared_mem_bytes: 0,
11937        };
11938        let __s_b1 = self.gpu.stream();
11939        let mut b1 = __s_b1.launch_builder(&f1);
11940        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
11941        unsafe {
11942            b1.launch(cfg1)?;
11943        }
11944        // pass 2: one block reduces NB partials -> token_out[0].
11945        let cfg2 = LaunchConfig {
11946            grid_dim: (1, 1, 1),
11947            block_dim: (256, 1, 1),
11948            shared_mem_bytes: 0,
11949        };
11950        let __s_b2 = self.gpu.stream();
11951        let mut b2 = __s_b2.launch_builder(&f2);
11952        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
11953        unsafe {
11954            b2.launch(cfg2)?;
11955        }
11956        Ok(())
11957    }
11958    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
11959    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
11960    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
11961    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
11962    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
11963    pub fn argmax_token_device_col(
11964        &self,
11965        logits: &CudaSlice<f32>,
11966        col: usize,
11967        n_vocab: usize,
11968        toks: &mut CudaSlice<u32>,
11969        out_idx: usize,
11970    ) -> Result<(), Box<dyn std::error::Error>> {
11971        let nb = ARGMAX_NB;
11972        let f1 = self.func("argmax_partial_f32");
11973        let f2 = self.func("argmax_final_f32");
11974        let mut guard = self.argmax_partials.lock().unwrap();
11975        if guard.is_none() {
11976            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
11977            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
11978            *guard = Some((pv, pi));
11979        }
11980        let (part_v, part_i) = guard.as_mut().unwrap();
11981        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
11982        let nv = n_vocab as i32;
11983        let nbi = nb as i32;
11984        let cfg1 = LaunchConfig {
11985            grid_dim: (nb as u32, 1, 1),
11986            block_dim: (256, 1, 1),
11987            shared_mem_bytes: 0,
11988        };
11989        let __s_b1 = self.gpu.stream();
11990        let mut b1 = __s_b1.launch_builder(&f1);
11991        b1.arg(&col_view)
11992            .arg(&mut *part_v)
11993            .arg(&mut *part_i)
11994            .arg(&nv);
11995        unsafe {
11996            b1.launch(cfg1)?;
11997        }
11998        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
11999        let cfg2 = LaunchConfig {
12000            grid_dim: (1, 1, 1),
12001            block_dim: (256, 1, 1),
12002            shared_mem_bytes: 0,
12003        };
12004        let __s_b2 = self.gpu.stream();
12005        let mut b2 = __s_b2.launch_builder(&f2);
12006        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
12007        unsafe {
12008            b2.launch(cfg2)?;
12009        }
12010        Ok(())
12011    }
12012    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
12013    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
12014        Ok(self.gpu.stream().clone_htod(v)?)
12015    }
12016    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
12017        let v = self.gpu.stream().clone_dtoh(d)?;
12018        self.gpu.stream().synchronize()?;
12019        Ok(v)
12020    }
12021
12022    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
12023        let v = self.gpu.stream().clone_dtoh(d)?;
12024        self.gpu.stream().synchronize()?;
12025        Ok(v)
12026    }
12027    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
12028    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
12029    /// contents change every step, the address must not, so a captured graph can read it).
12030    pub fn htod_u32_into(
12031        &self,
12032        dst: &mut CudaSlice<u32>,
12033        src: &[u32],
12034    ) -> Result<(), Box<dyn std::error::Error>> {
12035        let mut view = dst.slice_mut(0..src.len());
12036        self.gpu.stream().memcpy_htod(src, &mut view)?;
12037        Ok(())
12038    }
12039
12040    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
12041    /// table without changing the device address its reconcile kernel consumes.
12042    pub fn htod_i32_into(
12043        &self,
12044        dst: &mut CudaSlice<i32>,
12045        src: &[i32],
12046    ) -> Result<(), Box<dyn std::error::Error>> {
12047        let mut view = dst.slice_mut(0..src.len());
12048        self.gpu.stream().memcpy_htod(src, &mut view)?;
12049        Ok(())
12050    }
12051
12052    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
12053        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
12054        self.keep_if_capturing(&s);
12055        Ok(s)
12056    }
12057    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
12058    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
12059    #[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
12060    pub fn embed_gather_device_into(
12061        &self,
12062        embd: &CudaSlice<u8>,
12063        token_d: &CudaSlice<u32>,
12064        x_out: &mut CudaSlice<f32>,
12065        n_embd: usize,
12066        qtype: i32,
12067        row_bytes: usize,
12068    ) -> Result<(), Box<dyn std::error::Error>> {
12069        let f = self.func("embed_gather_u32");
12070        let cfg = LaunchConfig {
12071            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
12072            block_dim: (256, 1, 1),
12073            shared_mem_bytes: 0,
12074        };
12075        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
12076        let __s_b = self.gpu.stream();
12077        let mut b = __s_b.launch_builder(&f);
12078        b.arg(embd)
12079            .arg(token_d)
12080            .arg(x_out)
12081            .arg(&ne)
12082            .arg(&qt)
12083            .arg(&rb);
12084        unsafe {
12085            b.launch(cfg)?;
12086        }
12087        Ok(())
12088    }
12089    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
12090    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
12091        let v = self.gpu.stream().clone_dtoh(d)?;
12092        self.gpu.stream().synchronize()?;
12093        Ok(v[0])
12094    }
12095    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
12096    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
12097    /// the counter value after the throwaway capture warmups corrupt it.
12098    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
12099    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
12100    /// copy (fine at stream-idle boundaries, poison mid-round).
12101    pub fn i32_set_k(
12102        &self,
12103        dst: &mut CudaSlice<i32>,
12104        v: i32,
12105    ) -> Result<(), Box<dyn std::error::Error>> {
12106        let f = self.func("i32_set_k");
12107        let cfg = LaunchConfig {
12108            grid_dim: (1, 1, 1),
12109            block_dim: (1, 1, 1),
12110            shared_mem_bytes: 0,
12111        };
12112        let idx = 0i32;
12113        let __s_b = self.gpu.stream();
12114        let mut b = __s_b.launch_builder(&f);
12115        b.arg(dst).arg(&v).arg(&idx);
12116        unsafe {
12117            b.launch(cfg)?;
12118        }
12119        Ok(())
12120    }
12121
12122    pub fn set_i32_one(
12123        &self,
12124        d: &mut CudaSlice<i32>,
12125        v: i32,
12126    ) -> Result<(), Box<dyn std::error::Error>> {
12127        self.gpu.stream().memcpy_htod(&[v], d)?;
12128        Ok(())
12129    }
12130    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
12131    /// during priming / capture-state restore.
12132    pub fn set_u32_one(
12133        &self,
12134        d: &mut CudaSlice<u32>,
12135        v: u32,
12136    ) -> Result<(), Box<dyn std::error::Error>> {
12137        self.gpu.stream().memcpy_htod(&[v], d)?;
12138        Ok(())
12139    }
12140    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
12141    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
12142        let v = self.gpu.stream().clone_dtoh(d)?;
12143        self.gpu.stream().synchronize()?;
12144        Ok(v[0])
12145    }
12146    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
12147    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
12148        Ok(self.gpu.stream().clone_htod(bytes)?)
12149    }
12150    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
12151    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
12152    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
12153    #[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
12154    pub fn embed_gather_device(
12155        &self,
12156        embd: &CudaSlice<u8>,
12157        token_d: &CudaSlice<u32>,
12158        n_embd: usize,
12159        qtype: i32,
12160        row_bytes: usize,
12161    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12162        let f = self.func("embed_gather_u32");
12163        let mut x = self.alloc_uninit::<f32>(n_embd)?;
12164        let cfg = LaunchConfig {
12165            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
12166            block_dim: (256, 1, 1),
12167            shared_mem_bytes: 0,
12168        };
12169        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
12170        let __s_b = self.gpu.stream();
12171        let mut b = __s_b.launch_builder(&f);
12172        b.arg(embd)
12173            .arg(token_d)
12174            .arg(&mut x)
12175            .arg(&ne)
12176            .arg(&qt)
12177            .arg(&rb);
12178        unsafe {
12179            b.launch(cfg)?;
12180        }
12181        Ok(x)
12182    }
12183
12184    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
12185    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
12186    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
12187    #[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
12188    pub fn embed_gather_device_t(
12189        &self,
12190        embd: &CudaSlice<u8>,
12191        tokens: &[u32],
12192        n_embd: usize,
12193        qtype: i32,
12194        row_bytes: usize,
12195    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12196        let t = tokens.len();
12197        let tok_d = self.gpu.stream().clone_htod(tokens)?;
12198        let f = self.func("embed_gather_u32_t");
12199        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12200        let cfg = LaunchConfig {
12201            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12202            block_dim: (256, 1, 1),
12203            shared_mem_bytes: 0,
12204        };
12205        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12206        let __s_b = self.gpu.stream();
12207        let mut b = __s_b.launch_builder(&f);
12208        b.arg(embd)
12209            .arg(&tok_d)
12210            .arg(&mut x)
12211            .arg(&ne)
12212            .arg(&qt)
12213            .arg(&rb)
12214            .arg(&ti);
12215        unsafe {
12216            b.launch(cfg)?;
12217        }
12218        Ok(x)
12219    }
12220
12221    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
12222    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
12223    /// as embed_gather_device_t — bit-identical rows.
12224    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
12225    #[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
12226    pub fn embed_gather_device_tv(
12227        &self,
12228        embd: &CudaSlice<u8>,
12229        tok_v: &cudarc::driver::CudaView<u32>,
12230        t: usize,
12231        n_embd: usize,
12232        qtype: i32,
12233        row_bytes: usize,
12234    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12235        let f = self.func("embed_gather_u32_t");
12236        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12237        let cfg = LaunchConfig {
12238            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12239            block_dim: (256, 1, 1),
12240            shared_mem_bytes: 0,
12241        };
12242        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12243        let __s_b = self.gpu.stream();
12244        let mut b = __s_b.launch_builder(&f);
12245        b.arg(embd)
12246            .arg(tok_v)
12247            .arg(&mut x)
12248            .arg(&ne)
12249            .arg(&qt)
12250            .arg(&rb)
12251            .arg(&ti);
12252        unsafe {
12253            b.launch(cfg)?;
12254        }
12255        Ok(x)
12256    }
12257
12258    #[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
12259    pub fn embed_gather_device_td(
12260        &self,
12261        embd: &CudaSlice<u8>,
12262        tok_d: &CudaSlice<u32>,
12263        t: usize,
12264        n_embd: usize,
12265        qtype: i32,
12266        row_bytes: usize,
12267    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12268        let f = self.func("embed_gather_u32_t");
12269        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
12270        let cfg = LaunchConfig {
12271            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
12272            block_dim: (256, 1, 1),
12273            shared_mem_bytes: 0,
12274        };
12275        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
12276        let __s_b = self.gpu.stream();
12277        let mut b = __s_b.launch_builder(&f);
12278        b.arg(embd)
12279            .arg(tok_d)
12280            .arg(&mut x)
12281            .arg(&ne)
12282            .arg(&qt)
12283            .arg(&rb)
12284            .arg(&ti);
12285        unsafe {
12286            b.launch(cfg)?;
12287        }
12288        Ok(x)
12289    }
12290
12291    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
12292    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
12293    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
12294    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
12295    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
12296    #[inline]
12297    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
12298    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
12299        if self
12300            .capture_keep_on
12301            .load(std::sync::atomic::Ordering::Relaxed)
12302        {
12303            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
12304        }
12305    }
12306
12307    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
12308        &self,
12309        n: usize,
12310    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
12311        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12312        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
12313        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
12314        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
12315        // not cover engine-internal buffers). Debug-only: massive launch overhead.
12316        {
12317            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12318            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
12319                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
12320                use cudarc::driver::DevicePtrMut;
12321                let n_bytes = s.len() * std::mem::size_of::<T>();
12322                let stream = self.gpu.stream();
12323                let (p_, _g) = s.device_ptr_mut(&stream);
12324                unsafe {
12325                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
12326                        .result()?;
12327                }
12328            }
12329        }
12330        self.keep_if_capturing(&s);
12331        Ok(s)
12332    }
12333
12334    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
12335    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
12336    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
12337    /// consumers alloc through this (m=1 decode arms).
12338    pub fn uninit_q8_pair(
12339        &self,
12340        n: usize,
12341    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12342        Ok((
12343            self.alloc_uninit::<i8>(n)?,
12344            self.alloc_uninit::<f32>(n / 32)?,
12345        ))
12346    }
12347
12348    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12349        self.alloc_uninit::<f32>(n)
12350    }
12351
12352    /// i8 uninitialized scratch (same contract as `uninit`).
12353    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12354        self.alloc_uninit::<i8>(n)
12355    }
12356
12357    /// i32 uninitialized scratch (same contract as `uninit`) — the DSA indexer's position lists.
12358    pub fn uninit_i32(&self, n: usize) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
12359        self.alloc_uninit::<i32>(n)
12360    }
12361
12362    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
12363    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
12364    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
12365    #[allow(clippy::too_many_arguments)]
12366    pub fn rms_norm3(
12367        &self,
12368        x: &CudaSlice<f32>,
12369        w0: &CudaSlice<f32>,
12370        w1: &CudaSlice<f32>,
12371        w2: &CudaSlice<f32>,
12372        d0: &mut CudaSlice<f32>,
12373        d1: &mut CudaSlice<f32>,
12374        d2: &mut CudaSlice<f32>,
12375        ncols: usize,
12376        nrows: usize,
12377        eps: f32,
12378    ) -> Result<(), Box<dyn std::error::Error>> {
12379        let f = self.func("rms_norm3_f32");
12380        let cfg = LaunchConfig {
12381            grid_dim: (nrows as u32, 1, 1),
12382            block_dim: (rms_block(), 1, 1),
12383            shared_mem_bytes: 0,
12384        };
12385        let (nc, e) = (ncols as i32, eps);
12386        let __s_b = self.gpu.stream();
12387        let mut b = __s_b.launch_builder(&f);
12388        b.arg(x)
12389            .arg(w0)
12390            .arg(w1)
12391            .arg(w2)
12392            .arg(d0)
12393            .arg(d1)
12394            .arg(d2)
12395            .arg(&nc)
12396            .arg(&e);
12397        unsafe {
12398            b.launch(cfg)?;
12399        }
12400        Ok(())
12401    }
12402
12403    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
12404    #[allow(clippy::too_many_arguments)]
12405    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
12406    /// piggybacks on the same conditions.
12407    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
12408        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12409        *WARP_ON.get_or_init(|| {
12410            std::env::var("MEMRA_QKVNORM_W")
12411                .map(|v| v != "0")
12412                .unwrap_or(true)
12413        }) && ncols.is_multiple_of(4)
12414            && rows >= 64
12415    }
12416
12417    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
12418    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
12419    #[allow(clippy::too_many_arguments)]
12420    pub fn rms_norm_qkv_w4b(
12421        &self,
12422        q: &CudaSlice<f32>,
12423        k: &CudaSlice<f32>,
12424        v: &CudaSlice<f32>,
12425        wq: &CudaSlice<f32>,
12426        wk: &CudaSlice<f32>,
12427        wv: &CudaSlice<f32>,
12428        dq: &mut CudaSlice<f32>,
12429        dk: &mut CudaSlice<f32>,
12430        dv: &mut CudaSlice<f32>,
12431        dvb: &mut CudaSlice<u8>,
12432        ncols: usize,
12433        rq: usize,
12434        rk: usize,
12435        eps: f32,
12436        vf16: bool,
12437    ) -> Result<(), Box<dyn std::error::Error>> {
12438        assert!(ncols.is_multiple_of(4) && rq + 2 * rk >= 64);
12439        let f = self.func("rms_norm_qkv_w4b_f32");
12440        let rows = (rq + 2 * rk) as u32;
12441        let cfg = LaunchConfig {
12442            grid_dim: (rows.div_ceil(8), 1, 1),
12443            block_dim: (256, 1, 1),
12444            shared_mem_bytes: 0,
12445        };
12446        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
12447        let vf = vf16 as i32;
12448        let __s_b = self.gpu.stream();
12449        let mut b = __s_b.launch_builder(&f);
12450        b.arg(q)
12451            .arg(k)
12452            .arg(v)
12453            .arg(wq)
12454            .arg(wk)
12455            .arg(wv)
12456            .arg(dq)
12457            .arg(dk)
12458            .arg(dv)
12459            .arg(&mut *dvb)
12460            .arg(&nc)
12461            .arg(&rqi)
12462            .arg(&rki)
12463            .arg(&rvi)
12464            .arg(&e)
12465            .arg(&vf);
12466        unsafe {
12467            b.launch(cfg)?;
12468        }
12469        Ok(())
12470    }
12471
12472    #[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
12473    pub fn rms_norm_qkv(
12474        &self,
12475        q: &CudaSlice<f32>,
12476        k: &CudaSlice<f32>,
12477        v: &CudaSlice<f32>,
12478        wq: &CudaSlice<f32>,
12479        wk: &CudaSlice<f32>,
12480        wv: &CudaSlice<f32>,
12481        dq: &mut CudaSlice<f32>,
12482        dk: &mut CudaSlice<f32>,
12483        dv: &mut CudaSlice<f32>,
12484        ncols: usize,
12485        rq: usize,
12486        rk: usize,
12487        eps: f32,
12488    ) -> Result<(), Box<dyn std::error::Error>> {
12489        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
12490        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
12491        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
12492        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12493        let warp_on = *WARP_ON.get_or_init(|| {
12494            std::env::var("MEMRA_QKVNORM_W")
12495                .map(|v| v != "0")
12496                .unwrap_or(true)
12497        });
12498        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
12499        // replay numerics are untouched on every model; only prefill depth takes the new config.
12500        if warp_on && ncols.is_multiple_of(4) && rq + 2 * rk >= 64 {
12501            let f = self.func("rms_norm_qkv_w4_f32");
12502            let rows = (rq + 2 * rk) as u32;
12503            let cfg = LaunchConfig {
12504                grid_dim: (rows.div_ceil(8), 1, 1),
12505                block_dim: (256, 1, 1),
12506                shared_mem_bytes: 0,
12507            };
12508            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
12509            let __s_b = self.gpu.stream();
12510            let mut b = __s_b.launch_builder(&f);
12511            b.arg(q)
12512                .arg(k)
12513                .arg(v)
12514                .arg(wq)
12515                .arg(wk)
12516                .arg(wv)
12517                .arg(dq)
12518                .arg(dk)
12519                .arg(dv)
12520                .arg(&nc)
12521                .arg(&rqi)
12522                .arg(&rki)
12523                .arg(&rvi)
12524                .arg(&e);
12525            unsafe {
12526                b.launch(cfg)?;
12527            }
12528            return Ok(());
12529        }
12530        let f = self.func("rms_norm_qkv_f32");
12531        let grid = (rq + 2 * rk) as u32;
12532        let cfg = LaunchConfig {
12533            grid_dim: (grid, 1, 1),
12534            block_dim: (rms_block(), 1, 1),
12535            shared_mem_bytes: 0,
12536        };
12537        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
12538        let __s_b = self.gpu.stream();
12539        let mut b = __s_b.launch_builder(&f);
12540        b.arg(q)
12541            .arg(k)
12542            .arg(v)
12543            .arg(wq)
12544            .arg(wk)
12545            .arg(wv)
12546            .arg(dq)
12547            .arg(dk)
12548            .arg(dv)
12549            .arg(&nc)
12550            .arg(&rqi)
12551            .arg(&rki)
12552            .arg(&e);
12553        unsafe {
12554            b.launch(cfg)?;
12555        }
12556        Ok(())
12557    }
12558
12559    /// gemma4 fused pair of rms_norms over two different inputs (same width).
12560    #[allow(clippy::too_many_arguments)]
12561    pub fn rms_norm2x(
12562        &self,
12563        a: &CudaSlice<f32>,
12564        bb: &CudaSlice<f32>,
12565        wa: &CudaSlice<f32>,
12566        wb: &CudaSlice<f32>,
12567        da: &mut CudaSlice<f32>,
12568        db: &mut CudaSlice<f32>,
12569        ncols: usize,
12570        nrows: usize,
12571        eps: f32,
12572    ) -> Result<(), Box<dyn std::error::Error>> {
12573        let f = self.func("rms_norm2x_f32");
12574        let cfg = LaunchConfig {
12575            grid_dim: (2 * nrows as u32, 1, 1),
12576            block_dim: (rms_block(), 1, 1),
12577            shared_mem_bytes: 0,
12578        };
12579        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
12580        let __s_b = self.gpu.stream();
12581        let mut b = __s_b.launch_builder(&f);
12582        b.arg(a)
12583            .arg(bb)
12584            .arg(wa)
12585            .arg(wb)
12586            .arg(da)
12587            .arg(db)
12588            .arg(&nc)
12589            .arg(&nr)
12590            .arg(&e);
12591        unsafe {
12592            b.launch(cfg)?;
12593        }
12594        Ok(())
12595    }
12596
12597    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
12598    pub fn softcap(
12599        &self,
12600        y: &mut CudaSlice<f32>,
12601        cap: f32,
12602        n: usize,
12603    ) -> Result<(), Box<dyn std::error::Error>> {
12604        let f = self.func("softcap_f32");
12605        let cfg = LaunchConfig::for_num_elems(n as u32);
12606        let ni = n as i32;
12607        let __s_b = self.gpu.stream();
12608        let mut b = __s_b.launch_builder(&f);
12609        b.arg(y).arg(&cap).arg(&ni);
12610        unsafe {
12611            b.launch(cfg)?;
12612        }
12613        Ok(())
12614    }
12615
12616    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
12617    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
12618    pub fn mask_ids_rows(
12619        &self,
12620        y: &mut CudaSlice<f32>,
12621        ids: &CudaSlice<i32>,
12622        n_ids: usize,
12623        n_vocab: usize,
12624        t: usize,
12625    ) -> Result<(), Box<dyn std::error::Error>> {
12626        let f = self.func("mask_ids_rows_f32");
12627        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
12628        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
12629        let __s_b = self.gpu.stream();
12630        let mut b = __s_b.launch_builder(&f);
12631        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
12632        unsafe {
12633            b.launch(cfg)?;
12634        }
12635        Ok(())
12636    }
12637
12638    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
12639    #[allow(clippy::too_many_arguments)]
12640    pub fn add_scale_rms_norm(
12641        &self,
12642        a: &CudaSlice<f32>,
12643        b_in: &CudaSlice<f32>,
12644        c: f32,
12645        w: &CudaSlice<f32>,
12646        res: &mut CudaSlice<f32>,
12647        dst: &mut CudaSlice<f32>,
12648        ncols: usize,
12649        nrows: usize,
12650        eps: f32,
12651    ) -> Result<(), Box<dyn std::error::Error>> {
12652        let f = self.func("add_scale_rms_norm_f32");
12653        let cfg = LaunchConfig {
12654            grid_dim: (nrows as u32, 1, 1),
12655            block_dim: (rms_block(), 1, 1),
12656            shared_mem_bytes: 0,
12657        };
12658        let (nc, e2) = (ncols as i32, eps);
12659        let __s_b = self.gpu.stream();
12660        let mut b = __s_b.launch_builder(&f);
12661        b.arg(a)
12662            .arg(b_in)
12663            .arg(&c)
12664            .arg(w)
12665            .arg(res)
12666            .arg(dst)
12667            .arg(&nc)
12668            .arg(&e2);
12669        unsafe {
12670            b.launch(cfg)?;
12671        }
12672        Ok(())
12673    }
12674
12675    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
12676    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
12677    #[allow(clippy::too_many_arguments)]
12678    pub fn add_scale_rms_norm_q8_1(
12679        &self,
12680        a: &CudaSlice<f32>,
12681        b_in: &CudaSlice<f32>,
12682        c: f32,
12683        w: &CudaSlice<f32>,
12684        res: &mut CudaSlice<f32>,
12685        ncols: usize,
12686        nrows: usize,
12687        eps: f32,
12688    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12689        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12690        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12691        let (nc, e2) = (ncols as i32, eps);
12692        if Self::pdl_on() && Self::pdl_wb_on() {
12693            {
12694                use cudarc::driver::{DevicePtr, DevicePtrMut};
12695                let s = &self.gpu.stream();
12696                let (pa, _g0) = a.device_ptr(s);
12697                let (pb, _g1) = b_in.device_ptr(s);
12698                let (pw, _g2) = w.device_ptr(s);
12699                let (pr, _g3) = res.device_ptr_mut(s);
12700                let (pq, _g4) = out_q.device_ptr_mut(s);
12701                let (pd, _g5) = out_d.device_ptr_mut(s);
12702                let mut ps = [
12703                    &pa as *const _ as *mut std::ffi::c_void,
12704                    &pb as *const _ as *mut _,
12705                    &c as *const _ as *mut _,
12706                    &pw as *const _ as *mut _,
12707                    &pr as *const _ as *mut _,
12708                    &pq as *const _ as *mut _,
12709                    &pd as *const _ as *mut _,
12710                    &nc as *const _ as *mut _,
12711                    &e2 as *const _ as *mut _,
12712                ];
12713                unsafe {
12714                    self.launch_pdl(
12715                        "add_scale_rms_norm_q8_1",
12716                        (nrows as u32, 1, 1),
12717                        (rms_block(), 1, 1),
12718                        &mut ps,
12719                    )?;
12720                }
12721            }
12722            return Ok((out_q, out_d));
12723        }
12724        let f = self.func("add_scale_rms_norm_q8_1");
12725        let cfg = LaunchConfig {
12726            grid_dim: (nrows as u32, 1, 1),
12727            block_dim: (rms_block(), 1, 1),
12728            shared_mem_bytes: 0,
12729        };
12730        let __s_b = self.gpu.stream();
12731        let mut b = __s_b.launch_builder(&f);
12732        b.arg(a)
12733            .arg(b_in)
12734            .arg(&c)
12735            .arg(w)
12736            .arg(res)
12737            .arg(&mut out_q)
12738            .arg(&mut out_d)
12739            .arg(&nc)
12740            .arg(&e2);
12741        unsafe {
12742            b.launch(cfg)?;
12743        }
12744        Ok((out_q, out_d))
12745    }
12746
12747    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
12748    #[allow(clippy::too_many_arguments)]
12749    pub fn add_scale_rms_norm_q8_1_into(
12750        &self,
12751        a: &CudaSlice<f32>,
12752        b_in: &CudaSlice<f32>,
12753        c: f32,
12754        w: &CudaSlice<f32>,
12755        res: &mut CudaSlice<f32>,
12756        ncols: usize,
12757        nrows: usize,
12758        eps: f32,
12759        out_q: &mut CudaSlice<i8>,
12760        out_d: &mut CudaSlice<f32>,
12761    ) -> Result<(), Box<dyn std::error::Error>> {
12762        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
12763        let (nc, e2) = (ncols as i32, eps);
12764        if Self::pdl_on() && Self::pdl_wb_on() {
12765            use cudarc::driver::{DevicePtr, DevicePtrMut};
12766            let s = &self.gpu.stream();
12767            let (pa, _g0) = a.device_ptr(s);
12768            let (pb, _g1) = b_in.device_ptr(s);
12769            let (pw, _g2) = w.device_ptr(s);
12770            let (pr, _g3) = res.device_ptr_mut(s);
12771            let (pq, _g4) = out_q.device_ptr_mut(s);
12772            let (pd, _g5) = out_d.device_ptr_mut(s);
12773            let mut ps = [
12774                &pa as *const _ as *mut std::ffi::c_void,
12775                &pb as *const _ as *mut _,
12776                &c as *const _ as *mut _,
12777                &pw as *const _ as *mut _,
12778                &pr as *const _ as *mut _,
12779                &pq as *const _ as *mut _,
12780                &pd as *const _ as *mut _,
12781                &nc as *const _ as *mut _,
12782                &e2 as *const _ as *mut _,
12783            ];
12784            unsafe {
12785                self.launch_pdl(
12786                    "add_scale_rms_norm_q8_1",
12787                    (nrows as u32, 1, 1),
12788                    (rms_block(), 1, 1),
12789                    &mut ps,
12790                )?;
12791            }
12792            return Ok(());
12793        }
12794        let f = self.func("add_scale_rms_norm_q8_1");
12795        let cfg = LaunchConfig {
12796            grid_dim: (nrows as u32, 1, 1),
12797            block_dim: (rms_block(), 1, 1),
12798            shared_mem_bytes: 0,
12799        };
12800        let __s_b = self.gpu.stream();
12801        let mut b = __s_b.launch_builder(&f);
12802        b.arg(a)
12803            .arg(b_in)
12804            .arg(&c)
12805            .arg(w)
12806            .arg(res)
12807            .arg(&mut *out_q)
12808            .arg(&mut *out_d)
12809            .arg(&nc)
12810            .arg(&e2);
12811        unsafe {
12812            b.launch(cfg)?;
12813        }
12814        Ok(())
12815    }
12816
12817    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
12818    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
12819    #[allow(clippy::too_many_arguments)]
12820    pub fn rms_pre_add_scale_rms_norm_q8_1(
12821        &self,
12822        a: &CudaSlice<f32>,
12823        wa: &CudaSlice<f32>,
12824        b_in: &CudaSlice<f32>,
12825        c: f32,
12826        w: &CudaSlice<f32>,
12827        res: &mut CudaSlice<f32>,
12828        ncols: usize,
12829        nrows: usize,
12830        eps: f32,
12831    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12832        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12833        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12834        let (nc, e2) = (ncols as i32, eps);
12835        if Self::pdl_on() {
12836            {
12837                use cudarc::driver::{DevicePtr, DevicePtrMut};
12838                let s = &self.gpu.stream();
12839                let (pa, _g0) = a.device_ptr(s);
12840                let (pwa, _g1) = wa.device_ptr(s);
12841                let (pb, _g2) = b_in.device_ptr(s);
12842                let (pw, _g3) = w.device_ptr(s);
12843                let (pr, _g4) = res.device_ptr_mut(s);
12844                let (pq, _g5) = out_q.device_ptr_mut(s);
12845                let (pd, _g6) = out_d.device_ptr_mut(s);
12846                let mut ps = [
12847                    &pa as *const _ as *mut std::ffi::c_void,
12848                    &pwa as *const _ as *mut _,
12849                    &pb as *const _ as *mut _,
12850                    &c as *const _ as *mut _,
12851                    &pw as *const _ as *mut _,
12852                    &pr as *const _ as *mut _,
12853                    &pq as *const _ as *mut _,
12854                    &pd as *const _ as *mut _,
12855                    &nc as *const _ as *mut _,
12856                    &e2 as *const _ as *mut _,
12857                ];
12858                unsafe {
12859                    self.launch_pdl(
12860                        "rms_pre_add_scale_rms_norm_q8_1",
12861                        (nrows as u32, 1, 1),
12862                        (rms_block(), 1, 1),
12863                        &mut ps,
12864                    )?;
12865                }
12866            }
12867            return Ok((out_q, out_d));
12868        }
12869        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
12870        let cfg = LaunchConfig {
12871            grid_dim: (nrows as u32, 1, 1),
12872            block_dim: (rms_block(), 1, 1),
12873            shared_mem_bytes: 0,
12874        };
12875        let __s_b = self.gpu.stream();
12876        let mut b = __s_b.launch_builder(&f);
12877        b.arg(a)
12878            .arg(wa)
12879            .arg(b_in)
12880            .arg(&c)
12881            .arg(w)
12882            .arg(res)
12883            .arg(&mut out_q)
12884            .arg(&mut out_d)
12885            .arg(&nc)
12886            .arg(&e2);
12887        unsafe {
12888            b.launch(cfg)?;
12889        }
12890        Ok((out_q, out_d))
12891    }
12892
12893    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
12894    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
12895    pub fn gelu_tanh_mul_q8_1(
12896        &self,
12897        gate: &CudaSlice<f32>,
12898        up: &cudarc::driver::CudaView<f32>,
12899        act: &mut CudaSlice<f32>,
12900        ncols: usize,
12901        nrows: usize,
12902    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12903        debug_assert!(ncols.is_multiple_of(128));
12904        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
12905        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
12906        let nc = ncols as i32;
12907        if Self::pdl_on() {
12908            {
12909                use cudarc::driver::{DevicePtr, DevicePtrMut};
12910                let s = &self.gpu.stream();
12911                let (pg, _g0) = gate.device_ptr(s);
12912                let (pu, _g1) = up.device_ptr(s);
12913                let (pact, _g2) = act.device_ptr_mut(s);
12914                let (pq, _g3) = out_q.device_ptr_mut(s);
12915                let (pd, _g4) = out_d.device_ptr_mut(s);
12916                let mut ps = [
12917                    &pg as *const _ as *mut std::ffi::c_void,
12918                    &pu as *const _ as *mut _,
12919                    &pact as *const _ as *mut _,
12920                    &pq as *const _ as *mut _,
12921                    &pd as *const _ as *mut _,
12922                    &nc as *const _ as *mut _,
12923                ];
12924                unsafe {
12925                    self.launch_pdl(
12926                        "gelu_tanh_mul_q8_1",
12927                        (nrows as u32, 1, 1),
12928                        (rms_block(), 1, 1),
12929                        &mut ps,
12930                    )?;
12931                }
12932            }
12933            return Ok((out_q, out_d));
12934        }
12935        let f = self.func("gelu_tanh_mul_q8_1");
12936        let cfg = LaunchConfig {
12937            grid_dim: (nrows as u32, 1, 1),
12938            block_dim: (rms_block(), 1, 1),
12939            shared_mem_bytes: 0,
12940        };
12941        let __s_b = self.gpu.stream();
12942        let mut b = __s_b.launch_builder(&f);
12943        b.arg(gate)
12944            .arg(up)
12945            .arg(act)
12946            .arg(&mut out_q)
12947            .arg(&mut out_d)
12948            .arg(&nc);
12949        unsafe {
12950            b.launch(cfg)?;
12951        }
12952        Ok((out_q, out_d))
12953    }
12954
12955    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
12956    #[allow(clippy::too_many_arguments)]
12957    pub fn gelu_tanh_mul_q8_1_into(
12958        &self,
12959        gate: &CudaSlice<f32>,
12960        up: &cudarc::driver::CudaView<f32>,
12961        act: &mut CudaSlice<f32>,
12962        ncols: usize,
12963        nrows: usize,
12964        out_q: &mut CudaSlice<i8>,
12965        out_d: &mut CudaSlice<f32>,
12966    ) -> Result<(), Box<dyn std::error::Error>> {
12967        debug_assert!(ncols.is_multiple_of(128));
12968        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
12969        let nc = ncols as i32;
12970        if Self::pdl_on() {
12971            use cudarc::driver::{DevicePtr, DevicePtrMut};
12972            let s = &self.gpu.stream();
12973            let (pg, _g0) = gate.device_ptr(s);
12974            let (pu, _g1) = up.device_ptr(s);
12975            let (pact, _g2) = act.device_ptr_mut(s);
12976            let (pq, _g3) = out_q.device_ptr_mut(s);
12977            let (pd, _g4) = out_d.device_ptr_mut(s);
12978            let mut ps = [
12979                &pg as *const _ as *mut std::ffi::c_void,
12980                &pu as *const _ as *mut _,
12981                &pact as *const _ as *mut _,
12982                &pq as *const _ as *mut _,
12983                &pd as *const _ as *mut _,
12984                &nc as *const _ as *mut _,
12985            ];
12986            unsafe {
12987                self.launch_pdl(
12988                    "gelu_tanh_mul_q8_1",
12989                    (nrows as u32, 1, 1),
12990                    (rms_block(), 1, 1),
12991                    &mut ps,
12992                )?;
12993            }
12994            return Ok(());
12995        }
12996        let f = self.func("gelu_tanh_mul_q8_1");
12997        let cfg = LaunchConfig {
12998            grid_dim: (nrows as u32, 1, 1),
12999            block_dim: (rms_block(), 1, 1),
13000            shared_mem_bytes: 0,
13001        };
13002        let __s_b = self.gpu.stream();
13003        let mut b = __s_b.launch_builder(&f);
13004        b.arg(gate)
13005            .arg(up)
13006            .arg(&mut *act)
13007            .arg(&mut *out_q)
13008            .arg(&mut *out_d)
13009            .arg(&nc);
13010        unsafe {
13011            b.launch(cfg)?;
13012        }
13013        Ok(())
13014    }
13015
13016    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
13017    #[allow(clippy::too_many_arguments)]
13018    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
13019    pub fn add_rms_norm3_q8z(
13020        &self,
13021        a: &CudaSlice<f32>,
13022        b_in: &CudaSlice<f32>,
13023        w0: &CudaSlice<f32>,
13024        w1: &CudaSlice<f32>,
13025        w2: &CudaSlice<f32>,
13026        res: &mut CudaSlice<f32>,
13027        out1: &mut CudaSlice<f32>,
13028        ncols: usize,
13029        nrows: usize,
13030        eps: f32,
13031    ) -> Result<
13032        (
13033            (CudaSlice<i8>, CudaSlice<f32>),
13034            (CudaSlice<i8>, CudaSlice<f32>),
13035        ),
13036        Box<dyn std::error::Error>,
13037    > {
13038        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
13039        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13040        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
13041        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13042        let f = self.func("add_rms_norm3_q8z_f32");
13043        let cfg = LaunchConfig {
13044            grid_dim: (nrows as u32, 1, 1),
13045            block_dim: (rms_block(), 1, 1),
13046            shared_mem_bytes: 0,
13047        };
13048        let (nc, e2) = (ncols as i32, eps);
13049        let __s_b = self.gpu.stream();
13050        let mut b = __s_b.launch_builder(&f);
13051        b.arg(a)
13052            .arg(b_in)
13053            .arg(w0)
13054            .arg(w1)
13055            .arg(w2)
13056            .arg(res)
13057            .arg(&mut q0)
13058            .arg(&mut d0)
13059            .arg(out1)
13060            .arg(&mut q2)
13061            .arg(&mut d2)
13062            .arg(&nc)
13063            .arg(&e2);
13064        unsafe {
13065            b.launch(cfg)?;
13066        }
13067        Ok(((q0, d0), (q2, d2)))
13068    }
13069
13070    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
13071    #[allow(clippy::too_many_arguments)]
13072    pub fn add_rms_norm3(
13073        &self,
13074        a: &CudaSlice<f32>,
13075        b_in: &CudaSlice<f32>,
13076        w0: &CudaSlice<f32>,
13077        w1: &CudaSlice<f32>,
13078        w2: &CudaSlice<f32>,
13079        res: &mut CudaSlice<f32>,
13080        d0: &mut CudaSlice<f32>,
13081        d1: &mut CudaSlice<f32>,
13082        d2: &mut CudaSlice<f32>,
13083        ncols: usize,
13084        nrows: usize,
13085        eps: f32,
13086    ) -> Result<(), Box<dyn std::error::Error>> {
13087        let f = self.func("add_rms_norm3_f32");
13088        let cfg = LaunchConfig {
13089            grid_dim: (nrows as u32, 1, 1),
13090            block_dim: (rms_block(), 1, 1),
13091            shared_mem_bytes: 0,
13092        };
13093        let (nc, e2) = (ncols as i32, eps);
13094        let __s_b = self.gpu.stream();
13095        let mut b = __s_b.launch_builder(&f);
13096        b.arg(a)
13097            .arg(b_in)
13098            .arg(w0)
13099            .arg(w1)
13100            .arg(w2)
13101            .arg(res)
13102            .arg(d0)
13103            .arg(d1)
13104            .arg(d2)
13105            .arg(&nc)
13106            .arg(&e2);
13107        unsafe {
13108            b.launch(cfg)?;
13109        }
13110        Ok(())
13111    }
13112
13113    /// dst = (a + b) * c (residual add + layer scale, one launch).
13114    pub fn add_scale(
13115        &self,
13116        a: &CudaSlice<f32>,
13117        b_in: &CudaSlice<f32>,
13118        c: f32,
13119        dst: &mut CudaSlice<f32>,
13120        n: usize,
13121    ) -> Result<(), Box<dyn std::error::Error>> {
13122        let f = self.func("add_scale_f32");
13123        let cfg = LaunchConfig::for_num_elems(n as u32);
13124        let ni = n as i32;
13125        let __s_b = self.gpu.stream();
13126        let mut b = __s_b.launch_builder(&f);
13127        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
13128        unsafe {
13129            b.launch(cfg)?;
13130        }
13131        Ok(())
13132    }
13133
13134    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
13135    #[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
13136    pub fn layer_norm_bias(
13137        &self,
13138        x: &CudaSlice<f32>,
13139        w: &CudaSlice<f32>,
13140        b: &CudaSlice<f32>,
13141        dst: &mut CudaSlice<f32>,
13142        ncols: usize,
13143        nrows: usize,
13144        eps: f32,
13145    ) -> Result<(), Box<dyn std::error::Error>> {
13146        let f = self.func("layer_norm_bias_f32");
13147        let (nc, e) = (ncols as i32, eps);
13148        let cfg = LaunchConfig {
13149            grid_dim: (nrows as u32, 1, 1),
13150            block_dim: (256, 1, 1),
13151            shared_mem_bytes: 0,
13152        };
13153        let __s_b = self.gpu.stream();
13154        let mut lb = __s_b.launch_builder(&f);
13155        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
13156        unsafe {
13157            lb.launch(cfg)?;
13158        }
13159        Ok(())
13160    }
13161
13162    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
13163    pub fn gelu_tanh(
13164        &self,
13165        x: &CudaSlice<f32>,
13166        dst: &mut CudaSlice<f32>,
13167        n: usize,
13168    ) -> Result<(), Box<dyn std::error::Error>> {
13169        let f = self.func("gelu_tanh_f32");
13170        let ni = n as i64;
13171        let cfg = LaunchConfig {
13172            grid_dim: (n.div_ceil(256) as u32, 1, 1),
13173            block_dim: (256, 1, 1),
13174            shared_mem_bytes: 0,
13175        };
13176        let __s_b = self.gpu.stream();
13177        let mut lb = __s_b.launch_builder(&f);
13178        lb.arg(x).arg(&mut *dst).arg(&ni);
13179        unsafe {
13180            lb.launch(cfg)?;
13181        }
13182        Ok(())
13183    }
13184
13185    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
13186    pub fn row_softmax(
13187        &self,
13188        x: &mut CudaSlice<f32>,
13189        ncols: usize,
13190        nrows: usize,
13191    ) -> Result<(), Box<dyn std::error::Error>> {
13192        let f = self.func("row_softmax_f32");
13193        let nc = ncols as i32;
13194        let cfg = LaunchConfig {
13195            grid_dim: (nrows as u32, 1, 1),
13196            block_dim: (256, 1, 1),
13197            shared_mem_bytes: 0,
13198        };
13199        let __s_b = self.gpu.stream();
13200        let mut lb = __s_b.launch_builder(&f);
13201        lb.arg(&mut *x).arg(&nc);
13202        unsafe {
13203            lb.launch(cfg)?;
13204        }
13205        Ok(())
13206    }
13207
13208    pub fn rms_norm(
13209        &self,
13210        x: &CudaSlice<f32>,
13211        w: &CudaSlice<f32>,
13212        dst: &mut CudaSlice<f32>,
13213        ncols: usize,
13214        nrows: usize,
13215        eps: f32,
13216    ) -> Result<(), Box<dyn std::error::Error>> {
13217        let (nc, e) = (ncols as i32, eps);
13218        let kname = if Self::norm_ilp_on() {
13219            "rms_norm_f32_v2"
13220        } else {
13221            "rms_norm_f32"
13222        };
13223        if Self::pdl_on() && Self::pdl_wb_on() {
13224            use cudarc::driver::{DevicePtr, DevicePtrMut};
13225            let s = &self.gpu.stream();
13226            let (px, _g0) = x.device_ptr(s);
13227            let (pw, _g1) = w.device_ptr(s);
13228            let (pd, _g2) = dst.device_ptr_mut(s);
13229            let mut ps = [
13230                &px as *const _ as *mut std::ffi::c_void,
13231                &pw as *const _ as *mut _,
13232                &pd as *const _ as *mut _,
13233                &nc as *const _ as *mut _,
13234                &e as *const _ as *mut _,
13235            ];
13236            unsafe {
13237                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
13238            }
13239            return Ok(());
13240        }
13241        let f = self.func(kname);
13242        let cfg = LaunchConfig {
13243            grid_dim: (nrows as u32, 1, 1),
13244            block_dim: (rms_block(), 1, 1),
13245            shared_mem_bytes: 0,
13246        };
13247        let __s_b = self.gpu.stream();
13248        let mut b = __s_b.launch_builder(&f);
13249        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
13250        unsafe {
13251            b.launch(cfg)?;
13252        }
13253        Ok(())
13254    }
13255
13256    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
13257    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
13258    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
13259    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
13260    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
13261    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
13262    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
13263    pub fn rms_norm_decode(
13264        &self,
13265        x: &CudaSlice<f32>,
13266        w: &CudaSlice<f32>,
13267        dst: &mut CudaSlice<f32>,
13268        ncols: usize,
13269        nrows: usize,
13270        eps: f32,
13271    ) -> Result<(), Box<dyn std::error::Error>> {
13272        let f = self.func(if Self::norm_ilp_on() {
13273            "rms_norm_f32_v2"
13274        } else {
13275            "rms_norm_f32"
13276        });
13277        let cfg = LaunchConfig {
13278            grid_dim: (nrows as u32, 1, 1),
13279            block_dim: (1024, 1, 1),
13280            shared_mem_bytes: 0,
13281        };
13282        let (nc, e) = (ncols as i32, eps);
13283        let __s_b = self.gpu.stream();
13284        let mut b = __s_b.launch_builder(&f);
13285        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
13286        unsafe {
13287            b.launch(cfg)?;
13288        }
13289        Ok(())
13290    }
13291
13292    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
13293    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
13294    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
13295    pub fn rms_norm_q8_1(
13296        &self,
13297        x: &CudaSlice<f32>,
13298        w: &CudaSlice<f32>,
13299        ncols: usize,
13300        nrows: usize,
13301        eps: f32,
13302    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13303        let nblk = ncols / 32;
13304        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
13305        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
13306        let (nc, e) = (ncols as i32, eps);
13307        if Self::pdl_on() {
13308            {
13309                use cudarc::driver::{DevicePtr, DevicePtrMut};
13310                let s = &self.gpu.stream();
13311                let (px, _g0) = x.device_ptr(s);
13312                let (pw, _g1) = w.device_ptr(s);
13313                let (pq, _g2) = q.device_ptr_mut(s);
13314                let (pd, _g3) = d.device_ptr_mut(s);
13315                let mut ps = [
13316                    &px as *const _ as *mut std::ffi::c_void,
13317                    &pw as *const _ as *mut _,
13318                    &pq as *const _ as *mut _,
13319                    &pd as *const _ as *mut _,
13320                    &nc as *const _ as *mut _,
13321                    &e as *const _ as *mut _,
13322                ];
13323                unsafe {
13324                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
13325                }
13326            }
13327            return Ok((q, d));
13328        }
13329        let f = self.func("rms_norm_q8_1");
13330        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
13331        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
13332        let cfg = LaunchConfig {
13333            grid_dim: (nrows as u32, 1, 1),
13334            block_dim: (1024, 1, 1),
13335            shared_mem_bytes: 0,
13336        };
13337        let __s_b = self.gpu.stream();
13338        let mut b = __s_b.launch_builder(&f);
13339        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
13340        unsafe {
13341            b.launch(cfg)?;
13342        }
13343        Ok((q, d))
13344    }
13345
13346    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
13347    /// PDL arm), caller-owned outputs.
13348    #[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
13349    pub fn rms_norm_q8_1_into(
13350        &self,
13351        x: &CudaSlice<f32>,
13352        w: &CudaSlice<f32>,
13353        ncols: usize,
13354        nrows: usize,
13355        eps: f32,
13356        q: &mut CudaSlice<i8>,
13357        d: &mut CudaSlice<f32>,
13358    ) -> Result<(), Box<dyn std::error::Error>> {
13359        let nblk = ncols / 32;
13360        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
13361        let (nc, e) = (ncols as i32, eps);
13362        if Self::pdl_on() {
13363            use cudarc::driver::{DevicePtr, DevicePtrMut};
13364            let s = &self.gpu.stream();
13365            let (px, _g0) = x.device_ptr(s);
13366            let (pw, _g1) = w.device_ptr(s);
13367            let (pq, _g2) = q.device_ptr_mut(s);
13368            let (pd, _g3) = d.device_ptr_mut(s);
13369            let mut ps = [
13370                &px as *const _ as *mut std::ffi::c_void,
13371                &pw as *const _ as *mut _,
13372                &pq as *const _ as *mut _,
13373                &pd as *const _ as *mut _,
13374                &nc as *const _ as *mut _,
13375                &e as *const _ as *mut _,
13376            ];
13377            unsafe {
13378                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
13379            }
13380            return Ok(());
13381        }
13382        let f = self.func("rms_norm_q8_1");
13383        let cfg = LaunchConfig {
13384            grid_dim: (nrows as u32, 1, 1),
13385            block_dim: (1024, 1, 1),
13386            shared_mem_bytes: 0,
13387        };
13388        let __s_b = self.gpu.stream();
13389        let mut b = __s_b.launch_builder(&f);
13390        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
13391        unsafe {
13392            b.launch(cfg)?;
13393        }
13394        Ok(())
13395    }
13396
13397    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
13398    pub fn quantize_q8_1_into(
13399        &self,
13400        x: &CudaSlice<f32>,
13401        m: usize,
13402        in_f: usize,
13403        q: &mut CudaSlice<i8>,
13404        d: &mut CudaSlice<f32>,
13405    ) -> Result<(), Box<dyn std::error::Error>> {
13406        let nblk = in_f / 32;
13407        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
13408        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
13409        let (inf, mi) = (in_f as i32, m as i32);
13410        if Self::pdl_on() && Self::pdl_wb_on() {
13411            use cudarc::driver::{DevicePtr, DevicePtrMut};
13412            let s = &self.gpu.stream();
13413            let (px, _g0) = x.device_ptr(s);
13414            let (pq, _g1) = q.device_ptr_mut(s);
13415            let (pd, _g2) = d.device_ptr_mut(s);
13416            let mut ps = [
13417                &px as *const _ as *mut std::ffi::c_void,
13418                &pq as *const _ as *mut _,
13419                &pd as *const _ as *mut _,
13420                &inf as *const _ as *mut _,
13421                &mi as *const _ as *mut _,
13422            ];
13423            unsafe {
13424                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
13425            }
13426            return Ok(());
13427        }
13428        let f = self.func("quantize_q8_1");
13429        let __s_b = self.gpu.stream();
13430        let mut b = __s_b.launch_builder(&f);
13431        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
13432        unsafe {
13433            b.launch(cfg)?;
13434        }
13435        Ok(())
13436    }
13437
13438    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
13439    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
13440    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
13441    #[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
13442    pub fn add_rms_norm_q8_1(
13443        &self,
13444        a: &CudaSlice<f32>,
13445        b_in: &CudaSlice<f32>,
13446        w: &CudaSlice<f32>,
13447        res: &mut CudaSlice<f32>,
13448        ncols: usize,
13449        nrows: usize,
13450        eps: f32,
13451    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13452        let nblk = ncols / 32;
13453        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
13454        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
13455        let f = self.func("add_rms_norm_q8_1");
13456        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
13457        let cfg = LaunchConfig {
13458            grid_dim: (nrows as u32, 1, 1),
13459            block_dim: (1024, 1, 1),
13460            shared_mem_bytes: 0,
13461        };
13462        let (nc, e) = (ncols as i32, eps);
13463        let __s_bld = self.gpu.stream();
13464        let mut bld = __s_bld.launch_builder(&f);
13465        bld.arg(a)
13466            .arg(b_in)
13467            .arg(w)
13468            .arg(res)
13469            .arg(&mut q)
13470            .arg(&mut d)
13471            .arg(&nc)
13472            .arg(&e);
13473        unsafe {
13474            bld.launch(cfg)?;
13475        }
13476        Ok((q, d))
13477    }
13478
13479    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
13480    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
13481    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
13482    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
13483    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
13484    #[allow(clippy::too_many_arguments)]
13485    pub fn join_add_rms_norm_raw(
13486        &self,
13487        a0_raw: u64,
13488        a1_raw: u64,
13489        x: &CudaSlice<f32>,
13490        w: &CudaSlice<f32>,
13491        res: &mut CudaSlice<f32>,
13492        dst: &mut CudaSlice<f32>,
13493        ncols: usize,
13494        eps: f32,
13495    ) -> Result<(), Box<dyn std::error::Error>> {
13496        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
13497            return Err("join_add_rms_norm geometry".into());
13498        }
13499        let f = self.func("join_add_rms_norm_f32");
13500        let cfg = LaunchConfig {
13501            grid_dim: (1, 1, 1),
13502            block_dim: (rms_block(), 1, 1),
13503            shared_mem_bytes: 0,
13504        };
13505        let (nc, e) = (ncols as i32, eps);
13506        let __s_b = self.gpu.stream();
13507        let mut b = __s_b.launch_builder(&f);
13508        b.arg(&a0_raw)
13509            .arg(&a1_raw)
13510            .arg(x)
13511            .arg(w)
13512            .arg(&mut *res)
13513            .arg(&mut *dst)
13514            .arg(&nc)
13515            .arg(&e);
13516        unsafe {
13517            b.launch(cfg)?;
13518        }
13519        Ok(())
13520    }
13521
13522    #[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
13523    pub fn add_rms_norm(
13524        &self,
13525        a: &CudaSlice<f32>,
13526        b: &CudaSlice<f32>,
13527        w: &CudaSlice<f32>,
13528        res: &mut CudaSlice<f32>,
13529        dst: &mut CudaSlice<f32>,
13530        ncols: usize,
13531        nrows: usize,
13532        eps: f32,
13533    ) -> Result<(), Box<dyn std::error::Error>> {
13534        let (nc, e) = (ncols as i32, eps);
13535        let kname = if Self::norm_ilp_on() {
13536            "add_rms_norm_f32_v2"
13537        } else {
13538            "add_rms_norm_f32"
13539        };
13540        if Self::pdl_on() && Self::pdl_wb_on() {
13541            use cudarc::driver::{DevicePtr, DevicePtrMut};
13542            let s = &self.gpu.stream();
13543            let (pa, _g0) = a.device_ptr(s);
13544            let (pb, _g1) = b.device_ptr(s);
13545            let (pw, _g2) = w.device_ptr(s);
13546            let (pr, _g3) = res.device_ptr_mut(s);
13547            let (pd, _g4) = dst.device_ptr_mut(s);
13548            let mut ps = [
13549                &pa as *const _ as *mut std::ffi::c_void,
13550                &pb as *const _ as *mut _,
13551                &pw as *const _ as *mut _,
13552                &pr as *const _ as *mut _,
13553                &pd as *const _ as *mut _,
13554                &nc as *const _ as *mut _,
13555                &e as *const _ as *mut _,
13556            ];
13557            unsafe {
13558                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
13559            }
13560            return Ok(());
13561        }
13562        let f = self.func(kname);
13563        let cfg = LaunchConfig {
13564            grid_dim: (nrows as u32, 1, 1),
13565            block_dim: (rms_block(), 1, 1),
13566            shared_mem_bytes: 0,
13567        };
13568        let __s_b2 = self.gpu.stream();
13569        let mut b2 = __s_b2.launch_builder(&f);
13570        b2.arg(a)
13571            .arg(b)
13572            .arg(w)
13573            .arg(&mut *res)
13574            .arg(&mut *dst)
13575            .arg(&nc)
13576            .arg(&e);
13577        unsafe {
13578            b2.launch(cfg)?;
13579        }
13580        Ok(())
13581    }
13582
13583    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
13584    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
13585    #[allow(clippy::too_many_arguments)]
13586    pub fn rms_pre_add_rms_norm(
13587        &self,
13588        a: &CudaSlice<f32>,
13589        wa: &CudaSlice<f32>,
13590        b: &CudaSlice<f32>,
13591        w: &CudaSlice<f32>,
13592        res: &mut CudaSlice<f32>,
13593        dst: &mut CudaSlice<f32>,
13594        ncols: usize,
13595        nrows: usize,
13596        eps: f32,
13597    ) -> Result<(), Box<dyn std::error::Error>> {
13598        let f = self.func("rms_pre_add_rms_norm_f32");
13599        let cfg = LaunchConfig {
13600            grid_dim: (nrows as u32, 1, 1),
13601            block_dim: (rms_block(), 1, 1),
13602            shared_mem_bytes: 0,
13603        };
13604        let (nc, e) = (ncols as i32, eps);
13605        let __s_b2 = self.gpu.stream();
13606        let mut b2 = __s_b2.launch_builder(&f);
13607        b2.arg(a)
13608            .arg(wa)
13609            .arg(b)
13610            .arg(w)
13611            .arg(&mut *res)
13612            .arg(&mut *dst)
13613            .arg(&nc)
13614            .arg(&e);
13615        unsafe {
13616            b2.launch(cfg)?;
13617        }
13618        Ok(())
13619    }
13620
13621    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
13622    #[allow(clippy::too_many_arguments)]
13623    pub fn rms_pre_add_rms_norm_q8z(
13624        &self,
13625        a: &CudaSlice<f32>,
13626        wa: &CudaSlice<f32>,
13627        b: &CudaSlice<f32>,
13628        w: &CudaSlice<f32>,
13629        res: &mut CudaSlice<f32>,
13630        dst: &mut CudaSlice<f32>,
13631        ncols: usize,
13632        nrows: usize,
13633        eps: f32,
13634    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13635        debug_assert!(ncols.is_multiple_of(128));
13636        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
13637        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13638        let (nc, e) = (ncols as i32, eps);
13639        if Self::pdl_on() {
13640            {
13641                use cudarc::driver::{DevicePtr, DevicePtrMut};
13642                let s = &self.gpu.stream();
13643                let (pa, _g0) = a.device_ptr(s);
13644                let (pwa, _g1) = wa.device_ptr(s);
13645                let (pb, _g2) = b.device_ptr(s);
13646                let (pw, _g3) = w.device_ptr(s);
13647                let (pr, _g4) = res.device_ptr_mut(s);
13648                let (pdst, _g5) = dst.device_ptr_mut(s);
13649                let (pq, _g6) = out_q.device_ptr_mut(s);
13650                let (pd, _g7) = out_d.device_ptr_mut(s);
13651                let mut ps = [
13652                    &pa as *const _ as *mut std::ffi::c_void,
13653                    &pwa as *const _ as *mut _,
13654                    &pb as *const _ as *mut _,
13655                    &pw as *const _ as *mut _,
13656                    &pr as *const _ as *mut _,
13657                    &pdst as *const _ as *mut _,
13658                    &pq as *const _ as *mut _,
13659                    &pd as *const _ as *mut _,
13660                    &nc as *const _ as *mut _,
13661                    &e as *const _ as *mut _,
13662                ];
13663                unsafe {
13664                    self.launch_pdl(
13665                        "rms_pre_add_rms_norm_q8z_f32",
13666                        (nrows as u32, 1, 1),
13667                        (rms_block(), 1, 1),
13668                        &mut ps,
13669                    )?;
13670                }
13671            }
13672            return Ok((out_q, out_d));
13673        }
13674        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
13675        let cfg = LaunchConfig {
13676            grid_dim: (nrows as u32, 1, 1),
13677            block_dim: (rms_block(), 1, 1),
13678            shared_mem_bytes: 0,
13679        };
13680        let __s_b2 = self.gpu.stream();
13681        let mut b2 = __s_b2.launch_builder(&f);
13682        b2.arg(a)
13683            .arg(wa)
13684            .arg(b)
13685            .arg(w)
13686            .arg(&mut *res)
13687            .arg(&mut *dst)
13688            .arg(&mut out_q)
13689            .arg(&mut out_d)
13690            .arg(&nc)
13691            .arg(&e);
13692        unsafe {
13693            b2.launch(cfg)?;
13694        }
13695        Ok((out_q, out_d))
13696    }
13697
13698    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
13699    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
13700    /// body must stay attribute-free (the fused2_into precedent).
13701    #[allow(clippy::too_many_arguments)]
13702    pub fn rms_pre_add_rms_norm_q8z_into(
13703        &self,
13704        a: &CudaSlice<f32>,
13705        wa: &CudaSlice<f32>,
13706        b: &CudaSlice<f32>,
13707        w: &CudaSlice<f32>,
13708        res: &mut CudaSlice<f32>,
13709        dst: &mut CudaSlice<f32>,
13710        ncols: usize,
13711        nrows: usize,
13712        eps: f32,
13713        out_q: &mut CudaSlice<i8>,
13714        out_d: &mut CudaSlice<f32>,
13715    ) -> Result<(), Box<dyn std::error::Error>> {
13716        debug_assert!(ncols.is_multiple_of(128));
13717        let (nc, e) = (ncols as i32, eps);
13718        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
13719        let cfg = LaunchConfig {
13720            grid_dim: (nrows as u32, 1, 1),
13721            block_dim: (rms_block(), 1, 1),
13722            shared_mem_bytes: 0,
13723        };
13724        let __s_b = self.gpu.stream();
13725        let mut b2 = __s_b.launch_builder(&f);
13726        b2.arg(a)
13727            .arg(wa)
13728            .arg(b)
13729            .arg(w)
13730            .arg(&mut *res)
13731            .arg(&mut *dst)
13732            .arg(&mut *out_q)
13733            .arg(&mut *out_d)
13734            .arg(&nc)
13735            .arg(&e);
13736        unsafe {
13737            b2.launch(cfg)?;
13738        }
13739        Ok(())
13740    }
13741
13742    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
13743    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
13744    #[allow(clippy::too_many_arguments)]
13745    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
13746        &self,
13747        a: &CudaSlice<f32>,
13748        wa: &CudaSlice<f32>,
13749        b_in: &CudaSlice<f32>,
13750        c: f32,
13751        w: &CudaSlice<f32>,
13752        res: &mut CudaSlice<f32>,
13753        ncols: usize,
13754        nrows: usize,
13755        eps: f32,
13756        out_q: &mut CudaSlice<i8>,
13757        out_d: &mut CudaSlice<f32>,
13758    ) -> Result<(), Box<dyn std::error::Error>> {
13759        debug_assert!(ncols.is_multiple_of(128));
13760        let (nc, e2) = (ncols as i32, eps);
13761        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
13762        let cfg = LaunchConfig {
13763            grid_dim: (nrows as u32, 1, 1),
13764            block_dim: (rms_block(), 1, 1),
13765            shared_mem_bytes: 0,
13766        };
13767        let __s_b = self.gpu.stream();
13768        let mut b2 = __s_b.launch_builder(&f);
13769        b2.arg(a)
13770            .arg(wa)
13771            .arg(b_in)
13772            .arg(&c)
13773            .arg(w)
13774            .arg(&mut *res)
13775            .arg(&mut *out_q)
13776            .arg(&mut *out_d)
13777            .arg(&nc)
13778            .arg(&e2);
13779        unsafe {
13780            b2.launch(cfg)?;
13781        }
13782        Ok(())
13783    }
13784
13785    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
13786    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
13787    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
13788    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
13789    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
13790    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
13791    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
13792    pub fn g4_pnfold_on() -> bool {
13793        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13794        *ON.get_or_init(|| {
13795            std::env::var("MEMRA_G4_PNFOLD")
13796                .map(|v| v != "0")
13797                .unwrap_or(true)
13798        })
13799    }
13800
13801    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
13802    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
13803    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
13804    pub fn build_q4_out_concat3(
13805        &self,
13806        w0: &crate::model::GpuTensor,
13807        w1: &crate::model::GpuTensor,
13808        w2: &crate::model::GpuTensor,
13809    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
13810        use crate::model::GpuTensor;
13811        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
13812            match w {
13813                GpuTensor::Quant {
13814                    qtype,
13815                    row_bytes,
13816                    rp,
13817                    ..
13818                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
13819                _ => None,
13820            }
13821        };
13822        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
13823        else {
13824            return Ok(None);
13825        };
13826        if rb0 != rb1
13827            || rb0 != rb2
13828            || w0.in_features() != w1.in_features()
13829            || w0.in_features() != w2.in_features()
13830        {
13831            return Ok(None);
13832        }
13833        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
13834            match w {
13835                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
13836                _ => unreachable!(),
13837            }
13838        }
13839        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
13840        let total = rb0 * (o0 + o1 + o2);
13841        let mut cat = self.alloc_u8(total)?;
13842        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
13843        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
13844        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
13845        Ok(Some(GpuTensor::Quant {
13846            bytes: cat,
13847            qtype: QT_Q4_0,
13848            row_bytes: rb0,
13849            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
13850            scale: 1.0,
13851            rp: false,
13852            #[cfg(memra_cutlass)]
13853            cutlass: None,
13854            fp8: None,
13855            blk: None,
13856            rp4: None,
13857            f16: None,
13858        }))
13859    }
13860
13861    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
13862    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
13863    ///
13864    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
13865    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
13866    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
13867    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
13868    ///
13869    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
13870    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
13871    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
13872    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
13873    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
13874    ///
13875    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
13876    /// width. A future partial-rotary caller fails at its first launch with the geometry named
13877    /// instead of serving quietly wrong logits.
13878    fn full_width_rope_only(
13879        kernel: &str,
13880        n_rot: usize,
13881        head_dim: usize,
13882    ) -> Result<(), Box<dyn std::error::Error>> {
13883        if n_rot == head_dim {
13884            return Ok(());
13885        }
13886        Err(format!(
13887            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
13888             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
13889             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
13890             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
13891             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
13892        )
13893        .into())
13894    }
13895
13896    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
13897    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
13898    /// ([`Engine::full_width_rope_only`]).
13899    #[allow(clippy::too_many_arguments)]
13900    pub fn rms_norm_qkv_rope_cat(
13901        &self,
13902        qkv: &CudaSlice<f32>,
13903        wq: &CudaSlice<f32>,
13904        wk: &CudaSlice<f32>,
13905        wv: &CudaSlice<f32>,
13906        q: &mut CudaSlice<f32>,
13907        k: &mut CudaSlice<f32>,
13908        v: &mut CudaSlice<f32>,
13909        head_dim: usize,
13910        n_rot: usize,
13911        rq: usize,
13912        rk: usize,
13913        pos: &CudaSlice<i32>,
13914        nh_q: usize,
13915        nh_k: usize,
13916        base: f32,
13917        freq_scale: f32,
13918        ff: Option<&CudaSlice<f32>>,
13919        eps: f32,
13920    ) -> Result<(), Box<dyn std::error::Error>> {
13921        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
13922        let rows = rq + rk + rk;
13923        let theta_scale = base.powf(-2.0 / head_dim as f32);
13924        let (nc, rqi, rki, nhq, nhk) = (
13925            head_dim as i32,
13926            rq as i32,
13927            rk as i32,
13928            nh_q as i32,
13929            nh_k as i32,
13930        );
13931        if Self::pdl_on() {
13932            use cudarc::driver::{DevicePtr, DevicePtrMut};
13933            let s = &self.gpu.stream();
13934            let (pqkv, _g0) = qkv.device_ptr(s);
13935            let (pwq, _g1) = wq.device_ptr(s);
13936            let (pwk, _g2) = wk.device_ptr(s);
13937            let (pwv, _g3) = wv.device_ptr(s);
13938            let (pq, _g4) = q.device_ptr_mut(s);
13939            let (pk, _g5) = k.device_ptr_mut(s);
13940            let (pv, _g6) = v.device_ptr_mut(s);
13941            let (ppos, _g7) = pos.device_ptr(s);
13942            let (pff, _g8) = match ff {
13943                Some(t) => {
13944                    let (p, g) = t.device_ptr(s);
13945                    (p, Some(g))
13946                }
13947                None => (0, None),
13948            };
13949            let mut ps = [
13950                &pqkv as *const _ as *mut std::ffi::c_void,
13951                &pwq as *const _ as *mut _,
13952                &pwk as *const _ as *mut _,
13953                &pwv as *const _ as *mut _,
13954                &pq as *const _ as *mut _,
13955                &pk as *const _ as *mut _,
13956                &pv as *const _ as *mut _,
13957                &nc as *const _ as *mut _,
13958                &rqi as *const _ as *mut _,
13959                &rki as *const _ as *mut _,
13960                &ppos as *const _ as *mut _,
13961                &nhq as *const _ as *mut _,
13962                &nhk as *const _ as *mut _,
13963                &theta_scale as *const _ as *mut _,
13964                &freq_scale as *const _ as *mut _,
13965                &pff as *const _ as *mut _,
13966                &eps as *const _ as *mut _,
13967            ];
13968            unsafe {
13969                self.launch_pdl(
13970                    "rms_norm_qkv_rope_cat_f32",
13971                    (rows as u32, 1, 1),
13972                    (rms_block(), 1, 1),
13973                    &mut ps,
13974                )?;
13975            }
13976            return Ok(());
13977        }
13978        let f = self.func("rms_norm_qkv_rope_cat_f32");
13979        let cfg = LaunchConfig {
13980            grid_dim: (rows as u32, 1, 1),
13981            block_dim: (rms_block(), 1, 1),
13982            shared_mem_bytes: 0,
13983        };
13984        let __s_b = self.gpu.stream();
13985        let mut b = __s_b.launch_builder(&f);
13986        match ff {
13987            Some(t) => {
13988                b.arg(qkv)
13989                    .arg(wq)
13990                    .arg(wk)
13991                    .arg(wv)
13992                    .arg(&mut *q)
13993                    .arg(&mut *k)
13994                    .arg(&mut *v)
13995                    .arg(&nc)
13996                    .arg(&rqi)
13997                    .arg(&rki)
13998                    .arg(pos)
13999                    .arg(&nhq)
14000                    .arg(&nhk)
14001                    .arg(&theta_scale)
14002                    .arg(&freq_scale)
14003                    .arg(t)
14004                    .arg(&eps);
14005                unsafe {
14006                    b.launch(cfg)?;
14007                }
14008            }
14009            None => {
14010                let null: u64 = 0;
14011                b.arg(qkv)
14012                    .arg(wq)
14013                    .arg(wk)
14014                    .arg(wv)
14015                    .arg(&mut *q)
14016                    .arg(&mut *k)
14017                    .arg(&mut *v)
14018                    .arg(&nc)
14019                    .arg(&rqi)
14020                    .arg(&rki)
14021                    .arg(pos)
14022                    .arg(&nhq)
14023                    .arg(&nhk)
14024                    .arg(&theta_scale)
14025                    .arg(&freq_scale)
14026                    .arg(&null)
14027                    .arg(&eps);
14028                unsafe {
14029                    b.launch(cfg)?;
14030                }
14031            }
14032        }
14033        Ok(())
14034    }
14035
14036    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
14037    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
14038    /// ([`Engine::full_width_rope_only`]).
14039    #[allow(clippy::too_many_arguments)]
14040    pub fn rms_norm_qkv_rope(
14041        &self,
14042        q0: &CudaSlice<f32>,
14043        k0: &CudaSlice<f32>,
14044        v0: &CudaSlice<f32>,
14045        wq: &CudaSlice<f32>,
14046        wk: &CudaSlice<f32>,
14047        wv: &CudaSlice<f32>,
14048        q: &mut CudaSlice<f32>,
14049        k: &mut CudaSlice<f32>,
14050        v: &mut CudaSlice<f32>,
14051        head_dim: usize,
14052        n_rot: usize,
14053        rq: usize,
14054        rk: usize,
14055        pos: &CudaSlice<i32>,
14056        nh_q: usize,
14057        nh_k: usize,
14058        base: f32,
14059        freq_scale: f32,
14060        ff: Option<&CudaSlice<f32>>,
14061        eps: f32,
14062    ) -> Result<(), Box<dyn std::error::Error>> {
14063        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
14064        let f = self.func("rms_norm_qkv_rope_f32");
14065        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
14066        let cfg = LaunchConfig {
14067            grid_dim: (rows as u32, 1, 1),
14068            block_dim: (rms_block(), 1, 1),
14069            shared_mem_bytes: 0,
14070        };
14071        let theta_scale = base.powf(-2.0 / head_dim as f32);
14072        let (nc, rqi, rki, nhq, nhk) = (
14073            head_dim as i32,
14074            rq as i32,
14075            rk as i32,
14076            nh_q as i32,
14077            nh_k as i32,
14078        );
14079        let __s_b = self.gpu.stream();
14080        let mut b = __s_b.launch_builder(&f);
14081        match ff {
14082            Some(t) => {
14083                b.arg(q0)
14084                    .arg(k0)
14085                    .arg(v0)
14086                    .arg(wq)
14087                    .arg(wk)
14088                    .arg(wv)
14089                    .arg(&mut *q)
14090                    .arg(&mut *k)
14091                    .arg(&mut *v)
14092                    .arg(&nc)
14093                    .arg(&rqi)
14094                    .arg(&rki)
14095                    .arg(pos)
14096                    .arg(&nhq)
14097                    .arg(&nhk)
14098                    .arg(&theta_scale)
14099                    .arg(&freq_scale)
14100                    .arg(t)
14101                    .arg(&eps);
14102                unsafe {
14103                    b.launch(cfg)?;
14104                }
14105            }
14106            None => {
14107                let null: u64 = 0;
14108                b.arg(q0)
14109                    .arg(k0)
14110                    .arg(v0)
14111                    .arg(wq)
14112                    .arg(wk)
14113                    .arg(wv)
14114                    .arg(&mut *q)
14115                    .arg(&mut *k)
14116                    .arg(&mut *v)
14117                    .arg(&nc)
14118                    .arg(&rqi)
14119                    .arg(&rki)
14120                    .arg(pos)
14121                    .arg(&nhq)
14122                    .arg(&nhk)
14123                    .arg(&theta_scale)
14124                    .arg(&freq_scale)
14125                    .arg(&null)
14126                    .arg(&eps);
14127                unsafe {
14128                    b.launch(cfg)?;
14129                }
14130            }
14131        }
14132        Ok(())
14133    }
14134
14135    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
14136    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
14137    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
14138    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
14139    /// ([`Engine::full_width_rope_only`]).
14140    #[allow(clippy::too_many_arguments)]
14141    pub fn rms_norm_qkv_rope_append_dc(
14142        &self,
14143        q0: &CudaSlice<f32>,
14144        k0: &CudaSlice<f32>,
14145        v0: &CudaSlice<f32>,
14146        wq: &CudaSlice<f32>,
14147        wk: &CudaSlice<f32>,
14148        wv: &CudaSlice<f32>,
14149        q: &mut CudaSlice<f32>,
14150        k: &mut CudaSlice<f32>,
14151        v: &mut CudaSlice<f32>,
14152        head_dim: usize,
14153        n_rot: usize,
14154        rq: usize,
14155        rk: usize,
14156        pos: &CudaSlice<i32>,
14157        nh_q: usize,
14158        nh_k: usize,
14159        base: f32,
14160        freq_scale: f32,
14161        ff: Option<&CudaSlice<f32>>,
14162        eps: f32,
14163        kc: &mut CudaSlice<u8>,
14164        vc: &mut CudaSlice<u8>,
14165        t_dev: &CudaSlice<i32>,
14166        k_tok_bytes: usize,
14167        v_tok_bytes: usize,
14168        g: bool,
14169    ) -> Result<(), Box<dyn std::error::Error>> {
14170        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
14171        let rows = rq + rk + rk;
14172        let theta_scale = base.powf(-2.0 / head_dim as f32);
14173        let (nc, rqi, rki, nhq, nhk) = (
14174            head_dim as i32,
14175            rq as i32,
14176            rk as i32,
14177            nh_q as i32,
14178            nh_k as i32,
14179        );
14180        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
14181        if Self::pdl_on() && Self::pdl_wb_on() {
14182            use cudarc::driver::{DevicePtr, DevicePtrMut};
14183            let s = &self.gpu.stream();
14184            let (p0, _a0) = q0.device_ptr(s);
14185            let (p1, _a1) = k0.device_ptr(s);
14186            let (p2, _a2) = v0.device_ptr(s);
14187            let (pwq, _a3) = wq.device_ptr(s);
14188            let (pwk, _a4) = wk.device_ptr(s);
14189            let (pwv, _a5) = wv.device_ptr(s);
14190            let (pq, _a6) = q.device_ptr_mut(s);
14191            let (pk, _a7) = k.device_ptr_mut(s);
14192            let (pv, _a8) = v.device_ptr_mut(s);
14193            let (pp, _a9) = pos.device_ptr(s);
14194            let pff: u64 = match ff {
14195                Some(t) => {
14196                    let (p, _gg) = t.device_ptr(s);
14197                    p
14198                }
14199                None => 0,
14200            };
14201            let (pkc, _a10) = kc.device_ptr_mut(s);
14202            let (pvc, _a11) = vc.device_ptr_mut(s);
14203            let (pt, _a12) = t_dev.device_ptr(s);
14204            let mut ps = [
14205                &p0 as *const _ as *mut std::ffi::c_void,
14206                &p1 as *const _ as *mut _,
14207                &p2 as *const _ as *mut _,
14208                &pwq as *const _ as *mut _,
14209                &pwk as *const _ as *mut _,
14210                &pwv as *const _ as *mut _,
14211                &pq as *const _ as *mut _,
14212                &pk as *const _ as *mut _,
14213                &pv as *const _ as *mut _,
14214                &nc as *const _ as *mut _,
14215                &rqi as *const _ as *mut _,
14216                &rki as *const _ as *mut _,
14217                &pp as *const _ as *mut _,
14218                &nhq as *const _ as *mut _,
14219                &nhk as *const _ as *mut _,
14220                &theta_scale as *const _ as *mut _,
14221                &freq_scale as *const _ as *mut _,
14222                &pff as *const _ as *mut _,
14223                &eps as *const _ as *mut _,
14224                &pkc as *const _ as *mut _,
14225                &pvc as *const _ as *mut _,
14226                &pt as *const _ as *mut _,
14227                &ktb as *const _ as *mut _,
14228                &vtb as *const _ as *mut _,
14229            ];
14230            unsafe {
14231                self.launch_pdl_flash(
14232                    g,
14233                    "rms_norm_qkv_rope_append_dc_f32",
14234                    (rows as u32, 1, 1),
14235                    (rms_block(), 1, 1),
14236                    0,
14237                    &mut ps,
14238                )?;
14239            }
14240            return Ok(());
14241        }
14242        let f = if g {
14243            self.func_g("rms_norm_qkv_rope_append_dc_f32")
14244        } else {
14245            self.func("rms_norm_qkv_rope_append_dc_f32")
14246        };
14247        let cfg = LaunchConfig {
14248            grid_dim: (rows as u32, 1, 1),
14249            block_dim: (rms_block(), 1, 1),
14250            shared_mem_bytes: 0,
14251        };
14252        let __s_b = self.gpu.stream();
14253        let mut b = __s_b.launch_builder(&f);
14254        match ff {
14255            Some(t) => {
14256                b.arg(q0)
14257                    .arg(k0)
14258                    .arg(v0)
14259                    .arg(wq)
14260                    .arg(wk)
14261                    .arg(wv)
14262                    .arg(&mut *q)
14263                    .arg(&mut *k)
14264                    .arg(&mut *v)
14265                    .arg(&nc)
14266                    .arg(&rqi)
14267                    .arg(&rki)
14268                    .arg(pos)
14269                    .arg(&nhq)
14270                    .arg(&nhk)
14271                    .arg(&theta_scale)
14272                    .arg(&freq_scale)
14273                    .arg(t)
14274                    .arg(&eps)
14275                    .arg(&mut *kc)
14276                    .arg(&mut *vc)
14277                    .arg(t_dev)
14278                    .arg(&ktb)
14279                    .arg(&vtb);
14280                unsafe {
14281                    b.launch(cfg)?;
14282                }
14283            }
14284            None => {
14285                let null: u64 = 0;
14286                b.arg(q0)
14287                    .arg(k0)
14288                    .arg(v0)
14289                    .arg(wq)
14290                    .arg(wk)
14291                    .arg(wv)
14292                    .arg(&mut *q)
14293                    .arg(&mut *k)
14294                    .arg(&mut *v)
14295                    .arg(&nc)
14296                    .arg(&rqi)
14297                    .arg(&rki)
14298                    .arg(pos)
14299                    .arg(&nhq)
14300                    .arg(&nhk)
14301                    .arg(&theta_scale)
14302                    .arg(&freq_scale)
14303                    .arg(&null)
14304                    .arg(&eps)
14305                    .arg(&mut *kc)
14306                    .arg(&mut *vc)
14307                    .arg(t_dev)
14308                    .arg(&ktb)
14309                    .arg(&vtb);
14310                unsafe {
14311                    b.launch(cfg)?;
14312                }
14313            }
14314        }
14315        Ok(())
14316    }
14317
14318    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
14319    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
14320    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
14321    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
14322    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
14323    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
14324    /// `head_dim` ([`Engine::full_width_rope_only`]).
14325    #[allow(clippy::too_many_arguments)]
14326    pub fn rms_norm_qkv_rope_append(
14327        &self,
14328        q0: &CudaSlice<f32>,
14329        k0: &CudaSlice<f32>,
14330        v0: &CudaSlice<f32>,
14331        wq: &CudaSlice<f32>,
14332        wk: &CudaSlice<f32>,
14333        wv: &CudaSlice<f32>,
14334        q: &mut CudaSlice<f32>,
14335        k: &mut CudaSlice<f32>,
14336        v: &mut CudaSlice<f32>,
14337        head_dim: usize,
14338        n_rot: usize,
14339        rq: usize,
14340        rk: usize,
14341        pos: &CudaSlice<i32>,
14342        nh_q: usize,
14343        nh_k: usize,
14344        base: f32,
14345        freq_scale: f32,
14346        ff: Option<&CudaSlice<f32>>,
14347        eps: f32,
14348        kc: &mut CudaSlice<u8>,
14349        vc: &mut CudaSlice<u8>,
14350        t: usize,
14351        k_tok_bytes: usize,
14352        v_tok_bytes: usize,
14353        g: bool,
14354    ) -> Result<(), Box<dyn std::error::Error>> {
14355        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
14356        let rows = rq + rk + rk;
14357        let theta_scale = base.powf(-2.0 / head_dim as f32);
14358        let (nc, rqi, rki, nhq, nhk) = (
14359            head_dim as i32,
14360            rq as i32,
14361            rk as i32,
14362            nh_q as i32,
14363            nh_k as i32,
14364        );
14365        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
14366        let ti = t as i32;
14367        if Self::pdl_on() && Self::pdl_wb_on() {
14368            use cudarc::driver::{DevicePtr, DevicePtrMut};
14369            let s = &self.gpu.stream();
14370            let (p0, _a0) = q0.device_ptr(s);
14371            let (p1, _a1) = k0.device_ptr(s);
14372            let (p2, _a2) = v0.device_ptr(s);
14373            let (pwq, _a3) = wq.device_ptr(s);
14374            let (pwk, _a4) = wk.device_ptr(s);
14375            let (pwv, _a5) = wv.device_ptr(s);
14376            let (pq, _a6) = q.device_ptr_mut(s);
14377            let (pk, _a7) = k.device_ptr_mut(s);
14378            let (pv, _a8) = v.device_ptr_mut(s);
14379            let (pp, _a9) = pos.device_ptr(s);
14380            let pff: u64 = match ff {
14381                Some(t) => {
14382                    let (p, _gg) = t.device_ptr(s);
14383                    p
14384                }
14385                None => 0,
14386            };
14387            let (pkc, _a10) = kc.device_ptr_mut(s);
14388            let (pvc, _a11) = vc.device_ptr_mut(s);
14389            let mut ps = [
14390                &p0 as *const _ as *mut std::ffi::c_void,
14391                &p1 as *const _ as *mut _,
14392                &p2 as *const _ as *mut _,
14393                &pwq as *const _ as *mut _,
14394                &pwk as *const _ as *mut _,
14395                &pwv as *const _ as *mut _,
14396                &pq as *const _ as *mut _,
14397                &pk as *const _ as *mut _,
14398                &pv as *const _ as *mut _,
14399                &nc as *const _ as *mut _,
14400                &rqi as *const _ as *mut _,
14401                &rki as *const _ as *mut _,
14402                &pp as *const _ as *mut _,
14403                &nhq as *const _ as *mut _,
14404                &nhk as *const _ as *mut _,
14405                &theta_scale as *const _ as *mut _,
14406                &freq_scale as *const _ as *mut _,
14407                &pff as *const _ as *mut _,
14408                &eps as *const _ as *mut _,
14409                &pkc as *const _ as *mut _,
14410                &pvc as *const _ as *mut _,
14411                &ti as *const _ as *mut _,
14412                &ktb as *const _ as *mut _,
14413                &vtb as *const _ as *mut _,
14414            ];
14415            unsafe {
14416                self.launch_pdl_flash(
14417                    g,
14418                    "rms_norm_qkv_rope_append_f32",
14419                    (rows as u32, 1, 1),
14420                    (rms_block(), 1, 1),
14421                    0,
14422                    &mut ps,
14423                )?;
14424            }
14425            return Ok(());
14426        }
14427        let f = if g {
14428            self.func_g("rms_norm_qkv_rope_append_f32")
14429        } else {
14430            self.func("rms_norm_qkv_rope_append_f32")
14431        };
14432        let cfg = LaunchConfig {
14433            grid_dim: (rows as u32, 1, 1),
14434            block_dim: (rms_block(), 1, 1),
14435            shared_mem_bytes: 0,
14436        };
14437        let __s_b = self.gpu.stream();
14438        let mut b = __s_b.launch_builder(&f);
14439        let null: u64 = 0;
14440        b.arg(q0)
14441            .arg(k0)
14442            .arg(v0)
14443            .arg(wq)
14444            .arg(wk)
14445            .arg(wv)
14446            .arg(&mut *q)
14447            .arg(&mut *k)
14448            .arg(&mut *v)
14449            .arg(&nc)
14450            .arg(&rqi)
14451            .arg(&rki)
14452            .arg(pos)
14453            .arg(&nhq)
14454            .arg(&nhk)
14455            .arg(&theta_scale)
14456            .arg(&freq_scale);
14457        match ff {
14458            Some(t) => {
14459                b.arg(t);
14460            }
14461            None => {
14462                b.arg(&null);
14463            }
14464        }
14465        b.arg(&eps)
14466            .arg(&mut *kc)
14467            .arg(&mut *vc)
14468            .arg(&ti)
14469            .arg(&ktb)
14470            .arg(&vtb);
14471        unsafe {
14472            b.launch(cfg)?;
14473        }
14474        Ok(())
14475    }
14476
14477    pub fn add_q8_1(
14478        &self,
14479        a: &CudaSlice<f32>,
14480        b: &CudaSlice<f32>,
14481        res: &mut CudaSlice<f32>,
14482        ncols: usize,
14483        nrows: usize,
14484    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14485        debug_assert!(ncols.is_multiple_of(128));
14486        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14487        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14488        let f = self.func("add_q8_1_f32");
14489        let cfg = LaunchConfig {
14490            grid_dim: (nrows as u32, 1, 1),
14491            block_dim: (rms_block(), 1, 1),
14492            shared_mem_bytes: 0,
14493        };
14494        let nc = ncols as i32;
14495        let __s_b2 = self.gpu.stream();
14496        let mut b2 = __s_b2.launch_builder(&f);
14497        b2.arg(a)
14498            .arg(b)
14499            .arg(&mut *res)
14500            .arg(&mut out_q)
14501            .arg(&mut out_d)
14502            .arg(&nc);
14503        unsafe {
14504            b2.launch(cfg)?;
14505        }
14506        Ok((out_q, out_d))
14507    }
14508
14509    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
14510    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
14511    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
14512    #[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
14513    pub fn rms_pre_add_q8_1(
14514        &self,
14515        a: &CudaSlice<f32>,
14516        wa: &CudaSlice<f32>,
14517        b: &CudaSlice<f32>,
14518        res: &mut CudaSlice<f32>,
14519        ncols: usize,
14520        nrows: usize,
14521        eps: f32,
14522    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14523        debug_assert!(ncols.is_multiple_of(128));
14524        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14525        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14526        let f = self.func("rms_pre_add_q8_1_f32");
14527        let cfg = LaunchConfig {
14528            grid_dim: (nrows as u32, 1, 1),
14529            block_dim: (rms_block(), 1, 1),
14530            shared_mem_bytes: 0,
14531        };
14532        let (nc, ep) = (ncols as i32, eps);
14533        let __s_b2 = self.gpu.stream();
14534        let mut b2 = __s_b2.launch_builder(&f);
14535        b2.arg(a)
14536            .arg(wa)
14537            .arg(b)
14538            .arg(&mut *res)
14539            .arg(&mut out_q)
14540            .arg(&mut out_d)
14541            .arg(&nc)
14542            .arg(&ep);
14543        unsafe {
14544            b2.launch(cfg)?;
14545        }
14546        Ok((out_q, out_d))
14547    }
14548
14549    /// L2 norm per row (head_dim), no weight.
14550    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
14551    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
14552    pub fn l2_v2_on(ncols: usize) -> bool {
14553        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
14554    }
14555
14556    pub fn l2_norm_pp(
14557        &self,
14558        x: &CudaSlice<f32>,
14559        dst: &mut CudaSlice<f32>,
14560        dst16: Option<&mut CudaSlice<u8>>,
14561        ncols: usize,
14562        nrows: usize,
14563        eps: f32,
14564    ) -> Result<(), Box<dyn std::error::Error>> {
14565        if Self::l2_v2_on(ncols) {
14566            let f = self.func("l2_norm_pp_v2_f32");
14567            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
14568            let cfg = LaunchConfig {
14569                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
14570                block_dim: (256, 1, 1),
14571                shared_mem_bytes: 0,
14572            };
14573            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
14574            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
14575            let d16: u64 = match dst16 {
14576                Some(d) => self.addr_u8(d),
14577                None => 0,
14578            };
14579            let __s_b = self.gpu.stream();
14580            let mut b = __s_b.launch_builder(&f);
14581            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
14582            unsafe {
14583                b.launch(cfg)?;
14584            }
14585            return Ok(());
14586        }
14587        self.l2_norm(x, dst, ncols, nrows, eps)
14588    }
14589
14590    pub fn l2_norm(
14591        &self,
14592        x: &CudaSlice<f32>,
14593        dst: &mut CudaSlice<f32>,
14594        ncols: usize,
14595        nrows: usize,
14596        eps: f32,
14597    ) -> Result<(), Box<dyn std::error::Error>> {
14598        let f = self.func("l2_norm_f32");
14599        let cfg = LaunchConfig {
14600            grid_dim: (nrows as u32, 1, 1),
14601            block_dim: (256, 1, 1),
14602            shared_mem_bytes: 0,
14603        };
14604        let (nc, e) = (ncols as i32, eps);
14605        let __s_b = self.gpu.stream();
14606        let mut b = __s_b.launch_builder(&f);
14607        b.arg(x).arg(dst).arg(&nc).arg(&e);
14608        unsafe {
14609            b.launch(cfg)?;
14610        }
14611        Ok(())
14612    }
14613
14614    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
14615    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
14616    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
14617    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
14618    /// propagate through gdn_scan and flip argmax on marginal logits.
14619    pub fn l2_norm_decode(
14620        &self,
14621        x: &CudaSlice<f32>,
14622        dst: &mut CudaSlice<f32>,
14623        ncols: usize,
14624        nrows: usize,
14625        eps: f32,
14626    ) -> Result<(), Box<dyn std::error::Error>> {
14627        let f = self.func("l2_norm_f32");
14628        let cfg = LaunchConfig {
14629            grid_dim: (nrows as u32, 1, 1),
14630            block_dim: (32, 1, 1),
14631            shared_mem_bytes: 0,
14632        };
14633        let (nc, e) = (ncols as i32, eps);
14634        let __s_b = self.gpu.stream();
14635        let mut b = __s_b.launch_builder(&f);
14636        b.arg(x).arg(dst).arg(&nc).arg(&e);
14637        unsafe {
14638            b.launch(cfg)?;
14639        }
14640        Ok(())
14641    }
14642
14643    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
14644    #[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
14645    pub fn rope_neox(
14646        &self,
14647        x: &mut CudaSlice<f32>,
14648        pos: &CudaSlice<i32>,
14649        head_dim: usize,
14650        n_dims: usize,
14651        n_heads: usize,
14652        n_tokens: usize,
14653        freq_base: f32,
14654        freq_scale: f32,
14655    ) -> Result<(), Box<dyn std::error::Error>> {
14656        let f = self.func("rope_neox_f32");
14657        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14658        let grid = (n_heads * n_tokens) as u32;
14659        let cfg = LaunchConfig {
14660            grid_dim: (grid, 1, 1),
14661            block_dim: ((head_dim / 2) as u32, 1, 1),
14662            shared_mem_bytes: 0,
14663        };
14664        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14665        let __s_b = self.gpu.stream();
14666        let mut b = __s_b.launch_builder(&f);
14667        b.arg(x)
14668            .arg(pos)
14669            .arg(&hd)
14670            .arg(&nd)
14671            .arg(&nh)
14672            .arg(&theta_scale)
14673            .arg(&freq_scale);
14674        unsafe {
14675            b.launch(cfg)?;
14676        }
14677        Ok(())
14678    }
14679
14680    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
14681    #[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
14682    pub fn rope_neox_ff(
14683        &self,
14684        x: &mut CudaSlice<f32>,
14685        pos: &CudaSlice<i32>,
14686        head_dim: usize,
14687        n_dims: usize,
14688        n_heads: usize,
14689        n_tokens: usize,
14690        freq_base: f32,
14691        freq_scale: f32,
14692        ff: &CudaSlice<f32>,
14693    ) -> Result<(), Box<dyn std::error::Error>> {
14694        let f = self.func("rope_neox_ff_f32");
14695        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14696        let grid = (n_heads * n_tokens) as u32;
14697        let cfg = LaunchConfig {
14698            grid_dim: (grid, 1, 1),
14699            block_dim: ((head_dim / 2) as u32, 1, 1),
14700            shared_mem_bytes: 0,
14701        };
14702        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14703        let __s_b = self.gpu.stream();
14704        let mut b = __s_b.launch_builder(&f);
14705        b.arg(x)
14706            .arg(pos)
14707            .arg(&hd)
14708            .arg(&nd)
14709            .arg(&nh)
14710            .arg(&theta_scale)
14711            .arg(&freq_scale)
14712            .arg(ff);
14713        unsafe {
14714            b.launch(cfg)?;
14715        }
14716        Ok(())
14717    }
14718
14719    /// RoPE NEOX with per-dim freq factors AND the YaRN attention factor on cos/sin
14720    /// (qwen4_exp yarn lane — `rope_neox_ffm_f32`; ff = yarn_frequency_divisors, mscale =
14721    /// yarn_attention_factor). Identity inputs (ones, 1.0) reproduce `rope_neox` bit-for-bit.
14722    #[allow(clippy::too_many_arguments)]
14723    pub fn rope_neox_ffm(
14724        &self,
14725        x: &mut CudaSlice<f32>,
14726        pos: &CudaSlice<i32>,
14727        head_dim: usize,
14728        n_dims: usize,
14729        n_heads: usize,
14730        n_tokens: usize,
14731        freq_base: f32,
14732        freq_scale: f32,
14733        ff: &CudaSlice<f32>,
14734        mscale: f32,
14735    ) -> Result<(), Box<dyn std::error::Error>> {
14736        let f = self.func("rope_neox_ffm_f32");
14737        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14738        let grid = (n_heads * n_tokens) as u32;
14739        let cfg = LaunchConfig {
14740            grid_dim: (grid, 1, 1),
14741            block_dim: ((head_dim / 2) as u32, 1, 1),
14742            shared_mem_bytes: 0,
14743        };
14744        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
14745        let __s_b = self.gpu.stream();
14746        let mut b = __s_b.launch_builder(&f);
14747        b.arg(x)
14748            .arg(pos)
14749            .arg(&hd)
14750            .arg(&nd)
14751            .arg(&nh)
14752            .arg(&theta_scale)
14753            .arg(&freq_scale)
14754            .arg(ff)
14755            .arg(&mscale);
14756        unsafe {
14757            b.launch(cfg)?;
14758        }
14759        Ok(())
14760    }
14761
14762    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
14763    #[allow(clippy::too_many_arguments)]
14764    pub fn rope_neox2(
14765        &self,
14766        q: &mut CudaSlice<f32>,
14767        k: &mut CudaSlice<f32>,
14768        pos: &CudaSlice<i32>,
14769        head_dim: usize,
14770        n_dims: usize,
14771        nh_q: usize,
14772        nh_k: usize,
14773        n_tokens: usize,
14774        freq_base: f32,
14775        freq_scale: f32,
14776        ff: Option<&CudaSlice<f32>>,
14777    ) -> Result<(), Box<dyn std::error::Error>> {
14778        let f = self.func("rope_neox2_f32");
14779        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
14780        let grid = ((nh_q + nh_k) * n_tokens) as u32;
14781        let cfg = LaunchConfig {
14782            grid_dim: (grid, 1, 1),
14783            block_dim: ((head_dim / 2) as u32, 1, 1),
14784            shared_mem_bytes: 0,
14785        };
14786        let (hd, nd, nq, nk, nt) = (
14787            head_dim as i32,
14788            n_dims as i32,
14789            nh_q as i32,
14790            nh_k as i32,
14791            n_tokens as i32,
14792        );
14793        let __s_b = self.gpu.stream();
14794        let mut b = __s_b.launch_builder(&f);
14795        b.arg(q)
14796            .arg(k)
14797            .arg(pos)
14798            .arg(&hd)
14799            .arg(&nd)
14800            .arg(&nq)
14801            .arg(&nk)
14802            .arg(&nt)
14803            .arg(&theta_scale)
14804            .arg(&freq_scale);
14805        match ff {
14806            Some(ffv) => {
14807                b.arg(ffv);
14808                unsafe {
14809                    b.launch(cfg)?;
14810                }
14811            }
14812            None => {
14813                let null: u64 = 0;
14814                b.arg(&null);
14815                unsafe {
14816                    b.launch(cfg)?;
14817                }
14818            }
14819        }
14820        Ok(())
14821    }
14822
14823    /// gemma4 R1: dst = GELU_tanh(gate) * up.
14824    pub fn gelu_tanh_mul(
14825        &self,
14826        gate: &CudaSlice<f32>,
14827        up: &CudaSlice<f32>,
14828        dst: &mut CudaSlice<f32>,
14829        n: usize,
14830    ) -> Result<(), Box<dyn std::error::Error>> {
14831        let f = self.func("gelu_tanh_mul_f32");
14832        let cfg = LaunchConfig::for_num_elems(n as u32);
14833        let ni = n as i32;
14834        let __s_b = self.gpu.stream();
14835        let mut b = __s_b.launch_builder(&f);
14836        b.arg(gate).arg(up).arg(dst).arg(&ni);
14837        unsafe {
14838            b.launch(cfg)?;
14839        }
14840        Ok(())
14841    }
14842
14843    pub fn silu_mul(
14844        &self,
14845        gate: &CudaSlice<f32>,
14846        up: &CudaSlice<f32>,
14847        dst: &mut CudaSlice<f32>,
14848        n: usize,
14849    ) -> Result<(), Box<dyn std::error::Error>> {
14850        let f = self.func("silu_mul_f32");
14851        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
14852        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
14853        let ni = n as i32;
14854        let __s_b = self.gpu.stream();
14855        let mut b = __s_b.launch_builder(&f);
14856        b.arg(gate).arg(up).arg(dst).arg(&ni);
14857        unsafe {
14858            b.launch(cfg)?;
14859        }
14860        Ok(())
14861    }
14862
14863    /// SwiGLU twin using Memra's host-matching expf transcription.
14864    pub fn silu_mul_host_expf(
14865        &self,
14866        gate: &CudaSlice<f32>,
14867        up: &CudaSlice<f32>,
14868        dst: &mut CudaSlice<f32>,
14869        n: usize,
14870    ) -> Result<(), Box<dyn std::error::Error>> {
14871        let f = self.func("silu_mul_host_expf_f32");
14872        let cfg = LaunchConfig::for_num_elems(n as u32);
14873        let ni = n as i32;
14874        let __s_b = self.gpu.stream();
14875        let mut b = __s_b.launch_builder(&f);
14876        b.arg(gate).arg(up).arg(dst).arg(&ni);
14877        unsafe {
14878            b.launch(cfg)?;
14879        }
14880        Ok(())
14881    }
14882
14883    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
14884    pub fn silu_clamped_mul_host_expf(
14885        &self,
14886        gate: &CudaSlice<f32>,
14887        up: &CudaSlice<f32>,
14888        limit: f32,
14889        dst: &mut CudaSlice<f32>,
14890        n: usize,
14891    ) -> Result<(), Box<dyn std::error::Error>> {
14892        if !limit.is_finite() || limit <= 0.0 {
14893            return Err(
14894                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
14895            );
14896        }
14897        let f = self.func("silu_clamped_mul_host_expf_f32");
14898        let cfg = LaunchConfig::for_num_elems(n as u32);
14899        let ni = n as i32;
14900        let __s_b = self.gpu.stream();
14901        let mut b = __s_b.launch_builder(&f);
14902        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
14903        unsafe {
14904            b.launch(cfg)?;
14905        }
14906        Ok(())
14907    }
14908
14909    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
14910    /// for the down projection — kills the standalone convert pass. Bit-identical class.
14911    pub fn silu_mul_f16out(
14912        &self,
14913        gate: &CudaSlice<f32>,
14914        up: &CudaSlice<f32>,
14915        dst: &mut CudaSlice<f32>,
14916        dst16: &mut CudaSlice<u8>,
14917        n: usize,
14918    ) -> Result<(), Box<dyn std::error::Error>> {
14919        let f = self.func("silu_mul_f16out_f32");
14920        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
14921        let ni = n as i32;
14922        let __s_b = self.gpu.stream();
14923        let mut b = __s_b.launch_builder(&f);
14924        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
14925        unsafe {
14926            b.launch(cfg)?;
14927        }
14928        Ok(())
14929    }
14930
14931    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
14932    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
14933    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
14934    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
14935    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
14936    /// launches per dense FFN layer (the gate+up post-matmul scales).
14937    pub fn silu_mul_scaled(
14938        &self,
14939        gate: &CudaSlice<f32>,
14940        up: &CudaSlice<f32>,
14941        gs: f32,
14942        us: f32,
14943        dst: &mut CudaSlice<f32>,
14944        n: usize,
14945    ) -> Result<(), Box<dyn std::error::Error>> {
14946        let f = self.func("silu_mul_scaled_f32");
14947        let cfg = LaunchConfig::for_num_elems(n as u32);
14948        let ni = n as i32;
14949        let (gsf, usf) = (gs, us);
14950        let __s_b = self.gpu.stream();
14951        let mut b = __s_b.launch_builder(&f);
14952        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
14953        unsafe {
14954            b.launch(cfg)?;
14955        }
14956        Ok(())
14957    }
14958
14959    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
14960    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
14961    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
14962    #[allow(clippy::too_many_arguments)]
14963    pub fn swigluoai_mul_scaled(
14964        &self,
14965        gate: &CudaSlice<f32>,
14966        up: &CudaSlice<f32>,
14967        gs: f32,
14968        us: f32,
14969        alpha: f32,
14970        limit: f32,
14971        dst: &mut CudaSlice<f32>,
14972        n: usize,
14973    ) -> Result<(), Box<dyn std::error::Error>> {
14974        let f = self.func("swigluoai_mul_scaled_f32");
14975        let cfg = LaunchConfig::for_num_elems(n as u32);
14976        let ni = n as i32;
14977        let __s_b = self.gpu.stream();
14978        let mut b = __s_b.launch_builder(&f);
14979        b.arg(gate)
14980            .arg(up)
14981            .arg(&gs)
14982            .arg(&us)
14983            .arg(&alpha)
14984            .arg(&limit)
14985            .arg(dst)
14986            .arg(&ni);
14987        unsafe {
14988            b.launch(cfg)?;
14989        }
14990        Ok(())
14991    }
14992
14993    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
14994    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
14995    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
14996    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
14997    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
14998    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
14999    /// n must be a multiple of 32 (n_ff always is).
15000    pub fn silu_mul_scaled_q8_1(
15001        &self,
15002        gate: &CudaSlice<f32>,
15003        up: &CudaSlice<f32>,
15004        gs: f32,
15005        us: f32,
15006        n: usize,
15007    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15008        let f = self.func("silu_mul_scaled_q8_1");
15009        let nblk = n / 32;
15010        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
15011        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
15012        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
15013        let cfg = LaunchConfig::for_num_elems(n as u32);
15014        let (gsf, usf, ni) = (gs, us, n as i32);
15015        let __s_b = self.gpu.stream();
15016        let mut b = __s_b.launch_builder(&f);
15017        b.arg(gate)
15018            .arg(up)
15019            .arg(&gsf)
15020            .arg(&usf)
15021            .arg(&mut aq)
15022            .arg(&mut ad)
15023            .arg(&ni);
15024        unsafe {
15025            b.launch(cfg)?;
15026        }
15027        Ok((aq, ad))
15028    }
15029
15030    pub fn add(
15031        &self,
15032        a: &CudaSlice<f32>,
15033        b_in: &CudaSlice<f32>,
15034        dst: &mut CudaSlice<f32>,
15035        n: usize,
15036    ) -> Result<(), Box<dyn std::error::Error>> {
15037        let f = self.func("add_f32");
15038        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
15039        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15040        let ni = n as i32;
15041        let __s_bld = self.gpu.stream();
15042        let mut bld = __s_bld.launch_builder(&f);
15043        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
15044        unsafe {
15045            bld.launch(cfg)?;
15046        }
15047        Ok(())
15048    }
15049
15050    pub fn mul(
15051        &self,
15052        a: &CudaSlice<f32>,
15053        b_in: &CudaSlice<f32>,
15054        dst: &mut CudaSlice<f32>,
15055        n: usize,
15056    ) -> Result<(), Box<dyn std::error::Error>> {
15057        let f = self.func("mul_f32");
15058        let cfg = LaunchConfig::for_num_elems(n as u32);
15059        let ni = n as i32;
15060        let __s_bld = self.gpu.stream();
15061        let mut bld = __s_bld.launch_builder(&f);
15062        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
15063        unsafe {
15064            bld.launch(cfg)?;
15065        }
15066        Ok(())
15067    }
15068
15069    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
15070    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
15071    pub fn matmul(
15072        &self,
15073        w: &crate::model::GpuTensor,
15074        x: &CudaSlice<f32>,
15075        m: usize,
15076    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15077        use crate::model::GpuTensor;
15078        let in_f = w.in_features();
15079        let out_f = w.out_features();
15080        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
15081        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
15082        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
15083        // gives nothing). Quantize the activation once here then call the GEMM.
15084        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
15085        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
15086        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
15087        #[allow(non_snake_case)]
15088        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
15089        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
15090        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
15091            usize::MAX
15092        } else {
15093            16usize
15094        };
15095
15096        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
15097        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
15098        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
15099        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
15100        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
15101        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
15102        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
15103        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
15104        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
15105        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
15106        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
15107        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
15108        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
15109        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
15110        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
15111        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
15112        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
15113        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
15114        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
15115        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
15116        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
15117        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
15118        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
15119        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
15120        if m >= GEMM_M_THRESHOLD {
15121            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
15122                return Ok(y);
15123            }
15124            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
15125            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
15126            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
15127            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
15128            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
15129            // tile defaults differently by operand source.
15130            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15131                return Ok(y);
15132            }
15133            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
15134            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
15135            if let Some(y) = self.try_f16_gemm(w, x, m)? {
15136                return Ok(y);
15137            }
15138        }
15139        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
15140        // m threshold the rest of this method uses:
15141        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
15142        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
15143        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
15144        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
15145        //     across every tier by construction with no batched twin needed.
15146        //
15147        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
15148        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
15149        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
15150        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
15151        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
15152        // arms is what makes sure it never gets there.
15153        if let GpuTensor::Quant { qtype, .. } = w
15154            && *qtype == QT_F8_E4M3_BLK
15155        {
15156            if m >= GEMM_M_THRESHOLD
15157                && let Some(y) = self.try_e4m3_blk_prefill(w, x, m)?
15158            {
15159                return Ok(y);
15160            }
15161            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15162            if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
15163                return Ok(y);
15164            }
15165        }
15166        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
15167            return self.qmatvec_mmq(w, x, m);
15168        }
15169        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
15170            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15171            return self.qmatvec_gemm(w, &aq, &ad, m);
15172        }
15173        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
15174        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
15175        if m >= GEMM_M_THRESHOLD
15176            && let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)?
15177        {
15178            return Ok(y);
15179        }
15180        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
15181        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
15182        // to Stage-A f32-dequant (the correctness oracle path).
15183        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
15184        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
15185        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
15186        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
15187        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
15188        if m == 1
15189            && fast
15190            && let GpuTensor::Quant {
15191                bytes,
15192                qtype,
15193                row_bytes,
15194                rp,
15195                rp4,
15196                scale,
15197                ..
15198            } = w
15199            && self.mmvq_supports(*qtype)
15200        {
15201            // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
15202            // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
15203            // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
15204            let (bytes, rp) = match rp4 {
15205                Some(m4) => (m4, true),
15206                None => (bytes, *rp),
15207            };
15208            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15209            return self.qmatvec_mmvq(
15210                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
15211            );
15212        }
15213        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
15214        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
15215        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
15216        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
15217        // block below. MEMRA_NO_BATCHED -> per-m path.
15218        //
15219        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
15220        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
15221        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
15222        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
15223        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
15224        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
15225        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
15226        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
15227        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
15228        if (2..=16).contains(&m)
15229            && fast
15230            && std::env::var("MEMRA_NO_BATCHED").is_err()
15231            && (m <= 4 || Self::b8_enabled())
15232        {
15233            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
15234            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
15235            // is present (rp4) — the mirror pick below then routes to the _rp family.
15236            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
15237            // because the native e4m3 row layout is already aligned and needs no mirror.
15238            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
15239            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
15240            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
15241            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
15242            let m_ok = m <= 8
15243                || matches!(w, GpuTensor::Quant { qtype, .. }
15244                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
15245                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
15246            if m_ok
15247                && let GpuTensor::Quant {
15248                    bytes,
15249                    qtype,
15250                    row_bytes,
15251                    rp,
15252                    rp4,
15253                    ..
15254                } = w
15255                && self.batched_supports(*qtype)
15256                && self.mmvq_supports(*qtype)
15257            {
15258                let (bytes, rp) = match rp4 {
15259                    Some(m4) => (m4, true),
15260                    None => (bytes, *rp),
15261                };
15262                let mcols = Self::batched_mcols(m);
15263                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15264                let mut y = self.qmatvec_mmvq_batched(
15265                    bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
15266                )?;
15267                if let GpuTensor::Quant { scale, .. } = w
15268                    && *scale != 1.0
15269                {
15270                    self.scale_inplace(&mut y, *scale, m * out_f)?;
15271                }
15272                return Ok(y);
15273            }
15274        }
15275        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
15276        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
15277        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
15278        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
15279        // for this dtype, so the generic match below must never see it under `fast`.
15280        if fast
15281            && let GpuTensor::Quant {
15282                bytes,
15283                qtype,
15284                row_bytes,
15285                scale,
15286                ..
15287            } = w
15288            && *qtype == QT_F8_E4M3
15289        {
15290            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15291            return self.qmatvec_mmvq(
15292                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
15293            );
15294        }
15295        let mut y = match w {
15296            GpuTensor::Quant {
15297                bytes,
15298                qtype,
15299                row_bytes,
15300                ..
15301            } if fast && *qtype == QT_Q8_0 => {
15302                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15303            }
15304            GpuTensor::Quant {
15305                bytes,
15306                qtype,
15307                row_bytes,
15308                ..
15309            } if fast && *qtype == QT_Q4_K => {
15310                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15311            }
15312            GpuTensor::Quant {
15313                bytes,
15314                qtype,
15315                row_bytes,
15316                ..
15317            } if fast && *qtype == QT_Q6_K => {
15318                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15319            }
15320            GpuTensor::Quant {
15321                bytes,
15322                qtype,
15323                row_bytes,
15324                ..
15325            } if fast && *qtype == QT_Q5_K => {
15326                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15327            }
15328            GpuTensor::Quant {
15329                bytes,
15330                qtype,
15331                row_bytes,
15332                ..
15333            } if fast && *qtype == QT_Q3_K => {
15334                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15335            }
15336            GpuTensor::Quant {
15337                bytes,
15338                qtype,
15339                row_bytes,
15340                rp,
15341                ..
15342            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
15343                if *rp {
15344                    "qmatvec_nvfp4_dp4a_rp"
15345                } else {
15346                    "qmatvec_nvfp4_dp4a"
15347                },
15348                &bytes.slice(0..bytes.len()),
15349                x,
15350                m,
15351                in_f,
15352                out_f,
15353                *row_bytes,
15354            )?,
15355            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
15356            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
15357            // anomaly (research/kat-anomaly-20260802/).
15358            GpuTensor::Quant {
15359                bytes,
15360                qtype,
15361                row_bytes,
15362                ..
15363            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
15364                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
15365            }
15366            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
15367            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
15368            // without first writing the matching kernel, or func() will panic
15369            // "kernel ... not in any fatbin".
15370            GpuTensor::Quant {
15371                bytes,
15372                qtype,
15373                row_bytes,
15374                rp,
15375                ..
15376            } =>
15377            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
15378            // deq(row,j) form cannot address the planes; same value/product order).
15379            {
15380                self.qmatvec(
15381                    bytes,
15382                    x,
15383                    m,
15384                    in_f,
15385                    out_f,
15386                    if *rp && *qtype == QT_NVFP4 {
15387                        QT_NVFP4_RP
15388                    } else {
15389                        *qtype
15390                    },
15391                    *row_bytes,
15392                )?
15393            }
15394            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
15395            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
15396            // cuBLASLt f32 GEMV as the Float arm.
15397            GpuTensor::FloatBf16 { data, .. } => {
15398                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
15399                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
15400                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
15401                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
15402                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
15403                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
15404                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f.is_multiple_of(8) {
15405                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15406                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
15407                    y
15408                } else {
15409                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
15410                }
15411            }
15412        };
15413        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
15414        if let GpuTensor::Quant { scale, .. } = w
15415            && *scale != 1.0
15416        {
15417            self.scale_inplace(&mut y, *scale, m * out_f)?;
15418        }
15419        Ok(y)
15420    }
15421
15422    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
15423    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
15424    ///
15425    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
15426    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
15427    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
15428    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
15429    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
15430    /// path must not pay an env lookup for a flag that is off.
15431    pub fn stage_a_raw_needed() -> bool {
15432        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15433        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
15434    }
15435
15436    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
15437    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
15438    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
15439        use crate::model::GpuTensor;
15440        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
15441            return false;
15442        }
15443        match w {
15444            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
15445            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
15446            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
15447            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
15448            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
15449            // block class has no fused twin yet, so each of its projections takes its own launch.
15450            GpuTensor::Quant { qtype, .. } => {
15451                matches!(
15452                    *qtype,
15453                    QT_Q8_0
15454                        | QT_Q4_K
15455                        | QT_Q6_K
15456                        | QT_Q5_K
15457                        | QT_Q3_K
15458                        | QT_NVFP4
15459                        | QT_F8_E4M3
15460                        | QT_F8_E4M3_BLK
15461                        | QT_Q4_0
15462                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
15463            }
15464            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15465        }
15466    }
15467
15468    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
15469    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
15470    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
15471    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
15472    pub fn matmul_pre(
15473        &self,
15474        w: &crate::model::GpuTensor,
15475        aq: &CudaSlice<i8>,
15476        ad: &CudaSlice<f32>,
15477        x_fallback: &CudaSlice<f32>,
15478        m: usize,
15479    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15480        use crate::model::GpuTensor;
15481        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
15482        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
15483        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
15484        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
15485        // rc=30013 dig, 2026-07-31).
15486        let x_raw_ok = x_fallback.len() >= m * w.in_features();
15487        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
15488        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
15489        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
15490            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
15491                return Ok(y);
15492            }
15493            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
15494            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
15495            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
15496                return Ok(y);
15497            }
15498            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
15499            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
15500                return Ok(y);
15501            }
15502        }
15503        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
15504        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
15505        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
15506        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
15507        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
15508        if m >= 16
15509            && x_raw_ok
15510            && !self.verify_exact_on()
15511            && let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)?
15512        {
15513            return Ok(y);
15514        }
15515        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15516            return Ok(y);
15517        }
15518        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
15519        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
15520        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
15521        // aq/ad.
15522        if m >= 16
15523            && w.out_features() >= 128
15524            && self.mmq_supports(w)
15525            && !self.verify_exact_on()
15526            && x_raw_ok
15527        {
15528            return self.qmatvec_mmq(w, x_fallback, m);
15529        }
15530        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
15531        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
15532        if m >= 16
15533            && x_raw_ok
15534            && !self.verify_exact_on()
15535            && let Some(y) =
15536                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
15537        {
15538            return Ok(y);
15539        }
15540        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
15541        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
15542        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
15543            return self.qmatvec_gemm(w, aq, ad, m);
15544        }
15545        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
15546        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
15547        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
15548        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
15549        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
15550        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
15551        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
15552        // which reads `m * in_f` floats out of a 0-byte allocation ->
15553        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
15554        // it poisons the context, so every LATER request in that process fails with an unrelated
15555        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
15556        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
15557        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
15558        // dense artifact and left the arm with no working truth instrument.
15559        //
15560        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
15561        // strictly better than an illegal address surfacing later at an unrelated sync point, and
15562        // an oracle that cannot run must say so rather than corrupt the context it runs in.
15563        if !self.uses_q8_1_fast(w) {
15564            if !x_raw_ok {
15565                return Err(format!(
15566                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
15567                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
15568                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
15569                     activation (see Engine::rms_norm_decode, which is bit-identical to \
15570                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
15571                    x_fallback.len(),
15572                    m,
15573                    w.in_features(),
15574                    m * w.in_features()
15575                )
15576                .into());
15577            }
15578            return self.matmul(w, x_fallback, m);
15579        }
15580        let in_f = w.in_features();
15581        let out_f = w.out_features();
15582        let (bytes, qtype, row_bytes, scale, rp) = match w {
15583            GpuTensor::Quant {
15584                bytes,
15585                qtype,
15586                row_bytes,
15587                scale,
15588                rp,
15589                ..
15590            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15591            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
15592        };
15593        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
15594        // the dp4a/oracle tails below keep the raw GGUF bytes.
15595        let (mbytes, mrp) = match w {
15596            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15597            _ => (bytes, rp),
15598        };
15599        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
15600        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
15601        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
15602        if m == 1 && self.mmvq_supports(qtype) {
15603            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
15604        }
15605        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
15606        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
15607        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
15608        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
15609        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
15610        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
15611        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
15612        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
15613        // m=5..8 on the old per-m path (b8-tier-only seam).
15614        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
15615        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
15616        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
15617        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
15618            && std::env::var("MEMRA_NO_BATCHED").is_err()
15619            && (m <= 4 || Self::b8_enabled())
15620            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
15621            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
15622            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
15623            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
15624                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
15625        {
15626            let mcols = Self::batched_mcols(m);
15627            return self.qmatvec_mmvq_batched(
15628                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
15629            );
15630        }
15631        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
15632        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
15633        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
15634        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
15635        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
15636        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
15637            let (b2, r2) = if qtype == QT_Q4_0 {
15638                (mbytes, mrp)
15639            } else {
15640                (bytes, rp)
15641            };
15642            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
15643        }
15644        let name = match qtype {
15645            QT_Q8_0 => "qmatvec_q8_0_dp4a",
15646            QT_Q4_K => "qmatvec_q4_K_dp4a",
15647            QT_Q6_K => "qmatvec_q6_K_dp4a",
15648            QT_Q5_K => "qmatvec_q5_K_dp4a",
15649            QT_Q3_K => "qmatvec_q3_K_dp4a",
15650            QT_NVFP4 => {
15651                if rp {
15652                    "qmatvec_nvfp4_dp4a_rp"
15653                } else {
15654                    "qmatvec_nvfp4_dp4a"
15655                }
15656            }
15657            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15658            _ => unreachable!(),
15659        };
15660        let f = self.func(name);
15661        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15662        let cfg = LaunchConfig {
15663            grid_dim: (out_f as u32, m as u32, 1),
15664            block_dim: (128, 1, 1),
15665            shared_mem_bytes: 0,
15666        };
15667        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15668        let __s_b = self.gpu.stream();
15669        let mut b = __s_b.launch_builder(&f);
15670        b.arg(bytes)
15671            .arg(aq)
15672            .arg(ad)
15673            .arg(&mut y)
15674            .arg(&inf)
15675            .arg(&outf)
15676            .arg(&mi)
15677            .arg(&rb);
15678        unsafe {
15679            b.launch(cfg)?;
15680        }
15681        if scale != 1.0 {
15682            self.scale_inplace(&mut y, scale, m * out_f)?;
15683        }
15684        Ok(y)
15685    }
15686
15687    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
15688    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
15689    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
15690    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
15691    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
15692    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
15693    /// reduce as m=1); this method just forces that path unconditionally.
15694    pub fn matmul_decode_exact(
15695        &self,
15696        w: &crate::model::GpuTensor,
15697        x: &CudaSlice<f32>,
15698        m: usize,
15699    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15700        use crate::model::GpuTensor;
15701        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
15702        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
15703        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
15704        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
15705        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
15706        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
15707        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
15708        if let GpuTensor::Float { data, .. } = w {
15709            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
15710        }
15711        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
15712        // float linear (same n-independent reduction contract as the Float arm above).
15713        if let GpuTensor::FloatBf16 { data, .. } = w {
15714            let (in_f, out_f) = (w.in_features(), w.out_features());
15715            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
15716            // contract — the whole-weight f32 dequant disappears too).
15717            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
15718                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15719                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
15720                return Ok(y);
15721            }
15722            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
15723        }
15724        if !self.uses_q8_1_fast(w) {
15725            return self.matmul(w, x, m);
15726        }
15727        let in_f = w.in_features();
15728        let out_f = w.out_features();
15729        let (bytes, qtype, row_bytes, scale, rp) = match w {
15730            GpuTensor::Quant {
15731                bytes,
15732                qtype,
15733                row_bytes,
15734                scale,
15735                rp,
15736                ..
15737            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15738            _ => return self.matmul(w, x, m),
15739        };
15740        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
15741        // which does its own mirror pick).
15742        let (bytes, rp) = match w {
15743            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15744            _ => (bytes, rp),
15745        };
15746        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15747        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
15748        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
15749        // (token,row) by construction, which is exactly what this method exists to guarantee.
15750        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
15751            return Ok(y);
15752        }
15753        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
15754        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
15755        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
15756        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
15757        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
15758        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
15759        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
15760        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
15761        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
15762            && std::env::var("MEMRA_NO_BATCHED").is_err()
15763            && (m <= 4 || Self::b8_enabled())
15764            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
15765            // no mirror precondition, `rp` selects the layout only.
15766            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
15767                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
15768        {
15769            let mcols = Self::batched_mcols(m);
15770            return self.qmatvec_mmvq_batched(
15771                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
15772            );
15773        }
15774        if self.mmvq_supports(qtype) {
15775            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
15776            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
15777            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
15778        }
15779        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
15780        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
15781        self.matmul_pre(w, &aq, &ad, x, m)
15782    }
15783
15784    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
15785    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
15786    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
15787    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
15788    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
15789    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
15790    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
15791    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
15792    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
15793    pub fn matmul_decode_exact_pre(
15794        &self,
15795        w: &crate::model::GpuTensor,
15796        aq: &CudaSlice<i8>,
15797        ad: &CudaSlice<f32>,
15798        m: usize,
15799    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15800        use crate::model::GpuTensor;
15801        debug_assert!(
15802            self.uses_q8_1_fast(w),
15803            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
15804        );
15805        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
15806        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15807            return Ok(y);
15808        }
15809        let in_f = w.in_features();
15810        let out_f = w.out_features();
15811        let (bytes, qtype, row_bytes, scale, rp) = match w {
15812            GpuTensor::Quant {
15813                bytes,
15814                qtype,
15815                row_bytes,
15816                scale,
15817                rp,
15818                ..
15819            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15820            _ => {
15821                return Err(
15822                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
15823                );
15824            }
15825        };
15826        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
15827        let (bytes, rp) = match w {
15828            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15829            _ => (bytes, rp),
15830        };
15831        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
15832        if (2..=16).contains(&m)
15833            && self.batched_supports(qtype)
15834            && self.mmvq_supports(qtype)
15835            && std::env::var("MEMRA_NO_BATCHED").is_err()
15836            && (m <= 4 || Self::b8_enabled())
15837            && (m <= 8
15838                || qtype == QT_Q4_0
15839                || qtype == QT_Q6_K
15840                || qtype == QT_F8_E4M3
15841                || qtype == QT_NVFP4
15842                || qtype == QT_Q4_K
15843                || qtype == QT_Q5_K
15844                || qtype == QT_Q8_0)
15845        {
15846            let mcols = Self::batched_mcols(m);
15847            return self.qmatvec_mmvq_batched(
15848                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
15849            );
15850        }
15851        if self.mmvq_supports(qtype) {
15852            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
15853        }
15854        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
15855        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
15856        let x0 = self.zeros(0)?;
15857        self.matmul_pre(w, aq, ad, &x0, m)
15858    }
15859
15860    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
15861    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
15862    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
15863    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
15864    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
15865    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
15866    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
15867    /// per-tensor path.
15868    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
15869    pub fn matmul_decode_exact_dual_pre(
15870        &self,
15871        w0: &crate::model::GpuTensor,
15872        w1: &crate::model::GpuTensor,
15873        aq: &CudaSlice<i8>,
15874        ad: &CudaSlice<f32>,
15875        m: usize,
15876    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
15877    {
15878        use crate::model::GpuTensor;
15879        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15880        let on = *ON.get_or_init(|| {
15881            std::env::var("MEMRA_SPEC_DUAL_T")
15882                .map(|v| v != "0")
15883                .unwrap_or(true)
15884        });
15885        if !on
15886            || !(2..=7).contains(&m)
15887            || std::env::var("MEMRA_NO_BATCHED").is_ok()
15888            || !self.uses_q8_1_fast(w0)
15889            || !self.uses_q8_1_fast(w1)
15890        {
15891            return Ok(None);
15892        }
15893        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
15894        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
15895        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
15896        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
15897        if !self.mmvq_supports(QT_NVFP4) {
15898            return Ok(None);
15899        }
15900        let (in_f, out_f) = (w0.in_features(), w0.out_features());
15901        if w1.in_features() != in_f || w1.out_features() != out_f {
15902            return Ok(None);
15903        }
15904        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
15905            (
15906                GpuTensor::Quant {
15907                    bytes: b0,
15908                    qtype: q0,
15909                    row_bytes: rb0,
15910                    scale: s0,
15911                    rp: rp0,
15912                    rp4: None,
15913                    ..
15914                },
15915                GpuTensor::Quant {
15916                    bytes: b1,
15917                    qtype: q1,
15918                    row_bytes: rb1,
15919                    scale: s1,
15920                    rp: rp1,
15921                    rp4: None,
15922                    ..
15923                },
15924            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
15925                (b0, b1, *rb0, *s0, *s1, *rp0)
15926            }
15927            _ => return Ok(None),
15928        };
15929        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
15930        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
15931        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
15932        {
15933            return Ok(None);
15934        }
15935        let (y0, y1) =
15936            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
15937        Ok(Some(((y0, s0), (y1, s1))))
15938    }
15939
15940    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
15941    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
15942    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
15943    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
15944    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
15945    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
15946    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
15947    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
15948    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
15949    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
15950    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
15951    pub fn matmul_decode_exact_group4_pre(
15952        &self,
15953        ws: [&crate::model::GpuTensor; 4],
15954        aq: &CudaSlice<i8>,
15955        ad: &CudaSlice<f32>,
15956        m: usize,
15957    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15958        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15959        let on = *ON.get_or_init(|| {
15960            std::env::var("MEMRA_TK_GDN_GROUP")
15961                .map(|v| v != "0")
15962                .unwrap_or(true)
15963        });
15964        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
15965    }
15966
15967    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
15968    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
15969    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
15970    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
15971    pub fn matmul_decode_exact_group3_pre(
15972        &self,
15973        ws: [&crate::model::GpuTensor; 3],
15974        aq: &CudaSlice<i8>,
15975        ad: &CudaSlice<f32>,
15976        m: usize,
15977    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15978        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15979        let on = *ON.get_or_init(|| {
15980            std::env::var("MEMRA_TK_FA_GROUP")
15981                .map(|v| v != "0")
15982                .unwrap_or(true)
15983        });
15984        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
15985    }
15986
15987    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
15988    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
15989    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
15990    #[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
15991    fn matmul_decode_exact_group_pre(
15992        &self,
15993        ws: &[&crate::model::GpuTensor],
15994        aq: &CudaSlice<i8>,
15995        ad: &CudaSlice<f32>,
15996        m: usize,
15997        on: bool,
15998        tag: &'static str,
15999    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16000        use crate::model::GpuTensor;
16001        if !on
16002            || !(2..=16).contains(&m)
16003            || std::env::var("MEMRA_NO_BATCHED").is_ok()
16004            || (m > 4 && !Self::b8_enabled())
16005            || !self.mmvq_supports(QT_NVFP4)
16006            || !self.batched_supports(QT_NVFP4)
16007        {
16008            return Ok(None);
16009        }
16010        let in_f = ws[0].in_features();
16011        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
16012        for w in ws {
16013            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
16014                return Ok(None);
16015            }
16016            match w {
16017                GpuTensor::Quant {
16018                    bytes,
16019                    qtype,
16020                    scale,
16021                    rp: true,
16022                    rp4: None,
16023                    ..
16024                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
16025                    parts.push((bytes, w.out_features(), *scale));
16026                }
16027                _ => return Ok(None),
16028            }
16029        }
16030        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
16031        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16032        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16033        let mcols = if (5..=7).contains(&m) && b567 {
16034            m
16035        } else {
16036            Self::batched_mcols(m)
16037        };
16038        let kname: &'static str = match mcols {
16039            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
16040            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
16041            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
16042            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
16043            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
16044            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
16045            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
16046            _ => return Ok(None),
16047        };
16048        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
16049        // the second door's print on the slice-D battery — key the once-set by tag.
16050        if std::env::var("MEMRA_DEBUG").is_ok() {
16051            use std::sync::Mutex;
16052            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
16053            let mut seen = SEEN.lock().unwrap();
16054            if !seen.contains(&tag) {
16055                seen.push(tag);
16056                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
16057            }
16058        }
16059        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
16060        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
16061        let total: usize = parts.iter().map(|p| p.1).sum();
16062        let three = parts.len() == 3;
16063        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
16064        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
16065        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
16066        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
16067        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
16068        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
16069        let cfg = LaunchConfig {
16070            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16071            block_dim: (32, ROWS_PER_BLOCK, 1),
16072            shared_mem_bytes: 0,
16073        };
16074        let (inf, mi) = (in_f as i32, m as i32);
16075        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
16076        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
16077        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
16078        let s3 = if three { 1.0f32 } else { parts[3].2 };
16079        let w3 = if three { parts[0].0 } else { parts[3].0 };
16080        let f = self.func(kname);
16081        let __s_b = self.gpu.stream();
16082        let mut b = __s_b.launch_builder(&f);
16083        b.arg(parts[0].0)
16084            .arg(parts[1].0)
16085            .arg(parts[2].0)
16086            .arg(w3)
16087            .arg(aq)
16088            .arg(ad)
16089            .arg(&mut y0)
16090            .arg(&mut y1)
16091            .arg(&mut y2)
16092            .arg(&mut y3)
16093            .arg(&inf)
16094            .arg(&n0)
16095            .arg(&n1)
16096            .arg(&n2)
16097            .arg(&n3)
16098            .arg(&mi)
16099            .arg(&s0)
16100            .arg(&s1)
16101            .arg(&s2)
16102            .arg(&s3);
16103        unsafe {
16104            b.launch(cfg)?;
16105        }
16106        Ok(Some(if three {
16107            vec![y0, y1, y2]
16108        } else {
16109            vec![y0, y1, y2, y3]
16110        }))
16111    }
16112
16113    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
16114    /// launch computes both FFN projections of a verify batch — same activation, same shape,
16115    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
16116    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
16117    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
16118    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
16119    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
16120    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
16121    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
16122    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
16123    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
16124    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
16125    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
16126    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
16127    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
16128    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16129    pub fn matmul_decode_exact_dual(
16130        &self,
16131        w0: &crate::model::GpuTensor,
16132        w1: &crate::model::GpuTensor,
16133        x: &CudaSlice<f32>,
16134        m: usize,
16135    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
16136        use crate::model::GpuTensor;
16137        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16138        let on = *ON.get_or_init(|| {
16139            std::env::var("MEMRA_SPEC_DUAL_T")
16140                .map(|v| v != "0")
16141                .unwrap_or(true)
16142        });
16143        if !on
16144            || !(2..=4).contains(&m)
16145            || std::env::var("MEMRA_NO_BATCHED").is_ok()
16146            || !self.uses_q8_1_fast(w0)
16147            || !self.uses_q8_1_fast(w1)
16148        {
16149            return Ok(None);
16150        }
16151        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
16152        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
16153        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
16154        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
16155        if !self.mmvq_supports(QT_NVFP4) {
16156            return Ok(None);
16157        }
16158        let (in_f, out_f) = (w0.in_features(), w0.out_features());
16159        if w1.in_features() != in_f || w1.out_features() != out_f {
16160            return Ok(None);
16161        }
16162        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
16163            (
16164                GpuTensor::Quant {
16165                    bytes: b0,
16166                    qtype: q0,
16167                    row_bytes: rb0,
16168                    scale: s0,
16169                    rp: rp0,
16170                    rp4: None,
16171                    ..
16172                },
16173                GpuTensor::Quant {
16174                    bytes: b1,
16175                    qtype: q1,
16176                    row_bytes: rb1,
16177                    scale: s1,
16178                    rp: rp1,
16179                    rp4: None,
16180                    ..
16181                },
16182            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
16183                (b0, b1, *rb0, *s0, *s1, *rp0)
16184            }
16185            _ => return Ok(None),
16186        };
16187        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
16188        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
16189        if std::env::var("MEMRA_DEBUG").is_ok() {
16190            static ONCE: std::sync::Once = std::sync::Once::new();
16191            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
16192        }
16193        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16194        let (y0, y1) =
16195            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
16196        let mut y0 = y0;
16197        let mut y1 = y1;
16198        if s0 != 1.0 {
16199            self.scale_inplace(&mut y0, s0, m * out_f)?;
16200        }
16201        if s1 != 1.0 {
16202            self.scale_inplace(&mut y1, s1, m * out_f)?;
16203        }
16204        Ok(Some((y0, y1)))
16205    }
16206
16207    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
16208    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
16209    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
16210    /// twins (both buffers must be the repacked layout).
16211    #[allow(clippy::too_many_arguments)]
16212    #[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
16213    pub fn qmatvec_batched_dual_raw(
16214        &self,
16215        b0: &CudaSlice<u8>,
16216        b1: &CudaSlice<u8>,
16217        aq: &CudaSlice<i8>,
16218        ad: &CudaSlice<f32>,
16219        m: usize,
16220        in_f: usize,
16221        out_f: usize,
16222        row_bytes: usize,
16223        rp: bool,
16224    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16225        const ROWS_PER_BLOCK: u32 = 4;
16226        let mcols = Self::batched_mcols(m);
16227        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
16228        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
16229        let tiny_rp1 = rp
16230            && mcols == 4
16231            && out_f <= 128
16232            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
16233        let (name, rows_per_block) = if tiny_rp1 {
16234            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
16235        } else {
16236            match (mcols, rp, m) {
16237                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
16238                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
16239                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
16240                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
16241                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
16242                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
16243                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
16244                _ => {
16245                    return Err(
16246                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
16247                    );
16248                }
16249            }
16250        };
16251        let f = self.func(name);
16252        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
16253        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
16254        let cfg = LaunchConfig {
16255            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
16256            block_dim: (32, ROWS_PER_BLOCK, 1),
16257            shared_mem_bytes: 0,
16258        };
16259        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16260        let __s_b = self.gpu.stream();
16261        let mut b = __s_b.launch_builder(&f);
16262        b.arg(b0)
16263            .arg(b1)
16264            .arg(aq)
16265            .arg(ad)
16266            .arg(&mut y0)
16267            .arg(&mut y1)
16268            .arg(&inf)
16269            .arg(&outf)
16270            .arg(&mi)
16271            .arg(&rb);
16272        unsafe {
16273            b.launch(cfg)?;
16274        }
16275        Ok((y0, y1))
16276    }
16277
16278    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
16279    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
16280    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
16281    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
16282    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
16283    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
16284    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
16285    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
16286    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
16287    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
16288    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
16289    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16290    #[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
16291    pub fn matmul_pre_dual_noscale(
16292        &self,
16293        w0: &crate::model::GpuTensor,
16294        w1: &crate::model::GpuTensor,
16295        aq: &CudaSlice<i8>,
16296        ad: &CudaSlice<f32>,
16297        m: usize,
16298    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
16299    {
16300        use crate::model::GpuTensor;
16301        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
16302            return Ok(None);
16303        }
16304        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
16305        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
16306        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
16307        // would mix dispatch families across the pair — the exact class `q8_fused_params`
16308        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
16309        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
16310        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
16311        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
16312        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
16313        if !self.mmvq_supports(QT_NVFP4) {
16314            return Ok(None);
16315        }
16316        let (in_f, out_f) = (w0.in_features(), w0.out_features());
16317        if w1.in_features() != in_f || w1.out_features() != out_f {
16318            return Ok(None);
16319        }
16320        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
16321        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
16322        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
16323        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
16324        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
16325        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
16326        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
16327        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
16328        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
16329        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
16330        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
16331        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
16332        let no_mirror =
16333            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
16334        if self.q8_ffn_fuse2_on()
16335            && no_mirror(w0)
16336            && no_mirror(w1)
16337            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
16338        {
16339            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
16340            return Ok(Some(((y0, 1.0), (y1, 1.0))));
16341        }
16342        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
16343        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
16344        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
16345        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
16346        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
16347        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
16348        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
16349        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
16350        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
16351        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
16352            let (y0, y1) =
16353                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
16354            return Ok(Some(((y0, p0.3), (y1, p1.3))));
16355        }
16356        let (b0, q0, rb0, s0, rp0) = match w0 {
16357            GpuTensor::Quant {
16358                bytes,
16359                qtype,
16360                row_bytes,
16361                scale,
16362                rp,
16363                ..
16364            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16365            _ => return Ok(None),
16366        };
16367        let (b1, q1, rb1, s1, rp1) = match w1 {
16368            GpuTensor::Quant {
16369                bytes,
16370                qtype,
16371                row_bytes,
16372                scale,
16373                rp,
16374                ..
16375            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16376            _ => return Ok(None),
16377        };
16378        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
16379            return Ok(None);
16380        }
16381        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
16382        const RPW: u32 = 2;
16383        let rows_per_block = ROWS_PER_BLOCK * RPW;
16384        let f = self.func(if rp0 {
16385            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
16386        } else {
16387            "qmatvec_nvfp4_mmvq_dual_mr2"
16388        });
16389        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
16390        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
16391        let cfg = LaunchConfig {
16392            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
16393            block_dim: (32, ROWS_PER_BLOCK, 1),
16394            shared_mem_bytes: 0,
16395        };
16396        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
16397        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
16398        // yscale args stay 1.0 here (they exist for the single-tensor callers).
16399        let one = 1.0f32;
16400        let __s_b = self.gpu.stream();
16401        let mut b = __s_b.launch_builder(&f);
16402        b.arg(b0)
16403            .arg(b1)
16404            .arg(aq)
16405            .arg(ad)
16406            .arg(&mut y0)
16407            .arg(&mut y1)
16408            .arg(&inf)
16409            .arg(&outf)
16410            .arg(&mi)
16411            .arg(&rb)
16412            .arg(&one)
16413            .arg(&one);
16414        unsafe {
16415            b.launch(cfg)?;
16416        }
16417        Ok(Some(((y0, s0), (y1, s1))))
16418    }
16419
16420    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
16421    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
16422    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
16423    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
16424    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
16425    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
16426    /// back to the three singles.
16427    #[allow(clippy::too_many_arguments)]
16428    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16429    pub fn matmul_nvfp4_fused3(
16430        &self,
16431        w0: &crate::model::GpuTensor,
16432        w1: &crate::model::GpuTensor,
16433        w2: &crate::model::GpuTensor,
16434        aq: &CudaSlice<i8>,
16435        ad: &CudaSlice<f32>,
16436        m: usize,
16437    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
16438    {
16439        use crate::model::GpuTensor;
16440        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
16441        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
16442        // verbatim, weight rows read once for all m columns, bit-identical per
16443        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
16444        // segments would re-read the weight per row" note described the grid.y=m lift,
16445        // which this twin deliberately is NOT.
16446        if !self.mmvq_supports(QT_NVFP4)
16447            || !self.uses_q8_1_fast(w0)
16448            || !self.uses_q8_1_fast(w1)
16449            || !self.uses_q8_1_fast(w2)
16450        {
16451            return Ok(None);
16452        }
16453        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
16454        // door — same family and bit-identity law as the fused4 delegate above.
16455        if (9..=16).contains(&m) {
16456            return Ok(
16457                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
16458                    Some(mut ys) => {
16459                        let y2 = ys.pop().unwrap();
16460                        let y1 = ys.pop().unwrap();
16461                        let y0 = ys.pop().unwrap();
16462                        Some((y0, y1, y2))
16463                    }
16464                    None => None,
16465                },
16466            );
16467        }
16468        if !(1..=8).contains(&m) {
16469            return Ok(None);
16470        }
16471        if m > 1 {
16472            let in_f = w0.in_features();
16473            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
16474                || !self.batched_supports(QT_NVFP4)
16475                || std::env::var("MEMRA_NO_BATCHED").is_ok()
16476                || (m > 4 && !Self::b8_enabled())
16477                || !in_f.is_multiple_of(512)
16478                || in_f / 64 > 272
16479            {
16480                return Ok(None);
16481            }
16482        }
16483        let unpack = |w: &crate::model::GpuTensor| match w {
16484            GpuTensor::Quant {
16485                bytes,
16486                qtype,
16487                scale,
16488                rp,
16489                ..
16490            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16491            _ => None,
16492        };
16493        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
16494            return Ok(None);
16495        };
16496        let in_f = w0.in_features();
16497        if w1.in_features() != in_f || w2.in_features() != in_f {
16498            return Ok(None);
16499        }
16500        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
16501        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16502        const RPW: u32 = 2;
16503        let rows_pb = ROWS_PER_BLOCK * RPW;
16504        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16505        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16506        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16507        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
16508        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
16509        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16510        // only dereferenced for the launch-arg build inside this call.
16511        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
16512        if m > 1 {
16513            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
16514            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
16515                return Ok(None);
16516            }
16517            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
16518            let cfg = LaunchConfig {
16519                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
16520                block_dim: (32, ROWS_PER_BLOCK, 1),
16521                shared_mem_bytes: 0,
16522            };
16523            let __s_b = self.gpu.stream();
16524            let mut b = __s_b.launch_builder(&f);
16525            b.arg(b0)
16526                .arg(b1)
16527                .arg(b2)
16528                .arg(aq)
16529                .arg(ad)
16530                .arg(&mut y0)
16531                .arg(&mut y1)
16532                .arg(&mut y2)
16533                .arg(&inf)
16534                .arg(&oi0)
16535                .arg(&oi1)
16536                .arg(&oi2)
16537                .arg(&mi);
16538            unsafe {
16539                b.launch(cfg)?;
16540            }
16541            return Ok(Some((y0, y1, y2)));
16542        }
16543        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
16544        let cfg = LaunchConfig {
16545            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
16546            block_dim: (32, ROWS_PER_BLOCK, 1),
16547            shared_mem_bytes: 0,
16548        };
16549        let __s_b = self.gpu.stream();
16550        let mut b = __s_b.launch_builder(&f);
16551        b.arg(b0)
16552            .arg(b1)
16553            .arg(b2)
16554            .arg(aq)
16555            .arg(ad)
16556            .arg(&mut y0)
16557            .arg(&mut y1)
16558            .arg(&mut y2)
16559            .arg(&inf)
16560            .arg(&oi0)
16561            .arg(&oi1)
16562            .arg(&oi2)
16563            .arg(&mi)
16564            .arg(&p0.1)
16565            .arg(&p1.1)
16566            .arg(&p2.1);
16567        unsafe {
16568            b.launch(cfg)?;
16569        }
16570        Ok(Some((y0, y1, y2)))
16571    }
16572
16573    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
16574    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
16575    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
16576    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
16577    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
16578    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
16579    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
16580    /// same-binary interleaved A/B arm.
16581    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16582    pub fn matmul_nvfp4_fused2(
16583        &self,
16584        w0: &crate::model::GpuTensor,
16585        w1: &crate::model::GpuTensor,
16586        aq: &CudaSlice<i8>,
16587        ad: &CudaSlice<f32>,
16588        m: usize,
16589    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
16590        use crate::model::GpuTensor;
16591        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16592        let off =
16593            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
16594        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
16595        // read serves all m rows); the fused segments would re-read the weight per row.
16596        if off
16597            || m != 1
16598            || !self.mmvq_supports(QT_NVFP4)
16599            || !self.uses_q8_1_fast(w0)
16600            || !self.uses_q8_1_fast(w1)
16601        {
16602            return Ok(None);
16603        }
16604        let unpack = |w: &crate::model::GpuTensor| match w {
16605            GpuTensor::Quant {
16606                bytes,
16607                qtype,
16608                scale,
16609                rp,
16610                ..
16611            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16612            _ => None,
16613        };
16614        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
16615            return Ok(None);
16616        };
16617        let in_f = w0.in_features();
16618        if w1.in_features() != in_f {
16619            return Ok(None);
16620        }
16621        let (o0, o1) = (w0.out_features(), w1.out_features());
16622        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16623        const RPW: u32 = 2;
16624        let rows_pb = ROWS_PER_BLOCK * RPW;
16625        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16626        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
16627        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16628        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16629        let cfg = LaunchConfig {
16630            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
16631            block_dim: (32, ROWS_PER_BLOCK, 1),
16632            shared_mem_bytes: 0,
16633        };
16634        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
16635        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16636        // only dereferenced for the launch-arg build inside this call.
16637        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
16638        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
16639        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
16640        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
16641            {
16642                use cudarc::driver::{DevicePtr, DevicePtrMut};
16643                let s = &self.gpu.stream();
16644                let (pw0, _g0) = b0.device_ptr(s);
16645                let (pw1, _g1) = b1.device_ptr(s);
16646                let (paq, _g2) = aq.device_ptr(s);
16647                let (pad, _g3) = ad.device_ptr(s);
16648                let (py0, _g4) = y0.device_ptr_mut(s);
16649                let (py1, _g5) = y1.device_ptr_mut(s);
16650                let (s0, s1) = (p0.1, p1.1);
16651                let mut ps = [
16652                    &pw0 as *const _ as *mut std::ffi::c_void,
16653                    &pw1 as *const _ as *mut _,
16654                    &paq as *const _ as *mut _,
16655                    &pad as *const _ as *mut _,
16656                    &py0 as *const _ as *mut _,
16657                    &py1 as *const _ as *mut _,
16658                    &inf as *const _ as *mut _,
16659                    &oi0 as *const _ as *mut _,
16660                    &oi1 as *const _ as *mut _,
16661                    &mi as *const _ as *mut _,
16662                    &s0 as *const _ as *mut _,
16663                    &s1 as *const _ as *mut _,
16664                ];
16665                unsafe {
16666                    self.launch_pdl(
16667                        "qmatvec_nvfp4_mmvq_fused2_rp",
16668                        cfg.grid_dim,
16669                        cfg.block_dim,
16670                        &mut ps,
16671                    )?;
16672                }
16673            }
16674            return Ok(Some((y0, y1)));
16675        }
16676        let __s_b = self.gpu.stream();
16677        let mut b = __s_b.launch_builder(&f);
16678        b.arg(b0)
16679            .arg(b1)
16680            .arg(aq)
16681            .arg(ad)
16682            .arg(&mut y0)
16683            .arg(&mut y1)
16684            .arg(&inf)
16685            .arg(&oi0)
16686            .arg(&oi1)
16687            .arg(&mi)
16688            .arg(&p0.1)
16689            .arg(&p1.1);
16690        unsafe {
16691            b.launch(cfg)?;
16692        }
16693        Ok(Some((y0, y1)))
16694    }
16695
16696    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
16697    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
16698    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
16699    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
16700    pub fn matmul_nvfp4_fused2_into(
16701        &self,
16702        w0: &crate::model::GpuTensor,
16703        w1: &crate::model::GpuTensor,
16704        aq: &CudaSlice<i8>,
16705        ad: &CudaSlice<f32>,
16706        y0: &mut CudaSlice<f32>,
16707        y1: &mut CudaSlice<f32>,
16708    ) -> Result<bool, Box<dyn std::error::Error>> {
16709        use crate::model::GpuTensor;
16710        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16711        let off =
16712            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
16713        if off
16714            || !self.mmvq_supports(QT_NVFP4)
16715            || !self.uses_q8_1_fast(w0)
16716            || !self.uses_q8_1_fast(w1)
16717        {
16718            return Ok(false);
16719        }
16720        let unpack = |w: &crate::model::GpuTensor| match w {
16721            GpuTensor::Quant {
16722                bytes,
16723                qtype,
16724                scale,
16725                rp,
16726                ..
16727            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16728            _ => None,
16729        };
16730        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
16731            return Ok(false);
16732        };
16733        let in_f = w0.in_features();
16734        if w1.in_features() != in_f {
16735            return Ok(false);
16736        }
16737        let (o0, o1) = (w0.out_features(), w1.out_features());
16738        if y0.len() < o0 || y1.len() < o1 {
16739            return Ok(false);
16740        }
16741        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16742        const RPW: u32 = 2;
16743        let rows_pb = ROWS_PER_BLOCK * RPW;
16744        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16745        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
16746        let cfg = LaunchConfig {
16747            grid_dim: (nb(o0) + nb(o1), 1, 1),
16748            block_dim: (32, ROWS_PER_BLOCK, 1),
16749            shared_mem_bytes: 0,
16750        };
16751        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
16752        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16753        // only dereferenced for the launch-arg build inside this call.
16754        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
16755        let __s_b = self.gpu.stream();
16756        let mut b = __s_b.launch_builder(&f);
16757        b.arg(b0)
16758            .arg(b1)
16759            .arg(aq)
16760            .arg(ad)
16761            .arg(&mut *y0)
16762            .arg(&mut *y1)
16763            .arg(&inf)
16764            .arg(&oi0)
16765            .arg(&oi1)
16766            .arg(&mi)
16767            .arg(&p0.1)
16768            .arg(&p1.1);
16769        unsafe {
16770            b.launch(cfg)?;
16771        }
16772        Ok(true)
16773    }
16774
16775    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
16776    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
16777    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
16778    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
16779    #[allow(clippy::type_complexity)]
16780    #[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
16781    pub fn matmul_nvfp4_fused4(
16782        &self,
16783        w0: &crate::model::GpuTensor,
16784        w1: &crate::model::GpuTensor,
16785        w2: &crate::model::GpuTensor,
16786        w3: &crate::model::GpuTensor,
16787        aq: &CudaSlice<i8>,
16788        ad: &CudaSlice<f32>,
16789        m: usize,
16790    ) -> Result<
16791        Option<(
16792            CudaSlice<f32>,
16793            CudaSlice<f32>,
16794            CudaSlice<f32>,
16795            CudaSlice<f32>,
16796        )>,
16797        Box<dyn std::error::Error>,
16798    > {
16799        use crate::model::GpuTensor;
16800        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
16801        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
16802        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
16803        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
16804        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
16805        // Admission mirrors the singles' batched gates below.
16806        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
16807            || !self.mmvq_supports(QT_NVFP4)
16808            || !self.uses_q8_1_fast(w0)
16809            || !self.uses_q8_1_fast(w1)
16810            || !self.uses_q8_1_fast(w2)
16811            || !self.uses_q8_1_fast(w3)
16812        {
16813            return Ok(None);
16814        }
16815        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
16816        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
16817        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
16818        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
16819        if (9..=16).contains(&m) {
16820            return Ok(
16821                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
16822                    Some(mut ys) => {
16823                        let y3 = ys.pop().unwrap();
16824                        let y2 = ys.pop().unwrap();
16825                        let y1 = ys.pop().unwrap();
16826                        let y0 = ys.pop().unwrap();
16827                        Some((y0, y1, y2, y3))
16828                    }
16829                    None => None,
16830                },
16831            );
16832        }
16833        if !(1..=8).contains(&m) {
16834            return Ok(None);
16835        }
16836        if m > 1 {
16837            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
16838            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
16839            let in_f = w0.in_features();
16840            if !self.batched_supports(QT_NVFP4)
16841                || std::env::var("MEMRA_NO_BATCHED").is_ok()
16842                || (m > 4 && !Self::b8_enabled())
16843                || !in_f.is_multiple_of(512)
16844                || in_f / 64 > 272
16845            {
16846                return Ok(None);
16847            }
16848        }
16849        let unpack = |w: &crate::model::GpuTensor| match w {
16850            GpuTensor::Quant {
16851                bytes,
16852                qtype,
16853                scale,
16854                rp,
16855                ..
16856            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
16857            _ => None,
16858        };
16859        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
16860            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
16861        else {
16862            return Ok(None);
16863        };
16864        let in_f = w0.in_features();
16865        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
16866            return Ok(None);
16867        }
16868        let (o0, o1, o2, o3) = (
16869            w0.out_features(),
16870            w1.out_features(),
16871            w2.out_features(),
16872            w3.out_features(),
16873        );
16874        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
16875        const RPW: u32 = 2;
16876        let rows_pb = ROWS_PER_BLOCK * RPW;
16877        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
16878        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
16879        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
16880        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
16881        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
16882        let (inf, oi0, oi1, oi2, oi3, mi) = (
16883            in_f as i32,
16884            o0 as i32,
16885            o1 as i32,
16886            o2 as i32,
16887            o3 as i32,
16888            m as i32,
16889        );
16890        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
16891        // only dereferenced for the launch-arg build inside this call.
16892        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
16893        if m > 1 {
16894            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
16895            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
16896            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
16897                return Ok(None);
16898            }
16899            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
16900            let cfg = LaunchConfig {
16901                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
16902                block_dim: (32, ROWS_PER_BLOCK, 1),
16903                shared_mem_bytes: 0,
16904            };
16905            let __s_b = self.gpu.stream();
16906            let mut b = __s_b.launch_builder(&f);
16907            b.arg(b0)
16908                .arg(b1)
16909                .arg(b2)
16910                .arg(b3)
16911                .arg(aq)
16912                .arg(ad)
16913                .arg(&mut y0)
16914                .arg(&mut y1)
16915                .arg(&mut y2)
16916                .arg(&mut y3)
16917                .arg(&inf)
16918                .arg(&oi0)
16919                .arg(&oi1)
16920                .arg(&oi2)
16921                .arg(&oi3)
16922                .arg(&mi);
16923            unsafe {
16924                b.launch(cfg)?;
16925            }
16926            return Ok(Some((y0, y1, y2, y3)));
16927        }
16928        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
16929        let cfg = LaunchConfig {
16930            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
16931            block_dim: (32, ROWS_PER_BLOCK, 1),
16932            shared_mem_bytes: 0,
16933        };
16934        let __s_b = self.gpu.stream();
16935        let mut b = __s_b.launch_builder(&f);
16936        b.arg(b0)
16937            .arg(b1)
16938            .arg(b2)
16939            .arg(b3)
16940            .arg(aq)
16941            .arg(ad)
16942            .arg(&mut y0)
16943            .arg(&mut y1)
16944            .arg(&mut y2)
16945            .arg(&mut y3)
16946            .arg(&inf)
16947            .arg(&oi0)
16948            .arg(&oi1)
16949            .arg(&oi2)
16950            .arg(&oi3)
16951            .arg(&mi)
16952            .arg(&p0.1)
16953            .arg(&p1.1)
16954            .arg(&p2.1)
16955            .arg(&p3.1);
16956        unsafe {
16957            b.launch(cfg)?;
16958        }
16959        Ok(Some((y0, y1, y2, y3)))
16960    }
16961
16962    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
16963    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
16964    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
16965    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
16966    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
16967    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
16968    /// back to the per-tensor path.
16969    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
16970    pub fn matmul_q8_fused2(
16971        &self,
16972        w0: &crate::model::GpuTensor,
16973        w1: &crate::model::GpuTensor,
16974        aq: &CudaSlice<i8>,
16975        ad: &CudaSlice<f32>,
16976    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
16977        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
16978        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
16979        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
16980        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
16981        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
16982        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
16983            return Ok(Some(self.e4m3_fused2_core(
16984                p0.0,
16985                p1.0,
16986                aq,
16987                ad,
16988                w0.in_features(),
16989                p0.1,
16990                p1.1,
16991                p0.2,
16992                p0.3,
16993                p1.3,
16994            )?));
16995        }
16996        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
16997            return Ok(None);
16998        };
16999        Ok(Some(self.q8_fused2_core(
17000            p0.0,
17001            p1.0,
17002            aq,
17003            ad,
17004            w0.in_features(),
17005            p0.1,
17006            p1.1,
17007            p0.2,
17008        )?))
17009    }
17010
17011    #[allow(clippy::too_many_arguments)]
17012    fn q8_fused2_core(
17013        &self,
17014        b0: &CudaSlice<u8>,
17015        b1: &CudaSlice<u8>,
17016        aq: &CudaSlice<i8>,
17017        ad: &CudaSlice<f32>,
17018        in_f: usize,
17019        out0: usize,
17020        out1: usize,
17021        row_bytes: usize,
17022    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17023        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
17024        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
17025        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
17026        let f = self.func("qmatvec_q8_0_mmvq_fused2");
17027        let mut y0 = self.alloc_uninit::<f32>(out0)?;
17028        let mut y1 = self.alloc_uninit::<f32>(out1)?;
17029        let cfg = LaunchConfig {
17030            grid_dim: (nb0 + nb1, 1, 1),
17031            block_dim: (32, ROWS_PER_BLOCK, 1),
17032            shared_mem_bytes: 0,
17033        };
17034        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
17035        let __s_b = self.gpu.stream();
17036        let mut b = __s_b.launch_builder(&f);
17037        b.arg(b0)
17038            .arg(b1)
17039            .arg(aq)
17040            .arg(ad)
17041            .arg(&mut y0)
17042            .arg(&mut y1)
17043            .arg(&inf)
17044            .arg(&o0)
17045            .arg(&o1)
17046            .arg(&rbl);
17047        unsafe {
17048            b.launch(cfg)?;
17049        }
17050        Ok((y0, y1))
17051    }
17052
17053    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
17054    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
17055    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
17056    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
17057    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
17058    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17059    pub fn matmul_q8_fused2_x(
17060        &self,
17061        w0: &crate::model::GpuTensor,
17062        w1: &crate::model::GpuTensor,
17063        x: &CudaSlice<f32>,
17064    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17065        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
17066            return Ok(None);
17067        }
17068        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17069            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
17070            return Ok(Some(self.e4m3_fused2_core(
17071                p0.0,
17072                p1.0,
17073                &aq,
17074                &ad,
17075                w0.in_features(),
17076                p0.1,
17077                p1.1,
17078                p0.2,
17079                p0.3,
17080                p1.3,
17081            )?));
17082        }
17083        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
17084            return Ok(None);
17085        };
17086        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
17087        Ok(Some(self.q8_fused2_core(
17088            p0.0,
17089            p1.0,
17090            &aq,
17091            &ad,
17092            w0.in_features(),
17093            p0.1,
17094            p1.1,
17095            p0.2,
17096        )?))
17097    }
17098
17099    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
17100    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
17101    #[allow(clippy::too_many_arguments)]
17102    pub fn qmatvec_q8_fused2_raw(
17103        &self,
17104        b0: &CudaSlice<u8>,
17105        b1: &CudaSlice<u8>,
17106        x: &CudaSlice<f32>,
17107        in_f: usize,
17108        out0: usize,
17109        out1: usize,
17110        row_bytes: usize,
17111    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17112        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
17113        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
17114    }
17115
17116    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
17117    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
17118    /// (tensor,row) to three separate m=1 MMVQ launches.
17119    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
17120    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
17121    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17122    pub fn matmul_q4_fused3(
17123        &self,
17124        w0: &crate::model::GpuTensor,
17125        w1: &crate::model::GpuTensor,
17126        w2: &crate::model::GpuTensor,
17127        aq: &CudaSlice<i8>,
17128        ad: &CudaSlice<f32>,
17129    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17130    {
17131        use crate::model::GpuTensor;
17132        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17133            match w {
17134                GpuTensor::Quant {
17135                    qtype, row_bytes, ..
17136                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17137                _ => None,
17138            }
17139        };
17140        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
17141            return Ok(None);
17142        };
17143        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17144            return Ok(None);
17145        }
17146        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
17147        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
17148        // the separate matvecs (each routes its own rp).
17149        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17150            match w {
17151                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17152                    Some(m) => (m, true),
17153                    None => (bytes, *rp),
17154                },
17155                _ => unreachable!(),
17156            }
17157        }
17158        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17159        if rp0 != rp1 || rp1 != rp2 {
17160            return Ok(None);
17161        }
17162        let rp = rp0;
17163        let rpb: u32 = 4;
17164        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
17165        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
17166        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
17167        let mr1 = rp && Self::q40_mr1_on();
17168        let nb = |o: usize| {
17169            if mr1 {
17170                (o as u32).div_ceil(rpb)
17171            } else {
17172                (o as u32).div_ceil(2).div_ceil(rpb)
17173            }
17174        };
17175        let grid = nb(o0) + nb(o1) + nb(o2);
17176        let mut y0 = self.alloc_uninit::<f32>(o0)?;
17177        let mut y1 = self.alloc_uninit::<f32>(o1)?;
17178        let mut y2 = self.alloc_uninit::<f32>(o2)?;
17179        let f = self.func(if mr1 {
17180            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
17181        } else if rp {
17182            "qmatvec_q4_0_mmvq_fused3_rp"
17183        } else {
17184            "qmatvec_q4_0_mmvq_fused3"
17185        });
17186        let cfg = LaunchConfig {
17187            grid_dim: (grid, 1, 1),
17188            block_dim: (32, rpb, 1),
17189            shared_mem_bytes: 0,
17190        };
17191        let inf = w0.in_features() as i32;
17192        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
17193        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
17194        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
17195        // variant may take the programmatic-serialization launch.
17196        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17197            {
17198                use cudarc::driver::{DevicePtr, DevicePtrMut};
17199                let s = &self.gpu.stream();
17200                let (p0, _g0) = b0.device_ptr(s);
17201                let (p1, _g1) = b1.device_ptr(s);
17202                let (p2, _g2) = b2.device_ptr(s);
17203                let (paq, _g3) = aq.device_ptr(s);
17204                let (pad, _g4) = ad.device_ptr(s);
17205                let (py0, _g5) = y0.device_ptr_mut(s);
17206                let (py1, _g6) = y1.device_ptr_mut(s);
17207                let (py2, _g7) = y2.device_ptr_mut(s);
17208                let mut ps = [
17209                    &p0 as *const _ as *mut std::ffi::c_void,
17210                    &p1 as *const _ as *mut _,
17211                    &p2 as *const _ as *mut _,
17212                    &paq as *const _ as *mut _,
17213                    &pad as *const _ as *mut _,
17214                    &py0 as *const _ as *mut _,
17215                    &py1 as *const _ as *mut _,
17216                    &py2 as *const _ as *mut _,
17217                    &inf as *const _ as *mut _,
17218                    &oo0 as *const _ as *mut _,
17219                    &oo1 as *const _ as *mut _,
17220                    &oo2 as *const _ as *mut _,
17221                    &r0 as *const _ as *mut _,
17222                    &r1 as *const _ as *mut _,
17223                    &r2 as *const _ as *mut _,
17224                ];
17225                unsafe {
17226                    self.launch_pdl(
17227                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
17228                        (grid, 1, 1),
17229                        (32, rpb, 1),
17230                        &mut ps,
17231                    )?;
17232                }
17233            }
17234            return Ok(Some((y0, y1, y2)));
17235        }
17236        let __s_b = self.gpu.stream();
17237        let mut b = __s_b.launch_builder(&f);
17238        b.arg(b0)
17239            .arg(b1)
17240            .arg(b2)
17241            .arg(aq)
17242            .arg(ad)
17243            .arg(&mut y0)
17244            .arg(&mut y1)
17245            .arg(&mut y2)
17246            .arg(&inf)
17247            .arg(&oo0)
17248            .arg(&oo1)
17249            .arg(&oo2)
17250            .arg(&r0)
17251            .arg(&r1)
17252            .arg(&r2);
17253        unsafe {
17254            b.launch(cfg)?;
17255        }
17256        Ok(Some((y0, y1, y2)))
17257    }
17258
17259    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
17260    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
17261    #[allow(clippy::too_many_arguments)]
17262    pub fn matmul_q4_fused3_into(
17263        &self,
17264        w0: &crate::model::GpuTensor,
17265        w1: &crate::model::GpuTensor,
17266        w2: &crate::model::GpuTensor,
17267        aq: &CudaSlice<i8>,
17268        ad: &CudaSlice<f32>,
17269        y0: &mut CudaSlice<f32>,
17270        y1: &mut CudaSlice<f32>,
17271        y2: &mut CudaSlice<f32>,
17272    ) -> Result<bool, Box<dyn std::error::Error>> {
17273        use crate::model::GpuTensor;
17274        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17275            match w {
17276                GpuTensor::Quant {
17277                    qtype, row_bytes, ..
17278                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17279                _ => None,
17280            }
17281        };
17282        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
17283            return Ok(false);
17284        };
17285        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17286            return Ok(false);
17287        }
17288        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17289            match w {
17290                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17291                    Some(m) => (m, true),
17292                    None => (bytes, *rp),
17293                },
17294                _ => unreachable!(),
17295            }
17296        }
17297        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17298        if rp0 != rp1 || rp1 != rp2 {
17299            return Ok(false);
17300        }
17301        let rp = rp0;
17302        let rpb: u32 = 4;
17303        let mr1 = rp && Self::q40_mr1_on();
17304        let nb = |o: usize| {
17305            if mr1 {
17306                (o as u32).div_ceil(rpb)
17307            } else {
17308                (o as u32).div_ceil(2).div_ceil(rpb)
17309            }
17310        };
17311        let grid = nb(o0) + nb(o1) + nb(o2);
17312        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
17313        let f = self.func(if mr1 {
17314            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
17315        } else if rp {
17316            "qmatvec_q4_0_mmvq_fused3_rp"
17317        } else {
17318            "qmatvec_q4_0_mmvq_fused3"
17319        });
17320        let cfg = LaunchConfig {
17321            grid_dim: (grid, 1, 1),
17322            block_dim: (32, rpb, 1),
17323            shared_mem_bytes: 0,
17324        };
17325        let inf = w0.in_features() as i32;
17326        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
17327        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
17328        // PDL wave-A: identical to the owned twin (capture-lane parity).
17329        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17330            use cudarc::driver::{DevicePtr, DevicePtrMut};
17331            let s = &self.gpu.stream();
17332            let (p0, _g0) = b0.device_ptr(s);
17333            let (p1, _g1) = b1.device_ptr(s);
17334            let (p2, _g2) = b2.device_ptr(s);
17335            let (paq, _g3) = aq.device_ptr(s);
17336            let (pad, _g4) = ad.device_ptr(s);
17337            let (py0, _g5) = y0.device_ptr_mut(s);
17338            let (py1, _g6) = y1.device_ptr_mut(s);
17339            let (py2, _g7) = y2.device_ptr_mut(s);
17340            let mut ps = [
17341                &p0 as *const _ as *mut std::ffi::c_void,
17342                &p1 as *const _ as *mut _,
17343                &p2 as *const _ as *mut _,
17344                &paq as *const _ as *mut _,
17345                &pad as *const _ as *mut _,
17346                &py0 as *const _ as *mut _,
17347                &py1 as *const _ as *mut _,
17348                &py2 as *const _ as *mut _,
17349                &inf as *const _ as *mut _,
17350                &oo0 as *const _ as *mut _,
17351                &oo1 as *const _ as *mut _,
17352                &oo2 as *const _ as *mut _,
17353                &r0 as *const _ as *mut _,
17354                &r1 as *const _ as *mut _,
17355                &r2 as *const _ as *mut _,
17356            ];
17357            unsafe {
17358                self.launch_pdl(
17359                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
17360                    (grid, 1, 1),
17361                    (32, rpb, 1),
17362                    &mut ps,
17363                )?;
17364            }
17365            return Ok(true);
17366        }
17367        let __s_b = self.gpu.stream();
17368        let mut b = __s_b.launch_builder(&f);
17369        b.arg(b0)
17370            .arg(b1)
17371            .arg(b2)
17372            .arg(aq)
17373            .arg(ad)
17374            .arg(&mut *y0)
17375            .arg(&mut *y1)
17376            .arg(&mut *y2)
17377            .arg(&inf)
17378            .arg(&oo0)
17379            .arg(&oo1)
17380            .arg(&oo2)
17381            .arg(&r0)
17382            .arg(&r1)
17383            .arg(&r2);
17384        unsafe {
17385            b.launch(cfg)?;
17386        }
17387        Ok(true)
17388    }
17389
17390    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
17391    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17392    pub fn matmul_q4_fused2(
17393        &self,
17394        w0: &crate::model::GpuTensor,
17395        w1: &crate::model::GpuTensor,
17396        aq: &CudaSlice<i8>,
17397        ad: &CudaSlice<f32>,
17398    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17399        use crate::model::GpuTensor;
17400        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17401            match w {
17402                GpuTensor::Quant {
17403                    qtype, row_bytes, ..
17404                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17405                _ => None,
17406            }
17407        };
17408        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
17409            return Ok(None);
17410        };
17411        if w0.in_features() != w1.in_features() {
17412            return Ok(None);
17413        }
17414        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
17415        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17416            match w {
17417                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17418                    Some(m) => (m, true),
17419                    None => (bytes, *rp),
17420                },
17421                _ => unreachable!(),
17422            }
17423        }
17424        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17425        if rp0 != rp1 {
17426            return Ok(None);
17427        }
17428        let rp = rp0;
17429        let rpb: u32 = 4;
17430        // mr1 twin — see matmul_q4_fused3.
17431        let mr1 = rp && Self::q40_mr1_on();
17432        let nb = |o: usize| {
17433            if mr1 {
17434                (o as u32).div_ceil(rpb)
17435            } else {
17436                (o as u32).div_ceil(2).div_ceil(rpb)
17437            }
17438        };
17439        let grid = nb(o0) + nb(o1);
17440        let mut y0 = self.alloc_uninit::<f32>(o0)?;
17441        let mut y1 = self.alloc_uninit::<f32>(o1)?;
17442        let f = self.func(if mr1 {
17443            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
17444        } else if rp {
17445            "qmatvec_q4_0_mmvq_fused2_rp"
17446        } else {
17447            "qmatvec_q4_0_mmvq_fused2"
17448        });
17449        let cfg = LaunchConfig {
17450            grid_dim: (grid, 1, 1),
17451            block_dim: (32, rpb, 1),
17452            shared_mem_bytes: 0,
17453        };
17454        let inf = w0.in_features() as i32;
17455        let (oo0, oo1) = (o0 as i32, o1 as i32);
17456        let (r0, r1) = (rb0 as i64, rb1 as i64);
17457        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
17458        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17459            {
17460                use cudarc::driver::{DevicePtr, DevicePtrMut};
17461                let s = &self.gpu.stream();
17462                let (p0, _g0) = b0.device_ptr(s);
17463                let (p1, _g1) = b1.device_ptr(s);
17464                let (paq, _g2) = aq.device_ptr(s);
17465                let (pad, _g3) = ad.device_ptr(s);
17466                let (py0, _g4) = y0.device_ptr_mut(s);
17467                let (py1, _g5) = y1.device_ptr_mut(s);
17468                let mut ps = [
17469                    &p0 as *const _ as *mut std::ffi::c_void,
17470                    &p1 as *const _ as *mut _,
17471                    &paq as *const _ as *mut _,
17472                    &pad as *const _ as *mut _,
17473                    &py0 as *const _ as *mut _,
17474                    &py1 as *const _ as *mut _,
17475                    &inf as *const _ as *mut _,
17476                    &oo0 as *const _ as *mut _,
17477                    &oo1 as *const _ as *mut _,
17478                    &r0 as *const _ as *mut _,
17479                    &r1 as *const _ as *mut _,
17480                ];
17481                unsafe {
17482                    self.launch_pdl(
17483                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
17484                        (grid, 1, 1),
17485                        (32, rpb, 1),
17486                        &mut ps,
17487                    )?;
17488                }
17489            }
17490            return Ok(Some((y0, y1)));
17491        }
17492        let __s_b = self.gpu.stream();
17493        let mut b = __s_b.launch_builder(&f);
17494        b.arg(b0)
17495            .arg(b1)
17496            .arg(aq)
17497            .arg(ad)
17498            .arg(&mut y0)
17499            .arg(&mut y1)
17500            .arg(&inf)
17501            .arg(&oo0)
17502            .arg(&oo1)
17503            .arg(&r0)
17504            .arg(&r1);
17505        unsafe {
17506            b.launch(cfg)?;
17507        }
17508        Ok(Some((y0, y1)))
17509    }
17510
17511    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
17512    pub fn matmul_q4_fused2_into(
17513        &self,
17514        w0: &crate::model::GpuTensor,
17515        w1: &crate::model::GpuTensor,
17516        aq: &CudaSlice<i8>,
17517        ad: &CudaSlice<f32>,
17518        y0: &mut CudaSlice<f32>,
17519        y1: &mut CudaSlice<f32>,
17520    ) -> Result<bool, Box<dyn std::error::Error>> {
17521        use crate::model::GpuTensor;
17522        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17523            match w {
17524                GpuTensor::Quant {
17525                    qtype, row_bytes, ..
17526                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17527                _ => None,
17528            }
17529        };
17530        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
17531            return Ok(false);
17532        };
17533        if w0.in_features() != w1.in_features() {
17534            return Ok(false);
17535        }
17536        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17537            match w {
17538                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17539                    Some(m) => (m, true),
17540                    None => (bytes, *rp),
17541                },
17542                _ => unreachable!(),
17543            }
17544        }
17545        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17546        if rp0 != rp1 {
17547            return Ok(false);
17548        }
17549        let rp = rp0;
17550        let rpb: u32 = 4;
17551        let mr1 = rp && Self::q40_mr1_on();
17552        let nb = |o: usize| {
17553            if mr1 {
17554                (o as u32).div_ceil(rpb)
17555            } else {
17556                (o as u32).div_ceil(2).div_ceil(rpb)
17557            }
17558        };
17559        let grid = nb(o0) + nb(o1);
17560        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
17561        let f = self.func(if mr1 {
17562            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
17563        } else if rp {
17564            "qmatvec_q4_0_mmvq_fused2_rp"
17565        } else {
17566            "qmatvec_q4_0_mmvq_fused2"
17567        });
17568        let cfg = LaunchConfig {
17569            grid_dim: (grid, 1, 1),
17570            block_dim: (32, rpb, 1),
17571            shared_mem_bytes: 0,
17572        };
17573        let inf = w0.in_features() as i32;
17574        let (oo0, oo1) = (o0 as i32, o1 as i32);
17575        let (r0, r1) = (rb0 as i64, rb1 as i64);
17576        // PDL wave-A: identical to the owned twin (capture-lane parity).
17577        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
17578            use cudarc::driver::{DevicePtr, DevicePtrMut};
17579            let s = &self.gpu.stream();
17580            let (p0, _g0) = b0.device_ptr(s);
17581            let (p1, _g1) = b1.device_ptr(s);
17582            let (paq, _g2) = aq.device_ptr(s);
17583            let (pad, _g3) = ad.device_ptr(s);
17584            let (py0, _g4) = y0.device_ptr_mut(s);
17585            let (py1, _g5) = y1.device_ptr_mut(s);
17586            let mut ps = [
17587                &p0 as *const _ as *mut std::ffi::c_void,
17588                &p1 as *const _ as *mut _,
17589                &paq as *const _ as *mut _,
17590                &pad as *const _ as *mut _,
17591                &py0 as *const _ as *mut _,
17592                &py1 as *const _ as *mut _,
17593                &inf as *const _ as *mut _,
17594                &oo0 as *const _ as *mut _,
17595                &oo1 as *const _ as *mut _,
17596                &r0 as *const _ as *mut _,
17597                &r1 as *const _ as *mut _,
17598            ];
17599            unsafe {
17600                self.launch_pdl(
17601                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
17602                    (grid, 1, 1),
17603                    (32, rpb, 1),
17604                    &mut ps,
17605                )?;
17606            }
17607            return Ok(true);
17608        }
17609        let __s_b = self.gpu.stream();
17610        let mut b = __s_b.launch_builder(&f);
17611        b.arg(b0)
17612            .arg(b1)
17613            .arg(aq)
17614            .arg(ad)
17615            .arg(&mut *y0)
17616            .arg(&mut *y1)
17617            .arg(&inf)
17618            .arg(&oo0)
17619            .arg(&oo1)
17620            .arg(&r0)
17621            .arg(&r1);
17622        unsafe {
17623            b.launch(cfg)?;
17624        }
17625        Ok(true)
17626    }
17627
17628    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
17629    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
17630    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
17631    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
17632    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17633    pub fn matmul_q4_fused2_batched(
17634        &self,
17635        w0: &crate::model::GpuTensor,
17636        w1: &crate::model::GpuTensor,
17637        aq: &CudaSlice<i8>,
17638        ad: &CudaSlice<f32>,
17639        m: usize,
17640    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17641        use crate::model::GpuTensor;
17642        if !(2..=8).contains(&m) {
17643            return Ok(None);
17644        }
17645        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
17646            match w {
17647                GpuTensor::Quant {
17648                    qtype, row_bytes, ..
17649                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
17650                _ => None,
17651            }
17652        };
17653        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
17654            return Ok(None);
17655        };
17656        if w0.in_features() != w1.in_features() {
17657            return Ok(None);
17658        }
17659        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17660            match w {
17661                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17662                    Some(mr) => (mr, true),
17663                    None => (bytes, *rp),
17664                },
17665                _ => unreachable!(),
17666            }
17667        }
17668        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
17669        if !rp0 || !rp1 {
17670            return Ok(None);
17671        }
17672        let mcols = Self::batched_mcols(m);
17673        let rpb: u32 = 4;
17674        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
17675        let grid = nb(o0) + nb(o1);
17676        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17677        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17678        let f = self.func(match mcols {
17679            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
17680            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
17681            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
17682        });
17683        let cfg = LaunchConfig {
17684            grid_dim: (grid, 1, 1),
17685            block_dim: (32, rpb, 1),
17686            shared_mem_bytes: 0,
17687        };
17688        let inf = w0.in_features() as i32;
17689        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
17690        let rb = rb0 as i64;
17691        let __s_b = self.gpu.stream();
17692        let mut b = __s_b.launch_builder(&f);
17693        b.arg(b0)
17694            .arg(b1)
17695            .arg(aq)
17696            .arg(ad)
17697            .arg(&mut y0)
17698            .arg(&mut y1)
17699            .arg(&inf)
17700            .arg(&oo0)
17701            .arg(&oo1)
17702            .arg(&mi)
17703            .arg(&rb);
17704        unsafe {
17705            b.launch(cfg)?;
17706        }
17707        Ok(Some((y0, y1)))
17708    }
17709
17710    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
17711    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
17712    #[allow(clippy::too_many_arguments)]
17713    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17714    pub fn matmul_q4_fused3_batched(
17715        &self,
17716        w0: &crate::model::GpuTensor,
17717        w1: &crate::model::GpuTensor,
17718        w2: &crate::model::GpuTensor,
17719        aq: &CudaSlice<i8>,
17720        ad: &CudaSlice<f32>,
17721        m: usize,
17722    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17723    {
17724        use crate::model::GpuTensor;
17725        if !(2..=8).contains(&m) {
17726            return Ok(None);
17727        }
17728        let q4 = |w: &GpuTensor| -> Option<usize> {
17729            match w {
17730                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
17731                _ => None,
17732            }
17733        };
17734        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
17735            return Ok(None);
17736        };
17737        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
17738            return Ok(None);
17739        }
17740        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
17741            match w {
17742                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
17743                    Some(mr) => (mr, true),
17744                    None => (bytes, *rp),
17745                },
17746                _ => unreachable!(),
17747            }
17748        }
17749        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
17750        if !rp0 || !rp1 || !rp2 {
17751            return Ok(None);
17752        }
17753        let mcols = Self::batched_mcols(m);
17754        let rpb: u32 = 4;
17755        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
17756        let grid = nb(o0) + nb(o1) + nb(o2);
17757        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17758        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17759        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
17760        let f = self.func(match mcols {
17761            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
17762            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
17763            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
17764        });
17765        let cfg = LaunchConfig {
17766            grid_dim: (grid, 1, 1),
17767            block_dim: (32, rpb, 1),
17768            shared_mem_bytes: 0,
17769        };
17770        let inf = w0.in_features() as i32;
17771        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
17772        let rb = 0i64;
17773        let __s_b = self.gpu.stream();
17774        let mut b = __s_b.launch_builder(&f);
17775        b.arg(b0)
17776            .arg(b1)
17777            .arg(b2)
17778            .arg(aq)
17779            .arg(ad)
17780            .arg(&mut y0)
17781            .arg(&mut y1)
17782            .arg(&mut y2)
17783            .arg(&inf)
17784            .arg(&oo0)
17785            .arg(&oo1)
17786            .arg(&oo2)
17787            .arg(&mi)
17788            .arg(&rb);
17789        unsafe {
17790            b.launch(cfg)?;
17791        }
17792        Ok(Some((y0, y1, y2)))
17793    }
17794
17795    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17796    pub fn matmul_q8_fused3(
17797        &self,
17798        w0: &crate::model::GpuTensor,
17799        w1: &crate::model::GpuTensor,
17800        w2: &crate::model::GpuTensor,
17801        aq: &CudaSlice<i8>,
17802        ad: &CudaSlice<f32>,
17803    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17804    {
17805        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
17806        // are per-tensor FP8, so native residency without this arm meant three separate launches.
17807        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
17808            return Ok(Some(self.e4m3_fused3_core(
17809                p0.0,
17810                p1.0,
17811                p2.0,
17812                aq,
17813                ad,
17814                w0.in_features(),
17815                p0.1,
17816                p1.1,
17817                p2.1,
17818                p0.2,
17819                p0.3,
17820                p1.3,
17821                p2.3,
17822            )?));
17823        }
17824        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
17825            return Ok(None);
17826        };
17827        Ok(Some(self.q8_fused3_core(
17828            p0.0,
17829            p1.0,
17830            p2.0,
17831            aq,
17832            ad,
17833            w0.in_features(),
17834            p0.1,
17835            p1.1,
17836            p2.1,
17837            p0.2,
17838        )?))
17839    }
17840
17841    #[allow(clippy::too_many_arguments)]
17842    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17843    fn q8_fused3_core(
17844        &self,
17845        b0: &CudaSlice<u8>,
17846        b1: &CudaSlice<u8>,
17847        b2: &CudaSlice<u8>,
17848        aq: &CudaSlice<i8>,
17849        ad: &CudaSlice<f32>,
17850        in_f: usize,
17851        out0: usize,
17852        out1: usize,
17853        out2: usize,
17854        row_bytes: usize,
17855    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17856        const ROWS_PER_BLOCK: u32 = 4;
17857        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
17858        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
17859        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
17860        let f = self.func("qmatvec_q8_0_mmvq_fused3");
17861        let mut y0 = self.alloc_uninit::<f32>(out0)?;
17862        let mut y1 = self.alloc_uninit::<f32>(out1)?;
17863        let mut y2 = self.alloc_uninit::<f32>(out2)?;
17864        let cfg = LaunchConfig {
17865            grid_dim: (nb0 + nb1 + nb2, 1, 1),
17866            block_dim: (32, ROWS_PER_BLOCK, 1),
17867            shared_mem_bytes: 0,
17868        };
17869        let (inf, o0, o1, o2, rbl) = (
17870            in_f as i32,
17871            out0 as i32,
17872            out1 as i32,
17873            out2 as i32,
17874            row_bytes as i64,
17875        );
17876        let __s_b = self.gpu.stream();
17877        let mut b = __s_b.launch_builder(&f);
17878        b.arg(b0)
17879            .arg(b1)
17880            .arg(b2)
17881            .arg(aq)
17882            .arg(ad)
17883            .arg(&mut y0)
17884            .arg(&mut y1)
17885            .arg(&mut y2)
17886            .arg(&inf)
17887            .arg(&o0)
17888            .arg(&o1)
17889            .arg(&o2)
17890            .arg(&rbl);
17891        unsafe {
17892            b.launch(cfg)?;
17893        }
17894        Ok((y0, y1, y2))
17895    }
17896
17897    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
17898    #[allow(clippy::too_many_arguments)]
17899    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17900    pub fn qmatvec_q8_fused3_raw(
17901        &self,
17902        b0: &CudaSlice<u8>,
17903        b1: &CudaSlice<u8>,
17904        b2: &CudaSlice<u8>,
17905        x: &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        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
17913        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
17914    }
17915
17916    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
17917    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
17918    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
17919    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
17920    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
17921    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
17922    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
17923    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
17924    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
17925    /// twin must not introduce a batched program the reference path would not run).
17926    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17927    pub fn matmul_q8_fused2_t(
17928        &self,
17929        w0: &crate::model::GpuTensor,
17930        w1: &crate::model::GpuTensor,
17931        aq: &CudaSlice<i8>,
17932        ad: &CudaSlice<f32>,
17933        m: usize,
17934    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17935        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
17936        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
17937        // fuses too — same template body, still bit-identical to the two _b8 launches.
17938        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
17939            return Ok(None);
17940        }
17941        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
17942        // so the fused b8 launch would introduce a batched program the reference path would not run.
17943        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17944            if m > 4 && !Self::b8_enabled() {
17945                return Ok(None);
17946            }
17947            return Ok(Some(self.e4m3_fused2_t_core(
17948                p0.0,
17949                p1.0,
17950                aq,
17951                ad,
17952                m,
17953                w0.in_features(),
17954                p0.1,
17955                p1.1,
17956                p0.2,
17957                p0.3,
17958                p1.3,
17959            )?));
17960        }
17961        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
17962            return Ok(None);
17963        };
17964        Ok(Some(self.q8_fused2_t_core(
17965            p0.0,
17966            p1.0,
17967            aq,
17968            ad,
17969            m,
17970            w0.in_features(),
17971            p0.1,
17972            p1.1,
17973            p0.2,
17974        )?))
17975    }
17976
17977    #[allow(clippy::too_many_arguments)]
17978    fn q8_fused2_t_core(
17979        &self,
17980        b0: &CudaSlice<u8>,
17981        b1: &CudaSlice<u8>,
17982        aq: &CudaSlice<i8>,
17983        ad: &CudaSlice<f32>,
17984        m: usize,
17985        in_f: usize,
17986        out0: usize,
17987        out1: usize,
17988        row_bytes: usize,
17989    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17990        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
17991        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
17992        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
17993        let f = self.func(match Self::batched_mcols(m) {
17994            2 => "qmatvec_q8_0_mmvq_fused2_b2",
17995            4 => "qmatvec_q8_0_mmvq_fused2_b4",
17996            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
17997            _ => "qmatvec_q8_0_mmvq_fused2_b8",
17998        });
17999        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18000        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18001        let cfg = LaunchConfig {
18002            grid_dim: (nb0 + nb1, 1, 1),
18003            block_dim: (32, ROWS_PER_BLOCK, 1),
18004            shared_mem_bytes: 0,
18005        };
18006        let (inf, o0, o1, mi, rbl) = (
18007            in_f as i32,
18008            out0 as i32,
18009            out1 as i32,
18010            m as i32,
18011            row_bytes as i64,
18012        );
18013        let __s_b = self.gpu.stream();
18014        let mut b = __s_b.launch_builder(&f);
18015        b.arg(b0)
18016            .arg(b1)
18017            .arg(aq)
18018            .arg(ad)
18019            .arg(&mut y0)
18020            .arg(&mut y1)
18021            .arg(&inf)
18022            .arg(&o0)
18023            .arg(&o1)
18024            .arg(&mi)
18025            .arg(&rbl);
18026        unsafe {
18027            b.launch(cfg)?;
18028        }
18029        Ok((y0, y1))
18030    }
18031
18032    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
18033    /// q8_1 quant of the [m, in_f] activation), no env gating.
18034    #[allow(clippy::too_many_arguments)]
18035    pub fn qmatvec_q8_fused2_t_raw(
18036        &self,
18037        b0: &CudaSlice<u8>,
18038        b1: &CudaSlice<u8>,
18039        x: &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        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18047        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
18048    }
18049
18050    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
18051    /// `matmul_q8_fused2_t` with three ranges.
18052    #[allow(clippy::too_many_arguments)]
18053    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18054    pub fn matmul_q8_fused3_t(
18055        &self,
18056        w0: &crate::model::GpuTensor,
18057        w1: &crate::model::GpuTensor,
18058        w2: &crate::model::GpuTensor,
18059        aq: &CudaSlice<i8>,
18060        ad: &CudaSlice<f32>,
18061        m: usize,
18062    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18063    {
18064        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
18065            return Ok(None);
18066        }
18067        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
18068            return Ok(Some(self.e4m3_fused3_t_core(
18069                p0.0,
18070                p1.0,
18071                p2.0,
18072                aq,
18073                ad,
18074                m,
18075                w0.in_features(),
18076                p0.1,
18077                p1.1,
18078                p2.1,
18079                p0.2,
18080                p0.3,
18081                p1.3,
18082                p2.3,
18083            )?));
18084        }
18085        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
18086            return Ok(None);
18087        };
18088        Ok(Some(self.q8_fused3_t_core(
18089            p0.0,
18090            p1.0,
18091            p2.0,
18092            aq,
18093            ad,
18094            m,
18095            w0.in_features(),
18096            p0.1,
18097            p1.1,
18098            p2.1,
18099            p0.2,
18100        )?))
18101    }
18102
18103    #[allow(clippy::too_many_arguments)]
18104    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18105    fn q8_fused3_t_core(
18106        &self,
18107        b0: &CudaSlice<u8>,
18108        b1: &CudaSlice<u8>,
18109        b2: &CudaSlice<u8>,
18110        aq: &CudaSlice<i8>,
18111        ad: &CudaSlice<f32>,
18112        m: usize,
18113        in_f: usize,
18114        out0: usize,
18115        out1: usize,
18116        out2: usize,
18117        row_bytes: usize,
18118    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18119        const ROWS_PER_BLOCK: u32 = 4;
18120        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18121        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18122        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18123        let f = self.func(if Self::batched_mcols(m) == 2 {
18124            "qmatvec_q8_0_mmvq_fused3_b2"
18125        } else {
18126            "qmatvec_q8_0_mmvq_fused3_b4"
18127        });
18128        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18129        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18130        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
18131        let cfg = LaunchConfig {
18132            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18133            block_dim: (32, ROWS_PER_BLOCK, 1),
18134            shared_mem_bytes: 0,
18135        };
18136        let (inf, o0, o1, o2, mi, rbl) = (
18137            in_f as i32,
18138            out0 as i32,
18139            out1 as i32,
18140            out2 as i32,
18141            m as i32,
18142            row_bytes as i64,
18143        );
18144        let __s_b = self.gpu.stream();
18145        let mut b = __s_b.launch_builder(&f);
18146        b.arg(b0)
18147            .arg(b1)
18148            .arg(b2)
18149            .arg(aq)
18150            .arg(ad)
18151            .arg(&mut y0)
18152            .arg(&mut y1)
18153            .arg(&mut y2)
18154            .arg(&inf)
18155            .arg(&o0)
18156            .arg(&o1)
18157            .arg(&o2)
18158            .arg(&mi)
18159            .arg(&rbl);
18160        unsafe {
18161            b.launch(cfg)?;
18162        }
18163        Ok((y0, y1, y2))
18164    }
18165
18166    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
18167    #[allow(clippy::too_many_arguments)]
18168    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18169    pub fn qmatvec_q8_fused3_t_raw(
18170        &self,
18171        b0: &CudaSlice<u8>,
18172        b1: &CudaSlice<u8>,
18173        b2: &CudaSlice<u8>,
18174        x: &CudaSlice<f32>,
18175        m: usize,
18176        in_f: usize,
18177        out0: usize,
18178        out1: usize,
18179        out2: usize,
18180        row_bytes: usize,
18181    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18182        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18183        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
18184    }
18185
18186    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
18187    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
18188    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
18189    pub fn q8_ffn_fuse2_on(&self) -> bool {
18190        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18191        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
18192    }
18193
18194    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
18195    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
18196    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
18197    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
18198    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
18199    #[allow(clippy::type_complexity)]
18200    fn q8_fused_params<'w, const N: usize>(
18201        &self,
18202        ws: &[&'w crate::model::GpuTensor; N],
18203    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
18204        use crate::model::GpuTensor;
18205        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
18206            return None;
18207        }
18208        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
18209            return None;
18210        }
18211        let in_f = ws[0].in_features();
18212        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
18213        for (i, w) in ws.iter().enumerate() {
18214            match w {
18215                GpuTensor::Quant {
18216                    bytes,
18217                    qtype,
18218                    row_bytes,
18219                    scale,
18220                    ..
18221                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
18222                    out[i] = Some((bytes, w.out_features(), *row_bytes))
18223                }
18224                _ => return None,
18225            }
18226        }
18227        Some(out.map(|o| o.unwrap()))
18228    }
18229
18230    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
18231    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
18232    pub fn e4m3_dual_on(&self) -> bool {
18233        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18234        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
18235    }
18236
18237    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
18238    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
18239    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
18240    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
18241    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
18242    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
18243    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
18244    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
18245    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
18246    ///     Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
18247    ///     there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
18248    #[allow(clippy::type_complexity)]
18249    fn e4m3_fused_params<'w, const N: usize>(
18250        &self,
18251        ws: &[&'w crate::model::GpuTensor; N],
18252    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
18253        use crate::model::GpuTensor;
18254        if !self.e4m3_dual_on() {
18255            return None;
18256        }
18257        let in_f = ws[0].in_features();
18258        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
18259        for (i, w) in ws.iter().enumerate() {
18260            match w {
18261                GpuTensor::Quant {
18262                    bytes,
18263                    qtype,
18264                    row_bytes,
18265                    scale,
18266                    rp,
18267                    rp4,
18268                    ..
18269                } if *qtype == QT_F8_E4M3
18270                    && w.in_features() == in_f
18271                    && *row_bytes == in_f
18272                    && !*rp
18273                    && rp4.is_none() =>
18274                {
18275                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
18276                }
18277                _ => return None,
18278            }
18279        }
18280        Some(out.map(|o| o.unwrap()))
18281    }
18282
18283    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
18284    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
18285    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
18286    #[allow(clippy::too_many_arguments)]
18287    fn e4m3_fused2_core(
18288        &self,
18289        b0: &CudaSlice<u8>,
18290        b1: &CudaSlice<u8>,
18291        aq: &CudaSlice<i8>,
18292        ad: &CudaSlice<f32>,
18293        in_f: usize,
18294        out0: usize,
18295        out1: usize,
18296        row_bytes: usize,
18297        ws0: f32,
18298        ws1: f32,
18299    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18300        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18301        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18302        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18303        let f = self.func("qmatvec_e4m3_mmvq_fused2");
18304        let mut y0 = self.alloc_uninit::<f32>(out0)?;
18305        let mut y1 = self.alloc_uninit::<f32>(out1)?;
18306        let cfg = LaunchConfig {
18307            grid_dim: (nb0 + nb1, 1, 1),
18308            block_dim: (32, ROWS_PER_BLOCK, 1),
18309            shared_mem_bytes: 0,
18310        };
18311        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
18312        let __s_b = self.gpu.stream();
18313        let mut b = __s_b.launch_builder(&f);
18314        b.arg(b0)
18315            .arg(b1)
18316            .arg(aq)
18317            .arg(ad)
18318            .arg(&mut y0)
18319            .arg(&mut y1)
18320            .arg(&inf)
18321            .arg(&o0)
18322            .arg(&o1)
18323            .arg(&rbl)
18324            .arg(&ws0)
18325            .arg(&ws1);
18326        unsafe {
18327            b.launch(cfg)?;
18328        }
18329        Ok((y0, y1))
18330    }
18331
18332    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
18333    #[allow(clippy::too_many_arguments)]
18334    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18335    fn e4m3_fused3_core(
18336        &self,
18337        b0: &CudaSlice<u8>,
18338        b1: &CudaSlice<u8>,
18339        b2: &CudaSlice<u8>,
18340        aq: &CudaSlice<i8>,
18341        ad: &CudaSlice<f32>,
18342        in_f: usize,
18343        out0: usize,
18344        out1: usize,
18345        out2: usize,
18346        row_bytes: usize,
18347        ws0: f32,
18348        ws1: f32,
18349        ws2: f32,
18350    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18351        const ROWS_PER_BLOCK: u32 = 4;
18352        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18353        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18354        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18355        let f = self.func("qmatvec_e4m3_mmvq_fused3");
18356        let mut y0 = self.alloc_uninit::<f32>(out0)?;
18357        let mut y1 = self.alloc_uninit::<f32>(out1)?;
18358        let mut y2 = self.alloc_uninit::<f32>(out2)?;
18359        let cfg = LaunchConfig {
18360            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18361            block_dim: (32, ROWS_PER_BLOCK, 1),
18362            shared_mem_bytes: 0,
18363        };
18364        let (inf, o0, o1, o2, rbl) = (
18365            in_f as i32,
18366            out0 as i32,
18367            out1 as i32,
18368            out2 as i32,
18369            row_bytes as i64,
18370        );
18371        let __s_b = self.gpu.stream();
18372        let mut b = __s_b.launch_builder(&f);
18373        b.arg(b0)
18374            .arg(b1)
18375            .arg(b2)
18376            .arg(aq)
18377            .arg(ad)
18378            .arg(&mut y0)
18379            .arg(&mut y1)
18380            .arg(&mut y2)
18381            .arg(&inf)
18382            .arg(&o0)
18383            .arg(&o1)
18384            .arg(&o2)
18385            .arg(&rbl)
18386            .arg(&ws0)
18387            .arg(&ws1)
18388            .arg(&ws2);
18389        unsafe {
18390            b.launch(cfg)?;
18391        }
18392        Ok((y0, y1, y2))
18393    }
18394
18395    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
18396    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
18397    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
18398    #[allow(clippy::too_many_arguments)]
18399    fn e4m3_fused2_t_core(
18400        &self,
18401        b0: &CudaSlice<u8>,
18402        b1: &CudaSlice<u8>,
18403        aq: &CudaSlice<i8>,
18404        ad: &CudaSlice<f32>,
18405        m: usize,
18406        in_f: usize,
18407        out0: usize,
18408        out1: usize,
18409        row_bytes: usize,
18410        ws0: f32,
18411        ws1: f32,
18412    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18413        const ROWS_PER_BLOCK: u32 = 4;
18414        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18415        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18416        let f = self.func(match Self::batched_mcols(m) {
18417            2 => "qmatvec_e4m3_mmvq_fused2_b2",
18418            4 => "qmatvec_e4m3_mmvq_fused2_b4",
18419            _ => "qmatvec_e4m3_mmvq_fused2_b8",
18420        });
18421        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18422        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18423        let cfg = LaunchConfig {
18424            grid_dim: (nb0 + nb1, 1, 1),
18425            block_dim: (32, ROWS_PER_BLOCK, 1),
18426            shared_mem_bytes: 0,
18427        };
18428        let (inf, o0, o1, mi, rbl) = (
18429            in_f as i32,
18430            out0 as i32,
18431            out1 as i32,
18432            m as i32,
18433            row_bytes as i64,
18434        );
18435        let __s_b = self.gpu.stream();
18436        let mut b = __s_b.launch_builder(&f);
18437        b.arg(b0)
18438            .arg(b1)
18439            .arg(aq)
18440            .arg(ad)
18441            .arg(&mut y0)
18442            .arg(&mut y1)
18443            .arg(&inf)
18444            .arg(&o0)
18445            .arg(&o1)
18446            .arg(&mi)
18447            .arg(&rbl);
18448        unsafe {
18449            b.launch(cfg)?;
18450        }
18451        if ws0 != 1.0 {
18452            self.scale_inplace(&mut y0, ws0, m * out0)?;
18453        }
18454        if ws1 != 1.0 {
18455            self.scale_inplace(&mut y1, ws1, m * out1)?;
18456        }
18457        Ok((y0, y1))
18458    }
18459
18460    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
18461    #[allow(clippy::too_many_arguments)]
18462    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18463    fn e4m3_fused3_t_core(
18464        &self,
18465        b0: &CudaSlice<u8>,
18466        b1: &CudaSlice<u8>,
18467        b2: &CudaSlice<u8>,
18468        aq: &CudaSlice<i8>,
18469        ad: &CudaSlice<f32>,
18470        m: usize,
18471        in_f: usize,
18472        out0: usize,
18473        out1: usize,
18474        out2: usize,
18475        row_bytes: usize,
18476        ws0: f32,
18477        ws1: f32,
18478        ws2: f32,
18479    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18480        const ROWS_PER_BLOCK: u32 = 4;
18481        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18482        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18483        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
18484        let f = self.func(if Self::batched_mcols(m) == 2 {
18485            "qmatvec_e4m3_mmvq_fused3_b2"
18486        } else {
18487            "qmatvec_e4m3_mmvq_fused3_b4"
18488        });
18489        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
18490        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
18491        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
18492        let cfg = LaunchConfig {
18493            grid_dim: (nb0 + nb1 + nb2, 1, 1),
18494            block_dim: (32, ROWS_PER_BLOCK, 1),
18495            shared_mem_bytes: 0,
18496        };
18497        let (inf, o0, o1, o2, mi, rbl) = (
18498            in_f as i32,
18499            out0 as i32,
18500            out1 as i32,
18501            out2 as i32,
18502            m as i32,
18503            row_bytes as i64,
18504        );
18505        let __s_b = self.gpu.stream();
18506        let mut b = __s_b.launch_builder(&f);
18507        b.arg(b0)
18508            .arg(b1)
18509            .arg(b2)
18510            .arg(aq)
18511            .arg(ad)
18512            .arg(&mut y0)
18513            .arg(&mut y1)
18514            .arg(&mut y2)
18515            .arg(&inf)
18516            .arg(&o0)
18517            .arg(&o1)
18518            .arg(&o2)
18519            .arg(&mi)
18520            .arg(&rbl);
18521        unsafe {
18522            b.launch(cfg)?;
18523        }
18524        if ws0 != 1.0 {
18525            self.scale_inplace(&mut y0, ws0, m * out0)?;
18526        }
18527        if ws1 != 1.0 {
18528            self.scale_inplace(&mut y1, ws1, m * out1)?;
18529        }
18530        if ws2 != 1.0 {
18531            self.scale_inplace(&mut y2, ws2, m * out2)?;
18532        }
18533        Ok((y0, y1, y2))
18534    }
18535
18536    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
18537    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
18538    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
18539    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
18540    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
18541    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
18542    ///
18543    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
18544    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
18545    #[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
18546    pub fn qmatvec_e4m3_blk_mmvq(
18547        &self,
18548        bytes: &CudaSlice<u8>,
18549        aq: &CudaSlice<i8>,
18550        ad: &CudaSlice<f32>,
18551        scales: &CudaSlice<f32>,
18552        m: usize,
18553        in_f: usize,
18554        out_f: usize,
18555        row_bytes: usize,
18556        scale_cols: usize,
18557    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18558        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
18559        self.qmatvec_e4m3_blk_mmvq_into(
18560            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
18561        )?;
18562        Ok(y)
18563    }
18564
18565    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
18566    #[allow(clippy::too_many_arguments)]
18567    pub fn qmatvec_e4m3_blk_mmvq_into(
18568        &self,
18569        bytes: &CudaSlice<u8>,
18570        aq: &CudaSlice<i8>,
18571        ad: &CudaSlice<f32>,
18572        scales: &CudaSlice<f32>,
18573        m: usize,
18574        in_f: usize,
18575        out_f: usize,
18576        row_bytes: usize,
18577        scale_cols: usize,
18578        y: &mut CudaSlice<f32>,
18579    ) -> Result<(), Box<dyn std::error::Error>> {
18580        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18581        let f = self.func("qmatvec_e4m3_blk_mmvq");
18582        let cfg = LaunchConfig {
18583            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
18584            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
18585            shared_mem_bytes: 0,                // warp-only reduce
18586        };
18587        let (inf, outf, mi, rb, sc) = (
18588            in_f as i32,
18589            out_f as i32,
18590            m as i32,
18591            row_bytes as i64,
18592            scale_cols as i32,
18593        );
18594        let __s_b = self.gpu.stream();
18595        let mut b = __s_b.launch_builder(&f);
18596        b.arg(bytes)
18597            .arg(aq)
18598            .arg(ad)
18599            .arg(scales)
18600            .arg(&mut *y)
18601            .arg(&inf)
18602            .arg(&outf)
18603            .arg(&mi)
18604            .arg(&rb)
18605            .arg(&sc);
18606        unsafe {
18607            b.launch(cfg)?;
18608        }
18609        Ok(())
18610    }
18611
18612    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
18613    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
18614    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
18615    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
18616    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
18617    #[allow(clippy::too_many_arguments)]
18618    pub fn qmatvec_e4m3_blk_mmvq_batched(
18619        &self,
18620        bytes: &CudaSlice<u8>,
18621        aq: &CudaSlice<i8>,
18622        ad: &CudaSlice<f32>,
18623        scales: &CudaSlice<f32>,
18624        m: usize,
18625        in_f: usize,
18626        out_f: usize,
18627        row_bytes: usize,
18628        scale_cols: usize,
18629        mcols: usize,
18630    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18631        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18632        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
18633        let name = match mcols {
18634            2 => "qmatvec_e4m3_blk_mmvq_b2",
18635            4 => "qmatvec_e4m3_blk_mmvq_b4",
18636            8 => "qmatvec_e4m3_blk_mmvq_b8",
18637            16 => "qmatvec_e4m3_blk_mmvq_b16",
18638            _ => {
18639                return Err(
18640                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
18641                );
18642            }
18643        };
18644        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18645        let f = self.func(name);
18646        let cfg = LaunchConfig {
18647            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18648            block_dim: (32, ROWS_PER_BLOCK, 1),
18649            shared_mem_bytes: 0,
18650        };
18651        let (inf, outf, mi, rb, sc) = (
18652            in_f as i32,
18653            out_f as i32,
18654            m as i32,
18655            row_bytes as i64,
18656            scale_cols as i32,
18657        );
18658        let __s_b = self.gpu.stream();
18659        let mut b = __s_b.launch_builder(&f);
18660        b.arg(bytes)
18661            .arg(aq)
18662            .arg(ad)
18663            .arg(scales)
18664            .arg(&mut y)
18665            .arg(&inf)
18666            .arg(&outf)
18667            .arg(&mi)
18668            .arg(&rb)
18669            .arg(&sc);
18670        unsafe {
18671            b.launch(cfg)?;
18672        }
18673        Ok(y)
18674    }
18675
18676    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
18677    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
18678    #[allow(clippy::too_many_arguments)]
18679    pub fn qmatvec_e4m3_blk_batched_raw(
18680        &self,
18681        bytes: &CudaSlice<u8>,
18682        x: &CudaSlice<f32>,
18683        scales: &CudaSlice<f32>,
18684        m: usize,
18685        in_f: usize,
18686        out_f: usize,
18687        row_bytes: usize,
18688        scale_cols: usize,
18689        mcols: usize,
18690    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18691        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18692        self.qmatvec_e4m3_blk_mmvq_batched(
18693            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
18694        )
18695    }
18696
18697    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
18698    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
18699    #[allow(clippy::too_many_arguments)]
18700    pub fn qmatvec_e4m3_blk_mmvq_raw(
18701        &self,
18702        bytes: &CudaSlice<u8>,
18703        x: &CudaSlice<f32>,
18704        scales: &CudaSlice<f32>,
18705        m: usize,
18706        in_f: usize,
18707        out_f: usize,
18708        row_bytes: usize,
18709        scale_cols: usize,
18710    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18711        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18712        self.qmatvec_e4m3_blk_mmvq(
18713            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
18714        )
18715    }
18716
18717    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
18718    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
18719    #[allow(clippy::too_many_arguments)]
18720    pub fn qmatvec_e4m3_fused2_raw(
18721        &self,
18722        b0: &CudaSlice<u8>,
18723        b1: &CudaSlice<u8>,
18724        x: &CudaSlice<f32>,
18725        in_f: usize,
18726        out0: usize,
18727        out1: usize,
18728        row_bytes: usize,
18729        ws0: f32,
18730        ws1: f32,
18731    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18732        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18733        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
18734    }
18735
18736    #[allow(clippy::too_many_arguments)]
18737    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18738    pub fn qmatvec_e4m3_fused3_raw(
18739        &self,
18740        b0: &CudaSlice<u8>,
18741        b1: &CudaSlice<u8>,
18742        b2: &CudaSlice<u8>,
18743        x: &CudaSlice<f32>,
18744        in_f: usize,
18745        out0: usize,
18746        out1: usize,
18747        out2: usize,
18748        row_bytes: usize,
18749        ws0: f32,
18750        ws1: f32,
18751        ws2: f32,
18752    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18753        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18754        self.e4m3_fused3_core(
18755            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
18756        )
18757    }
18758
18759    #[allow(clippy::too_many_arguments)]
18760    pub fn qmatvec_e4m3_fused2_t_raw(
18761        &self,
18762        b0: &CudaSlice<u8>,
18763        b1: &CudaSlice<u8>,
18764        x: &CudaSlice<f32>,
18765        m: usize,
18766        in_f: usize,
18767        out0: usize,
18768        out1: usize,
18769        row_bytes: usize,
18770        ws0: f32,
18771        ws1: f32,
18772    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18773        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18774        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
18775    }
18776
18777    #[allow(clippy::too_many_arguments)]
18778    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18779    pub fn qmatvec_e4m3_fused3_t_raw(
18780        &self,
18781        b0: &CudaSlice<u8>,
18782        b1: &CudaSlice<u8>,
18783        b2: &CudaSlice<u8>,
18784        x: &CudaSlice<f32>,
18785        m: usize,
18786        in_f: usize,
18787        out0: usize,
18788        out1: usize,
18789        out2: usize,
18790        row_bytes: usize,
18791        ws0: f32,
18792        ws1: f32,
18793        ws2: f32,
18794    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18795        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
18796        self.e4m3_fused3_t_core(
18797            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
18798        )
18799    }
18800
18801    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
18802    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
18803    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
18804    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
18805    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
18806    ///
18807    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
18808    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
18809    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
18810    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
18811    fn try_e4m3_blk_pre(
18812        &self,
18813        w: &crate::model::GpuTensor,
18814        aq: &CudaSlice<i8>,
18815        ad: &CudaSlice<f32>,
18816        m: usize,
18817    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
18818        use crate::model::GpuTensor;
18819        if let GpuTensor::Quant {
18820            bytes,
18821            qtype,
18822            row_bytes,
18823            blk: Some(g),
18824            ..
18825        } = w
18826            && *qtype == QT_F8_E4M3_BLK
18827        {
18828            // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
18829            // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
18830            // below, so the decode-exactness contract is preserved at every width. Gated by
18831            // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
18832            // one rollback door covers every dtype's batched tier.
18833            if (2..=16).contains(&m)
18834                && std::env::var("MEMRA_NO_BATCHED").is_err()
18835                && (m <= 4 || Self::b8_enabled())
18836            {
18837                let mcols = Self::batched_mcols(m);
18838                return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
18839                    bytes,
18840                    aq,
18841                    ad,
18842                    &g.scales,
18843                    m,
18844                    w.in_features(),
18845                    w.out_features(),
18846                    *row_bytes,
18847                    g.cols,
18848                    mcols,
18849                )?));
18850            }
18851            return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
18852                bytes,
18853                aq,
18854                ad,
18855                &g.scales,
18856                m,
18857                w.in_features(),
18858                w.out_features(),
18859                *row_bytes,
18860                g.cols,
18861            )?));
18862        }
18863        Ok(None)
18864    }
18865
18866    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
18867    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
18868    ///
18869    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
18870    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
18871    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
18872    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
18873    /// prefill keeps the floor's arithmetic and the floor's kernels.
18874    ///
18875    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
18876    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
18877    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
18878    /// (projection, prefill call) and frees immediately.
18879    ///
18880    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
18881    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
18882    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
18883    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
18884    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
18885    /// single-variable comparison instead of a two-variable one.
18886    ///
18887    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
18888    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
18889    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
18890    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
18891    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
18892    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
18893    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
18894    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
18895    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
18896    ///
18897    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
18898    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
18899    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
18900    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
18901    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
18902    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
18903    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
18904    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
18905    /// because v2's denominator had its slab already resident while this class's floor must build it
18906    /// every call; same tile, opposite sign, because the question changed.
18907    ///
18908    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
18909    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
18910    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
18911    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
18912    fn try_e4m3_blk_prefill(
18913        &self,
18914        w: &crate::model::GpuTensor,
18915        x: &CudaSlice<f32>,
18916        m: usize,
18917    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
18918        use crate::model::GpuTensor;
18919        let GpuTensor::Quant {
18920            bytes,
18921            qtype,
18922            blk: Some(g),
18923            ..
18924        } = w
18925        else {
18926            return Ok(None);
18927        };
18928        if *qtype != QT_F8_E4M3_BLK {
18929            return Ok(None);
18930        }
18931        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
18932        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
18933        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
18934        // through to the dequant below when they do, never silently produce nothing.
18935        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
18936            return Ok(Some(y));
18937        }
18938        let (in_f, out_f) = (w.in_features(), w.out_features());
18939        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
18940        let tmp = GpuTensor::Quant {
18941            bytes: slab,
18942            qtype: QT_Q8_0,
18943            row_bytes: in_f / 32 * 34,
18944            ne: vec![in_f as u64, out_f as u64],
18945            scale: 1.0,
18946            rp: false,
18947            #[cfg(memra_cutlass)]
18948            cutlass: None,
18949            fp8: None,
18950            blk: None,
18951            f16: None,
18952            rp4: None,
18953        };
18954        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
18955        Ok(Some(self.matmul(&tmp, x, m)?))
18956    }
18957
18958    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18959    pub fn matmul_pre_noscale(
18960        &self,
18961        w: &crate::model::GpuTensor,
18962        aq: &CudaSlice<i8>,
18963        ad: &CudaSlice<f32>,
18964        m: usize,
18965    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
18966        use crate::model::GpuTensor;
18967        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
18968        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
18969        // rather than let the tail below refuse and cost the caller a re-dispatch.
18970        if m == 1
18971            && let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)?
18972        {
18973            return Ok(Some((y, 1.0)));
18974        }
18975        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
18976        if m != 1 || !self.uses_q8_1_fast(w) {
18977            return Ok(None);
18978        }
18979        let in_f = w.in_features();
18980        let out_f = w.out_features();
18981        let (bytes, qtype, row_bytes, scale, rp) = match w {
18982            GpuTensor::Quant {
18983                bytes,
18984                qtype,
18985                row_bytes,
18986                scale,
18987                rp,
18988                ..
18989            } => (bytes, *qtype, *row_bytes, *scale, *rp),
18990            _ => return Ok(None),
18991        };
18992        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
18993        if self.mmvq_supports(qtype) {
18994            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
18995            let (mbytes, mrp) = match w {
18996                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
18997                _ => (bytes, rp),
18998            };
18999            let y = self.qmatvec_mmvq(
19000                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
19001            )?;
19002            return Ok(Some((y, scale)));
19003        }
19004        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
19005        let name = match qtype {
19006            QT_Q8_0 => "qmatvec_q8_0_dp4a",
19007            QT_Q4_K => "qmatvec_q4_K_dp4a",
19008            QT_Q6_K => "qmatvec_q6_K_dp4a",
19009            QT_Q5_K => "qmatvec_q5_K_dp4a",
19010            QT_Q3_K => "qmatvec_q3_K_dp4a",
19011            QT_NVFP4 => {
19012                if rp {
19013                    "qmatvec_nvfp4_dp4a_rp"
19014                } else {
19015                    "qmatvec_nvfp4_dp4a"
19016                }
19017            }
19018            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
19019            _ => return Ok(None),
19020        };
19021        let f = self.func(name);
19022        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19023        let cfg = LaunchConfig {
19024            grid_dim: (out_f as u32, m as u32, 1),
19025            block_dim: (128, 1, 1),
19026            shared_mem_bytes: 0,
19027        };
19028        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19029        let __s_b = self.gpu.stream();
19030        let mut b = __s_b.launch_builder(&f);
19031        b.arg(bytes)
19032            .arg(aq)
19033            .arg(ad)
19034            .arg(&mut y)
19035            .arg(&inf)
19036            .arg(&outf)
19037            .arg(&mi)
19038            .arg(&rb);
19039        unsafe {
19040            b.launch(cfg)?;
19041        }
19042        Ok(Some((y, scale)))
19043    }
19044
19045    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
19046    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
19047    pub fn mmvq_supports(&self, qtype: i32) -> bool {
19048        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
19049        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
19050        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
19051        // is a pure function of the dtype — the decode-parity law holds under every env.
19052        if qtype == QT_F8_E4M3 {
19053            return true;
19054        }
19055        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
19056            return false;
19057        }
19058        matches!(
19059            qtype,
19060            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
19061        )
19062    }
19063
19064    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
19065    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
19066    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
19067    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
19068    #[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
19069    pub fn qmatvec_mmvq(
19070        &self,
19071        bytes: &CudaSlice<u8>,
19072        aq: &CudaSlice<i8>,
19073        ad: &CudaSlice<f32>,
19074        m: usize,
19075        in_f: usize,
19076        out_f: usize,
19077        qtype: i32,
19078        row_bytes: usize,
19079        scale: f32,
19080        rp: bool,
19081    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19082        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
19083        self.qmatvec_mmvq_into(
19084            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
19085        )?;
19086        Ok(y)
19087    }
19088
19089    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
19090    #[allow(clippy::too_many_arguments)]
19091    #[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
19092    pub fn qmatvec_mmvq_into(
19093        &self,
19094        bytes: &CudaSlice<u8>,
19095        aq: &CudaSlice<i8>,
19096        ad: &CudaSlice<f32>,
19097        m: usize,
19098        in_f: usize,
19099        out_f: usize,
19100        qtype: i32,
19101        row_bytes: usize,
19102        scale: f32,
19103        rp: bool,
19104        y: &mut CudaSlice<f32>,
19105    ) -> Result<(), Box<dyn std::error::Error>> {
19106        debug_assert!(y.len() >= m * out_f);
19107        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
19108        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
19109        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
19110        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
19111        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
19112        if qtype == QT_Q8_0
19113            && rp
19114            && m == 1
19115            && out_f >= 64
19116            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
19117            && {
19118                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19119                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
19120            }
19121        {
19122            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
19123            let cfg = LaunchConfig {
19124                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
19125                block_dim: (32, 2, 1),
19126                shared_mem_bytes: 0,
19127            };
19128            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
19129            let __s_b = self.gpu.stream();
19130            let mut b = __s_b.launch_builder(&f);
19131            b.arg(bytes)
19132                .arg(aq)
19133                .arg(ad)
19134                .arg(&mut *y)
19135                .arg(&inf)
19136                .arg(&outf)
19137                .arg(&mi)
19138                .arg(&rb);
19139            unsafe {
19140                b.launch(cfg)?;
19141            }
19142            if scale != 1.0 {
19143                self.scale_inplace(y, scale, out_f)?;
19144            }
19145            return Ok(());
19146        }
19147        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
19148        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
19149        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
19150        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
19151        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
19152        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
19153        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
19154        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
19155        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
19156            2
19157        } else {
19158            1
19159        };
19160        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
19161        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
19162        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
19163        // valid-window interleaved, bit-identical per row — same dot program).
19164        if m == 1 && qtype == QT_Q4_0 {
19165            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
19166            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
19167            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
19168            mr = *Q40MR.get_or_init(|| {
19169                std::env::var("MEMRA_Q40_MR")
19170                    .ok()
19171                    .and_then(|v| v.parse().ok())
19172                    .unwrap_or(1)
19173            });
19174        }
19175        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
19176        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
19177        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
19178        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
19179        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
19180        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
19181        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
19182        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
19183        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
19184        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
19185        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
19186        let q5_force = q5_mode.as_deref() == Some("2");
19187        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
19188        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
19189        let q5_il = qtype == QT_Q5_K
19190            && m == 1
19191            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
19192        if q5_il && !q5_force && out_f > 65536 {
19193            mr = 1;
19194        }
19195        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
19196        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
19197        if qtype == QT_Q4_0 && rp && mr != 1 {
19198            mr = 2;
19199        }
19200        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
19201        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
19202        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
19203        if qtype == QT_Q8_0 && rp {
19204            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
19205            mr = *Q80MR.get_or_init(|| {
19206                std::env::var("MEMRA_Q80_MR")
19207                    .ok()
19208                    .and_then(|v| v.parse().ok())
19209                    .unwrap_or(1)
19210            });
19211        }
19212        let name = match (qtype, mr, rp) {
19213            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
19214            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
19215            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
19216            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
19217            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
19218            (QT_Q5_K, 2, _) => {
19219                if q5_il {
19220                    "qmatvec_q5_K_mmvq_mr2_il"
19221                } else {
19222                    "qmatvec_q5_K_mmvq_mr2"
19223                }
19224            }
19225            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
19226            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
19227            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
19228            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
19229            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
19230            (QT_Q8_0, _, true)
19231                if in_f.is_multiple_of(1024) && {
19232                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19233                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
19234                } =>
19235            {
19236                "qmatvec_q8_0_mmvq_rpca"
19237            }
19238            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
19239            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
19240            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
19241            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
19242            // reach a GGUF-layout kernel or vice versa.
19243            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
19244            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
19245            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
19246            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
19247            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
19248            (QT_Q5_K, _, _) => {
19249                if q5_il {
19250                    "qmatvec_q5_K_mmvq_il"
19251                } else {
19252                    "qmatvec_q5_K_mmvq"
19253                }
19254            }
19255            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
19256            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
19257            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
19258            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
19259        };
19260        let f = self.func(name);
19261        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
19262        let rows_per_block = ROWS_PER_BLOCK * mr;
19263        let cfg = LaunchConfig {
19264            grid_dim: (
19265                (out_f as u32 + rows_per_block - 1) / rows_per_block,
19266                m as u32,
19267                1,
19268            ),
19269            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
19270            shared_mem_bytes: 0,                // warp-only reduce at m=1
19271        };
19272        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19273        let __s_b = self.gpu.stream();
19274        let mut b = __s_b.launch_builder(&f);
19275        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
19276        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
19277        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
19278        // weight_scale). Other mmvq kernels keep the 8-arg signature.
19279        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
19280            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
19281            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
19282            if Self::pdl_on()
19283                && Self::pdl_mmvq_on()
19284                && Self::pdl_nvfp4q8_on()
19285                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
19286            {
19287                use cudarc::driver::{DevicePtr, DevicePtrMut};
19288                let s = &self.gpu.stream();
19289                let (pw, _g0) = bytes.device_ptr(s);
19290                let (paq, _g1) = aq.device_ptr(s);
19291                let (pad, _g2) = ad.device_ptr(s);
19292                let (py, _g3) = y.device_ptr_mut(s);
19293                let mut ps = [
19294                    &pw as *const _ as *mut std::ffi::c_void,
19295                    &paq as *const _ as *mut _,
19296                    &pad as *const _ as *mut _,
19297                    &py as *const _ as *mut _,
19298                    &inf as *const _ as *mut _,
19299                    &outf as *const _ as *mut _,
19300                    &mi as *const _ as *mut _,
19301                    &rb as *const _ as *mut _,
19302                    &scale as *const _ as *mut _,
19303                ];
19304                unsafe {
19305                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
19306                }
19307                return Ok(());
19308            }
19309            b.arg(bytes)
19310                .arg(aq)
19311                .arg(ad)
19312                .arg(&mut *y)
19313                .arg(&inf)
19314                .arg(&outf)
19315                .arg(&mi)
19316                .arg(&rb)
19317                .arg(&scale);
19318            unsafe {
19319                b.launch(cfg)?;
19320            }
19321        } else if Self::pdl_on()
19322            && Self::pdl_mmvq_on()
19323            && (matches!(
19324                name,
19325                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
19326            ) || (Self::pdl_nvfp4q8_on()
19327                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
19328        {
19329            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
19330            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
19331            // names may take this launch (unmarked kernels would read unordered).
19332            {
19333                use cudarc::driver::{DevicePtr, DevicePtrMut};
19334                let s = &self.gpu.stream();
19335                let (pw, _g0) = bytes.device_ptr(s);
19336                let (paq, _g1) = aq.device_ptr(s);
19337                let (pad, _g2) = ad.device_ptr(s);
19338                let (py, _g3) = y.device_ptr_mut(s);
19339                let mut ps = [
19340                    &pw as *const _ as *mut std::ffi::c_void,
19341                    &paq as *const _ as *mut _,
19342                    &pad as *const _ as *mut _,
19343                    &py as *const _ as *mut _,
19344                    &inf as *const _ as *mut _,
19345                    &outf as *const _ as *mut _,
19346                    &mi as *const _ as *mut _,
19347                    &rb as *const _ as *mut _,
19348                ];
19349                unsafe {
19350                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
19351                }
19352            }
19353            if scale != 1.0 {
19354                self.scale_inplace(y, scale, m * out_f)?;
19355            }
19356        } else {
19357            b.arg(bytes)
19358                .arg(aq)
19359                .arg(ad)
19360                .arg(&mut *y)
19361                .arg(&inf)
19362                .arg(&outf)
19363                .arg(&mi)
19364                .arg(&rb);
19365            unsafe {
19366                b.launch(cfg)?;
19367            }
19368            if scale != 1.0 {
19369                self.scale_inplace(y, scale, m * out_f)?;
19370            }
19371        }
19372        Ok(())
19373    }
19374
19375    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
19376    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
19377    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
19378    #[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
19379    pub fn qmatvec_mmvq_raw(
19380        &self,
19381        bytes: &CudaSlice<u8>,
19382        x: &CudaSlice<f32>,
19383        m: usize,
19384        in_f: usize,
19385        out_f: usize,
19386        qtype: i32,
19387        row_bytes: usize,
19388        rp: bool,
19389    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19390        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19391        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
19392    }
19393
19394    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
19395    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
19396    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
19397    pub fn batched_supports(&self, qtype: i32) -> bool {
19398        matches!(
19399            qtype,
19400            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
19401        )
19402    }
19403
19404    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
19405    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
19406    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
19407    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
19408    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
19409    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
19410    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
19411    pub fn iq_fast_enabled() -> bool {
19412        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19413        *ON.get_or_init(|| {
19414            std::env::var("MEMRA_IQ_FAST")
19415                .map(|v| v != "0")
19416                .unwrap_or(true)
19417        })
19418    }
19419
19420    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
19421    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
19422    pub fn b8_enabled() -> bool {
19423        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19424        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
19425    }
19426
19427    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
19428    pub fn batched_mcols(m: usize) -> usize {
19429        if m == 2 {
19430            2
19431        } else if m <= 4 {
19432            4
19433        } else if m <= 8 {
19434            8
19435        } else {
19436            16
19437        }
19438    }
19439
19440    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
19441    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
19442    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
19443    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
19444    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
19445        Some(match (qtype, mcols) {
19446            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
19447            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
19448            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
19449            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
19450            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
19451            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
19452            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
19453            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
19454            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
19455            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
19456            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
19457            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
19458            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
19459            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
19460            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
19461            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
19462            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
19463            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
19464            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
19465            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
19466            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
19467            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
19468            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
19469            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
19470            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
19471            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
19472            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
19473            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
19474            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
19475            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
19476            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
19477            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
19478            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
19479            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
19480            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
19481            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
19482            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
19483            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
19484            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
19485            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
19486            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
19487            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
19488            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
19489            _ => return None,
19490        })
19491    }
19492
19493    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
19494    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
19495    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
19496    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
19497    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
19498    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
19499    ///
19500    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
19501    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
19502    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
19503    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
19504    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
19505    /// msweep on all six 27B shapes (2026-07-03):
19506    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
19507    ///          it applies for b4 (-3..-14%), never loses;
19508    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
19509    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
19510    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
19511    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
19512    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
19513    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
19514    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
19515    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
19516    /// b2: in_f>=6144 -> r2, else base.
19517    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
19518    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
19519    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
19520    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
19521    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
19522    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
19523    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
19524    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
19525    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
19526    /// Device SM count (cached) — grid-fill policy input.
19527    pub fn sm_count(&self) -> i32 {
19528        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
19529        *SMS.get_or_init(|| {
19530            use cudarc::driver::sys::CUdevice_attribute_enum as A;
19531            self.gpu
19532                .ctx
19533                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
19534                .unwrap_or(82)
19535        })
19536    }
19537
19538    #[allow(clippy::too_many_arguments)]
19539    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19540    #[allow(clippy::if_same_then_else)] // allow: a fallback mapping table; distinct inputs deliberately share a target arm
19541    pub fn batched_variant(
19542        &self,
19543        _m: usize,
19544        in_f: usize,
19545        out_f: usize,
19546        qtype: i32,
19547        row_bytes: usize,
19548        mcols: usize,
19549        rp: bool,
19550    ) -> &'static str {
19551        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
19552        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
19553        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
19554        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
19555        if qtype == QT_Q8_0 {
19556            return if rp { "rp" } else { "base" };
19557        }
19558        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19559        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
19560            Ok("base") => "base",
19561            Ok("pf") => "pf",
19562            Ok("r2") => "r2",
19563            Ok("r2w8") => "r2w8",
19564            Ok("pfr2") => "pfr2",
19565            Ok("ca") => "ca",
19566            Ok("car2") => "car2",
19567            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
19568            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
19569            Ok("rp") => "rp",
19570            Ok("rpr2") => "rpr2",
19571            Ok("rpr2w8") => "rpr2w8",
19572            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
19573            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
19574            Ok("rpca") => "rpca",
19575            Ok("rpcar2") => "rpcar2",
19576            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
19577            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
19578            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
19579            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
19580            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
19581            // bit-identical to the decode path — measurement corpus ONLY, never auto).
19582            Ok("rpsc") => "rpsc",
19583            Ok("rpms") => "rpms",
19584            Ok("rpmsc") => "rpmsc",
19585            Ok("rpks") => "rpks",
19586            Ok("rpksc") => "rpksc",
19587            _ => "auto",
19588        });
19589        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
19590        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
19591        // shapes qualify; anything else falls back to the register variants.
19592        let ca_ok = qtype == QT_NVFP4 && row_bytes.is_multiple_of(16) && in_f.is_multiple_of(1024);
19593        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
19594        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
19595        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
19596        // forced MEMRA_MMVQ_BV values still work).
19597        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19598        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
19599        let sc_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(256) && (in_f / 64 <= 272);
19600        let ks_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(512) && (in_f / 64 <= 272);
19601        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
19602        let sms = *SMS.get_or_init(|| {
19603            use cudarc::driver::sys::CUdevice_attribute_enum as A;
19604            self.gpu
19605                .ctx
19606                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
19607                .unwrap_or(82)
19608        });
19609        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
19610        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
19611        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
19612        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
19613        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
19614        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
19615        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
19616        // AUTO RULE = the measured winners table (differs from NVFP4's!):
19617        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
19618        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
19619        //     r2 1258us) — kernels kept behind the force seam for the corpus;
19620        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
19621        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
19622        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
19623        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
19624        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
19625        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
19626        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
19627        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
19628        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
19629        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
19630        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
19631        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19632        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
19633            Ok("base") => "base",
19634            Ok("r2") => "r2",
19635            Ok("r2w8") => "r2w8",
19636            _ => "auto",
19637        });
19638        let variant: &'static str = if qtype == QT_Q4_0 {
19639            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
19640            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
19641            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
19642            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
19643            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
19644                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
19645                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
19646                // + syncs cost more than the stalls, bank-pad made no difference);
19647                // register load-ahead flat (nvcc already reorders). The b-tier limiter
19648                // is still unidentified — see the jsonl row.
19649                Ok("base") => "base",
19650                Ok("r2") => "r2",
19651                Ok("ms") => "ms",
19652                Ok("sm") => "sm",
19653                Ok("la") => "la",
19654                _ => "auto",
19655            });
19656            let v = if q40 != "auto" {
19657                q40
19658            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
19659                "r2"
19660            } else {
19661                "base"
19662            };
19663            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
19664            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
19665            // and the limiter is the per-column activation load chain (long_scoreboard
19666            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
19667            if rp {
19668                match v {
19669                    "ms" => "r2ms_rp",
19670                    "sm" => "r2sm_rp",
19671                    "la" => "r2la_rp",
19672                    "r2" => "r2_rp",
19673                    _ => "rp",
19674                }
19675            } else if matches!(v, "ms" | "sm" | "la") {
19676                "r2"
19677            } else {
19678                v
19679            }
19680        } else if qtype != QT_NVFP4 && !kq_r2 {
19681            "base"
19682        } else if kq_r2 && rp {
19683            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
19684            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
19685            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
19686            "rp"
19687        } else if kq_r2 {
19688            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
19689            // mcols != 4 forced r2w8 falls to unbounded r2.
19690            if kq_bv != "auto" {
19691                if kq_bv == "r2w8" && mcols != 4 {
19692                    "r2"
19693                } else {
19694                    kq_bv
19695                }
19696            } else if bv != "auto" {
19697                match bv {
19698                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
19699                    "r2w8" | "rpr2w8" => {
19700                        if mcols != 4 {
19701                            "r2"
19702                        } else {
19703                            "r2w8"
19704                        }
19705                    }
19706                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
19707                }
19708            } else {
19709                #[allow(clippy::manual_div_ceil)]
19710                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19711                let blocks = (out_f + 7) / 8;
19712                let waves = blocks as f64 / (7 * sms as usize) as f64;
19713                let filled = blocks >= 4 * sms as usize;
19714                let use_r2 = if qtype == QT_Q4_K {
19715                    filled
19716                } else {
19717                    waves >= 2.0
19718                };
19719                if use_r2 { "r2" } else { "base" }
19720            }
19721        } else if bv != "auto" {
19722            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
19723            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
19724            // unsupported (shape, mcols) combos fall back to pf/r2.
19725            // On rp buffers, forced legacy names map to their rp twins (layout law).
19726            let v = if bv == "r2w8" && mcols == 2 {
19727                "r2"
19728            } else if bv == "ca" && (!ca_ok || mcols == 8) {
19729                "pf"
19730            } else if bv == "car2" && (!ca_ok || mcols == 8) {
19731                "r2"
19732            } else if bv == "pfr2" && mcols == 8 {
19733                "r2"
19734            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
19735                "rpr2"
19736            }
19737            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
19738            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
19739                if mcols == 8 { "rpr2w8" } else { "rpr2" }
19740            } else if bv == "rpcar2" && mcols == 2 {
19741                "rpca"
19742            }
19743            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
19744            // (rpms has no smem and no alignment need — always valid on rp buffers).
19745            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
19746                "rpr2"
19747            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
19748                "rpr2"
19749            } else {
19750                bv
19751            };
19752            if rp {
19753                match v {
19754                    "base" | "pf" | "ca" | "rp" => "rp",
19755                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
19756                    "r2w8" | "rpr2w8" => {
19757                        if mcols == 2 {
19758                            "rpr2"
19759                        } else {
19760                            "rpr2w8"
19761                        }
19762                    }
19763                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
19764                }
19765            } else {
19766                v
19767            }
19768        } else if mcols == 8 {
19769            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
19770            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
19771            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
19772            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
19773            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
19774            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
19775            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
19776            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
19777            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
19778            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
19779            if rp {
19780                if sc_ok { "rpsc" } else { "rpr2w8" }
19781            } else {
19782                "r2w8"
19783            }
19784        } else if mcols >= 4 {
19785            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
19786            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
19787            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
19788            #[allow(clippy::manual_div_ceil)]
19789            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19790            let blocks = (out_f + 7) / 8;
19791            let r7 = 7 * sms as usize;
19792            let r8 = 8 * sms as usize;
19793            let waves = blocks as f64 / r7 as f64;
19794            let filled = blocks >= 4 * sms as usize;
19795            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
19796            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
19797            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
19798            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
19799                // the extra residency drops the INTEGER wave count -> the straggler wave a
19800                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
19801                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
19802                if rp { "rpr2w8" } else { "r2w8" }
19803            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
19804                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
19805                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
19806                if rp { "rpr2" } else { "r2" }
19807            } else {
19808                // fractional straggler-wave window with no crossing, or grid too small to fill
19809                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
19810                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
19811                if rp { "rp" } else { "pf" }
19812            }
19813        } else if in_f >= 6144 {
19814            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
19815            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
19816            // stays.
19817            if rp { "rpr2" } else { "r2" }
19818        } else if rp {
19819            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
19820            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
19821            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
19822            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
19823            #[allow(clippy::manual_div_ceil)]
19824            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
19825            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
19826            if sc_ok && (0.9..=1.1).contains(&waves) {
19827                "rpsc"
19828            } else {
19829                "rp"
19830            }
19831        } else {
19832            "base"
19833        };
19834        variant
19835    }
19836
19837    #[allow(clippy::too_many_arguments)]
19838    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19839    #[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
19840    pub fn qmatvec_mmvq_batched(
19841        &self,
19842        bytes: &CudaSlice<u8>,
19843        aq: &CudaSlice<i8>,
19844        ad: &CudaSlice<f32>,
19845        m: usize,
19846        in_f: usize,
19847        out_f: usize,
19848        qtype: i32,
19849        row_bytes: usize,
19850        mcols: usize,
19851        scale: f32,
19852        rp: bool,
19853    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19854        const ROWS_PER_BLOCK: u32 = 4;
19855        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
19856        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
19857        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
19858        // weight keeps its rp-layout kernel family regardless of the override.
19859        let forced: Option<&'static str> = {
19860            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
19861            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
19862                .as_deref()
19863                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
19864        };
19865        let variant = match forced {
19866            Some(v) if !rp || v.contains("rp") => v,
19867            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
19868        };
19869        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
19870            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
19871        })?;
19872        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
19873        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
19874        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
19875        let variant = if mcols == 16 {
19876            if rp { "rp" } else { "base" }
19877        } else {
19878            variant
19879        };
19880        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
19881        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
19882        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
19883        // per-(token,row) chain (columns c >= m never execute in either form) ->
19884        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
19885        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
19886        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19887        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
19888        if b567
19889            && qtype == QT_NVFP4
19890            && rp
19891            && mcols == 8
19892            && (5..=7).contains(&m)
19893            && matches!(variant, "rpsc" | "rpr2w8")
19894        {
19895            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
19896            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
19897            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19898            let cfg = LaunchConfig {
19899                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
19900                block_dim: (32, ROWS_PER_BLOCK, 1),
19901                shared_mem_bytes: 0,
19902            };
19903            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19904            let __s_b = self.gpu.stream();
19905            let mut b = __s_b.launch_builder(&f);
19906            b.arg(bytes)
19907                .arg(aq)
19908                .arg(ad)
19909                .arg(&mut y)
19910                .arg(&inf)
19911                .arg(&outf)
19912                .arg(&mi)
19913                .arg(&rb);
19914            unsafe {
19915                b.launch(cfg)?;
19916            }
19917            if scale != 1.0 {
19918                self.scale_inplace(&mut y, scale, m * out_f)?;
19919            }
19920            return Ok(y);
19921        }
19922        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
19923            "base" => (base_name.into(), ROWS_PER_BLOCK),
19924            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
19925            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
19926            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
19927            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
19928            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
19929            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
19930            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
19931            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
19932            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
19933            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
19934            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
19935            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
19936            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
19937            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
19938        };
19939        debug_assert!(
19940            !rp || name.contains("_rp"),
19941            "rp weight dispatched to a GGUF-layout kernel"
19942        );
19943        let f = self.func(&name);
19944        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19945        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
19946        let smem = if name.contains("_r2sm_rp") {
19947            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
19948        } else {
19949            0
19950        };
19951        let cfg = LaunchConfig {
19952            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
19953            block_dim: (32, ROWS_PER_BLOCK, 1),
19954            shared_mem_bytes: smem,
19955        };
19956        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
19957        let __s_b = self.gpu.stream();
19958        let mut b = __s_b.launch_builder(&f);
19959        b.arg(bytes)
19960            .arg(aq)
19961            .arg(ad)
19962            .arg(&mut y)
19963            .arg(&inf)
19964            .arg(&outf)
19965            .arg(&mi)
19966            .arg(&rb);
19967        unsafe {
19968            b.launch(cfg)?;
19969        }
19970        if scale != 1.0 {
19971            self.scale_inplace(&mut y, scale, m * out_f)?;
19972        }
19973        Ok(y)
19974    }
19975
19976    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
19977    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
19978    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
19979    #[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
19980    pub fn qmatvec_batched_raw(
19981        &self,
19982        bytes: &CudaSlice<u8>,
19983        x: &CudaSlice<f32>,
19984        m: usize,
19985        in_f: usize,
19986        out_f: usize,
19987        qtype: i32,
19988        row_bytes: usize,
19989        mcols: usize,
19990        rp: bool,
19991    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19992        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19993        self.qmatvec_mmvq_batched(
19994            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
19995        )
19996    }
19997
19998    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
19999    #[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
20000    pub fn qmatvec_nvfp4_batched_raw(
20001        &self,
20002        bytes: &CudaSlice<u8>,
20003        x: &CudaSlice<f32>,
20004        m: usize,
20005        in_f: usize,
20006        out_f: usize,
20007        row_bytes: usize,
20008        mcols: usize,
20009        rp: bool,
20010    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20011        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
20012    }
20013
20014    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
20015    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
20016    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
20017    fn try_fp4_gemm(
20018        &self,
20019        w: &crate::model::GpuTensor,
20020        x: &CudaSlice<f32>,
20021        m: usize,
20022        in_f: usize,
20023        out_f: usize,
20024    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20025        use crate::model::GpuTensor;
20026        if cfg!(memra_portable_cuda) {
20027            return Ok(None);
20028        }
20029        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which ONLY the sm_120a fatbin contains:
20030        // cu/qmatvec_gemm.cu omits it on portable builds (MEMRA_PORTABLE_CUDA) AND on sm_100a
20031        // (build.rs passes -DMEMRA_DISABLE_NATIVE_FP4=1 there — the mxf4 block-scale MMA is an
20032        // sm_120a instruction encoding). Refuse at the door on EVERY build that lacks it. The
20033        // portable refusal alone was an enumeration, not a property: a 100a build is not
20034        // portable, so `MEMRA_FP4=1` sailed past it into Engine::func's "kernel not in any
20035        // fatbin" panic — found by the 100a fatbin-lookup census, lane/glm5-b200-prep-20260901
20036        // (same enumeration-vs-property class as the 2026-08-23 stub-polarity fixes in build.rs).
20037        if std::env::var("MEMRA_FP4").is_ok() {
20038            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
20039            assert!(
20040                konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a"),
20041                "MEMRA_FP4 forces the native mxf4 block-scale GEMM (qmatvec_gemm_nvfp4_fp4), \
20042                 which only the sm_120a fatbin contains — this is an sm_{} build. Unset \
20043                 MEMRA_FP4; the W4A8 int8 path is the correct default for NVFP4 weights.",
20044                env!("MEMRA_BUILT_CUDA_ARCH")
20045            );
20046        }
20047        if std::env::var("MEMRA_FP4").is_err() {
20048            return Ok(None);
20049        }
20050        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
20051        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
20052        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
20053        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
20054        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
20055        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
20056        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
20057        // for the common no-macro-scale case.
20058        #[cfg(memra_cutlass)]
20059        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
20060            if let GpuTensor::Quant {
20061                bytes,
20062                qtype,
20063                scale,
20064                row_bytes,
20065                cutlass,
20066                ..
20067            } = w
20068            {
20069                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
20070                    if let Some(cw) = cutlass {
20071                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
20072                        let y = self.cutlass_fp4_gemm(
20073                            &cw.b_packed,
20074                            &cw.sfb_swizzled,
20075                            x,
20076                            *scale,
20077                            m,
20078                            out_f,
20079                            in_f,
20080                        )?;
20081                        return Ok(Some(y));
20082                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
20083                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
20084                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
20085                        // (the load-time repack ~doubles it) — needed for models that don't fit the
20086                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
20087                        let (b_packed, sfb_sw) =
20088                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
20089                        let y =
20090                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
20091                        return Ok(Some(y));
20092                    }
20093                }
20094            }
20095        }
20096        if let GpuTensor::Quant {
20097            bytes,
20098            qtype,
20099            row_bytes,
20100            scale,
20101            rp,
20102            ..
20103        } = w
20104        {
20105            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
20106            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
20107            if *qtype == QT_NVFP4 && in_f.is_multiple_of(64) && !*rp {
20108                let y =
20109                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
20110                return Ok(Some(y));
20111            }
20112        }
20113        Ok(None)
20114    }
20115
20116    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
20117    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
20118    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
20119    #[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
20120    pub fn rms_norm_f16out(
20121        &self,
20122        x: &CudaSlice<f32>,
20123        w: &CudaSlice<f32>,
20124        dst: &mut CudaSlice<f32>,
20125        dst16: &mut CudaSlice<u8>,
20126        ncols: usize,
20127        nrows: usize,
20128        eps: f32,
20129    ) -> Result<(), Box<dyn std::error::Error>> {
20130        let f = self.func("rms_norm_f16out_f32");
20131        let cfg = LaunchConfig {
20132            grid_dim: (nrows as u32, 1, 1),
20133            block_dim: (rms_block(), 1, 1),
20134            shared_mem_bytes: 0,
20135        };
20136        let (nc, e) = (ncols as i32, eps);
20137        let __s_b = self.gpu.stream();
20138        let mut b = __s_b.launch_builder(&f);
20139        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
20140        unsafe {
20141            b.launch(cfg)?;
20142        }
20143        Ok(())
20144    }
20145
20146    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
20147    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
20148    #[allow(clippy::too_many_arguments)]
20149    pub fn add_rms_norm_f16out(
20150        &self,
20151        a: &CudaSlice<f32>,
20152        b: &CudaSlice<f32>,
20153        w: &CudaSlice<f32>,
20154        res: &mut CudaSlice<f32>,
20155        dst: &mut CudaSlice<f32>,
20156        dst16: &mut CudaSlice<u8>,
20157        ncols: usize,
20158        nrows: usize,
20159        eps: f32,
20160    ) -> Result<(), Box<dyn std::error::Error>> {
20161        let f = self.func("add_rms_norm_f16out_f32");
20162        let cfg = LaunchConfig {
20163            grid_dim: (nrows as u32, 1, 1),
20164            block_dim: (rms_block(), 1, 1),
20165            shared_mem_bytes: 0,
20166        };
20167        let (nc, e) = (ncols as i32, eps);
20168        let __s_lb = self.gpu.stream();
20169        let mut lb = __s_lb.launch_builder(&f);
20170        lb.arg(a)
20171            .arg(b)
20172            .arg(w)
20173            .arg(res)
20174            .arg(dst)
20175            .arg(dst16)
20176            .arg(&nc)
20177            .arg(&e);
20178        unsafe {
20179            lb.launch(cfg)?;
20180        }
20181        Ok(())
20182    }
20183
20184    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
20185    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
20186    pub fn matmul_group_xh(
20187        &self,
20188        ws: &[&crate::model::GpuTensor],
20189        x: &CudaSlice<f32>,
20190        xh: &CudaSlice<u8>,
20191        m: usize,
20192    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20193        let mut out = Vec::with_capacity(ws.len());
20194        let in_f = ws[0].in_features();
20195        for w in ws {
20196            if w.in_features() == in_f
20197                && m >= 16
20198                && !self.verify_exact_on()
20199                && let Some(y) = self.try_f16_gemm_pre(w, xh, m)?
20200            {
20201                out.push(y);
20202                continue;
20203            }
20204            out.push(self.matmul(w, x, m)?);
20205        }
20206        Ok(out)
20207    }
20208
20209    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
20210    /// GDN steps). Layouts [T, H].
20211    pub fn gdn_pad_mask(
20212        &self,
20213        beta: &mut CudaSlice<f32>,
20214        g_log: &mut CudaSlice<f32>,
20215        len_d: &CudaSlice<i32>,
20216        h: usize,
20217        t: usize,
20218    ) -> Result<(), Box<dyn std::error::Error>> {
20219        let f = self.func("gdn_pad_mask_f32");
20220        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
20221        let (hi, ti) = (h as i32, t as i32);
20222        let __s_b = self.gpu.stream();
20223        let mut b = __s_b.launch_builder(&f);
20224        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
20225        unsafe {
20226            b.launch(cfg)?;
20227        }
20228        Ok(())
20229    }
20230
20231    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
20232    /// gather for the padded prime graph's h_seed/hlast.
20233    pub fn row_gather_dev(
20234        &self,
20235        src: &CudaSlice<f32>,
20236        dst: &mut CudaSlice<f32>,
20237        len_d: &CudaSlice<i32>,
20238        ncols: usize,
20239    ) -> Result<(), Box<dyn std::error::Error>> {
20240        let f = self.func("row_gather_dev_f32");
20241        let cfg = LaunchConfig::for_num_elems(ncols as u32);
20242        let nc = ncols as i32;
20243        let __s_b = self.gpu.stream();
20244        let mut b = __s_b.launch_builder(&f);
20245        b.arg(src).arg(dst).arg(len_d).arg(&nc);
20246        unsafe {
20247            b.launch(cfg)?;
20248        }
20249        Ok(())
20250    }
20251
20252    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
20253    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
20254    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
20255    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
20256    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
20257    /// different in_f) falls back to its own `matmul` — behavior unchanged.
20258    pub fn matmul_group(
20259        &self,
20260        ws: &[&crate::model::GpuTensor],
20261        x: &CudaSlice<f32>,
20262        m: usize,
20263    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20264        use crate::model::GpuTensor;
20265        let mut out = Vec::with_capacity(ws.len());
20266        let any_mirror = ws
20267            .iter()
20268            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
20269        if m >= 16 && any_mirror && !self.verify_exact_on() {
20270            let in_f = ws[0].in_features();
20271            let xh = self.f16_act(x, m * in_f, in_f)?;
20272            for w in ws {
20273                if w.in_features() == in_f
20274                    && let Some(y) = self.try_f16_gemm_pre(w, &xh, m)?
20275                {
20276                    out.push(y);
20277                    continue;
20278                }
20279                out.push(self.matmul(w, x, m)?);
20280            }
20281            return Ok(out);
20282        }
20283        for w in ws {
20284            out.push(self.matmul(w, x, m)?);
20285        }
20286        Ok(out)
20287    }
20288
20289    /// Cross-request grouped matmul (task #13): run ONE projection group over the
20290    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
20291    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
20292    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
20293    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
20294    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
20295    pub fn matmul_group_multi(
20296        &self,
20297        ws: &[&crate::model::GpuTensor],
20298        xs: &[&CudaSlice<f32>],
20299        ms: &[usize],
20300    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
20301        assert_eq!(xs.len(), ms.len());
20302        let in_f = ws[0].in_features();
20303        let total: usize = ms.iter().sum();
20304        let mut xcat = self.uninit(total * in_f)?;
20305        let mut off = 0usize;
20306        for (x, &m) in xs.iter().zip(ms) {
20307            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
20308            off += m;
20309        }
20310        let ys = self.matmul_group(ws, &xcat, total)?;
20311        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
20312        for (w, y) in ws.iter().zip(ys) {
20313            let out_f = w.out_features();
20314            let mut off = 0usize;
20315            for (s, &m) in ms.iter().enumerate() {
20316                let mut ys_s = self.uninit(m * out_f)?;
20317                let src = y.slice(off * out_f..(off + m) * out_f);
20318                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
20319                out[s].push(ys_s);
20320                off += m;
20321            }
20322        }
20323        Ok(out)
20324    }
20325
20326    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
20327    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
20328    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
20329    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
20330    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
20331    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
20332    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
20333    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
20334    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
20335    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
20336        use crate::model::GpuTensor;
20337        if !legacy_quant_gemm_allowed(
20338            cfg!(memra_portable_cuda),
20339            cfg!(memra_hopper_mma),
20340            std::env::var_os("MEMRA_NO_GEMM").is_some(),
20341        ) {
20342            return false;
20343        }
20344        match w {
20345            GpuTensor::Quant { qtype, .. } => {
20346                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
20347                    || (*qtype == QT_NVFP4 && w.in_features().is_multiple_of(64))
20348            }
20349            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
20350        }
20351    }
20352
20353    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
20354    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
20355    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
20356    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
20357    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
20358    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
20359    #[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
20360    pub fn qmatvec_gemm(
20361        &self,
20362        w: &crate::model::GpuTensor,
20363        aq: &CudaSlice<i8>,
20364        ad: &CudaSlice<f32>,
20365        m: usize,
20366    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20367        use crate::model::GpuTensor;
20368        let in_f = w.in_features();
20369        let out_f = w.out_features();
20370        let (bytes, qtype, row_bytes, scale, rp) = match w {
20371            GpuTensor::Quant {
20372                bytes,
20373                qtype,
20374                row_bytes,
20375                scale,
20376                rp,
20377                ..
20378            } => (bytes, *qtype, *row_bytes, *scale, *rp),
20379            _ => unreachable!("gemm_supports guaranteed Quant"),
20380        };
20381        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
20382        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
20383        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
20384        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
20385        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
20386        if cfg!(memra_hopper_mma)
20387            && qtype == QT_Q8_0
20388            && out_f.is_multiple_of(64)
20389            && wgmma_gemm_enabled()
20390            && let GpuTensor::Quant { rp4: Some(m4), .. } = w
20391        {
20392            let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
20393            if scale != 1.0 {
20394                self.scale_inplace(&mut y, scale, m * out_f)?;
20395            }
20396            return Ok(y);
20397        }
20398        let name = match qtype {
20399            QT_Q8_0 => "qmatvec_gemm_q8_0",
20400            QT_Q4_K => "qmatvec_gemm_q4_K",
20401            QT_Q4_0 => {
20402                if rp {
20403                    "qmatvec_gemm_q4_0_rp"
20404                } else {
20405                    "qmatvec_gemm_q4_0"
20406                }
20407            }
20408            QT_Q5_K => "qmatvec_gemm_q5_K",
20409            QT_Q6_K => "qmatvec_gemm_q6_K",
20410            QT_NVFP4 => {
20411                if rp {
20412                    "qmatvec_gemm_nvfp4_rp"
20413                } else {
20414                    "qmatvec_gemm_nvfp4"
20415                }
20416            }
20417            _ => unreachable!(),
20418        };
20419        let f = self.func(name);
20420        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
20421        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
20422        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
20423        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
20424        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
20425        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
20426        let k1_tile = if is_k1 {
20427            k1_launch_override().unwrap_or((128, 128, 8))
20428        } else {
20429            (128, 128, 8)
20430        };
20431        let (bm, bn): (u32, u32) = if is_k1 {
20432            (k1_tile.0, k1_tile.1)
20433        } else {
20434            (64, 256)
20435        };
20436        let warps: u32 = if is_k1 {
20437            k1_tile.2
20438        } else {
20439            match qtype {
20440                QT_NVFP4 => 8,
20441                _ => 4,
20442            }
20443        };
20444        let cfg = LaunchConfig {
20445            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
20446            block_dim: (32, warps, 1),
20447            shared_mem_bytes: 0,
20448        };
20449        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20450        let __s_b = self.gpu.stream();
20451        let mut b = __s_b.launch_builder(&f);
20452        b.arg(bytes)
20453            .arg(aq)
20454            .arg(ad)
20455            .arg(&mut y)
20456            .arg(&inf)
20457            .arg(&outf)
20458            .arg(&mi)
20459            .arg(&rb);
20460        unsafe {
20461            b.launch(cfg)?;
20462        }
20463        if scale != 1.0 {
20464            self.scale_inplace(&mut y, scale, m * out_f)?;
20465        }
20466        Ok(y)
20467    }
20468
20469    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
20470    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
20471    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
20472    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
20473    #[allow(clippy::too_many_arguments)]
20474    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20475    #[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
20476    pub fn qmatvec_gemm_raw(
20477        &self,
20478        bytes: &CudaSlice<u8>,
20479        x: &CudaSlice<f32>,
20480        m: usize,
20481        in_f: usize,
20482        out_f: usize,
20483        qtype: i32,
20484        row_bytes: usize,
20485    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20486        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20487        let name = match qtype {
20488            QT_Q8_0 => "qmatvec_gemm_q8_0",
20489            QT_Q4_K => "qmatvec_gemm_q4_K",
20490            QT_Q4_0 => "qmatvec_gemm_q4_0",
20491            QT_Q5_K => "qmatvec_gemm_q5_K",
20492            QT_Q6_K => "qmatvec_gemm_q6_K",
20493            QT_NVFP4 => "qmatvec_gemm_nvfp4",
20494            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
20495            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
20496        };
20497        let f = self.func(name);
20498        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
20499        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
20500        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
20501        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
20502        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
20503        let k1_tile = if is_k1 {
20504            k1_launch_override().unwrap_or((128, 128, 8))
20505        } else {
20506            (128, 128, 8)
20507        };
20508        let (bm, bn): (u32, u32) = if is_k1 {
20509            (k1_tile.0, k1_tile.1)
20510        } else {
20511            (64, 256)
20512        };
20513        let warps: u32 = if is_k1 {
20514            k1_tile.2
20515        } else {
20516            match qtype {
20517                QT_NVFP4 | QT_NVFP4_RP => 8,
20518                _ => 4,
20519            }
20520        };
20521        let cfg = LaunchConfig {
20522            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
20523            block_dim: (32, warps, 1),
20524            shared_mem_bytes: 0,
20525        };
20526        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20527        let __s_b = self.gpu.stream();
20528        let mut b = __s_b.launch_builder(&f);
20529        b.arg(bytes)
20530            .arg(&aq)
20531            .arg(&ad)
20532            .arg(&mut y)
20533            .arg(&inf)
20534            .arg(&outf)
20535            .arg(&mi)
20536            .arg(&rb);
20537        unsafe {
20538            b.launch(cfg)?;
20539        }
20540        Ok(y)
20541    }
20542
20543    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
20544    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
20545    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
20546    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
20547    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
20548    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
20549    pub fn qmatvec_gemm_q8_0_wgmma_raw(
20550        &self,
20551        rp4: &CudaSlice<u8>,
20552        aq: &CudaSlice<i8>,
20553        ad: &CudaSlice<f32>,
20554        m: usize,
20555        in_f: usize,
20556        out_f: usize,
20557    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20558        assert!(
20559            out_f.is_multiple_of(64) && in_f.is_multiple_of(32),
20560            "wgmma GEMM needs out_f%64==0, in_f%32==0"
20561        );
20562        let f = self.func("qmatvec_gemm_q8_0_wgmma");
20563        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
20564        let cfg = LaunchConfig {
20565            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
20566            block_dim: (128, 1, 1),
20567            shared_mem_bytes: 0,
20568        };
20569        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
20570        let __s_b = self.gpu.stream();
20571        let mut b = __s_b.launch_builder(&f);
20572        b.arg(rp4)
20573            .arg(aq)
20574            .arg(ad)
20575            .arg(&mut y)
20576            .arg(&inf)
20577            .arg(&outf)
20578            .arg(&mi);
20579        unsafe {
20580            b.launch(cfg)?;
20581        }
20582        Ok(y)
20583    }
20584
20585    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
20586    pub fn scale_inplace(
20587        &self,
20588        y: &mut CudaSlice<f32>,
20589        s: f32,
20590        n: usize,
20591    ) -> Result<(), Box<dyn std::error::Error>> {
20592        let f = self.func("scale_f32");
20593        let cfg = LaunchConfig::for_num_elems(n as u32);
20594        let (sf, ni) = (s, n as i32);
20595        let __s_b = self.gpu.stream();
20596        let mut b = __s_b.launch_builder(&f);
20597        b.arg(y).arg(&sf).arg(&ni);
20598        unsafe {
20599            b.launch(cfg)?;
20600        }
20601        Ok(())
20602    }
20603
20604    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
20605    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
20606    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
20607    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
20608    pub fn bf16_to_f32(
20609        &self,
20610        data: &cudarc::driver::CudaView<'_, u8>,
20611        n: usize,
20612    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20613        let mut out = self.alloc_uninit::<f32>(n)?;
20614        let f = self.func("bf16_to_f32");
20615        let cfg = LaunchConfig::for_num_elems(n as u32);
20616        let ni = n as i32;
20617        let __s_b = self.gpu.stream();
20618        let mut b = __s_b.launch_builder(&f);
20619        b.arg(data).arg(&mut out).arg(&ni);
20620        unsafe {
20621            b.launch(cfg)?;
20622        }
20623        Ok(out)
20624    }
20625
20626    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
20627    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
20628    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
20629    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
20630    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
20631    /// calls, the spec-verify contract) vs plain linear.
20632    #[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
20633    fn linear_bf16_chunked(
20634        &self,
20635        x: &CudaSlice<f32>,
20636        data: &CudaSlice<u8>,
20637        m: usize,
20638        in_f: usize,
20639        out_f: usize,
20640        exact: bool,
20641        canonical_chunk_rows: Option<usize>,
20642    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20643        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
20644        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
20645        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
20646        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20647        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20648        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
20649        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
20650        let started = timing.then(std::time::Instant::now);
20651        let result =
20652            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
20653        if let Some(started) = started {
20654            use std::sync::atomic::Ordering;
20655            self.stream().synchronize()?;
20656            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
20657                + started.elapsed().as_nanos() as u64;
20658            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
20659                + (in_f * out_f * 2) as u64;
20660            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
20661            if calls.is_multiple_of(1024) {
20662                eprintln!(
20663                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
20664                     weight_gb={:.2}",
20665                    ns as f64 / 1.0e6,
20666                    ns as f64 / calls as f64 / 1.0e3,
20667                    wb as f64 / 1.0e9,
20668                );
20669            }
20670        }
20671        result
20672    }
20673
20674    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
20675    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
20676    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
20677    /// numeric-class doors (DEV_ROUTES precedent).
20678    pub(crate) fn bf16_mmv_on() -> bool {
20679        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20680        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
20681    }
20682
20683    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
20684    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
20685    fn matvec_bf16(
20686        &self,
20687        data: &CudaSlice<u8>,
20688        x: &CudaSlice<f32>,
20689        in_f: usize,
20690        out_f: usize,
20691    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20692        if data.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) {
20693            return Err(format!(
20694                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
20695                data.len(),
20696                x.len()
20697            )
20698            .into());
20699        }
20700        let mut y = self.alloc_uninit::<f32>(out_f)?;
20701        let f = self.func("matvec_bf16_f32acc");
20702        let cfg = LaunchConfig {
20703            grid_dim: (out_f as u32, 1, 1),
20704            block_dim: (mmv_block(), 1, 1),
20705            shared_mem_bytes: 0,
20706        };
20707        let ini = in_f as i32;
20708        let __s_bld = self.gpu.stream();
20709        let mut bld = __s_bld.launch_builder(&f);
20710        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
20711        unsafe {
20712            bld.launch(cfg)?;
20713        }
20714        Ok(y)
20715    }
20716
20717    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
20718    /// launches, a position upload, and the rope launch; the position is read directly from
20719    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
20720    #[allow(clippy::too_many_arguments)]
20721    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
20722    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
20723    /// Bit-identical to the split kernels; requires head_dim == 128 and
20724    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
20725    #[allow(clippy::too_many_arguments)]
20726    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
20727    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
20728    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
20729    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
20730    #[allow(clippy::too_many_arguments)]
20731    pub fn qk_norm_rope_append_inc_dcw_rows(
20732        &self,
20733        q_raw_t: &CudaSlice<f32>,
20734        k_raw_t: &CudaSlice<f32>,
20735        v_raw_t: &CudaSlice<f32>,
20736        qw: &CudaSlice<f32>,
20737        kw: &CudaSlice<f32>,
20738        q_out_t: &mut CudaSlice<f32>,
20739        k_out_t: &mut CudaSlice<f32>,
20740        tab: &CudaSlice<u64>,
20741        pos_t: &CudaSlice<i32>,
20742        same_session: bool,
20743        t: usize,
20744        kv_dim_k: usize,
20745        kv_dim_v: usize,
20746        k_tok_bytes: usize,
20747        v_tok_bytes: usize,
20748        head_dim: usize,
20749        n_dims: usize,
20750        nh_q: usize,
20751        nh_k: usize,
20752        eps: f32,
20753        freq_base: f32,
20754        freq_scale: f32,
20755        ff: Option<&CudaSlice<f32>>,
20756    ) -> Result<(), Box<dyn std::error::Error>> {
20757        if head_dim != 128
20758            || kv_dim_v != kv_dim_k
20759            || kv_dim_k != nh_k * head_dim
20760            || t == 0
20761            || t > 32
20762            || tab.len() < t * 6
20763            || pos_t.len() < t
20764            || q_raw_t.len() < t * nh_q * head_dim
20765            || k_raw_t.len() < t * nh_k * head_dim
20766            || v_raw_t.len() < t * kv_dim_v
20767            || q_out_t.len() < t * nh_q * head_dim
20768            || k_out_t.len() < t * nh_k * head_dim
20769        {
20770            return Err(format!(
20771                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
20772                 nh_q={nh_q} nh_k={nh_k}"
20773            )
20774            .into());
20775        }
20776        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
20777        let same_t: i32 = if same_session { t as i32 } else { 0 };
20778        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
20779        let cfg = LaunchConfig {
20780            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
20781            block_dim: (128, 1, 1),
20782            shared_mem_bytes: 0,
20783        };
20784        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
20785        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20786        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
20787        let null: u64 = 0;
20788        let __s_b = self.gpu.stream();
20789        let mut b = __s_b.launch_builder(&f);
20790        b.arg(q_raw_t)
20791            .arg(k_raw_t)
20792            .arg(v_raw_t)
20793            .arg(qw)
20794            .arg(kw)
20795            .arg(q_out_t)
20796            .arg(k_out_t)
20797            .arg(tab)
20798            .arg(pos_t)
20799            .arg(&same_t)
20800            .arg(&kvk)
20801            .arg(&kvv)
20802            .arg(&ktb)
20803            .arg(&vtb)
20804            .arg(&hd)
20805            .arg(&nd)
20806            .arg(&nq)
20807            .arg(&nk)
20808            .arg(&eps)
20809            .arg(&theta_scale)
20810            .arg(&freq_scale);
20811        match ff {
20812            Some(freqs) => {
20813                b.arg(freqs);
20814            }
20815            None => {
20816                b.arg(&null);
20817            }
20818        }
20819        unsafe {
20820            b.launch(cfg)?;
20821        }
20822        Ok(())
20823    }
20824
20825    #[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
20826    pub fn qk_norm_rope_append_inc_dcw(
20827        &self,
20828        q_raw: &CudaSlice<f32>,
20829        k_raw: &CudaSlice<f32>,
20830        v_raw: &CudaSlice<f32>,
20831        qw: &CudaSlice<f32>,
20832        kw: &CudaSlice<f32>,
20833        q_out: &mut CudaSlice<f32>,
20834        k_out: &mut CudaSlice<f32>,
20835        pos: &CudaSlice<i32>,
20836        k_plane: &mut CudaSlice<u8>,
20837        v_plane: &mut CudaSlice<u8>,
20838        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
20839        // (single) writer, exactly like the split append+inc pair it replaces.
20840        len_dev: &CudaSlice<i32>,
20841        base_dev: Option<&CudaSlice<i32>>,
20842        done_ctr: &mut CudaSlice<u32>,
20843        kv_dim_k: usize,
20844        kv_dim_v: usize,
20845        k_tok_bytes: usize,
20846        v_tok_bytes: usize,
20847        head_dim: usize,
20848        n_dims: usize,
20849        nh_q: usize,
20850        nh_k: usize,
20851        eps: f32,
20852        freq_base: f32,
20853        freq_scale: f32,
20854        ff: Option<&CudaSlice<f32>>,
20855    ) -> Result<(), Box<dyn std::error::Error>> {
20856        if head_dim != 128
20857            || kv_dim_v != kv_dim_k
20858            || kv_dim_k != nh_k * head_dim
20859            || q_raw.len() < nh_q * head_dim
20860            || k_raw.len() < nh_k * head_dim
20861            || v_raw.len() < kv_dim_v
20862            || q_out.len() < nh_q * head_dim
20863            || k_out.len() < nh_k * head_dim
20864            || pos.is_empty()
20865            || done_ctr.is_empty()
20866        {
20867            return Err(format!(
20868                "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}"
20869            )
20870            .into());
20871        }
20872        let f = self.func("qk_norm_rope_append_inc_dcw");
20873        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
20874        let cfg = LaunchConfig {
20875            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
20876            block_dim: (128, 1, 1),
20877            shared_mem_bytes: 0,
20878        };
20879        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
20880        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20881        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
20882        let null: u64 = 0;
20883        let __s_b = self.gpu.stream();
20884        let mut b = __s_b.launch_builder(&f);
20885        b.arg(q_raw)
20886            .arg(k_raw)
20887            .arg(v_raw)
20888            .arg(qw)
20889            .arg(kw)
20890            .arg(q_out)
20891            .arg(k_out)
20892            .arg(pos)
20893            .arg(&mut *k_plane)
20894            .arg(&mut *v_plane)
20895            .arg(len_dev);
20896        match base_dev {
20897            Some(base) => {
20898                b.arg(base);
20899            }
20900            None => {
20901                b.arg(&null);
20902            }
20903        }
20904        b.arg(&mut *done_ctr)
20905            .arg(&kvk)
20906            .arg(&kvv)
20907            .arg(&ktb)
20908            .arg(&vtb)
20909            .arg(&hd)
20910            .arg(&nd)
20911            .arg(&nq)
20912            .arg(&eps)
20913            .arg(&theta_scale)
20914            .arg(&freq_scale);
20915        match ff {
20916            Some(freqs) => {
20917                b.arg(freqs);
20918            }
20919            None => {
20920                b.arg(&null);
20921            }
20922        }
20923        unsafe {
20924            b.launch(cfg)?;
20925        }
20926        Ok(())
20927    }
20928
20929    #[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
20930    pub fn qk_norm_rope_into(
20931        &self,
20932        q_raw: &CudaSlice<f32>,
20933        k_raw: &CudaSlice<f32>,
20934        qw: &CudaSlice<f32>,
20935        kw: &CudaSlice<f32>,
20936        q_out: &mut CudaSlice<f32>,
20937        k_out: &mut CudaSlice<f32>,
20938        pos: &CudaSlice<i32>,
20939        head_dim: usize,
20940        n_dims: usize,
20941        nh_q: usize,
20942        nh_k: usize,
20943        eps: f32,
20944        freq_base: f32,
20945        freq_scale: f32,
20946        ff: Option<&CudaSlice<f32>>,
20947    ) -> Result<(), Box<dyn std::error::Error>> {
20948        if head_dim > 512
20949            || q_raw.len() < nh_q * head_dim
20950            || k_raw.len() < nh_k * head_dim
20951            || q_out.len() < nh_q * head_dim
20952            || k_out.len() < nh_k * head_dim
20953            || qw.len() < head_dim
20954            || kw.len() < head_dim
20955            || pos.is_empty()
20956        {
20957            return Err(format!(
20958                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
20959            )
20960            .into());
20961        }
20962        let f = self.func("qk_norm_rope_f32");
20963        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
20964        let cfg = LaunchConfig {
20965            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
20966            block_dim: (128, 1, 1),
20967            shared_mem_bytes: 0,
20968        };
20969        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
20970        let __s_b = self.gpu.stream();
20971        let mut b = __s_b.launch_builder(&f);
20972        b.arg(q_raw)
20973            .arg(k_raw)
20974            .arg(qw)
20975            .arg(kw)
20976            .arg(q_out)
20977            .arg(k_out)
20978            .arg(pos)
20979            .arg(&hd)
20980            .arg(&nd)
20981            .arg(&nq)
20982            .arg(&eps)
20983            .arg(&theta_scale)
20984            .arg(&freq_scale);
20985        match ff {
20986            Some(ffv) => {
20987                b.arg(ffv);
20988                unsafe {
20989                    b.launch(cfg)?;
20990                }
20991            }
20992            None => {
20993                let null: u64 = 0;
20994                b.arg(&null);
20995                unsafe {
20996                    b.launch(cfg)?;
20997                }
20998            }
20999        }
21000        Ok(())
21001    }
21002
21003    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
21004    /// launch computes a rank's whole O partial from its four canonical column blocks.
21005    #[allow(clippy::too_many_arguments)]
21006    pub fn matvec_f32_b4_into(
21007        &self,
21008        w: [&CudaSlice<f32>; 4],
21009        x: &CudaSlice<f32>,
21010        y: &mut CudaSlice<f32>,
21011        block_cols: usize,
21012        out_f: usize,
21013    ) -> Result<(), Box<dyn std::error::Error>> {
21014        if !block_cols.is_multiple_of(4)
21015            || x.len() < 4 * block_cols
21016            || y.len() < out_f
21017            || w.iter().any(|w| w.len() != out_f * block_cols)
21018        {
21019            return Err(format!(
21020                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
21021                x.len()
21022            )
21023            .into());
21024        }
21025        let f = self.func("matvec_f32_b4");
21026        let cfg = LaunchConfig {
21027            grid_dim: (out_f as u32, 1, 1),
21028            block_dim: (128, 1, 1),
21029            shared_mem_bytes: 0,
21030        };
21031        let (bc, of) = (block_cols as i32, out_f as i32);
21032        let __s_b = self.gpu.stream();
21033        let mut b = __s_b.launch_builder(&f);
21034        b.arg(w[0])
21035            .arg(w[1])
21036            .arg(w[2])
21037            .arg(w[3])
21038            .arg(x)
21039            .arg(y)
21040            .arg(&bc)
21041            .arg(&of);
21042        unsafe {
21043            b.launch(cfg)?;
21044        }
21045        Ok(())
21046    }
21047
21048    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
21049    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
21050    pub fn axpy_rows_seq_into(
21051        &self,
21052        x: &CudaSlice<f32>,
21053        w: &CudaSlice<f32>,
21054        y: &mut CudaSlice<f32>,
21055        width: usize,
21056        n_rows: usize,
21057    ) -> Result<(), Box<dyn std::error::Error>> {
21058        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
21059            return Err(format!(
21060                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
21061                x.len(),
21062                w.len(),
21063                y.len()
21064            )
21065            .into());
21066        }
21067        let f = self.func("axpy_rows_seq_f32");
21068        let cfg = LaunchConfig::for_num_elems(width as u32);
21069        let (wi, nr) = (width as i32, n_rows as i32);
21070        let __s_b = self.gpu.stream();
21071        let mut b = __s_b.launch_builder(&f);
21072        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
21073        unsafe {
21074            b.launch(cfg)?;
21075        }
21076        Ok(())
21077    }
21078
21079    /// Token-major sequential weighted row sums. Each token reduces exactly `slots` rows in
21080    /// canonical route order.
21081    pub fn axpy_rows_seq_tokens_into(
21082        &self,
21083        x: &CudaSlice<f32>,
21084        w: &CudaSlice<f32>,
21085        y: &mut CudaSlice<f32>,
21086        width: usize,
21087        slots: usize,
21088        tokens: usize,
21089    ) -> Result<(), Box<dyn std::error::Error>> {
21090        let rows = slots
21091            .checked_mul(tokens)
21092            .ok_or("axpy_rows_seq_tokens row count overflow")?;
21093        if x.len() < rows * width || w.len() < rows || y.len() < tokens * width {
21094            return Err(format!(
21095                "axpy_rows_seq_tokens geometry x={} w={} y={} width={width} \
21096                 slots={slots} tokens={tokens}",
21097                x.len(),
21098                w.len(),
21099                y.len()
21100            )
21101            .into());
21102        }
21103        let f = self.func("axpy_rows_seq_tokens_f32");
21104        let block = 256u32;
21105        let cfg = LaunchConfig {
21106            grid_dim: ((width as u32).div_ceil(block), tokens as u32, 1),
21107            block_dim: (block, 1, 1),
21108            shared_mem_bytes: 0,
21109        };
21110        let (wi, sl, tk) = (width as i32, slots as i32, tokens as i32);
21111        let __s_b = self.gpu.stream();
21112        let mut b = __s_b.launch_builder(&f);
21113        b.arg(x).arg(w).arg(y).arg(&wi).arg(&sl).arg(&tk);
21114        unsafe {
21115            b.launch(cfg)?;
21116        }
21117        Ok(())
21118    }
21119
21120    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
21121    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
21122    /// exact sequential FP chain of the base kernel over that window.
21123    #[allow(clippy::too_many_arguments)]
21124    pub fn axpy_rows_seq_md_off_into(
21125        &self,
21126        x: &CudaSlice<f32>,
21127        w_route: &CudaSlice<f32>,
21128        md: &CudaSlice<f32>,
21129        sel: &CudaSlice<i32>,
21130        y: &mut CudaSlice<f32>,
21131        width: usize,
21132        n_rows: usize,
21133        row0: usize,
21134    ) -> Result<(), Box<dyn std::error::Error>> {
21135        if x.len() < (row0 + n_rows) * width
21136            || w_route.len() < row0 + n_rows
21137            || sel.len() < row0 + n_rows
21138            || y.len() < width
21139        {
21140            return Err(format!(
21141                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
21142                 rows={n_rows} row0={row0}",
21143                x.len(),
21144                w_route.len(),
21145                sel.len(),
21146                y.len()
21147            )
21148            .into());
21149        }
21150        let f = self.func("axpy_rows_seq_md_off_f32");
21151        let cfg = LaunchConfig::for_num_elems(width as u32);
21152        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
21153        let __s_b = self.gpu.stream();
21154        let mut b = __s_b.launch_builder(&f);
21155        b.arg(x)
21156            .arg(w_route)
21157            .arg(md)
21158            .arg(sel)
21159            .arg(y)
21160            .arg(&wi)
21161            .arg(&nr)
21162            .arg(&r0);
21163        unsafe {
21164            b.launch(cfg)?;
21165        }
21166        Ok(())
21167    }
21168
21169    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
21170    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
21171    #[allow(clippy::too_many_arguments)]
21172    pub fn axpy_rows_seq_md_into(
21173        &self,
21174        x: &CudaSlice<f32>,
21175        w_route: &CudaSlice<f32>,
21176        md: &CudaSlice<f32>,
21177        sel: &CudaSlice<i32>,
21178        y: &mut CudaSlice<f32>,
21179        width: usize,
21180        n_rows: usize,
21181    ) -> Result<(), Box<dyn std::error::Error>> {
21182        if x.len() < n_rows * width
21183            || w_route.len() < n_rows
21184            || sel.len() < n_rows
21185            || y.len() < width
21186        {
21187            return Err(format!(
21188                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
21189                x.len(),
21190                w_route.len(),
21191                sel.len(),
21192                y.len()
21193            )
21194            .into());
21195        }
21196        let f = self.func("axpy_rows_seq_md_f32");
21197        let cfg = LaunchConfig::for_num_elems(width as u32);
21198        let (wi, nr) = (width as i32, n_rows as i32);
21199        let __s_b = self.gpu.stream();
21200        let mut b = __s_b.launch_builder(&f);
21201        b.arg(x)
21202            .arg(w_route)
21203            .arg(md)
21204            .arg(sel)
21205            .arg(y)
21206            .arg(&wi)
21207            .arg(&nr);
21208        unsafe {
21209            b.launch(cfg)?;
21210        }
21211        Ok(())
21212    }
21213
21214    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
21215    #[allow(clippy::too_many_arguments)]
21216    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
21217    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
21218    /// land column-major-of-rows: yq[c*out_q + row] etc.
21219    #[allow(clippy::too_many_arguments)]
21220    pub fn matvec_bf16_qkvg_tcol_into(
21221        &self,
21222        wq: &CudaSlice<u8>,
21223        wk: &CudaSlice<u8>,
21224        wv: &CudaSlice<u8>,
21225        wg: &CudaSlice<u8>,
21226        x_t: &CudaSlice<f32>,
21227        yq: &mut CudaSlice<f32>,
21228        yk: &mut CudaSlice<f32>,
21229        yv: &mut CudaSlice<f32>,
21230        yg: &mut CudaSlice<f32>,
21231        in_f: usize,
21232        out_q: usize,
21233        out_kv: usize,
21234        out_g: usize,
21235        t: usize,
21236    ) -> Result<(), Box<dyn std::error::Error>> {
21237        if t == 0
21238            || t > 8
21239            || !in_f.is_multiple_of(8)
21240            || x_t.len() < t * in_f
21241            || yq.len() < t * out_q
21242            || yk.len() < t * out_kv
21243            || yv.len() < t * out_kv
21244            || (out_g > 0 && yg.len() < t * out_g)
21245        {
21246            return Err("matvec_bf16_qkvg_tcol geometry".into());
21247        }
21248        let grid = out_q + 2 * out_kv + out_g;
21249        let cfg = LaunchConfig {
21250            grid_dim: (grid as u32, 1, 1),
21251            block_dim: (mmv_block(), 1, 1),
21252            shared_mem_bytes: 0,
21253        };
21254        let (ini, oq, okv, og, ti) = (
21255            in_f as i32,
21256            out_q as i32,
21257            out_kv as i32,
21258            out_g as i32,
21259            t as i32,
21260        );
21261        let __s_b = self.gpu.stream();
21262        // One runtime-T program for every live width. The compile-time 2/4/8 twins are
21263        // retained in the fatbin as research controls, but dispatching them by the current
21264        // batch width changes kernels inside a request when peers arrive or retire. That is
21265        // a load-history numeric-program switch, and their pre-twin TOKFP receipts did not
21266        // qualify it (Hermes `64fa2b55baf0d887`).
21267        let f = self.func("matvec_bf16_qkvg_tcol");
21268        let mut b = __s_b.launch_builder(&f);
21269        b.arg(wq)
21270            .arg(wk)
21271            .arg(wv)
21272            .arg(wg)
21273            .arg(x_t)
21274            .arg(yq)
21275            .arg(yk)
21276            .arg(yv)
21277            .arg(yg)
21278            .arg(&ini)
21279            .arg(&oq)
21280            .arg(&okv)
21281            .arg(&og)
21282            .arg(&ti);
21283        unsafe {
21284            b.launch(cfg)?;
21285        }
21286        Ok(())
21287    }
21288
21289    #[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
21290    pub fn matvec_bf16_qkvg_into(
21291        &self,
21292        wq: &CudaSlice<u8>,
21293        wk: &CudaSlice<u8>,
21294        wv: &CudaSlice<u8>,
21295        wg: &CudaSlice<u8>,
21296        x: &CudaSlice<f32>,
21297        yq: &mut CudaSlice<f32>,
21298        yk: &mut CudaSlice<f32>,
21299        yv: &mut CudaSlice<f32>,
21300        yg: &mut CudaSlice<f32>,
21301        in_f: usize,
21302        out_q: usize,
21303        out_kv: usize,
21304        out_g: usize,
21305    ) -> Result<(), Box<dyn std::error::Error>> {
21306        if !in_f.is_multiple_of(8)
21307            || wq.len() != out_q * in_f * 2
21308            || wk.len() != out_kv * in_f * 2
21309            || wv.len() != out_kv * in_f * 2
21310            || wg.len() < out_g * in_f * 2
21311            || x.len() < in_f
21312            || yq.len() < out_q
21313            || yk.len() < out_kv
21314            || yv.len() < out_kv
21315            || (out_g > 0 && yg.len() < out_g)
21316        {
21317            return Err(format!(
21318                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
21319            )
21320            .into());
21321        }
21322        let f = self.func("matvec_bf16_qkvg");
21323        let cfg = LaunchConfig {
21324            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
21325            block_dim: (mmv_block(), 1, 1),
21326            shared_mem_bytes: 0,
21327        };
21328        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
21329        let __s_b = self.gpu.stream();
21330        let mut b = __s_b.launch_builder(&f);
21331        b.arg(wq)
21332            .arg(wk)
21333            .arg(wv)
21334            .arg(wg)
21335            .arg(x)
21336            .arg(yq)
21337            .arg(yk)
21338            .arg(yv)
21339            .arg(yg)
21340            .arg(&inf)
21341            .arg(&oq)
21342            .arg(&okv)
21343            .arg(&og);
21344        unsafe {
21345            b.launch(cfg)?;
21346        }
21347        Ok(())
21348    }
21349
21350    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
21351    pub fn matvec_bf16_b4_into(
21352        &self,
21353        w: [&CudaSlice<u8>; 4],
21354        x: &CudaSlice<f32>,
21355        y: &mut CudaSlice<f32>,
21356        block_cols: usize,
21357        out_f: usize,
21358    ) -> Result<(), Box<dyn std::error::Error>> {
21359        if !block_cols.is_multiple_of(8)
21360            || x.len() < 4 * block_cols
21361            || y.len() < out_f
21362            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
21363        {
21364            return Err(format!(
21365                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
21366                x.len()
21367            )
21368            .into());
21369        }
21370        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
21371        // bit-identical per row (the second row's stream hides the first's reduce tail).
21372        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21373        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
21374        let f = self.func(if x2 {
21375            "matvec_bf16_b4_x2"
21376        } else {
21377            "matvec_bf16_b4"
21378        });
21379        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
21380        let cfg = LaunchConfig {
21381            grid_dim: (grid as u32, 1, 1),
21382            block_dim: (mmv_block(), 1, 1),
21383            shared_mem_bytes: 0,
21384        };
21385        let (bc, of) = (block_cols as i32, out_f as i32);
21386        let __s_b = self.gpu.stream();
21387        let mut b = __s_b.launch_builder(&f);
21388        b.arg(w[0])
21389            .arg(w[1])
21390            .arg(w[2])
21391            .arg(w[3])
21392            .arg(x)
21393            .arg(y)
21394            .arg(&bc)
21395            .arg(&of);
21396        unsafe {
21397            b.launch(cfg)?;
21398        }
21399        Ok(())
21400    }
21401
21402    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
21403    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
21404    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
21405    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
21406    /// t=1 program).
21407    pub fn matvec_bf16_b4_tcol_into(
21408        &self,
21409        w: [&CudaSlice<u8>; 4],
21410        x_t: &CudaSlice<f32>,
21411        y_t: &mut CudaSlice<f32>,
21412        block_cols: usize,
21413        out_f: usize,
21414        t: usize,
21415    ) -> Result<(), Box<dyn std::error::Error>> {
21416        if !block_cols.is_multiple_of(8)
21417            || t == 0
21418            || t > 8
21419            || x_t.len() < t * 4 * block_cols
21420            || y_t.len() < t * out_f
21421            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
21422        {
21423            return Err(format!(
21424                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
21425                x_t.len()
21426            )
21427            .into());
21428        }
21429        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
21430            return Err(
21431                "b4 tcol verify is qualified against the plain b4 kernel only \
21432                        (MEMRA_B4_X2=1 is a different t=1 program)"
21433                    .into(),
21434            );
21435        }
21436        // Keep one runtime-T program at every live width. Compile-time twins remain research
21437        // controls only; selecting them from the changing batch width switches programs
21438        // mid-request.
21439        let cfg = LaunchConfig {
21440            grid_dim: (out_f as u32, 1, 1),
21441            block_dim: (mmv_block(), 1, 1),
21442            shared_mem_bytes: 0,
21443        };
21444        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
21445        let __s_b = self.gpu.stream();
21446        let f = self.func("matvec_bf16_b4_tcol");
21447        let mut b = __s_b.launch_builder(&f);
21448        b.arg(w[0])
21449            .arg(w[1])
21450            .arg(w[2])
21451            .arg(w[3])
21452            .arg(x_t)
21453            .arg(y_t)
21454            .arg(&bc)
21455            .arg(&of)
21456            .arg(&ti);
21457        unsafe {
21458            b.launch(cfg)?;
21459        }
21460        Ok(())
21461    }
21462
21463    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
21464    /// q8_0 row bytes for an `in_f`-wide weight row: one 34-byte block per 32 weights.
21465    pub fn q8_0_row_bytes(in_f: usize) -> usize {
21466        in_f / 32 * 34
21467    }
21468
21469    /// Encode a resident bf16 weight slab into its q8_0 mirror (MEMRA_STEP_TP_W8). Runs once
21470    /// per matrix at load; the block program is the one `quant_K_block` writes for the KV
21471    /// cache, so the two formats cannot drift apart.
21472    pub fn encode_q8_0_from_bf16(
21473        &self,
21474        w_bf16: &CudaSlice<u8>,
21475        out: &mut CudaSlice<u8>,
21476        in_f: usize,
21477        out_f: usize,
21478    ) -> Result<(), Box<dyn std::error::Error>> {
21479        if !in_f.is_multiple_of(32)
21480            || w_bf16.len() < in_f * out_f * 2
21481            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
21482        {
21483            return Err(format!(
21484                "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
21485                w_bf16.len(),
21486                out.len()
21487            )
21488            .into());
21489        }
21490        let f = self.func("encode_q8_0_rows_from_bf16");
21491        // Flat 1D grid of (row, 32-block) pairs, 4 pairs per block: rows on grid.y would cap
21492        // at 65535 and the LM head has 128896 rows.
21493        const PAIRS_PER_BLOCK: u32 = 4;
21494        let pairs = (out_f * (in_f / 32)) as u64;
21495        let cfg = LaunchConfig {
21496            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
21497            block_dim: (32, PAIRS_PER_BLOCK, 1),
21498            shared_mem_bytes: 0,
21499        };
21500        let (ini, outi) = (in_f as i32, out_f as i32);
21501        let __s_b = self.gpu.stream();
21502        let mut b = __s_b.launch_builder(&f);
21503        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
21504        unsafe {
21505            b.launch(cfg)?;
21506        }
21507        Ok(())
21508    }
21509
21510    /// ROW-RANGE-VIEW twin of `encode_q8_0_from_bf16`. Identical kernel, identical launch
21511    /// geometry, identical per-row program: only the operand type differs, because the split
21512    /// decode paths hold their rows as a `CudaView` of the resident slab, not as an owned slab.
21513    pub fn encode_q8_0_from_bf16_view(
21514        &self,
21515        w_bf16: &cudarc::driver::CudaView<'_, u8>,
21516        out: &mut CudaSlice<u8>,
21517        in_f: usize,
21518        out_f: usize,
21519    ) -> Result<(), Box<dyn std::error::Error>> {
21520        if !in_f.is_multiple_of(32)
21521            || w_bf16.len() < in_f * out_f * 2
21522            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
21523        {
21524            return Err(format!(
21525                "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
21526                w_bf16.len(),
21527                out.len()
21528            )
21529            .into());
21530        }
21531        let f = self.func("encode_q8_0_rows_from_bf16");
21532        const PAIRS_PER_BLOCK: u32 = 4;
21533        let pairs = (out_f * (in_f / 32)) as u64;
21534        let cfg = LaunchConfig {
21535            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
21536            block_dim: (32, PAIRS_PER_BLOCK, 1),
21537            shared_mem_bytes: 0,
21538        };
21539        let (ini, outi) = (in_f as i32, out_f as i32);
21540        let __s_b = self.gpu.stream();
21541        let mut b = __s_b.launch_builder(&f);
21542        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
21543        unsafe {
21544            b.launch(cfg)?;
21545        }
21546        Ok(())
21547    }
21548
21549    /// Fused q8_0 QKV against a q8_1 activation (MEMRA_STEP_TP_W8): one launch over the
21550    /// stacked q/k/v rows, each row running the exact `qmatvec_q8_0_mmvq_rp` per-row program.
21551    /// Bit-identical to three per-matrix mmvq calls; it exists because those three launches
21552    /// plus the activation quantize measured SLOWER than the bf16 fused kernel.
21553    #[allow(clippy::too_many_arguments)]
21554    pub fn qmatvec_q8_0_qkv_rp_into(
21555        &self,
21556        wq: &CudaSlice<u8>,
21557        wk: &CudaSlice<u8>,
21558        wv: &CudaSlice<u8>,
21559        aq: &CudaSlice<i8>,
21560        ad: &CudaSlice<f32>,
21561        yq: &mut CudaSlice<f32>,
21562        yk: &mut CudaSlice<f32>,
21563        yv: &mut CudaSlice<f32>,
21564        in_f: usize,
21565        out_q: usize,
21566        out_kv: usize,
21567    ) -> Result<(), Box<dyn std::error::Error>> {
21568        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
21569        let rows = out_q + 2 * out_kv;
21570        let nblk = in_f / 32;
21571        if !in_f.is_multiple_of(32)
21572            || aq.len() < in_f
21573            || ad.len() < nblk
21574            || yq.len() < out_q
21575            || yk.len() < out_kv
21576            || yv.len() < out_kv
21577            || wq.len() < out_q * nblk * 34
21578            || wk.len() < out_kv * nblk * 34
21579            || wv.len() < out_kv * nblk * 34
21580        {
21581            return Err(
21582                format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
21583            );
21584        }
21585        let f = self.func("qmatvec_q8_0_qkv_rp");
21586        let cfg = LaunchConfig {
21587            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21588            block_dim: (32, ROWS_PER_BLOCK, 1),
21589            shared_mem_bytes: 0,
21590        };
21591        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
21592        let __s_b = self.gpu.stream();
21593        let mut b = __s_b.launch_builder(&f);
21594        b.arg(wq)
21595            .arg(wk)
21596            .arg(wv)
21597            .arg(aq)
21598            .arg(ad)
21599            .arg(yq)
21600            .arg(yk)
21601            .arg(yv)
21602            .arg(&ini)
21603            .arg(&oq)
21604            .arg(&okv);
21605        unsafe {
21606            b.launch(cfg)?;
21607        }
21608        Ok(())
21609    }
21610
21611    /// Fused q8_0 O projection over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8): one
21612    /// launch, one warp per output row, per-block reduce then add — the same shape
21613    /// `matvec_bf16_b4` uses, against a q8_1 activation.
21614    #[allow(clippy::too_many_arguments)]
21615    pub fn qmatvec_q8_0_b4_rp_into(
21616        &self,
21617        w: [&CudaSlice<u8>; 4],
21618        aq: &CudaSlice<i8>,
21619        ad: &CudaSlice<f32>,
21620        y: &mut CudaSlice<f32>,
21621        block_cols: usize,
21622        out_f: usize,
21623    ) -> Result<(), Box<dyn std::error::Error>> {
21624        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
21625        let nblk = block_cols / 32;
21626        if !block_cols.is_multiple_of(32)
21627            || aq.len() < 4 * block_cols
21628            || ad.len() < 4 * nblk
21629            || y.len() < out_f
21630            || w.iter().any(|p| p.len() < out_f * nblk * 34)
21631        {
21632            return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
21633        }
21634        let f = self.func("qmatvec_q8_0_b4_rp");
21635        let cfg = LaunchConfig {
21636            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21637            block_dim: (32, ROWS_PER_BLOCK, 1),
21638            shared_mem_bytes: 0,
21639        };
21640        let (bc, of) = (block_cols as i32, out_f as i32);
21641        let __s_b = self.gpu.stream();
21642        let mut b = __s_b.launch_builder(&f);
21643        b.arg(w[0])
21644            .arg(w[1])
21645            .arg(w[2])
21646            .arg(w[3])
21647            .arg(aq)
21648            .arg(ad)
21649            .arg(y)
21650            .arg(&bc)
21651            .arg(&of);
21652        unsafe {
21653            b.launch(cfg)?;
21654        }
21655        Ok(())
21656    }
21657
21658    /// T-column twin of `matvec_bf16_via_q8_mirror`: one q8 launch over all t rows, sharing the
21659    /// same pointer-keyed mirror cache and a t-wide q8_1 activation.
21660    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
21661    fn matvec_bf16_via_q8_mirror_t(
21662        &self,
21663        data: &CudaSlice<u8>,
21664        x: &CudaSlice<f32>,
21665        y: &mut CudaSlice<f32>,
21666        in_f: usize,
21667        out_f: usize,
21668        t: usize,
21669    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
21670        use cudarc::driver::DevicePtr;
21671        let key = {
21672            let s = self.gpu.stream();
21673            let (p, _g) = data.device_ptr(&s);
21674            (p, in_f as u32, out_f as u32)
21675        };
21676        {
21677            let mut mirrors = self
21678                .w8_mirrors
21679                .lock()
21680                .map_err(|_| "w8 mirror map is poisoned")?;
21681            if !mirrors.contains_key(&key) {
21682                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
21683                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
21684                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
21685                mirrors.insert(key, planar);
21686            }
21687        }
21688        let nblk = in_f / 32;
21689        // The t-wide activation scratch is keyed by (in_f, t-cap) so a wider walk regrows it.
21690        let akey = in_f * 64 + t.min(32);
21691        {
21692            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21693            if !act.contains_key(&akey) {
21694                let aq = self.alloc_i8_uninit(32 * in_f)?;
21695                let ad = self.alloc_uninit::<f32>(32 * nblk)?;
21696                act.insert(akey, (aq, ad));
21697            }
21698            let (aq, ad) = act.get_mut(&akey).expect("just inserted");
21699            self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
21700        }
21701        let mirrors = self
21702            .w8_mirrors
21703            .lock()
21704            .map_err(|_| "w8 mirror map is poisoned")?;
21705        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21706        let mirror = mirrors.get(&key).expect("built above");
21707        let (aq, ad) = act.get(&akey).expect("built above");
21708        const ROWS_PER_BLOCK: u32 = 4;
21709        let (ini, of) = (in_f as i32, out_f as i32);
21710        // MEMRA_Q8T_WONCE=1: the weight-once twin — one row grid, each weight int4 loaded once
21711        // and dotted against all t columns. The `_t` form re-streams the shared weights per
21712        // column through __ldcs (measured 1.43-1.67x a single-column call for 2 columns).
21713        if q8t_wonce_on() && t <= 32 {
21714            let f = self.func(if t <= 8 {
21715                "qmatvec_q8_0_rows_tw"
21716            } else {
21717                "qmatvec_q8_0_rows_tw32"
21718            });
21719            let cfg = LaunchConfig {
21720                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21721                block_dim: (32, ROWS_PER_BLOCK, 1),
21722                shared_mem_bytes: 0,
21723            };
21724            let ti = t as i32;
21725            let __s_b = self.gpu.stream();
21726            let mut b = __s_b.launch_builder(&f);
21727            b.arg(mirror)
21728                .arg(aq)
21729                .arg(ad)
21730                .arg(&mut *y)
21731                .arg(&ini)
21732                .arg(&of)
21733                .arg(&ti);
21734            unsafe {
21735                b.launch(cfg)?;
21736            }
21737            return Ok(Some(()));
21738        }
21739        let f = self.func("qmatvec_q8_0_rows_t");
21740        let cfg = LaunchConfig {
21741            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
21742            block_dim: (32, ROWS_PER_BLOCK, 1),
21743            shared_mem_bytes: 0,
21744        };
21745        let __s_b = self.gpu.stream();
21746        let mut b = __s_b.launch_builder(&f);
21747        b.arg(mirror)
21748            .arg(aq)
21749            .arg(ad)
21750            .arg(&mut *y)
21751            .arg(&ini)
21752            .arg(&of);
21753        unsafe {
21754            b.launch(cfg)?;
21755        }
21756        Ok(Some(()))
21757    }
21758
21759    /// Get-or-build this bf16 weight's q8_0 mirror and run the GEMV through it. Returns
21760    /// `None` when the shape has no mirror form, so the caller falls back to bf16.
21761    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
21762    fn matvec_bf16_via_q8_mirror(
21763        &self,
21764        data: &CudaSlice<u8>,
21765        x: &CudaSlice<f32>,
21766        y: &mut CudaSlice<f32>,
21767        in_f: usize,
21768        out_f: usize,
21769    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
21770        use cudarc::driver::DevicePtr;
21771        let key = {
21772            let s = self.gpu.stream();
21773            let (p, _g) = data.device_ptr(&s);
21774            (p, in_f as u32, out_f as u32)
21775        };
21776        {
21777            let mut mirrors = self
21778                .w8_mirrors
21779                .lock()
21780                .map_err(|_| "w8 mirror map is poisoned")?;
21781            if !mirrors.contains_key(&key) {
21782                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
21783                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
21784                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
21785                mirrors.insert(key, planar);
21786                // Which weights this half actually covers is not obvious from the call graph:
21787                // the head and the shared expert may reach the GPU through the rows fast path
21788                // or the fused dual-silu launcher instead of here. One line per mirror answers
21789                // that without a profiler (the hybrid half measured +0.1% and this is how we
21790                // find out whether it even fired).
21791                if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
21792                    eprintln!(
21793                        "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
21794                        mirrors.len()
21795                    );
21796                }
21797            }
21798        }
21799        let nblk = in_f / 32;
21800        {
21801            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21802            if !act.contains_key(&in_f) {
21803                let aq = self.alloc_uninit::<i8>(in_f)?;
21804                let ad = self.alloc_uninit::<f32>(nblk)?;
21805                act.insert(in_f, (aq, ad));
21806            }
21807            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
21808            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
21809        }
21810        let mirrors = self
21811            .w8_mirrors
21812            .lock()
21813            .map_err(|_| "w8 mirror map is poisoned")?;
21814        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21815        let mirror = mirrors.get(&key).expect("built above");
21816        let (aq, ad) = act.get(&in_f).expect("built above");
21817        self.qmatvec_mmvq_into(
21818            mirror,
21819            aq,
21820            ad,
21821            1,
21822            in_f,
21823            out_f,
21824            QT_Q8_0,
21825            Self::q8_0_row_bytes(in_f),
21826            1.0,
21827            true,
21828            y,
21829        )?;
21830        Ok(Some(()))
21831    }
21832
21833    /// T-column q8_0 QKV for the VERIFY walk (MEMRA_STEP_TP_W8). nsys put the bf16 twin
21834    /// `matvec_bf16_qkvg_tcol` at 12.3% of spec GPU time and `matvec_bf16_b4_tcol` at 24.8%:
21835    /// the W8 door had replaced only the decode kernels, so 37% of the verify still streamed
21836    /// bf16. Bit-identical to `t` separate `qmatvec_q8_0_qkv_rp` calls.
21837    #[allow(clippy::too_many_arguments)]
21838    pub fn qmatvec_q8_0_qkv_rp_t_into(
21839        &self,
21840        wq: &CudaSlice<u8>,
21841        wk: &CudaSlice<u8>,
21842        wv: &CudaSlice<u8>,
21843        aq: &CudaSlice<i8>,
21844        ad: &CudaSlice<f32>,
21845        yq: &mut CudaSlice<f32>,
21846        yk: &mut CudaSlice<f32>,
21847        yv: &mut CudaSlice<f32>,
21848        in_f: usize,
21849        out_q: usize,
21850        out_kv: usize,
21851        t: usize,
21852    ) -> Result<(), Box<dyn std::error::Error>> {
21853        const ROWS_PER_BLOCK: u32 = 4;
21854        let rows = out_q + 2 * out_kv;
21855        let nblk = in_f / 32;
21856        if !in_f.is_multiple_of(32)
21857            || t == 0
21858            || aq.len() < t * in_f
21859            || ad.len() < t * nblk
21860            || yq.len() < t * out_q
21861            || yk.len() < t * out_kv
21862            || yv.len() < t * out_kv
21863        {
21864            return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
21865        }
21866        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
21867        // MEMRA_Q8T_WONCE=1: weight-once twin — see qmatvec.cu's `_tw` block for why the `_t`
21868        // form re-streams the fully-shared QKV weights per column (__ldcs + column grid axis;
21869        // measured 1.67x a single-column call for 2 columns).
21870        if q8t_wonce_on() && t <= 32 {
21871            let f = self.func(if t <= 8 {
21872                "qmatvec_q8_0_qkv_rp_tw"
21873            } else {
21874                "qmatvec_q8_0_qkv_rp_tw32"
21875            });
21876            let cfg = LaunchConfig {
21877                grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21878                block_dim: (32, ROWS_PER_BLOCK, 1),
21879                shared_mem_bytes: 0,
21880            };
21881            let ti = t as i32;
21882            let __s_b = self.gpu.stream();
21883            let mut b = __s_b.launch_builder(&f);
21884            b.arg(wq)
21885                .arg(wk)
21886                .arg(wv)
21887                .arg(aq)
21888                .arg(ad)
21889                .arg(yq)
21890                .arg(yk)
21891                .arg(yv)
21892                .arg(&ini)
21893                .arg(&oq)
21894                .arg(&okv)
21895                .arg(&ti);
21896            unsafe {
21897                b.launch(cfg)?;
21898            }
21899            return Ok(());
21900        }
21901        let f = self.func("qmatvec_q8_0_qkv_rp_t");
21902        let cfg = LaunchConfig {
21903            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
21904            block_dim: (32, ROWS_PER_BLOCK, 1),
21905            shared_mem_bytes: 0,
21906        };
21907        let __s_b = self.gpu.stream();
21908        let mut b = __s_b.launch_builder(&f);
21909        b.arg(wq)
21910            .arg(wk)
21911            .arg(wv)
21912            .arg(aq)
21913            .arg(ad)
21914            .arg(yq)
21915            .arg(yk)
21916            .arg(yv)
21917            .arg(&ini)
21918            .arg(&oq)
21919            .arg(&okv);
21920        unsafe {
21921            b.launch(cfg)?;
21922        }
21923        Ok(())
21924    }
21925
21926    /// T-column q8_0 o_proj over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8, verify walk).
21927    /// Bit-identical to `t` separate `qmatvec_q8_0_b4_rp` calls.
21928    #[allow(clippy::too_many_arguments)]
21929    pub fn qmatvec_q8_0_b4_rp_t_into(
21930        &self,
21931        w: [&CudaSlice<u8>; 4],
21932        aq: &CudaSlice<i8>,
21933        ad: &CudaSlice<f32>,
21934        y: &mut CudaSlice<f32>,
21935        block_cols: usize,
21936        out_f: usize,
21937        t: usize,
21938    ) -> Result<(), Box<dyn std::error::Error>> {
21939        const ROWS_PER_BLOCK: u32 = 4;
21940        let nblk = block_cols / 32;
21941        if !block_cols.is_multiple_of(32)
21942            || t == 0
21943            || aq.len() < t * 4 * block_cols
21944            || ad.len() < t * 4 * nblk
21945            || y.len() < t * out_f
21946        {
21947            return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
21948        }
21949        let (bc, of) = (block_cols as i32, out_f as i32);
21950        // MEMRA_Q8T_WONCE=1: weight-once twin (see qmatvec.cu; `_t` measured 1.43x for 2 columns
21951        // on fully-shared o_proj weights).
21952        if q8t_wonce_on() && t <= 32 {
21953            let f = self.func(if t <= 8 {
21954                "qmatvec_q8_0_b4_rp_tw"
21955            } else {
21956                "qmatvec_q8_0_b4_rp_tw32"
21957            });
21958            let cfg = LaunchConfig {
21959                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
21960                block_dim: (32, ROWS_PER_BLOCK, 1),
21961                shared_mem_bytes: 0,
21962            };
21963            let ti = t as i32;
21964            let __s_b = self.gpu.stream();
21965            let mut b = __s_b.launch_builder(&f);
21966            b.arg(w[0])
21967                .arg(w[1])
21968                .arg(w[2])
21969                .arg(w[3])
21970                .arg(aq)
21971                .arg(ad)
21972                .arg(y)
21973                .arg(&bc)
21974                .arg(&of)
21975                .arg(&ti);
21976            unsafe {
21977                b.launch(cfg)?;
21978            }
21979            return Ok(());
21980        }
21981        let f = self.func("qmatvec_q8_0_b4_rp_t");
21982        let cfg = LaunchConfig {
21983            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
21984            block_dim: (32, ROWS_PER_BLOCK, 1),
21985            shared_mem_bytes: 0,
21986        };
21987        let __s_b = self.gpu.stream();
21988        let mut b = __s_b.launch_builder(&f);
21989        b.arg(w[0])
21990            .arg(w[1])
21991            .arg(w[2])
21992            .arg(w[3])
21993            .arg(aq)
21994            .arg(ad)
21995            .arg(y)
21996            .arg(&bc)
21997            .arg(&of);
21998        unsafe {
21999            b.launch(cfg)?;
22000        }
22001        Ok(())
22002    }
22003
22004    /// MEMRA_W8_VIEW: the q8_0 mirror for a bf16 GEMV whose weight is a ROW-RANGE VIEW.
22005    /// `MEMRA_W8_HYBRID` hangs off `matvec_bf16_into`, and the two split decode paths pinned in
22006    /// the step37 serving env send only their HI half there: HEAD_SPLIT runs
22007    /// `rank1.matvec_bf16_into(head_hi)` beside `e.matvec_bf16_view_into(head_lo)`, and
22008    /// SHEXP_OVERLAP does the same with the shared-expert down rows. The view launcher had no
22009    /// mirror, so the lo half kept streaming 2 B/w while its twin ran at 1.0625, and because the
22010    /// halves execute CONCURRENTLY on the two cards the critical path is the SLOW half.
22011    /// NUMERIC CLASS: identical to the rest of `MEMRA_STEP_TP_W8`, so it carries that argmax
22012    /// acceptance and that maxdiff class, not a new one. Default OFF until measured.
22013    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
22014    fn matvec_bf16_view_via_q8_mirror(
22015        &self,
22016        data: &cudarc::driver::CudaView<'_, u8>,
22017        x: &CudaSlice<f32>,
22018        y: &mut CudaSlice<f32>,
22019        in_f: usize,
22020        out_f: usize,
22021    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
22022        use cudarc::driver::DevicePtr;
22023        let key = {
22024            let s = self.gpu.stream();
22025            let (p, _g) = data.device_ptr(&s);
22026            (p, in_f as u32, out_f as u32)
22027        };
22028        {
22029            let mut mirrors = self
22030                .w8_mirrors
22031                .lock()
22032                .map_err(|_| "w8 mirror map is poisoned")?;
22033            if !mirrors.contains_key(&key) {
22034                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
22035                self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
22036                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
22037                mirrors.insert(key, planar);
22038                // Unconditional, once per distinct shape: a door with no announce cannot be read
22039                // in BOTH directions, and this lane was already burned once by a sweep that
22040                // inferred "never engages" from a log line that did not exist in the tree.
22041                eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
22042            }
22043        }
22044        let nblk = in_f / 32;
22045        {
22046            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
22047            if !act.contains_key(&in_f) {
22048                let aq = self.alloc_uninit::<i8>(in_f)?;
22049                let ad = self.alloc_uninit::<f32>(nblk)?;
22050                act.insert(in_f, (aq, ad));
22051            }
22052            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
22053            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
22054        }
22055        let mirrors = self
22056            .w8_mirrors
22057            .lock()
22058            .map_err(|_| "w8 mirror map is poisoned")?;
22059        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
22060        let mirror = mirrors.get(&key).expect("built above");
22061        let (aq, ad) = act.get(&in_f).expect("built above");
22062        self.qmatvec_mmvq_into(
22063            mirror,
22064            aq,
22065            ad,
22066            1,
22067            in_f,
22068            out_f,
22069            QT_Q8_0,
22070            Self::q8_0_row_bytes(in_f),
22071            1.0,
22072            true,
22073            y,
22074        )?;
22075        Ok(Some(()))
22076    }
22077
22078    pub fn matvec_bf16_into(
22079        &self,
22080        data: &CudaSlice<u8>,
22081        x: &CudaSlice<f32>,
22082        y: &mut CudaSlice<f32>,
22083        in_f: usize,
22084        out_f: usize,
22085    ) -> Result<(), Box<dyn std::error::Error>> {
22086        if data.len() != in_f * out_f * 2
22087            || x.len() < in_f
22088            || !in_f.is_multiple_of(8)
22089            || y.len() < out_f
22090        {
22091            return Err(format!(
22092                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22093                data.len(),
22094                x.len(),
22095                y.len()
22096            )
22097            .into());
22098        }
22099        // MEMRA_STEP_TP_W8, hybrid half: route this GEMV through a q8_0 mirror of the same
22100        // weight. Covers exactly the bf16 GEMVs that are NOT in a TP resident bank — the LM
22101        // head (324.4 -> 163.7 us measured), the shared-expert down rows (13.0 -> 5.6 us) and
22102        // the dense-FFN layers. Same numeric class as the QKV/o_proj arms (int8 dp4a with
22103        // per-32 scales), so it rides the same argmax acceptance; the bf16 slab stays resident
22104        // for prefill. The mirror builds on first use and is keyed by the slab's pointer.
22105        if step_tp_w8_on()
22106            && w8_hybrid_on()
22107            && in_f.is_multiple_of(32)
22108            && out_f >= 64
22109            && let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)?
22110        {
22111            return Ok(());
22112        }
22113        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
22114        // block, exact f32acc per-row program — cures the 1-iteration latency
22115        // starvation (shexp down measured 420GB/s at in_f=1280).
22116        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22117        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
22118            && in_f <= 2048;
22119        if x4 {
22120            let f = self.func("matvec_bf16_f32acc_x4");
22121            let cfg = LaunchConfig {
22122                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
22123                block_dim: (mmv_block(), 1, 1),
22124                shared_mem_bytes: 0,
22125            };
22126            let (ini, outi) = (in_f as i32, out_f as i32);
22127            let __s_b = self.gpu.stream();
22128            let mut b = __s_b.launch_builder(&f);
22129            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
22130            unsafe {
22131                b.launch(cfg)?;
22132            }
22133            return Ok(());
22134        }
22135        let f = self.func("matvec_bf16_f32acc");
22136        let cfg = LaunchConfig {
22137            grid_dim: (out_f as u32, 1, 1),
22138            block_dim: (mmv_block(), 1, 1),
22139            shared_mem_bytes: 0,
22140        };
22141        let ini = in_f as i32;
22142        let __s_b = self.gpu.stream();
22143        let mut b = __s_b.launch_builder(&f);
22144        b.arg(data).arg(x).arg(y).arg(&ini);
22145        unsafe {
22146            b.launch(cfg)?;
22147        }
22148        Ok(())
22149    }
22150
22151    /// BF16 matvec over activation/output views. Automatic TP4 attention keeps each rank's
22152    /// O-projection input and canonical partial inside persistent slabs, so copying either view
22153    /// into a temporary allocation would give back the bandwidth and allocator win that TP is
22154    /// meant to provide.
22155    pub fn matvec_bf16_views_into(
22156        &self,
22157        data: &CudaSlice<u8>,
22158        x: &cudarc::driver::CudaView<'_, f32>,
22159        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
22160        in_f: usize,
22161        out_f: usize,
22162    ) -> Result<(), Box<dyn std::error::Error>> {
22163        if data.len() != in_f * out_f * 2
22164            || x.len() < in_f
22165            || !in_f.is_multiple_of(8)
22166            || y.len() < out_f
22167        {
22168            return Err(format!(
22169                "matvec_bf16_views_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22170                data.len(),
22171                x.len(),
22172                y.len()
22173            )
22174            .into());
22175        }
22176        let f = self.func("matvec_bf16_f32acc");
22177        let cfg = LaunchConfig {
22178            grid_dim: (out_f as u32, 1, 1),
22179            block_dim: (mmv_block(), 1, 1),
22180            shared_mem_bytes: 0,
22181        };
22182        let ini = in_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);
22186        unsafe {
22187            b.launch(cfg)?;
22188        }
22189        Ok(())
22190    }
22191
22192    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
22193    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
22194    pub fn matvec_bf16_view_into(
22195        &self,
22196        data: &cudarc::driver::CudaView<'_, u8>,
22197        x: &CudaSlice<f32>,
22198        y: &mut CudaSlice<f32>,
22199        in_f: usize,
22200        out_f: usize,
22201    ) -> Result<(), Box<dyn std::error::Error>> {
22202        if data.len() != in_f * out_f * 2
22203            || x.len() < in_f
22204            || !in_f.is_multiple_of(8)
22205            || y.len() < out_f
22206        {
22207            return Err(format!(
22208                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
22209                data.len(),
22210                x.len(),
22211                y.len()
22212            )
22213            .into());
22214        }
22215        if w8_view_on()
22216            && step_tp_w8_on()
22217            && w8_hybrid_on()
22218            && in_f.is_multiple_of(32)
22219            && out_f >= 64
22220            && let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)?
22221        {
22222            return Ok(());
22223        }
22224        let f = self.func("matvec_bf16_f32acc");
22225        let cfg = LaunchConfig {
22226            grid_dim: (out_f as u32, 1, 1),
22227            block_dim: (mmv_block(), 1, 1),
22228            shared_mem_bytes: 0,
22229        };
22230        let ini = in_f as i32;
22231        let __s_b = self.gpu.stream();
22232        let mut b = __s_b.launch_builder(&f);
22233        b.arg(data).arg(x).arg(y).arg(&ini);
22234        unsafe {
22235            b.launch(cfg)?;
22236        }
22237        Ok(())
22238    }
22239
22240    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
22241    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
22242    pub fn matvec_bf16_raw_out(
22243        &self,
22244        w: &CudaSlice<u8>,
22245        x: &CudaSlice<f32>,
22246        y_raw: u64,
22247        in_f: usize,
22248        out_f: usize,
22249    ) -> Result<(), Box<dyn std::error::Error>> {
22250        if w.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) || y_raw == 0 {
22251            return Err("matvec_bf16_raw_out geometry".into());
22252        }
22253        let f = self.func("matvec_bf16_f32acc");
22254        let cfg = LaunchConfig {
22255            grid_dim: (out_f as u32, 1, 1),
22256            block_dim: (mmv_block(), 1, 1),
22257            shared_mem_bytes: 0,
22258        };
22259        let ini = in_f as i32;
22260        let __s_b = self.gpu.stream();
22261        let mut b = __s_b.launch_builder(&f);
22262        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
22263        unsafe {
22264            b.launch(cfg)?;
22265        }
22266        Ok(())
22267    }
22268
22269    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
22270    /// UVA pointers so the caller passes persistent-static rows without holding locks).
22271    /// Exact per-element sequence of the split add + add_scaled_rows pair.
22272    pub fn add3_raw(
22273        &self,
22274        a: &CudaSlice<f32>,
22275        b: &CudaSlice<f32>,
22276        sh_raw: u64,
22277        scale_raw: u64,
22278        dst: &mut CudaSlice<f32>,
22279        n: usize,
22280    ) -> Result<(), Box<dyn std::error::Error>> {
22281        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
22282            return Err("add3_raw geometry".into());
22283        }
22284        let f = self.func("add3_f32");
22285        let cfg = LaunchConfig {
22286            grid_dim: ((n as u32).div_ceil(256), 1, 1),
22287            block_dim: (256, 1, 1),
22288            shared_mem_bytes: 0,
22289        };
22290        let ni = n as i32;
22291        let __s_b = self.gpu.stream();
22292        let mut bld = __s_b.launch_builder(&f);
22293        bld.arg(a)
22294            .arg(b)
22295            .arg(&sh_raw)
22296            .arg(&scale_raw)
22297            .arg(dst)
22298            .arg(&ni);
22299        unsafe {
22300            bld.launch(cfg)?;
22301        }
22302        Ok(())
22303    }
22304
22305    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
22306    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
22307    pub fn matvec_bf16_down_addscale_into(
22308        &self,
22309        w: &CudaSlice<u8>,
22310        x: &CudaSlice<f32>,
22311        scale: &CudaSlice<f32>,
22312        dst: &mut CudaSlice<f32>,
22313        in_f: usize,
22314        out_f: usize,
22315    ) -> Result<(), Box<dyn std::error::Error>> {
22316        if w.len() != in_f * out_f * 2
22317            || x.len() < in_f
22318            || !in_f.is_multiple_of(8)
22319            || dst.len() < out_f
22320            || scale.is_empty()
22321        {
22322            return Err("matvec_bf16_down_addscale geometry".into());
22323        }
22324        let f = self.func("matvec_bf16_down_addscale");
22325        let cfg = LaunchConfig {
22326            grid_dim: (out_f as u32, 1, 1),
22327            block_dim: (mmv_block(), 1, 1),
22328            shared_mem_bytes: 0,
22329        };
22330        let ini = in_f as i32;
22331        let __s_b = self.gpu.stream();
22332        let mut b = __s_b.launch_builder(&f);
22333        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
22334        unsafe {
22335            b.launch(cfg)?;
22336        }
22337        Ok(())
22338    }
22339
22340    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
22341    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
22342    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
22343    #[allow(clippy::too_many_arguments)]
22344    pub fn matvec_bf16_dual_silu_rows_into(
22345        &self,
22346        wg: &CudaSlice<u8>,
22347        wu: &CudaSlice<u8>,
22348        x: &CudaSlice<f32>,
22349        act: &mut CudaSlice<f32>,
22350        in_f: usize,
22351        out_f: usize,
22352        limit: Option<f32>,
22353        t: usize,
22354    ) -> Result<(), Box<dyn std::error::Error>> {
22355        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
22356            return Err("matvec_bf16_dual_silu_rows geometry".into());
22357        }
22358        let f = self.func("matvec_bf16_dual_silu_rows");
22359        let cfg = LaunchConfig {
22360            grid_dim: (out_f as u32, t as u32, 1),
22361            block_dim: (mmv_block(), 1, 1),
22362            shared_mem_bytes: 0,
22363        };
22364        let (ini, outi) = (in_f as i32, out_f as i32);
22365        let lim = limit.unwrap_or(0.0);
22366        let __s_b = self.gpu.stream();
22367        let mut b = __s_b.launch_builder(&f);
22368        b.arg(wg)
22369            .arg(wu)
22370            .arg(x)
22371            .arg(&mut *act)
22372            .arg(&ini)
22373            .arg(&outi)
22374            .arg(&lim);
22375        unsafe {
22376            b.launch(cfg)?;
22377        }
22378        Ok(())
22379    }
22380
22381    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
22382    pub fn matvec_bf16_rows_into(
22383        &self,
22384        w: &CudaSlice<u8>,
22385        x: &CudaSlice<f32>,
22386        y: &mut CudaSlice<f32>,
22387        in_f: usize,
22388        out_f: usize,
22389        t: usize,
22390    ) -> Result<(), Box<dyn std::error::Error>> {
22391        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || !in_f.is_multiple_of(8)
22392        {
22393            return Err("matvec_bf16_rows geometry".into());
22394        }
22395        // MEMRA_STEP_TP_W8 + MEMRA_W8_HYBRID, t > 1: the VERIFY walk's shexp/dense rows land
22396        // here too (`matvec_bf16_f32acc_x4_rows` was 78 launches/round at 56.5 us in a spec
22397        // capture, ~162 ms of GPU over 37 rounds), and the t==1 gate below skipped them. The
22398        // t-column q8 kernel is bit-identical to t single-row calls.
22399        if (2..=32).contains(&t)
22400            && step_tp_w8_on()
22401            && w8_hybrid_on()
22402            && in_f.is_multiple_of(32)
22403            && out_f >= 64
22404            && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
22405        {
22406            return Ok(());
22407        }
22408        // MEMRA_STEP_TP_W8: the LM head reaches the device HERE, not through
22409        // matvec_bf16_into — the W8 trace showed the hybrid half building mirrors only for
22410        // in_f=1280 out_f=4096 (the shared-expert down rows, which SHEXP_OVERLAP already
22411        // hides, hence its +0.1%). Route the t=1 decode row through the q8 mirror; wider t
22412        // (the verify walk) keeps bf16 so the prefill class is untouched.
22413        if t == 1
22414            && step_tp_w8_on()
22415            && w8_hybrid_on()
22416            && in_f.is_multiple_of(32)
22417            && out_f >= 64
22418            && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
22419        {
22420            return Ok(());
22421        }
22422        // MEMRA_BF16_TCOLS_WIDE (lane/glm5-matvec door T, default ON since 2026-08-31): t=2..=16 rides the
22423        // weight-once t-column class instead of the grid.y=t per-token weight re-read below.
22424        // Placed AFTER the W8-mirror intercepts (their precedence unchanged). Bit-identical
22425        // per (row, token) to the _rows kernel by the tcols class's standing construction
22426        // (order-pinned per-token chains + the identical red[256] tree); the motivating call
22427        // is the DFlash2 drafter's t=15 block-head matmul, which re-read the 1.269 GB lm
22428        // head 15x per spec round. Rollback seam: unset or =0 falls through unchanged.
22429        if (2..=16).contains(&t) && bf16_tcols_wide_on() {
22430            if BF16_TCOLS_WIDE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
22431                eprintln!(
22432                    "[bf16-tcols-wide] engaged: t={t} in_f={in_f} out_f={out_f} rides the \
22433                     weight-once tcols class (MEMRA_BF16_TCOLS_WIDE=1)"
22434                );
22435            }
22436            if t <= 8 {
22437                return self.matvec_bf16_tcols_into(w, x, y, in_f, out_f, t);
22438            }
22439            return self.matvec_bf16_tcols16_into(w, x, y, in_f, out_f, t);
22440        }
22441        let f = self.func("matvec_bf16_f32acc_x4_rows");
22442        let cfg = LaunchConfig {
22443            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
22444            block_dim: (mmv_block(), 1, 1),
22445            shared_mem_bytes: 0,
22446        };
22447        let (ini, outi) = (in_f as i32, out_f as i32);
22448        let __s_b = self.gpu.stream();
22449        let mut b = __s_b.launch_builder(&f);
22450        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
22451        unsafe {
22452            b.launch(cfg)?;
22453        }
22454        Ok(())
22455    }
22456
22457    /// T-COLUMN twin of the bf16 rows matvec (lane/glm5-verify-batch, the
22458    /// varlen-batched-cores pattern): one block owns 4 output rows for ALL t tokens, so the
22459    /// weight pack is read ONCE and reused across tokens — vs the `_rows` twin's grid.y=t
22460    /// per-token weight re-read. Per-(row,token) BIT-IDENTICAL to the t=1 program by
22461    /// construction (order-pinned single-chain accumulators, identical shared-tree reduce
22462    /// per token — LAW:vl-bit-identity-order-pinning); the `glm5_verify_batch_gpu` tcols
22463    /// bit-gate holds it. t is bounded by the kernel's MEMRA_BF16_TCOLS_MAX = 8.
22464    pub fn matvec_bf16_tcols_into(
22465        &self,
22466        w: &CudaSlice<u8>,
22467        x: &CudaSlice<f32>,
22468        y: &mut CudaSlice<f32>,
22469        in_f: usize,
22470        out_f: usize,
22471        t: usize,
22472    ) -> Result<(), Box<dyn std::error::Error>> {
22473        if x.len() < t * in_f
22474            || y.len() < t * out_f
22475            || !(2..=8).contains(&t)
22476            || !in_f.is_multiple_of(8)
22477        {
22478            return Err("matvec_bf16_tcols geometry".into());
22479        }
22480        // MEMRA_BF16_TCOLS_X1 (lane/glm5-matvec door X, default ON since 2026-08-31): one row per block
22481        // (grid.x = out_f) — 4x the wave count on the ~one-wave trunk grids (census: same
22482        // kernel runs 59% of peak at 512..2048 blocks, 80% at 38720). Per-row body and
22483        // reduce tree verbatim — bit-identical per (row, token). Rollback: unset or =0.
22484        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the chosen grid
22485        // form takes its `_rf` fused-reduce-tail twin — one barrier sequence shared by the t
22486        // columns plus intra-warp shuffles at the identical pairing (9t -> 3 barriers per
22487        // block). Composes with door X (grid choice first, tail twin second). Requires a
22488        // power-of-two block (the fused tail must pass exactly through s=32); any other
22489        // MEMRA_MMV_BLOCK falls through to the standing tree. Rollback: unset or =0.
22490        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
22491        if rf
22492            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
22493                == 0
22494        {
22495            eprintln!(
22496                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
22497                 shared across the t token columns + intra-warp shuffles at the identical \
22498                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
22499            );
22500        }
22501        let x1 = bf16_tcols_x1_on();
22502        let (fname, grid_x) = match (x1, rf) {
22503            (true, true) => ("matvec_bf16_f32acc_x1_tcols_rf", out_f as u32),
22504            (true, false) => ("matvec_bf16_f32acc_x1_tcols", out_f as u32),
22505            (false, true) => ("matvec_bf16_f32acc_x4_tcols_rf", out_f.div_ceil(4) as u32),
22506            (false, false) => ("matvec_bf16_f32acc_x4_tcols", out_f.div_ceil(4) as u32),
22507        };
22508        if x1 && BF16_TCOLS_X1_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
22509            eprintln!(
22510                "[bf16-tcols-x1] engaged: one-row-per-block tcols grid \
22511                 (MEMRA_BF16_TCOLS_X1=1)"
22512            );
22513        }
22514        let f = self.func(fname);
22515        let cfg = LaunchConfig {
22516            grid_dim: (grid_x, 1, 1),
22517            block_dim: (mmv_block(), 1, 1),
22518            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
22519        };
22520        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22521        let __s_b = self.gpu.stream();
22522        let mut b = __s_b.launch_builder(&f);
22523        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22524        unsafe {
22525            b.launch(cfg)?;
22526        }
22527        Ok(())
22528    }
22529
22530    /// WIDE-T twin of [`Self::matvec_bf16_tcols_into`] (lane/glm5-matvec door T,
22531    /// `MEMRA_BF16_TCOLS_WIDE`): t = 9..=16 through the SEPARATE `..._tcols16` kernel — its
22532    /// acc[16] register footprint never touches the priced t<=8 class (the qmatvec `_tw32`
22533    /// acc-sizing lesson). Bit-identical per (row, token) to the t=1 program by the same
22534    /// order-pinned construction; gated by `glm5_matvec_doors_gpu`.
22535    pub fn matvec_bf16_tcols16_into(
22536        &self,
22537        w: &CudaSlice<u8>,
22538        x: &CudaSlice<f32>,
22539        y: &mut CudaSlice<f32>,
22540        in_f: usize,
22541        out_f: usize,
22542        t: usize,
22543    ) -> Result<(), Box<dyn std::error::Error>> {
22544        if x.len() < t * in_f
22545            || y.len() < t * out_f
22546            || !(9..=16).contains(&t)
22547            || !in_f.is_multiple_of(8)
22548        {
22549            return Err("matvec_bf16_tcols16 geometry".into());
22550        }
22551        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the wide-t twin
22552        // takes its `_rf` fused tail too — the drafter head's t=15 is the extreme case (135
22553        // barriers -> 6 per block). Same power-of-two block guard as the t<=8 dispatch.
22554        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
22555        if rf
22556            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
22557                == 0
22558        {
22559            eprintln!(
22560                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
22561                 shared across the t token columns + intra-warp shuffles at the identical \
22562                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
22563            );
22564        }
22565        let f = self.func(if rf {
22566            "matvec_bf16_f32acc_x4_tcols16_rf"
22567        } else {
22568            "matvec_bf16_f32acc_x4_tcols16"
22569        });
22570        let cfg = LaunchConfig {
22571            grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
22572            block_dim: (mmv_block(), 1, 1),
22573            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
22574        };
22575        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22576        let __s_b = self.gpu.stream();
22577        let mut b = __s_b.launch_builder(&f);
22578        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22579        unsafe {
22580            b.launch(cfg)?;
22581        }
22582        Ok(())
22583    }
22584
22585    /// GATE-ONLY launcher for door R arms no route dispatches (`glm5_matvec_doors_gpu`):
22586    /// the shifted-pairing RED twin (`matvec_bf16_f32acc_x1_tcols_rf_redshift`, the arm that
22587    /// proves the bit bar can see an association change) and the `_rf` twins at t=1 (the
22588    /// routed launchers refuse t<2; the door-R bar covers t=1..=16, so the degenerate
22589    /// column-loop bounds are gated here). `kernel` is an ALLOWLIST, not a name proxy.
22590    #[allow(clippy::too_many_arguments)]
22591    pub fn matvec_bf16_tcols_gate_kernel_into(
22592        &self,
22593        kernel: &str,
22594        w: &CudaSlice<u8>,
22595        x: &CudaSlice<f32>,
22596        y: &mut CudaSlice<f32>,
22597        in_f: usize,
22598        out_f: usize,
22599        t: usize,
22600    ) -> Result<(), Box<dyn std::error::Error>> {
22601        let (grid_x, t_max) = match kernel {
22602            "matvec_bf16_f32acc_x1_tcols_rf" | "matvec_bf16_f32acc_x1_tcols_rf_redshift" => {
22603                (out_f as u32, 8usize)
22604            }
22605            "matvec_bf16_f32acc_x4_tcols_rf" => (out_f.div_ceil(4) as u32, 8usize),
22606            "matvec_bf16_f32acc_x4_tcols16_rf" => (out_f.div_ceil(4) as u32, 16usize),
22607            _ => return Err("matvec_bf16_tcols_gate_kernel_into: unknown kernel".into()),
22608        };
22609        if x.len() < t * in_f
22610            || y.len() < t * out_f
22611            || !(1..=t_max).contains(&t)
22612            || !in_f.is_multiple_of(8)
22613            || !mmv_block().is_power_of_two()
22614        {
22615            return Err("matvec_bf16_tcols_gate_kernel geometry".into());
22616        }
22617        let f = self.func(kernel);
22618        let cfg = LaunchConfig {
22619            grid_dim: (grid_x, 1, 1),
22620            block_dim: (mmv_block(), 1, 1),
22621            shared_mem_bytes: (t as u32) * mmv_block() * 4,
22622        };
22623        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
22624        let __s_b = self.gpu.stream();
22625        let mut b = __s_b.launch_builder(&f);
22626        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
22627        unsafe {
22628            b.launch(cfg)?;
22629        }
22630        Ok(())
22631    }
22632
22633    /// DECODE-EXACT matmul for the glm5 verify-batch walk (lane/glm5-verify-batch): the
22634    /// exact `matmul_decode_exact` dispatch with ONE addition — FloatBf16 weights at
22635    /// t=2..=8 under `MEMRA_BF16_MMV` ride the tcols twin above (weight read once for all
22636    /// t rows). Refused back to `matmul_decode_exact` whenever the t=1 decode chain would
22637    /// ride the W8 q8-mirror class instead of the bf16 rows kernel (the decode-parity law:
22638    /// the m>1 class must equal the m=1 class). Every class stays per-row bit-exact vs
22639    /// the t=1 chain; only the verify-batch walk calls this.
22640    pub fn matmul_rows_exact(
22641        &self,
22642        w: &crate::model::GpuTensor,
22643        x: &CudaSlice<f32>,
22644        m: usize,
22645    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22646        use crate::model::GpuTensor;
22647        if let GpuTensor::FloatBf16 { data, .. } = w
22648            && (2..=8).contains(&m)
22649            && Self::bf16_mmv_on()
22650            && w.in_features().is_multiple_of(8)
22651            && !(step_tp_w8_on() && w8_hybrid_on())
22652        {
22653            let (in_f, out_f) = (w.in_features(), w.out_features());
22654            // Door W: rows-exact is verify-walk-only by contract, so its y is a pooled
22655            // draw (vws_uninit == alloc_uninit with the door off).
22656            let mut y = self.vws_uninit(m * out_f)?;
22657            self.matvec_bf16_tcols_into(data, x, &mut y, in_f, out_f, m)?;
22658            return Ok(y);
22659        }
22660        self.matmul_decode_exact(w, x, m)
22661    }
22662
22663    #[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
22664    pub fn matvec_bf16_dual_silu_into(
22665        &self,
22666        wg: &CudaSlice<u8>,
22667        wu: &CudaSlice<u8>,
22668        x: &CudaSlice<f32>,
22669        act: &mut CudaSlice<f32>,
22670        in_f: usize,
22671        out_f: usize,
22672        limit: Option<f32>,
22673    ) -> Result<(), Box<dyn std::error::Error>> {
22674        if wg.len() != in_f * out_f * 2
22675            || wu.len() != in_f * out_f * 2
22676            || x.len() < in_f
22677            || !in_f.is_multiple_of(8)
22678            || act.len() < out_f
22679        {
22680            return Err("matvec_bf16_dual_silu geometry".into());
22681        }
22682        let f = self.func("matvec_bf16_dual_silu");
22683        let cfg = LaunchConfig {
22684            grid_dim: (out_f as u32, 1, 1),
22685            block_dim: (mmv_block(), 1, 1),
22686            shared_mem_bytes: 0,
22687        };
22688        let (ini, outi) = (in_f as i32, out_f as i32);
22689        let lim = limit.unwrap_or(0.0);
22690        let __s_b = self.gpu.stream();
22691        let mut b = __s_b.launch_builder(&f);
22692        b.arg(wg)
22693            .arg(wu)
22694            .arg(x)
22695            .arg(act)
22696            .arg(&ini)
22697            .arg(&outi)
22698            .arg(&lim);
22699        unsafe {
22700            b.launch(cfg)?;
22701        }
22702        Ok(())
22703    }
22704
22705    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
22706    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
22707    #[allow(clippy::too_many_arguments)]
22708    pub fn matvec_bf16_dual_view_into(
22709        &self,
22710        wg: &cudarc::driver::CudaView<'_, u8>,
22711        wu: &cudarc::driver::CudaView<'_, u8>,
22712        x: &CudaSlice<f32>,
22713        yg: &mut CudaSlice<f32>,
22714        yu: &mut CudaSlice<f32>,
22715        in_f: usize,
22716        out_f: usize,
22717    ) -> Result<(), Box<dyn std::error::Error>> {
22718        if wg.len() != in_f * out_f * 2
22719            || wu.len() != in_f * out_f * 2
22720            || x.len() < in_f
22721            || !in_f.is_multiple_of(8)
22722            || yg.len() < out_f
22723            || yu.len() < out_f
22724        {
22725            return Err(format!(
22726                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
22727                wg.len(),
22728                wu.len(),
22729                x.len()
22730            )
22731            .into());
22732        }
22733        let f = self.func("matvec_bf16_dual");
22734        let cfg = LaunchConfig {
22735            grid_dim: ((2 * out_f) as u32, 1, 1),
22736            block_dim: (mmv_block(), 1, 1),
22737            shared_mem_bytes: 0,
22738        };
22739        let (ini, outi) = (in_f as i32, out_f as i32);
22740        let __s_b = self.gpu.stream();
22741        let mut b = __s_b.launch_builder(&f);
22742        b.arg(wg)
22743            .arg(wu)
22744            .arg(x)
22745            .arg(yg)
22746            .arg(yu)
22747            .arg(&ini)
22748            .arg(&outi);
22749        unsafe {
22750            b.launch(cfg)?;
22751        }
22752        Ok(())
22753    }
22754
22755    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
22756    #[allow(clippy::too_many_arguments)]
22757    pub fn matvec_bf16_dual_into(
22758        &self,
22759        wg: &CudaSlice<u8>,
22760        wu: &CudaSlice<u8>,
22761        x: &CudaSlice<f32>,
22762        yg: &mut CudaSlice<f32>,
22763        yu: &mut CudaSlice<f32>,
22764        in_f: usize,
22765        out_f: usize,
22766    ) -> Result<(), Box<dyn std::error::Error>> {
22767        if wg.len() != in_f * out_f * 2
22768            || wu.len() != in_f * out_f * 2
22769            || x.len() < in_f
22770            || !in_f.is_multiple_of(8)
22771            || yg.len() < out_f
22772            || yu.len() < out_f
22773        {
22774            return Err(format!(
22775                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
22776                wg.len(),
22777                wu.len(),
22778                x.len()
22779            )
22780            .into());
22781        }
22782        let f = self.func("matvec_bf16_dual");
22783        let cfg = LaunchConfig {
22784            grid_dim: ((2 * out_f) as u32, 1, 1),
22785            block_dim: (mmv_block(), 1, 1),
22786            shared_mem_bytes: 0,
22787        };
22788        let (ini, outi) = (in_f as i32, out_f as i32);
22789        let __s_b = self.gpu.stream();
22790        let mut b = __s_b.launch_builder(&f);
22791        b.arg(wg)
22792            .arg(wu)
22793            .arg(x)
22794            .arg(yg)
22795            .arg(yu)
22796            .arg(&ini)
22797            .arg(&outi);
22798        unsafe {
22799            b.launch(cfg)?;
22800        }
22801        Ok(())
22802    }
22803
22804    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
22805    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
22806    #[allow(dead_code)] // allow: base form of the matvec_bf16_dual_* family; kept as the reference entry point
22807    pub(crate) fn matvec_bf16_dual(
22808        &self,
22809        wg: &CudaSlice<u8>,
22810        wu: &CudaSlice<u8>,
22811        x: &CudaSlice<f32>,
22812        in_f: usize,
22813        out_f: usize,
22814    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22815        if wg.len() != in_f * out_f * 2
22816            || wu.len() != in_f * out_f * 2
22817            || x.len() < in_f
22818            || !in_f.is_multiple_of(8)
22819        {
22820            return Err(format!(
22821                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
22822                wg.len(),
22823                wu.len(),
22824                x.len()
22825            )
22826            .into());
22827        }
22828        let mut yg = self.alloc_uninit::<f32>(out_f)?;
22829        let mut yu = self.alloc_uninit::<f32>(out_f)?;
22830        let f = self.func("matvec_bf16_dual");
22831        let cfg = LaunchConfig {
22832            grid_dim: ((2 * out_f) as u32, 1, 1),
22833            block_dim: (mmv_block(), 1, 1),
22834            shared_mem_bytes: 0,
22835        };
22836        let (ini, outi) = (in_f as i32, out_f as i32);
22837        let __s_b = self.gpu.stream();
22838        let mut b = __s_b.launch_builder(&f);
22839        b.arg(wg)
22840            .arg(wu)
22841            .arg(x)
22842            .arg(&mut yg)
22843            .arg(&mut yu)
22844            .arg(&ini)
22845            .arg(&outi);
22846        unsafe {
22847            b.launch(cfg)?;
22848        }
22849        Ok((yg, yu))
22850    }
22851
22852    #[allow(clippy::too_many_arguments)]
22853    #[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
22854    fn linear_bf16_chunked_inner(
22855        &self,
22856        x: &CudaSlice<f32>,
22857        data: &CudaSlice<u8>,
22858        m: usize,
22859        in_f: usize,
22860        out_f: usize,
22861        exact: bool,
22862        canonical_chunk_rows: Option<usize>,
22863    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22864        const CHUNK_BYTES: usize = 256 << 20;
22865        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
22866        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
22867        if m == 1
22868            && !exact
22869            && canonical_chunk_rows.is_none()
22870            && in_f.is_multiple_of(8)
22871            && Self::bf16_mmv_on()
22872        {
22873            return self.matvec_bf16(data, x, in_f, out_f);
22874        }
22875        // MEMRA_PP_BF16: prefill on the RESIDENT bf16 bytes through cuBLASLt tensor cores.
22876        // Below this door the whole weight is dequanted to f32 and multiplied without tensor
22877        // cores — the step37 prime's 14x gap to vLLM. `exact` and canonical-chunk callers are
22878        // numerical programs with their own equality gates and are left alone.
22879        if m >= 16
22880            && !exact
22881            && canonical_chunk_rows.is_none()
22882            && data.len() == in_f * out_f * 2
22883            && crate::f16_ffi::pp_bf16_enabled()
22884        {
22885            // None = cuBLASLt declined this shape (it announced which one); fall through to the
22886            // f32 dequant GEMM below, which is always correct.
22887            if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
22888                return Ok(y);
22889            }
22890        }
22891        let row_bytes = in_f
22892            .checked_mul(std::mem::size_of::<f32>())
22893            .ok_or("BF16 chunk row byte count overflow")?;
22894        if row_bytes == 0 || out_f == 0 {
22895            return Err("BF16 chunk dimensions must be nonzero".into());
22896        }
22897        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
22898        let chunk_rows = match canonical_chunk_rows {
22899            Some(0) => {
22900                return Err("canonical BF16 chunk rows must be nonzero".into());
22901            }
22902            Some(rows) if rows > max_chunk_rows => {
22903                return Err(format!(
22904                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
22905                )
22906                .into());
22907            }
22908            Some(rows) if out_f % rows != 0 => {
22909                return Err(format!(
22910                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
22911                )
22912                .into());
22913            }
22914            Some(rows) => rows,
22915            None => max_chunk_rows,
22916        };
22917        if chunk_rows >= out_f {
22918            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
22919            return if exact {
22920                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
22921            } else {
22922                self.linear(x, &wf32, m, in_f, out_f)
22923            };
22924        }
22925        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22926        let mut r0 = 0usize;
22927        while r0 < out_f {
22928            let rows = chunk_rows.min(out_f - r0);
22929            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
22930            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
22931            let yc = if exact {
22932                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
22933            } else {
22934                self.linear(x, &wf32, m, in_f, rows)?
22935            };
22936            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
22937            for mi in 0..m {
22938                let src = yc.slice(mi * rows..(mi + 1) * rows);
22939                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
22940                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
22941            }
22942            r0 += rows;
22943        }
22944        Ok(y)
22945    }
22946
22947    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
22948    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
22949    /// chunked BF16 numerical program instead of re-encoding the weight.
22950    pub fn linear_bf16_resident(
22951        &self,
22952        x: &CudaSlice<f32>,
22953        data: &CudaSlice<u8>,
22954        m: usize,
22955        in_f: usize,
22956        out_f: usize,
22957    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22958        if data.len() != in_f * out_f * 2 {
22959            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
22960        }
22961        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
22962    }
22963
22964    /// Execute a resident BF16 projection as fixed-width output-row chunks.
22965    ///
22966    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
22967    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
22968    /// model topology rather than the active rank count.
22969    pub fn linear_bf16_resident_canonical_rows(
22970        &self,
22971        x: &CudaSlice<f32>,
22972        data: &CudaSlice<u8>,
22973        m: usize,
22974        in_f: usize,
22975        out_f: usize,
22976        canonical_chunk_rows: usize,
22977    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22978        if data.len() != in_f * out_f * 2 {
22979            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
22980        }
22981        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
22982    }
22983
22984    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
22985    ///
22986    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
22987    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
22988    pub fn linear_f32_resident_canonical_rows(
22989        &self,
22990        x: &CudaSlice<f32>,
22991        data: &CudaSlice<f32>,
22992        m: usize,
22993        in_f: usize,
22994        out_f: usize,
22995        canonical_chunk_rows: usize,
22996    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22997        self.linear_f32_resident_canonical_rows_inner(
22998            x,
22999            data,
23000            m,
23001            in_f,
23002            out_f,
23003            canonical_chunk_rows,
23004            false,
23005        )
23006    }
23007
23008    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
23009    ///
23010    /// The projection shapes and values are identical to
23011    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
23012    /// changes, replacing one device copy per token with one placement kernel per output chunk.
23013    pub fn linear_f32_resident_canonical_rows_strided(
23014        &self,
23015        x: &CudaSlice<f32>,
23016        data: &CudaSlice<f32>,
23017        m: usize,
23018        in_f: usize,
23019        out_f: usize,
23020        canonical_chunk_rows: usize,
23021    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23022        self.linear_f32_resident_canonical_rows_inner(
23023            x,
23024            data,
23025            m,
23026            in_f,
23027            out_f,
23028            canonical_chunk_rows,
23029            true,
23030        )
23031    }
23032
23033    #[allow(clippy::too_many_arguments)]
23034    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23035    #[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
23036    fn linear_f32_resident_canonical_rows_inner(
23037        &self,
23038        x: &CudaSlice<f32>,
23039        data: &CudaSlice<f32>,
23040        m: usize,
23041        in_f: usize,
23042        out_f: usize,
23043        canonical_chunk_rows: usize,
23044        strided_output: bool,
23045    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23046        if data.len() != in_f * out_f {
23047            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
23048        }
23049        if canonical_chunk_rows == 0
23050            || canonical_chunk_rows > out_f
23051            || out_f % canonical_chunk_rows != 0
23052        {
23053            return Err(format!(
23054                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
23055            )
23056            .into());
23057        }
23058        if canonical_chunk_rows == out_f {
23059            return self.linear(x, data, m, in_f, out_f);
23060        }
23061
23062        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
23063        let input = x.slice(0..x.len());
23064        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
23065            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
23066            if m == 1 {
23067                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
23068                self.linear_device_into(
23069                    &input,
23070                    &weights,
23071                    &mut destination,
23072                    1,
23073                    in_f,
23074                    canonical_chunk_rows,
23075                )?;
23076                continue;
23077            }
23078            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
23079            if strided_output {
23080                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
23081            } else {
23082                for token in 0..m {
23083                    let source = chunk
23084                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
23085                    let mut destination =
23086                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
23087                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
23088                }
23089            }
23090        }
23091        Ok(y)
23092    }
23093
23094    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
23095    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
23096    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
23097    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
23098    #[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
23099    pub fn linear_f32_resident_canonical_rows_t1_into(
23100        &self,
23101        x: &CudaSlice<f32>,
23102        data: &CudaSlice<f32>,
23103        y: &mut CudaSlice<f32>,
23104        in_f: usize,
23105        out_f: usize,
23106        canonical_chunk_rows: usize,
23107    ) -> Result<(), Box<dyn std::error::Error>> {
23108        if data.len() != in_f * out_f {
23109            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
23110        }
23111        if y.len() != out_f || x.len() != in_f {
23112            return Err(format!(
23113                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
23114                x.len(),
23115                y.len()
23116            )
23117            .into());
23118        }
23119        if canonical_chunk_rows == 0
23120            || canonical_chunk_rows > out_f
23121            || out_f % canonical_chunk_rows != 0
23122        {
23123            return Err(format!(
23124                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
23125            )
23126            .into());
23127        }
23128        let input = x.slice(0..x.len());
23129        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
23130            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
23131            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
23132            self.linear_device_into(
23133                &input,
23134                &weights,
23135                &mut destination,
23136                1,
23137                in_f,
23138                canonical_chunk_rows,
23139            )?;
23140        }
23141        Ok(())
23142    }
23143
23144    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
23145    /// without the allocation, for workspace-resident operands.
23146    pub fn linear_t1_into(
23147        &self,
23148        x: &cudarc::driver::CudaView<'_, f32>,
23149        w: &cudarc::driver::CudaView<'_, f32>,
23150        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
23151        in_f: usize,
23152        out_f: usize,
23153    ) -> Result<(), Box<dyn std::error::Error>> {
23154        self.linear_device_into(x, w, y, 1, in_f, out_f)
23155    }
23156
23157    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
23158    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
23159    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
23160    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
23161    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
23162    /// router/shexp sites and matmul_decode_exact's Float arm.
23163    pub fn linear_decode_exact(
23164        &self,
23165        x: &CudaSlice<f32>,
23166        w: &CudaSlice<f32>,
23167        m_tokens: usize,
23168        in_f: usize,
23169        out_f: usize,
23170    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23171        if m_tokens == 1 {
23172            return self.linear(x, w, 1, in_f, out_f);
23173        }
23174        let xv = self.view(x, m_tokens * in_f);
23175        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
23176        for t in 0..m_tokens {
23177            let row = xv.slice(t * in_f..(t + 1) * in_f);
23178            let mut xr = self.alloc_uninit::<f32>(in_f)?;
23179            self.copy_view_into(&mut xr, 0, &row, in_f)?;
23180            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
23181            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
23182        }
23183        Ok(y)
23184    }
23185
23186    pub fn linear(
23187        &self,
23188        x: &CudaSlice<f32>,
23189        w: &CudaSlice<f32>,
23190        m_tokens: usize,
23191        in_f: usize,
23192        out_f: usize,
23193    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23194        self.linear_device(x, w, m_tokens, in_f, out_f)
23195    }
23196
23197    fn linear_device<I>(
23198        &self,
23199        x: &I,
23200        w: &I,
23201        m_tokens: usize,
23202        in_f: usize,
23203        out_f: usize,
23204    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
23205    where
23206        I: cudarc::driver::DevicePtr<f32>,
23207    {
23208        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
23209        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
23210        Ok(c)
23211    }
23212
23213    fn linear_device_into<I, O>(
23214        &self,
23215        x: &I,
23216        w: &I,
23217        c: &mut O,
23218        m_tokens: usize,
23219        in_f: usize,
23220        out_f: usize,
23221    ) -> Result<(), Box<dyn std::error::Error>>
23222    where
23223        I: cudarc::driver::DevicePtr<f32>,
23224        O: cudarc::driver::DevicePtrMut<f32>,
23225    {
23226        use cudarc::cublaslt::{Matmul, MatmulConfig};
23227        let cfg = MatmulConfig {
23228            transa: true,
23229            transb: false,
23230            transc: false,
23231            m: out_f as u64,
23232            n: m_tokens as u64,
23233            k: in_f as u64,
23234            alpha: 1.0,
23235            lda: in_f as i64,
23236            ldb: in_f as i64,
23237            beta: 0.0,
23238            ldc: out_f as i64,
23239            stride_a: None,
23240            stride_b: None,
23241            stride_c: None,
23242            stride_bias: None,
23243            batch_size: None,
23244        };
23245        let blas = self.gpu.blas();
23246        unsafe {
23247            blas.matmul(cfg, w, x, c, None, None)?;
23248        }
23249        Ok(())
23250    }
23251
23252    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
23253    ///
23254    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
23255    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
23256    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
23257    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
23258    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
23259    /// launch error mid-request.
23260    #[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
23261    pub fn sdpa_naive(
23262        &self,
23263        q: &CudaSlice<f32>,
23264        k: &CudaSlice<f32>,
23265        v: &CudaSlice<f32>,
23266        o: &mut CudaSlice<f32>,
23267        head_dim: usize,
23268        n_head: usize,
23269        n_head_kv: usize,
23270        t: usize,
23271        t_kv: usize,
23272        scale: f32,
23273        causal: bool,
23274    ) -> Result<(), Box<dyn std::error::Error>> {
23275        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
23276            return self.sdpa_naive_gmem(
23277                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23278            );
23279        }
23280        let f = self.func("sdpa_naive_f32");
23281        let cfg = LaunchConfig {
23282            grid_dim: (n_head as u32, t as u32, 1),
23283            block_dim: (128, 1, 1),
23284            shared_mem_bytes: (t_kv * 4) as u32,
23285        };
23286        let (hd, nh, nhkv, ti, tkvi, cz) = (
23287            head_dim as i32,
23288            n_head as i32,
23289            n_head_kv as i32,
23290            t as i32,
23291            t_kv as i32,
23292            causal as i32,
23293        );
23294        let __s_b = self.gpu.stream();
23295        let mut b = __s_b.launch_builder(&f);
23296        b.arg(q)
23297            .arg(k)
23298            .arg(v)
23299            .arg(o)
23300            .arg(&hd)
23301            .arg(&nh)
23302            .arg(&nhkv)
23303            .arg(&ti)
23304            .arg(&tkvi)
23305            .arg(&scale)
23306            .arg(&cz);
23307        unsafe {
23308            b.launch(cfg)?;
23309        }
23310        Ok(())
23311    }
23312
23313    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
23314    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
23315    /// of dynamic shared memory: identical loop structure and reduction order, so the output
23316    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
23317    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
23318    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
23319    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
23320    /// T==T_kv caller cannot silently allocate tens of GB.
23321    #[allow(clippy::too_many_arguments)]
23322    pub fn sdpa_naive_gmem(
23323        &self,
23324        q: &CudaSlice<f32>,
23325        k: &CudaSlice<f32>,
23326        v: &CudaSlice<f32>,
23327        o: &mut CudaSlice<f32>,
23328        head_dim: usize,
23329        n_head: usize,
23330        n_head_kv: usize,
23331        t: usize,
23332        t_kv: usize,
23333        scale: f32,
23334        causal: bool,
23335    ) -> Result<(), Box<dyn std::error::Error>> {
23336        let ws_len = n_head
23337            .checked_mul(t)
23338            .and_then(|x| x.checked_mul(t_kv))
23339            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
23340        let ws_bytes = ws_len
23341            .checked_mul(std::mem::size_of::<f32>())
23342            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
23343        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
23344            return Err(format!(
23345                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
23346                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
23347                 needs a tiled/flash kernel, not the naive oracle"
23348            )
23349            .into());
23350        }
23351        let mut scores = self.uninit(ws_len)?;
23352        let f = self.func("sdpa_naive_gmem_f32");
23353        let cfg = LaunchConfig {
23354            grid_dim: (n_head as u32, t as u32, 1),
23355            block_dim: (128, 1, 1),
23356            shared_mem_bytes: 0,
23357        };
23358        let (hd, nh, nhkv, ti, tkvi, cz) = (
23359            head_dim as i32,
23360            n_head as i32,
23361            n_head_kv as i32,
23362            t as i32,
23363            t_kv as i32,
23364            causal as i32,
23365        );
23366        let __s_b = self.gpu.stream();
23367        let mut b = __s_b.launch_builder(&f);
23368        b.arg(q)
23369            .arg(k)
23370            .arg(v)
23371            .arg(o)
23372            .arg(&mut scores)
23373            .arg(&hd)
23374            .arg(&nh)
23375            .arg(&nhkv)
23376            .arg(&ti)
23377            .arg(&tkvi)
23378            .arg(&scale)
23379            .arg(&cz);
23380        unsafe {
23381            b.launch(cfg)?;
23382        }
23383        Ok(())
23384    }
23385
23386    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
23387    /// bidirectional image islands. `span_id` labels each absolute kv position
23388    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
23389    /// reproducing the reference's non-causal image batch. window 0 = no window.
23390    #[allow(clippy::too_many_arguments)]
23391    pub fn sdpa_naive_island(
23392        &self,
23393        q: &CudaSlice<f32>,
23394        k: &CudaSlice<f32>,
23395        v: &CudaSlice<f32>,
23396        o: &mut CudaSlice<f32>,
23397        span_id: &CudaSlice<i32>,
23398        head_dim: usize,
23399        n_head: usize,
23400        n_head_kv: usize,
23401        t: usize,
23402        t_kv: usize,
23403        scale: f32,
23404        window: usize,
23405    ) -> Result<(), Box<dyn std::error::Error>> {
23406        let f = self.func("sdpa_naive_island_f32");
23407        let cfg = LaunchConfig {
23408            grid_dim: (n_head as u32, t as u32, 1),
23409            block_dim: (128, 1, 1),
23410            shared_mem_bytes: (t_kv * 4) as u32,
23411        };
23412        let (hd, nh, nhkv, ti, tkvi, wi) = (
23413            head_dim as i32,
23414            n_head as i32,
23415            n_head_kv as i32,
23416            t as i32,
23417            t_kv as i32,
23418            window as i32,
23419        );
23420        let __s_b = self.gpu.stream();
23421        let mut b = __s_b.launch_builder(&f);
23422        b.arg(q)
23423            .arg(k)
23424            .arg(v)
23425            .arg(o)
23426            .arg(span_id)
23427            .arg(&hd)
23428            .arg(&nh)
23429            .arg(&nhkv)
23430            .arg(&ti)
23431            .arg(&tkvi)
23432            .arg(&scale)
23433            .arg(&wi);
23434        unsafe {
23435            b.launch(cfg)?;
23436        }
23437        Ok(())
23438    }
23439
23440    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
23441    #[allow(clippy::too_many_arguments)]
23442    pub fn sdpa_naive_w(
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        window: usize,
23456    ) -> Result<(), Box<dyn std::error::Error>> {
23457        let f = self.func("sdpa_naive_w_f32");
23458        let cfg = LaunchConfig {
23459            grid_dim: (n_head as u32, t as u32, 1),
23460            block_dim: (128, 1, 1),
23461            shared_mem_bytes: (t_kv * 4) as u32,
23462        };
23463        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
23464            head_dim as i32,
23465            n_head as i32,
23466            n_head_kv as i32,
23467            t as i32,
23468            t_kv as i32,
23469            causal as i32,
23470            window as i32,
23471        );
23472        let __s_b = self.gpu.stream();
23473        let mut b = __s_b.launch_builder(&f);
23474        b.arg(q)
23475            .arg(k)
23476            .arg(v)
23477            .arg(o)
23478            .arg(&hd)
23479            .arg(&nh)
23480            .arg(&nhkv)
23481            .arg(&ti)
23482            .arg(&tkvi)
23483            .arg(&scale)
23484            .arg(&cz)
23485            .arg(&wi);
23486        unsafe {
23487            b.launch(cfg)?;
23488        }
23489        Ok(())
23490    }
23491
23492    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
23493    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
23494    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
23495    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
23496    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
23497    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
23498    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
23499    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
23500    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
23501    #[allow(clippy::too_many_arguments)]
23502    pub fn sdpa_naive_w_lo(
23503        &self,
23504        q: &CudaSlice<f32>,
23505        k: &CudaSlice<f32>,
23506        v: &CudaSlice<f32>,
23507        o: &mut CudaSlice<f32>,
23508        head_dim: usize,
23509        n_head: usize,
23510        n_head_kv: usize,
23511        t: usize,
23512        t_kv: usize,
23513        scale: f32,
23514        causal: bool,
23515        window: usize,
23516    ) -> Result<(), Box<dyn std::error::Error>> {
23517        let kv_lo = if window > 0 {
23518            (t_kv - t + 1).saturating_sub(window)
23519        } else {
23520            0
23521        };
23522        let smem = (t_kv - kv_lo) * 4;
23523        if smem > 48 * 1024 {
23524            return Err(format!(
23525                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
23526                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
23527                 a window this wide needs the multi-pass long-ctx kernel"
23528            )
23529            .into());
23530        }
23531        let f = self.func("sdpa_naive_w_lo_f32");
23532        let cfg = LaunchConfig {
23533            grid_dim: (n_head as u32, t as u32, 1),
23534            block_dim: (128, 1, 1),
23535            shared_mem_bytes: smem as u32,
23536        };
23537        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
23538            head_dim as i32,
23539            n_head as i32,
23540            n_head_kv as i32,
23541            t as i32,
23542            t_kv as i32,
23543            causal as i32,
23544            window as i32,
23545            kv_lo as i32,
23546        );
23547        let __s_b = self.gpu.stream();
23548        let mut b = __s_b.launch_builder(&f);
23549        b.arg(q)
23550            .arg(k)
23551            .arg(v)
23552            .arg(o)
23553            .arg(&hd)
23554            .arg(&nh)
23555            .arg(&nhkv)
23556            .arg(&ti)
23557            .arg(&tkvi)
23558            .arg(&scale)
23559            .arg(&cz)
23560            .arg(&wi)
23561            .arg(&lo);
23562        unsafe {
23563            b.launch(cfg)?;
23564        }
23565        Ok(())
23566    }
23567
23568    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
23569    #[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
23570    pub fn sdpa_naive_view(
23571        &self,
23572        q: &CudaSlice<f32>,
23573        k: &cudarc::driver::CudaView<f32>,
23574        v: &cudarc::driver::CudaView<f32>,
23575        o: &mut CudaSlice<f32>,
23576        head_dim: usize,
23577        n_head: usize,
23578        n_head_kv: usize,
23579        t: usize,
23580        t_kv: usize,
23581        scale: f32,
23582        causal: bool,
23583    ) -> Result<(), Box<dyn std::error::Error>> {
23584        let f = self.func("sdpa_naive_f32");
23585        let cfg = LaunchConfig {
23586            grid_dim: (n_head as u32, t as u32, 1),
23587            block_dim: (128, 1, 1),
23588            shared_mem_bytes: (t_kv * 4) as u32,
23589        };
23590        let (hd, nh, nhkv, ti, tkvi, cz) = (
23591            head_dim as i32,
23592            n_head as i32,
23593            n_head_kv as i32,
23594            t as i32,
23595            t_kv as i32,
23596            causal as i32,
23597        );
23598        let __s_b = self.gpu.stream();
23599        let mut b = __s_b.launch_builder(&f);
23600        b.arg(q)
23601            .arg(k)
23602            .arg(v)
23603            .arg(o)
23604            .arg(&hd)
23605            .arg(&nh)
23606            .arg(&nhkv)
23607            .arg(&ti)
23608            .arg(&tkvi)
23609            .arg(&scale)
23610            .arg(&cz);
23611        unsafe {
23612            b.launch(cfg)?;
23613        }
23614        Ok(())
23615    }
23616
23617    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
23618    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
23619    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
23620    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
23621    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
23622    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
23623    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
23624    #[allow(clippy::too_many_arguments)]
23625    pub fn fa_dequant_kv_view_f32(
23626        &self,
23627        k: &cudarc::driver::CudaView<u8>,
23628        v: &cudarc::driver::CudaView<u8>,
23629        kf: &mut CudaSlice<f32>,
23630        vf: &mut CudaSlice<f32>,
23631        kv_dim_k: usize,
23632        kv_dim_v: usize,
23633        t_kv: usize,
23634        k_tok_bytes: usize,
23635        v_tok_bytes: usize,
23636        g: bool,
23637    ) -> Result<(), Box<dyn std::error::Error>> {
23638        let f = if g {
23639            self.func_g("fa_dequant_kv_ws_f32")
23640        } else {
23641            self.func("fa_dequant_kv_ws_f32")
23642        };
23643        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
23644        #[allow(clippy::manual_div_ceil)]
23645        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23646        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23647        let cfg = LaunchConfig {
23648            grid_dim: (nblk.max(1), 1, 1),
23649            block_dim: (256, 1, 1),
23650            shared_mem_bytes: 0,
23651        };
23652        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
23653        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23654        let __s_b = self.gpu.stream();
23655        let mut b = __s_b.launch_builder(&f);
23656        b.arg(k)
23657            .arg(v)
23658            .arg(&mut *kf)
23659            .arg(&mut *vf)
23660            .arg(&kdk)
23661            .arg(&kdv)
23662            .arg(&tkvi)
23663            .arg(&ktb)
23664            .arg(&vtb);
23665        unsafe {
23666            b.launch(cfg)?;
23667        }
23668        Ok(())
23669    }
23670
23671    #[allow(clippy::too_many_arguments)]
23672    pub fn sdpa_naive_quantized_view(
23673        &self,
23674        q: &CudaSlice<f32>,
23675        k: &cudarc::driver::CudaView<u8>,
23676        v: &cudarc::driver::CudaView<u8>,
23677        o: &mut CudaSlice<f32>,
23678        head_dim: usize,
23679        n_head: usize,
23680        n_head_kv: usize,
23681        t: usize,
23682        t_kv: usize,
23683        scale: f32,
23684        causal: bool,
23685        k_tok_bytes: usize,
23686        v_tok_bytes: usize,
23687    ) -> Result<(), Box<dyn std::error::Error>> {
23688        let kv_dim = n_head_kv * head_dim;
23689        let mut kf = self.uninit(t_kv * kv_dim)?;
23690        let mut vf = self.uninit(t_kv * kv_dim)?;
23691        let f = self.func("fa_dequant_kv_ws_f32");
23692        let total = (2 * t_kv * kv_dim) as u64;
23693        #[allow(clippy::manual_div_ceil)]
23694        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23695        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23696        let cfg = LaunchConfig {
23697            grid_dim: (nblk.max(1), 1, 1),
23698            block_dim: (256, 1, 1),
23699            shared_mem_bytes: 0,
23700        };
23701        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
23702        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
23703        let __s_b = self.gpu.stream();
23704        let mut b = __s_b.launch_builder(&f);
23705        b.arg(k)
23706            .arg(v)
23707            .arg(&mut kf)
23708            .arg(&mut vf)
23709            .arg(&kv_dim_i)
23710            .arg(&kv_dim_i)
23711            .arg(&t_kv_i)
23712            .arg(&k_tok_bytes_i)
23713            .arg(&v_tok_bytes_i);
23714        unsafe { b.launch(cfg)? };
23715        self.sdpa_naive(
23716            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23717        )
23718    }
23719
23720    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
23721    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
23722    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
23723    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
23724    /// unwindowed function above and produces bit-identical output at window == 0.
23725    ///
23726    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
23727    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
23728    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
23729    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
23730    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
23731    #[allow(clippy::too_many_arguments)]
23732    pub fn sdpa_naive_w_quantized_view(
23733        &self,
23734        q: &CudaSlice<f32>,
23735        k: &cudarc::driver::CudaView<u8>,
23736        v: &cudarc::driver::CudaView<u8>,
23737        o: &mut CudaSlice<f32>,
23738        head_dim: usize,
23739        n_head: usize,
23740        n_head_kv: usize,
23741        t: usize,
23742        t_kv: usize,
23743        scale: f32,
23744        causal: bool,
23745        window: usize,
23746        k_tok_bytes: usize,
23747        v_tok_bytes: usize,
23748    ) -> Result<(), Box<dyn std::error::Error>> {
23749        let kv_dim = n_head_kv * head_dim;
23750        let mut kf = self.uninit(t_kv * kv_dim)?;
23751        let mut vf = self.uninit(t_kv * kv_dim)?;
23752        let f = self.func("fa_dequant_kv_ws_f32");
23753        let total = (2 * t_kv * kv_dim) as u64;
23754        #[allow(clippy::manual_div_ceil)]
23755        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23756        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
23757        let cfg = LaunchConfig {
23758            grid_dim: (nblk.max(1), 1, 1),
23759            block_dim: (256, 1, 1),
23760            shared_mem_bytes: 0,
23761        };
23762        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
23763        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
23764        let __s_b = self.gpu.stream();
23765        let mut b = __s_b.launch_builder(&f);
23766        b.arg(k)
23767            .arg(v)
23768            .arg(&mut kf)
23769            .arg(&mut vf)
23770            .arg(&kv_dim_i)
23771            .arg(&kv_dim_i)
23772            .arg(&t_kv_i)
23773            .arg(&k_tok_bytes_i)
23774            .arg(&v_tok_bytes_i);
23775        unsafe { b.launch(cfg)? };
23776        self.sdpa_naive_w(
23777            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
23778        )
23779    }
23780
23781    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
23782    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
23783    /// Q/K/V/O [head_dim, n_head(_kv), T].
23784    #[allow(clippy::too_many_arguments)]
23785    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23786    #[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
23787    pub fn fa_prefill(
23788        &self,
23789        q: &CudaSlice<f32>,
23790        k: &CudaSlice<f32>,
23791        v: &CudaSlice<f32>,
23792        o: &mut CudaSlice<f32>,
23793        head_dim: usize,
23794        n_head: usize,
23795        n_head_kv: usize,
23796        t: usize,
23797        t_kv: usize,
23798        scale: f32,
23799        causal: bool,
23800    ) -> Result<(), Box<dyn std::error::Error>> {
23801        if portable_mma_gated() {
23802            return self.sdpa_naive(
23803                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
23804            );
23805        }
23806        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
23807        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
23808        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
23809        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
23810        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
23811        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
23812        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
23813        let fa3_on = head_dim == 256
23814            && causal
23815            && t == t_kv
23816            && match std::env::var("MEMRA_FA3").as_deref() {
23817                Ok("0") => false,
23818                // The force arm consults the arch now: the bf16 stage below calls
23819                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
23820                // a portable build. Refuse at the switch, not at the lookup.
23821                Ok("1") => {
23822                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
23823                    true
23824                }
23825                _ => cfg!(memra_hopper_mma),
23826            };
23827        if fa3_on {
23828            let n = t * n_head * head_dim;
23829            let nkv = t * n_head_kv * head_dim;
23830            let mut q16 = self.alloc_u8_uninit(n * 2)?;
23831            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
23832            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
23833            self.f32_to_bf16_into(q, &mut q16, n)?;
23834            self.f32_to_bf16_into(k, &mut k16, nkv)?;
23835            self.f32_to_bf16_into(v, &mut v16, nkv)?;
23836            let rc = {
23837                use cudarc::driver::{DevicePtr, DevicePtrMut};
23838                let stream = self.gpu.stream();
23839                let (qp, _g1) = q16.device_ptr(&stream);
23840                let (kp, _g2) = k16.device_ptr(&stream);
23841                let (vp, _g3) = v16.device_ptr(&stream);
23842                let (op, _g4) = o.device_ptr_mut(&stream);
23843                unsafe {
23844                    memra_fa3_prefill(
23845                        qp as *const core::ffi::c_void,
23846                        kp as *const core::ffi::c_void,
23847                        vp as *const core::ffi::c_void,
23848                        op as *mut f32,
23849                        t as i32,
23850                        n_head as i32,
23851                        n_head_kv as i32,
23852                        head_dim as i32,
23853                        scale,
23854                        stream.cu_stream() as *mut core::ffi::c_void,
23855                    )
23856                }
23857            };
23858            if rc != 0 {
23859                return Err(format!("memra_fa3_prefill rc={rc}").into());
23860            }
23861            return Ok(());
23862        }
23863        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
23864        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
23865        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
23866        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
23867        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23868        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
23869        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
23870            const BLOCK_Q: usize = 64;
23871            const BKX: usize = 32;
23872            let f = self.func("fa_prefill_bf16_p1");
23873            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
23874                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
23875            use cudarc::driver::sys::CUfunction_attribute_enum as A;
23876            f.set_attribute(
23877                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23878                shmem as i32,
23879            )?;
23880            let cfg = LaunchConfig {
23881                grid_dim: (
23882                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
23883                    n_head as u32,
23884                    1,
23885                ),
23886                block_dim: (32, 4, 1),
23887                shared_mem_bytes: shmem,
23888            };
23889            let (hd, nh, nhkv, ti, tkvi, cz) = (
23890                head_dim as i32,
23891                n_head as i32,
23892                n_head_kv as i32,
23893                t as i32,
23894                t_kv as i32,
23895                causal as i32,
23896            );
23897            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
23898            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
23899            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
23900            let __s_b = self.gpu.stream();
23901            let mut b = __s_b.launch_builder(&f);
23902            b.arg(&qb)
23903                .arg(&kb)
23904                .arg(&vb)
23905                .arg(o)
23906                .arg(&hd)
23907                .arg(&nh)
23908                .arg(&nhkv)
23909                .arg(&ti)
23910                .arg(&tkvi)
23911                .arg(&scale)
23912                .arg(&cz);
23913            unsafe {
23914                b.launch(cfg)?;
23915            }
23916            return Ok(());
23917        }
23918        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
23919        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
23920        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
23921        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
23922        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
23923        const BK: usize = 32;
23924        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
23925        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
23926        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
23927        let (block_q, warps, w2_sfx): (usize, u32, &str) =
23928            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
23929        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
23930        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
23931        // other head_dims to sdpa_naive before reaching here.
23932        let hd_sfx = fa_hd_suffix(head_dim)?;
23933        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
23934        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
23935        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
23936        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
23937        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
23938        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
23939        let (kb16, vb16) = if bf16kv {
23940            let n = t_kv * n_head_kv * head_dim;
23941            let mut kb = self.alloc_u8_uninit(n * 2)?;
23942            let mut vb = self.alloc_u8_uninit(n * 2)?;
23943            let fcv = self.func("f32_to_bf16_bulk");
23944            let ni = n as i64;
23945            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
23946            let __s_b = self.gpu.stream();
23947            let mut b = __s_b.launch_builder(&fcv);
23948            b.arg(k).arg(&mut kb).arg(&ni);
23949            unsafe {
23950                b.launch(cfgc)?;
23951            }
23952            let __s_b = self.gpu.stream();
23953            let mut b = __s_b.launch_builder(&fcv);
23954            b.arg(v).arg(&mut vb).arg(&ni);
23955            unsafe {
23956                b.launch(cfgc)?;
23957            }
23958            (Some(kb), Some(vb))
23959        } else {
23960            (None, None)
23961        };
23962        let f = self.func(&if bf16kv {
23963            format!("fa_prefill_bf16kv_pp{hd_sfx}")
23964        } else {
23965            format!(
23966                "fa_prefill_f32{}{}{hd_sfx}",
23967                if floor { "" } else { "_pp" },
23968                if floor { "" } else { w2_sfx }
23969            )
23970        });
23971        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
23972        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
23973        let kv_stages = if bf16kv { 2 } else { 1 };
23974        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
23975            + 4 * (block_q * BK + 2 * block_q)) as u32;
23976        use cudarc::driver::sys::CUfunction_attribute_enum as A;
23977        f.set_attribute(
23978            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23979            shmem as i32,
23980        )?;
23981        let cfg = LaunchConfig {
23982            grid_dim: (
23983                (t as u32 + block_q as u32 - 1) / block_q as u32,
23984                n_head as u32,
23985                1,
23986            ),
23987            block_dim: (32, warps, 1),
23988            shared_mem_bytes: shmem,
23989        };
23990        let (hd, nh, nhkv, ti, tkvi, cz) = (
23991            head_dim as i32,
23992            n_head as i32,
23993            n_head_kv as i32,
23994            t as i32,
23995            t_kv as i32,
23996            causal as i32,
23997        );
23998        let __s_b = self.gpu.stream();
23999        let mut b = __s_b.launch_builder(&f);
24000        b.arg(q);
24001        match (&kb16, &vb16) {
24002            (Some(kb), Some(vb)) => {
24003                b.arg(kb).arg(vb);
24004            }
24005            _ => {
24006                b.arg(k).arg(v);
24007            }
24008        }
24009        b.arg(o)
24010            .arg(&hd)
24011            .arg(&nh)
24012            .arg(&nhkv)
24013            .arg(&ti)
24014            .arg(&tkvi)
24015            .arg(&scale)
24016            .arg(&cz);
24017        unsafe {
24018            b.launch(cfg)?;
24019        }
24020        Ok(())
24021    }
24022
24023    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
24024    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
24025    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
24026    #[allow(clippy::too_many_arguments)]
24027    pub fn fa_prefill_w(
24028        &self,
24029        q: &CudaSlice<f32>,
24030        k: &CudaSlice<f32>,
24031        v: &CudaSlice<f32>,
24032        o: &mut CudaSlice<f32>,
24033        head_dim: usize,
24034        n_head: usize,
24035        n_head_kv: usize,
24036        t: usize,
24037        t_kv: usize,
24038        scale: f32,
24039        causal: bool,
24040        window: usize,
24041    ) -> Result<(), Box<dyn std::error::Error>> {
24042        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
24043        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
24044        if portable_mma_gated() {
24045            return self.sdpa_naive_w(
24046                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
24047            );
24048        }
24049        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
24050        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
24051        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
24052        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24053        let faw_f32 =
24054            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
24055        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
24056        self.fa_prefill_w_arm(
24057            q,
24058            k,
24059            v,
24060            o,
24061            head_dim,
24062            n_head,
24063            n_head_kv,
24064            t,
24065            t_kv,
24066            scale,
24067            causal,
24068            window,
24069            floor || faw_f32,
24070            floor,
24071        )
24072    }
24073
24074    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
24075    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
24076    #[allow(clippy::too_many_arguments)]
24077    #[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
24078    pub fn fa_prefill_w_pre(
24079        &self,
24080        qb: &CudaSlice<u8>,
24081        kb: &CudaSlice<u8>,
24082        vb: &CudaSlice<u8>,
24083        o: &mut CudaSlice<f32>,
24084        head_dim: usize,
24085        n_head: usize,
24086        n_head_kv: usize,
24087        t: usize,
24088        t_kv: usize,
24089        scale: f32,
24090        causal: bool,
24091        window: usize,
24092        v_f16: bool,
24093    ) -> Result<(), Box<dyn std::error::Error>> {
24094        const BLOCK_Q: usize = 64;
24095        const BK: usize = 32;
24096        debug_assert_eq!(head_dim, 256);
24097        let hp = fa_f16pv_on()
24098            && faw_hp_on()
24099            && n_head.is_multiple_of(2)
24100            && (n_head / n_head_kv).is_multiple_of(2);
24101        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
24102        if hp {
24103            const BLOCK_QH: usize = 32;
24104            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
24105            // else re-encode through the pooled scratch (stream-ordered reuse).
24106            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
24107            let vh: &CudaSlice<u8> = if v_f16 {
24108                vb
24109            } else {
24110                let n = t_kv * n_head_kv * head_dim;
24111                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
24112                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
24113                }
24114                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
24115                vguard.as_ref().unwrap()
24116            };
24117            let f = self.func("fa_prefill_w_bf16_p1h2");
24118            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
24119            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24120            f.set_attribute(
24121                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24122                shmem as i32,
24123            )?;
24124            let cfg = LaunchConfig {
24125                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
24126                block_dim: (32, 4, 1),
24127                shared_mem_bytes: shmem,
24128            };
24129            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24130                head_dim as i32,
24131                n_head as i32,
24132                n_head_kv as i32,
24133                t as i32,
24134                t_kv as i32,
24135                causal as i32,
24136                window as i32,
24137            );
24138            let __s_b = self.gpu.stream();
24139            let mut b = __s_b.launch_builder(&f);
24140            b.arg(qb)
24141                .arg(kb)
24142                .arg(vh)
24143                .arg(o)
24144                .arg(&hd)
24145                .arg(&nh)
24146                .arg(&nhkv)
24147                .arg(&ti)
24148                .arg(&tkvi)
24149                .arg(&scale)
24150                .arg(&cz)
24151                .arg(&wi);
24152            unsafe {
24153                b.launch(cfg)?;
24154            }
24155            return Ok(());
24156        }
24157        let f = self.func("fa_prefill_w_bf16_p1");
24158        let shmem =
24159            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24160        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24161        f.set_attribute(
24162            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24163            shmem as i32,
24164        )?;
24165        let cfg = LaunchConfig {
24166            grid_dim: (
24167                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24168                n_head as u32,
24169                1,
24170            ),
24171            block_dim: (32, 4, 1),
24172            shared_mem_bytes: shmem,
24173        };
24174        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24175            head_dim as i32,
24176            n_head as i32,
24177            n_head_kv as i32,
24178            t as i32,
24179            t_kv as i32,
24180            causal as i32,
24181            window as i32,
24182        );
24183        let __s_b = self.gpu.stream();
24184        let mut b = __s_b.launch_builder(&f);
24185        b.arg(qb)
24186            .arg(kb)
24187            .arg(vb)
24188            .arg(o)
24189            .arg(&hd)
24190            .arg(&nh)
24191            .arg(&nhkv)
24192            .arg(&ti)
24193            .arg(&tkvi)
24194            .arg(&scale)
24195            .arg(&cz)
24196            .arg(&wi);
24197        unsafe {
24198            b.launch(cfg)?;
24199        }
24200        Ok(())
24201    }
24202
24203    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
24204    #[allow(clippy::too_many_arguments)]
24205    #[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
24206    pub fn fa_prefill_w_arm(
24207        &self,
24208        q: &CudaSlice<f32>,
24209        k: &CudaSlice<f32>,
24210        v: &CudaSlice<f32>,
24211        o: &mut CudaSlice<f32>,
24212        head_dim: usize,
24213        n_head: usize,
24214        n_head_kv: usize,
24215        t: usize,
24216        t_kv: usize,
24217        scale: f32,
24218        causal: bool,
24219        window: usize,
24220        f32_stage: bool,
24221        floor: bool,
24222    ) -> Result<(), Box<dyn std::error::Error>> {
24223        const BLOCK_Q: usize = 64;
24224        const BK: usize = 32;
24225        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
24226        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
24227        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
24228        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
24229        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24230        let p1 = !floor
24231            && !f32_stage
24232            && *P1_ON.get_or_init(|| {
24233                std::env::var("MEMRA_FAW_P1")
24234                    .map(|v| v != "0")
24235                    .unwrap_or(true)
24236            });
24237        let hp = p1
24238            && fa_f16pv_on()
24239            && faw_hp_on()
24240            && n_head.is_multiple_of(2)
24241            && (n_head / n_head_kv).is_multiple_of(2);
24242        if hp {
24243            const BLOCK_QH: usize = 32;
24244            let f = self.func("fa_prefill_w_bf16_p1h2");
24245            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
24246            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24247            f.set_attribute(
24248                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24249                shmem as i32,
24250            )?;
24251            let cfg = LaunchConfig {
24252                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
24253                block_dim: (32, 4, 1),
24254                shared_mem_bytes: shmem,
24255            };
24256            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24257                head_dim as i32,
24258                n_head as i32,
24259                n_head_kv as i32,
24260                t as i32,
24261                t_kv as i32,
24262                causal as i32,
24263                window as i32,
24264            );
24265            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24266            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24267            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
24268            let __s_b = self.gpu.stream();
24269            let mut b = __s_b.launch_builder(&f);
24270            b.arg(&qb)
24271                .arg(&kb)
24272                .arg(&vh)
24273                .arg(o)
24274                .arg(&hd)
24275                .arg(&nh)
24276                .arg(&nhkv)
24277                .arg(&ti)
24278                .arg(&tkvi)
24279                .arg(&scale)
24280                .arg(&cz)
24281                .arg(&wi);
24282            unsafe {
24283                b.launch(cfg)?;
24284            }
24285            return Ok(());
24286        }
24287        if p1 {
24288            let f = self.func("fa_prefill_w_bf16_p1");
24289            let shmem =
24290                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24291            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24292            f.set_attribute(
24293                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24294                shmem as i32,
24295            )?;
24296            let cfg = LaunchConfig {
24297                grid_dim: (
24298                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24299                    n_head as u32,
24300                    1,
24301                ),
24302                block_dim: (32, 4, 1),
24303                shared_mem_bytes: shmem,
24304            };
24305            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24306                head_dim as i32,
24307                n_head as i32,
24308                n_head_kv as i32,
24309                t as i32,
24310                t_kv as i32,
24311                causal as i32,
24312                window as i32,
24313            );
24314            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24315            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24316            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24317            let __s_b = self.gpu.stream();
24318            let mut b = __s_b.launch_builder(&f);
24319            b.arg(&qb)
24320                .arg(&kb)
24321                .arg(&vb)
24322                .arg(o)
24323                .arg(&hd)
24324                .arg(&nh)
24325                .arg(&nhkv)
24326                .arg(&ti)
24327                .arg(&tkvi)
24328                .arg(&scale)
24329                .arg(&cz)
24330                .arg(&wi);
24331            unsafe {
24332                b.launch(cfg)?;
24333            }
24334            return Ok(());
24335        }
24336        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
24337        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
24338        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24339        let g4 = !floor
24340            && !f32_stage
24341            && n_head_kv == 1
24342            && n_head.is_multiple_of(4)
24343            && *G4_ON.get_or_init(|| {
24344                std::env::var("MEMRA_FAW_G4")
24345                    .map(|v| v != "0")
24346                    .unwrap_or(true)
24347            });
24348        if g4 {
24349            const SP_M: usize = 16;
24350            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
24351            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
24352            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24353            let o2 = *O2_ON.get_or_init(|| {
24354                std::env::var("MEMRA_FAW_O2")
24355                    .map(|v| v != "0")
24356                    .unwrap_or(true)
24357            });
24358            let f = self.func(if o2 {
24359                "fa_prefill_w_bf16_g4o2"
24360            } else {
24361                "fa_prefill_w_bf16_g4"
24362            });
24363            let shmem = if o2 {
24364                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
24365            } else {
24366                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
24367                    as u32
24368            };
24369            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24370            f.set_attribute(
24371                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24372                shmem as i32,
24373            )?;
24374            let cfg = LaunchConfig {
24375                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
24376                block_dim: (32, 4, 1),
24377                shared_mem_bytes: shmem,
24378            };
24379            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24380                head_dim as i32,
24381                n_head as i32,
24382                n_head_kv as i32,
24383                t as i32,
24384                t_kv as i32,
24385                causal as i32,
24386                window as i32,
24387            );
24388            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24389            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24390            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24391            let __s_b = self.gpu.stream();
24392            let mut b = __s_b.launch_builder(&f);
24393            b.arg(&qb)
24394                .arg(&kb)
24395                .arg(&vb)
24396                .arg(o)
24397                .arg(&hd)
24398                .arg(&nh)
24399                .arg(&nhkv)
24400                .arg(&ti)
24401                .arg(&tkvi)
24402                .arg(&scale)
24403                .arg(&cz)
24404                .arg(&wi);
24405            unsafe {
24406                b.launch(cfg)?;
24407            }
24408            return Ok(());
24409        }
24410        let f = self.func(if floor {
24411            "fa_prefill_w_f32"
24412        } else if f32_stage {
24413            "fa_prefill_w_f32_pp"
24414        } else {
24415            "fa_prefill_w_bf16_pp"
24416        });
24417        let shmem =
24418            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
24419        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24420        f.set_attribute(
24421            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24422            shmem as i32,
24423        )?;
24424        let cfg = LaunchConfig {
24425            grid_dim: (
24426                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24427                n_head as u32,
24428                1,
24429            ),
24430            block_dim: (32, 4, 1),
24431            shared_mem_bytes: shmem,
24432        };
24433        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
24434            head_dim as i32,
24435            n_head as i32,
24436            n_head_kv as i32,
24437            t as i32,
24438            t_kv as i32,
24439            causal as i32,
24440            window as i32,
24441        );
24442        if f32_stage {
24443            let __s_b = self.gpu.stream();
24444            let mut b = __s_b.launch_builder(&f);
24445            b.arg(q)
24446                .arg(k)
24447                .arg(v)
24448                .arg(o)
24449                .arg(&hd)
24450                .arg(&nh)
24451                .arg(&nhkv)
24452                .arg(&ti)
24453                .arg(&tkvi)
24454                .arg(&scale)
24455                .arg(&cz)
24456                .arg(&wi);
24457            unsafe {
24458                b.launch(cfg)?;
24459            }
24460        } else {
24461            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24462            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24463            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24464            let __s_b = self.gpu.stream();
24465            let mut b = __s_b.launch_builder(&f);
24466            b.arg(&qb)
24467                .arg(&kb)
24468                .arg(&vb)
24469                .arg(o)
24470                .arg(&hd)
24471                .arg(&nh)
24472                .arg(&nhkv)
24473                .arg(&ti)
24474                .arg(&tkvi)
24475                .arg(&scale)
24476                .arg(&cz)
24477                .arg(&wi);
24478            unsafe {
24479                b.launch(cfg)?;
24480            }
24481        }
24482        Ok(())
24483    }
24484
24485    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
24486    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
24487    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
24488    #[allow(clippy::too_many_arguments)]
24489    pub fn fa_prefill_hd512(
24490        &self,
24491        q: &CudaSlice<f32>,
24492        k: &CudaSlice<f32>,
24493        v: &CudaSlice<f32>,
24494        o: &mut CudaSlice<f32>,
24495        head_dim: usize,
24496        n_head: usize,
24497        n_head_kv: usize,
24498        t: usize,
24499        t_kv: usize,
24500        scale: f32,
24501        causal: bool,
24502    ) -> Result<(), Box<dyn std::error::Error>> {
24503        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
24504        if portable_mma_gated() {
24505            return self.sdpa_naive(
24506                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
24507            );
24508        }
24509        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
24510        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
24511        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
24512        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
24513        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
24514        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24515        let f32_stage =
24516            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
24517        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
24518        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
24519        // Own numeric config (partial-sum order) — battery-gated.
24520        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24521        let sp = !f32_stage
24522            && *SP_ON.get_or_init(|| {
24523                std::env::var("MEMRA_FA512_SP")
24524                    .map(|v| v != "0")
24525                    .unwrap_or(true)
24526            });
24527        self.fa_prefill_hd512_arm(
24528            q,
24529            k,
24530            v,
24531            o,
24532            head_dim,
24533            n_head,
24534            n_head_kv,
24535            t,
24536            t_kv,
24537            scale,
24538            causal,
24539            f32_stage,
24540            sp,
24541            sp && fa_f16pv_on(),
24542        )
24543    }
24544
24545    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
24546    #[allow(clippy::too_many_arguments)]
24547    pub fn fa_prefill_hd512_pre(
24548        &self,
24549        qb: &CudaSlice<u8>,
24550        kb: &CudaSlice<u8>,
24551        vb: &CudaSlice<u8>,
24552        o: &mut CudaSlice<f32>,
24553        head_dim: usize,
24554        n_head: usize,
24555        n_head_kv: usize,
24556        t: usize,
24557        t_kv: usize,
24558        scale: f32,
24559        causal: bool,
24560        v_f16: bool,
24561    ) -> Result<(), Box<dyn std::error::Error>> {
24562        debug_assert_eq!(head_dim, 512);
24563        const SP_M: usize = 16;
24564        const BKS: usize = 32;
24565        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
24566        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
24567        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
24568        let f16pv = fa_f16pv_on();
24569        let nw = if f16pv { fa512_wide_warps() } else { 2 };
24570        let hp = f16pv
24571            && fa512_hp_on()
24572            && n_head.is_multiple_of(2)
24573            && (n_head / n_head_kv).is_multiple_of(2);
24574        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
24575        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
24576        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
24577            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
24578            let n = t_kv * n_head_kv * head_dim;
24579            let need = n * 2;
24580            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
24581                *vguard = Some(self.alloc_uninit::<u8>(need)?);
24582            }
24583            let dst = vguard.as_mut().unwrap();
24584            self.bf16_to_f16_into(vb, n, dst)?;
24585            vguard.as_ref().unwrap()
24586        } else {
24587            vb
24588        };
24589        let f = self.func(if hp {
24590            "fa_prefill_bf16_hd512_sp16h2"
24591        } else {
24592            match (f16pv, nw) {
24593                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
24594                (true, _) => "fa_prefill_bf16_hd512_sp16",
24595                _ => "fa_prefill_bf16_hd512_sp",
24596            }
24597        });
24598        let (nwarp, npart) = if hp {
24599            (4usize, 4usize)
24600        } else if nw > 2 {
24601            (nw, nw)
24602        } else {
24603            (2, 1)
24604        };
24605        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
24606        let shmem = if hp {
24607            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
24608                as u32
24609        } else {
24610            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
24611                + 4 * (npart * SP_M * BKS + SP_M)) as u32
24612        };
24613        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24614        f.set_attribute(
24615            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24616            shmem as i32,
24617        )?;
24618        let grid_y = if hp {
24619            (n_head / 2) as u32
24620        } else {
24621            n_head as u32
24622        };
24623        let cfg = LaunchConfig {
24624            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
24625            block_dim: (32, nwarp as u32, 1),
24626            shared_mem_bytes: shmem,
24627        };
24628        let (hd, nh, nhkv, ti, tkvi, cz) = (
24629            head_dim as i32,
24630            n_head as i32,
24631            n_head_kv as i32,
24632            t as i32,
24633            t_kv as i32,
24634            causal as i32,
24635        );
24636        let __s_b = self.gpu.stream();
24637        let mut b = __s_b.launch_builder(&f);
24638        b.arg(qb)
24639            .arg(kb)
24640            .arg(vref)
24641            .arg(o)
24642            .arg(&hd)
24643            .arg(&nh)
24644            .arg(&nhkv)
24645            .arg(&ti)
24646            .arg(&tkvi)
24647            .arg(&scale)
24648            .arg(&cz);
24649        unsafe {
24650            b.launch(cfg)?;
24651        }
24652        Ok(())
24653    }
24654
24655    /// Absorbed-form MLA prefill attention over a DSA-GATHERED index list, on tensor cores —
24656    /// the MEMRA_MLA_TC_PREFILL kernel (`fa_mla_gathered_bf16`, cu/flash_attn.cu). One CTA per
24657    /// (query, 16-head band); the query's index list is shared across heads (the DSA indexer
24658    /// mixes heads BEFORE top-k), which is exactly what gives the MMA its m axis. V is K
24659    /// (NoPE latent rows), so the kernel is `kv_rank == 512, d_rope == 0` ONLY and this
24660    /// launcher refuses anything else rather than approximate.
24661    #[allow(clippy::too_many_arguments)]
24662    pub fn mla_attn_gathered_tc(
24663        &self,
24664        q_lat_bf: &CudaSlice<u8>,   // [t_q, n_head, 512] bf16
24665        cache_bf: &CudaSlice<u8>,   // [t_kv, 512] bf16 latent rows
24666        idx: &CudaSlice<i32>,       // [t_q, width], ascending, -1 trailing
24667        o_lat: &mut CudaSlice<f32>, // [t_q, n_head, 512] f32
24668        n_head: usize,
24669        kv_rank: usize,
24670        t_q: usize,
24671        width: usize,
24672        scale: f32,
24673    ) -> Result<(), Box<dyn std::error::Error>> {
24674        if kv_rank != 512 {
24675            return Err(format!(
24676                "mla_attn_gathered_tc is stamped at kv_rank 512 (the glm5_next latent width); \
24677                 got {kv_rank} — the caller's door must fall back to the f32 gathered kernel"
24678            )
24679            .into());
24680        }
24681        if t_q == 0 || n_head == 0 {
24682            return Ok(());
24683        }
24684        const SP_M: usize = 16;
24685        const BKS: usize = 32;
24686        const HD: usize = 512;
24687        let f = self.func("fa_mla_gathered_bf16");
24688        // sQ + sK (V aliases K) + sP bf16, sS + sL f32, sIdx i32.
24689        let shmem =
24690            (2 * (SP_M * HD + BKS * HD + SP_M * BKS) + 4 * (SP_M * BKS + SP_M) + 4 * BKS) as u32;
24691        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24692        f.set_attribute(
24693            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24694            shmem as i32,
24695        )?;
24696        let cfg = LaunchConfig {
24697            grid_dim: (t_q as u32, (n_head as u32).div_ceil(SP_M as u32), 1),
24698            block_dim: (32, 2, 1),
24699            shared_mem_bytes: shmem,
24700        };
24701        let (nh, tq, w) = (n_head as i32, t_q as i32, width as i32);
24702        let __s_b = self.gpu.stream();
24703        let mut b = __s_b.launch_builder(&f);
24704        b.arg(q_lat_bf)
24705            .arg(cache_bf)
24706            .arg(idx)
24707            .arg(o_lat)
24708            .arg(&nh)
24709            .arg(&tq)
24710            .arg(&w)
24711            .arg(&scale);
24712        unsafe {
24713            b.launch(cfg)?;
24714        }
24715        Ok(())
24716    }
24717
24718    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
24719    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
24720    #[allow(clippy::too_many_arguments)]
24721    #[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
24722    pub fn fa_prefill_hd512_arm(
24723        &self,
24724        q: &CudaSlice<f32>,
24725        k: &CudaSlice<f32>,
24726        v: &CudaSlice<f32>,
24727        o: &mut CudaSlice<f32>,
24728        head_dim: usize,
24729        n_head: usize,
24730        n_head_kv: usize,
24731        t: usize,
24732        t_kv: usize,
24733        scale: f32,
24734        causal: bool,
24735        f32_stage: bool,
24736        sp: bool,
24737        f16pv: bool,
24738    ) -> Result<(), Box<dyn std::error::Error>> {
24739        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
24740        if sp && !f32_stage {
24741            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
24742            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
24743            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
24744            const SP_M: usize = 16;
24745            const BKS: usize = 32;
24746            let nw = if f16pv { fa512_wide_warps() } else { 2 };
24747            let hp = f16pv
24748                && fa512_hp_on()
24749                && n_head.is_multiple_of(2)
24750                && (n_head / n_head_kv).is_multiple_of(2);
24751            let f = self.func(if hp {
24752                "fa_prefill_bf16_hd512_sp16h2"
24753            } else {
24754                match (f16pv, nw) {
24755                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
24756                    (true, _) => "fa_prefill_bf16_hd512_sp16",
24757                    _ => "fa_prefill_bf16_hd512_sp",
24758                }
24759            });
24760            let (nwarp, npart) = if hp {
24761                (4usize, 4usize)
24762            } else if nw > 2 {
24763                (nw, nw)
24764            } else {
24765                (2, 1)
24766            };
24767            let shmem = if hp {
24768                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
24769                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
24770            } else {
24771                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
24772                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
24773            };
24774            use cudarc::driver::sys::CUfunction_attribute_enum as A;
24775            f.set_attribute(
24776                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24777                shmem as i32,
24778            )?;
24779            let grid_y = if hp {
24780                (n_head / 2) as u32
24781            } else {
24782                n_head as u32
24783            };
24784            let cfg = LaunchConfig {
24785                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
24786                block_dim: (32, nwarp as u32, 1),
24787                shared_mem_bytes: shmem,
24788            };
24789            let (hd, nh, nhkv, ti, tkvi, cz) = (
24790                head_dim as i32,
24791                n_head as i32,
24792                n_head_kv as i32,
24793                t as i32,
24794                t_kv as i32,
24795                causal as i32,
24796            );
24797            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24798            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24799            let vb = if f16pv {
24800                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
24801            } else {
24802                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
24803            };
24804            let __s_b = self.gpu.stream();
24805            let mut b = __s_b.launch_builder(&f);
24806            b.arg(&qb)
24807                .arg(&kb)
24808                .arg(&vb)
24809                .arg(o)
24810                .arg(&hd)
24811                .arg(&nh)
24812                .arg(&nhkv)
24813                .arg(&ti)
24814                .arg(&tkvi)
24815                .arg(&scale)
24816                .arg(&cz);
24817            unsafe {
24818                b.launch(cfg)?;
24819            }
24820            return Ok(());
24821        }
24822        const BLOCK_Q: usize = 32;
24823        const BK: usize = 32;
24824        const HALF: usize = 256;
24825        let f = self.func(if f32_stage {
24826            "fa_prefill_f32_hd512"
24827        } else {
24828            "fa_prefill_bf16_hd512"
24829        });
24830        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
24831        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
24832            + 4 * BLOCK_Q) as u32;
24833        use cudarc::driver::sys::CUfunction_attribute_enum as A;
24834        f.set_attribute(
24835            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24836            shmem as i32,
24837        )?;
24838        let cfg = LaunchConfig {
24839            grid_dim: (
24840                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
24841                n_head as u32,
24842                2,
24843            ),
24844            block_dim: (32, 2, 1),
24845            shared_mem_bytes: shmem,
24846        };
24847        let (hd, nh, nhkv, ti, tkvi, cz) = (
24848            head_dim as i32,
24849            n_head as i32,
24850            n_head_kv as i32,
24851            t as i32,
24852            t_kv as i32,
24853            causal as i32,
24854        );
24855        if f32_stage {
24856            let __s_b = self.gpu.stream();
24857            let mut b = __s_b.launch_builder(&f);
24858            b.arg(q)
24859                .arg(k)
24860                .arg(v)
24861                .arg(o)
24862                .arg(&hd)
24863                .arg(&nh)
24864                .arg(&nhkv)
24865                .arg(&ti)
24866                .arg(&tkvi)
24867                .arg(&scale)
24868                .arg(&cz);
24869            unsafe {
24870                b.launch(cfg)?;
24871            }
24872        } else {
24873            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
24874            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
24875            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
24876            let __s_b = self.gpu.stream();
24877            let mut b = __s_b.launch_builder(&f);
24878            b.arg(&qb)
24879                .arg(&kb)
24880                .arg(&vb)
24881                .arg(o)
24882                .arg(&hd)
24883                .arg(&nh)
24884                .arg(&nhkv)
24885                .arg(&ti)
24886                .arg(&tkvi)
24887                .arg(&scale)
24888                .arg(&cz);
24889            unsafe {
24890                b.launch(cfg)?;
24891            }
24892        }
24893        Ok(())
24894    }
24895
24896    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
24897    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
24898    /// separate f32_to_bf16 the FA entries would run).
24899    #[allow(clippy::too_many_arguments)]
24900    pub fn rope_neox2_bf16e(
24901        &self,
24902        q: &mut CudaSlice<f32>,
24903        k: &mut CudaSlice<f32>,
24904        qb: &mut CudaSlice<u8>,
24905        kb: &mut CudaSlice<u8>,
24906        pos: &CudaSlice<i32>,
24907        head_dim: usize,
24908        n_dims: usize,
24909        nh_q: usize,
24910        nh_k: usize,
24911        n_tokens: usize,
24912        base: f32,
24913        freq_scale: f32,
24914        ff: Option<&CudaSlice<f32>>,
24915    ) -> Result<(), Box<dyn std::error::Error>> {
24916        let f = self.func("rope_neox2_bf16e_f32");
24917        let rows = ((nh_q + nh_k) * n_tokens) as u32;
24918        let cfg = LaunchConfig {
24919            grid_dim: (rows, 1, 1),
24920            block_dim: ((head_dim / 2) as u32, 1, 1),
24921            shared_mem_bytes: 0,
24922        };
24923        let theta_scale = base.powf(-2.0 / n_dims as f32);
24924        let (hd, nd, nhq, nhk, nt) = (
24925            head_dim as i32,
24926            n_dims as i32,
24927            nh_q as i32,
24928            nh_k as i32,
24929            n_tokens as i32,
24930        );
24931        let __s_b = self.gpu.stream();
24932        let mut b = __s_b.launch_builder(&f);
24933        match ff {
24934            Some(t) => {
24935                b.arg(&mut *q)
24936                    .arg(&mut *k)
24937                    .arg(&mut *qb)
24938                    .arg(&mut *kb)
24939                    .arg(pos)
24940                    .arg(&hd)
24941                    .arg(&nd)
24942                    .arg(&nhq)
24943                    .arg(&nhk)
24944                    .arg(&nt)
24945                    .arg(&theta_scale)
24946                    .arg(&freq_scale)
24947                    .arg(t);
24948                unsafe {
24949                    b.launch(cfg)?;
24950                }
24951            }
24952            None => {
24953                let null: u64 = 0;
24954                b.arg(&mut *q)
24955                    .arg(&mut *k)
24956                    .arg(&mut *qb)
24957                    .arg(&mut *kb)
24958                    .arg(pos)
24959                    .arg(&hd)
24960                    .arg(&nd)
24961                    .arg(&nhq)
24962                    .arg(&nhk)
24963                    .arg(&nt)
24964                    .arg(&theta_scale)
24965                    .arg(&freq_scale)
24966                    .arg(&null);
24967                unsafe {
24968                    b.launch(cfg)?;
24969                }
24970            }
24971        }
24972        Ok(())
24973    }
24974
24975    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
24976    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
24977    pub fn f32_to_bf16(
24978        &self,
24979        x: &CudaSlice<f32>,
24980        n: usize,
24981    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
24982        assert!(
24983            n.is_multiple_of(4),
24984            "f32_to_bf16 requires n % 4 == 0, got {n}"
24985        );
24986        let mut y = self.alloc_uninit::<u8>(n * 2)?;
24987        let f = self.func("f32_to_bf16_flat");
24988        let n_i = n as i64;
24989        let cfg = LaunchConfig {
24990            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
24991            block_dim: (256, 1, 1),
24992            shared_mem_bytes: 0,
24993        };
24994        let __s_b = self.gpu.stream();
24995        let mut b = __s_b.launch_builder(&f);
24996        b.arg(x).arg(&mut y).arg(&n_i);
24997        unsafe {
24998            b.launch(cfg)?;
24999        }
25000        Ok(y)
25001    }
25002
25003    pub fn f32_to_f16(
25004        &self,
25005        x: &CudaSlice<f32>,
25006        n: usize,
25007    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25008        assert!(
25009            n.is_multiple_of(4),
25010            "f32_to_f16 requires n % 4 == 0, got {n}"
25011        );
25012        let mut y = self.alloc_uninit::<u8>(n * 2)?;
25013        let f = self.func("f32_to_f16_flat");
25014        let n_i = n as i64;
25015        let cfg = LaunchConfig {
25016            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
25017            block_dim: (256, 1, 1),
25018            shared_mem_bytes: 0,
25019        };
25020        let __s_b = self.gpu.stream();
25021        let mut b = __s_b.launch_builder(&f);
25022        b.arg(x).arg(&mut y).arg(&n_i);
25023        unsafe {
25024            b.launch(cfg)?;
25025        }
25026        Ok(y)
25027    }
25028
25029    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
25030    pub fn bf16_to_f16(
25031        &self,
25032        xb: &CudaSlice<u8>,
25033        n: usize,
25034    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25035        let mut y = self.alloc_uninit::<u8>(n * 2)?;
25036        self.bf16_to_f16_into(xb, n, &mut y)?;
25037        Ok(y)
25038    }
25039
25040    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
25041    pub fn bf16_to_f16_into(
25042        &self,
25043        xb: &CudaSlice<u8>,
25044        n: usize,
25045        y: &mut CudaSlice<u8>,
25046    ) -> Result<(), Box<dyn std::error::Error>> {
25047        assert!(
25048            n.is_multiple_of(2),
25049            "bf16_to_f16 requires n % 2 == 0, got {n}"
25050        );
25051        assert!(y.len() >= n * 2);
25052        let f = self.func("bf16_to_f16_flat");
25053        let n2 = (n / 2) as i64;
25054        let cfg = LaunchConfig {
25055            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
25056            block_dim: (256, 1, 1),
25057            shared_mem_bytes: 0,
25058        };
25059        let __s_b = self.gpu.stream();
25060        let mut b = __s_b.launch_builder(&f);
25061        b.arg(xb).arg(y).arg(&n2);
25062        unsafe {
25063            b.launch(cfg)?;
25064        }
25065        Ok(())
25066    }
25067
25068    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
25069    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
25070    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
25071    /// head_dim in {256, 128}, bf16kv lane on.
25072    #[allow(clippy::too_many_arguments)]
25073    pub fn fa_prefill_vl8(
25074        &self,
25075        seqs: &[FaSeqVl],
25076        head_dim: usize,
25077        n_head: usize,
25078        n_head_kv: usize,
25079        scale: f32,
25080    ) -> Result<(), Box<dyn std::error::Error>> {
25081        const BK: usize = 32;
25082        let b = seqs.len();
25083        assert!((1..=8).contains(&b));
25084        let mut packed = [FaSeqVl::default(); 8];
25085        packed[..b].copy_from_slice(seqs);
25086        let v = FaVl8(packed);
25087        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25088        let ept = (n_head_kv * head_dim) as i32;
25089        {
25090            let f = self.func("fa_mirror_vl");
25091            let max_n = (max_t as i64) * ept as i64;
25092            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25093            for which in 0..2i32 {
25094                let cfg = LaunchConfig {
25095                    grid_dim: (blocks, 1, b as u32),
25096                    block_dim: (256, 1, 1),
25097                    shared_mem_bytes: 0,
25098                };
25099                let __s_lb = self.gpu.stream();
25100                let mut lb = __s_lb.launch_builder(&f);
25101                lb.arg(&v).arg(&ept).arg(&which);
25102                unsafe {
25103                    lb.launch(cfg)?;
25104                }
25105            }
25106        }
25107        let hd_sfx = fa_hd_suffix(head_dim)?;
25108        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
25109        let block_q = 64usize;
25110        let kv_stages = 2usize;
25111        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
25112            + 4 * (block_q * BK + 2 * block_q)) as u32;
25113        use cudarc::driver::sys::CUfunction_attribute_enum as A;
25114        f.set_attribute(
25115            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25116            shmem as i32,
25117        )?;
25118        let cfg = LaunchConfig {
25119            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
25120            block_dim: (32, 4, 1),
25121            shared_mem_bytes: shmem,
25122        };
25123        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
25124        let __s_lb = self.gpu.stream();
25125        let mut lb = __s_lb.launch_builder(&f);
25126        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
25127        unsafe {
25128            lb.launch(cfg)?;
25129        }
25130        Ok(())
25131    }
25132
25133    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
25134    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
25135    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
25136    #[allow(clippy::too_many_arguments)]
25137    pub fn attn_pre_vl8(
25138        &self,
25139        seqs: &[AttnPreVl],
25140        wq: &CudaSlice<f32>,
25141        wk: &CudaSlice<f32>,
25142        head_dim: usize,
25143        rope_dims: usize,
25144        n_head: usize,
25145        n_head_kv: usize,
25146        eps: f32,
25147        freq_base: f32,
25148        freq_scale: f32,
25149        kv_dim_k: usize,
25150        kv_dim_v: usize,
25151        k_tok_bytes: usize,
25152        v_tok_bytes: usize,
25153    ) -> Result<(), Box<dyn std::error::Error>> {
25154        let b = seqs.len();
25155        assert!((1..=8).contains(&b));
25156        let mut packed = [AttnPreVl::default(); 8];
25157        packed[..b].copy_from_slice(seqs);
25158        let v = AttnPreVl8(packed);
25159        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25160        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
25161        {
25162            let f = self.func("q_gate_split_vl");
25163            let n = max_t * (n_head * head_dim) as u32;
25164            let cfg = LaunchConfig {
25165                grid_dim: (n.div_ceil(256), 1, b as u32),
25166                block_dim: (256, 1, 1),
25167                shared_mem_bytes: 0,
25168            };
25169            let __s_lb = self.gpu.stream();
25170            let mut lb = __s_lb.launch_builder(&f);
25171            lb.arg(&v).arg(&hd).arg(&nh);
25172            unsafe {
25173                lb.launch(cfg)?;
25174            }
25175        }
25176        {
25177            let f = self.func("attn_rms_vl");
25178            let cfg = LaunchConfig {
25179                grid_dim: (max_t * n_head as u32, 2, b as u32),
25180                block_dim: (rms_block(), 1, 1),
25181                shared_mem_bytes: 0,
25182            };
25183            let __s_lb = self.gpu.stream();
25184            let mut lb = __s_lb.launch_builder(&f);
25185            lb.arg(&v)
25186                .arg(wq)
25187                .arg(wk)
25188                .arg(&hd)
25189                .arg(&nh)
25190                .arg(&nhkv)
25191                .arg(&eps);
25192            unsafe {
25193                lb.launch(cfg)?;
25194            }
25195        }
25196        {
25197            let f = self.func("attn_rope_vl");
25198            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
25199            let nd = rope_dims as i32;
25200            let cfg = LaunchConfig {
25201                grid_dim: (max_t * n_head as u32, 2, b as u32),
25202                block_dim: ((head_dim / 2) as u32, 1, 1),
25203                shared_mem_bytes: 0,
25204            };
25205            let __s_lb = self.gpu.stream();
25206            let mut lb = __s_lb.launch_builder(&f);
25207            lb.arg(&v)
25208                .arg(&hd)
25209                .arg(&nd)
25210                .arg(&nh)
25211                .arg(&nhkv)
25212                .arg(&theta_scale)
25213                .arg(&freq_scale);
25214            unsafe {
25215                lb.launch(cfg)?;
25216            }
25217        }
25218        {
25219            let f = self.func("append_kv_vl");
25220            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
25221            let cfg = LaunchConfig {
25222                grid_dim: (nblk, max_t, b as u32),
25223                block_dim: (32, 1, 1),
25224                shared_mem_bytes: 0,
25225            };
25226            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
25227            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25228            let __s_lb = self.gpu.stream();
25229            let mut lb = __s_lb.launch_builder(&f);
25230            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
25231            unsafe {
25232                lb.launch(cfg)?;
25233            }
25234        }
25235        Ok(())
25236    }
25237
25238    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
25239    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
25240    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
25241    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
25242    #[allow(clippy::too_many_arguments)]
25243    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25244    #[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
25245    pub fn fa_prefill_view(
25246        &self,
25247        q: &CudaSlice<f32>,
25248        k: &cudarc::driver::CudaView<u8>,
25249        v: &cudarc::driver::CudaView<u8>,
25250        o: &mut CudaSlice<f32>,
25251        head_dim: usize,
25252        n_head: usize,
25253        n_head_kv: usize,
25254        t: usize,
25255        t_kv: usize,
25256        scale: f32,
25257        causal: bool,
25258        k_tok_bytes: usize,
25259        v_tok_bytes: usize,
25260        g: bool,
25261    ) -> Result<(), Box<dyn std::error::Error>> {
25262        if portable_mma_gated() {
25263            return self.sdpa_naive_quantized_view(
25264                q,
25265                k,
25266                v,
25267                o,
25268                head_dim,
25269                n_head,
25270                n_head_kv,
25271                t,
25272                t_kv,
25273                scale,
25274                causal,
25275                k_tok_bytes,
25276                v_tok_bytes,
25277            );
25278        }
25279        const BLOCK_Q: usize = 64;
25280        const BK: usize = 32;
25281        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
25282        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
25283        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
25284        let f = if g {
25285            self.func_g(&name)
25286        } else {
25287            self.func(&name)
25288        };
25289        let shmem =
25290            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
25291        use cudarc::driver::sys::CUfunction_attribute_enum as A;
25292        f.set_attribute(
25293            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25294            shmem as i32,
25295        )?;
25296        let cfg = LaunchConfig {
25297            grid_dim: (
25298                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25299                n_head as u32,
25300                1,
25301            ),
25302            block_dim: (32, 4, 1),
25303            shared_mem_bytes: shmem,
25304        };
25305        let (hd, nh, nhkv, ti, tkvi, cz) = (
25306            head_dim as i32,
25307            n_head as i32,
25308            n_head_kv as i32,
25309            t as i32,
25310            t_kv as i32,
25311            causal as i32,
25312        );
25313        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25314        let __s_b = self.gpu.stream();
25315        let mut b = __s_b.launch_builder(&f);
25316        b.arg(q)
25317            .arg(k)
25318            .arg(v)
25319            .arg(o)
25320            .arg(&hd)
25321            .arg(&nh)
25322            .arg(&nhkv)
25323            .arg(&ti)
25324            .arg(&tkvi)
25325            .arg(&scale)
25326            .arg(&cz)
25327            .arg(&ktb)
25328            .arg(&vtb);
25329        unsafe {
25330            b.launch(cfg)?;
25331        }
25332        Ok(())
25333    }
25334
25335    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
25336    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
25337    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
25338    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
25339    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
25340    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
25341    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
25342    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
25343    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
25344    #[allow(clippy::too_many_arguments)]
25345    #[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
25346    pub fn fa_prefill_view_ws(
25347        &self,
25348        q: &CudaSlice<f32>,
25349        k: &cudarc::driver::CudaView<u8>,
25350        v: &cudarc::driver::CudaView<u8>,
25351        o: &mut CudaSlice<f32>,
25352        head_dim: usize,
25353        n_head: usize,
25354        n_head_kv: usize,
25355        t: usize,
25356        t_kv: usize,
25357        scale: f32,
25358        causal: bool,
25359        k_tok_bytes: usize,
25360        v_tok_bytes: usize,
25361        g: bool,
25362    ) -> Result<(), Box<dyn std::error::Error>> {
25363        if portable_mma_gated() {
25364            return self.sdpa_naive_quantized_view(
25365                q,
25366                k,
25367                v,
25368                o,
25369                head_dim,
25370                n_head,
25371                n_head_kv,
25372                t,
25373                t_kv,
25374                scale,
25375                causal,
25376                k_tok_bytes,
25377                v_tok_bytes,
25378            );
25379        }
25380        const BLOCK_Q: usize = 64;
25381        const BK: usize = 32;
25382        let kv_dim_k = n_head_kv * head_dim;
25383        let kv_dim_v = n_head_kv * head_dim;
25384        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
25385        let v_ws_bytes = t_kv * kv_dim_v * 2;
25386        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
25387        let mut guard = self.prime_deqw_ws.lock().unwrap();
25388        let need_grow = match guard.as_ref() {
25389            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
25390            None => true,
25391        };
25392        if need_grow {
25393            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
25394            let (ck, cv) = guard
25395                .as_ref()
25396                .map(|(a, b)| (a.len(), b.len()))
25397                .unwrap_or((0, 0));
25398            *guard = Some((
25399                self.alloc_u8(grow(ck, k_ws_bytes))?,
25400                self.alloc_u8(grow(cv, v_ws_bytes))?,
25401            ));
25402        }
25403        let (kw, vw) = guard.as_mut().unwrap();
25404        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
25405        {
25406            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
25407            let f = if g {
25408                self.func_g("fa_dequant_kv_ws_bf16")
25409            } else {
25410                self.func("fa_dequant_kv_ws_bf16")
25411            };
25412            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
25413            #[allow(clippy::manual_div_ceil)]
25414            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25415            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
25416            let cfg = LaunchConfig {
25417                grid_dim: (nblk.max(1), 1, 1),
25418                block_dim: (256, 1, 1),
25419                shared_mem_bytes: 0,
25420            };
25421            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
25422            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25423            let __s_b = self.gpu.stream();
25424            let mut b = __s_b.launch_builder(&f);
25425            b.arg(k)
25426                .arg(v)
25427                .arg(&mut *kw)
25428                .arg(&mut *vw)
25429                .arg(&kdk)
25430                .arg(&kdv)
25431                .arg(&tkvi)
25432                .arg(&ktb)
25433                .arg(&vtb);
25434            unsafe {
25435                b.launch(cfg)?;
25436            }
25437        }
25438        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
25439        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
25440        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
25441        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
25442        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
25443        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
25444        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
25445        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
25446            .map(|v| v != "0")
25447            .unwrap_or(true);
25448        {
25449            let hd_sfx = fa_hd_suffix(head_dim)?;
25450            let f = self.func(&format!(
25451                "fa_prefill_qw{}{hd_sfx}",
25452                if db { "_db" } else { "" }
25453            ));
25454            let shmem = if db {
25455                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
25456                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
25457            } else {
25458                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
25459            };
25460            use cudarc::driver::sys::CUfunction_attribute_enum as A;
25461            f.set_attribute(
25462                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25463                shmem as i32,
25464            )?;
25465            let cfg = LaunchConfig {
25466                grid_dim: (
25467                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25468                    n_head as u32,
25469                    1,
25470                ),
25471                block_dim: (32, 4, 1),
25472                shared_mem_bytes: shmem,
25473            };
25474            let (hd, nh, nhkv, ti, tkvi, cz) = (
25475                head_dim as i32,
25476                n_head as i32,
25477                n_head_kv as i32,
25478                t as i32,
25479                t_kv as i32,
25480                causal as i32,
25481            );
25482            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
25483            let __s_b = self.gpu.stream();
25484            let mut b = __s_b.launch_builder(&f);
25485            b.arg(q)
25486                .arg(&*kw)
25487                .arg(&*vw)
25488                .arg(o)
25489                .arg(&hd)
25490                .arg(&nh)
25491                .arg(&nhkv)
25492                .arg(&ti)
25493                .arg(&tkvi)
25494                .arg(&scale)
25495                .arg(&cz)
25496                .arg(&kdk)
25497                .arg(&kdv);
25498            unsafe {
25499                b.launch(cfg)?;
25500            }
25501        }
25502        Ok(())
25503    }
25504
25505    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
25506    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
25507    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
25508    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
25509    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
25510    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
25511    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
25512    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
25513    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
25514    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
25515    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
25516    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
25517    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
25518    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
25519    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
25520    #[allow(clippy::too_many_arguments)]
25521    #[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
25522    pub fn fa_prefill_view_ws_w_hd128(
25523        &self,
25524        q: &CudaSlice<f32>,
25525        k: &cudarc::driver::CudaView<u8>,
25526        v: &cudarc::driver::CudaView<u8>,
25527        o: &mut CudaSlice<f32>,
25528        head_dim: usize,
25529        n_head: usize,
25530        n_head_kv: usize,
25531        t: usize,
25532        t_kv: usize,
25533        scale: f32,
25534        causal: bool,
25535        window: usize,
25536        k_tok_bytes: usize,
25537        v_tok_bytes: usize,
25538    ) -> Result<(), Box<dyn std::error::Error>> {
25539        assert_eq!(
25540            head_dim, 128,
25541            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
25542        );
25543        if portable_mma_gated() {
25544            return self.sdpa_naive_w_quantized_view(
25545                q,
25546                k,
25547                v,
25548                o,
25549                head_dim,
25550                n_head,
25551                n_head_kv,
25552                t,
25553                t_kv,
25554                scale,
25555                causal,
25556                window,
25557                k_tok_bytes,
25558                v_tok_bytes,
25559            );
25560        }
25561        const BLOCK_Q: usize = 64;
25562        const BK: usize = 32;
25563        let kv_dim_k = n_head_kv * head_dim;
25564        let kv_dim_v = n_head_kv * head_dim;
25565        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
25566        let v_ws_bytes = t_kv * kv_dim_v * 2;
25567        let mut guard = self.prime_deqw_ws.lock().unwrap();
25568        let need_grow = match guard.as_ref() {
25569            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
25570            None => true,
25571        };
25572        if need_grow {
25573            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
25574            let (ck, cv) = guard
25575                .as_ref()
25576                .map(|(a, b)| (a.len(), b.len()))
25577                .unwrap_or((0, 0));
25578            *guard = Some((
25579                self.alloc_u8(grow(ck, k_ws_bytes))?,
25580                self.alloc_u8(grow(cv, v_ws_bytes))?,
25581            ));
25582        }
25583        let (kw, vw) = guard.as_mut().unwrap();
25584        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
25585        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
25586        {
25587            let f = self.func("fa_dequant_kv_ws_bf16");
25588            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
25589            #[allow(clippy::manual_div_ceil)]
25590            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
25591            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
25592            let cfg = LaunchConfig {
25593                grid_dim: (nblk.max(1), 1, 1),
25594                block_dim: (256, 1, 1),
25595                shared_mem_bytes: 0,
25596            };
25597            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
25598            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25599            let __s_b = self.gpu.stream();
25600            let mut b = __s_b.launch_builder(&f);
25601            b.arg(k)
25602                .arg(v)
25603                .arg(&mut *kw)
25604                .arg(&mut *vw)
25605                .arg(&kdk)
25606                .arg(&kdv)
25607                .arg(&tkvi)
25608                .arg(&ktb)
25609                .arg(&vtb);
25610            unsafe {
25611                b.launch(cfg)?;
25612            }
25613        }
25614        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
25615        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
25616            .map(|v| v != "0")
25617            .unwrap_or(true);
25618        {
25619            let f = self.func(if db {
25620                "fa_prefill_qw_db_w_hd128"
25621            } else {
25622                "fa_prefill_qw_w_hd128"
25623            });
25624            let shmem = if db {
25625                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
25626            } else {
25627                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
25628            };
25629            use cudarc::driver::sys::CUfunction_attribute_enum as A;
25630            f.set_attribute(
25631                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25632                shmem as i32,
25633            )?;
25634            let cfg = LaunchConfig {
25635                grid_dim: (
25636                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
25637                    n_head as u32,
25638                    1,
25639                ),
25640                block_dim: (32, 4, 1),
25641                shared_mem_bytes: shmem,
25642            };
25643            let (hd, nh, nhkv, ti, tkvi, cz) = (
25644                head_dim as i32,
25645                n_head as i32,
25646                n_head_kv as i32,
25647                t as i32,
25648                t_kv as i32,
25649                causal as i32,
25650            );
25651            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
25652            let __s_b = self.gpu.stream();
25653            let mut b = __s_b.launch_builder(&f);
25654            b.arg(q)
25655                .arg(&*kw)
25656                .arg(&*vw)
25657                .arg(o)
25658                .arg(&hd)
25659                .arg(&nh)
25660                .arg(&nhkv)
25661                .arg(&ti)
25662                .arg(&tkvi)
25663                .arg(&scale)
25664                .arg(&cz)
25665                .arg(&kdk)
25666                .arg(&kdv)
25667                .arg(&wnd);
25668            unsafe {
25669                b.launch(cfg)?;
25670            }
25671        }
25672        Ok(())
25673    }
25674
25675    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
25676    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
25677    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
25678    #[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
25679    pub fn fa_decode(
25680        &self,
25681        q: &CudaSlice<f32>,
25682        k: &cudarc::driver::CudaView<u8>,
25683        v: &cudarc::driver::CudaView<u8>,
25684        o: &mut CudaSlice<f32>,
25685        head_dim: usize,
25686        n_head: usize,
25687        n_head_kv: usize,
25688        t_kv: usize,
25689        scale: f32,
25690        k_tok_bytes: usize,
25691        v_tok_bytes: usize,
25692    ) -> Result<(), Box<dyn std::error::Error>> {
25693        self.fa_decode_kvmod(
25694            q,
25695            k,
25696            v,
25697            o,
25698            head_dim,
25699            n_head,
25700            n_head_kv,
25701            t_kv,
25702            scale,
25703            k_tok_bytes,
25704            v_tok_bytes,
25705            false,
25706        )
25707    }
25708
25709    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
25710    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
25711    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
25712    #[allow(clippy::too_many_arguments)]
25713    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
25714    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
25715    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
25716    #[allow(clippy::too_many_arguments)]
25717    #[allow(clippy::too_many_arguments)]
25718    fn fa_decode_scalar_unified(
25719        &self,
25720        q: &cudarc::driver::CudaView<f32>,
25721        k: &cudarc::driver::CudaView<u8>,
25722        v: &cudarc::driver::CudaView<u8>,
25723        o: &mut cudarc::driver::CudaViewMut<f32>,
25724        head_dim: usize,
25725        n_head: usize,
25726        n_head_kv: usize,
25727        t_kv_host: usize,
25728        t_kv_dev: Option<&CudaSlice<i32>>,
25729        scale: f32,
25730        n_splits: usize,
25731        split_keys: usize,
25732        k_tok_bytes: usize,
25733        v_tok_bytes: usize,
25734        g: bool,
25735        part_o: &mut CudaSlice<f32>,
25736        part_m: &mut CudaSlice<f32>,
25737        part_l: &mut CudaSlice<f32>,
25738        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
25739    ) -> Result<(), Box<dyn std::error::Error>> {
25740        let f = if g {
25741            self.func_g("fa_decode_f32")
25742        } else {
25743            self.fa_func("fa_decode_f32", head_dim)
25744        };
25745        let cfg = LaunchConfig {
25746            grid_dim: (n_head as u32, n_splits as u32, 1),
25747            block_dim: (head_dim as u32, 1, 1),
25748            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
25749        };
25750        let (hd, nh, nhkv, nsp) = (
25751            head_dim as i32,
25752            n_head as i32,
25753            n_head_kv as i32,
25754            n_splits as i32,
25755        );
25756        let (ktb, vtb, tkvi, ski) = (
25757            k_tok_bytes as i64,
25758            v_tok_bytes as i64,
25759            t_kv_host as i32,
25760            split_keys as i32,
25761        );
25762        let __s_b = self.gpu.stream();
25763        let mut b = __s_b.launch_builder(&f);
25764        match t_kv_dev {
25765            Some(d) => {
25766                b.arg(q)
25767                    .arg(k)
25768                    .arg(v)
25769                    .arg(&mut *part_o)
25770                    .arg(&mut *part_m)
25771                    .arg(&mut *part_l)
25772                    .arg(&hd)
25773                    .arg(&nh)
25774                    .arg(&nhkv)
25775                    .arg(&tkvi)
25776                    .arg(d)
25777                    .arg(&scale)
25778                    .arg(&nsp)
25779                    .arg(&ski)
25780                    .arg(&ktb)
25781                    .arg(&vtb);
25782                unsafe {
25783                    b.launch(cfg)?;
25784                }
25785            }
25786            None => {
25787                let null: u64 = 0;
25788                b.arg(q)
25789                    .arg(k)
25790                    .arg(v)
25791                    .arg(&mut *part_o)
25792                    .arg(&mut *part_m)
25793                    .arg(&mut *part_l)
25794                    .arg(&hd)
25795                    .arg(&nh)
25796                    .arg(&nhkv)
25797                    .arg(&tkvi)
25798                    .arg(&null)
25799                    .arg(&scale)
25800                    .arg(&nsp)
25801                    .arg(&ski)
25802                    .arg(&ktb)
25803                    .arg(&vtb);
25804                unsafe {
25805                    b.launch(cfg)?;
25806                }
25807            }
25808        }
25809        let cfg2 = LaunchConfig {
25810            grid_dim: (n_head as u32, 1, 1),
25811            block_dim: (head_dim as u32, 1, 1),
25812            shared_mem_bytes: 0,
25813        };
25814        if let Some((oq, od)) = q8_out {
25815            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
25816            let fc = if g {
25817                self.func_g("fa_decode_combine_q8_1")
25818            } else {
25819                self.fa_func("fa_decode_combine_q8_1", head_dim)
25820            };
25821            let __s_b2 = self.gpu.stream();
25822            let mut b2 = __s_b2.launch_builder(&fc);
25823            b2.arg(&*part_o)
25824                .arg(&*part_m)
25825                .arg(&*part_l)
25826                .arg(oq)
25827                .arg(od)
25828                .arg(&hd)
25829                .arg(&nh)
25830                .arg(&nsp);
25831            unsafe {
25832                b2.launch(cfg2)?;
25833            }
25834            return Ok(());
25835        }
25836        let fc = if g {
25837            self.func_g("fa_decode_combine_f32")
25838        } else {
25839            self.fa_func("fa_decode_combine_f32", head_dim)
25840        };
25841        let __s_b2 = self.gpu.stream();
25842        let mut b2 = __s_b2.launch_builder(&fc);
25843        b2.arg(&*part_o)
25844            .arg(&*part_m)
25845            .arg(&*part_l)
25846            .arg(o)
25847            .arg(&hd)
25848            .arg(&nh)
25849            .arg(&nsp);
25850        unsafe {
25851            b2.launch(cfg2)?;
25852        }
25853        Ok(())
25854    }
25855
25856    #[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
25857    pub fn fa_decode_kvmod(
25858        &self,
25859        q: &CudaSlice<f32>,
25860        k: &cudarc::driver::CudaView<u8>,
25861        v: &cudarc::driver::CudaView<u8>,
25862        o: &mut CudaSlice<f32>,
25863        head_dim: usize,
25864        n_head: usize,
25865        n_head_kv: usize,
25866        t_kv: usize,
25867        scale: f32,
25868        k_tok_bytes: usize,
25869        v_tok_bytes: usize,
25870        g: bool,
25871    ) -> Result<(), Box<dyn std::error::Error>> {
25872        let q_view = q.as_view();
25873        let mut o_view = o.as_view_mut();
25874        self.fa_decode_kvmod_view(
25875            &q_view,
25876            k,
25877            v,
25878            &mut o_view,
25879            head_dim,
25880            n_head,
25881            n_head_kv,
25882            t_kv,
25883            scale,
25884            k_tok_bytes,
25885            v_tok_bytes,
25886            g,
25887        )
25888    }
25889
25890    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
25891    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
25892    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
25893    /// per-session KV view and FA launch.
25894    #[allow(clippy::too_many_arguments)]
25895    #[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
25896    pub fn fa_decode_kvmod_view(
25897        &self,
25898        q: &cudarc::driver::CudaView<f32>,
25899        k: &cudarc::driver::CudaView<u8>,
25900        v: &cudarc::driver::CudaView<u8>,
25901        o: &mut cudarc::driver::CudaViewMut<f32>,
25902        head_dim: usize,
25903        n_head: usize,
25904        n_head_kv: usize,
25905        t_kv: usize,
25906        scale: f32,
25907        k_tok_bytes: usize,
25908        v_tok_bytes: usize,
25909        g: bool,
25910    ) -> Result<(), Box<dyn std::error::Error>> {
25911        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
25912        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
25913        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
25914        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
25915        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
25916        //
25917        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
25918        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
25919        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
25920        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
25921        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
25922        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
25923        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
25924        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
25925        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
25926        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
25927        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
25928        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
25929        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
25930        // fall to the exact scalar there instead of the broken register arm.
25931        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
25932        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
25933        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
25934        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
25935        if g && head_dim == 256 && !fa_v4_at(t_kv) {
25936            fa_vec = false;
25937        }
25938        let sp = fa_split_keys(t_kv, n_head_kv);
25939        let n_splits = if fa_vec {
25940            ((t_kv + sp - 1) / sp).max(1)
25941        } else {
25942            ((t_kv + 255) / 256).max(1)
25943        };
25944        let o_len = n_head * n_splits * head_dim;
25945        let ml_len = n_head * n_splits;
25946        let mut part_guard = self.fa_part_pool.lock().unwrap();
25947        if part_guard
25948            .as_ref()
25949            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
25950            .unwrap_or(true)
25951        {
25952            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
25953            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
25954            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
25955            // later live allocations land at those addresses, and the next graph REPLAY writes
25956            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
25957            // output corruption began the burst after the trunk's t_kv growth first realloc'd
25958            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
25959            // the baked addresses alive (single-stream: eager writes the new buffers, replays
25960            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
25961            // (total retired < final size).
25962            let old = part_guard.take();
25963            let (co, cm) = old
25964                .as_ref()
25965                .map(|pp| (pp.0.len(), pp.1.len()))
25966                .unwrap_or((0, 0));
25967            if let Some(old) = old {
25968                self.fa_part_retired.lock().unwrap().push(old);
25969            }
25970            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
25971                eprintln!(
25972                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
25973                    co, o_len, cm, ml_len
25974                );
25975            }
25976            *part_guard =
25977                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
25978        }
25979        let pg = part_guard.as_mut().unwrap();
25980        self.gpu
25981            .stream()
25982            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
25983        self.gpu
25984            .stream()
25985            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
25986        self.gpu
25987            .stream()
25988            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
25989        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
25990        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
25991        let (hd, nh, nhkv, tkvi, nsp) = (
25992            head_dim as i32,
25993            n_head as i32,
25994            n_head_kv as i32,
25995            t_kv as i32,
25996            n_splits as i32,
25997        );
25998        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25999        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
26000        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
26001        // silently truncating the accumulator.
26002        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
26003        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
26004        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
26005        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
26006        // 178.4 -> 173.7 when 512 rode vec unconditionally).
26007        let fa512_min = fa512_min_tkv();
26008        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
26009        // g-module keeps the v4 pick (its class is not the depth-decay class).
26010        let deep = fa_vec
26011            && head_dim == 256
26012            && fa_v4_at(t_kv)
26013            && !g
26014            && fa_deep_at(t_kv)
26015            && !matches!(fa_v4_mode(), "noB3" | "stage");
26016        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
26017            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
26018            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
26019            let gqa = (n_head / n_head_kv).max(1) as u32;
26020            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
26021            (
26022                fv,
26023                LaunchConfig {
26024                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26025                    block_dim: (32, gqa, 1),
26026                    shared_mem_bytes: 0,
26027                },
26028            )
26029        } else if fa_vec && head_dim <= 256 {
26030            let gqa = (n_head / n_head_kv).max(1) as u32;
26031            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
26032            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
26033            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
26034            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
26035            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
26036            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
26037            // dequant each tile ONCE per block.
26038            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
26039            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
26040            // there by 12x — latency, not bandwidth, rules small KV).
26041            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26042            let smem_tkv = *SMEM_TKV.get_or_init(|| {
26043                std::env::var("MEMRA_FA_SMEM_TKV")
26044                    .ok()
26045                    .and_then(|v| v.parse().ok())
26046                    .unwrap_or_else(|| {
26047                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
26048                    })
26049            });
26050            if fa_v4_at(t_kv) && head_dim == 256 {
26051                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
26052                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
26053                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
26054                let v4name = match fa_v4_mode() {
26055                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
26056                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
26057                    _ if deep => "fa_decode_vec_q_v4_deep",
26058                    _ => "fa_decode_vec_q_v4",
26059                };
26060                let fv = if g {
26061                    self.func_g(v4name)
26062                } else {
26063                    self.func(v4name)
26064                };
26065                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
26066                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
26067                let shmem = (if deep { 12160 } else { 11520 }
26068                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
26069                use cudarc::driver::sys::CUfunction_attribute_enum as A;
26070                fv.set_attribute(
26071                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26072                    shmem as i32,
26073                )?;
26074                (
26075                    fv,
26076                    LaunchConfig {
26077                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26078                        block_dim: (32, gqa, 1),
26079                        shared_mem_bytes: shmem,
26080                    },
26081                )
26082            } else if fa_v3_active(head_dim) {
26083                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
26084                // smem = sV only (half of v2's).
26085                let fv = if g {
26086                    self.func_g("fa_decode_vec_q_v3")
26087                } else {
26088                    self.func("fa_decode_vec_q_v3")
26089                };
26090                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
26091                (
26092                    fv,
26093                    LaunchConfig {
26094                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26095                        block_dim: (32, gqa, 1),
26096                        shared_mem_bytes: shmem,
26097                    },
26098                )
26099            } else if fa_v2_on() {
26100                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
26101                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
26102                // partials; same 32KB sK+sV tile as the smem twin.
26103                let fv = if g {
26104                    self.func_g("fa_decode_vec_q_v2")
26105                } else {
26106                    self.func("fa_decode_vec_q_v2")
26107                };
26108                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
26109                (
26110                    fv,
26111                    LaunchConfig {
26112                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26113                        block_dim: (32, gqa, 1),
26114                        shared_mem_bytes: shmem,
26115                    },
26116                )
26117            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
26118            {
26119                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
26120                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
26121                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
26122                let fv = if g {
26123                    self.func_g("fa_decode_vec_q_smem")
26124                } else {
26125                    self.func("fa_decode_vec_q_smem")
26126                };
26127                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
26128                use cudarc::driver::sys::CUfunction_attribute_enum as A;
26129                fv.set_attribute(
26130                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26131                    shmem as i32,
26132                )?;
26133                (
26134                    fv,
26135                    LaunchConfig {
26136                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26137                        block_dim: (32, gqa, 1),
26138                        shared_mem_bytes: shmem,
26139                    },
26140                )
26141            } else {
26142                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
26143                // dequant, zero dynamic shared memory.
26144                let fv = if g {
26145                    self.func_g("fa_decode_vec_q")
26146                } else {
26147                    self.func("fa_decode_vec_q")
26148                };
26149                (
26150                    fv,
26151                    LaunchConfig {
26152                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
26153                        block_dim: (32, gqa, 1),
26154                        shared_mem_bytes: 0,
26155                    },
26156                )
26157            }
26158        } else {
26159            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
26160            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
26161            return self.fa_decode_scalar_unified(
26162                q,
26163                k,
26164                v,
26165                o,
26166                head_dim,
26167                n_head,
26168                n_head_kv,
26169                t_kv,
26170                None,
26171                scale,
26172                n_splits,
26173                if fa_vec { sp } else { 256 },
26174                k_tok_bytes,
26175                v_tok_bytes,
26176                g,
26177                part_o,
26178                part_m,
26179                part_l,
26180                None,
26181            );
26182        };
26183        let __s_b = self.gpu.stream();
26184        let mut b = __s_b.launch_builder(&f);
26185        b.arg(q)
26186            .arg(k)
26187            .arg(v)
26188            .arg(&mut *part_o)
26189            .arg(&mut *part_m)
26190            .arg(&mut *part_l)
26191            .arg(&hd)
26192            .arg(&nh)
26193            .arg(&nhkv)
26194            .arg(&tkvi)
26195            .arg(&scale)
26196            .arg(&nsp)
26197            .arg(&ktb)
26198            .arg(&vtb);
26199        unsafe {
26200            b.launch(cfg)?;
26201        }
26202        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
26203        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
26204        let (fc, cfg2) = (
26205            if g {
26206                self.func_g("fa_decode_combine_f32")
26207            } else {
26208                self.fa_func("fa_decode_combine_f32", head_dim)
26209            },
26210            LaunchConfig {
26211                grid_dim: (n_head as u32, 1, 1),
26212                block_dim: (head_dim as u32, 1, 1),
26213                shared_mem_bytes: 0,
26214            },
26215        );
26216        let __s_b2 = self.gpu.stream();
26217        let mut b2 = __s_b2.launch_builder(&fc);
26218        b2.arg(&*part_o)
26219            .arg(&*part_m)
26220            .arg(&*part_l)
26221            .arg(o)
26222            .arg(&hd)
26223            .arg(&nh)
26224            .arg(&nsp);
26225        unsafe {
26226            b2.launch(cfg2)?;
26227        }
26228        Ok(())
26229    }
26230
26231    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
26232    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
26233    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
26234    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
26235    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
26236    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
26237    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
26238    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
26239    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
26240    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
26241    #[allow(clippy::too_many_arguments)]
26242    pub fn fa_decode_batch_seqs_v4(
26243        &self,
26244        q: &CudaSlice<f32>,
26245        kv_ptrs: &cudarc::driver::CudaView<u64>,
26246        pos_seq: &CudaSlice<i32>,
26247        o: &mut CudaSlice<f32>,
26248        head_dim: usize,
26249        n_head: usize,
26250        n_head_kv: usize,
26251        b_n: usize,
26252        t_kv_max: usize,
26253        scale: f32,
26254        split_keys: usize,
26255        k_tok_bytes: usize,
26256        v_tok_bytes: usize,
26257    ) -> Result<(), Box<dyn std::error::Error>> {
26258        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
26259        #[allow(clippy::manual_div_ceil)]
26260        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26261        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
26262        let o_len = b_n * n_head * n_splits_max * head_dim;
26263        let ml_len = b_n * n_head * n_splits_max;
26264        let mut part_guard = self.fa_part_pool.lock().unwrap();
26265        if part_guard
26266            .as_ref()
26267            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26268            .unwrap_or(true)
26269        {
26270            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26271            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26272            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26273            // later live allocations land at those addresses, and the next graph REPLAY writes
26274            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26275            // output corruption began the burst after the trunk's t_kv growth first realloc'd
26276            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26277            // the baked addresses alive (single-stream: eager writes the new buffers, replays
26278            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26279            // (total retired < final size).
26280            let old = part_guard.take();
26281            let (co, cm) = old
26282                .as_ref()
26283                .map(|pp| (pp.0.len(), pp.1.len()))
26284                .unwrap_or((0, 0));
26285            if let Some(old) = old {
26286                self.fa_part_retired.lock().unwrap().push(old);
26287            }
26288            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26289                eprintln!(
26290                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26291                    co, o_len, cm, ml_len
26292                );
26293            }
26294            *part_guard =
26295                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
26296        }
26297        let pg = part_guard.as_mut().unwrap();
26298        self.gpu
26299            .stream()
26300            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
26301        self.gpu
26302            .stream()
26303            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
26304        self.gpu
26305            .stream()
26306            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
26307        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
26308        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
26309        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
26310        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26311        let gqa = (n_head / n_head_kv).max(1) as u32;
26312        let f = self.func("fa_decode_vec_q_seqs_v4");
26313        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
26314        let shmem = (11520 + 32 * head_dim * 2) as u32;
26315        use cudarc::driver::sys::CUfunction_attribute_enum as A;
26316        f.set_attribute(
26317            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26318            shmem as i32,
26319        )?;
26320        let cfg = LaunchConfig {
26321            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
26322            block_dim: (32, gqa, 1),
26323            shared_mem_bytes: shmem,
26324        };
26325        {
26326            let __s_b = self.gpu.stream();
26327            let mut b = __s_b.launch_builder(&f);
26328            b.arg(q)
26329                .arg(kv_ptrs)
26330                .arg(pos_seq)
26331                .arg(&mut *part_o)
26332                .arg(&mut *part_m)
26333                .arg(&mut *part_l)
26334                .arg(&hd)
26335                .arg(&nh)
26336                .arg(&nhkv)
26337                .arg(&scale)
26338                .arg(&nspm)
26339                .arg(&spk)
26340                .arg(&ktb)
26341                .arg(&vtb);
26342            unsafe {
26343                b.launch(cfg)?;
26344            }
26345        }
26346        let fc = self.func("fa_decode_combine_seqs");
26347        let cfg2 = LaunchConfig {
26348            grid_dim: (n_head as u32, b_n as u32, 1),
26349            block_dim: (head_dim as u32, 1, 1),
26350            shared_mem_bytes: 0,
26351        };
26352        let __s_b2 = self.gpu.stream();
26353        let mut b2 = __s_b2.launch_builder(&fc);
26354        b2.arg(&*part_o)
26355            .arg(&*part_m)
26356            .arg(&*part_l)
26357            .arg(o)
26358            .arg(&hd)
26359            .arg(&nh)
26360            .arg(pos_seq)
26361            .arg(&nspm)
26362            .arg(&spk);
26363        unsafe {
26364            b2.launch(cfg2)?;
26365        }
26366        Ok(())
26367    }
26368
26369    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
26370    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
26371    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
26372    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
26373    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
26374    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
26375    #[allow(clippy::too_many_arguments)]
26376    pub fn append_kv_quantized_seqs(
26377        &self,
26378        k_rows: &CudaSlice<f32>,
26379        v_rows: &CudaSlice<f32>,
26380        kv_ptrs: &cudarc::driver::CudaView<u64>,
26381        pos_seq: &CudaSlice<i32>,
26382        b_n: usize,
26383        kv_dim_k: usize,
26384        kv_dim_v: usize,
26385        k_tok_bytes: usize,
26386        v_tok_bytes: usize,
26387    ) -> Result<(), Box<dyn std::error::Error>> {
26388        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
26389        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
26390        let cfg = LaunchConfig {
26391            grid_dim: (nblk, b_n as u32, 1),
26392            block_dim: (32, 1, 1),
26393            shared_mem_bytes: 0,
26394        };
26395        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
26396        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26397        let __s_b = self.gpu.stream();
26398        let mut b = __s_b.launch_builder(&f);
26399        b.arg(k_rows)
26400            .arg(v_rows)
26401            .arg(kv_ptrs)
26402            .arg(pos_seq)
26403            .arg(&kdk)
26404            .arg(&kdv)
26405            .arg(&ktb)
26406            .arg(&vtb);
26407        unsafe {
26408            b.launch(cfg)?;
26409        }
26410        Ok(())
26411    }
26412
26413    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
26414    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
26415    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
26416    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
26417    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
26418    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
26419        std::env::var("MEMRA_NO_FA_VEC").is_err()
26420            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
26421            && base_len + 1 >= fa_vec_min_tkv()
26422            && head_dim <= 256
26423            && head_dim.is_multiple_of(32)
26424    }
26425
26426    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
26427    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
26428    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
26429    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
26430    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
26431    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
26432    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
26433    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
26434    #[allow(clippy::too_many_arguments)]
26435    pub fn fa_decode_rows(
26436        &self,
26437        q: &CudaSlice<f32>,
26438        k: &cudarc::driver::CudaView<u8>,
26439        v: &cudarc::driver::CudaView<u8>,
26440        o: &mut CudaSlice<f32>,
26441        head_dim: usize,
26442        n_head: usize,
26443        n_head_kv: usize,
26444        base_len: usize,
26445        t: usize,
26446        scale: f32,
26447        k_tok_bytes: usize,
26448        v_tok_bytes: usize,
26449        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
26450        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
26451        // keep the host arg. None is a bug for hd512 (asserted below).
26452        base_dev: Option<(&CudaSlice<i32>, i32)>,
26453        // K and V planes hold the same values (gemma globals, wv:=wk): pick
26454        // the _kv twin — V plane never read, value rides the q8_0 key dq.
26455        kv_shared: bool,
26456        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
26457        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
26458        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
26459        g: bool,
26460        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
26461        // (hd512 path) — the standalone quantize launch folds away.
26462        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
26463    ) -> Result<(), Box<dyn std::error::Error>> {
26464        debug_assert!(
26465            base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim.is_multiple_of(32)
26466        );
26467        let t_kv_max = base_len + t; // LAST row's key bound
26468        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
26469        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
26470        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
26471        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
26472        // (parity law), so the partition is freely tunable — verify and decode move together.
26473        if head_dim == 512 {
26474            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26475            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
26476            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
26477            let v = *SP512.get_or_init(|| {
26478                std::env::var("MEMRA_FA_SP512")
26479                    .ok()
26480                    .and_then(|x| x.parse().ok())
26481                    .unwrap_or(0)
26482            });
26483            sp = if v >= 8 {
26484                v
26485            } else {
26486                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
26487            };
26488        }
26489        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
26490        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26491        let gqa = (n_head / n_head_kv).max(1) as u32;
26492        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
26493        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
26494        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
26495        // the different partition changes the combine's FP order (greedy tie flips at depth;
26496        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
26497        // consecutive rows by their OWN ladder value and launch once per group — each row then
26498        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
26499        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
26500        // sp override is t_kv-independent by construction).
26501        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
26502        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
26503            groups.push((0, t, sp));
26504        } else {
26505            let mut r0 = 0usize;
26506            while r0 < t {
26507                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
26508                let mut r1 = r0 + 1;
26509                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
26510                    r1 += 1;
26511                }
26512                groups.push((r0, r1 - r0, sp_g));
26513                r0 = r1;
26514            }
26515        }
26516        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
26517        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
26518        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
26519        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26520        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
26521            std::env::var("MEMRA_FA_SMEM_TKV")
26522                .ok()
26523                .and_then(|v| v.parse().ok())
26524                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
26525        });
26526        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
26527        let v3 = fa_v3_active(head_dim);
26528        let smem_rows =
26529            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
26530        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
26531        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
26532        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
26533        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
26534        let _ = kv_shared;
26535        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
26536        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
26537        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
26538        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
26539        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
26540        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
26541        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
26542        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
26543        // (kv_head, split) stages its tile once and loops the rows over it — kills the
26544        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
26545        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
26546        // shared by every hd512 caller through this wrapper (decode+verify flip together;
26547        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
26548        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
26549        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
26550        // not unpack-bound; jsonl 2026-07-14.
26551        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26552        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
26553        let tb512 = head_dim == 512
26554            && sp <= 32
26555            && n_head / n_head_kv.max(1) <= 16
26556            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
26557        let fname = if tb512 {
26558            "fa_decode_vec_q_rows_v4_512_tb"
26559        } else if i2 {
26560            "fa_decode_vec_q_rows_dpl16_i2"
26561        } else if head_dim == 512 {
26562            "fa_decode_vec_q_rows_dpl16"
26563        }
26564        // gemma globals (parity law)
26565        else if v4 {
26566            "fa_decode_vec_q_rows_v4"
26567        } else if v3 {
26568            "fa_decode_vec_q_rows_v3"
26569        } else if fa_v2_on() {
26570            "fa_decode_vec_q_rows_v2"
26571        } else if smem_rows {
26572            "fa_decode_vec_q_rows_smem"
26573        } else {
26574            "fa_decode_vec_q_rows"
26575        };
26576        let f = if head_dim == 512 {
26577            self.fa_func(fname, head_dim)
26578        } else if g {
26579            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
26580            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
26581            // g-module rows against decode's g-module v4 — different programs, short-VG
26582            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
26583            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
26584            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
26585            // dq macros are format-aware.
26586            self.func_g(if smem_rows {
26587                "fa_decode_vec_q_rows"
26588            } else {
26589                fname
26590            })
26591        } else {
26592            self.func(fname)
26593        };
26594        let shmem = if tb512 {
26595            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
26596            let gk = Self::gkv_on();
26597            let sh =
26598                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
26599            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26600            f.set_attribute(
26601                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26602                sh as i32,
26603            )?;
26604            sh
26605        } else if v4 || v3 || smem_rows || fa_v2_on() {
26606            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
26607            let sh = (if v4 {
26608                11520 + 32 * head_dim * if g { 1 } else { 2 }
26609            } else if v3 {
26610                32 * head_dim * 2
26611            } else {
26612                2 * 32 * head_dim * 2
26613            }) as u32;
26614            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26615            f.set_attribute(
26616                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26617                sh as i32,
26618            )?;
26619            sh
26620        } else {
26621            0
26622        };
26623        // Per-GROUP launches (single group in the common case — identical to the pre-fix
26624        // single launch there): each group gets its own partials (the rows kernel indexes
26625        // partials by its LOCAL grid.z row) and q/o row-offset views.
26626        for &(r0, t_g, sp_g) in &groups {
26627            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
26628            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
26629            let base_i = (base_len + r0) as i32;
26630            let o_len = t_g * n_head * n_splits_g * head_dim;
26631            let ml_len = t_g * n_head * n_splits_g;
26632            let mut part_guard = self.fa_part_pool.lock().unwrap();
26633            if part_guard
26634                .as_ref()
26635                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26636                .unwrap_or(true)
26637            {
26638                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26639                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26640                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26641                // later live allocations land at those addresses, and the next graph REPLAY writes
26642                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26643                // output corruption began the burst after the trunk's t_kv growth first realloc'd
26644                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26645                // the baked addresses alive (single-stream: eager writes the new buffers, replays
26646                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26647                // (total retired < final size).
26648                let old = part_guard.take();
26649                let (co, cm) = old
26650                    .as_ref()
26651                    .map(|pp| (pp.0.len(), pp.1.len()))
26652                    .unwrap_or((0, 0));
26653                if let Some(old) = old {
26654                    self.fa_part_retired.lock().unwrap().push(old);
26655                }
26656                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26657                    eprintln!(
26658                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26659                        co, o_len, cm, ml_len
26660                    );
26661                }
26662                *part_guard =
26663                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
26664            }
26665            let pg = part_guard.as_mut().unwrap();
26666            self.gpu
26667                .stream()
26668                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
26669            self.gpu
26670                .stream()
26671                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
26672            self.gpu
26673                .stream()
26674                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
26675            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
26676            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
26677            let qv = self.view(q, t * n_head * head_dim);
26678            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
26679            let cfg = LaunchConfig {
26680                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
26681                block_dim: (32, gqa, 1),
26682                shared_mem_bytes: shmem,
26683            };
26684            {
26685                let __s_b = self.gpu.stream();
26686                let mut b = __s_b.launch_builder(&f);
26687                if tb512 {
26688                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
26689                    let (bd, plus) =
26690                        base_dev.expect("hd512 rows twin requires a device base counter");
26691                    let plus_g = plus + r0 as i32;
26692                    let nr = t_g as i32;
26693                    if Self::pdl_on() && Self::pdl_wb_on() {
26694                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
26695                        use cudarc::driver::{DevicePtr, DevicePtrMut};
26696                        let s = &self.gpu.stream();
26697                        let (pq, _b0) = q_g.device_ptr(s);
26698                        let (pk, _b1) = k.device_ptr(s);
26699                        let (pv, _b2) = v.device_ptr(s);
26700                        let (po, _b3) = part_o.device_ptr_mut(s);
26701                        let (pm, _b4) = part_m.device_ptr_mut(s);
26702                        let (pl, _b5) = part_l.device_ptr_mut(s);
26703                        let (pb, _b6) = bd.device_ptr(s);
26704                        let mut ps = [
26705                            &pq as *const _ as *mut std::ffi::c_void,
26706                            &pk as *const _ as *mut _,
26707                            &pv as *const _ as *mut _,
26708                            &po as *const _ as *mut _,
26709                            &pm as *const _ as *mut _,
26710                            &pl as *const _ as *mut _,
26711                            &hd as *const _ as *mut _,
26712                            &nh as *const _ as *mut _,
26713                            &nhkv as *const _ as *mut _,
26714                            &pb as *const _ as *mut _,
26715                            &plus_g as *const _ as *mut _,
26716                            &scale as *const _ as *mut _,
26717                            &nspm as *const _ as *mut _,
26718                            &spk as *const _ as *mut _,
26719                            &ktb as *const _ as *mut _,
26720                            &vtb as *const _ as *mut _,
26721                            &nr as *const _ as *mut _,
26722                        ];
26723                        unsafe {
26724                            self.launch_pdl_flash(
26725                                Self::gkv_on(),
26726                                "fa_decode_vec_q_rows_v4_512_tb",
26727                                (n_head_kv as u32, n_splits_g as u32, 1),
26728                                (32, gqa, 1),
26729                                shmem,
26730                                &mut ps,
26731                            )?;
26732                        }
26733                    } else {
26734                        let cfg_tb = LaunchConfig {
26735                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
26736                            block_dim: (32, gqa, 1),
26737                            shared_mem_bytes: shmem,
26738                        };
26739                        b.arg(&q_g)
26740                            .arg(k)
26741                            .arg(v)
26742                            .arg(&mut *part_o)
26743                            .arg(&mut *part_m)
26744                            .arg(&mut *part_l)
26745                            .arg(&hd)
26746                            .arg(&nh)
26747                            .arg(&nhkv)
26748                            .arg(bd)
26749                            .arg(&plus_g)
26750                            .arg(&scale)
26751                            .arg(&nspm)
26752                            .arg(&spk)
26753                            .arg(&ktb)
26754                            .arg(&vtb)
26755                            .arg(&nr);
26756                        unsafe {
26757                            b.launch(cfg_tb)?;
26758                        }
26759                    }
26760                } else if head_dim == 512 {
26761                    let (bd, plus) =
26762                        base_dev.expect("hd512 rows twin requires a device base counter");
26763                    let plus_g = plus + r0 as i32;
26764                    b.arg(&q_g)
26765                        .arg(k)
26766                        .arg(v)
26767                        .arg(&mut *part_o)
26768                        .arg(&mut *part_m)
26769                        .arg(&mut *part_l)
26770                        .arg(&hd)
26771                        .arg(&nh)
26772                        .arg(&nhkv)
26773                        .arg(bd)
26774                        .arg(&plus_g)
26775                        .arg(&scale)
26776                        .arg(&nspm)
26777                        .arg(&spk)
26778                        .arg(&ktb)
26779                        .arg(&vtb);
26780                    unsafe {
26781                        b.launch(cfg)?;
26782                    }
26783                } else {
26784                    b.arg(&q_g)
26785                        .arg(k)
26786                        .arg(v)
26787                        .arg(&mut *part_o)
26788                        .arg(&mut *part_m)
26789                        .arg(&mut *part_l)
26790                        .arg(&hd)
26791                        .arg(&nh)
26792                        .arg(&nhkv)
26793                        .arg(&base_i)
26794                        .arg(&scale)
26795                        .arg(&nspm)
26796                        .arg(&spk)
26797                        .arg(&ktb)
26798                        .arg(&vtb);
26799                    unsafe {
26800                        b.launch(cfg)?;
26801                    }
26802                }
26803            }
26804            let cfg2 = LaunchConfig {
26805                grid_dim: (n_head as u32, t_g as u32, 1),
26806                block_dim: (head_dim as u32, 1, 1),
26807                shared_mem_bytes: 0,
26808            };
26809            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
26810            if head_dim == 512 {
26811                // device-len combine (shared by verify/eager/graph — parity by symbol): the
26812                // per-row n_splits derives from the SAME counter the rows kernel read.
26813                let (bd, plus) = base_dev.unwrap();
26814                let plus_g = plus + r0 as i32;
26815                if let Some((oq, od)) = q8_out.as_mut() {
26816                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
26817                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
26818                    if Self::pdl_on() && Self::pdl_wb_on() {
26819                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
26820                        use cudarc::driver::{DevicePtr, DevicePtrMut};
26821                        let s = &self.gpu.stream();
26822                        let (po, _g0) = part_o.device_ptr(s);
26823                        let (pm, _g1) = part_m.device_ptr(s);
26824                        let (pl, _g2) = part_l.device_ptr(s);
26825                        let (pq, _g3) = oq.device_ptr_mut(s);
26826                        let (pd, _g4) = od.device_ptr_mut(s);
26827                        let (pb, _g5) = bd.device_ptr(s);
26828                        let mut ps = [
26829                            &po as *const _ as *mut std::ffi::c_void,
26830                            &pm as *const _ as *mut _,
26831                            &pl as *const _ as *mut _,
26832                            &pq as *const _ as *mut _,
26833                            &pd as *const _ as *mut _,
26834                            &hd as *const _ as *mut _,
26835                            &nh as *const _ as *mut _,
26836                            &pb as *const _ as *mut _,
26837                            &plus_g as *const _ as *mut _,
26838                            &nspm as *const _ as *mut _,
26839                            &spk as *const _ as *mut _,
26840                        ];
26841                        unsafe {
26842                            self.launch_pdl_flash(
26843                                Self::gkv_on(),
26844                                "fa_decode_combine_rows_dc_q8_1",
26845                                cfg2.grid_dim,
26846                                cfg2.block_dim,
26847                                0,
26848                                &mut ps,
26849                            )?;
26850                        }
26851                        continue;
26852                    }
26853                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
26854                    let __s_b2 = self.gpu.stream();
26855                    let mut b2 = __s_b2.launch_builder(&fc);
26856                    b2.arg(&*part_o)
26857                        .arg(&*part_m)
26858                        .arg(&*part_l)
26859                        .arg(&mut **oq)
26860                        .arg(&mut **od)
26861                        .arg(&hd)
26862                        .arg(&nh)
26863                        .arg(bd)
26864                        .arg(&plus_g)
26865                        .arg(&nspm)
26866                        .arg(&spk);
26867                    unsafe {
26868                        b2.launch(cfg2)?;
26869                    }
26870                    continue;
26871                }
26872                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
26873                let __s_b2 = self.gpu.stream();
26874                let mut b2 = __s_b2.launch_builder(&fc);
26875                b2.arg(&*part_o)
26876                    .arg(&*part_m)
26877                    .arg(&*part_l)
26878                    .arg(&mut o_g)
26879                    .arg(&hd)
26880                    .arg(&nh)
26881                    .arg(bd)
26882                    .arg(&plus_g)
26883                    .arg(&nspm)
26884                    .arg(&spk);
26885                unsafe {
26886                    b2.launch(cfg2)?;
26887                }
26888            } else {
26889                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
26890                // leave the caller's pair unwritten (consumer would read garbage).
26891                assert!(
26892                    q8_out.is_none(),
26893                    "rows q8 emit requires the hd512 dc combine"
26894                );
26895                let fc = self.func("fa_decode_combine_rows");
26896                let __s_b2 = self.gpu.stream();
26897                let mut b2 = __s_b2.launch_builder(&fc);
26898                b2.arg(&*part_o)
26899                    .arg(&*part_m)
26900                    .arg(&*part_l)
26901                    .arg(&mut o_g)
26902                    .arg(&hd)
26903                    .arg(&nh)
26904                    .arg(&base_i)
26905                    .arg(&nspm)
26906                    .arg(&spk);
26907                unsafe {
26908                    b2.launch(cfg2)?;
26909                }
26910            }
26911        }
26912        Ok(())
26913    }
26914
26915    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
26916    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
26917    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
26918    #[allow(clippy::too_many_arguments)]
26919    pub fn fa_decode_rows_w(
26920        &self,
26921        q: &CudaSlice<f32>,
26922        k: &cudarc::driver::CudaView<u8>,
26923        v: &cudarc::driver::CudaView<u8>,
26924        o: &mut CudaSlice<f32>,
26925        head_dim: usize,
26926        n_head: usize,
26927        n_head_kv: usize,
26928        base_dev: &CudaSlice<i32>,
26929        base_plus: i32,
26930        t: usize,
26931        scale: f32,
26932        window: usize,
26933        k_tok_bytes: usize,
26934        v_tok_bytes: usize,
26935        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
26936    ) -> Result<(), Box<dyn std::error::Error>> {
26937        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
26938        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
26939        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
26940        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
26941        debug_assert!(head_dim == 256);
26942        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
26943        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
26944        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
26945        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
26946        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
26947        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
26948        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
26949        let sp = {
26950            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26951            let v = *SPW.get_or_init(|| {
26952                std::env::var("MEMRA_FA_SPW")
26953                    .ok()
26954                    .and_then(|x| x.parse().ok())
26955                    .unwrap_or(0)
26956            });
26957            if v >= 8 {
26958                v
26959            } else {
26960                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
26961            }
26962        };
26963        #[allow(clippy::manual_div_ceil)]
26964        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26965        let n_splits_max = (window + sp - 1) / sp;
26966        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
26967        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
26968        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26969        let gqa = (n_head / n_head_kv).max(1) as u32;
26970        let o_len = t * n_head * n_splits_max * head_dim;
26971        let ml_len = t * n_head * n_splits_max;
26972        let mut part_guard = self.fa_part_pool.lock().unwrap();
26973        if part_guard
26974            .as_ref()
26975            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
26976            .unwrap_or(true)
26977        {
26978            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
26979            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
26980            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
26981            // later live allocations land at those addresses, and the next graph REPLAY writes
26982            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
26983            // output corruption began the burst after the trunk's t_kv growth first realloc'd
26984            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
26985            // the baked addresses alive (single-stream: eager writes the new buffers, replays
26986            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
26987            // (total retired < final size).
26988            let old = part_guard.take();
26989            let (co, cm) = old
26990                .as_ref()
26991                .map(|pp| (pp.0.len(), pp.1.len()))
26992                .unwrap_or((0, 0));
26993            if let Some(old) = old {
26994                self.fa_part_retired.lock().unwrap().push(old);
26995            }
26996            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
26997                eprintln!(
26998                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
26999                    co, o_len, cm, ml_len
27000                );
27001            }
27002            *part_guard =
27003                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27004        }
27005        let pg = part_guard.as_mut().unwrap();
27006        self.gpu
27007            .stream()
27008            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27009        self.gpu
27010            .stream()
27011            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27012        self.gpu
27013            .stream()
27014            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27015        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27016        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
27017        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
27018        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
27019        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
27020        // floor (deep-ctx broadcast win); register twin between.
27021        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
27022        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
27023            std::env::var("MEMRA_FA_SMEM_TKV")
27024                .ok()
27025                .and_then(|v| v.parse().ok())
27026                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
27027        });
27028        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
27029        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
27030        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
27031        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
27032        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
27033        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27034        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
27035        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
27036        // per (lane, format-module) keeps parity structural; the old register-i2 detour
27037        // (-33%) is retired.
27038        let wg = Self::wkv_on();
27039        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
27040        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
27041        let sp2 =
27042            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
27043        if sp2 {
27044            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
27045            if Self::pdl_on() && Self::pdl_wb_on() {
27046                // wave-B2b: flavor mirrors wg.
27047                use cudarc::driver::{DevicePtr, DevicePtrMut};
27048                let s = &self.gpu.stream();
27049                let (pq, _b0) = q.device_ptr(s);
27050                let (pk, _b1) = k.device_ptr(s);
27051                let (pv, _b2) = v.device_ptr(s);
27052                let (po, _b3) = part_o.device_ptr_mut(s);
27053                let (pm, _b4) = part_m.device_ptr_mut(s);
27054                let (pl, _b5) = part_l.device_ptr_mut(s);
27055                let (pb, _b6) = base_dev.device_ptr(s);
27056                let mut ps = [
27057                    &pq as *const _ as *mut std::ffi::c_void,
27058                    &pk as *const _ as *mut _,
27059                    &pv as *const _ as *mut _,
27060                    &po as *const _ as *mut _,
27061                    &pm as *const _ as *mut _,
27062                    &pl as *const _ as *mut _,
27063                    &hd as *const _ as *mut _,
27064                    &nh as *const _ as *mut _,
27065                    &nhkv as *const _ as *mut _,
27066                    &pb as *const _ as *mut _,
27067                    &base_plus as *const _ as *mut _,
27068                    &scale as *const _ as *mut _,
27069                    &nspm as *const _ as *mut _,
27070                    &spk as *const _ as *mut _,
27071                    &ktb as *const _ as *mut _,
27072                    &vtb as *const _ as *mut _,
27073                    &wini as *const _ as *mut _,
27074                ];
27075                unsafe {
27076                    self.launch_pdl_flash(
27077                        wg,
27078                        "fa_decode_vec_q_rows_v4_w_sp",
27079                        (n_head_kv as u32, n_splits_max as u32, t as u32),
27080                        (32, gqa + 1, 1),
27081                        sh,
27082                        &mut ps,
27083                    )?;
27084                }
27085            } else {
27086                let f = if wg {
27087                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
27088                } else {
27089                    self.func("fa_decode_vec_q_rows_v4_w_sp")
27090                };
27091                f.set_attribute(
27092                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27093                    sh as i32,
27094                )?;
27095                let cfg = LaunchConfig {
27096                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27097                    block_dim: (32, gqa + 1, 1),
27098                    shared_mem_bytes: sh,
27099                };
27100                let __s_b = self.gpu.stream();
27101                let mut b = __s_b.launch_builder(&f);
27102                b.arg(q)
27103                    .arg(k)
27104                    .arg(v)
27105                    .arg(&mut *part_o)
27106                    .arg(&mut *part_m)
27107                    .arg(&mut *part_l)
27108                    .arg(&hd)
27109                    .arg(&nh)
27110                    .arg(&nhkv)
27111                    .arg(base_dev)
27112                    .arg(&base_plus)
27113                    .arg(&scale)
27114                    .arg(&nspm)
27115                    .arg(&spk)
27116                    .arg(&ktb)
27117                    .arg(&vtb)
27118                    .arg(&wini);
27119                unsafe {
27120                    b.launch(cfg)?;
27121                }
27122            }
27123        } else {
27124            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
27125                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
27126                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
27127                use cudarc::driver::{DevicePtr, DevicePtrMut};
27128                let s = &self.gpu.stream();
27129                let (pq, _b0) = q.device_ptr(s);
27130                let (pk, _b1) = k.device_ptr(s);
27131                let (pv, _b2) = v.device_ptr(s);
27132                let (po, _b3) = part_o.device_ptr_mut(s);
27133                let (pm, _b4) = part_m.device_ptr_mut(s);
27134                let (pl, _b5) = part_l.device_ptr_mut(s);
27135                let (pb, _b6) = base_dev.device_ptr(s);
27136                let mut ps = [
27137                    &pq as *const _ as *mut std::ffi::c_void,
27138                    &pk as *const _ as *mut _,
27139                    &pv as *const _ as *mut _,
27140                    &po as *const _ as *mut _,
27141                    &pm as *const _ as *mut _,
27142                    &pl as *const _ as *mut _,
27143                    &hd as *const _ as *mut _,
27144                    &nh as *const _ as *mut _,
27145                    &nhkv as *const _ as *mut _,
27146                    &pb as *const _ as *mut _,
27147                    &base_plus as *const _ as *mut _,
27148                    &scale as *const _ as *mut _,
27149                    &nspm as *const _ as *mut _,
27150                    &spk as *const _ as *mut _,
27151                    &ktb as *const _ as *mut _,
27152                    &vtb as *const _ as *mut _,
27153                    &wini as *const _ as *mut _,
27154                ];
27155                unsafe {
27156                    self.launch_pdl_flash(
27157                        wg,
27158                        "fa_decode_vec_q_rows_v4_w",
27159                        (n_head_kv as u32, n_splits_max as u32, t as u32),
27160                        (32, gqa, 1),
27161                        sh,
27162                        &mut ps,
27163                    )?;
27164                }
27165            } else {
27166                let pick = |name: &str| {
27167                    if wg {
27168                        self.func_g(name)
27169                    } else {
27170                        self.func(name)
27171                    }
27172                };
27173                let (f, sh) = if fa_v4_at(window) {
27174                    let f = pick("fa_decode_vec_q_rows_v4_w");
27175                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
27176                } else if smem_tkv > 0 && window >= smem_tkv {
27177                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
27178                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
27179                    (
27180                        pick("fa_decode_vec_q_rows_smem_w"),
27181                        (2 * 32 * head_dim * 2) as u32,
27182                    )
27183                } else {
27184                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
27185                };
27186                f.set_attribute(
27187                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27188                    sh as i32,
27189                )?;
27190                let cfg = LaunchConfig {
27191                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27192                    block_dim: (32, gqa, 1),
27193                    shared_mem_bytes: sh,
27194                };
27195                let __s_b = self.gpu.stream();
27196                let mut b = __s_b.launch_builder(&f);
27197                b.arg(q)
27198                    .arg(k)
27199                    .arg(v)
27200                    .arg(&mut *part_o)
27201                    .arg(&mut *part_m)
27202                    .arg(&mut *part_l)
27203                    .arg(&hd)
27204                    .arg(&nh)
27205                    .arg(&nhkv)
27206                    .arg(base_dev)
27207                    .arg(&base_plus)
27208                    .arg(&scale)
27209                    .arg(&nspm)
27210                    .arg(&spk)
27211                    .arg(&ktb)
27212                    .arg(&vtb)
27213                    .arg(&wini);
27214                unsafe {
27215                    b.launch(cfg)?;
27216                }
27217            }
27218        }
27219        let cfg2 = LaunchConfig {
27220            grid_dim: (n_head as u32, t as u32, 1),
27221            block_dim: (head_dim as u32, 1, 1),
27222            shared_mem_bytes: 0,
27223        };
27224        if let Some((oq, od)) = q8_out {
27225            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
27226            // consumes the pair directly; the standalone quantize launch folds away.
27227            if Self::pdl_on() && Self::pdl_wb_on() {
27228                // wave-B2: flavor mirrors the builder's wg choice.
27229                use cudarc::driver::{DevicePtr, DevicePtrMut};
27230                let s = &self.gpu.stream();
27231                let (po, _g0) = part_o.device_ptr(s);
27232                let (pm, _g1) = part_m.device_ptr(s);
27233                let (pl, _g2) = part_l.device_ptr(s);
27234                let (pq, _g3) = oq.device_ptr_mut(s);
27235                let (pd, _g4) = od.device_ptr_mut(s);
27236                let mut ps = [
27237                    &po as *const _ as *mut std::ffi::c_void,
27238                    &pm as *const _ as *mut _,
27239                    &pl as *const _ as *mut _,
27240                    &pq as *const _ as *mut _,
27241                    &pd as *const _ as *mut _,
27242                    &hd as *const _ as *mut _,
27243                    &nh as *const _ as *mut _,
27244                    &nspm as *const _ as *mut _,
27245                    &spk as *const _ as *mut _,
27246                    &wini as *const _ as *mut _,
27247                ];
27248                unsafe {
27249                    self.launch_pdl_flash(
27250                        wg,
27251                        "fa_decode_combine_rows_w_q8_1",
27252                        cfg2.grid_dim,
27253                        cfg2.block_dim,
27254                        0,
27255                        &mut ps,
27256                    )?;
27257                }
27258                return Ok(());
27259            }
27260            let fc = if wg {
27261                self.func_g("fa_decode_combine_rows_w_q8_1")
27262            } else {
27263                self.func("fa_decode_combine_rows_w_q8_1")
27264            };
27265            let __s_b2 = self.gpu.stream();
27266            let mut b2 = __s_b2.launch_builder(&fc);
27267            b2.arg(&*part_o)
27268                .arg(&*part_m)
27269                .arg(&*part_l)
27270                .arg(oq)
27271                .arg(od)
27272                .arg(&hd)
27273                .arg(&nh)
27274                .arg(&nspm)
27275                .arg(&spk)
27276                .arg(&wini);
27277            unsafe {
27278                b2.launch(cfg2)?;
27279            }
27280            return Ok(());
27281        }
27282        let fc = if wg {
27283            self.func_g("fa_decode_combine_rows_w")
27284        } else {
27285            self.func("fa_decode_combine_rows_w")
27286        };
27287        let __s_b2 = self.gpu.stream();
27288        let mut b2 = __s_b2.launch_builder(&fc);
27289        b2.arg(&*part_o)
27290            .arg(&*part_m)
27291            .arg(&*part_l)
27292            .arg(o)
27293            .arg(&hd)
27294            .arg(&nh)
27295            .arg(&nspm)
27296            .arg(&spk)
27297            .arg(&wini);
27298        unsafe {
27299            b2.launch(cfg2)?;
27300        }
27301        Ok(())
27302    }
27303
27304    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
27305    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
27306    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
27307    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
27308    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
27309    #[allow(clippy::too_many_arguments)]
27310    pub fn fa_decode_rows_dc(
27311        &self,
27312        q: &CudaSlice<f32>,
27313        k: &cudarc::driver::CudaView<u8>,
27314        v: &cudarc::driver::CudaView<u8>,
27315        o: &mut CudaSlice<f32>,
27316        head_dim: usize,
27317        n_head: usize,
27318        n_head_kv: usize,
27319        base_dev: &CudaSlice<i32>,
27320        t_kv_upper: usize,
27321        t: usize,
27322        scale: f32,
27323        k_tok_bytes: usize,
27324        v_tok_bytes: usize,
27325        base_plus: i32,
27326        g: bool,
27327    ) -> Result<(), Box<dyn std::error::Error>> {
27328        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
27329        assert!(
27330            v4 || fa_v3_active(head_dim),
27331            "stream fa rows requires the v3 or v4 lane"
27332        );
27333        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
27334        if v4 {
27335            let sp = fa_split_keys(t_kv_upper, n_head_kv);
27336            #[allow(clippy::manual_div_ceil)]
27337            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27338            let n_splits_max = (t_kv_upper + sp - 1) / sp;
27339            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27340            let (nspm, spk) = (n_splits_max as i32, sp as i32);
27341            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27342            let gqa = (n_head / n_head_kv).max(1) as u32;
27343            let o_len = t * n_head * n_splits_max * head_dim;
27344            let ml_len = t * n_head * n_splits_max;
27345            let mut part_guard = self.fa_part_pool.lock().unwrap();
27346            if part_guard
27347                .as_ref()
27348                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27349                .unwrap_or(true)
27350            {
27351                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27352                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27353                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27354                // later live allocations land at those addresses, and the next graph REPLAY writes
27355                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27356                // output corruption began the burst after the trunk's t_kv growth first realloc'd
27357                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27358                // the baked addresses alive (single-stream: eager writes the new buffers, replays
27359                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27360                // (total retired < final size).
27361                let old = part_guard.take();
27362                let (co, cm) = old
27363                    .as_ref()
27364                    .map(|pp| (pp.0.len(), pp.1.len()))
27365                    .unwrap_or((0, 0));
27366                if let Some(old) = old {
27367                    self.fa_part_retired.lock().unwrap().push(old);
27368                }
27369                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27370                    eprintln!(
27371                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27372                        co, o_len, cm, ml_len
27373                    );
27374                }
27375                *part_guard =
27376                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27377            }
27378            let pg = part_guard.as_mut().unwrap();
27379            self.gpu
27380                .stream()
27381                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27382            self.gpu
27383                .stream()
27384                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27385            self.gpu
27386                .stream()
27387                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27388            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27389            let f = if g {
27390                self.func_g("fa_decode_vec_q_rows_v4_dc")
27391            } else {
27392                self.func("fa_decode_vec_q_rows_v4_dc")
27393            };
27394            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
27395            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27396            f.set_attribute(
27397                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27398                sh as i32,
27399            )?;
27400            let cfg = LaunchConfig {
27401                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27402                block_dim: (32, gqa, 1),
27403                shared_mem_bytes: sh,
27404            };
27405            let __s_b = self.gpu.stream();
27406            let mut b = __s_b.launch_builder(&f);
27407            b.arg(q)
27408                .arg(k)
27409                .arg(v)
27410                .arg(&mut *part_o)
27411                .arg(&mut *part_m)
27412                .arg(&mut *part_l)
27413                .arg(&hd)
27414                .arg(&nh)
27415                .arg(&nhkv)
27416                .arg(base_dev)
27417                .arg(&base_plus)
27418                .arg(&scale)
27419                .arg(&nspm)
27420                .arg(&spk)
27421                .arg(&ktb)
27422                .arg(&vtb);
27423            unsafe {
27424                b.launch(cfg)?;
27425            }
27426            let fc = self.func("fa_decode_combine_rows_dc");
27427            let cfg2 = LaunchConfig {
27428                grid_dim: (n_head as u32, t as u32, 1),
27429                block_dim: (head_dim as u32, 1, 1),
27430                shared_mem_bytes: 0,
27431            };
27432            let __s_b2 = self.gpu.stream();
27433            let mut b2 = __s_b2.launch_builder(&fc);
27434            b2.arg(&*part_o)
27435                .arg(&*part_m)
27436                .arg(&*part_l)
27437                .arg(o)
27438                .arg(&hd)
27439                .arg(&nh)
27440                .arg(base_dev)
27441                .arg(&base_plus)
27442                .arg(&nspm)
27443                .arg(&spk);
27444            unsafe {
27445                b2.launch(cfg2)?;
27446            }
27447            return Ok(());
27448        }
27449        let sp = fa_split_keys(t_kv_upper, n_head_kv);
27450        #[allow(clippy::manual_div_ceil)]
27451        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27452        let n_splits_max = (t_kv_upper + sp - 1) / sp;
27453        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27454        let (nspm, spk) = (n_splits_max as i32, sp as i32);
27455        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27456        let gqa = (n_head / n_head_kv).max(1) as u32;
27457        let o_len = t * n_head * n_splits_max * head_dim;
27458        let ml_len = t * n_head * n_splits_max;
27459        let mut part_guard = self.fa_part_pool.lock().unwrap();
27460        if part_guard
27461            .as_ref()
27462            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27463            .unwrap_or(true)
27464        {
27465            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27466            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27467            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27468            // later live allocations land at those addresses, and the next graph REPLAY writes
27469            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27470            // output corruption began the burst after the trunk's t_kv growth first realloc'd
27471            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27472            // the baked addresses alive (single-stream: eager writes the new buffers, replays
27473            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27474            // (total retired < final size).
27475            let old = part_guard.take();
27476            let (co, cm) = old
27477                .as_ref()
27478                .map(|pp| (pp.0.len(), pp.1.len()))
27479                .unwrap_or((0, 0));
27480            if let Some(old) = old {
27481                self.fa_part_retired.lock().unwrap().push(old);
27482            }
27483            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27484                eprintln!(
27485                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27486                    co, o_len, cm, ml_len
27487                );
27488            }
27489            *part_guard =
27490                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27491        }
27492        let pg = part_guard.as_mut().unwrap();
27493        self.gpu
27494            .stream()
27495            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27496        self.gpu
27497            .stream()
27498            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27499        self.gpu
27500            .stream()
27501            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27502        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27503        let f = self.func("fa_decode_vec_q_rows_v3_dc");
27504        let sh = (32 * head_dim * 2) as u32;
27505        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27506        f.set_attribute(
27507            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27508            sh as i32,
27509        )?;
27510        let cfg = LaunchConfig {
27511            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
27512            block_dim: (32, gqa, 1),
27513            shared_mem_bytes: sh,
27514        };
27515        let __s_b = self.gpu.stream();
27516        let mut b = __s_b.launch_builder(&f);
27517        b.arg(q)
27518            .arg(k)
27519            .arg(v)
27520            .arg(&mut *part_o)
27521            .arg(&mut *part_m)
27522            .arg(&mut *part_l)
27523            .arg(&hd)
27524            .arg(&nh)
27525            .arg(&nhkv)
27526            .arg(base_dev)
27527            .arg(&scale)
27528            .arg(&nspm)
27529            .arg(&spk)
27530            .arg(&ktb)
27531            .arg(&vtb);
27532        unsafe {
27533            b.launch(cfg)?;
27534        }
27535        let fc = self.func("fa_decode_combine_rows_dc");
27536        let cfg2 = LaunchConfig {
27537            grid_dim: (n_head as u32, t as u32, 1),
27538            block_dim: (head_dim as u32, 1, 1),
27539            shared_mem_bytes: 0,
27540        };
27541        let plus0 = 0i32;
27542        let __s_b2 = self.gpu.stream();
27543        let mut b2 = __s_b2.launch_builder(&fc);
27544        b2.arg(&*part_o)
27545            .arg(&*part_m)
27546            .arg(&*part_l)
27547            .arg(o)
27548            .arg(&hd)
27549            .arg(&nh)
27550            .arg(base_dev)
27551            .arg(&plus0)
27552            .arg(&nspm)
27553            .arg(&spk);
27554        unsafe {
27555            b2.launch(cfg2)?;
27556        }
27557        Ok(())
27558    }
27559
27560    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
27561    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
27562    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
27563    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
27564    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
27565    ///
27566    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
27567    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
27568    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
27569    /// grouping (different but mathematically-equal log-sum-exp merge).
27570    #[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
27571    pub fn fa_decode_dc(
27572        &self,
27573        q: &CudaSlice<f32>,
27574        k: &cudarc::driver::CudaView<u8>,
27575        v: &cudarc::driver::CudaView<u8>,
27576        o: &mut CudaSlice<f32>,
27577        head_dim: usize,
27578        n_head: usize,
27579        n_head_kv: usize,
27580        t_kv_dev: &CudaSlice<i32>,
27581        bucket_max: usize,
27582        scale: f32,
27583        k_tok_bytes: usize,
27584        v_tok_bytes: usize,
27585        g: bool,
27586    ) -> Result<(), Box<dyn std::error::Error>> {
27587        self.fa_decode_dc_q8(
27588            q,
27589            k,
27590            v,
27591            o,
27592            head_dim,
27593            n_head,
27594            n_head_kv,
27595            t_kv_dev,
27596            bucket_max,
27597            scale,
27598            k_tok_bytes,
27599            v_tok_bytes,
27600            g,
27601            None,
27602        )
27603    }
27604
27605    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
27606    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
27607    #[allow(clippy::too_many_arguments)]
27608    #[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
27609    pub fn fa_decode_dc_q8(
27610        &self,
27611        q: &CudaSlice<f32>,
27612        k: &cudarc::driver::CudaView<u8>,
27613        v: &cudarc::driver::CudaView<u8>,
27614        o: &mut CudaSlice<f32>,
27615        head_dim: usize,
27616        n_head: usize,
27617        n_head_kv: usize,
27618        t_kv_dev: &CudaSlice<i32>,
27619        bucket_max: usize,
27620        scale: f32,
27621        k_tok_bytes: usize,
27622        v_tok_bytes: usize,
27623        g: bool,
27624        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
27625    ) -> Result<(), Box<dyn std::error::Error>> {
27626        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
27627        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
27628        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
27629        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
27630        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
27631        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
27632        // 2026-07-12).
27633        let mut fa_vec =
27634            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
27635        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
27636            fa_vec = false;
27637        } // mirror kvmod/geom
27638        let sp = fa_split_keys(bucket_max, n_head_kv);
27639        let n_splits = if fa_vec {
27640            ((bucket_max + sp - 1) / sp).max(1)
27641        } else {
27642            ((bucket_max + 255) / 256).max(1)
27643        };
27644        let o_len = n_head * n_splits * head_dim;
27645        let ml_len = n_head * n_splits;
27646        let mut part_guard = self.fa_part_pool.lock().unwrap();
27647        if part_guard
27648            .as_ref()
27649            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
27650            .unwrap_or(true)
27651        {
27652            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
27653            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
27654            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
27655            // later live allocations land at those addresses, and the next graph REPLAY writes
27656            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
27657            // output corruption began the burst after the trunk's t_kv growth first realloc'd
27658            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
27659            // the baked addresses alive (single-stream: eager writes the new buffers, replays
27660            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
27661            // (total retired < final size).
27662            let old = part_guard.take();
27663            let (co, cm) = old
27664                .as_ref()
27665                .map(|pp| (pp.0.len(), pp.1.len()))
27666                .unwrap_or((0, 0));
27667            if let Some(old) = old {
27668                self.fa_part_retired.lock().unwrap().push(old);
27669            }
27670            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
27671                eprintln!(
27672                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
27673                    co, o_len, cm, ml_len
27674                );
27675            }
27676            *part_guard =
27677                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
27678        }
27679        let pg = part_guard.as_mut().unwrap();
27680        self.gpu
27681            .stream()
27682            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
27683        self.gpu
27684            .stream()
27685            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
27686        self.gpu
27687            .stream()
27688            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
27689        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
27690        let (hd, nh, nhkv, nsp) = (
27691            head_dim as i32,
27692            n_head as i32,
27693            n_head_kv as i32,
27694            n_splits as i32,
27695        );
27696        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27697        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
27698        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
27699        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
27700        let deep = fa_vec
27701            && head_dim == 256
27702            && fa_v4_at(bucket_max)
27703            && !g
27704            && fa_deep_at(bucket_max)
27705            && !matches!(fa_v4_mode(), "noB3" | "stage");
27706        let (f, cfg) = if fa_vec
27707            && head_dim == 512
27708            && bucket_max >= {
27709                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
27710                *FA512_MIN_DC.get_or_init(|| {
27711                    std::env::var("MEMRA_FA512_MIN")
27712                        .ok()
27713                        .and_then(|v| v.parse().ok())
27714                        .unwrap_or(512)
27715                })
27716            } {
27717            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
27718            let gqa = (n_head / n_head_kv).max(1) as u32;
27719            (
27720                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
27721                LaunchConfig {
27722                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27723                    block_dim: (32, gqa, 1),
27724                    shared_mem_bytes: 0,
27725                },
27726            )
27727        } else if fa_vec && head_dim == 512 {
27728            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
27729            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
27730            let q_view = q.as_view();
27731            let mut o_view = o.as_view_mut();
27732            return self.fa_decode_scalar_unified(
27733                &q_view,
27734                k,
27735                v,
27736                &mut o_view,
27737                head_dim,
27738                n_head,
27739                n_head_kv,
27740                0,
27741                Some(t_kv_dev),
27742                scale,
27743                n_splits,
27744                sp,
27745                k_tok_bytes,
27746                v_tok_bytes,
27747                g,
27748                &mut *part_o,
27749                &mut *part_m,
27750                &mut *part_l,
27751                q8_out,
27752            );
27753        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
27754            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
27755            // incl the g-module route + raw-e4m3 sV sizing.
27756            let gqa = (n_head / n_head_kv).max(1) as u32;
27757            let fv = if g {
27758                self.func_g("fa_decode_vec_q_v4_dc")
27759            } else if deep {
27760                self.func("fa_decode_vec_q_v4_deep_dc")
27761            } else {
27762                self.func("fa_decode_vec_q_v4_dc")
27763            };
27764            let shmem =
27765                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
27766            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27767            fv.set_attribute(
27768                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27769                shmem as i32,
27770            )?;
27771            (
27772                fv,
27773                LaunchConfig {
27774                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27775                    block_dim: (32, gqa, 1),
27776                    shared_mem_bytes: shmem,
27777                },
27778            )
27779        } else if fa_vec && fa_v3_active(head_dim) {
27780            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
27781            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
27782            let gqa = (n_head / n_head_kv).max(1) as u32;
27783            let fv = if g {
27784                self.func_g("fa_decode_vec_q_v3_dc")
27785            } else {
27786                self.func("fa_decode_vec_q_v3_dc")
27787            };
27788            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
27789            (
27790                fv,
27791                LaunchConfig {
27792                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27793                    block_dim: (32, gqa, 1),
27794                    shared_mem_bytes: shmem,
27795                },
27796            )
27797        } else if fa_vec && fa_v2_on() {
27798            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
27799            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
27800            // a numeric config; eager, rows-verify and graph all switch together).
27801            let gqa = (n_head / n_head_kv).max(1) as u32;
27802            let fv = if g {
27803                self.func_g("fa_decode_vec_q_v2_dc")
27804            } else {
27805                self.func("fa_decode_vec_q_v2_dc")
27806            };
27807            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
27808            (
27809                fv,
27810                LaunchConfig {
27811                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27812                    block_dim: (32, gqa, 1),
27813                    shared_mem_bytes: shmem,
27814                },
27815            )
27816        } else if fa_vec {
27817            let gqa = (n_head / n_head_kv).max(1) as u32;
27818            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
27819            let fv = if g {
27820                self.func_g("fa_decode_vec_q_dc")
27821            } else {
27822                self.func("fa_decode_vec_q_dc")
27823            };
27824            (
27825                fv,
27826                LaunchConfig {
27827                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
27828                    block_dim: (32, gqa, 1),
27829                    shared_mem_bytes: 0,
27830                },
27831            )
27832        } else {
27833            let q_view = q.as_view();
27834            let mut o_view = o.as_view_mut();
27835            return self.fa_decode_scalar_unified(
27836                &q_view,
27837                k,
27838                v,
27839                &mut o_view,
27840                head_dim,
27841                n_head,
27842                n_head_kv,
27843                0,
27844                Some(t_kv_dev),
27845                scale,
27846                n_splits,
27847                if fa_vec { sp } else { 256 },
27848                k_tok_bytes,
27849                v_tok_bytes,
27850                g,
27851                &mut *part_o,
27852                &mut *part_m,
27853                &mut *part_l,
27854                q8_out,
27855            );
27856        };
27857        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
27858        let __s_b = self.gpu.stream();
27859        let mut b = __s_b.launch_builder(&f);
27860        b.arg(q)
27861            .arg(k)
27862            .arg(v)
27863            .arg(&mut *part_o)
27864            .arg(&mut *part_m)
27865            .arg(&mut *part_l)
27866            .arg(&hd)
27867            .arg(&nh)
27868            .arg(&nhkv)
27869            .arg(t_kv_dev)
27870            .arg(&scale)
27871            .arg(&nsp)
27872            .arg(&ski)
27873            .arg(&ktb)
27874            .arg(&vtb);
27875        unsafe {
27876            b.launch(cfg)?;
27877        }
27878        let cfg2 = LaunchConfig {
27879            grid_dim: (n_head as u32, 1, 1),
27880            block_dim: (head_dim as u32, 1, 1),
27881            shared_mem_bytes: 0,
27882        };
27883        if let Some((oq, od)) = q8_out {
27884            let fc = if g {
27885                self.func_g("fa_decode_combine_q8_1")
27886            } else {
27887                self.fa_func("fa_decode_combine_q8_1", head_dim)
27888            };
27889            let __s_b2 = self.gpu.stream();
27890            let mut b2 = __s_b2.launch_builder(&fc);
27891            b2.arg(&*part_o)
27892                .arg(&*part_m)
27893                .arg(&*part_l)
27894                .arg(oq)
27895                .arg(od)
27896                .arg(&hd)
27897                .arg(&nh)
27898                .arg(&nsp);
27899            unsafe {
27900                b2.launch(cfg2)?;
27901            }
27902            return Ok(());
27903        }
27904        let fc = if g {
27905            self.func_g("fa_decode_combine_f32")
27906        } else {
27907            self.fa_func("fa_decode_combine_f32", head_dim)
27908        };
27909        let __s_b2 = self.gpu.stream();
27910        let mut b2 = __s_b2.launch_builder(&fc);
27911        b2.arg(&*part_o)
27912            .arg(&*part_m)
27913            .arg(&*part_l)
27914            .arg(o)
27915            .arg(&hd)
27916            .arg(&nh)
27917            .arg(&nsp);
27918        unsafe {
27919            b2.launch(cfg2)?;
27920        }
27921        Ok(())
27922    }
27923
27924    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
27925    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
27926    /// at equal rows.
27927    #[allow(clippy::too_many_arguments)]
27928    pub fn append_kv_quantized_dcw(
27929        &self,
27930        k_row: &CudaSlice<f32>,
27931        v_row: &CudaSlice<f32>,
27932        kc: &mut CudaSlice<u8>,
27933        vc: &mut CudaSlice<u8>,
27934        len_dev: &CudaSlice<i32>,
27935        base_dev: Option<&CudaSlice<i32>>,
27936        kv_dim_k: usize,
27937        kv_dim_v: usize,
27938        k_tok_bytes: usize,
27939        v_tok_bytes: usize,
27940    ) -> Result<(), Box<dyn std::error::Error>> {
27941        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
27942        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
27943        let cfg = LaunchConfig {
27944            grid_dim: (nblk, 1, 1),
27945            block_dim: (32, 1, 1),
27946            shared_mem_bytes: 0,
27947        };
27948        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
27949        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27950        let null: u64 = 0;
27951        let __s_b = self.gpu.stream();
27952        let mut b = __s_b.launch_builder(&f);
27953        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
27954        match base_dev {
27955            Some(base) => {
27956                b.arg(base);
27957            }
27958            None => {
27959                b.arg(&null);
27960            }
27961        }
27962        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
27963        unsafe {
27964            b.launch(cfg)?;
27965        }
27966        Ok(())
27967    }
27968
27969    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
27970    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
27971        let f = self.func("inc_i32");
27972        let cfg = LaunchConfig {
27973            grid_dim: (1, 1, 1),
27974            block_dim: (1, 1, 1),
27975            shared_mem_bytes: 0,
27976        };
27977        let __s_b = self.gpu.stream();
27978        let mut b = __s_b.launch_builder(&f);
27979        b.arg(counter);
27980        unsafe {
27981            b.launch(cfg)?;
27982        }
27983        Ok(())
27984    }
27985
27986    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
27987    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
27988    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
27989    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
27990    /// kernel class on this lane); callers keep eager below the vec floor and for any other
27991    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
27992    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
27993    /// alive across bucket growth.
27994    #[allow(clippy::too_many_arguments)]
27995    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
27996    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
27997    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
27998    /// THE ONE PLACE THE FA PARTIAL POOL IS ALLOCATED.
27999    ///
28000    /// Eight call sites grow this pool and all eight retire-on-grow correctly, but only ONE
28001    /// of them carried the `[fa-pool] grow` receipt, so that receipt under-reported grows by
28002    /// seven eighths and no grow could honestly be dated against a request. Routing every
28003    /// grower through here makes the count real. The receipt names the site so a ladder can
28004    /// be attributed, and stays bounded so a pathological ladder cannot flood a serving log.
28005    ///
28006    /// `MEMRA_FA_PART_ZERO=1` (DEFAULT OFF, diagnostic only) zeroes the fresh buffers. A grow
28007    /// hands every subsequent launch three UNINITIALIZED banks; if the poison is a combine
28008    /// reading a partial bank its producer never wrote, that makes every row and every head
28009    /// non-finite at once, which is the shape the level-2 bad-row bitmap reports at the
28010    /// global-attention join.
28011    ///
28012    /// READ IT IN ONE DIRECTION ONLY. Zeroed banks carry m = 0.0, not NEG_INF, so the
28013    /// empty-split no-op guard never engages: a bank that is entirely unwritten still
28014    /// combines to L = 0 and O/L = 0/0 = NaN. So **silence under this arm convicts the pool;
28015    /// continued trapping acquits nothing**, because only the PARTIALLY unwritten class (real
28016    /// splits beside stale zeroed ones) goes quiet. Discriminator, never a fix, and never a
28017    /// serving arm: where it does go quiet the output is still wrong, it just looks plausible.
28018    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28019    fn fa_part_alloc(
28020        &self,
28021        o_len: usize,
28022        ml_len: usize,
28023        co: usize,
28024        cm: usize,
28025    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
28026        static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
28027        let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
28028        if n < 64 {
28029            eprintln!(
28030                "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
28031                self.ctx().ordinal(),
28032                fa_part_zero_on()
28033            );
28034        }
28035        let mut po = self.alloc_uninit::<f32>(o_len)?;
28036        let mut pm = self.alloc_uninit::<f32>(ml_len)?;
28037        let mut pl = self.alloc_uninit::<f32>(ml_len)?;
28038        if fa_part_zero_on() {
28039            self.gpu.stream().memset_zeros(&mut po)?;
28040            self.gpu.stream().memset_zeros(&mut pm)?;
28041            self.gpu.stream().memset_zeros(&mut pl)?;
28042        }
28043        Ok((po, pm, pl))
28044    }
28045
28046    fn fa_part_pool_grow(
28047        &self,
28048        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
28049        o_len: usize,
28050        ml_len: usize,
28051    ) -> Result<(), Box<dyn std::error::Error>> {
28052        if part_guard
28053            .as_ref()
28054            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28055            .unwrap_or(true)
28056        {
28057            let old = part_guard.take();
28058            let (co, cm) = old
28059                .as_ref()
28060                .map(|pp| (pp.0.len(), pp.1.len()))
28061                .unwrap_or((0, 0));
28062            if let Some(old) = old {
28063                self.fa_part_retired.lock().unwrap().push(old);
28064            }
28065            // GROW RECEIPT. This pool is grow-only, retires-on-grow and never frees, and every
28066            // FA decode/verify launch in the process reads and writes it. A grow is therefore a
28067            // process-lifetime EVENT — new addresses, a retired buffer kept alive forever, and
28068            // a different partial layout — and it is invisible in every log we have. The step37
28069            // spec fault is clean for the first two or three requests of a process and then
28070            // poisons trunk layer 20 (research: MEMRA_SPEC_NAN_SCAN), which is exactly the
28071            // shape a mid-life pool grow would produce, so the grows have to be datable
28072            // against the requests. Cap raised from 8 after the first run measured FOUR
28073            // grows per device (380928 -> 761856 -> 1523712 -> 3047424): with two devices the
28074            // 8 slots were spent before any grow could be dated against a request, which was
28075            // the entire point of the receipt. Still bounded so a pathological ladder cannot
28076            // flood a serving log.
28077            *part_guard =
28078                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28079        }
28080        Ok(())
28081    }
28082
28083    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
28084    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
28085    pub fn fa_dcw_pool_ensure(
28086        &self,
28087        head_dim: usize,
28088        n_head: usize,
28089        n_head_kv: usize,
28090        bucket_max: usize,
28091    ) -> Result<(), Box<dyn std::error::Error>> {
28092        let sp = fa_split_keys(bucket_max, n_head_kv);
28093        #[allow(clippy::manual_div_ceil)]
28094        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28095        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28096        let o_len = n_head * n_splits * head_dim;
28097        let ml_len = n_head * n_splits;
28098        let mut part_guard = self.fa_part_pool.lock().unwrap();
28099        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
28100    }
28101
28102    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
28103    /// appended; one launch walks the KV stream once with two query rows (per-row causal
28104    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
28105    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
28106    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
28107    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
28108    /// outputs (the head gate fuses into the combine as in the t=1 path).
28109    #[allow(clippy::too_many_arguments)]
28110    pub fn fa_decode_dcw2(
28111        &self,
28112        q2: &CudaSlice<f32>,
28113        k_ring: &cudarc::driver::CudaView<u8>,
28114        v_ring: &cudarc::driver::CudaView<u8>,
28115        o2: &mut CudaSlice<f32>,
28116        head_dim: usize,
28117        n_head: usize,
28118        n_head_kv: usize,
28119        len_dev: &CudaSlice<i32>,
28120        base_dev: Option<&CudaSlice<i32>>,
28121        window: usize,
28122        bucket_max: usize,
28123        scale: f32,
28124        k_tok_bytes: usize,
28125        v_tok_bytes: usize,
28126        gate2: &CudaSlice<f32>,
28127    ) -> Result<(), Box<dyn std::error::Error>> {
28128        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
28129        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
28130            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
28131        }
28132        let sp = fa_split_keys(bucket_max, n_head_kv);
28133        #[allow(clippy::manual_div_ceil)]
28134        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28135        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28136        // Partials for BOTH rows: row-major halves.
28137        let o_len = 2 * n_head * n_splits * head_dim;
28138        let ml_len = 2 * n_head * n_splits;
28139        let mut part_guard = self.fa_part_pool.lock().unwrap();
28140        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28141        let pg = part_guard.as_mut().unwrap();
28142        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28143        let (hd, nh, nhkv, nsp) = (
28144            head_dim as i32,
28145            n_head as i32,
28146            n_head_kv as i32,
28147            n_splits as i32,
28148        );
28149        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28150        let (ski, win) = (sp as i32, window as i32);
28151        let gqa = (n_head / n_head_kv).max(1) as u32;
28152        let smem = (32 * head_dim * 2) as u32;
28153        let f = self.func("fa_decode_vec_q_v3_dcw2");
28154        let cfg = LaunchConfig {
28155            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28156            block_dim: (32, gqa, 1),
28157            shared_mem_bytes: smem,
28158        };
28159        let null: u64 = 0;
28160        {
28161            let __s_b = self.gpu.stream();
28162            let mut b = __s_b.launch_builder(&f);
28163            b.arg(q2)
28164                .arg(k_ring)
28165                .arg(v_ring)
28166                .arg(&mut *part_o)
28167                .arg(&mut *part_m)
28168                .arg(&mut *part_l)
28169                .arg(&hd)
28170                .arg(&nh)
28171                .arg(&nhkv)
28172                .arg(len_dev);
28173            match base_dev {
28174                Some(base) => {
28175                    b.arg(base);
28176                }
28177                None => {
28178                    b.arg(&null);
28179                }
28180            }
28181            b.arg(&win)
28182                .arg(&scale)
28183                .arg(&nsp)
28184                .arg(&ski)
28185                .arg(&ktb)
28186                .arg(&vtb);
28187            unsafe {
28188                b.launch(cfg)?;
28189            }
28190        }
28191        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
28192        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
28193        // one launch covers both rows with the exact t=1 program per (row, head).
28194        let fc = {
28195            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28196            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28197                self.func("fa_decode_combine_gate_f32_s")
28198            } else {
28199                self.func("fa_decode_combine_gate_f32")
28200            }
28201        };
28202        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
28203        let nh2 = (2 * n_head) as i32;
28204        let cfg2 = LaunchConfig {
28205            grid_dim: ((2 * n_head) as u32, 1, 1),
28206            block_dim: (head_dim as u32, 1, 1),
28207            shared_mem_bytes: if combine_shared {
28208                (2 * n_splits * 4) as u32
28209            } else {
28210                0
28211            },
28212        };
28213        let __s_b2 = self.gpu.stream();
28214        let mut b2 = __s_b2.launch_builder(&fc);
28215        b2.arg(&*part_o)
28216            .arg(&*part_m)
28217            .arg(&*part_l)
28218            .arg(gate2)
28219            .arg(o2)
28220            .arg(&hd)
28221            .arg(&nh2)
28222            .arg(&nsp);
28223        unsafe {
28224            b2.launch(cfg2)?;
28225        }
28226        Ok(())
28227    }
28228
28229    /// T-ROW dcw decode attention over a per-row session table (the per-session
28230    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
28231    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
28232    /// program verbatim with that row's ring/len/base and its own split geometry, so each
28233    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
28234    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
28235    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
28236    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
28237    #[allow(clippy::too_many_arguments)]
28238    pub fn fa_decode_dcw_rows(
28239        &self,
28240        q_rows: &CudaSlice<f32>,
28241        tab: &CudaSlice<u64>,
28242        o_rows: &mut CudaSlice<f32>,
28243        t: usize,
28244        head_dim: usize,
28245        n_head: usize,
28246        n_head_kv: usize,
28247        window: usize,
28248        max_ns: usize,
28249        scale: f32,
28250        k_tok_bytes: usize,
28251        v_tok_bytes: usize,
28252        gate_rows: &CudaSlice<f32>,
28253    ) -> Result<(), Box<dyn std::error::Error>> {
28254        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
28255            || head_dim > 256
28256            || !head_dim.is_multiple_of(32)
28257            || !fa_v3_on()
28258        {
28259            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
28260        }
28261        if fa_sm_count() < 128
28262            || std::env::var("MEMRA_FA_SPLIT").is_ok()
28263            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
28264            || std::env::var("MEMRA_FA_SP16").is_ok()
28265        {
28266            return Err(
28267                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
28268                 (or a <128-SM rig) keep the per-row path"
28269                    .into(),
28270            );
28271        }
28272        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
28273            return Err("fa_decode_dcw_rows geometry".into());
28274        }
28275        let o_len = t * n_head * max_ns * head_dim;
28276        let ml_len = t * n_head * max_ns;
28277        let mut part_guard = self.fa_part_pool.lock().unwrap();
28278        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28279        let pg = part_guard.as_mut().unwrap();
28280        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28281        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
28282        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28283        let (win, mns) = (window as i32, max_ns as i32);
28284        let gqa = (n_head / n_head_kv).max(1) as u32;
28285        let smem = (32 * head_dim * 2) as u32;
28286        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
28287        let cfg = LaunchConfig {
28288            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
28289            block_dim: (32, gqa, 1),
28290            shared_mem_bytes: smem,
28291        };
28292        {
28293            let __s_b = self.gpu.stream();
28294            let mut b = __s_b.launch_builder(&f);
28295            b.arg(q_rows)
28296                .arg(tab)
28297                .arg(&mut *part_o)
28298                .arg(&mut *part_m)
28299                .arg(&mut *part_l)
28300                .arg(&hd)
28301                .arg(&nh)
28302                .arg(&nhkv)
28303                .arg(&win)
28304                .arg(&scale)
28305                .arg(&mns)
28306                .arg(&ktb)
28307                .arg(&vtb);
28308            unsafe {
28309                b.launch(cfg)?;
28310            }
28311        }
28312        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
28313        // row r head h reads its own partial bank; splits past a row's ns_eff carry
28314        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
28315        let fc = {
28316            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28317            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28318                self.func("fa_decode_combine_gate_f32_s")
28319            } else {
28320                self.func("fa_decode_combine_gate_f32")
28321            }
28322        };
28323        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
28324        let nht = (t * n_head) as i32;
28325        let cfg2 = LaunchConfig {
28326            grid_dim: ((t * n_head) as u32, 1, 1),
28327            block_dim: (head_dim as u32, 1, 1),
28328            shared_mem_bytes: if combine_shared {
28329                (2 * max_ns * 4) as u32
28330            } else {
28331                0
28332            },
28333        };
28334        let __s_b2 = self.gpu.stream();
28335        let mut b2 = __s_b2.launch_builder(&fc);
28336        b2.arg(&*part_o)
28337            .arg(&*part_m)
28338            .arg(&*part_l)
28339            .arg(gate_rows)
28340            .arg(o_rows)
28341            .arg(&hd)
28342            .arg(&nht)
28343            .arg(&mns);
28344        unsafe {
28345            b2.launch(cfg2)?;
28346        }
28347        Ok(())
28348    }
28349
28350    #[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
28351    pub fn fa_decode_dcw(
28352        &self,
28353        q: &CudaSlice<f32>,
28354        k_ring: &cudarc::driver::CudaView<u8>,
28355        v_ring: &cudarc::driver::CudaView<u8>,
28356        o: &mut CudaSlice<f32>,
28357        head_dim: usize,
28358        n_head: usize,
28359        n_head_kv: usize,
28360        len_dev: &CudaSlice<i32>,
28361        base_dev: Option<&CudaSlice<i32>>,
28362        window: usize,
28363        bucket_max: usize,
28364        scale: f32,
28365        k_tok_bytes: usize,
28366        v_tok_bytes: usize,
28367        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
28368        // one launch saved); `o` then receives the GATED output and the caller skips its
28369        // attn_head_gate call.
28370        fused_gate: Option<&CudaSlice<f32>>,
28371    ) -> Result<(), Box<dyn std::error::Error>> {
28372        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
28373        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
28374            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"
28375                .into());
28376        }
28377        let sp = fa_split_keys(bucket_max, n_head_kv);
28378        #[allow(clippy::manual_div_ceil)]
28379        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28380        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
28381        let o_len = n_head * n_splits * head_dim;
28382        let ml_len = n_head * n_splits;
28383        let mut part_guard = self.fa_part_pool.lock().unwrap();
28384        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
28385        let pg = part_guard.as_mut().unwrap();
28386        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
28387        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
28388        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
28389        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
28390        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28391        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
28392        // finds the attention children BY their three-memset signature and updates the
28393        // memset widths per bucket — capturing without them silently kills retargeting
28394        // (battery-v8 token drift, 2026-08-21).
28395        let memset_on = *MEMSET_ON
28396            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
28397            || crate::tp::token_graph_building();
28398        if memset_on {
28399            self.gpu
28400                .stream()
28401                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28402            self.gpu
28403                .stream()
28404                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28405            self.gpu
28406                .stream()
28407                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28408        }
28409        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28410        let (hd, nh, nhkv, nsp) = (
28411            head_dim as i32,
28412            n_head as i32,
28413            n_head_kv as i32,
28414            n_splits as i32,
28415        );
28416        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28417        let (ski, win) = (sp as i32, window as i32);
28418        let gqa = (n_head / n_head_kv).max(1) as u32;
28419        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
28420        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
28421        // see fa_dec_v3_walk_u). Same launch geometry.
28422        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28423        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
28424        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
28425            Ok("2") => 2,
28426            Ok("1") => 1,
28427            _ => 0,
28428        });
28429        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
28430        // permission-blocked in this container and the module params are not exposed, so this
28431        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
28432        // prints cumulative cycle shares every 430 launches.
28433        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28434        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
28435        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
28436            std::sync::Mutex::new(None);
28437        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
28438        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
28439        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
28440        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28441        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
28442            && (n_head / n_head_kv).is_multiple_of(2)
28443            && (n_head / n_head_kv) >= 2;
28444        let f = if fprof {
28445            self.func("fa_decode_vec_q_v3_dcw_prof")
28446        } else if hs2 {
28447            self.func("fa_decode_vec_q_v3_dcw_hs2")
28448        } else if hoist == 2 {
28449            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
28450            self.func("fa_decode_vec_q_v3_dcw_hc")
28451        } else if hoist == 1 {
28452            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
28453            self.func("fa_decode_vec_q_v3_dcw_h")
28454        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
28455            self.func("fa_decode_vec_q_v3_dcw_u8")
28456        } else {
28457            self.func("fa_decode_vec_q_v3_dcw")
28458        };
28459        let cfg = LaunchConfig {
28460            grid_dim: if hs2 {
28461                ((2 * n_head_kv) as u32, n_splits as u32, 1)
28462            } else {
28463                (n_head_kv as u32, n_splits as u32, 1)
28464            },
28465            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
28466            shared_mem_bytes: smem,
28467        };
28468        let null: u64 = 0;
28469        let __s_b = self.gpu.stream();
28470        let mut b = __s_b.launch_builder(&f);
28471        b.arg(q)
28472            .arg(k_ring)
28473            .arg(v_ring)
28474            .arg(&mut *part_o)
28475            .arg(&mut *part_m)
28476            .arg(&mut *part_l)
28477            .arg(&hd)
28478            .arg(&nh)
28479            .arg(&nhkv)
28480            .arg(len_dev);
28481        match base_dev {
28482            Some(base) => {
28483                b.arg(base);
28484            }
28485            None => {
28486                b.arg(&null);
28487            }
28488        }
28489        b.arg(&win)
28490            .arg(&scale)
28491            .arg(&nsp)
28492            .arg(&ski)
28493            .arg(&ktb)
28494            .arg(&vtb);
28495        if fprof {
28496            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
28497            if guard
28498                .as_ref()
28499                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
28500            {
28501                *guard = Some((self.ctx().ordinal(), self.htod_u64(&[0u64; 8])?));
28502            }
28503            let (_, buf) = guard.as_mut().expect("armed above");
28504            b.arg(&*buf);
28505            unsafe {
28506                b.launch(cfg)?;
28507            }
28508            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
28509            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
28510            if n.is_multiple_of(430) {
28511                self.stream().synchronize()?;
28512                let h = self.dtoh_u64(buf)?;
28513                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
28514                let tot: u64 = h[..6].iter().sum();
28515                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
28516                for (i, name) in phases.iter().enumerate() {
28517                    let pct = if tot > 0 {
28518                        h[i] as f64 / tot as f64 * 100.0
28519                    } else {
28520                        0.0
28521                    };
28522                    line.push_str(&format!(" {name}={pct:.1}%"));
28523                }
28524                if h[6] > 0 {
28525                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
28526                }
28527                eprintln!("{line}");
28528            }
28529        } else {
28530            unsafe {
28531                b.launch(cfg)?;
28532            }
28533        }
28534        let mut combine_shared = false;
28535        let fc = if fused_gate.is_some() {
28536            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
28537            // n_splits-deep dependent global load chain every thread used to walk twice).
28538            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28539            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
28540                combine_shared = true;
28541                self.func("fa_decode_combine_gate_f32_s")
28542            } else {
28543                self.func("fa_decode_combine_gate_f32")
28544            }
28545        } else {
28546            self.fa_func("fa_decode_combine_f32", head_dim)
28547        };
28548        let cfg2 = LaunchConfig {
28549            grid_dim: (n_head as u32, 1, 1),
28550            block_dim: (head_dim as u32, 1, 1),
28551            shared_mem_bytes: if combine_shared {
28552                (2 * n_splits * 4) as u32
28553            } else {
28554                0
28555            },
28556        };
28557        let __s_b2 = self.gpu.stream();
28558        let mut b2 = __s_b2.launch_builder(&fc);
28559        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
28560        if let Some(gate_row) = fused_gate {
28561            b2.arg(gate_row);
28562        }
28563        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
28564        unsafe {
28565            b2.launch(cfg2)?;
28566        }
28567        Ok(())
28568    }
28569
28570    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
28571    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
28572    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
28573    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
28574    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
28575    #[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
28576    pub fn fa_geom_eager(
28577        &self,
28578        t_kv: usize,
28579        head_dim: usize,
28580        n_head_kv: usize,
28581        g: bool,
28582    ) -> (bool, usize) {
28583        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
28584        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
28585        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
28586        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
28587        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
28588        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
28589        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
28590        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
28591        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
28592        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
28593        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim.is_multiple_of(32));
28594        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
28595        // family; everything else falls to the g-module scalar.
28596        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
28597        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
28598        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
28599        if g && head_dim == 256 && !fa_v4_at(t_kv) {
28600            fa_vec = false;
28601        }
28602        let sp = fa_split_keys(t_kv, n_head_kv);
28603        let n_splits = if fa_vec {
28604            ((t_kv + sp - 1) / sp).max(1)
28605        } else {
28606            ((t_kv + 255) / 256).max(1)
28607        };
28608        (fa_vec, n_splits)
28609    }
28610
28611    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
28612    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
28613    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
28614    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
28615    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
28616    pub fn fa_bucket_key(
28617        &self,
28618        t_kv: usize,
28619        head_dim: usize,
28620        n_head_kv: usize,
28621        g: bool,
28622    ) -> (bool, usize) {
28623        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
28624    }
28625
28626    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
28627    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
28628    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
28629    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
28630    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
28631    /// device data) — every per-step varying scalar must come from a device counter. Returns the
28632    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
28633    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
28634    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
28635    /// replays (transients returning to the pool get reused by unrelated work and corrupt
28636    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
28637    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28638    pub fn capture_graph_retained<F>(
28639        &self,
28640        step: F,
28641    ) -> Result<
28642        (
28643            cudarc::driver::CudaGraph,
28644            Vec<Box<dyn std::any::Any + Send>>,
28645        ),
28646        Box<dyn std::error::Error>,
28647    >
28648    where
28649        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28650    {
28651        use cudarc::driver::sys::CUgraphInstantiate_flags;
28652        self.capture_graph_retained_flags(
28653            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28654            step,
28655        )
28656    }
28657
28658    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
28659    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
28660    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
28661    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
28662    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28663    pub fn capture_graph_retained_flags<F>(
28664        &self,
28665        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
28666        mut step: F,
28667    ) -> Result<
28668        (
28669            cudarc::driver::CudaGraph,
28670            Vec<Box<dyn std::any::Any + Send>>,
28671        ),
28672        Box<dyn std::error::Error>,
28673    >
28674    where
28675        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28676    {
28677        use cudarc::driver::sys::CUstreamCaptureMode;
28678        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
28679        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
28680        // while the capture region is open become dead copy NODES replayed every launch
28681        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
28682        // warmup runs allocate the same transient sequence at the same pool addresses, so
28683        // retaining the warmup clones preserves the draft-graph fix without polluting the
28684        // captured graph.
28685        self.capture_keep.lock().unwrap().clear();
28686        let was_tracking = self.gpu.ctx.is_event_tracking();
28687        if was_tracking {
28688            unsafe {
28689                self.gpu.ctx.disable_event_tracking();
28690            }
28691        }
28692        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28693            self.capture_keep_on
28694                .store(true, std::sync::atomic::Ordering::Relaxed);
28695            let w = (|| {
28696                step(self)?;
28697                step(self)
28698            })();
28699            self.capture_keep_on
28700                .store(false, std::sync::atomic::Ordering::Relaxed);
28701            w?;
28702            self.gpu.stream().synchronize()?;
28703            self.gpu
28704                .stream()
28705                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28706            let r = step(self);
28707            let g = self.gpu.stream().end_capture(flags);
28708            r?;
28709            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28710            graph.upload()?;
28711            Ok(graph)
28712        };
28713        let result = run();
28714        self.capture_keep_on
28715            .store(false, std::sync::atomic::Ordering::Relaxed);
28716        if was_tracking {
28717            unsafe {
28718                self.gpu.ctx.enable_event_tracking();
28719            }
28720        }
28721        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
28722        Ok((result?, keeper))
28723    }
28724
28725    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
28726    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
28727    /// alloc-free with persistent operands, and their bodies carry device side effects
28728    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
28729    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
28730    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28731    pub fn capture_graph_retained_nowarm<F>(
28732        &self,
28733        mut step: F,
28734    ) -> Result<
28735        (
28736            cudarc::driver::CudaGraph,
28737            Vec<Box<dyn std::any::Any + Send>>,
28738        ),
28739        Box<dyn std::error::Error>,
28740    >
28741    where
28742        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28743    {
28744        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
28745        let was_tracking = self.gpu.ctx.is_event_tracking();
28746        if was_tracking {
28747            unsafe {
28748                self.gpu.ctx.disable_event_tracking();
28749            }
28750        }
28751        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28752            self.gpu.stream().synchronize()?;
28753            self.gpu
28754                .stream()
28755                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28756            let r = step(self);
28757            let g = self.gpu.stream().end_capture(
28758                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28759            );
28760            r?;
28761            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28762            graph.upload()?;
28763            Ok(graph)
28764        };
28765        let result = run();
28766        if was_tracking {
28767            unsafe {
28768                self.gpu.ctx.enable_event_tracking();
28769            }
28770        }
28771        Ok((result?, Vec::new()))
28772    }
28773
28774    pub fn capture_graph<F>(
28775        &self,
28776        mut step: F,
28777    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
28778    where
28779        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
28780    {
28781        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
28782        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
28783        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
28784        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
28785        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
28786        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
28787        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
28788        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
28789        let was_tracking = self.gpu.ctx.is_event_tracking();
28790        if was_tracking {
28791            unsafe {
28792                self.gpu.ctx.disable_event_tracking();
28793            }
28794        }
28795        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
28796        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
28797        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
28798        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
28799        // measure that scan's real cost on the generic path. Diagnostic door only; the
28800        // default stays AUTO_FREE until a measured A/B justifies moving it.
28801        let iflag = {
28802            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
28803            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
28804                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
28805                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
28806                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
28807                Ok("priority") => {
28808                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
28809                }
28810                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
28811            })
28812        };
28813        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
28814        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
28815        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
28816        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
28817        // eager step executions and are node-count-invariant. Printing the split bounds the
28818        // refactor's ceiling instead of assuming it.
28819        let ct = {
28820            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28821            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
28822        };
28823        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
28824        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
28825        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
28826        // chased, and node-count-invariant, so no capture-body refactor could touch it.
28827        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
28828        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
28829        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
28830        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
28831        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
28832        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
28833        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
28834        // grow and never frees, resident counters/scratch, cache set in place), and the
28835        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
28836        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
28837        // settling and pool mapping. Arbitrated adversarially, not by taste:
28838        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
28839        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
28840        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
28841        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
28842        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
28843        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
28844        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
28845        let warmups = {
28846            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
28847            *W.get_or_init(|| {
28848                std::env::var("MEMRA_GRAPH_WARMUPS")
28849                    .ok()
28850                    .and_then(|v| v.parse().ok())
28851                    .filter(|n| *n >= 1)
28852                    .unwrap_or(1)
28853            })
28854        };
28855        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
28856            let t_w = std::time::Instant::now();
28857            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
28858            for _ in 0..warmups {
28859                step(self)?;
28860            }
28861            self.gpu.stream().synchronize()?;
28862            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
28863            // capture the third run.
28864            let t_c = std::time::Instant::now();
28865            self.gpu
28866                .stream()
28867                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
28868            // If the body errors mid-capture, end the capture before propagating so the stream isn't
28869            // left in a capturing state.
28870            let r = step(self);
28871            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
28872            let t_i = std::time::Instant::now();
28873            let g = self.gpu.stream().end_capture(iflag);
28874            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
28875            r?;
28876            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
28877            let t_u = std::time::Instant::now();
28878            graph.upload()?;
28879            if ct {
28880                println!(
28881                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
28882                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
28883                    t_u.elapsed().as_secs_f64() * 1e3
28884                );
28885            }
28886            Ok(graph)
28887        };
28888        let result = run();
28889        if was_tracking {
28890            unsafe {
28891                self.gpu.ctx.enable_event_tracking();
28892            }
28893        }
28894        result
28895    }
28896
28897    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
28898    #[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
28899    pub fn gdn_scan_s128_view(
28900        &self,
28901        q: &CudaSlice<f32>,
28902        k: &CudaSlice<f32>,
28903        v: &CudaSlice<f32>,
28904        g: &CudaSlice<f32>,
28905        beta: &CudaSlice<f32>,
28906        state_in: &cudarc::driver::CudaView<f32>,
28907        state_out: &mut cudarc::driver::CudaViewMut<f32>,
28908        o: &mut CudaSlice<f32>,
28909        n_head: usize,
28910        t: usize,
28911        scale: f32,
28912    ) -> Result<(), Box<dyn std::error::Error>> {
28913        let f = self.func("gdn_scan_s128");
28914        const S_V: u32 = 128;
28915        const WARP: u32 = 32;
28916        const COLS: u32 = 4;
28917        let cfg = LaunchConfig {
28918            grid_dim: (n_head as u32, 1, S_V / COLS),
28919            block_dim: (WARP, COLS, 1),
28920            shared_mem_bytes: 0,
28921        };
28922        let (h, ti) = (n_head as i32, t as i32);
28923        let __s_b = self.gpu.stream();
28924        let mut b = __s_b.launch_builder(&f);
28925        b.arg(q)
28926            .arg(k)
28927            .arg(v)
28928            .arg(g)
28929            .arg(beta)
28930            .arg(state_in)
28931            .arg(state_out)
28932            .arg(o)
28933            .arg(&h)
28934            .arg(&ti)
28935            .arg(&scale);
28936        unsafe {
28937            b.launch(cfg)?;
28938        }
28939        Ok(())
28940    }
28941
28942    /// conv1d where the input is a CudaView (resident conv state assembled in place).
28943    #[allow(clippy::too_many_arguments)]
28944    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
28945    #[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
28946    pub fn ssm_conv1d_view(
28947        &self,
28948        x: &cudarc::driver::CudaView<f32>,
28949        w: &CudaSlice<f32>,
28950        y: &mut CudaSlice<f32>,
28951        conv_dim: usize,
28952        t: usize,
28953        d_conv: usize,
28954        silu: bool,
28955    ) -> Result<(), Box<dyn std::error::Error>> {
28956        let f = self.func("ssm_conv1d_silu_f32");
28957        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
28958        let cfg = LaunchConfig {
28959            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
28960            block_dim: (256, 1, 1),
28961            shared_mem_bytes: 0,
28962        };
28963        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
28964        let __s_b = self.gpu.stream();
28965        let mut b = __s_b.launch_builder(&f);
28966        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
28967        unsafe {
28968            b.launch(cfg)?;
28969        }
28970        Ok(())
28971    }
28972
28973    /// Depthwise causal conv1d + optional SiLU.
28974    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
28975    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
28976    /// FUSED prefill conv (token-major input, zero left-state): replaces
28977    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
28978    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
28979    #[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
28980    pub fn ssm_conv1d_tm(
28981        &self,
28982        qkv_tm: &CudaSlice<f32>,
28983        w: &CudaSlice<f32>,
28984        y: &mut CudaSlice<f32>,
28985        conv_dim: usize,
28986        t: usize,
28987        d_conv: usize,
28988    ) -> Result<(), Box<dyn std::error::Error>> {
28989        let f = self.func("ssm_conv1d_tm_f32");
28990        let cfg = LaunchConfig {
28991            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
28992            block_dim: (256, 1, 1),
28993            shared_mem_bytes: 0,
28994        };
28995        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
28996        let __s_b = self.gpu.stream();
28997        let mut b = __s_b.launch_builder(&f);
28998        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
28999        unsafe {
29000            b.launch(cfg)?;
29001        }
29002        Ok(())
29003    }
29004
29005    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
29006    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
29007    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
29008    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
29009    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
29010    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
29011    /// columns; the final ring == what T sequential decode ring rolls leave).
29012    #[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
29013    pub fn ssm_conv1d_tm_state(
29014        &self,
29015        qkv_tm: &CudaSlice<f32>,
29016        conv_state: &mut CudaSlice<f32>,
29017        w: &CudaSlice<f32>,
29018        y: &mut CudaSlice<f32>,
29019        conv_dim: usize,
29020        t: usize,
29021        d_conv: usize,
29022    ) -> Result<(), Box<dyn std::error::Error>> {
29023        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
29024    }
29025
29026    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
29027    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
29028    #[allow(clippy::too_many_arguments)]
29029    #[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
29030    pub fn ssm_conv1d_tm_state_pad(
29031        &self,
29032        qkv_tm: &CudaSlice<f32>,
29033        conv_state: &mut CudaSlice<f32>,
29034        w: &CudaSlice<f32>,
29035        y: &mut CudaSlice<f32>,
29036        conv_dim: usize,
29037        t: usize,
29038        d_conv: usize,
29039        pad_len: Option<&CudaSlice<i32>>,
29040    ) -> Result<(), Box<dyn std::error::Error>> {
29041        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
29042        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
29043        // the window kernel both read the pre-roll ring; the roll launches after both) — but
29044        // cloning first keeps the ordering trivially correct under any future stream split.
29045        let ring_old = if t < d_conv - 1 {
29046            Some(self.clone_dtod(conv_state)?)
29047        } else {
29048            None
29049        };
29050        {
29051            let f = self.func("ssm_conv1d_tm_state_f32");
29052            let cfg = LaunchConfig {
29053                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29054                block_dim: (256, 1, 1),
29055                shared_mem_bytes: 0,
29056            };
29057            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29058            let __s_b = self.gpu.stream();
29059            let mut b = __s_b.launch_builder(&f);
29060            b.arg(qkv_tm)
29061                .arg(&*conv_state)
29062                .arg(w)
29063                .arg(y)
29064                .arg(&cd)
29065                .arg(&ti)
29066                .arg(&dc);
29067            unsafe {
29068                b.launch(cfg)?;
29069            }
29070        }
29071        match (ring_old, pad_len) {
29072            (None, Some(len_d)) => {
29073                let f = self.func("ssm_conv_ring_update_dev_f32");
29074                let n = conv_dim * (d_conv - 1);
29075                let cfg = LaunchConfig::for_num_elems(n as u32);
29076                let (cd, dc) = (conv_dim as i32, d_conv as i32);
29077                let __s_b = self.gpu.stream();
29078                let mut b = __s_b.launch_builder(&f);
29079                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
29080                unsafe {
29081                    b.launch(cfg)?;
29082                }
29083            }
29084            (None, None) => {
29085                let f = self.func("ssm_conv_ring_update_f32");
29086                let n = conv_dim * (d_conv - 1);
29087                let cfg = LaunchConfig::for_num_elems(n as u32);
29088                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29089                let __s_b = self.gpu.stream();
29090                let mut b = __s_b.launch_builder(&f);
29091                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
29092                unsafe {
29093                    b.launch(cfg)?;
29094                }
29095            }
29096            (Some(old), _) => {
29097                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
29098            }
29099        }
29100        Ok(())
29101    }
29102
29103    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
29104    #[allow(clippy::too_many_arguments)]
29105    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29106    #[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
29107    pub fn ssm_conv1d_tm_state_pad_v(
29108        &self,
29109        qkv_tm: &cudarc::driver::CudaView<f32>,
29110        conv_state: &mut CudaSlice<f32>,
29111        w: &CudaSlice<f32>,
29112        y: &mut CudaSlice<f32>,
29113        conv_dim: usize,
29114        t: usize,
29115        d_conv: usize,
29116        pad_len: Option<&CudaSlice<i32>>,
29117    ) -> Result<(), Box<dyn std::error::Error>> {
29118        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
29119        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
29120        // the window kernel both read the pre-roll ring; the roll launches after both) — but
29121        // cloning first keeps the ordering trivially correct under any future stream split.
29122        let ring_old = if t < d_conv - 1 {
29123            Some(self.clone_dtod(conv_state)?)
29124        } else {
29125            None
29126        };
29127        {
29128            let f = self.func("ssm_conv1d_tm_state_f32");
29129            let cfg = LaunchConfig {
29130                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29131                block_dim: (256, 1, 1),
29132                shared_mem_bytes: 0,
29133            };
29134            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29135            let __s_b = self.gpu.stream();
29136            let mut b = __s_b.launch_builder(&f);
29137            b.arg(qkv_tm)
29138                .arg(&*conv_state)
29139                .arg(w)
29140                .arg(y)
29141                .arg(&cd)
29142                .arg(&ti)
29143                .arg(&dc);
29144            unsafe {
29145                b.launch(cfg)?;
29146            }
29147        }
29148        match (ring_old, pad_len) {
29149            (None, Some(len_d)) => {
29150                let f = self.func("ssm_conv_ring_update_dev_f32");
29151                let n = conv_dim * (d_conv - 1);
29152                let cfg = LaunchConfig::for_num_elems(n as u32);
29153                let (cd, dc) = (conv_dim as i32, d_conv as i32);
29154                let __s_b = self.gpu.stream();
29155                let mut b = __s_b.launch_builder(&f);
29156                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
29157                unsafe {
29158                    b.launch(cfg)?;
29159                }
29160            }
29161            (None, None) => {
29162                let f = self.func("ssm_conv_ring_update_f32");
29163                let n = conv_dim * (d_conv - 1);
29164                let cfg = LaunchConfig::for_num_elems(n as u32);
29165                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29166                let __s_b = self.gpu.stream();
29167                let mut b = __s_b.launch_builder(&f);
29168                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
29169                unsafe {
29170                    b.launch(cfg)?;
29171                }
29172            }
29173            (Some(_), _) => unreachable!(
29174                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
29175            ),
29176        }
29177        Ok(())
29178    }
29179
29180    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
29181    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
29182    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
29183    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
29184    pub fn ssm_conv_ring_rebuild(
29185        &self,
29186        qkv_tm: &CudaSlice<f32>,
29187        ring_old: &CudaSlice<f32>,
29188        conv_state: &mut CudaSlice<f32>,
29189        conv_dim: usize,
29190        tc: usize,
29191        d_conv: usize,
29192    ) -> Result<(), Box<dyn std::error::Error>> {
29193        let f = self.func("ssm_conv_ring_rebuild_f32");
29194        let n = conv_dim * (d_conv - 1);
29195        let cfg = LaunchConfig::for_num_elems(n as u32);
29196        let (cd, ti, dc) = (conv_dim as i32, tc 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)
29200            .arg(ring_old)
29201            .arg(conv_state)
29202            .arg(&cd)
29203            .arg(&ti)
29204            .arg(&dc);
29205        unsafe {
29206            b.launch(cfg)?;
29207        }
29208        Ok(())
29209    }
29210
29211    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
29212    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
29213    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
29214    /// the argmax + run-spec gates are the authority.
29215    #[allow(clippy::too_many_arguments)]
29216    pub fn gdn_prep_decode(
29217        &self,
29218        conv_out: &CudaSlice<f32>,
29219        beta_raw: &CudaSlice<f32>,
29220        alpha: &CudaSlice<f32>,
29221        dt_bias: &CudaSlice<f32>,
29222        a: &CudaSlice<f32>,
29223        q_l2: &mut CudaSlice<f32>,
29224        k_l2: &mut CudaSlice<f32>,
29225        v_g: &mut CudaSlice<f32>,
29226        beta: &mut CudaSlice<f32>,
29227        g_log: &mut CudaSlice<f32>,
29228        d_state: usize,
29229        num_v: usize,
29230        num_k: usize,
29231        key_dim: usize,
29232        eps: f32,
29233    ) -> Result<(), Box<dyn std::error::Error>> {
29234        let f = self.func("gdn_prep_decode_f32");
29235        let cfg = LaunchConfig {
29236            grid_dim: (num_v as u32, 1, 1),
29237            block_dim: (32, 4, 1),
29238            shared_mem_bytes: 0,
29239        };
29240        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
29241        let __s_b = self.gpu.stream();
29242        let mut b = __s_b.launch_builder(&f);
29243        b.arg(conv_out)
29244            .arg(beta_raw)
29245            .arg(alpha)
29246            .arg(dt_bias)
29247            .arg(a)
29248            .arg(q_l2)
29249            .arg(k_l2)
29250            .arg(v_g)
29251            .arg(beta)
29252            .arg(g_log)
29253            .arg(&ds)
29254            .arg(&nv)
29255            .arg(&nk)
29256            .arg(&kd)
29257            .arg(&eps);
29258        unsafe {
29259            b.launch(cfg)?;
29260        }
29261        Ok(())
29262    }
29263
29264    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
29265    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
29266    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
29267    #[allow(clippy::too_many_arguments)]
29268    #[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
29269    pub fn ssm_conv1d_gdn(
29270        &self,
29271        qkv_tm: &CudaSlice<f32>,
29272        w: &CudaSlice<f32>,
29273        q_g: &mut CudaSlice<f32>,
29274        k_g: &mut CudaSlice<f32>,
29275        v_g: &mut CudaSlice<f32>,
29276        conv_dim: usize,
29277        t: usize,
29278        d_conv: usize,
29279        d_state: usize,
29280        num_v: usize,
29281        num_k: usize,
29282        key_dim: usize,
29283    ) -> Result<(), Box<dyn std::error::Error>> {
29284        let f = self.func("ssm_conv1d_gdn_f32");
29285        let cfg = LaunchConfig {
29286            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29287            block_dim: (256, 1, 1),
29288            shared_mem_bytes: 0,
29289        };
29290        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29291        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
29292        let __s_b = self.gpu.stream();
29293        let mut b = __s_b.launch_builder(&f);
29294        b.arg(qkv_tm)
29295            .arg(w)
29296            .arg(q_g)
29297            .arg(k_g)
29298            .arg(v_g)
29299            .arg(&cd)
29300            .arg(&ti)
29301            .arg(&dc)
29302            .arg(&ds)
29303            .arg(&nv)
29304            .arg(&nk)
29305            .arg(&kd);
29306        unsafe {
29307            b.launch(cfg)?;
29308        }
29309        Ok(())
29310    }
29311
29312    #[allow(clippy::too_many_arguments)]
29313    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
29314    #[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
29315    pub fn ssm_conv1d(
29316        &self,
29317        x: &CudaSlice<f32>,
29318        w: &CudaSlice<f32>,
29319        y: &mut CudaSlice<f32>,
29320        conv_dim: usize,
29321        t: usize,
29322        d_conv: usize,
29323        silu: bool,
29324    ) -> Result<(), Box<dyn std::error::Error>> {
29325        let f = self.func("ssm_conv1d_silu_f32");
29326        let cfg = LaunchConfig {
29327            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
29328            block_dim: (256, 1, 1),
29329            shared_mem_bytes: 0,
29330        };
29331        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
29332        let __s_b = self.gpu.stream();
29333        let mut b = __s_b.launch_builder(&f);
29334        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
29335        unsafe {
29336            b.launch(cfg)?;
29337        }
29338        Ok(())
29339    }
29340
29341    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
29342    /// o:[128,H,T]. Single sequence.
29343    #[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
29344    pub fn gdn_scan_s128(
29345        &self,
29346        q: &CudaSlice<f32>,
29347        k: &CudaSlice<f32>,
29348        v: &CudaSlice<f32>,
29349        g: &CudaSlice<f32>,
29350        beta: &CudaSlice<f32>,
29351        state_in: &CudaSlice<f32>,
29352        state_out: &mut CudaSlice<f32>,
29353        o: &mut CudaSlice<f32>,
29354        n_head: usize,
29355        t: usize,
29356        scale: f32,
29357    ) -> Result<(), Box<dyn std::error::Error>> {
29358        let f = self.func("gdn_scan_s128");
29359        const S_V: u32 = 128;
29360        const WARP: u32 = 32;
29361        const COLS_PER_BLOCK: u32 = 4;
29362        let cfg = LaunchConfig {
29363            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
29364            block_dim: (WARP, COLS_PER_BLOCK, 1),
29365            shared_mem_bytes: 0,
29366        };
29367        let (h, ti) = (n_head as i32, t as i32);
29368        let __s_b = self.gpu.stream();
29369        let mut b = __s_b.launch_builder(&f);
29370        b.arg(q)
29371            .arg(k)
29372            .arg(v)
29373            .arg(g)
29374            .arg(beta)
29375            .arg(state_in)
29376            .arg(state_out)
29377            .arg(o)
29378            .arg(&h)
29379            .arg(&ti)
29380            .arg(&scale);
29381        unsafe {
29382            b.launch(cfg)?;
29383        }
29384        Ok(())
29385    }
29386
29387    // ==== B2' batched decode state ops (decode_batch.rs) ====
29388    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
29389    // Bodies are the single-seq kernels per sequence — bit-identical per row.
29390
29391    #[allow(clippy::too_many_arguments)]
29392    #[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
29393    pub fn ssm_conv1d_fused_decode_b(
29394        &self,
29395        qkv_cols: &CudaSlice<f32>,
29396        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
29397        w: &CudaSlice<f32>,
29398        conv_outs: &mut CudaSlice<f32>,
29399        conv_dim: usize,
29400        d_conv: usize,
29401        b_n: usize,
29402    ) -> Result<(), Box<dyn std::error::Error>> {
29403        let f = self.func("ssm_conv1d_fused_decode_b_f32");
29404        let cfg = LaunchConfig {
29405            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
29406            block_dim: (256, 1, 1),
29407            shared_mem_bytes: 0,
29408        };
29409        let (cd, dc) = (conv_dim as i32, d_conv as i32);
29410        let __s_b = self.gpu.stream();
29411        let mut b = __s_b.launch_builder(&f);
29412        b.arg(qkv_cols)
29413            .arg(conv_state_ptrs)
29414            .arg(w)
29415            .arg(conv_outs)
29416            .arg(&cd)
29417            .arg(&dc);
29418        unsafe {
29419            b.launch(cfg)?;
29420        }
29421        Ok(())
29422    }
29423
29424    #[allow(clippy::too_many_arguments)]
29425    pub fn gdn_prep_decode_b(
29426        &self,
29427        conv_outs: &CudaSlice<f32>,
29428        beta_raws: &CudaSlice<f32>,
29429        alphas: &CudaSlice<f32>,
29430        dt_bias: &CudaSlice<f32>,
29431        a: &CudaSlice<f32>,
29432        q_l2: &mut CudaSlice<f32>,
29433        k_l2: &mut CudaSlice<f32>,
29434        v_g: &mut CudaSlice<f32>,
29435        beta: &mut CudaSlice<f32>,
29436        g_log: &mut CudaSlice<f32>,
29437        d_state: usize,
29438        num_v: usize,
29439        num_k: usize,
29440        key_dim: usize,
29441        eps: f32,
29442        conv_dim: usize,
29443        b_n: usize,
29444    ) -> Result<(), Box<dyn std::error::Error>> {
29445        let f = self.func("gdn_prep_decode_b_f32");
29446        let cfg = LaunchConfig {
29447            grid_dim: (num_v as u32, 1, b_n as u32),
29448            block_dim: (32, 4, 1),
29449            shared_mem_bytes: 0,
29450        };
29451        let (ds, nv, nk, kd, cd) = (
29452            d_state as i32,
29453            num_v as i32,
29454            num_k as i32,
29455            key_dim as i32,
29456            conv_dim as i32,
29457        );
29458        let __s_b = self.gpu.stream();
29459        let mut b = __s_b.launch_builder(&f);
29460        b.arg(conv_outs)
29461            .arg(beta_raws)
29462            .arg(alphas)
29463            .arg(dt_bias)
29464            .arg(a)
29465            .arg(q_l2)
29466            .arg(k_l2)
29467            .arg(v_g)
29468            .arg(beta)
29469            .arg(g_log)
29470            .arg(&ds)
29471            .arg(&nv)
29472            .arg(&nk)
29473            .arg(&kd)
29474            .arg(&eps)
29475            .arg(&cd);
29476        unsafe {
29477            b.launch(cfg)?;
29478        }
29479        Ok(())
29480    }
29481
29482    #[allow(clippy::too_many_arguments)]
29483    pub fn gdn_scan_s128_batched(
29484        &self,
29485        q: &CudaSlice<f32>,
29486        k: &CudaSlice<f32>,
29487        v: &CudaSlice<f32>,
29488        g: &CudaSlice<f32>,
29489        beta: &CudaSlice<f32>,
29490        state_in_ptrs: &cudarc::driver::CudaView<u64>,
29491        state_out_ptrs: &cudarc::driver::CudaView<u64>,
29492        o: &mut CudaSlice<f32>,
29493        n_head: usize,
29494        b_n: usize,
29495        scale: f32,
29496    ) -> Result<(), Box<dyn std::error::Error>> {
29497        let f = self.func("gdn_scan_s128_b");
29498        const S_V: u32 = 128;
29499        const WARP: u32 = 32;
29500        const COLS_PER_BLOCK: u32 = 4;
29501        let cfg = LaunchConfig {
29502            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
29503            block_dim: (WARP, COLS_PER_BLOCK, 1),
29504            shared_mem_bytes: 0,
29505        };
29506        let h = n_head as i32;
29507        let __s_b = self.gpu.stream();
29508        let mut b = __s_b.launch_builder(&f);
29509        b.arg(q)
29510            .arg(k)
29511            .arg(v)
29512            .arg(g)
29513            .arg(beta)
29514            .arg(state_in_ptrs)
29515            .arg(state_out_ptrs)
29516            .arg(o)
29517            .arg(&h)
29518            .arg(&scale);
29519        unsafe {
29520            b.launch(cfg)?;
29521        }
29522        Ok(())
29523    }
29524
29525    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
29526    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
29527    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
29528    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
29529    /// numeric class; only the pointer arithmetic moved host-side.
29530    #[allow(clippy::too_many_arguments)]
29531    #[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
29532    pub fn ssm_conv1d_fused_decode_b_view(
29533        &self,
29534        qkv_cols: &cudarc::driver::CudaView<f32>,
29535        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
29536        w: &CudaSlice<f32>,
29537        conv_outs: &mut CudaSlice<f32>,
29538        conv_dim: usize,
29539        d_conv: usize,
29540        b_n: usize,
29541    ) -> Result<(), Box<dyn std::error::Error>> {
29542        let f = self.func("ssm_conv1d_fused_decode_b_f32");
29543        let cfg = LaunchConfig {
29544            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
29545            block_dim: (256, 1, 1),
29546            shared_mem_bytes: 0,
29547        };
29548        let (cd, dc) = (conv_dim as i32, d_conv as i32);
29549        let __s_b = self.gpu.stream();
29550        let mut b = __s_b.launch_builder(&f);
29551        b.arg(qkv_cols)
29552            .arg(conv_state_ptrs)
29553            .arg(w)
29554            .arg(conv_outs)
29555            .arg(&cd)
29556            .arg(&dc);
29557        unsafe {
29558            b.launch(cfg)?;
29559        }
29560        Ok(())
29561    }
29562
29563    #[allow(clippy::too_many_arguments)]
29564    pub fn gdn_prep_decode_b_view(
29565        &self,
29566        conv_outs: &CudaSlice<f32>,
29567        beta_raws: &cudarc::driver::CudaView<f32>,
29568        alphas: &cudarc::driver::CudaView<f32>,
29569        dt_bias: &CudaSlice<f32>,
29570        a: &CudaSlice<f32>,
29571        q_l2: &mut CudaSlice<f32>,
29572        k_l2: &mut CudaSlice<f32>,
29573        v_g: &mut CudaSlice<f32>,
29574        beta: &mut CudaSlice<f32>,
29575        g_log: &mut CudaSlice<f32>,
29576        d_state: usize,
29577        num_v: usize,
29578        num_k: usize,
29579        key_dim: usize,
29580        eps: f32,
29581        conv_dim: usize,
29582        b_n: usize,
29583    ) -> Result<(), Box<dyn std::error::Error>> {
29584        let f = self.func("gdn_prep_decode_b_f32");
29585        let cfg = LaunchConfig {
29586            grid_dim: (num_v as u32, 1, b_n as u32),
29587            block_dim: (32, 4, 1),
29588            shared_mem_bytes: 0,
29589        };
29590        let (ds, nv, nk, kd, cd) = (
29591            d_state as i32,
29592            num_v as i32,
29593            num_k as i32,
29594            key_dim as i32,
29595            conv_dim as i32,
29596        );
29597        let __s_b = self.gpu.stream();
29598        let mut b = __s_b.launch_builder(&f);
29599        b.arg(conv_outs)
29600            .arg(beta_raws)
29601            .arg(alphas)
29602            .arg(dt_bias)
29603            .arg(a)
29604            .arg(q_l2)
29605            .arg(k_l2)
29606            .arg(v_g)
29607            .arg(beta)
29608            .arg(g_log)
29609            .arg(&ds)
29610            .arg(&nv)
29611            .arg(&nk)
29612            .arg(&kd)
29613            .arg(&eps)
29614            .arg(&cd);
29615        unsafe {
29616            b.launch(cfg)?;
29617        }
29618        Ok(())
29619    }
29620
29621    #[allow(clippy::too_many_arguments)]
29622    pub fn gdn_scan_s128_batched_view(
29623        &self,
29624        q: &CudaSlice<f32>,
29625        k: &CudaSlice<f32>,
29626        v: &CudaSlice<f32>,
29627        g: &CudaSlice<f32>,
29628        beta: &CudaSlice<f32>,
29629        state_in_ptrs: &cudarc::driver::CudaView<u64>,
29630        state_out_ptrs: &cudarc::driver::CudaView<u64>,
29631        o: &mut cudarc::driver::CudaViewMut<f32>,
29632        n_head: usize,
29633        b_n: usize,
29634        scale: f32,
29635    ) -> Result<(), Box<dyn std::error::Error>> {
29636        let f = self.func("gdn_scan_s128_b");
29637        const S_V: u32 = 128;
29638        const WARP: u32 = 32;
29639        const COLS_PER_BLOCK: u32 = 4;
29640        let cfg = LaunchConfig {
29641            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
29642            block_dim: (WARP, COLS_PER_BLOCK, 1),
29643            shared_mem_bytes: 0,
29644        };
29645        let h = n_head as i32;
29646        let __s_b = self.gpu.stream();
29647        let mut b = __s_b.launch_builder(&f);
29648        b.arg(q)
29649            .arg(k)
29650            .arg(v)
29651            .arg(g)
29652            .arg(beta)
29653            .arg(state_in_ptrs)
29654            .arg(state_out_ptrs)
29655            .arg(o)
29656            .arg(&h)
29657            .arg(&scale);
29658        unsafe {
29659            b.launch(cfg)?;
29660        }
29661        Ok(())
29662    }
29663
29664    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
29665    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
29666    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
29667    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
29668    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
29669    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
29670    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
29671    /// identity law); prime_cache/forward/forward_last are the only callers.
29672    pub fn gdn_chunked_enabled() -> bool {
29673        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
29674        *E.get_or_init(|| {
29675            std::env::var("MEMRA_GDN_CHUNKED")
29676                .map(|v| v != "0")
29677                .unwrap_or(true)
29678        })
29679    }
29680
29681    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
29682    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
29683    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
29684    /// of 32 in [32, 128] (kernel row mappings require it).
29685    pub fn gdn_chunk_size() -> usize {
29686        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29687        *C.get_or_init(|| {
29688            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
29689                .ok()
29690                .and_then(|v| v.parse().ok())
29691                .unwrap_or(32);
29692            c.clamp(32, 128) / 32 * 32
29693        })
29694    }
29695
29696    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
29697    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
29698    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
29699    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
29700    #[allow(clippy::too_many_arguments)]
29701    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
29702    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
29703    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
29704    #[allow(clippy::too_many_arguments)]
29705    pub fn gdn_chunk_k123(
29706        &self,
29707        q: &CudaSlice<f32>,
29708        k: &CudaSlice<f32>,
29709        v: &CudaSlice<f32>,
29710        g: &CudaSlice<f32>,
29711        beta: &CudaSlice<f32>,
29712        wb16: Option<&mut CudaSlice<u8>>,
29713        n_head: usize,
29714        t: usize,
29715        c: usize,
29716        hk: usize,
29717        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
29718    ) -> Result<
29719        (
29720            CudaSlice<f32>,
29721            CudaSlice<f32>,
29722            CudaSlice<f32>,
29723            CudaSlice<f32>,
29724        ),
29725        Box<dyn std::error::Error>,
29726    > {
29727        const D: usize = 128;
29728        let h = n_head;
29729        #[allow(clippy::manual_div_ceil)]
29730        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29731        let nc = (t + c - 1) / c;
29732        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
29733        let mut gcum = self.uninit(t * h)?;
29734        let mut a = self.uninit(nc * h * c * c)?;
29735        let mut p = self.uninit(nc * h * c * c)?;
29736        let mut u = self.uninit(nc * h * c * D)?;
29737        let mut w = self.uninit(nc * h * c * D)?;
29738        {
29739            // K1
29740            let f = self.func("gdn_chunk_cumgate_f32");
29741            let cfg = LaunchConfig {
29742                grid_dim: (nc as u32, h as u32, 1),
29743                block_dim: (32, 1, 1),
29744                shared_mem_bytes: 0,
29745            };
29746            let __s_b = self.gpu.stream();
29747            let mut b = __s_b.launch_builder(&f);
29748            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
29749            unsafe {
29750                b.launch(cfg)?;
29751            }
29752        }
29753        if let Some((qb, kb, pb)) = k2w {
29754            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
29755            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
29756            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
29757            let f = self.func("gdn_k2_wgmma");
29758            let cfg = LaunchConfig {
29759                grid_dim: (nc as u32, h as u32, 1),
29760                block_dim: (128, 1, 1),
29761                shared_mem_bytes: 0,
29762            };
29763            let hki = hk as i32;
29764            let __s_b = self.gpu.stream();
29765            let mut b = __s_b.launch_builder(&f);
29766            b.arg(qb)
29767                .arg(kb)
29768                .arg(&gcum)
29769                .arg(beta)
29770                .arg(&mut a)
29771                .arg(&mut *pb)
29772                .arg(&hi)
29773                .arg(&ti)
29774                .arg(&ci)
29775                .arg(&hki);
29776            unsafe {
29777                b.launch(cfg)?;
29778            }
29779        } else if c <= 64 && !portable_mma_gated() {
29780            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
29781            let f = self.func("gdn_chunk_attn_f32");
29782            f.set_attribute(
29783                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29784                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
29785            )?;
29786            #[allow(clippy::manual_div_ceil)]
29787            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29788            let jt = ((c + 31) / 32) as u32;
29789            let cfg = LaunchConfig {
29790                grid_dim: (nc as u32, h as u32, jt),
29791                block_dim: (256, 1, 1),
29792                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
29793            };
29794            let hki = hk as i32;
29795            let __s_b = self.gpu.stream();
29796            let mut b = __s_b.launch_builder(&f);
29797            b.arg(q)
29798                .arg(k)
29799                .arg(&gcum)
29800                .arg(beta)
29801                .arg(&mut a)
29802                .arg(&mut p)
29803                .arg(&hi)
29804                .arg(&ti)
29805                .arg(&ci)
29806                .arg(&hki);
29807            unsafe {
29808                b.launch(cfg)?;
29809            }
29810        } else {
29811            // K2 generic (C = 128, or the portable target's low-smem fallback)
29812            assert!(
29813                hk == h,
29814                "generic K2 is broadcast-only (de-broadcast rides C==32)"
29815            );
29816            let f = self.func("gdn_chunk_attn_g_f32");
29817            let cfg = LaunchConfig {
29818                grid_dim: (nc as u32, h as u32, 1),
29819                block_dim: (32, 8, 1),
29820                shared_mem_bytes: 0,
29821            };
29822            let __s_b = self.gpu.stream();
29823            let mut b = __s_b.launch_builder(&f);
29824            b.arg(q)
29825                .arg(k)
29826                .arg(&gcum)
29827                .arg(beta)
29828                .arg(&mut a)
29829                .arg(&mut p)
29830                .arg(&hi)
29831                .arg(&ti)
29832                .arg(&ci);
29833            unsafe {
29834                b.launch(cfg)?;
29835            }
29836        }
29837        {
29838            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
29839            let cfg = LaunchConfig {
29840                grid_dim: (nc as u32, h as u32, 1),
29841                block_dim: (256, 1, 1),
29842                shared_mem_bytes: 0,
29843            };
29844            match c {
29845                32 | 64 => {
29846                    let f = self.func(if c == 32 {
29847                        "gdn_chunk_solve32_f32"
29848                    } else {
29849                        "gdn_chunk_solve64_f32"
29850                    });
29851                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
29852                    let wb: u64 = match wb16 {
29853                        Some(d) => self.addr_u8(d),
29854                        None => 0,
29855                    };
29856                    let hki = hk as i32;
29857                    let __s_b = self.gpu.stream();
29858                    let mut b = __s_b.launch_builder(&f);
29859                    b.arg(v)
29860                        .arg(k)
29861                        .arg(&a)
29862                        .arg(&gcum)
29863                        .arg(&mut u)
29864                        .arg(&mut w)
29865                        .arg(&wb)
29866                        .arg(&hi)
29867                        .arg(&ti)
29868                        .arg(&hki);
29869                    unsafe {
29870                        b.launch(cfg)?;
29871                    }
29872                }
29873                _ => {
29874                    assert!(hk == h, "generic K3 is broadcast-only");
29875                    let f = self.func("gdn_chunk_solve_f32");
29876                    let __s_b = self.gpu.stream();
29877                    let mut b = __s_b.launch_builder(&f);
29878                    b.arg(v)
29879                        .arg(k)
29880                        .arg(&a)
29881                        .arg(&gcum)
29882                        .arg(&mut u)
29883                        .arg(&mut w)
29884                        .arg(&hi)
29885                        .arg(&ti)
29886                        .arg(&ci);
29887                    unsafe {
29888                        b.launch(cfg)?;
29889                    }
29890                }
29891            }
29892        }
29893        Ok((gcum, p, u, w))
29894    }
29895
29896    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
29897    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
29898    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
29899    pub fn gdn_db_on() -> bool {
29900        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
29901    }
29902
29903    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
29904    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
29905    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
29906    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
29907    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
29908    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
29909    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
29910    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
29911    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
29912        !portable_mma_gated()
29913            && c == 32
29914            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
29915                Ok("1") => true,
29916                Ok("0") => false,
29917                _ => gdn_mma_default_on(),
29918            }
29919    }
29920
29921    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
29922    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
29923    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
29924    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
29925    /// force would silently produce garbage. Required since the sm_120a mma default
29926    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
29927    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
29928        cfg!(memra_hopper_mma)
29929            && self.gdn_mma_enabled(c)
29930            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
29931    }
29932
29933    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
29934    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
29935    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
29936    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
29937    #[allow(clippy::too_many_arguments)]
29938    #[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
29939    pub fn ssm_conv1d_gdn_state_pad(
29940        &self,
29941        qkv_tm: &cudarc::driver::CudaView<f32>,
29942        conv_state: &mut CudaSlice<f32>,
29943        w: &CudaSlice<f32>,
29944        q_g: &mut CudaSlice<f32>,
29945        k_g: &mut CudaSlice<f32>,
29946        v_g: &mut CudaSlice<f32>,
29947        conv_dim: usize,
29948        t: usize,
29949        d_conv: usize,
29950        d_state: usize,
29951        num_v: usize,
29952        num_k: usize,
29953        key_dim: usize,
29954        hk: usize,
29955        pad_len: Option<&CudaSlice<i32>>,
29956    ) -> Result<(), Box<dyn std::error::Error>> {
29957        assert!(
29958            t >= d_conv - 1,
29959            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
29960        );
29961        {
29962            let f = self.func("ssm_conv1d_gdn_state_f32");
29963            let cfg = LaunchConfig {
29964                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
29965                block_dim: (256, 1, 1),
29966                shared_mem_bytes: 0,
29967            };
29968            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
29969            let (ds, nv, nk, kd, hki) = (
29970                d_state as i32,
29971                num_v as i32,
29972                num_k as i32,
29973                key_dim as i32,
29974                hk as i32,
29975            );
29976            let __s_b = self.gpu.stream();
29977            let mut b = __s_b.launch_builder(&f);
29978            b.arg(qkv_tm)
29979                .arg(&*conv_state)
29980                .arg(w)
29981                .arg(q_g)
29982                .arg(k_g)
29983                .arg(v_g)
29984                .arg(&cd)
29985                .arg(&ti)
29986                .arg(&dc)
29987                .arg(&ds)
29988                .arg(&nv)
29989                .arg(&nk)
29990                .arg(&kd)
29991                .arg(&hki);
29992            unsafe {
29993                b.launch(cfg)?;
29994            }
29995        }
29996        match pad_len {
29997            Some(len_d) => {
29998                let f = self.func("ssm_conv_ring_update_dev_f32");
29999                let n = conv_dim * (d_conv - 1);
30000                let cfg = LaunchConfig::for_num_elems(n as u32);
30001                let (cd, dc) = (conv_dim as i32, d_conv as i32);
30002                let __s_b = self.gpu.stream();
30003                let mut b = __s_b.launch_builder(&f);
30004                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
30005                unsafe {
30006                    b.launch(cfg)?;
30007                }
30008            }
30009            None => {
30010                let f = self.func("ssm_conv_ring_update_f32");
30011                let n = conv_dim * (d_conv - 1);
30012                let cfg = LaunchConfig::for_num_elems(n as u32);
30013                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
30014                let __s_b = self.gpu.stream();
30015                let mut b = __s_b.launch_builder(&f);
30016                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
30017                unsafe {
30018                    b.launch(cfg)?;
30019                }
30020            }
30021        }
30022        Ok(())
30023    }
30024
30025    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
30026    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
30027    /// K2/K3 can write them.
30028    pub fn gdn_chunk_alloc(
30029        &self,
30030        n_head: usize,
30031        t: usize,
30032        c: usize,
30033        hk: usize,
30034    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
30035        const D: usize = 128;
30036        assert!(
30037            c == 32,
30038            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
30039        );
30040        let h = n_head;
30041        #[allow(clippy::manual_div_ceil)]
30042        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30043        let nc = (t + c - 1) / c;
30044        Ok(GdnChunkBufs {
30045            gcum: self.uninit(t * h)?,
30046            a: self.uninit(nc * h * c * c)?,
30047            p: self.uninit(nc * h * c * c)?,
30048            u: self.uninit(nc * h * c * D)?,
30049            w: self.uninit(nc * h * c * D)?,
30050            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
30051            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
30052            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
30053            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
30054            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
30055            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
30056            o: self.uninit(D * h * t)?,
30057            t,
30058            nc,
30059        })
30060    }
30061
30062    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
30063    pub fn f32_to_bf16_v(
30064        &self,
30065        x: &cudarc::driver::CudaView<f32>,
30066        dst: &mut CudaSlice<u8>,
30067        n: usize,
30068    ) -> Result<(), Box<dyn std::error::Error>> {
30069        let f = self.func("f32_to_bf16_bulk");
30070        let ni = n as i64;
30071        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
30072        let __s_b = self.gpu.stream();
30073        let mut b = __s_b.launch_builder(&f);
30074        b.arg(x).arg(dst).arg(&ni);
30075        unsafe {
30076            b.launch(cfg)?;
30077        }
30078        Ok(())
30079    }
30080
30081    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
30082    pub fn f32_to_bf16_into(
30083        &self,
30084        x: &CudaSlice<f32>,
30085        dst: &mut CudaSlice<u8>,
30086        n: usize,
30087    ) -> Result<(), Box<dyn std::error::Error>> {
30088        let f = self.func("f32_to_bf16_bulk");
30089        let ni = n as i64;
30090        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
30091        let __s_b = self.gpu.stream();
30092        let mut b = __s_b.launch_builder(&f);
30093        b.arg(x).arg(dst).arg(&ni);
30094        unsafe {
30095            b.launch(cfg)?;
30096        }
30097        Ok(())
30098    }
30099
30100    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
30101    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
30102    pub fn gdn_chunk_k123_vl8(
30103        &self,
30104        seqs: &[GdnSeqVl],
30105        n_head: usize,
30106        hk: usize,
30107        wq: Option<&GdnWVl8>,
30108    ) -> Result<(), Box<dyn std::error::Error>> {
30109        let b = seqs.len();
30110        assert!((1..=8).contains(&b), "gdn_chunk_k123_vl8: 1..=8 sequences");
30111        let mut packed = [GdnSeqVl::default(); 8];
30112        packed[..b].copy_from_slice(seqs);
30113        let v = GdnVl8(packed);
30114        let (hi, ci) = (n_head as i32, 32i32);
30115        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
30116        {
30117            let f = self.func("gdn_chunk_cumgate_vl");
30118            let cfg = LaunchConfig {
30119                grid_dim: (max_nc, n_head as u32, b as u32),
30120                block_dim: (32, 1, 1),
30121                shared_mem_bytes: 0,
30122            };
30123            let __s_lb = self.gpu.stream();
30124            let mut lb = __s_lb.launch_builder(&f);
30125            lb.arg(&v).arg(&hi).arg(&ci);
30126            unsafe {
30127                lb.launch(cfg)?;
30128            }
30129        }
30130        let hki = hk as i32;
30131        if let Some(w) = wq {
30132            // K2-wgmma vl twin (writes A + pre-masked Pb16)
30133            let f = self.func("gdn_k2_wgmma_vl");
30134            let cfg = LaunchConfig {
30135                grid_dim: (max_nc, n_head as u32, b as u32),
30136                block_dim: (128, 1, 1),
30137                shared_mem_bytes: 0,
30138            };
30139            let __s_lb = self.gpu.stream();
30140            let mut lb = __s_lb.launch_builder(&f);
30141            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
30142            unsafe {
30143                lb.launch(cfg)?;
30144            }
30145        } else {
30146            let f = self.func("gdn_chunk_attn_vl");
30147            f.set_attribute(
30148                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30149                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
30150            )?;
30151            let cfg = LaunchConfig {
30152                grid_dim: (max_nc, n_head as u32, b as u32),
30153                block_dim: (256, 1, 1),
30154                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
30155            };
30156            let __s_lb = self.gpu.stream();
30157            let mut lb = __s_lb.launch_builder(&f);
30158            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30159            unsafe {
30160                lb.launch(cfg)?;
30161            }
30162        }
30163        {
30164            let f = self.func("gdn_chunk_solve32_vl");
30165            let cfg = LaunchConfig {
30166                grid_dim: (max_nc, n_head as u32, b as u32),
30167                block_dim: (256, 1, 1),
30168                shared_mem_bytes: 0,
30169            };
30170            let __s_lb = self.gpu.stream();
30171            let mut lb = __s_lb.launch_builder(&f);
30172            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30173            unsafe {
30174                lb.launch(cfg)?;
30175            }
30176        }
30177        Ok(())
30178    }
30179
30180    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
30181    /// fused gate-prep, 5 launches for every sequence (per-element math identical
30182    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
30183    #[allow(clippy::too_many_arguments)]
30184    pub fn gdn_prep_vl8(
30185        &self,
30186        seqs: &[GdnPrepVl],
30187        conv_w: &CudaSlice<f32>,
30188        dt_bias: &CudaSlice<f32>,
30189        a: &CudaSlice<f32>,
30190        conv_dim: usize,
30191        d_conv: usize,
30192        d_state: usize,
30193        num_v: usize,
30194        num_k: usize,
30195        key_dim: usize,
30196        hk: usize,
30197        eps: f32,
30198    ) -> Result<(), Box<dyn std::error::Error>> {
30199        let b = seqs.len();
30200        assert!((1..=8).contains(&b));
30201        let mut packed = [GdnPrepVl::default(); 8];
30202        packed[..b].copy_from_slice(seqs);
30203        let v = GdnPrepVl8(packed);
30204        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
30205        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
30206        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
30207        assert!(
30208            conv_fuse || hk == num_v,
30209            "de-broadcast requires the fused conv"
30210        );
30211        if conv_fuse {
30212            let f = self.func("ssm_conv1d_gdn_state_vl");
30213            let cfg = LaunchConfig {
30214                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
30215                block_dim: (256, 1, 1),
30216                shared_mem_bytes: 0,
30217            };
30218            let (dsi, nvi, nki, kdi, hki) = (
30219                d_state as i32,
30220                num_v as i32,
30221                num_k as i32,
30222                key_dim as i32,
30223                hk as i32,
30224            );
30225            let __s_lb = self.gpu.stream();
30226            let mut lb = __s_lb.launch_builder(&f);
30227            lb.arg(&v)
30228                .arg(conv_w)
30229                .arg(&cdi)
30230                .arg(&dci)
30231                .arg(&dsi)
30232                .arg(&nvi)
30233                .arg(&nki)
30234                .arg(&kdi)
30235                .arg(&hki);
30236            unsafe {
30237                lb.launch(cfg)?;
30238            }
30239        } else {
30240            let f = self.func("ssm_conv1d_tm_state_vl");
30241            let cfg = LaunchConfig {
30242                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
30243                block_dim: (256, 1, 1),
30244                shared_mem_bytes: 0,
30245            };
30246            let __s_lb = self.gpu.stream();
30247            let mut lb = __s_lb.launch_builder(&f);
30248            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
30249            unsafe {
30250                lb.launch(cfg)?;
30251            }
30252        }
30253        {
30254            let f = self.func("ssm_conv_ring_update_vl");
30255            let n = (conv_dim * (d_conv - 1)) as u32;
30256            let cfg = LaunchConfig {
30257                grid_dim: (n.div_ceil(256), 1, b as u32),
30258                block_dim: (256, 1, 1),
30259                shared_mem_bytes: 0,
30260            };
30261            let __s_lb = self.gpu.stream();
30262            let mut lb = __s_lb.launch_builder(&f);
30263            lb.arg(&v).arg(&cdi).arg(&dci);
30264            unsafe {
30265                lb.launch(cfg)?;
30266            }
30267        }
30268        if !conv_fuse {
30269            let f = self.func("qkv_to_gdn_repack_vl");
30270            let n = max_t * (num_v * d_state) as u32;
30271            let cfg = LaunchConfig {
30272                grid_dim: (n.div_ceil(256), 1, b as u32),
30273                block_dim: (256, 1, 1),
30274                shared_mem_bytes: 0,
30275            };
30276            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
30277            let __s_lb = self.gpu.stream();
30278            let mut lb = __s_lb.launch_builder(&f);
30279            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
30280            unsafe {
30281                lb.launch(cfg)?;
30282            }
30283        }
30284        if Self::l2_v2_on(d_state) {
30285            let f = self.func("gdn_l2_v2_vl");
30286            let cfg = LaunchConfig {
30287                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
30288                block_dim: (256, 1, 1),
30289                shared_mem_bytes: 0,
30290            };
30291            let (dsi, nvi) = (d_state as i32, hk as i32);
30292            let __s_lb = self.gpu.stream();
30293            let mut lb = __s_lb.launch_builder(&f);
30294            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
30295            unsafe {
30296                lb.launch(cfg)?;
30297            }
30298        } else {
30299            let f = self.func("gdn_l2_vl");
30300            let cfg = LaunchConfig {
30301                grid_dim: (max_t * hk as u32, 2, b as u32),
30302                block_dim: (256, 1, 1),
30303                shared_mem_bytes: 0,
30304            };
30305            let (dsi, nvi) = (d_state as i32, hk as i32);
30306            let __s_lb = self.gpu.stream();
30307            let mut lb = __s_lb.launch_builder(&f);
30308            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
30309            unsafe {
30310                lb.launch(cfg)?;
30311            }
30312        }
30313        {
30314            let f = self.func("gdn_gate_prep_vl");
30315            let n = max_t * num_v as u32;
30316            let cfg = LaunchConfig {
30317                grid_dim: (n.div_ceil(256), 1, b as u32),
30318                block_dim: (256, 1, 1),
30319                shared_mem_bytes: 0,
30320            };
30321            let nvi = num_v as i32;
30322            let __s_lb = self.gpu.stream();
30323            let mut lb = __s_lb.launch_builder(&f);
30324            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
30325            unsafe {
30326                lb.launch(cfg)?;
30327            }
30328        }
30329        Ok(())
30330    }
30331
30332    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
30333    pub fn gdn_mirror_vl8(
30334        &self,
30335        seqs: &[GdnSeqVl],
30336        n_head: usize,
30337        which: i32,
30338        hk: usize,
30339    ) -> Result<(), Box<dyn std::error::Error>> {
30340        let b = seqs.len();
30341        assert!((1..=8).contains(&b));
30342        let mut packed = [GdnSeqVl::default(); 8];
30343        packed[..b].copy_from_slice(seqs);
30344        let v = GdnVl8(packed);
30345        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
30346        let max_n = seqs
30347            .iter()
30348            .map(|s| {
30349                if which == 0 {
30350                    s.t as i64 * ept as i64
30351                } else {
30352                    s.nc as i64 * ept as i64 * 32
30353                }
30354            })
30355            .max()
30356            .unwrap();
30357        let f = self.func("gdn_mirror_vl");
30358        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
30359        let cfg = LaunchConfig {
30360            grid_dim: (blocks, 1, b as u32),
30361            block_dim: (256, 1, 1),
30362            shared_mem_bytes: 0,
30363        };
30364        let __s_lb = self.gpu.stream();
30365        let mut lb = __s_lb.launch_builder(&f);
30366        lb.arg(&v).arg(&ept).arg(&which);
30367        unsafe {
30368            lb.launch(cfg)?;
30369        }
30370        Ok(())
30371    }
30372
30373    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
30374    pub fn gdn_tail_vl8(
30375        &self,
30376        seqs: &[GdnPrepVl],
30377        norm_w: &CudaSlice<f32>,
30378        d_state: usize,
30379        num_v: usize,
30380        eps: f32,
30381    ) -> Result<(), Box<dyn std::error::Error>> {
30382        let b = seqs.len();
30383        assert!((1..=8).contains(&b));
30384        let mut packed = [GdnPrepVl::default(); 8];
30385        packed[..b].copy_from_slice(seqs);
30386        let v = GdnPrepVl8(packed);
30387        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
30388        let f = self.func("gated_rmsnorm_f16out_vl");
30389        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
30390        let cfg = LaunchConfig {
30391            grid_dim: (max_t * num_v as u32, 1, b as u32),
30392            block_dim: (128, 1, 1),
30393            shared_mem_bytes: 0,
30394        };
30395        let (dsi, nvi) = (d_state as i32, num_v as i32);
30396        let __s_lb = self.gpu.stream();
30397        let mut lb = __s_lb.launch_builder(&f);
30398        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
30399        unsafe {
30400            lb.launch(cfg)?;
30401        }
30402        Ok(())
30403    }
30404
30405    /// Raw device address helpers for the varlen by-value arg struct (single-stream
30406    /// launches; every buffer outlives the call — the f16 FFI discipline).
30407    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
30408        use cudarc::driver::DevicePtr;
30409        let s = self.gpu.stream();
30410        let (p, _g) = x.device_ptr(&s);
30411        p
30412    }
30413    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
30414        use cudarc::driver::DevicePtrMut;
30415        let s = self.gpu.stream();
30416        let (p, _g) = x.device_ptr_mut(&s);
30417        p
30418    }
30419    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
30420        use cudarc::driver::DevicePtr;
30421        let s = self.gpu.stream();
30422        let (p, _g) = x.device_ptr(&s);
30423        p
30424    }
30425    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
30426        use cudarc::driver::DevicePtr;
30427        let s = self.gpu.stream();
30428        let (p, _g) = x.device_ptr(&s);
30429        p
30430    }
30431
30432    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
30433    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
30434    /// launches, so this is strictly bit-gateable against them).
30435    pub fn gdn_chunk_vl8(
30436        &self,
30437        seqs: &[GdnSeqVl],
30438        n_head: usize,
30439        scale: f32,
30440        hk: usize,
30441        wq: Option<&GdnWVl8>,
30442    ) -> Result<(), Box<dyn std::error::Error>> {
30443        const NSPLIT: u32 = 4;
30444        let b = seqs.len();
30445        assert!((1..=8).contains(&b), "gdn_chunk_vl8: 1..=8 sequences");
30446        let mut packed = [GdnSeqVl::default(); 8];
30447        packed[..b].copy_from_slice(seqs);
30448        let v = GdnVl8(packed);
30449        let (hi, ci) = (n_head as i32, 32i32);
30450        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
30451        let hki = hk as i32;
30452        if let Some(w) = wq {
30453            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
30454            let f = self.func("gdn_k45_wgmma_vl");
30455            let cfg = LaunchConfig {
30456                grid_dim: (n_head as u32, NSPLIT, b as u32),
30457                block_dim: (256, 1, 1),
30458                shared_mem_bytes: 0,
30459            };
30460            let __s_lb = self.gpu.stream();
30461            let mut lb = __s_lb.launch_builder(&f);
30462            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
30463            unsafe {
30464                lb.launch(cfg)?;
30465            }
30466            let _ = max_nc;
30467            return Ok(());
30468        }
30469        {
30470            let f = self.func("gdn_chunk_state_mma_vl");
30471            let cfg = LaunchConfig {
30472                grid_dim: (n_head as u32, NSPLIT, b as u32),
30473                block_dim: (256, 1, 1),
30474                shared_mem_bytes: 0,
30475            };
30476            let __s_lb = self.gpu.stream();
30477            let mut lb = __s_lb.launch_builder(&f);
30478            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
30479            unsafe {
30480                lb.launch(cfg)?;
30481            }
30482        }
30483        {
30484            let f = self.func("gdn_chunk_output_mma_vl");
30485            let cfg = LaunchConfig {
30486                grid_dim: (max_nc, n_head as u32, b as u32),
30487                block_dim: (256, 1, 1),
30488                shared_mem_bytes: 0,
30489            };
30490            let __s_lb = self.gpu.stream();
30491            let mut lb = __s_lb.launch_builder(&f);
30492            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
30493            unsafe {
30494                lb.launch(cfg)?;
30495            }
30496        }
30497        Ok(())
30498    }
30499    #[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
30500    pub fn gdn_scan_chunked(
30501        &self,
30502        q: &CudaSlice<f32>,
30503        k: &CudaSlice<f32>,
30504        v: &CudaSlice<f32>,
30505        g: &CudaSlice<f32>,
30506        beta: &CudaSlice<f32>,
30507        kb16_pre: Option<&CudaSlice<u8>>,
30508        qb16_pre: Option<&CudaSlice<u8>>,
30509        state_in: &CudaSlice<f32>,
30510        state_out: &mut CudaSlice<f32>,
30511        o: &mut CudaSlice<f32>,
30512        n_head: usize,
30513        t: usize,
30514        scale: f32,
30515        c: usize,
30516        hk: usize,
30517    ) -> Result<(), Box<dyn std::error::Error>> {
30518        const D: usize = 128;
30519        const NSPLIT: u32 = 4;
30520        assert!(
30521            (1..=128).contains(&c),
30522            "gdn_scan_chunked: C must be in 1..=128"
30523        );
30524        let h = n_head;
30525        #[allow(clippy::manual_div_ceil)]
30526        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30527        let nc = (t + c - 1) / c;
30528        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
30529        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
30530        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
30531        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
30532        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
30533        let gdn_mma_pre = !portable_mma_gated()
30534            && c == 32
30535            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
30536                Ok("1") => true,
30537                Ok("0") => false,
30538                _ => gdn_mma_default_on(),
30539            };
30540        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
30541            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
30542        } else {
30543            None
30544        };
30545        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
30546        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
30547        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
30548        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
30549        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
30550            && gdn_mma_pre
30551            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
30552        let nk = t * hk * D;
30553        let mut kb16_local: Option<CudaSlice<u8>> = None;
30554        if gdn_mma_pre && kb16_pre.is_none() {
30555            let mut kb = self.alloc_u8_uninit(nk * 2)?;
30556            let f = self.func("f32_to_bf16_bulk");
30557            let n2 = nk as i64;
30558            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
30559            let __s_b = self.gpu.stream();
30560            let mut b = __s_b.launch_builder(&f);
30561            b.arg(k).arg(&mut kb).arg(&n2);
30562            unsafe {
30563                b.launch(cfg2)?;
30564            }
30565            kb16_local = Some(kb);
30566        }
30567        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
30568        if let Some(kb) = kb16_pre {
30569            assert!(kb.len() >= nk * 2, "kb16_pre too small");
30570        }
30571        let mut qb16: Option<CudaSlice<u8>> = None;
30572        let mut pb16: Option<CudaSlice<u8>> = None;
30573        if gdn_wgmma_pre {
30574            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
30575            // the standalone bulk cvt only serves callers without the prep mirror.
30576            if qb16_pre.is_none() {
30577                let mut qb = self.alloc_u8_uninit(nk * 2)?;
30578                let f = self.func("f32_to_bf16_bulk");
30579                let n2 = nk as i64;
30580                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
30581                let __s_b = self.gpu.stream();
30582                let mut b = __s_b.launch_builder(&f);
30583                b.arg(q).arg(&mut qb).arg(&n2);
30584                unsafe {
30585                    b.launch(cfg2)?;
30586                }
30587                qb16 = Some(qb);
30588            } else if let Some(qb) = qb16_pre {
30589                assert!(qb.len() >= nk * 2, "qb16_pre too small");
30590            }
30591            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
30592        }
30593        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
30594        let k2w = if gdn_wgmma_pre {
30595            Some((
30596                *qb16_ref0.as_ref().unwrap(),
30597                *kb16_ref0.as_ref().unwrap(),
30598                pb16.as_mut().unwrap(),
30599            ))
30600        } else {
30601            None
30602        };
30603        let (gcum, p, u, w) =
30604            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
30605        let _ = &w;
30606        let mut y = self.uninit(nc * h * c * D)?;
30607        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
30608        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
30609        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
30610        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
30611        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
30612        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
30613        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
30614        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
30615        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
30616        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
30617        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
30618        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
30619        // sites must agree or the pre-work arms while the scan takes the scalar route.
30620        let gdn_mma = !portable_mma_gated()
30621            && c == 32
30622            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
30623                Ok("1") => true,
30624                Ok("0") => false,
30625                _ => gdn_mma_default_on(),
30626            };
30627        if gdn_mma {
30628            let wb16 = wb16_pre
30629                .take()
30630                .expect("mma path pre-allocates wb16 (K3 store fold)");
30631            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
30632            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
30633            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
30634            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
30635            // pass runs inside the persistent-M kernel; Y and Ssnap are never
30636            // materialized. New numeric class (gk folds into k^T instead of ys) —
30637            // explicit opt-in until the state-carry battery promotes it. Env read per
30638            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
30639            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
30640            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
30641            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
30642            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
30643            if gdn_wgmma_pre {
30644                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
30645                let qb16 = qb16_ref0.unwrap();
30646                let pb16 = pb16.as_ref().unwrap();
30647                {
30648                    let f = self.func("gdn_k45_wgmma");
30649                    let cfg = LaunchConfig {
30650                        grid_dim: (h as u32, 4, 1),
30651                        block_dim: (256, 1, 1),
30652                        shared_mem_bytes: 0,
30653                    };
30654                    let hki = hk as i32;
30655                    let __s_b = self.gpu.stream();
30656                    let mut b = __s_b.launch_builder(&f);
30657                    b.arg(kb16_ref)
30658                        .arg(&gcum)
30659                        .arg(beta)
30660                        .arg(&u)
30661                        .arg(&wb16)
30662                        .arg(qb16)
30663                        .arg(pb16)
30664                        .arg(o)
30665                        .arg(&scale)
30666                        .arg(state_in)
30667                        .arg(&mut *state_out)
30668                        .arg(&hi)
30669                        .arg(&ti)
30670                        .arg(&ci)
30671                        .arg(&hki);
30672                    unsafe {
30673                        b.launch(cfg)?;
30674                    }
30675                }
30676                return Ok(());
30677            }
30678            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
30679            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
30680            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
30681            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
30682            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
30683            {
30684                let f = self.func("gdn_chunk_state_mma");
30685                let cfg = LaunchConfig {
30686                    grid_dim: (h as u32, NSPLIT, 1),
30687                    block_dim: (256, 1, 1),
30688                    shared_mem_bytes: 0,
30689                };
30690                let hki = hk as i32;
30691                let __s_b = self.gpu.stream();
30692                let mut b = __s_b.launch_builder(&f);
30693                b.arg(kb16_ref)
30694                    .arg(&gcum)
30695                    .arg(beta)
30696                    .arg(&u)
30697                    .arg(&wb16)
30698                    .arg(&mut y16)
30699                    .arg(&mut ssnap16)
30700                    .arg(state_in)
30701                    .arg(&mut *state_out)
30702                    .arg(&hi)
30703                    .arg(&ti)
30704                    .arg(&ci)
30705                    .arg(&hki);
30706                unsafe {
30707                    b.launch(cfg)?;
30708                }
30709            }
30710            {
30711                // K5-mma (bf16 St/Y consumers)
30712                let f = self.func("gdn_chunk_output_mma");
30713                #[allow(clippy::manual_div_ceil)]
30714                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30715                let jt = ((c + 31) / 32) as u32;
30716                let cfg = LaunchConfig {
30717                    grid_dim: (nc as u32, h as u32, jt),
30718                    block_dim: (256, 1, 1),
30719                    shared_mem_bytes: 0,
30720                };
30721                let hki = hk as i32;
30722                let __s_b = self.gpu.stream();
30723                let mut b = __s_b.launch_builder(&f);
30724                b.arg(q)
30725                    .arg(&gcum)
30726                    .arg(&p)
30727                    .arg(&y16)
30728                    .arg(&ssnap16)
30729                    .arg(o)
30730                    .arg(&hi)
30731                    .arg(&ti)
30732                    .arg(&ci)
30733                    .arg(&scale)
30734                    .arg(&hki);
30735                unsafe {
30736                    b.launch(cfg)?;
30737                }
30738            }
30739            return Ok(());
30740        }
30741        {
30742            // K4 (sequential over chunks inside; blocks col-partition the state)
30743            let f = self.func("gdn_chunk_state_f32");
30744            let cfg = LaunchConfig {
30745                grid_dim: (h as u32, NSPLIT, 1),
30746                block_dim: (256, 1, 1),
30747                shared_mem_bytes: 0,
30748            };
30749            let __s_b = self.gpu.stream();
30750            let mut b = __s_b.launch_builder(&f);
30751            b.arg(k)
30752                .arg(&gcum)
30753                .arg(beta)
30754                .arg(&u)
30755                .arg(&w)
30756                .arg(&mut y)
30757                .arg(&mut ssnap)
30758                .arg(state_in)
30759                .arg(&mut *state_out)
30760                .arg(&hi)
30761                .arg(&ti)
30762                .arg(&ci);
30763            unsafe {
30764                b.launch(cfg)?;
30765            }
30766        }
30767        {
30768            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
30769            let f = self.func("gdn_chunk_output_f32");
30770            #[allow(clippy::manual_div_ceil)]
30771            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30772            let jt = ((c + 31) / 32) as u32;
30773            let cfg = LaunchConfig {
30774                grid_dim: (nc as u32, h as u32, jt),
30775                block_dim: (256, 1, 1),
30776                shared_mem_bytes: 0,
30777            };
30778            let __s_b = self.gpu.stream();
30779            let mut b = __s_b.launch_builder(&f);
30780            b.arg(q)
30781                .arg(&gcum)
30782                .arg(&p)
30783                .arg(&y)
30784                .arg(&ssnap)
30785                .arg(o)
30786                .arg(&hi)
30787                .arg(&ti)
30788                .arg(&ci)
30789                .arg(&scale);
30790            unsafe {
30791                b.launch(cfg)?;
30792            }
30793        }
30794        Ok(())
30795    }
30796
30797    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
30798    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
30799    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
30800    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
30801    ///
30802    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
30803    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
30804    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
30805    #[allow(clippy::too_many_arguments)]
30806    #[allow(clippy::too_many_arguments)]
30807    pub fn gdn_scan_prefill(
30808        &self,
30809        q: &CudaSlice<f32>,
30810        k: &CudaSlice<f32>,
30811        v: &CudaSlice<f32>,
30812        g: &CudaSlice<f32>,
30813        beta: &CudaSlice<f32>,
30814        kb16_pre: Option<&CudaSlice<u8>>,
30815        qb16_pre: Option<&CudaSlice<u8>>,
30816        state_in: &CudaSlice<f32>,
30817        state_out: &mut CudaSlice<f32>,
30818        o: &mut CudaSlice<f32>,
30819        n_head: usize,
30820        t: usize,
30821        scale: f32,
30822        hk: usize,
30823    ) -> Result<(), Box<dyn std::error::Error>> {
30824        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
30825            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
30826            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
30827        }
30828        if Self::gdn_chunked_enabled() && t >= 16 {
30829            self.gdn_scan_chunked(
30830                q,
30831                k,
30832                v,
30833                g,
30834                beta,
30835                kb16_pre,
30836                qb16_pre,
30837                state_in,
30838                state_out,
30839                o,
30840                n_head,
30841                t,
30842                scale,
30843                Self::gdn_chunk_size(),
30844                hk,
30845            )
30846        } else {
30847            assert!(
30848                hk == n_head,
30849                "s128 scan is broadcast-only (prep guarantees by predicate)"
30850            );
30851            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
30852        }
30853    }
30854
30855    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
30856    #[allow(clippy::too_many_arguments)]
30857    fn gdn_scan_diff(
30858        &self,
30859        q: &CudaSlice<f32>,
30860        k: &CudaSlice<f32>,
30861        v: &CudaSlice<f32>,
30862        g: &CudaSlice<f32>,
30863        beta: &CudaSlice<f32>,
30864        state_in: &CudaSlice<f32>,
30865        state_out: &mut CudaSlice<f32>,
30866        o: &mut CudaSlice<f32>,
30867        n_head: usize,
30868        t: usize,
30869        scale: f32,
30870    ) -> Result<(), Box<dyn std::error::Error>> {
30871        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
30872        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
30873        let mut o_c = self.uninit(o.len())?;
30874        let mut st_c = self.uninit(state_out.len())?;
30875        self.gdn_scan_chunked(
30876            q,
30877            k,
30878            v,
30879            g,
30880            beta,
30881            None,
30882            None,
30883            state_in,
30884            &mut st_c,
30885            &mut o_c,
30886            n_head,
30887            t,
30888            scale,
30889            Self::gdn_chunk_size(),
30890            n_head,
30891        )?;
30892        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
30893        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
30894        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
30895        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
30896            let mut max_abs = 0f32;
30897            let mut max_rel = 0f32;
30898            let mut sum_rel = 0f64;
30899            for (x, y) in a.iter().zip(b) {
30900                let ad = (x - y).abs();
30901                let rel = ad / x.abs().max(y.abs()).max(1e-3);
30902                if ad > max_abs {
30903                    max_abs = ad;
30904                }
30905                if rel > max_rel {
30906                    max_rel = rel;
30907                }
30908                sum_rel += rel as f64;
30909            }
30910            (max_abs, max_rel, sum_rel / a.len() as f64)
30911        };
30912        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
30913        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
30914        println!(
30915            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
30916                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
30917            Self::gdn_chunk_size()
30918        );
30919        Ok(())
30920    }
30921
30922    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
30923    pub fn gdn_glog(
30924        &self,
30925        alpha: &CudaSlice<f32>,
30926        dt_bias: &CudaSlice<f32>,
30927        a: &CudaSlice<f32>,
30928        g_log: &mut CudaSlice<f32>,
30929        n_head: usize,
30930        t: usize,
30931    ) -> Result<(), Box<dyn std::error::Error>> {
30932        let f = self.func("gdn_glog_f32");
30933        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
30934        let (h, ti) = (n_head as i32, t as i32);
30935        let __s_b = self.gpu.stream();
30936        let mut b = __s_b.launch_builder(&f);
30937        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
30938        unsafe {
30939            b.launch(cfg)?;
30940        }
30941        Ok(())
30942    }
30943
30944    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
30945    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
30946    pub fn sigmoid_v(
30947        &self,
30948        x: &cudarc::driver::CudaView<f32>,
30949        y: &mut CudaSlice<f32>,
30950        n: usize,
30951    ) -> Result<(), Box<dyn std::error::Error>> {
30952        let f = self.func("sigmoid_f32");
30953        let cfg = LaunchConfig::for_num_elems(n as u32);
30954        let ni = n as i32;
30955        let __s_b = self.gpu.stream();
30956        let mut b = __s_b.launch_builder(&f);
30957        b.arg(x).arg(y).arg(&ni);
30958        unsafe {
30959            b.launch(cfg)?;
30960        }
30961        Ok(())
30962    }
30963
30964    pub fn gdn_glog_v(
30965        &self,
30966        alpha: &cudarc::driver::CudaView<f32>,
30967        dt_bias: &CudaSlice<f32>,
30968        a: &CudaSlice<f32>,
30969        g_log: &mut CudaSlice<f32>,
30970        n_head: usize,
30971        t: usize,
30972    ) -> Result<(), Box<dyn std::error::Error>> {
30973        let f = self.func("gdn_glog_f32");
30974        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
30975        let (h, ti) = (n_head as i32, t as i32);
30976        let __s_b = self.gpu.stream();
30977        let mut b = __s_b.launch_builder(&f);
30978        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
30979        unsafe {
30980            b.launch(cfg)?;
30981        }
30982        Ok(())
30983    }
30984
30985    pub fn sigmoid(
30986        &self,
30987        x: &CudaSlice<f32>,
30988        y: &mut CudaSlice<f32>,
30989        n: usize,
30990    ) -> Result<(), Box<dyn std::error::Error>> {
30991        let f = self.func("sigmoid_f32");
30992        let cfg = LaunchConfig::for_num_elems(n as u32);
30993        let ni = n as i32;
30994        let __s_b = self.gpu.stream();
30995        let mut b = __s_b.launch_builder(&f);
30996        b.arg(x).arg(y).arg(&ni);
30997        unsafe {
30998            b.launch(cfg)?;
30999        }
31000        Ok(())
31001    }
31002
31003    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
31004    /// (replaces sigmoid + mul + convert). Bit-identical class.
31005    pub fn sig_mul_f16out(
31006        &self,
31007        a: &CudaSlice<f32>,
31008        g: &CudaSlice<f32>,
31009        dst: &mut CudaSlice<f32>,
31010        dst16: &mut CudaSlice<u8>,
31011        n: usize,
31012    ) -> Result<(), Box<dyn std::error::Error>> {
31013        let f = self.func("sig_mul_f16out_f32");
31014        let cfg = LaunchConfig::for_num_elems(n as u32);
31015        let ni = n as i32;
31016        let __s_b = self.gpu.stream();
31017        let mut b = __s_b.launch_builder(&f);
31018        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
31019        unsafe {
31020            b.launch(cfg)?;
31021        }
31022        Ok(())
31023    }
31024
31025    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
31026    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
31027    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
31028    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
31029    ///
31030    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
31031    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
31032    /// applies the wrong number of distinct gate values.
31033    #[allow(clippy::too_many_arguments)]
31034    pub fn attn_head_gate(
31035        &self,
31036        a: &CudaSlice<f32>,
31037        g: &CudaSlice<f32>,
31038        dst: &mut CudaSlice<f32>,
31039        dst16: Option<&mut CudaSlice<u8>>,
31040        head_dim: usize,
31041        n_head: usize,
31042        t: usize,
31043    ) -> Result<(), Box<dyn std::error::Error>> {
31044        let f = self.func("attn_head_gate_f32");
31045        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
31046        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
31047        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
31048        let d16: u64 = match dst16 {
31049            Some(d) => self.addr_u8(d),
31050            None => 0,
31051        };
31052        let __s_b = self.gpu.stream();
31053        let mut b = __s_b.launch_builder(&f);
31054        b.arg(a)
31055            .arg(g)
31056            .arg(dst)
31057            .arg(&d16)
31058            .arg(&hd)
31059            .arg(&nh)
31060            .arg(&ti);
31061        unsafe {
31062            b.launch(cfg)?;
31063        }
31064        Ok(())
31065    }
31066
31067    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
31068    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
31069    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
31070    ///
31071    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
31072    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
31073    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
31074    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
31075    #[allow(clippy::too_many_arguments)]
31076    pub fn swiglu_clamped_mul_scaled(
31077        &self,
31078        gate: &CudaSlice<f32>,
31079        up: &CudaSlice<f32>,
31080        gs: f32,
31081        us: f32,
31082        limit: f32,
31083        dst: &mut CudaSlice<f32>,
31084        n: usize,
31085    ) -> Result<(), Box<dyn std::error::Error>> {
31086        debug_assert!(
31087            limit > 1e-6,
31088            "swiglu_clamped needs a live limit; use silu_mul_scaled"
31089        );
31090        let f = self.func("swiglu_clamped_mul_scaled_f32");
31091        let cfg = LaunchConfig::for_num_elems(n as u32);
31092        let ni = n as i32;
31093        let __s_b = self.gpu.stream();
31094        let mut b = __s_b.launch_builder(&f);
31095        b.arg(gate)
31096            .arg(up)
31097            .arg(&gs)
31098            .arg(&us)
31099            .arg(&limit)
31100            .arg(dst)
31101            .arg(&ni);
31102        unsafe {
31103            b.launch(cfg)?;
31104        }
31105        Ok(())
31106    }
31107
31108    /// glm5_next PRE-clamped SwiGLU: `dst = silu(min(gate*gs, limit)) * clamp(up*us, +-limit)`.
31109    /// The gate clamp is BEFORE silu and one-sided — vendor `Glm5NextTextMLP.forward` /
31110    /// `Glm5NextTextExperts._apply_gate`, one `swiglu_limit` shared by the dense MLP, the routed
31111    /// experts and the shared expert on every layer.
31112    ///
31113    /// This is NOT `swiglu_clamped_mul_scaled` (step35 clamps the silu OUTPUT) and NOT
31114    /// `swigluoai_mul_scaled` (alpha-swish plus a `1 +` linear term). Same caller contract as the
31115    /// post-clamp sibling: `limit > 1e-6`, else the plain `silu_mul_scaled` path.
31116    #[allow(clippy::too_many_arguments)]
31117    pub fn swiglu_preclamped_mul_scaled(
31118        &self,
31119        gate: &CudaSlice<f32>,
31120        up: &CudaSlice<f32>,
31121        gs: f32,
31122        us: f32,
31123        limit: f32,
31124        dst: &mut CudaSlice<f32>,
31125        n: usize,
31126    ) -> Result<(), Box<dyn std::error::Error>> {
31127        debug_assert!(
31128            limit > 1e-6,
31129            "swiglu_preclamped needs a live limit; use silu_mul_scaled"
31130        );
31131        let f = self.func("swiglu_preclamped_mul_scaled_f32");
31132        let cfg = LaunchConfig::for_num_elems(n as u32);
31133        let ni = n as i32;
31134        let __s_b = self.gpu.stream();
31135        let mut b = __s_b.launch_builder(&f);
31136        b.arg(gate)
31137            .arg(up)
31138            .arg(&gs)
31139            .arg(&us)
31140            .arg(&limit)
31141            .arg(dst)
31142            .arg(&ni);
31143        unsafe {
31144            b.launch(cfg)?;
31145        }
31146        Ok(())
31147    }
31148
31149    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
31150    #[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
31151    pub fn gated_rmsnorm(
31152        &self,
31153        o: &CudaSlice<f32>,
31154        w: &CudaSlice<f32>,
31155        z: &CudaSlice<f32>,
31156        dst: &mut CudaSlice<f32>,
31157        ncols: usize,
31158        nrows: usize,
31159        eps: f32,
31160    ) -> Result<(), Box<dyn std::error::Error>> {
31161        let f = self.func("gated_rmsnorm_f32");
31162        let cfg = LaunchConfig {
31163            grid_dim: (nrows as u32, 1, 1),
31164            block_dim: (128, 1, 1),
31165            shared_mem_bytes: 0,
31166        };
31167        let (nc, e) = (ncols as i32, eps);
31168        let __s_b = self.gpu.stream();
31169        let mut b = __s_b.launch_builder(&f);
31170        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
31171        unsafe {
31172            b.launch(cfg)?;
31173        }
31174        Ok(())
31175    }
31176
31177    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
31178    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
31179    #[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
31180    pub fn gated_rmsnorm_f16out(
31181        &self,
31182        o: &CudaSlice<f32>,
31183        w: &CudaSlice<f32>,
31184        z: &CudaSlice<f32>,
31185        dst: &mut CudaSlice<f32>,
31186        dst16: &mut CudaSlice<u8>,
31187        ncols: usize,
31188        nrows: usize,
31189        eps: f32,
31190    ) -> Result<(), Box<dyn std::error::Error>> {
31191        let f = self.func("gated_rmsnorm_f16out_f32");
31192        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
31193        let cfg = LaunchConfig {
31194            grid_dim: (nrows as u32, 1, 1),
31195            block_dim: (128, 1, 1),
31196            shared_mem_bytes: 0,
31197        };
31198        let (nc, e) = (ncols as i32, eps);
31199        let __s_b = self.gpu.stream();
31200        let mut b = __s_b.launch_builder(&f);
31201        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
31202        unsafe {
31203            b.launch(cfg)?;
31204        }
31205        Ok(())
31206    }
31207
31208    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
31209    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
31210    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
31211    #[allow(clippy::too_many_arguments)]
31212    pub fn add_rms_norm_zq8(
31213        &self,
31214        a: &CudaSlice<f32>,
31215        b_in: &CudaSlice<f32>,
31216        w: &CudaSlice<f32>,
31217        res: &mut CudaSlice<f32>,
31218        z: &mut CudaSlice<f32>,
31219        ncols: usize,
31220        nrows: usize,
31221        eps: f32,
31222    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
31223        assert!(ncols.is_multiple_of(32));
31224        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
31225        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
31226        let f = self.func("add_rms_norm_zq8");
31227        let cfg = LaunchConfig {
31228            grid_dim: (nrows as u32, 1, 1),
31229            block_dim: (1024, 1, 1),
31230            shared_mem_bytes: 0,
31231        };
31232        let (nc, ep) = (ncols as i32, eps);
31233        let __s_b = self.gpu.stream();
31234        let mut b = __s_b.launch_builder(&f);
31235        b.arg(a)
31236            .arg(b_in)
31237            .arg(w)
31238            .arg(res)
31239            .arg(z)
31240            .arg(&mut q)
31241            .arg(&mut d)
31242            .arg(&nc)
31243            .arg(&ep);
31244        unsafe {
31245            b.launch(cfg)?;
31246        }
31247        Ok((q, d))
31248    }
31249
31250    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
31251    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
31252    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
31253    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
31254    #[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
31255    pub fn gated_rmsnorm_zv(
31256        &self,
31257        o: &CudaSlice<f32>,
31258        w: &CudaSlice<f32>,
31259        z: &cudarc::driver::CudaView<f32>,
31260        dst: &mut CudaSlice<f32>,
31261        ncols: usize,
31262        nrows: usize,
31263        eps: f32,
31264    ) -> Result<(), Box<dyn std::error::Error>> {
31265        let f = self.func("gated_rmsnorm_f32");
31266        let cfg = LaunchConfig {
31267            grid_dim: (nrows as u32, 1, 1),
31268            block_dim: (128, 1, 1),
31269            shared_mem_bytes: 0,
31270        };
31271        let (nc, e) = (ncols as i32, eps);
31272        let __s_b = self.gpu.stream();
31273        let mut b = __s_b.launch_builder(&f);
31274        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
31275        unsafe {
31276            b.launch(cfg)?;
31277        }
31278        Ok(())
31279    }
31280
31281    #[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
31282    pub fn gated_rmsnorm_f16out_zv(
31283        &self,
31284        o: &CudaSlice<f32>,
31285        w: &CudaSlice<f32>,
31286        z: &cudarc::driver::CudaView<f32>,
31287        dst: &mut CudaSlice<f32>,
31288        dst16: &mut CudaSlice<u8>,
31289        ncols: usize,
31290        nrows: usize,
31291        eps: f32,
31292    ) -> Result<(), Box<dyn std::error::Error>> {
31293        let f = self.func("gated_rmsnorm_f16out_f32");
31294        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
31295        let cfg = LaunchConfig {
31296            grid_dim: (nrows as u32, 1, 1),
31297            block_dim: (128, 1, 1),
31298            shared_mem_bytes: 0,
31299        };
31300        let (nc, e) = (ncols as i32, eps);
31301        let __s_b = self.gpu.stream();
31302        let mut b = __s_b.launch_builder(&f);
31303        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
31304        unsafe {
31305            b.launch(cfg)?;
31306        }
31307        Ok(())
31308    }
31309
31310    pub fn gated_rmsnorm_q8_1(
31311        &self,
31312        o: &CudaSlice<f32>,
31313        w: &CudaSlice<f32>,
31314        z: &CudaSlice<f32>,
31315        ncols: usize,
31316        nrows: usize,
31317        eps: f32,
31318    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
31319        assert!(ncols.is_multiple_of(32));
31320        let f = self.func("gated_rmsnorm_q8_1");
31321        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
31322        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
31323        let cfg = LaunchConfig {
31324            grid_dim: (nrows as u32, 1, 1),
31325            block_dim: (128, 1, 1),
31326            shared_mem_bytes: 0,
31327        };
31328        let (nc, ep) = (ncols as i32, eps);
31329        let __s_b = self.gpu.stream();
31330        let mut b = __s_b.launch_builder(&f);
31331        b.arg(o)
31332            .arg(w)
31333            .arg(z)
31334            .arg(&mut out_q)
31335            .arg(&mut out_d)
31336            .arg(&nc)
31337            .arg(&ep);
31338        unsafe {
31339            b.launch(cfg)?;
31340        }
31341        Ok((out_q, out_d))
31342    }
31343
31344    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
31345    pub fn transpose(
31346        &self,
31347        inp: &CudaSlice<f32>,
31348        rows: usize,
31349        cols: usize,
31350    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31351        let f = self.func("transpose_f32");
31352        let mut out = self.zeros(rows * cols)?;
31353        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
31354        let (r, c) = (rows as i32, cols as i32);
31355        let __s_b = self.gpu.stream();
31356        let mut b = __s_b.launch_builder(&f);
31357        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
31358        unsafe {
31359            b.launch(cfg)?;
31360        }
31361        Ok(out)
31362    }
31363
31364    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
31365    pub fn repeat_heads(
31366        &self,
31367        inp: &CudaSlice<f32>,
31368        out: &mut CudaSlice<f32>,
31369        head_dim: usize,
31370        n_in: usize,
31371        n_out: usize,
31372        t: usize,
31373    ) -> Result<(), Box<dyn std::error::Error>> {
31374        let f = self.func("repeat_heads_f32");
31375        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
31376        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
31377        let __s_b = self.gpu.stream();
31378        let mut b = __s_b.launch_builder(&f);
31379        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
31380        unsafe {
31381            b.launch(cfg)?;
31382        }
31383        Ok(())
31384    }
31385
31386    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
31387    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
31388    ///
31389    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
31390    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
31391    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
31392    pub fn q_gate_split(
31393        &self,
31394        qf: &CudaSlice<f32>,
31395        q_out: &mut CudaSlice<f32>,
31396        gate_out: &mut CudaSlice<f32>,
31397        head_dim: usize,
31398        n_head: usize,
31399        t: usize,
31400    ) -> Result<(), Box<dyn std::error::Error>> {
31401        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
31402        let out_need = head_dim * n_head * t;
31403        if q_out.len() < out_need || gate_out.len() < out_need {
31404            return Err(format!(
31405                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
31406                q_out.len(),
31407                gate_out.len()
31408            )
31409            .into());
31410        }
31411        let f = self.func("q_gate_split_f32");
31412        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
31413        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
31414        let __s_b = self.gpu.stream();
31415        let mut b = __s_b.launch_builder(&f);
31416        b.arg(qf)
31417            .arg(q_out)
31418            .arg(gate_out)
31419            .arg(&hd)
31420            .arg(&nh)
31421            .arg(&ti);
31422        unsafe {
31423            b.launch(cfg)?;
31424        }
31425        Ok(())
31426    }
31427
31428    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
31429    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
31430    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
31431    #[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
31432    pub fn qkv_to_gdn_repack(
31433        &self,
31434        conv_out: &CudaSlice<f32>,
31435        q_g: &mut CudaSlice<f32>,
31436        k_g: &mut CudaSlice<f32>,
31437        v_g: &mut CudaSlice<f32>,
31438        d_state: usize,
31439        num_v: usize,
31440        num_k: usize,
31441        key_dim: usize,
31442        t: usize,
31443    ) -> Result<(), Box<dyn std::error::Error>> {
31444        let f = self.func("qkv_to_gdn_repack_f32");
31445        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
31446        let (ds, nv, nk, kd, ti) = (
31447            d_state as i32,
31448            num_v as i32,
31449            num_k as i32,
31450            key_dim as i32,
31451            t as i32,
31452        );
31453        let __s_b = self.gpu.stream();
31454        let mut b = __s_b.launch_builder(&f);
31455        b.arg(conv_out)
31456            .arg(q_g)
31457            .arg(k_g)
31458            .arg(v_g)
31459            .arg(&ds)
31460            .arg(&nv)
31461            .arg(&nk)
31462            .arg(&kd)
31463            .arg(&ti);
31464        unsafe {
31465            b.launch(cfg)?;
31466        }
31467        Ok(())
31468    }
31469
31470    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
31471    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
31472    pub fn conv_left_pad(
31473        &self,
31474        src: &CudaSlice<f32>,
31475        dst: &mut CudaSlice<f32>,
31476        conv_dim: usize,
31477        t: usize,
31478        pad: usize,
31479    ) -> Result<(), Box<dyn std::error::Error>> {
31480        let f = self.func("conv_left_pad_f32");
31481        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
31482        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
31483        let __s_b = self.gpu.stream();
31484        let mut b = __s_b.launch_builder(&f);
31485        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
31486        unsafe {
31487            b.launch(cfg)?;
31488        }
31489        Ok(())
31490    }
31491
31492    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
31493    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
31494    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
31495    pub fn conv_assemble_and_roll(
31496        &self,
31497        qkv_col: &CudaSlice<f32>,
31498        conv_state: &mut CudaSlice<f32>,
31499        conv_in: &mut CudaSlice<f32>,
31500        conv_dim: usize,
31501        pad: usize,
31502    ) -> Result<(), Box<dyn std::error::Error>> {
31503        let f = self.func("conv_assemble_and_roll_f32");
31504        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
31505        let (cd, p) = (conv_dim as i32, pad as i32);
31506        let __s_b = self.gpu.stream();
31507        let mut b = __s_b.launch_builder(&f);
31508        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
31509        unsafe {
31510            b.launch(cfg)?;
31511        }
31512        Ok(())
31513    }
31514
31515    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
31516    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
31517    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
31518    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
31519    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
31520    pub fn ssm_conv1d_fused_decode(
31521        &self,
31522        qkv_col: &CudaSlice<f32>,
31523        conv_state: &mut CudaSlice<f32>,
31524        w: &CudaSlice<f32>,
31525        conv_out: &mut CudaSlice<f32>,
31526        conv_dim: usize,
31527        d_conv: usize,
31528    ) -> Result<(), Box<dyn std::error::Error>> {
31529        let f = self.func("ssm_conv1d_fused_decode_f32");
31530        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
31531        let (cd, dc) = (conv_dim as i32, d_conv as i32);
31532        let __s_b = self.gpu.stream();
31533        let mut b = __s_b.launch_builder(&f);
31534        b.arg(qkv_col)
31535            .arg(conv_state)
31536            .arg(w)
31537            .arg(conv_out)
31538            .arg(&cd)
31539            .arg(&dc);
31540        unsafe {
31541            b.launch(cfg)?;
31542        }
31543        Ok(())
31544    }
31545
31546    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
31547    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
31548    pub fn slice_range(
31549        &self,
31550        src: &CudaSlice<f32>,
31551        start: usize,
31552        len: usize,
31553    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31554        let host = self.gpu.stream().clone_dtoh(src)?;
31555        self.gpu.stream().synchronize()?;
31556        self.htod(&host[start..start + len])
31557    }
31558}
31559
31560#[cfg(test)]
31561mod target_dispatch_tests {
31562    use super::legacy_quant_gemm_allowed;
31563
31564    #[test]
31565    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
31566        // sm_120a native lane
31567        assert!(legacy_quant_gemm_allowed(false, false, false));
31568        assert!(!legacy_quant_gemm_allowed(false, false, true));
31569        // pure portable lane (sm_89): gated
31570        assert!(!legacy_quant_gemm_allowed(true, false, false));
31571        assert!(!legacy_quant_gemm_allowed(true, false, true));
31572        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
31573        assert!(legacy_quant_gemm_allowed(true, true, false));
31574        assert!(!legacy_quant_gemm_allowed(true, true, true));
31575    }
31576
31577    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
31578    #[test]
31579    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
31580        assert!(!legacy_quant_gemm_allowed(
31581            cfg!(memra_portable_cuda),
31582            cfg!(memra_hopper_mma),
31583            false
31584        ));
31585    }
31586
31587    #[cfg(memra_hopper_mma)]
31588    #[test]
31589    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
31590        assert!(legacy_quant_gemm_allowed(
31591            cfg!(memra_portable_cuda),
31592            cfg!(memra_hopper_mma),
31593            false
31594        ));
31595        assert!(super::portable_mma_gated() == false);
31596    }
31597}
31598
31599/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
31600/// inherent methods (inherent methods win name resolution, so no recursion).
31601impl memra_kv::KvDev for Engine {
31602    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31603        Engine::zeros(self, n)
31604    }
31605    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31606        Engine::uninit(self, n)
31607    }
31608    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
31609        Engine::alloc_u8(self, n)
31610    }
31611    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
31612        Engine::htod_i32(self, v)
31613    }
31614    fn clone_dtod(
31615        &self,
31616        src: &CudaSlice<f32>,
31617    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
31618        Engine::clone_dtod(self, src)
31619    }
31620    fn copy_into(
31621        &self,
31622        dst: &mut CudaSlice<f32>,
31623        off: usize,
31624        src: &CudaSlice<f32>,
31625        len: usize,
31626    ) -> Result<(), Box<dyn std::error::Error>> {
31627        Engine::copy_into(self, dst, off, src, len)
31628    }
31629    fn copy_range_into(
31630        &self,
31631        dst: &mut CudaSlice<f32>,
31632        dst_off: usize,
31633        src: &CudaSlice<f32>,
31634        src_off: usize,
31635        len: usize,
31636    ) -> Result<(), Box<dyn std::error::Error>> {
31637        Engine::copy_range_into(self, dst, dst_off, src, src_off, len)
31638    }
31639    fn set_i32_one(
31640        &self,
31641        d: &mut CudaSlice<i32>,
31642        v: i32,
31643    ) -> Result<(), Box<dyn std::error::Error>> {
31644        Engine::set_i32_one(self, d, v)
31645    }
31646}
31647
31648#[cfg(test)]
31649mod fused_gate_bounds_tests {
31650    use super::*;
31651
31652    /// The fused `[q|gate]` split's read-site guard, on the device.
31653    ///
31654    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
31655    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
31656    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
31657    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
31658    /// `FusedQGateExtent` before the launch.
31659    ///
31660    /// Catch demonstration for this test (guard temporarily removed, then restored):
31661    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
31662    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
31663    /// the call returns `Err`. Receipt in the lane report.
31664    #[test]
31665    #[ignore = "requires a CUDA GPU"]
31666    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
31667        let e = Engine::new(0).unwrap();
31668        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
31669        let fused = 2 * head_dim * n_head * t;
31670        let out_n = head_dim * n_head * t;
31671
31672        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
31673        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
31674        let mut q = e.uninit(out_n).unwrap();
31675        let mut gate = e.uninit(out_n).unwrap();
31676        let err = e
31677            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
31678            .expect_err("half-width wq must be refused, not read past")
31679            .to_string();
31680        assert!(err.contains("NO fused gate"), "{err}");
31681        assert!(err.contains(&format!("{fused}")), "{err}");
31682
31683        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
31684        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
31685        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
31686        let wide = e.htod(&host).unwrap();
31687        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
31688            .expect("full-width wq splits");
31689        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
31690        for tok in 0..t {
31691            for hh in 0..n_head {
31692                for d in 0..head_dim {
31693                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
31694                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
31695                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
31696                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
31697                }
31698            }
31699        }
31700
31701        // undersized destinations are refused too (the other half of the extent contract)
31702        let mut small = e.uninit(out_n - 1).unwrap();
31703        assert!(
31704            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
31705                .is_err()
31706        );
31707    }
31708}
31709
31710/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
31711/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
31712/// any launch, so the refusal is testable without a device.
31713#[cfg(test)]
31714mod fused_rope_width_tests {
31715    use super::Engine;
31716
31717    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
31718    /// safetensors route derives the same), which is why the fusion is legal there today.
31719    #[test]
31720    fn full_width_is_accepted() {
31721        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
31722        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
31723        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
31724    }
31725
31726    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
31727    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
31728    ///
31729    /// ```text
31730    /// attention.key_length     512   rope.dimension_count     512   (global class)
31731    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
31732    /// ```
31733    ///
31734    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
31735    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
31736    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
31737    /// instead of a silently over-rotated head.
31738    #[test]
31739    fn gemma4_official_artifact_widths_pass() {
31740        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
31741        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
31742    }
31743
31744    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
31745    /// with no `n_dims`, silently rotating the pass-through band.
31746    #[test]
31747    fn partial_rotary_is_refused_with_the_geometry_named() {
31748        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
31749        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
31750            .expect_err("partial rotary must refuse");
31751        let msg = err.to_string();
31752        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
31753        assert!(msg.contains("n_rot 64"), "{msg}");
31754        assert!(msg.contains("head_dim 256"), "{msg}");
31755        assert!(
31756            msg.contains("64..256"),
31757            "names the band it would corrupt: {msg}"
31758        );
31759        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
31760        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
31761        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
31762        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
31763    }
31764}