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 std::sync::{Arc, Mutex};
4use cudarc::driver::{CudaContext, CudaStream, CudaModule, CudaFunction, CudaSlice, LaunchConfig, PushKernelArg};
5use cudarc::nvrtc::Ptx;
6
7pub use memra_gguf;
8pub use memra_runtime;
9
10pub mod model;
11pub mod forward;
12pub mod hybrid;
13pub mod hybrid_forward;
14/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
15/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
16pub mod cache {
17    pub use memra_kv::*;
18}
19pub mod decode;
20pub mod decode_batch;
21/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
22/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
23/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
24pub mod mla;
25pub mod pp;
26pub mod spec;
27pub mod gemma_spec;
28pub mod round_stream;
29pub mod graph_update;
30pub mod dflash;
31pub mod eagle;
32pub use memra_sampling as sampler;
33
34/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
35/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
36/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
37/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
38/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
39///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
40///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
41///                     stream sync per projection (round-47 ledgered defect).
42///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
43///                     construction, zero syncs, f32 C with the act row-scale folded in.
44/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
45/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
46/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
47/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
48/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
49/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
50///
51/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
52/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
53/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
54/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
55/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
56/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
57/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
58/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
59///
60/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
61/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
62/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
63/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
64/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
65/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
66/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
67///
68/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
69/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
70/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
71/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
72/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
73/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
74/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
75/// the k-quant-only admission survives as the rollback seam, not the default.
76/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
77pub fn moe_f16g_mode() -> u8 {
78    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
79    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
80        Ok("0") => 0,
81        Ok("2") => 2,
82        Ok("3") => 3,
83        Ok(_) => 1,
84        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
85        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
86        Err(_) => 2,
87    })
88}
89/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
90/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
91/// (shape_sel, cross) for the FFI:
92///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
93///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
94///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
95///                         back to 32x64 in-launcher when the device/in_f can't take it).
96///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
97///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
98///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
99///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
100///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
101///                         verdict was stale).
102pub fn moe_f16g_sk_params() -> (i32, i32) {
103    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
104    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
105        Ok("0") => (-1, 0),
106        Ok("32") => (0, i32::MAX),
107        Ok("128") => (0, 1),
108        _ => {
109            let cross = std::env::var("MEMRA_F16G_SK_CROSS").ok()
110                .and_then(|v| v.parse().ok()).unwrap_or(64);
111            (0, cross)
112        }
113    })
114}
115/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
116/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
117/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
118/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
119/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
120/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
121/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
122/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
123/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
124/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
125pub fn moe_f16g_direct_on(qtype: i32) -> bool {
126    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
127    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
128        Ok("0") => 0,
129        Ok("kq") => 1,
130        _ => 2,
131    });
132    match m {
133        0 => false,
134        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
135        _ => true,
136    }
137}
138/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
139/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
140/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
141/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
142/// stage under q35's routing skew. Bit-identical to every other sk form by construction
143/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
144/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
145/// tail. in_f % 64 != 0 falls back in-launcher.
146pub fn moe_f16g_tail_on() -> bool {
147    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
148    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
149}
150
151/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
152/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
153/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
154/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
155/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
156/// still opens this door for A/B.
157pub fn moe_f16g_gemma_on() -> bool {
158    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
159    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
160}
161
162/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
163/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
164/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
165pub fn moe_fuse_actq_on() -> bool {
166    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
167    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
168}
169
170/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
171/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
172/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
173/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
174/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
175/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
176/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
177/// verify already use (dispatch parity, one router kernel for every t).
178/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
179pub fn router_prefill_exact_on() -> bool {
180    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
181    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
182}
183
184pub fn router_kernel_on() -> bool {
185    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
186    *ON.get_or_init(|| {
187        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
188        if !on { eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)"); }
189        on
190    })
191}
192
193/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
194/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
195/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
196/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
197/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
198/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
199/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
200/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
201/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
202/// seam, perf-only: bits are equal by the kernel-check gate).
203/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
204/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
205/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
206/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
207pub const ROUTER_BATCH_MIN_T: usize = 8;
208pub fn router_batch_on() -> bool {
209    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
211}
212mod cpu_experts;
213pub mod moe_cache;
214pub mod spill;
215mod spill_pread;
216#[cfg(memra_cutlass)]
217pub mod cutlass_ffi;
218pub mod mmq_ffi;
219pub mod f16_ffi;
220pub mod prime_graph;
221pub mod fp8_ffi;
222
223// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
224// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
225// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
226// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
227// broke every machine that wasn't the build machine. Same bytes, same module image;
228// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
229const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
230const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
231const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
232const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
233const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
234const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
235/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
236const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
237
238/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
239/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
240/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
241/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
242/// compile-time default (zero behavior change).
243fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
244    assert!(!(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
245            "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane");
246    match std::env::var("MEMRA_GEMM_FATBIN") {
247        Ok(path) => std::borrow::Cow::Owned(
248            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}"))),
249        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
250    }
251}
252
253/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
254/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
255/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
256/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
257/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
258/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
259pub(crate) const fn portable_mma_gated() -> bool {
260    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
261}
262
263/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
264/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
265/// in a pure helper so the dispatch guard can be regression-tested without constructing an
266/// Engine or allocating a GPU tensor.
267const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
268    (!portable_cuda || hopper_mma) && !no_gemm
269}
270
271// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
272// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
273// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
274// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
275// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
276// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
277// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
278const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
279const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
280const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
281const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
282const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
283
284/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
285/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
286pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
287
288/// The flash_attn fatbin matching the selected KV formats.
289fn flash_fatbin_bytes() -> &'static [u8] {
290    match kv_cache_formats() {
291        ("q8_0", "q5_1") => FLASH_FATBIN,
292        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
293        ("q8_0", "fp8")  => FLASH_FATBIN_VF8,
294        ("fp8",  "q5_1") => FLASH_FATBIN_KF8,
295        ("fp8",  "q4_0") => FLASH_FATBIN_KF8VQ4,
296        ("fp8",  "fp8")  => FLASH_FATBIN_KF8VF8,
297        other => unreachable!("kv_cache_formats returned {other:?}"),
298    }
299}
300
301/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
302/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
303/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
304/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
305/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
306/// defaults (zero behavior change).
307fn k1_launch_override() -> Option<(u32, u32, u32)> {
308    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
309    *K1.get_or_init(|| {
310        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
311        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
312        match p.as_slice() { [bm, bn, w] => Some((*bm, *bn, *w)), _ => None }
313    })
314}
315
316/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
317/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
318/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
319/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
320/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
321/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
322pub(crate) fn wgmma_gemm_enabled() -> bool {
323    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
324    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
325}
326
327/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
328/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
329/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
330/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
331/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
332/// the split count changes the combine's FP summation order, and the spec verify's batched forward
333/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
334/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
335/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
336/// adaptive retries (any retry MUST pass run-spec self-consistency first).
337/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
338/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
339/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
340/// between eager decode and the verify (the spec-exactness law).
341pub const FA_VEC_MIN_TKV: usize = 96;
342/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
343/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
344/// which moves the crossover — sweep per model, adopt per the battery.
345pub fn fa_vec_min_tkv() -> usize {
346    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
347    *V.get_or_init(|| std::env::var("MEMRA_FA_VEC_MIN").ok()
348        .and_then(|v| v.parse().ok())
349        .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)))
350}
351
352/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
353/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
354/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
355///
356/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
357/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
358/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
359/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
360/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
361/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
362pub fn fa_f16pv_on() -> bool {
363    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
364    *ON.get_or_init(|| std::env::var("MEMRA_FA_F16PV").map(|v| v != "0")
365        .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err()))
366}
367
368/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
369/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
370/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
371pub fn fa512_hp_on() -> bool {
372    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
373    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
374}
375
376/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
377/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
378/// accumulation. Even n_head and even GQA group required (guarded per call).
379pub fn faw_hp_on() -> bool {
380    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
382}
383
384/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
385/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
386/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
387pub fn fa512_wide_warps() -> usize {
388    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
390        Ok("1") => 4, _ => 2,
391    })
392}
393
394/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
395/// and the gemma global-layer rows/parity call sites.
396pub fn fa512_min_tkv() -> usize {
397    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
398    *FA512_MIN.get_or_init(|| std::env::var("MEMRA_FA512_MIN").ok()
399        .and_then(|v| v.parse().ok()).unwrap_or(512))
400}
401/// Per-model crossover default, set at model load BEFORE the first decode (per-model
402/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
403/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
404pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
405    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
406/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
407/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
408/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
409pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize =
410    std::sync::atomic::AtomicUsize::new(32);
411/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
412/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
413/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
414/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
415/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
416pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
417    std::sync::atomic::AtomicBool::new(false);
418/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
419/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
420/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
421/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
422/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
423/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
424pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
425    std::sync::atomic::AtomicBool::new(true);
426pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
427    std::sync::atomic::AtomicUsize::new(16);
428/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
429/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
430/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
431/// latency-bound at 256 threads — 7us/launch measured).
432pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
433/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
434pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
435/// Per-model stream-k override for SPEC serving (-1 = unset → env/default; 0 = force
436/// tiling; 1 = force sk). Set by generate_spec_gemma per model tier — the sk autotune's
437/// per-process kernel coin made 12B-class spec cells bimodal, while the 26B's drafter
438/// measures BETTER under sk's fold order (2026-07-27). mmq_ffi reads this before the env.
439pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
440/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
441/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
442pub use memra_kv::KV_FP8_FORCE;
443pub(crate) fn rms_block() -> u32 {
444    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
445    *V.get_or_init(|| std::env::var("MEMRA_RMS_BLOCK").ok()
446        .and_then(|v| v.parse().ok())
447        .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)))
448}
449
450pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
451    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
452    if let Some(forced) = *S.get_or_init(|| {
453        std::env::var("MEMRA_FA_SPLIT").ok().and_then(|v| v.parse().ok())
454            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
455    }) { return forced; }
456    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
457    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
458    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
459    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
460    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
461    //
462    // SM-AWARE SHORT-CTX RUNG (2026-07-06 g7e): the 32-key rung was tuned on the 82-SM 5090.
463    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
464    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on g7e (N=1 sweep + N=3
465    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
466    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
467    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
468    // rig-divergence law: this branch is measured on 188 SMs only).
469    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
470    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
471    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
472    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
473    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
474        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1") {
475        return if t_kv <= 8192 { 16 } else if t_kv <= 16384 { 64 } else { 128 };
476    }
477    let big_rig = fa_sm_count() >= 128;
478    if big_rig {
479        let _ = n_head_kv;
480        if t_kv <= 2048 { 16 } else if t_kv <= 16384 { 64 } else { 128 }
481    } else if n_head_kv <= 4 {
482        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
483        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
484        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
485        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
486        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
487        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
488        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
489        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
490        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
491        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
492        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
493        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
494        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
495        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
496        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
497        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
498        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
499        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
500        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
501        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
502        if t_kv <= 512 { 8 } else if t_kv <= 16384 { 64 } else { 128 }
503    } else {
504        if t_kv <= 8192 { 32 } else if t_kv <= 16384 { 64 } else { 128 }
505    }
506}
507
508/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
509/// same attribute Engine::batched_variant reads).
510fn fa_sm_count() -> i32 {
511    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
512    *N.get_or_init(|| {
513        cudarc::driver::result::init().ok();
514        cudarc::driver::result::device::get(0)
515            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
516                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
517            .unwrap_or(82)
518    })
519}
520
521/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
522/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
523/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
524fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
525    match head_dim {
526        256 => Ok(""),
527        128 => Ok("_hd128"),
528        d => Err(format!("fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
529                          callers must gate to sdpa_naive").into()),
530    }
531}
532
533/// Quant type codes matching qmatvec.cu QType enum.
534pub const QT_Q8_0: i32 = 0;
535pub const QT_Q4_K: i32 = 1;
536pub const QT_Q6_K: i32 = 2;
537pub const QT_Q5_K: i32 = 3;
538pub const QT_Q3_K: i32 = 4;
539pub const QT_IQ4_XS: i32 = 5;
540pub const QT_IQ3_S: i32 = 6;
541pub const QT_NVFP4: i32 = 7;
542/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
543/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
544/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
545/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
546/// — ONE weight copy total, no Q8_0 re-encode duplicate.
547pub const QT_F8_E4M3: i32 = 10;
548/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
549/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
550pub const QT_NVFP4_RP: i32 = 9;
551/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
552pub const QT_F32: i32 = 8;
553pub const QT_BF16: i32 = 11;
554pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
555/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
556/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
557/// dp4a/MMQ implementation exists.
558pub const QT_Q2_K: i32 = 13;
559
560/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
561pub struct Engine {
562    pub gpu: memra_runtime::Gpu,
563    module: Arc<CudaModule>,
564    hybrid: Arc<CudaModule>,
565    qmatvec: Arc<CudaModule>,
566    flash: Arc<CudaModule>,
567    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
568    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
569    /// Lazy: loaded on first global-format use; None until then.
570    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
571    gemm: Arc<CudaModule>,
572    router: Arc<CudaModule>,
573    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
574    sample: Arc<CudaModule>,
575    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
576        /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
577    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
578    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
579    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
580    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
581    /// the single largest block. The cache still owns every address for its full lifetime.
582    moe_cache_layout: Mutex<Option<Vec<usize>>>,
583    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
584    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
585    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
586    /// verify between replays) reuse their addresses and the replay reads/writes live memory
587    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
588    capture_keep_on: std::sync::atomic::AtomicBool,
589    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
590    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
591    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
592    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
593    verify_exact: std::sync::atomic::AtomicBool,
594    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
595    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
596    pub copy_stream: Arc<CudaStream>,
597    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
598    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
599    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
600    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
601    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
602    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
603    #[cfg(memra_cutlass)]
604    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
605    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
606    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
607    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
608    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
609    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
610    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
611    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
612    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
613    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
614    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
615    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
616    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
617    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
618    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
619    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
620    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
621    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
622    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
623    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
624    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
625    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
626    /// before capture under the generate_graph tracking-off window so it carries no events).
627    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
628    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
629    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
630    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
631    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
632    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
633    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
634    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
635    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
636    router_stage: Mutex<Option<PinnedStage>>,
637}
638
639/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
640/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
641/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
642/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
643/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
644/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
645/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
646/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
647/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
648fn fa_v2_on() -> bool {
649    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
650    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
651    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
652    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
653    // + graph bit-identity green on all three models.
654    std::env::var("MEMRA_FA_V2").map(|v| v != "0").unwrap_or(true)
655}
656
657/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
658/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
659/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
660/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
661/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
662/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
663/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
664fn fa_v3_on() -> bool {
665    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
666    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
667    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
668    std::env::var("MEMRA_FA_V3").map(|v| v != "0").unwrap_or(true)
669}
670
671/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
672/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
673/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
674/// predicate so the twins can never diverge.
675fn fa_v4_mode() -> &'static str {
676    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
677    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
678}
679fn fa_v4_on() -> bool { fa_v4_mode() != "0" }   // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
680/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
681/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
682/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
683/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
684/// stays kernel-family-identical to decode at the same t_kv.
685/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
686/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
687pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
688    std::sync::atomic::AtomicUsize::new(1024);
689pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
690    std::sync::atomic::AtomicUsize::new(usize::MAX);
691pub fn fa_v4_at_pub(t_kv: usize) -> bool { fa_v4_at(t_kv) }
692fn fa_v4_at(t_kv: usize) -> bool {
693    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
694    let mx = *M.get_or_init(|| std::env::var("MEMRA_FA_V4_MAX").ok()
695        .and_then(|v| v.parse().ok())
696        .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)));
697    fa_v4_on() && t_kv < mx
698}
699/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
700/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
701/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
702/// (same split partition, same softmax/accumulation order, same partials/combine) and only
703/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
704/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
705/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
706/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
707/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
708/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
709/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
710/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
711/// within one process (the v2/v3 pattern).
712pub const FA_DEEP_MIN_DEFAULT: usize = 0;
713fn fa_deep_at(t_kv: usize) -> bool {
714    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") { return false; }
715    let min = std::env::var("MEMRA_FA_DEEP_MIN").ok().and_then(|v| v.parse().ok())
716        .unwrap_or(FA_DEEP_MIN_DEFAULT);
717    t_kv >= min
718}
719/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
720pub fn fa_deep_at_pub(t_kv: usize) -> bool { fa_deep_at(t_kv) }
721
722fn fa_v3_active(head_dim: usize) -> bool {
723    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
724    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
725    fa_v3_on() && head_dim % 128 == 0 && kv_cache_formats() == ("q8_0", "q5_1")
726        && !Engine::kv_fp8_on()
727}
728
729/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
730/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
731/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
732/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
733/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
734/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
735/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
736pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
737    std::env::var("MEMRA_NO_FA_VEC").is_err()
738        && t_kv >= fa_vec_min_tkv()
739        && head_dim == 256
740        && fa_v4_at(t_kv)
741        && !matches!(fa_v4_mode(), "noB3" | "stage")
742        && !Engine::kv_fp8_on()
743}
744/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
745pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize { fa_split_keys(t_kv, n_head_kv) }
746
747/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
748/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
749/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
750/// so we allocate through `result::malloc_host` with flags=0 directly.
751struct PinnedStage {
752    ptr: *mut u8,
753    cap: usize,
754}
755unsafe impl Send for PinnedStage {}
756impl PinnedStage {
757    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
758        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
759        Ok(PinnedStage { ptr, cap })
760    }
761}
762impl Drop for PinnedStage {
763    fn drop(&mut self) {
764        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
765    }
766}
767
768/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
769/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
770pub const ARGMAX_NB: usize = 256;
771
772/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
773pub(crate) use memra_fa3_vl as fa3_vl_raw;
774
775unsafe extern "C" {
776    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
777    fn memra_fa3_prefill(q16: *const core::ffi::c_void, k16: *const core::ffi::c_void,
778                        v16: *const core::ffi::c_void, o: *mut f32,
779                        t: i32, h: i32, hkv: i32, d: i32, scale: f32,
780                        stream: *mut core::ffi::c_void) -> i32;
781    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
782    pub(crate) fn memra_fa3_vl(q16s: *const *const core::ffi::c_void, k16s: *const *const core::ffi::c_void,
783                   v16s: *const *const core::ffi::c_void, os: *const *mut f32,
784                   ts: *const i32, b: i32, h: i32, hkv: i32, d: i32, scale: f32,
785                   stream: *mut core::ffi::c_void) -> i32;
786}
787
788/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
789/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
790/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
791/// (slots are never re-allocated), so passing raw values is stable across the launch.
792#[repr(C)]
793#[derive(Clone, Copy)]
794pub struct WPtr8(pub [u64; 8]);
795unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
796
797/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
798/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
799/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
800/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
801#[repr(C)]
802#[derive(Clone, Copy, Default)]
803pub struct GdnSeqVl {
804    pub kb16: u64, pub gcum: u64, pub beta: u64, pub u: u64, pub wb16: u64,
805    pub y: u64, pub ssnap: u64, pub state_in: u64, pub state_out: u64,
806    pub q: u64, pub p: u64, pub o: u64,
807    pub k: u64, pub v: u64, pub g: u64, pub a: u64, pub w: u64,
808    pub t: i32, pub nc: i32,
809}
810unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
811#[repr(C)]
812#[derive(Clone, Copy)]
813pub struct GdnVl8(pub [GdnSeqVl; 8]);
814unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
815
816/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
817/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
818#[repr(C)]
819#[derive(Clone, Copy, Default)]
820pub struct GdnWVl { pub qb16: u64, pub pb16: u64 }
821unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
822#[repr(C)]
823#[derive(Clone, Copy)]
824pub struct GdnWVl8(pub [GdnWVl; 8]);
825unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
826
827/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
828#[repr(C)]
829#[derive(Clone, Copy, Default)]
830pub struct GdnPrepVl {
831    pub qkv: u64, pub conv_state: u64, pub conv_out: u64,
832    pub q_g: u64, pub k_g: u64, pub v_g: u64,
833    pub q_l2: u64, pub k_l2: u64,
834    pub beta_raw: u64, pub alpha: u64, pub beta: u64, pub g_log: u64,
835    pub o: u64, pub z: u64, pub gn: u64, pub gn16: u64,
836    pub kb16: u64,
837    pub qb16: u64,
838    pub t: i32, pub pad: i32,
839}
840unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
841#[repr(C)]
842#[derive(Clone, Copy)]
843pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
844unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
845
846/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
847#[repr(C)]
848#[derive(Clone, Copy, Default)]
849pub struct FaSeqVl {
850    pub q: u64, pub k16: u64, pub v16: u64, pub o: u64, pub kf: u64, pub vf: u64,
851    pub t: i32, pub pad: i32,
852}
853unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
854#[repr(C)]
855#[derive(Clone, Copy)]
856pub struct FaVl8(pub [FaSeqVl; 8]);
857unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
858
859/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
860#[repr(C)]
861#[derive(Clone, Copy, Default)]
862pub struct AttnPreVl {
863    pub qf: u64, pub kf: u64, pub vf: u64,
864    pub q: u64, pub gate: u64, pub qn: u64, pub kn: u64,
865    pub kc: u64, pub vc: u64,
866    pub t: i32, pub pad: i32,
867}
868unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
869#[repr(C)]
870#[derive(Clone, Copy)]
871pub struct AttnPreVl8(pub [AttnPreVl; 8]);
872unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
873
874/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
875/// varlen K1-K5 chain fills them).
876pub struct GdnChunkBufs {
877    pub gcum: CudaSlice<f32>,
878    pub a: CudaSlice<f32>,
879    pub p: CudaSlice<f32>,
880    pub u: CudaSlice<f32>,
881    pub w: CudaSlice<f32>,
882    pub kb16: CudaSlice<u8>,
883    pub wb16: CudaSlice<u8>,
884    pub y16: CudaSlice<u8>,
885    pub ssnap16: CudaSlice<u8>,
886    pub qb16: CudaSlice<u8>,
887    pub pb16: CudaSlice<u8>,
888    pub o: CudaSlice<f32>,
889    pub t: usize,
890    pub nc: usize,
891}
892
893/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
894#[repr(C)]
895#[derive(Clone, Copy)]
896pub struct F32x8(pub [f32; 8]);
897unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
898
899/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
900/// process. Bench binaries read it right after the call to print gen-only throughput without the
901/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
902pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
903
904impl Engine {
905    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
906        let gpu = memra_runtime::Gpu::new(ordinal)?;
907        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
908        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
909        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
910        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
911            use cudarc::driver::sys::CUdevice_attribute_enum as A;
912            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
913                .and_then(|d| unsafe { Ok((
914                    cudarc::driver::result::device::get_attribute(d, A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)?,
915                    cudarc::driver::result::device::get_attribute(d, A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)?)) })
916                .unwrap_or((0, 0));
917            let built = env!("MEMRA_BUILT_CUDA_ARCH");
918            let ok = matches!((built, maj, min),
919                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9));
920            if !ok {
921                return Err(format!(
922                    "memra was built for sm_{built} but device {ordinal} reports compute \
923                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
924                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass.").into());
925            }
926        }
927        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
928        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
929        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
930        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
931        unsafe {
932            use cudarc::driver::sys;
933            let dev: sys::CUdevice = ordinal as sys::CUdevice;
934            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
935            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
936                let mut thresh: u64 = u64::MAX;
937                let _ = sys::cuMemPoolSetAttribute(
938                    pool,
939                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
940                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
941                );
942            }
943        }
944        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
945        let hybrid = gpu.ctx.load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
946        let qmatvec = gpu.ctx.load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
947        let flash = gpu.ctx.load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
948        let gemm = gpu.ctx.load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
949        let router = gpu.ctx.load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
950        let sample = gpu.ctx.load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
951        let copy_stream = gpu.ctx.new_stream()?;
952        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
953        // cudarc is in multi-stream mode (main stream +
954        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
955        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
956        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
957        // (~7 ms/tok host time, measured nsys 2026-07-04 g7e), and +4.6% measured on 27B decode —
958        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
959        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
960        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
961        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
962        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
963        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
964        // implicit event tracking.
965        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
966        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
967        if std::env::var("MEMRA_EVT").map(|v| v == "1").unwrap_or(false) {
968            // escape hatch: keep cudarc's implicit cross-stream event tracking.
969        } else {
970            unsafe { gpu.ctx.disable_event_tracking(); }
971        }
972        Ok(Self { gpu, module, hybrid, qmatvec, flash, flash_g: std::sync::OnceLock::new(), gemm, router, sample,
973                  moe_cache: Mutex::new(None),
974                  moe_cache_layout: Mutex::new(None),
975                  copy_stream,
976                  capture_keep_on: std::sync::atomic::AtomicBool::new(false),
977                  verify_exact: std::sync::atomic::AtomicBool::new(false),
978                  capture_keep: Mutex::new(Vec::new()),
979                  argmax_partials: Mutex::new(None),
980                  prime_deqw_ws: Mutex::new(None),
981                  router_stage: Mutex::new(None),
982                  fp8_scratch: Mutex::new(None),
983                  fa_vf16_scratch: Mutex::new(None),
984                  fa_part_pool: Mutex::new(None),
985                  fa_part_retired: Mutex::new(Vec::new()),
986                  fn_cache: Mutex::new(Default::default()),
987                  f16_scratch: Mutex::new(None),
988                  #[cfg(memra_cutlass)]
989                  cutlass_scratch: Mutex::new(None) })
990    }
991
992    pub fn ctx(&self) -> &Arc<CudaContext> { &self.gpu.ctx }
993    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
994    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
995    pub fn stream(&self) -> Arc<CudaStream> { self.gpu.stream() }
996    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
997    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
998    pub fn gkv_on() -> bool {
999        memra_kv::gkv_on()
1000    }
1001
1002    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1003    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1004    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1005    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1006    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1007    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1008    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1009    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1010    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1011    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1012    /// ON for both — no acceptance cost measured.
1013    pub fn wkv_on() -> bool {
1014        memra_kv::wkv_on()
1015    }
1016
1017    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1018    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1019    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1020    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1021    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1022    pub fn kv_fp8_on() -> bool {
1023        memra_kv::kv_fp8_on()
1024    }
1025
1026    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1027    /// when the fp8-globals arm is on; everything else from the default flash module.
1028    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1029        if head_dim == 512 && Self::gkv_on() { self.func_g(name) } else { self.func(name) }
1030    }
1031
1032    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1033    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1034    /// per-format fatbins; fall back to the base modules for those.
1035    fn func_g(&self, name: &str) -> CudaFunction {
1036        let m = self.flash_g.get_or_init(|| {
1037            self.gpu.ctx.load_module(cudarc::nvrtc::Ptx::from_binary(FLASH_FATBIN_KF8VF8.to_vec()))
1038                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1039        });
1040        let key = format!("g:{name}");
1041        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) { return f.clone(); }
1042        let f = match m.load_function(name) {
1043            Ok(f) => f,
1044            Err(_) => self.func(name),
1045        };
1046        self.fn_cache.lock().unwrap().insert(key, f.clone());
1047        f
1048    }
1049
1050    fn func(&self, name: &str) -> CudaFunction {
1051        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1052        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1053        if let Some(f) = self.fn_cache.lock().unwrap().get(name) { return f.clone(); }
1054        let f = self.module.load_function(name)
1055            .or_else(|_| self.hybrid.load_function(name))
1056            .or_else(|_| self.qmatvec.load_function(name))
1057            .or_else(|_| self.flash.load_function(name))
1058            .or_else(|_| self.gemm.load_function(name))
1059            .or_else(|_| self.router.load_function(name))
1060            .or_else(|_| self.sample.load_function(name))
1061            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1062        self.fn_cache.lock().unwrap().insert(name.to_string(), f.clone());
1063        f
1064    }
1065
1066    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1067    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1068    pub fn scatter_trim_logits(&self, src: &CudaSlice<f32>, d2t: &CudaSlice<u32>,
1069                               dst: &mut CudaSlice<f32>, d_vocab: usize, n_vocab: usize)
1070                               -> Result<(), Box<dyn std::error::Error>> {
1071        let f1 = self.func("scatter_trim_logits_f32");
1072        let f2 = self.func("scatter_trim_logits_pass2_f32");
1073        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1074        let cfg1 = LaunchConfig { grid_dim: (256, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1075        let __s_b1 = self.gpu.stream();
1076        let mut b1 = __s_b1.launch_builder(&f1);
1077        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1078        unsafe { b1.launch(cfg1)?; }
1079        let cfg2 = LaunchConfig { grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1080        let __s_b2 = self.gpu.stream();
1081        let mut b2 = __s_b2.launch_builder(&f2);
1082        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1083        unsafe { b2.launch(cfg2)?; }
1084        Ok(())
1085    }
1086
1087    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1088    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1089
1090    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1091    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1092    #[allow(clippy::too_many_arguments)]
1093    pub fn filter_stats(&self, x: &CudaSlice<f32>, row_stride: usize, rows: &CudaSlice<i32>,
1094                        out_th: &mut CudaSlice<f32>, out_z: &mut CudaSlice<f32>,
1095                        out_max: &mut CudaSlice<f32>, n: usize, nrow: usize,
1096                        temp: f32, top_k: i32, top_p: f32, min_p: f32)
1097                        -> Result<(), Box<dyn std::error::Error>> {
1098        let f = self.func("filter_stats_f32");
1099        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1100        let cfg = LaunchConfig { grid_dim: (nrow as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
1101        let __s_b = self.gpu.stream();
1102        let mut b = __s_b.launch_builder(&f);
1103        b.arg(x).arg(&rs).arg(rows).arg(&mut *out_th).arg(&mut *out_z).arg(&mut *out_max)
1104         .arg(&ni).arg(&nr).arg(&temp).arg(&top_k).arg(&top_p).arg(&min_p);
1105        unsafe { b.launch(cfg)?; }
1106        Ok(())
1107    }
1108
1109    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1110    #[allow(clippy::too_many_arguments)]
1111    pub fn softmax_gather_filtered(&self, x: &CudaSlice<f32>, row_stride: usize,
1112                                   ids: &CudaSlice<u32>, rows: &CudaSlice<i32>,
1113                                   th: &CudaSlice<f32>, z: &CudaSlice<f32>,
1114                                   out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32)
1115                                   -> Result<(), Box<dyn std::error::Error>> {
1116        let f = self.func("softmax_gather_filtered_f32");
1117        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1118        let cfg = LaunchConfig { grid_dim: (npair as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1119        let __s_b = self.gpu.stream();
1120        let mut b = __s_b.launch_builder(&f);
1121        b.arg(x).arg(&rs).arg(ids).arg(rows).arg(th).arg(z).arg(&mut *out).arg(&ni).arg(&np).arg(&temp);
1122        unsafe { b.launch(cfg)?; }
1123        Ok(())
1124    }
1125
1126    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1127    #[allow(clippy::too_many_arguments)]
1128    pub fn residual_sample_filtered(&self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize,
1129                                    temp: f32, seed: u64, stream_pos: u32,
1130                                    p_stats: (f32, f32, f32), q_stats: (f32, f32, f32),
1131                                    out_tok: &mut CudaSlice<u32>)
1132                                    -> Result<(), Box<dyn std::error::Error>> {
1133        let f = self.func("residual_sample_filtered_f32");
1134        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1135        let has_q: i32 = q.is_some() as i32;
1136        let qbuf = q.unwrap_or(p);
1137        let (pm, pth, pz) = p_stats; let (qm, qth, qz) = q_stats;
1138        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
1139        let __s_b = self.gpu.stream();
1140        let mut b = __s_b.launch_builder(&f);
1141        b.arg(p).arg(qbuf).arg(&has_q).arg(&ni).arg(&temp).arg(&slo).arg(&shi).arg(&stream_pos)
1142         .arg(&pm).arg(&pth).arg(&pz).arg(&qm).arg(&qth).arg(&qz).arg(&mut *out_tok);
1143        unsafe { b.launch(cfg)?; }
1144        Ok(())
1145    }
1146
1147    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1148    #[allow(clippy::too_many_arguments)]
1149    pub fn gumbel_perturb_filtered(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1150                                   seed: u64, stream_pos: u32, temp: f32, row_max: f32, th: f32)
1151                                   -> Result<(), Box<dyn std::error::Error>> {
1152        let f = self.func("gumbel_perturb_filtered_f32");
1153        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1154        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1155        let __s_b = self.gpu.stream();
1156        let mut b = __s_b.launch_builder(&f);
1157        b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp).arg(&row_max).arg(&th);
1158        unsafe { b.launch(cfg)?; }
1159        Ok(())
1160    }
1161
1162    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1163    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1164    /// filtered rejection sampling exact for the penalized target.
1165    #[allow(clippy::too_many_arguments)]
1166    pub fn penalize_logits(&self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize,
1167                           rep: f32, freq: f32, present: f32, n: usize)
1168                           -> Result<(), Box<dyn std::error::Error>> {
1169        if n_hist == 0 { return Ok(()); }
1170        let f = self.func("penalize_logits_f32");
1171        let (nh, ni) = (n_hist as i32, n as i32);
1172        let cfg = LaunchConfig { grid_dim: (n_hist.div_ceil(128) as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
1173        let __s_b = self.gpu.stream();
1174        let mut b = __s_b.launch_builder(&f);
1175        b.arg(&mut *x).arg(hist).arg(&nh).arg(&rep).arg(&freq).arg(&present).arg(&ni);
1176        unsafe { b.launch(cfg)?; }
1177        Ok(())
1178    }
1179
1180    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1181    #[allow(clippy::too_many_arguments)]
1182    pub fn penalize_logits_rows(&self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize,
1183                                rep: f32, freq: f32, present: f32, n: usize, nrow: usize)
1184                                -> Result<(), Box<dyn std::error::Error>> {
1185        if n_hist == 0 || nrow == 0 { return Ok(()); }
1186        let f = self.func("penalize_logits_rows_f32");
1187        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1188        let cfg = LaunchConfig { grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
1189        let __s_b = self.gpu.stream();
1190        let mut b = __s_b.launch_builder(&f);
1191        b.arg(&mut *x).arg(hist).arg(&nh).arg(&rep).arg(&freq).arg(&present).arg(&ni).arg(&nr);
1192        unsafe { b.launch(cfg)?; }
1193        Ok(())
1194    }
1195
1196    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1197    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1198    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1199    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1200    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1201    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1202    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1203    pub fn wpf_level() -> u32 {
1204        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1205        *ON.get_or_init(|| std::env::var("MEMRA_WPF").ok()
1206            .and_then(|v| v.parse().ok()).unwrap_or(1))
1207    }
1208
1209    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1210    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1211    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1212    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1213    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1214    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1215    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1216    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1217    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1218    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1219    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1220    pub fn set_verify_exact(&self, on: bool) {
1221        self.verify_exact.store(on, std::sync::atomic::Ordering::Relaxed);
1222    }
1223    pub(crate) fn verify_exact_on(&self) -> bool {
1224        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1225    }
1226
1227    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1228    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1229    pub fn qkv_append_on() -> bool {
1230        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1231        *ON.get_or_init(|| std::env::var("MEMRA_QKV_APPEND").map(|v| v != "0").unwrap_or(true))
1232    }
1233
1234    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1235    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1236    pub fn pdl_wb_on() -> bool {
1237        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1238        *ON.get_or_init(|| std::env::var("MEMRA_PDL_WB").map(|v| v != "0").unwrap_or(true))
1239    }
1240
1241    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1242    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1243    /// per-model no-harm bisect knob.
1244    pub fn pdl_mmvq_on() -> bool {
1245        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1246        *ON.get_or_init(|| std::env::var("MEMRA_PDL_MMVQ").map(|v| v != "0").unwrap_or(true))
1247    }
1248
1249    pub fn pdl_on() -> bool {
1250        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1251        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1252    }
1253
1254    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1255    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1256    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1257    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1258    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1259    fn q40_mr1_on() -> bool {
1260        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1261        match *Q40MR.get_or_init(|| std::env::var("MEMRA_Q40_MR").ok()
1262            .and_then(|v| v.parse().ok())) {
1263            Some(v) => v == 1,
1264            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1265        }
1266    }
1267
1268    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1269    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1270    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1271    /// writes wrong bytes silently.
1272    fn pdl_func_flash(&self, g: bool, name: &'static str)
1273        -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1274        use cudarc::driver::sys as cu;
1275        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1276        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1277        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1278        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1279        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1280        // this engine's CUcontext; single-context runs behave exactly as before.
1281        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1282            std::sync::Mutex::new(None);
1283        static FNS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool, &'static str), usize>>> =
1284            std::sync::Mutex::new(None);
1285        let ctx_key = self.ctx().cu_ctx() as usize;
1286        if let Some(&f) = FNS.lock().unwrap().get_or_insert_with(Default::default)
1287            .get(&(ctx_key, g, name)) { return Ok(f as cu::CUfunction); }
1288        let module = {
1289            let mut mods = MODS.lock().unwrap();
1290            let map = mods.get_or_insert_with(Default::default);
1291            match map.get(&(ctx_key, g)) {
1292                Some(&m) => m,
1293                None => {
1294                    let m = self.pdl_load_module_in_ctx(
1295                        if g { FLASH_FATBIN_KF8VF8 } else { FLASH_FATBIN })?;
1296                    map.insert((ctx_key, g), m);
1297                    m
1298                }
1299            }
1300        };
1301        let cname = std::ffi::CString::new(name)?;
1302        let mut f: cu::CUfunction = std::ptr::null_mut();
1303        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1304        if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into()); }
1305        FNS.lock().unwrap().get_or_insert_with(Default::default)
1306            .insert((ctx_key, g, name), f as usize);
1307        Ok(f)
1308    }
1309
1310    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
1311    /// the module to the thread's CURRENT context — a remote-stage engine must not
1312    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
1313    /// current context before returning.
1314    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1315        use cudarc::driver::sys as cu;
1316        let mut prev: cu::CUcontext = std::ptr::null_mut();
1317        unsafe { cu::cuCtxGetCurrent(&mut prev).result()?; }
1318        self.ctx().bind_to_thread()?;
1319        let mut m: cu::CUmodule = std::ptr::null_mut();
1320        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1321        let restore = if prev.is_null() { cu::CUresult::CUDA_SUCCESS }
1322                      else { unsafe { cu::cuCtxSetCurrent(prev) } };
1323        if r != cu::CUresult::CUDA_SUCCESS {
1324            return Err(format!("pdl module load: {r:?}").into());
1325        }
1326        if restore != cu::CUresult::CUDA_SUCCESS {
1327            return Err(format!("pdl module load: ctx restore {restore:?}").into());
1328        }
1329        Ok(m as usize)
1330    }
1331
1332    fn pdl_func(&self, name: &'static str) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1333        use cudarc::driver::sys as cu;
1334        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
1335        // are context-scoped; key everything by this engine's CUcontext).
1336        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1337            std::sync::Mutex::new(None);
1338        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
1339        // duplicate module, loaded lazily on the first kernels-module miss.
1340        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1341            std::sync::Mutex::new(None);
1342        static FNS: std::sync::Mutex<Option<std::collections::HashMap<(usize, &'static str), usize>>> =
1343            std::sync::Mutex::new(None);
1344        let ctx_key = self.ctx().cu_ctx() as usize;
1345        if let Some(&f) = FNS.lock().unwrap().get_or_insert_with(Default::default)
1346            .get(&(ctx_key, name)) { return Ok(f as cu::CUfunction); }
1347        let module = {
1348            let mut mods = MODULES.lock().unwrap();
1349            let map = mods.get_or_insert_with(Default::default);
1350            match map.get(&ctx_key) {
1351                Some(&m) => m,
1352                None => {
1353                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
1354                    map.insert(ctx_key, m);
1355                    m
1356                }
1357            }
1358        };
1359        let cname = std::ffi::CString::new(name)?;
1360        let mut f: cu::CUfunction = std::ptr::null_mut();
1361        let mut r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1362        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1363            let qmodule = {
1364                let mut mods = QMODULES.lock().unwrap();
1365                let map = mods.get_or_insert_with(Default::default);
1366                match map.get(&ctx_key) {
1367                    Some(&m) => m,
1368                    None => {
1369                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1370                        map.insert(ctx_key, m);
1371                        m
1372                    }
1373                }
1374            };
1375            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1376        }
1377        if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("pdl_func {name}: {r:?}").into()); }
1378        FNS.lock().unwrap().get_or_insert_with(Default::default)
1379            .insert((ctx_key, name), f as usize);
1380        Ok(f)
1381    }
1382
1383    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
1384    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
1385    ///
1386    /// # Safety
1387    /// `params` must match the kernel's exact parameter list (order, types, count) —
1388    /// a mismatch corrupts the launch silently.
1389    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
1390    /// builder path's fa_func/func_g choice exactly).
1391    ///
1392    /// # Safety
1393    /// Same contract as `launch_pdl`.
1394    unsafe fn launch_pdl_flash(&self, g: bool, name: &'static str, grid: (u32, u32, u32),
1395                               block: (u32, u32, u32), smem: u32,
1396                               params: &mut [*mut std::ffi::c_void])
1397                               -> Result<(), Box<dyn std::error::Error>> {
1398        use cudarc::driver::sys as cu;
1399        let f = self.pdl_func_flash(g, name)?;
1400        if smem > 0 {
1401            // mirror the builder path's opt-in ceiling (idempotent host-side set).
1402            let r = unsafe { cu::cuFuncSetAttribute(f,
1403                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1404                smem as i32) };
1405            if r != cu::CUresult::CUDA_SUCCESS {
1406                return Err(format!("pdl smem attr {name}: {r:?}").into());
1407            }
1408        }
1409        let mut attr = cu::CUlaunchAttribute {
1410            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1411            pad: [0; 4],
1412            value: cu::CUlaunchAttributeValue { programmaticStreamSerializationAllowed: 1 },
1413        };
1414        let cfg = cu::CUlaunchConfig {
1415            gridDimX: grid.0, gridDimY: grid.1, gridDimZ: grid.2,
1416            blockDimX: block.0, blockDimY: block.1, blockDimZ: block.2,
1417            sharedMemBytes: smem, hStream: self.gpu.stream().cu_stream(),
1418            attrs: &mut attr, numAttrs: 1,
1419        };
1420        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1421        if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("launch_pdl_flash {name}: {r:?}").into()); }
1422        Ok(())
1423    }
1424
1425    unsafe fn launch_pdl(&self, name: &'static str, grid: (u32, u32, u32), block: (u32, u32, u32),
1426                         params: &mut [*mut std::ffi::c_void])
1427                         -> Result<(), Box<dyn std::error::Error>> {
1428        use cudarc::driver::sys as cu;
1429        let f = self.pdl_func(name)?;
1430        let mut attr = cu::CUlaunchAttribute {
1431            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1432            pad: [0; 4],
1433            value: cu::CUlaunchAttributeValue { programmaticStreamSerializationAllowed: 1 },
1434        };
1435        let cfg = cu::CUlaunchConfig {
1436            gridDimX: grid.0, gridDimY: grid.1, gridDimZ: grid.2,
1437            blockDimX: block.0, blockDimY: block.1, blockDimZ: block.2,
1438            sharedMemBytes: 0, hStream: self.gpu.stream().cu_stream(),
1439            attrs: &mut attr, numAttrs: 1,
1440        };
1441        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1442        if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("launch_pdl {name}: {r:?}").into()); }
1443        Ok(())
1444    }
1445
1446    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
1447    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
1448    pub fn prefetch_weight_l2(&self, w: &crate::model::GpuTensor)
1449                              -> Result<(), Box<dyn std::error::Error>> {
1450        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
1451            let p = rp4.as_ref().unwrap_or(bytes);
1452            self.prefetch_l2(p, p.len())?;
1453        }
1454        Ok(())
1455    }
1456
1457    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
1458    /// by the DEVICE token id at tok[idx] into f32.
1459    pub fn gather_row_bf16(&self, table: &CudaSlice<u8>, tok: &CudaSlice<u32>, idx: usize,
1460                           dst: &mut CudaSlice<f32>, ncols: usize)
1461                           -> Result<(), Box<dyn std::error::Error>> {
1462        let f = self.func("gather_row_bf16_f32");
1463        let cfg = LaunchConfig { grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
1464                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1465        let (nc, ix) = (ncols as i32, idx as i32);
1466        let __s_b = self.gpu.stream();
1467        let mut b = __s_b.launch_builder(&f);
1468        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
1469        unsafe { b.launch(cfg)?; }
1470        Ok(())
1471    }
1472
1473    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
1474    pub fn add_row_inplace(&self, logits: &mut CudaSlice<f32>, bias: &CudaSlice<f32>,
1475                           n: usize, row_off: usize)
1476                           -> Result<(), Box<dyn std::error::Error>> {
1477        let f = self.func("add_row_inplace_f32");
1478        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1),
1479                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1480        let (ni, off) = (n as i32, row_off as i64);
1481        let __s_b = self.gpu.stream();
1482        let mut b = __s_b.launch_builder(&f);
1483        b.arg(logits).arg(bias).arg(&ni).arg(&off);
1484        unsafe { b.launch(cfg)?; }
1485        Ok(())
1486    }
1487
1488    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
1489    pub fn prefetch_l2(&self, p: &CudaSlice<u8>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1490        let f = self.func("prefetch_l2_bytes");
1491        let lines = n.div_ceil(128);
1492        let ni = n as i64;
1493        let cfg = LaunchConfig { grid_dim: (lines.div_ceil(256) as u32, 1, 1),
1494                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1495        let __s_b = self.gpu.stream();
1496        let mut b = __s_b.launch_builder(&f);
1497        b.arg(p).arg(&ni);
1498        unsafe { b.launch(cfg)?; }
1499        Ok(())
1500    }
1501
1502    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
1503    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
1504    pub fn router_gemv(&self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize,
1505                       n_experts: usize, t: usize)
1506                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1507        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
1508        // stream differs) — too small to justify a numeric config change; deleted.
1509        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
1510        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
1511        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
1512        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
1513            Ok("0") => false,
1514            Ok(_) => true,
1515            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1516        };
1517        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
1518        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
1519        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
1520        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
1521        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
1522        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
1523        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
1524        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
1525        // (perf-only, bits equal).
1526        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
1527        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
1528    }
1529
1530    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
1531    /// force both forms; `batch` requires `w8`).
1532    pub fn router_gemv_form(&self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize,
1533                            n_experts: usize, t: usize, w8: bool, batch: bool)
1534                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1535        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
1536        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
1537        let f = if batch { self.func("router_gemv_f32_w8_batch") }
1538                else if w8 { self.func("router_gemv_f32_w8") }
1539                else { self.func("router_gemv_f32") };
1540        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
1541        let cfg = if batch {
1542            LaunchConfig { grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
1543                           block_dim: (32, 8, 1), shared_mem_bytes: 0 }
1544        } else {
1545            LaunchConfig { grid_dim: (n_experts as u32, t as u32, 1),
1546                           block_dim: (32, if w8 { 8 } else { 1 }, 1), shared_mem_bytes: 0 }
1547        };
1548        let __s_b = self.gpu.stream();
1549        let mut b = __s_b.launch_builder(&f);
1550        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
1551        unsafe { b.launch(cfg)?; }
1552        Ok(y)
1553    }
1554
1555    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
1556    pub fn rows_permute(&self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>, nrows: usize,
1557                        ncols: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1558        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
1559        let f = self.func("rows_permute_f32");
1560        let (nc, nr) = (ncols as i32, nrows as i32);
1561        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (256, 1, 1),
1562                                 shared_mem_bytes: 0 };
1563        let __s_b = self.gpu.stream();
1564        let mut b = __s_b.launch_builder(&f);
1565        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
1566        unsafe { b.launch(cfg)?; }
1567        Ok(dst)
1568    }
1569
1570    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
1571    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
1572    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
1573    /// decode chain and the small-t spec-verify chain match per row by construction.
1574    pub fn sigmoid_dot_rows(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, n_embd: usize,
1575                            t: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1576        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
1577        // config; same class as MEMRA_ROUTER_V2).
1578        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1579        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
1580            let gs = self.linear(x, w, t, n_embd, 1)?;
1581            let mut g = self.uninit(t)?;
1582            self.sigmoid(&gs, &mut g, t)?;
1583            return Ok(g);
1584        }
1585        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
1586        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
1587        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
1588        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
1589        // flags doctrine; this per-token form serves every t.
1590        let mut g = self.alloc_uninit::<f32>(t)?;
1591        let f = self.func("sigmoid_dot_rows_f32");
1592        let (ne, ti) = (n_embd as i32, t as i32);
1593        let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (32, 8, 1),
1594                                 shared_mem_bytes: 0 };
1595        let __s_b = self.gpu.stream();
1596        let mut b = __s_b.launch_builder(&f);
1597        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
1598        unsafe { b.launch(cfg)?; }
1599        Ok(g)
1600    }
1601
1602    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
1603    pub fn spec_rollback_stream(&self, len_ptrs: &CudaSlice<u64>, pos_start: &CudaSlice<i32>,
1604                                acc: &CudaSlice<u32>, base: usize, n_rows: usize)
1605                                -> Result<(), Box<dyn std::error::Error>> {
1606        let f = self.func("spec_rollback_stream");
1607        let (b, nr) = (base as i32, n_rows as i32);
1608        let cfg = LaunchConfig { grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
1609                                 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
1610        let __s_bl = self.gpu.stream();
1611        let mut bl = __s_bl.launch_builder(&f);
1612        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
1613        unsafe { bl.launch(cfg)?; }
1614        Ok(())
1615    }
1616
1617    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
1618    pub fn plain_tok_ring(&self, vam: &CudaSlice<u32>, pos_start: &CudaSlice<i32>,
1619                          base: usize, ring: &mut CudaSlice<u32>)
1620                          -> Result<(), Box<dyn std::error::Error>> {
1621        let f = self.func("plain_tok_ring");
1622        let (b, cap) = (base as i32, ring.len() as i32);
1623        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1624        let __s_bl = self.gpu.stream();
1625        let mut bl = __s_bl.launch_builder(&f);
1626        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
1627        unsafe { bl.launch(cfg)?; }
1628        Ok(())
1629    }
1630
1631    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
1632    pub fn spec_ring_commit(&self, vtok: &CudaSlice<u32>, acc: &CudaSlice<u32>,
1633                            brk: &CudaSlice<u32>, ring: &mut CudaSlice<u32>,
1634                            pend: &mut CudaSlice<u32>)
1635                            -> Result<(), Box<dyn std::error::Error>> {
1636        let f = self.func("spec_ring_commit");
1637        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1638        let __s_b = self.gpu.stream();
1639        let mut b = __s_b.launch_builder(&f);
1640        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
1641        unsafe { b.launch(cfg)?; }
1642        Ok(())
1643    }
1644    pub fn i32_copy_add(&self, src: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, delta: i32)
1645                        -> Result<(), Box<dyn std::error::Error>> {
1646        let f = self.func("i32_copy_add");
1647        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1648        let __s_b = self.gpu.stream();
1649        let mut b = __s_b.launch_builder(&f);
1650        b.arg(src).arg(dst).arg(&delta);
1651        unsafe { b.launch(cfg)?; }
1652        Ok(())
1653    }
1654    pub fn u32_copy(&self, src: &CudaSlice<u32>, dst: &mut CudaSlice<u32>)
1655                    -> Result<(), Box<dyn std::error::Error>> {
1656        let f = self.func("u32_copy");
1657        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1658        let __s_b = self.gpu.stream();
1659        let mut b = __s_b.launch_builder(&f);
1660        b.arg(src).arg(dst);
1661        unsafe { b.launch(cfg)?; }
1662        Ok(())
1663    }
1664
1665    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
1666    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
1667    /// caps acceptance exactly like drafting fewer tokens).
1668    pub fn spec_adapt_k(&self, acc: &CudaSlice<u32>, brk: &mut CudaSlice<u32>,
1669                        floor: usize, cap: usize)
1670                        -> Result<(), Box<dyn std::error::Error>> {
1671        let f = self.func("spec_adapt_k");
1672        let (fl, cp) = (floor as i32, cap as i32);
1673        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1674        let __s_b = self.gpu.stream();
1675        let mut b = __s_b.launch_builder(&f);
1676        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
1677        unsafe { b.launch(cfg)?; }
1678        Ok(())
1679    }
1680
1681    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
1682    pub fn spec_accept_greedy_dc(&self, preds: &CudaSlice<u32>, vtok: &CudaSlice<u32>,
1683                                 last_pred: &CudaSlice<u32>, brk: &CudaSlice<u32>,
1684                                 out: &mut CudaSlice<u32>)
1685                                 -> Result<(), Box<dyn std::error::Error>> {
1686        let f = self.func("spec_accept_greedy_dc");
1687        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1688        let __s_b = self.gpu.stream();
1689        let mut b = __s_b.launch_builder(&f);
1690        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
1691        unsafe { b.launch(cfg)?; }
1692        Ok(())
1693    }
1694
1695    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
1696    pub fn pos_iota(&self, pos0: &CudaSlice<i32>, out: &mut CudaSlice<i32>, t: usize)
1697                    -> Result<(), Box<dyn std::error::Error>> {
1698        let f = self.func("pos_iota_i32");
1699        let ti = t as i32;
1700        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (t.max(1) as u32, 1, 1),
1701                                 shared_mem_bytes: 0 };
1702        let __s_b = self.gpu.stream();
1703        let mut b = __s_b.launch_builder(&f);
1704        b.arg(pos0).arg(out).arg(&ti);
1705        unsafe { b.launch(cfg)?; }
1706        Ok(())
1707    }
1708    #[allow(clippy::too_many_arguments)]
1709    pub fn append_kv_quantized_rows_dc(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
1710                                       kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
1711                                       t0_dev: &CudaSlice<i32>, t: usize,
1712                                       kv_dim_k: usize, kv_dim_v: usize,
1713                                       k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
1714                                       -> Result<(), Box<dyn std::error::Error>> {
1715        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc") }
1716                else { self.func("append_quantize_kv_q8_0_q5_1_rows_dc") };
1717        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
1718        let cfg = LaunchConfig { grid_dim: (nblk, t as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1719        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
1720        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
1721        let __s_b = self.gpu.stream();
1722        let mut b = __s_b.launch_builder(&f);
1723        b.arg(k_rows).arg(v_rows).arg(kc).arg(vc).arg(t0_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
1724        unsafe { b.launch(cfg)?; }
1725        Ok(())
1726    }
1727
1728    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
1729    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
1730    #[allow(clippy::too_many_arguments)]
1731    pub fn append_kv_quantized_row_dc_inc(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
1732                                          kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
1733                                          t0_dev: &mut CudaSlice<i32>,
1734                                          kv_dim_k: usize, kv_dim_v: usize,
1735                                          k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
1736                                          -> Result<(), Box<dyn std::error::Error>> {
1737        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc") }
1738                else { self.func("append_quantize_kv_q8_0_q5_1_dc_inc") };
1739        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
1740        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (nthreads, 1, 1),
1741                                 shared_mem_bytes: 0 };
1742        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
1743        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
1744        let __s_b = self.gpu.stream();
1745        let mut b = __s_b.launch_builder(&f);
1746        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(t0_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
1747        unsafe { b.launch(cfg)?; }
1748        Ok(())
1749    }
1750
1751    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
1752    pub fn pack_tok_p(&self, tok: &CudaSlice<u32>, p: &CudaSlice<f32>, out: &mut CudaSlice<u32>,
1753                      slot: usize) -> Result<(), Box<dyn std::error::Error>> {
1754        let f = self.func("pack_tok_p");
1755        let sl = slot as i32;
1756        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1757        let __s_b = self.gpu.stream();
1758        let mut b = __s_b.launch_builder(&f);
1759        b.arg(tok).arg(p).arg(out).arg(&sl);
1760        unsafe { b.launch(cfg)?; }
1761        Ok(())
1762    }
1763    pub fn tok_map_u32(&self, tok: &mut CudaSlice<u32>, map: &CudaSlice<u32>)
1764                       -> Result<(), Box<dyn std::error::Error>> {
1765        let f = self.func("tok_map_u32");
1766        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1767        let __s_b = self.gpu.stream();
1768        let mut b = __s_b.launch_builder(&f);
1769        b.arg(tok).arg(map);
1770        unsafe { b.launch(cfg)?; }
1771        Ok(())
1772    }
1773
1774    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
1775    #[allow(clippy::too_many_arguments)]
1776    pub fn spec_assemble_verify(&self, tokp: &CudaSlice<u32>, pend: &CudaSlice<u32>,
1777                                d2t: Option<&CudaSlice<u32>>, vtok: &mut CudaSlice<u32>,
1778                                brk: &mut CudaSlice<u32>, p_min: f32, k: usize, pmin0: bool)
1779                                -> Result<(), Box<dyn std::error::Error>> {
1780        let f = self.func("spec_assemble_verify");
1781        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
1782        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1783        let __s_b = self.gpu.stream();
1784        let mut b = __s_b.launch_builder(&f);
1785        match d2t {
1786            Some(m) => { b.arg(tokp).arg(pend).arg(m).arg(vtok).arg(brk).arg(&p_min).arg(&ki).arg(&pm);
1787                         unsafe { b.launch(cfg)?; } }
1788            None => { let null: u64 = 0;
1789                      b.arg(tokp).arg(pend).arg(&null).arg(vtok).arg(brk).arg(&p_min).arg(&ki).arg(&pm);
1790                      unsafe { b.launch(cfg)?; } }
1791        }
1792        Ok(())
1793    }
1794
1795    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
1796    #[allow(clippy::too_many_arguments)]
1797    pub fn ssm_conv_ring_rebuild_dc(&self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>,
1798                                    conv_state: &mut CudaSlice<f32>, conv_dim: usize,
1799                                    acc: &CudaSlice<u32>, base: usize, t_v: usize, d_conv: usize)
1800                                    -> Result<(), Box<dyn std::error::Error>> {
1801        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
1802        let n = conv_dim * (d_conv - 1);
1803        let cfg = LaunchConfig::for_num_elems(n as u32);
1804        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
1805        let __s_b = self.gpu.stream();
1806        let mut b = __s_b.launch_builder(&f);
1807        b.arg(qkv_tm).arg(ring_old).arg(conv_state).arg(&cd).arg(acc).arg(&b0).arg(&tv).arg(&dc);
1808        unsafe { b.launch(cfg)?; }
1809        Ok(())
1810    }
1811    #[allow(clippy::too_many_arguments)]
1812    pub fn gdn_scan_s128_dc(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
1813                            g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
1814                            state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
1815                            n_head: usize, acc: &CudaSlice<u32>, base: usize, t_v: usize,
1816                            scale: f32)
1817                            -> Result<(), Box<dyn std::error::Error>> {
1818        let f = self.func("gdn_scan_s128_dc");
1819        const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
1820        let cfg = LaunchConfig {
1821            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
1822            block_dim: (WARP, COLS_PER_BLOCK, 1),
1823            shared_mem_bytes: 0,
1824        };
1825        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
1826        let __s_b = self.gpu.stream();
1827        let mut b = __s_b.launch_builder(&f);
1828        b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o)
1829         .arg(&h).arg(acc).arg(&b0).arg(&tv).arg(&scale);
1830        unsafe { b.launch(cfg)?; }
1831        Ok(())
1832    }
1833
1834    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
1835    pub fn spec_rollback_kv(&self, len_ptrs: &CudaSlice<u64>, saved: &CudaSlice<i32>,
1836                            acc: &CudaSlice<u32>, base: usize, n_layer: usize)
1837                            -> Result<(), Box<dyn std::error::Error>> {
1838        let f = self.func("spec_rollback_kv");
1839        let (b, nl) = (base as i32, n_layer as i32);
1840        let cfg = LaunchConfig { grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
1841                                 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
1842        let __s_bl = self.gpu.stream();
1843        let mut bl = __s_bl.launch_builder(&f);
1844        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
1845        unsafe { bl.launch(cfg)?; }
1846        Ok(())
1847    }
1848
1849    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
1850    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
1851    pub fn spec_seed_gather(&self, vx: &CudaSlice<f32>, fill_prev: &CudaSlice<f32>,
1852                            acc: &CudaSlice<u32>, h_seed: &mut CudaSlice<f32>,
1853                            base: usize, n_embd: usize)
1854                            -> Result<(), Box<dyn std::error::Error>> {
1855        let f = self.func("spec_seed_gather");
1856        let (b, ne) = (base as i32, n_embd as i32);
1857        let cfg = LaunchConfig { grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
1858                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1859        let __s_bl = self.gpu.stream();
1860        let mut bl = __s_bl.launch_builder(&f);
1861        bl.arg(vx).arg(fill_prev).arg(acc).arg(h_seed).arg(&b).arg(&ne);
1862        unsafe { bl.launch(cfg)?; }
1863        Ok(())
1864    }
1865
1866
1867    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
1868    pub fn spec_accept_greedy(&self, preds: &CudaSlice<u32>, draft: &CudaSlice<u32>,
1869                              last_pred: u32, base: usize, k_round: usize,
1870                              out: &mut CudaSlice<u32>)
1871                              -> Result<(), Box<dyn std::error::Error>> {
1872        let f = self.func("spec_accept_greedy");
1873        let (b, k) = (base as i32, k_round as i32);
1874        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1875        let __s_bl = self.gpu.stream();
1876        let mut bl = __s_bl.launch_builder(&f);
1877        bl.arg(preds).arg(draft).arg(&last_pred).arg(&b).arg(&k).arg(out);
1878        unsafe { bl.launch(cfg)?; }
1879        Ok(())
1880    }
1881
1882    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
1883    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
1884    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
1885
1886    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
1887    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
1888    pub fn gumbel_perturb(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1889                          seed: u64, stream_pos: u32, temp: f32)
1890                          -> Result<(), Box<dyn std::error::Error>> {
1891        let f = self.func("gumbel_perturb_f32");
1892        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1893        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1894        let __s_b = self.gpu.stream();
1895        let mut b = __s_b.launch_builder(&f);
1896        b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp);
1897        unsafe { b.launch(cfg)?; }
1898        Ok(())
1899    }
1900
1901    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
1902    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
1903    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
1904    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
1905    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
1906    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
1907    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
1908    pub fn mask_logits_col(&self, logits: &mut CudaSlice<f32>, mask: &CudaSlice<u32>,
1909                           col: usize, n: usize, mask_words: usize)
1910                           -> Result<(), Box<dyn std::error::Error>> {
1911        let f = self.func("mask_logits_f32");
1912        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
1913        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
1914                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1915        let __s_b = self.gpu.stream();
1916        let mut b = __s_b.launch_builder(&f);
1917        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
1918        unsafe { b.launch(cfg)?; }
1919        Ok(())
1920    }
1921
1922    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
1923    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
1924    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
1925    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
1926    /// (the lane index is the in-row position; `col` only moves the input pointer). That
1927    /// pointer-invariance IS the serving isolation contract for sampled rows.
1928    pub fn gumbel_perturb_col(&self, x: &CudaSlice<f32>, col: usize, y: &mut CudaSlice<f32>,
1929                              n: usize, seed: u64, stream_pos: u32, temp: f32)
1930                              -> Result<(), Box<dyn std::error::Error>> {
1931        let f = self.func("gumbel_perturb_f32");
1932        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1933        let col_view = x.slice(col * n..(col + 1) * n);
1934        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1935        let __s_b = self.gpu.stream();
1936        let mut b = __s_b.launch_builder(&f);
1937        b.arg(&col_view).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp);
1938        unsafe { b.launch(cfg)?; }
1939        Ok(())
1940    }
1941
1942    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
1943    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
1944    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
1945    /// reads it (counter is data, not state — graph-replay-safe).
1946    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
1947        let f = self.func("memra_sctr_inc");
1948        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
1949        let __s_b = self.gpu.stream();
1950        let mut b = __s_b.launch_builder(&f);
1951        b.arg(&mut *ctr);
1952        unsafe { b.launch(cfg)?; }
1953        Ok(())
1954    }
1955
1956    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
1957    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
1958    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
1959    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
1960    pub fn gumbel_perturb_ctr(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1961                              seed: u64, ctr: &CudaSlice<u32>, temp: f32)
1962                              -> Result<(), Box<dyn std::error::Error>> {
1963        let f = self.func("gumbel_perturb_ctr_f32");
1964        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1965        let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1966        let __s_b = self.gpu.stream();
1967        let mut b = __s_b.launch_builder(&f);
1968        b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(ctr).arg(&temp);
1969        unsafe { b.launch(cfg)?; }
1970        Ok(())
1971    }
1972
1973    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
1974    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
1975    /// (smallest-index tie-break — matches the argmax-gate contract).
1976    pub fn softmax_gather(&self, x: &CudaSlice<f32>, row_stride: usize,
1977                          ids: &CudaSlice<u32>, rows: &CudaSlice<i32>,
1978                          out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32)
1979                          -> Result<(), Box<dyn std::error::Error>> {
1980        let f = self.func("softmax_gather_f32");
1981        let (ni, rs) = (n as i32, row_stride as i64);
1982        let np = npair as i32;
1983        let cfg = LaunchConfig { grid_dim: (npair as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1984        let __s_b = self.gpu.stream();
1985        let mut b = __s_b.launch_builder(&f);
1986        b.arg(x).arg(&rs).arg(ids).arg(rows).arg(&mut *out).arg(&ni).arg(&np).arg(&temp);
1987        unsafe { b.launch(cfg)?; }
1988        Ok(())
1989    }
1990
1991    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
1992    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
1993    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
1994    pub fn residual_sample(&self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize,
1995                           temp: f32, seed: u64, stream_pos: u32,
1996                           out_tok: &mut CudaSlice<u32>)
1997                           -> Result<(), Box<dyn std::error::Error>> {
1998        let f = self.func("residual_sample_f32");
1999        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2000        let nth = 1024u32;
2001        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (nth, 1, 1), shared_mem_bytes: 0 };
2002        let has_q: i32 = q.is_some() as i32;
2003        let qbuf = q.unwrap_or(p);   // dummy when absent; kernel gates on has_q
2004        let __s_b = self.gpu.stream();
2005        let mut b = __s_b.launch_builder(&f);
2006        b.arg(p).arg(qbuf).arg(&has_q).arg(&ni).arg(&temp).arg(&slo).arg(&shi).arg(&stream_pos)
2007         .arg(&mut *out_tok);
2008        unsafe { b.launch(cfg)?; }
2009        Ok(())
2010    }
2011
2012    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
2013    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
2014    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
2015    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
2016    pub fn with_moe_cache<R>(&self, max_block_bytes: usize,
2017                             f: impl FnOnce(&mut crate::moe_cache::MoeSlotCache, &Engine) -> Result<R, Box<dyn std::error::Error>>)
2018                             -> Result<R, Box<dyn std::error::Error>> {
2019        let mut guard = self.moe_cache.lock().unwrap();
2020        if guard.is_none() {
2021            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
2022        }
2023        let cache = guard.as_mut().unwrap();
2024                f(cache, self)
2025    }
2026
2027    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
2028    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
2029    pub fn freeze_moe_cache(&self) {
2030        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
2031            cache.freeze();
2032        }
2033    }
2034
2035    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
2036    /// Never constructs a cache.
2037    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
2038        self.moe_cache
2039            .lock()
2040            .unwrap()
2041            .as_ref()
2042            .map(crate::moe_cache::MoeSlotCache::export_residency)
2043    }
2044
2045    pub(crate) fn moe_cache_frozen(&self) -> bool {
2046        self.moe_cache
2047            .lock()
2048            .unwrap()
2049            .as_ref()
2050            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
2051    }
2052
2053    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
2054    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
2055    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
2056    /// while leaving the profiling warmup's established batched behavior untouched.
2057    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
2058    /// tokenwise arm anyway.)
2059    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
2060        crate::cpu_experts::configured()
2061            && self.moe_cache_frozen()
2062            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
2063    }
2064
2065    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
2066    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
2067        assert!(
2068            self.moe_cache.lock().unwrap().is_none(),
2069            "MoE cache layout configured after cache construction"
2070        );
2071        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
2072    }
2073
2074    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
2075        self.moe_cache_layout.lock().unwrap().clone()
2076    }
2077
2078    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
2079    pub fn moe_cache_enabled() -> bool {
2080        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
2081 }
2082
2083    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
2084    /// Returns None if the cache was never built (disabled or no MoE forward ran).
2085    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
2086        let guard = self.moe_cache.lock().unwrap();
2087        guard.as_ref()            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
2088    }
2089
2090    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
2091    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
2092    /// callers compare a before/after snapshot around a decode window.
2093    pub fn cpu_expert_stats(
2094        &self,
2095    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
2096        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
2097    }
2098
2099    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
2100    /// the backend tail that resident-GPU expert work did not hide.
2101    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
2102        crate::cpu_experts::predictor_stats()
2103    }
2104
2105    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
2106        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
2107    }
2108
2109    /// CPU-routed expert selections grouped by how many of their three projections were already
2110    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
2111    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
2112        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
2113    }
2114
2115    /// Positioned-read proof-backend counters:
2116    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
2117    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
2118
2119        let guard = self.moe_cache.lock().unwrap();
2120        guard.as_ref().and_then(|cache| cache.pread_stats()).map(|stats| (
2121            stats.reads,
2122            stats.bytes,
2123            stats.read_errors,
2124            stats.short_reads,
2125            stats.fallbacks,
2126            stats.buffer_waits,
2127            stats.ring_full,
2128        ))
2129    }
2130
2131    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
2132    pub fn moe_cache_reset_counters(&self) {
2133        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() { c.reset_counters(); }
2134    }
2135
2136    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2137        Ok(self.gpu.stream().clone_htod(v)?)
2138    }
2139
2140    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
2141    /// past the final q4_0 block through their aligned window — the bytes never reach a
2142    /// result (funnelshift discards them) but must be mapped memory.
2143    pub fn htod_bytes_padded(&self, v: &[u8], pad: usize)
2144                             -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2145        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
2146        {
2147            let mut view = d.slice_mut(0..v.len());
2148            self.gpu.stream().memcpy_htod(v, &mut view)?;
2149        }
2150        Ok(d)
2151    }
2152
2153    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
2154    pub fn copy_into(&self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize)
2155                     -> Result<(), Box<dyn std::error::Error>> {
2156        let mut view = dst.slice_mut(off..off + len);
2157        self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2158        Ok(())
2159    }
2160
2161    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
2162    /// u8 twin of copy_into (D2D byte-range copy at an offset).
2163    pub fn copy_u8_into(&self, dst: &mut CudaSlice<u8>, off: usize, src: &CudaSlice<u8>, len: usize)
2164                        -> Result<(), Box<dyn std::error::Error>> {
2165        let mut view = dst.slice_mut(off..off + len);
2166        self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2167        Ok(())
2168    }
2169
2170    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
2171    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
2172    pub fn htod_u8_into(&self, dst: &mut CudaSlice<u8>, off: usize, src: &[u8])
2173                        -> Result<(), Box<dyn std::error::Error>> {
2174        let mut view = dst.slice_mut(off..off + src.len());
2175        self.gpu.stream().memcpy_htod(src, &mut view)?;
2176        Ok(())
2177    }
2178
2179    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
2180        b.slice(0..len)
2181    }
2182
2183    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
2184    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
2185    pub fn view_u8_range<'a>(&self, b: &'a CudaSlice<u8>, start: usize, end: usize)
2186                             -> cudarc::driver::CudaView<'a, u8> {
2187        b.slice(start..end)
2188    }
2189    pub fn view_u8<'a>(&self, b: &'a CudaSlice<u8>, len: usize) -> cudarc::driver::CudaView<'a, u8> {
2190        b.slice(0..len)
2191    }
2192
2193    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
2194    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
2195    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
2196    pub fn append_kv_quantized(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
2197                               kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize,
2198                               kv_dim_k: usize, kv_dim_v: usize,
2199                               k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2200                               -> Result<(), Box<dyn std::error::Error>> {
2201        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1") } else { self.func("append_quantize_kv_q8_0_q5_1") };
2202        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2203        let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2204        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
2205        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2206        let __s_b = self.gpu.stream();
2207        let mut b = __s_b.launch_builder(&f);
2208        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(&ti).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2209        unsafe { b.launch(cfg)?; }
2210        Ok(())
2211    }
2212
2213    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
2214    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
2215    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
2216    pub fn append_kv_quantized_dc(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
2217                                  kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t_dev: &CudaSlice<i32>,
2218                                  kv_dim_k: usize, kv_dim_v: usize,
2219                                  k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2220                               -> Result<(), Box<dyn std::error::Error>> {
2221        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2222        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2223        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2224        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
2225        if Self::pdl_on() && Self::pdl_wb_on() {
2226            use cudarc::driver::{DevicePtr, DevicePtrMut};
2227            let s = &self.gpu.stream();
2228            let (pk, _g0) = k_row.device_ptr(s); let (pv, _g1) = v_row.device_ptr(s);
2229            let (pkc, _g2) = kc.device_ptr_mut(s); let (pvc, _g3) = vc.device_ptr_mut(s);
2230            let (pt, _g4) = t_dev.device_ptr(s);
2231            let mut ps = [
2232                &pk as *const _ as *mut std::ffi::c_void, &pv as *const _ as *mut _,
2233                &pkc as *const _ as *mut _, &pvc as *const _ as *mut _,
2234                &pt as *const _ as *mut _, &kdk as *const _ as *mut _,
2235                &kdv as *const _ as *mut _, &ktb as *const _ as *mut _,
2236                &vtb as *const _ as *mut _,
2237            ];
2238            unsafe { self.launch_pdl_flash(g, "append_quantize_kv_q8_0_q5_1_dc",
2239                                           (nblk, 1, 1), (32, 1, 1), 0, &mut ps)?; }
2240            return Ok(());
2241        }
2242        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_dc") } else { self.func("append_quantize_kv_q8_0_q5_1_dc") };
2243        let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2244        let __s_b = self.gpu.stream();
2245        let mut b = __s_b.launch_builder(&f);
2246        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(t_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2247        unsafe { b.launch(cfg)?; }
2248        Ok(())
2249    }
2250
2251    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
2252    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
2253    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
2254    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
2255    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
2256    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
2257    #[allow(clippy::too_many_arguments)]
2258    pub fn append_kv_quantized_rows(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
2259                                    kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
2260                                    t0: usize, t: usize, kv_dim_k: usize, kv_dim_v: usize,
2261                                    k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2262                               -> Result<(), Box<dyn std::error::Error>> {
2263        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
2264            for i in 0..t {
2265                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
2266                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
2267                self.append_kv_quantized_view(&k_row, &v_row, kc, vc, t0 + i,
2268                                              kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes, g)?;
2269            }
2270            return Ok(());
2271        }
2272        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_rows") } else { self.func("append_quantize_kv_q8_0_q5_1_rows") };
2273        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2274        let cfg = LaunchConfig { grid_dim: (nblk, t as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2275        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
2276        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2277        let __s_b = self.gpu.stream();
2278        let mut b = __s_b.launch_builder(&f);
2279        b.arg(k_rows).arg(v_rows).arg(kc).arg(vc).arg(&t0i).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2280        unsafe { b.launch(cfg)?; }
2281        Ok(())
2282    }
2283
2284    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
2285    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
2286    /// later, inside a captured graph) without a host round-trip.
2287    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
2288        let f = self.func("inc_i32");
2289        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
2290        let __s_b = self.gpu.stream();
2291        let mut b = __s_b.launch_builder(&f);
2292        b.arg(p);
2293        unsafe { b.launch(cfg)?; }
2294        Ok(())
2295    }
2296
2297    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
2298    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
2299    pub fn append_kv_quantized_view(&self, k_row: &cudarc::driver::CudaView<f32>,
2300                                    v_row: &cudarc::driver::CudaView<f32>,
2301                                    kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize,
2302                                    kv_dim_k: usize, kv_dim_v: usize,
2303                                    k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2304                                    -> Result<(), Box<dyn std::error::Error>> {
2305        let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1") }
2306                else { self.func("append_quantize_kv_q8_0_q5_1") };
2307        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2308        let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2309        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
2310        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2311        let __s_b = self.gpu.stream();
2312        let mut b = __s_b.launch_builder(&f);
2313        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(&ti).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2314        unsafe { b.launch(cfg)?; }
2315        Ok(())
2316    }
2317
2318    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
2319    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
2320    pub fn copy_view_into(&self, dst: &mut CudaSlice<f32>, off: usize,
2321                          src: &cudarc::driver::CudaView<f32>, len: usize)
2322                          -> Result<(), Box<dyn std::error::Error>> {
2323        let mut view = dst.slice_mut(off..off + len);
2324        self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2325        Ok(())
2326    }
2327
2328    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
2329    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
2330    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
2331    pub fn clone_dtod(&self, src: &CudaSlice<f32>) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2332        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
2333        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
2334        Ok(dst)
2335    }
2336
2337    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
2338    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
2339    pub fn dtod_copy_view(&self, src: &cudarc::driver::CudaView<f32>, dst: &mut CudaSlice<f32>)
2340                          -> Result<(), Box<dyn std::error::Error>> {
2341        self.gpu.stream().memcpy_dtod(src, dst)?;
2342        Ok(())
2343    }
2344
2345    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
2346    pub fn dtod_copy_view_i8(&self, src: &cudarc::driver::CudaView<i8>, dst: &mut CudaSlice<i8>)
2347                             -> Result<(), Box<dyn std::error::Error>> {
2348        self.gpu.stream().memcpy_dtod(src, dst)?;
2349        Ok(())
2350    }
2351
2352    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
2353    pub fn dtod_copy_into(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, offset: usize)
2354                          -> Result<(), Box<dyn std::error::Error>> {
2355        let n = src.len();
2356        let mut dv = dst.slice_mut(offset..offset + n);
2357        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
2358        Ok(())
2359    }
2360
2361    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
2362    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
2363        self.alloc_uninit::<i8>(n)
2364    }
2365
2366    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
2367    pub fn qmatvec(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize,
2368                   qtype: i32, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2369        let f = self.func("qmatvec_f32");
2370        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
2371        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2372        let (inf, outf, mi, qt, rb) = (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
2373        let __s_b = self.gpu.stream();
2374        let mut b = __s_b.launch_builder(&f);
2375        b.arg(w).arg(x).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qt).arg(&rb);
2376        unsafe { b.launch(cfg)?; }
2377        Ok(y)
2378    }
2379
2380    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
2381    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2382        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
2383        self.keep_if_capturing(&s);
2384        Ok(s)
2385    }
2386
2387    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
2388    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
2389    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
2390    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2391        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
2392        self.keep_if_capturing(&s);
2393        Ok(s)
2394    }
2395
2396    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
2397    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
2398    pub fn memset_zeros_view(&self, dst: &mut cudarc::driver::CudaViewMut<f32>)
2399                             -> Result<(), Box<dyn std::error::Error>> {
2400        self.gpu.stream().memset_zeros(dst)?;
2401        Ok(())
2402    }
2403
2404    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
2405    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
2406    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
2407    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
2408    /// stream would require an event).
2409    pub fn stage_expert(&self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize)
2410                        -> Result<(), Box<dyn std::error::Error>> {
2411        let mut dst = scratch.slice_mut(off..off + host_bytes.len());  // CudaViewMut<u8>
2412        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?;            // accepts &[u8] HostSlice src
2413        Ok(())
2414    }
2415
2416    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
2417    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
2418    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
2419    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
2420    /// One CTA per token row, 256 threads (one per expert).
2421    pub fn moe_router_topk(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
2422                           -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2423        let f = self.func("moe_router_topk_f32");
2424        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;  // kernel fully overwrites
2425        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;    // kernel fully overwrites
2426        let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2427                                 shared_mem_bytes: 0 };
2428        let (ne, nu) = (n_expert as i32, n_used as i32);
2429        let __s_b = self.gpu.stream();
2430        let mut b = __s_b.launch_builder(&f);
2431        b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu);
2432        unsafe { b.launch(cfg)?; }
2433        Ok((sel_idx, sel_w))
2434    }
2435
2436    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
2437    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
2438    pub fn moe_router_topk_scaled(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize,
2439                                  n_used: usize, ex_scale: &CudaSlice<f32>)
2440                                  -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2441        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
2442        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
2443        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
2444        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
2445        let f = self.func("moe_router_topk_scaled_f32");
2446        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
2447        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
2448        let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2449                                 shared_mem_bytes: 0 };
2450        let (ne, nu) = (n_expert as i32, n_used as i32);
2451        let __s_b = self.gpu.stream();
2452        let mut b = __s_b.launch_builder(&f);
2453        b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu).arg(ex_scale);
2454        unsafe { b.launch(cfg)?; }
2455        Ok((sel_idx, sel_w))
2456    }
2457
2458    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
2459    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
2460    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
2461    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
2462    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
2463    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
2464    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
2465    pub fn moe_router_topk_host(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
2466                                -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
2467        let f = self.func("moe_router_topk_f32");
2468        let n = t * n_used;
2469        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
2470        let mut sel_w = self.alloc_uninit::<f32>(n)?;
2471        let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2472                                 shared_mem_bytes: 0 };
2473        let (ne, nu) = (n_expert as i32, n_used as i32);
2474        let __s_b = self.gpu.stream();
2475        let mut b = __s_b.launch_builder(&f);
2476        b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu);
2477        unsafe { b.launch(cfg)?; }
2478        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
2479        let bytes = n * 8;
2480        let mut guard = self.router_stage.lock().unwrap();
2481        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
2482            *guard = Some(PinnedStage::new(bytes.max(4096))?);
2483        }
2484        let stage = guard.as_mut().unwrap();
2485        let (si, sw) = unsafe {
2486            (std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
2487             std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n))
2488        };
2489        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;   // async (pinned dst)
2490        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;     // async (pinned dst)
2491        self.gpu.stream().synchronize()?;               // ONE sync for both
2492        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
2493    }
2494
2495    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
2496    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
2497    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
2498    pub fn stage_expert_async(&self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize)
2499                              -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
2500        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
2501        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
2502        Ok(self.copy_stream.record_event(None)?)
2503    }
2504
2505    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
2506    pub fn compute_wait(&self, ev: &cudarc::driver::CudaEvent) -> Result<(), Box<dyn std::error::Error>> {
2507        self.gpu.stream().wait(ev)?;
2508        Ok(())
2509    }
2510
2511    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
2512    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
2513    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
2514    /// CudaView base+offset pointer is honored by the launch arg.
2515    pub fn qmatvec_view(&self, w: &CudaSlice<u8>, range: std::ops::Range<usize>,
2516                        x: &cudarc::driver::CudaView<f32>, m: usize, in_f: usize, out_f: usize,
2517                        qtype: i32, row_bytes: usize)
2518                        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2519        let f = self.func("qmatvec_f32");
2520        let wv = w.slice(range);  // CudaView<u8>, offset honored
2521        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
2522        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2523        let (inf, outf, mi, qt, rb) = (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
2524        let __s_b = self.gpu.stream();
2525        let mut b = __s_b.launch_builder(&f);
2526        b.arg(&wv).arg(x).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qt).arg(&rb);
2527        unsafe { b.launch(cfg)?; }
2528        Ok(y)
2529    }
2530
2531    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
2532    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
2533    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
2534    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
2535    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
2536    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
2537    #[allow(clippy::too_many_arguments)]
2538    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
2539    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
2540    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
2541    pub fn moe_gate_up_silu8_q8(&self, gp: WPtr8, up: WPtr8,
2542                                aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2543                                in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32,
2544                                rb_g: usize, rb_u: usize)
2545                                -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2546        let f = self.func("moe_gate_up_silu8_q8");
2547        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
2548        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2549                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2550        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
2551        let __s_b = self.gpu.stream();
2552        let mut b = __s_b.launch_builder(&f);
2553        b.arg(&gp).arg(&up).arg(aq).arg(ad).arg(&mut act)
2554         .arg(&inf).arg(&nff).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2555        unsafe { b.launch(cfg)?; }
2556        Ok(act)
2557    }
2558
2559    #[allow(clippy::too_many_arguments)]
2560    pub fn moe_down8_fma_q8(&self, dp: WPtr8, w: F32x8,
2561                            aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
2562                            dst: &mut cudarc::driver::CudaViewMut<f32>,
2563                            in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize)
2564                            -> Result<(), Box<dyn std::error::Error>> {
2565        let f = self.func("moe_down8_fma_q8");
2566        let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
2567                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2568        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
2569        let __s_b = self.gpu.stream();
2570        let mut b = __s_b.launch_builder(&f);
2571        b.arg(&dp).arg(&w).arg(aq2).arg(ad2).arg(dst)
2572         .arg(&inf).arg(&outf).arg(&nu).arg(&qt).arg(&rbi);
2573        unsafe { b.launch(cfg)?; }
2574        Ok(())
2575    }
2576
2577    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
2578    pub fn qmatvec_expert_q8(&self, w: &CudaSlice<u8>, range: std::ops::Range<usize>,
2579                             aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
2580                             in_f: usize, out_f: usize, qtype: i32, row_bytes: usize)
2581                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2582        let f = self.func("qmatvec_expert_q8");
2583        let wv = w.slice(range);
2584        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
2585        const ROWS: u32 = 4;   // MEMRA_MMVQ_ROWS
2586        let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
2587                                 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2588        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
2589        let __s_b = self.gpu.stream();
2590        let mut b = __s_b.launch_builder(&f);
2591        b.arg(&wv).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qtype).arg(&rbi);
2592        unsafe { b.launch(cfg)?; }
2593        Ok(y)
2594    }
2595
2596    pub fn moe_gate_up_silu8(&self, gp: WPtr8, up: WPtr8, x: &cudarc::driver::CudaView<f32>,
2597                             in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32,
2598                             rb_g: usize, rb_u: usize)
2599                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2600        let f = self.func("moe_gate_up_silu8_f32");
2601        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;  // fully overwritten
2602        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2603                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2604        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
2605        let __s_b = self.gpu.stream();
2606        let mut b = __s_b.launch_builder(&f);
2607        b.arg(&gp).arg(&up).arg(x).arg(&mut act)
2608         .arg(&inf).arg(&nff).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2609        unsafe { b.launch(cfg)?; }
2610        Ok(act)
2611    }
2612
2613    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
2614    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
2615    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
2616    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
2617    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
2618    #[allow(clippy::too_many_arguments)]
2619    pub fn moe_down8_fma_into(&self, dp: WPtr8, w: F32x8, act: &CudaSlice<f32>,
2620                              dst: &mut cudarc::driver::CudaViewMut<f32>,
2621                              in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize)
2622                              -> Result<(), Box<dyn std::error::Error>> {
2623        let f = self.func("moe_down8_fma_f32");
2624        let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
2625                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2626        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
2627        let __s_b = self.gpu.stream();
2628        let mut b = __s_b.launch_builder(&f);
2629        b.arg(&dp).arg(&w).arg(act).arg(dst).arg(&inf).arg(&outf).arg(&nu).arg(&qt).arg(&rbv);
2630        unsafe { b.launch(cfg)?; }
2631        Ok(())
2632    }
2633
2634    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
2635    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
2636    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
2637    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
2638    #[allow(clippy::too_many_arguments)]
2639    /// dp4a q8 twin of the _dev pair (resident-experts arc).
2640    ///
2641    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
2642    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
2643    /// down's FMA chain stays slot-ordered serial). Seams:
2644    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
2645    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
2646    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
2647    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
2648    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
2649    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
2650    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
2651    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
2652    ///                       only) | w8h2 (h2 x slot-parallel)
2653    #[allow(clippy::too_many_arguments)]
2654    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
2655    #[allow(clippy::too_many_arguments)]
2656    pub fn moe_pairs_matvec_q8(&self, table: &CudaSlice<u64>, proj: i32,
2657                               pair_tok: &CudaSlice<i32>, pair_ex: &CudaSlice<i32>,
2658                               aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2659                               in_f: usize, out_f: usize, n_expert: usize, n_pairs: usize,
2660                               qtype: i32, row_bytes: usize)
2661                               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2662        let f = self.func("moe_pairs_matvec_q8");
2663        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2664        const ROWS: u32 = 4;
2665        let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
2666                                 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2667        let (inf, outf, ne, np, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2668                                        n_pairs as i32, row_bytes as i64);
2669        let __s_b = self.gpu.stream();
2670        let mut b = __s_b.launch_builder(&f);
2671        b.arg(table).arg(&proj).arg(pair_tok).arg(pair_ex).arg(aq).arg(ad).arg(&mut y)
2672         .arg(&inf).arg(&outf).arg(&ne).arg(&np).arg(&qtype).arg(&rbi);
2673        unsafe { b.launch(cfg)?; }
2674        Ok(y)
2675    }
2676
2677    /// Expert-major pair matvec (weight-reuse across each expert's token group).
2678    #[allow(clippy::too_many_arguments)]
2679    pub fn moe_pairs_matvec_q8_em(&self, table: &CudaSlice<u64>, proj: i32,
2680                                  ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>,
2681                                  ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>,
2682                                  aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2683                                  in_f: usize, out_f: usize, n_expert: usize, n_active: usize,
2684                                  n_pairs: usize, qtype: i32, row_bytes: usize)
2685                                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2686        let f = self.func("moe_pairs_matvec_q8_em");
2687        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2688        const ROWS: u32 = 4;
2689        let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
2690                                 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2691        let (inf, outf, ne, na, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2692                                        n_active as i32, row_bytes as i64);
2693        let __s_b = self.gpu.stream();
2694        let mut b = __s_b.launch_builder(&f);
2695        b.arg(table).arg(&proj).arg(ex_ids).arg(ex_off).arg(ex_pairs).arg(pair_tok)
2696         .arg(aq).arg(ad).arg(&mut y)
2697         .arg(&inf).arg(&outf).arg(&ne).arg(&na).arg(&qtype).arg(&rbi);
2698        unsafe { b.launch(cfg)?; }
2699        Ok(y)
2700    }
2701
2702    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
2703    // weight group once per (row,group) then dp4a's across the expert's token group.
2704    #[allow(clippy::too_many_arguments)]
2705    pub fn moe_pairs_matvec_q8_dec(&self, table: &CudaSlice<u64>, proj: i32,
2706                                   ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>,
2707                                   ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>,
2708                                   aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2709                                   in_f: usize, out_f: usize, n_expert: usize, n_active: usize,
2710                                   n_pairs: usize, qtype: i32, row_bytes: usize)
2711                                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2712        let f = self.func("moe_pairs_matvec_q8_dec");
2713        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2714        const ROWS: u32 = 4;
2715        let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
2716                                 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2717        let (inf, outf, ne, na, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2718                                        n_active as i32, row_bytes as i64);
2719        let __s_b = self.gpu.stream();
2720        let mut b = __s_b.launch_builder(&f);
2721        b.arg(table).arg(&proj).arg(ex_ids).arg(ex_off).arg(ex_pairs).arg(pair_tok)
2722         .arg(aq).arg(ad).arg(&mut y)
2723         .arg(&inf).arg(&outf).arg(&ne).arg(&na).arg(&qtype).arg(&rbi);
2724        unsafe { b.launch(cfg)?; }
2725        Ok(y)
2726    }
2727
2728    pub fn moe_pairs_gelu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize)
2729                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2730        let f = self.func("moe_pairs_gelu_mul");
2731        let mut act = self.alloc_uninit::<f32>(n)?;
2732        let cfg = LaunchConfig::for_num_elems(n as u32);
2733        let nl = n as i64;
2734        let __s_b = self.gpu.stream();
2735        let mut b = __s_b.launch_builder(&f);
2736        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
2737        unsafe { b.launch(cfg)?; }
2738        Ok(act)
2739    }
2740
2741    pub fn moe_pairs_silu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize)
2742                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2743        let f = self.func("moe_pairs_silu_mul");
2744        let mut act = self.alloc_uninit::<f32>(n)?;
2745        let cfg = LaunchConfig::for_num_elems(n as u32);
2746        let nl = n as i64;
2747        let __s_b = self.gpu.stream();
2748        let mut b = __s_b.launch_builder(&f);
2749        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
2750        unsafe { b.launch(cfg)?; }
2751        Ok(act)
2752    }
2753
2754    #[allow(clippy::too_many_arguments)]
2755    pub fn moe_pairs_scatter(&self, y_down: &CudaSlice<f32>, pair_w: &CudaSlice<f32>,
2756                             tok_pair_off: &CudaSlice<i32>, tok_pair_ids: &CudaSlice<i32>,
2757                             moe_out: &mut CudaSlice<f32>, t: usize, n_embd: usize)
2758                             -> Result<(), Box<dyn std::error::Error>> {
2759        let f = self.func("moe_pairs_scatter");
2760        let cfg = LaunchConfig { grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
2761                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2762        let ne = n_embd as i32;
2763        let __s_b = self.gpu.stream();
2764        let mut b = __s_b.launch_builder(&f);
2765        b.arg(y_down).arg(pair_w).arg(tok_pair_off).arg(tok_pair_ids).arg(moe_out).arg(&ne);
2766        unsafe { b.launch(cfg)?; }
2767        Ok(())
2768    }
2769
2770    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
2771    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
2772    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
2773    #[allow(clippy::too_many_arguments)]
2774    pub fn moe_gate_up_gelu8_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
2775                                    aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2776                                    in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2777                                    qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2778                                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2779        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
2780        let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
2781                                        rb_g as i64, rb_u as i64);
2782        let f = self.func("moe_gate_up_gelu8_dev_q8");
2783        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2784                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2785        let __s_b = self.gpu.stream();
2786        let mut b = __s_b.launch_builder(&f);
2787        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2788         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2789        unsafe { b.launch(cfg)?; }
2790        Ok(act)
2791    }
2792
2793    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
2794    #[allow(clippy::too_many_arguments)]
2795    pub fn moe_gate_up_gelu8_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2796                                         aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize,
2797                                         in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2798                                         qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2799                                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2800        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
2801        let (inf, nff, ne, rbg, rbu, nu) = (in_f as i32, n_ff as i32, n_expert as i32,
2802                                            rb_g as i64, rb_u as i64, n_used as i32);
2803        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
2804        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, t as u32),
2805                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2806        let __s_b = self.gpu.stream();
2807        let mut b = __s_b.launch_builder(&f);
2808        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2809         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu);
2810        unsafe { b.launch(cfg)?; }
2811        Ok(act)
2812    }
2813
2814    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
2815    #[allow(clippy::too_many_arguments)]
2816    pub fn moe_gate_up_gelu8_dev_q8_csr(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2817                                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, n_pairs: usize,
2818                                        in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2819                                        qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2820                                        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2821        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
2822        let (inf, nff, ne, rbg, rbu, nu, npi) = (in_f as i32, n_ff as i32, n_expert as i32,
2823                                                 rb_g as i64, rb_u as i64, n_used as i32,
2824                                                 n_pairs as i32);
2825        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
2826        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_pairs as u32, 1),
2827                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2828        let __s_b = self.gpu.stream();
2829        let mut b = __s_b.launch_builder(&f);
2830        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2831         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(&npi);
2832        unsafe { b.launch(cfg)?; }
2833        Ok(act)
2834    }
2835
2836    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
2837    #[allow(clippy::too_many_arguments)]
2838    pub fn moe_down8_fma_dev_q8_rows_g(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2839                                       w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
2840                                       dst: &mut CudaSlice<f32>, t: usize,
2841                                       in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
2842                                       qt: i32, rb: usize)
2843                                       -> Result<(), Box<dyn std::error::Error>> {
2844        let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
2845                                        n_expert as i32, rb as i64);
2846        let f = self.func("moe_down8_fma_dev_q8_rows_g");
2847        let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, t as u32),
2848                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2849        let __s_b = self.gpu.stream();
2850        let mut b = __s_b.launch_builder(&f);
2851        b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
2852         .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
2853        unsafe { b.launch(cfg)?; }
2854        Ok(())
2855    }
2856
2857    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
2858    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
2859    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
2860    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
2861        let (out_f, in_f) = (2048usize, 2816usize);
2862        let nblk = in_f / 32;
2863        let mut seed = 0x9E3779B97F4A7C15u64;
2864        let mut rng = move || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); (seed >> 33) as u8 };
2865        let mut w = vec![0u8; out_f * nblk * 18];
2866        for b in w.iter_mut() { *b = rng(); }
2867        for r in 0..out_f {
2868            for g in 0..nblk {
2869                let off = (r * nblk + g) * 18;
2870                w[off] = 0x00; w[off + 1] = 0x2C;   // sane half d
2871            }
2872        }
2873        let qplane = out_f * nblk * 16;
2874        let mut wrp = vec![0u8; w.len()];
2875        for r in 0..out_f {
2876            for g in 0..nblk {
2877                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
2878                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
2879                    .copy_from_slice(&src[0..2]);
2880                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
2881            }
2882        }
2883        let w_d = self.htod_bytes(&w)?;
2884        let wrp_d = self.htod_bytes(&wrp)?;
2885        let mut aq = vec![0i8; m * in_f];
2886        for v in aq.iter_mut() { *v = rng() as i8; }
2887        let aq_d = self.htod_i8(&aq)?;
2888        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
2889        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
2890        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
2891        const RPB: u32 = 4;
2892        let cfg = LaunchConfig { grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
2893                                 block_dim: (32, RPB, 1), shared_mem_bytes: 0 };
2894        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
2895        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
2896        let fb = self.func("qmatvec_q4_0_mmvq_b4");
2897        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
2898        {
2899            let __s_b = self.gpu.stream();
2900            let mut b = __s_b.launch_builder(&fb);
2901            b.arg(&w_d).arg(&aq_d).arg(&ad_d).arg(&mut y0).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
2902            unsafe { b.launch(cfg)?; }
2903            let __s_b = self.gpu.stream();
2904            let mut b = __s_b.launch_builder(&fr);
2905            b.arg(&wrp_d).arg(&aq_d).arg(&ad_d).arg(&mut y1).arg(&inf).arg(&outf).arg(&mi).arg(&qp);
2906            unsafe { b.launch(cfg)?; }
2907        }
2908        self.gpu.stream().synchronize()?;
2909        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
2910        let nd = h0.iter().zip(&h1).filter(|(a, b)| a.to_bits() != b.to_bits()).count();
2911        if nd != 0 { return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into()); }
2912        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
2913            self.gpu.stream().synchronize()?;
2914            let t0 = std::time::Instant::now();
2915            for _ in 0..500 {
2916                if rp {
2917                    let __s_b = self.gpu.stream();
2918                    let mut b = __s_b.launch_builder(&fr);
2919                    b.arg(&wrp_d).arg(&aq_d).arg(&ad_d).arg(&mut y1)
2920                     .arg(&inf).arg(&outf).arg(&mi).arg(&qp);
2921                    unsafe { b.launch(cfg)?; }
2922                } else {
2923                    let __s_b = self.gpu.stream();
2924                    let mut b = __s_b.launch_builder(&fb);
2925                    b.arg(&w_d).arg(&aq_d).arg(&ad_d).arg(&mut y0)
2926                     .arg(&inf).arg(&outf).arg(&mi).arg(&rb);
2927                    unsafe { b.launch(cfg)?; }
2928                }
2929            }
2930            self.gpu.stream().synchronize()?;
2931            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
2932        };
2933        let _ = time(false)?; let _ = time(true)?;   // warm
2934        Ok((time(false)?, time(true)?))
2935    }
2936
2937    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
2938    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
2939    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
2940    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
2941    pub fn build_q4_rp4(&self, t: &mut crate::model::GpuTensor)
2942                        -> Result<(), Box<dyn std::error::Error>> {
2943        use crate::model::GpuTensor;
2944        let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
2945        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 { return Ok(()); }
2946        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
2947        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 { return Ok(()); }
2948        let nblk = in_f / 32;
2949        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
2950        let f = self.func("q4_0_split_rp_build");
2951        let n = (out_f * nblk) as i32;
2952        let cfg = LaunchConfig { grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
2953                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2954        let (of, nb) = (out_f as i32, nblk as i32);
2955        let _ = n;
2956        let __s_b = self.gpu.stream();
2957        let mut b = __s_b.launch_builder(&f);
2958        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
2959        unsafe { b.launch(cfg)?; }
2960        *rp4 = Some(dst);
2961        Ok(())
2962    }
2963
2964    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
2965    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
2966    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
2967    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
2968    pub fn build_q8_rp4(&self, t: &mut crate::model::GpuTensor)
2969                        -> Result<(), Box<dyn std::error::Error>> {
2970        use crate::model::GpuTensor;
2971        let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
2972        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 { return Ok(()); }
2973        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
2974        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 { return Ok(()); }
2975        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
2976        Ok(())
2977    }
2978
2979    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
2980    /// mirror without a GpuTensor (same kernel the loader path above uses).
2981    pub fn build_q8_rp4_raw(&self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize)
2982                            -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2983        assert!(in_f % 32 == 0);
2984        let nblk = in_f / 32;
2985        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
2986        let f = self.func("q8_0_split_rp_build");
2987        let cfg = LaunchConfig { grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
2988                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2989        let (of, nb) = (out_f as i32, nblk as i32);
2990        let __s_b = self.gpu.stream();
2991        let mut b = __s_b.launch_builder(&f);
2992        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
2993        unsafe { b.launch(cfg)?; }
2994        Ok(dst)
2995    }
2996
2997    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
2998    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
2999    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
3000    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
3001    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
3002    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
3003    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
3004    pub fn build_q4k_rp4(&self, t: &mut crate::model::GpuTensor)
3005                         -> Result<(), Box<dyn std::error::Error>> {
3006        use crate::model::GpuTensor;
3007        let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
3008        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 { return Ok(()); }
3009        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
3010        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 { return Ok(()); }
3011        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
3012        Ok(())
3013    }
3014
3015    pub fn build_q6k_rp4(&self, t: &mut crate::model::GpuTensor)
3016                         -> Result<(), Box<dyn std::error::Error>> {
3017        use crate::model::GpuTensor;
3018        let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
3019        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 { return Ok(()); }
3020        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
3021        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 { return Ok(()); }
3022        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
3023        Ok(())
3024    }
3025
3026    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
3027    pub fn build_kq_rp4_raw(&self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, qtype: i32)
3028                            -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3029        assert!(in_f % 256 == 0);
3030        let nsbk = in_f / 256;
3031        let (sb_bytes, kname) = match qtype {
3032            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
3033            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
3034            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
3035        };
3036        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
3037        let f = self.func(kname);
3038        let cfg = LaunchConfig { grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
3039                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3040        let (of, nb) = (out_f as i32, nsbk as i32);
3041        let __s_b = self.gpu.stream();
3042        let mut b = __s_b.launch_builder(&f);
3043        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
3044        unsafe { b.launch(cfg)?; }
3045        Ok(dst)
3046    }
3047
3048    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
3049    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
3050    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
3051    pub fn kqrp_enabled() -> bool {
3052        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3053        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
3054            Ok("0") => false,
3055            Ok(_) => true,
3056            Err(_) => cfg!(memra_hopper_mma),
3057        })
3058    }
3059
3060    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
3061    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
3062    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
3063    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
3064    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
3065    pub fn build_q4_rp_swap(&self, t: &mut crate::model::GpuTensor)
3066                            -> Result<bool, Box<dyn std::error::Error>> {
3067        self.build_q4_rp4(t)?;
3068        self.gpu.stream().synchronize()?;   // build kernel reads the GGUF bytes — drain BEFORE dropping them
3069        use crate::model::GpuTensor;
3070        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else { return Ok(false) };
3071        match rp4.take() {
3072            Some(split) => {
3073                *bytes = split;   // the GGUF-layout buffer drops here
3074                *rp = true;
3075                Ok(true)
3076            }
3077            None => Ok(false),
3078        }
3079    }
3080
3081    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
3082    pub fn q4rp_enabled() -> bool {
3083        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3084        *ON.get_or_init(|| std::env::var("MEMRA_Q4RP").map(|v| v != "0").unwrap_or(true))
3085    }
3086
3087    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
3088    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
3089    pub fn copy_rows_strided(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
3090                             row_elems: usize, n_rows: usize, src_stride: usize, src_off: usize)
3091                             -> Result<(), Box<dyn std::error::Error>> {
3092        let f = self.func("copy_rows_strided_f32");
3093        let cfg = LaunchConfig { grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
3094                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3095        let (re, nr) = (row_elems as i32, n_rows as i32);
3096        let (st, off) = (src_stride as i64, src_off as i64);
3097        let __s_b = self.gpu.stream();
3098        let mut b = __s_b.launch_builder(&f);
3099        b.arg(src).arg(&mut *dst).arg(&re).arg(&nr).arg(&st).arg(&off);
3100        unsafe { b.launch(cfg)?; }
3101        Ok(())
3102    }
3103
3104    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
3105    pub fn u32_set_k(&self, dst: &mut CudaSlice<u32>, v: u32, idx: usize)
3106                     -> Result<(), Box<dyn std::error::Error>> {
3107        let f = self.func("u32_set_k");
3108        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
3109        let ii = idx as i32;
3110        let __s_b = self.gpu.stream();
3111        let mut b = __s_b.launch_builder(&f);
3112        b.arg(dst).arg(&v).arg(&ii);
3113        unsafe { b.launch(cfg)?; }
3114        Ok(())
3115    }
3116
3117    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
3118    pub fn i32_add_k(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
3119        let f = self.func("i32_add_k");
3120        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3121        let __s_b = self.gpu.stream();
3122        let mut b = __s_b.launch_builder(&f);
3123        b.arg(d).arg(&v);
3124        unsafe { b.launch(cfg)?; }
3125        Ok(())
3126    }
3127
3128    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
3129    pub fn i32_iota_from(&self, ctr: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, n: usize)
3130                         -> Result<(), Box<dyn std::error::Error>> {
3131        let f = self.func("i32_iota_from");
3132        let cfg = LaunchConfig::for_num_elems(n as u32);
3133        let ni = n as i32;
3134        let __s_b = self.gpu.stream();
3135        let mut b = __s_b.launch_builder(&f);
3136        b.arg(ctr).arg(dst).arg(&ni);
3137        unsafe { b.launch(cfg)?; }
3138        Ok(())
3139    }
3140
3141    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
3142    pub fn u32_map_k(&self, buf: &mut CudaSlice<u32>, map: &CudaSlice<u32>, idx: usize)
3143                     -> Result<(), Box<dyn std::error::Error>> {
3144        let f = self.func("u32_map_k");
3145        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
3146        let ii = idx as i32;
3147        let __s_b = self.gpu.stream();
3148        let mut b = __s_b.launch_builder(&f);
3149        b.arg(buf).arg(map).arg(&ii);
3150        unsafe { b.launch(cfg)?; }
3151        Ok(())
3152    }
3153
3154    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
3155    #[allow(clippy::too_many_arguments)]
3156    pub fn u32_pack2(&self, a: &CudaSlice<u32>, off_a: usize, n1: usize,
3157                     b_in: &CudaSlice<u32>, n2: usize, out: &mut CudaSlice<u32>)
3158                     -> Result<(), Box<dyn std::error::Error>> {
3159        let f = self.func("u32_pack2");
3160        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
3161        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
3162        let __s_b = self.gpu.stream();
3163        let mut b = __s_b.launch_builder(&f);
3164        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
3165        unsafe { b.launch(cfg)?; }
3166        Ok(())
3167    }
3168
3169    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
3170    pub fn moe_w_exscale(&self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>,
3171                         s: &CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
3172        let f = self.func("moe_w_exscale");
3173        let cfg = LaunchConfig::for_num_elems(n as u32);
3174        let ni = n as i32;
3175        let __s_b = self.gpu.stream();
3176        let mut b = __s_b.launch_builder(&f);
3177        b.arg(w).arg(sel).arg(s).arg(&ni);
3178        unsafe { b.launch(cfg)?; }
3179        Ok(())
3180    }
3181
3182    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
3183    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
3184    pub fn moe_w_scale_by_expert(&self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>,
3185                                 macros: &CudaSlice<f32>, n_expert: usize, n: usize)
3186                                 -> Result<(), Box<dyn std::error::Error>> {
3187        let f = self.func("moe_w_scale_by_expert");
3188        let cfg = LaunchConfig { grid_dim: (n.div_ceil(64) as u32, 1, 1),
3189                                 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
3190        let (ne, nn) = (n_expert as i32, n as i32);
3191        let __s_b = self.gpu.stream();
3192        let mut b = __s_b.launch_builder(&f);
3193        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
3194        unsafe { b.launch(cfg)?; }
3195        Ok(())
3196    }
3197
3198    pub fn moe_gate_up_silu8_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3199                                    aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3200                                    in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3201                                    qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3202                                    macros: &CudaSlice<f32>)
3203                                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3204        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
3205        let (mode, wpb) = GU.get_or_init(|| {
3206            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
3207            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB").ok()
3208                .and_then(|v| v.parse().ok()).unwrap_or(4u32).clamp(1, 16);
3209            (mode, wpb)
3210        });
3211        let (mode, wpb) = (mode.as_str(), *wpb);
3212        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
3213        let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3214                                        rb_g as i64, rb_u as i64);
3215        let (f, cfg) = match mode {
3216            "1" | "2" | "4" => {
3217                let rpw: u32 = mode.parse().unwrap();
3218                let f = self.func(match rpw { 1 => "moe_gate_up_silu8_dev_q8_r1",
3219                                              2 => "moe_gate_up_silu8_dev_q8_r2",
3220                                              _ => "moe_gate_up_silu8_dev_q8_r4" });
3221                let rows_per_block = (rpw * wpb) as usize;
3222                let gx = n_ff.div_ceil(rows_per_block) as u32;
3223                (f, LaunchConfig { grid_dim: (gx, n_used as u32, 1),
3224                                   block_dim: (32, wpb, 1), shared_mem_bytes: 0 })
3225            }
3226            "j8" if n_used <= 32 => (self.func("moe_gate_up_silu8_dev_q8_j8"),
3227                     LaunchConfig { grid_dim: (n_ff as u32, 1, 1),
3228                                    block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3229            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
3230            "vsm2" => {
3231                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
3232                let sh = (rb_g + rb_u) as u32;
3233                use cudarc::driver::sys::CUfunction_attribute_enum as A;
3234                f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
3235                (f, LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3236                                   block_dim: (32, 1, 1), shared_mem_bytes: sh })
3237            }
3238            "vsm" => {
3239                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
3240                let sh = (rb_g + rb_u) as u32;
3241                use cudarc::driver::sys::CUfunction_attribute_enum as A;
3242                f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
3243                (f, LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3244                                   block_dim: (32, 1, 1), shared_mem_bytes: sh })
3245            }
3246            "sg" => (self.func("moe_gate_up_silu8_dev_q8_sg"),
3247                     LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3248                                    block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3249            "j8sg" if n_used <= 32 => (self.func("moe_gate_up_silu8_dev_q8_j8sg"),
3250                     LaunchConfig { grid_dim: (n_ff as u32, 1, 1),
3251                                    block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3252            "u64" if in_f == 2048 => (self.func("moe_gate_up_silu8_dev_q8_u64"),
3253                     LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3254                                    block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3255            "gs4" if in_f == 2048 => (self.func("moe_gate_up_silu8_dev_q8_gs4"),
3256                     LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3257                                    block_dim: (32, 4, 1), shared_mem_bytes: 0 }),
3258            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
3259            "v" | "" => (self.func("moe_gate_up_silu8_dev_q8_v"),
3260                    LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3261                                   block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3262            "s2" => (self.func("moe_gate_up_silu8_dev_q8_s2"),
3263                     LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3264                                    block_dim: (32, 2, 1), shared_mem_bytes: 0 }),
3265            "s2z" => {
3266                let rz = wpb.min(16);        // s2z smem tile is [16][2]
3267                (self.func("moe_gate_up_silu8_dev_q8_s2z"),
3268                 LaunchConfig { grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
3269                                block_dim: (32, 2, rz), shared_mem_bytes: 0 })
3270            }
3271            _ => (self.func("moe_gate_up_silu8_dev_q8"),
3272                  LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3273                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3274        };
3275        let __s_b = self.gpu.stream();
3276        let mut b = __s_b.launch_builder(&f);
3277        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3278         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(macros);
3279        unsafe { b.launch(cfg)?; }
3280        Ok(act)
3281    }
3282
3283    #[allow(clippy::too_many_arguments)]
3284    pub fn moe_down8_fma_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3285                                w: &cudarc::driver::CudaView<f32>,
3286                                aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3287                                dst: &mut cudarc::driver::CudaViewMut<f32>,
3288                                in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3289                                qt: i32, rb: usize)
3290                                -> Result<(), Box<dyn std::error::Error>> {
3291        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
3292        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
3293        let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3294                                        n_expert as i32, rb as i64);
3295        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
3296        // the h2 twins are nsb==16 (in_f==512) shape-gated.
3297        let (f, cfg) = match mode.as_str() {
3298            m @ ("1" | "2" | "4") if n_used <= 8 => {
3299                let rpw: usize = m.parse().unwrap();
3300                let f = self.func(match rpw { 1 => "moe_down8_fma_dev_q8_w8r1",
3301                                              2 => "moe_down8_fma_dev_q8_w8r2",
3302                                              _ => "moe_down8_fma_dev_q8_w8r4" });
3303                (f, LaunchConfig { grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
3304                                   block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3305            }
3306            "h2" if in_f == 512 => (self.func("moe_down8_fma_dev_q8_h2"),
3307                LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3308                               block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3309            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
3310            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
3311            "" if in_f == 704 && n_used <= 8 =>
3312                (self.func("moe_down8_fma_dev_q8_w8r2"),
3313                 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3314                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3315            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
3316            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
3317            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
3318            "w8h2v" | "" if in_f == 512 && n_used <= 8 =>
3319                (self.func("moe_down8_fma_dev_q8_w8h2v"),
3320                 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3321                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3322            "w8h2r2v" if in_f == 512 && n_used <= 8 =>
3323                (self.func("moe_down8_fma_dev_q8_w8h2r2v"),
3324                 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3325                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3326            "w8h2r2" if in_f == 512 && n_used <= 8 =>
3327                (self.func("moe_down8_fma_dev_q8_w8h2r2"),
3328                 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3329                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3330            "w8h2" if in_f == 512 && n_used <= 8 =>
3331                (self.func("moe_down8_fma_dev_q8_w8h2"),
3332                 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3333                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3334            _ => (self.func("moe_down8_fma_dev_q8"),
3335                  LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3336                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3337        };
3338        let __s_b = self.gpu.stream();
3339        let mut b = __s_b.launch_builder(&f);
3340        b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3341         .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3342        unsafe { b.launch(cfg)?; }
3343        Ok(())
3344    }
3345
3346    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
3347    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
3348    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
3349    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
3350    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
3351    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
3352    #[allow(clippy::too_many_arguments)]
3353    pub fn moe_gate_up_silu8_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3354                                         aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize,
3355                                         in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3356                                         qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3357                                         macros: &CudaSlice<f32>)
3358                                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3359        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
3360        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
3361        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, t as u32),
3362                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3363        let (inf, nff, ne, nu, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3364                                            n_used as i32, rb_g as i64, rb_u as i64);
3365        let __s_b = self.gpu.stream();
3366        let mut b = __s_b.launch_builder(&f);
3367        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3368         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(macros);
3369        unsafe { b.launch(cfg)?; }
3370        Ok(act)
3371    }
3372
3373    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
3374    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
3375    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
3376    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
3377    #[allow(clippy::too_many_arguments)]
3378    pub fn moe_down8_fma_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3379                                     w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3380                                     dst: &mut CudaSlice<f32>, t: usize,
3381                                     in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3382                                     qt: i32, rb: usize)
3383                                     -> Result<(), Box<dyn std::error::Error>> {
3384        assert!(in_f == 512 && n_used <= 8, "down rows twin is w8h2v shape-gated");
3385        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
3386        let cfg = LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
3387                                 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 };
3388        let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3389                                        n_expert as i32, rb as i64);
3390        let __s_b = self.gpu.stream();
3391        let mut b = __s_b.launch_builder(&f);
3392        b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3393         .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3394        unsafe { b.launch(cfg)?; }
3395        Ok(())
3396    }
3397
3398    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
3399    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
3400    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
3401    #[allow(clippy::too_many_arguments)]
3402    pub fn moe_gate_up_silu8_dev_q8_csr(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3403                                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3404                                        n_pairs: usize, in_f: usize, n_ff: usize, n_used: usize,
3405                                        n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
3406                                        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3407        let f = self.func("moe_gate_up_silu8_dev_q8_csr_iq4");
3408        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
3409        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_pairs as u32, 1),
3410                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3411        let (inf, nff, ne, nu, npi, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3412                                                 n_used as i32, n_pairs as i32, rb_g as i64, rb_u as i64);
3413        let __s_b = self.gpu.stream();
3414        let mut b = __s_b.launch_builder(&f);
3415        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3416         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(&npi);
3417        unsafe { b.launch(cfg)?; }
3418        Ok(act)
3419    }
3420
3421
3422    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
3423    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
3424    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
3425    #[allow(clippy::too_many_arguments)]
3426    pub fn moe_down8_fma_dev_q8_variant(&self, variant: &str, table: &CudaSlice<u64>,
3427                                        sel: &cudarc::driver::CudaView<i32>,
3428                                        w: &cudarc::driver::CudaView<f32>,
3429                                        aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3430                                        dst: &mut cudarc::driver::CudaViewMut<f32>,
3431                                        in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3432                                        qt: i32, rb: usize)
3433                                        -> Result<(), Box<dyn std::error::Error>> {
3434        let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3435                                        n_expert as i32, rb as i64);
3436        let (f, cfg) = match variant {
3437            "w8h2" | "w8h2v" => {
3438                (self.func(if variant == "w8h2" { "moe_down8_fma_dev_q8_w8h2" }
3439                           else { "moe_down8_fma_dev_q8_w8h2v" }),
3440                 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3441                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3442            }
3443            "w8h2r2" | "w8h2r2v" => {
3444                (self.func(if variant == "w8h2r2" { "moe_down8_fma_dev_q8_w8h2r2" }
3445                           else { "moe_down8_fma_dev_q8_w8h2r2v" }),
3446                 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3447                                block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3448            }
3449            _ => (self.func("moe_down8_fma_dev_q8"),
3450                  LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3451                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3452        };
3453        let __s_b = self.gpu.stream();
3454        let mut b = __s_b.launch_builder(&f);
3455        b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3456         .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3457        unsafe { b.launch(cfg)?; }
3458        Ok(())
3459    }
3460
3461    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
3462    #[allow(clippy::too_many_arguments)]
3463    pub fn moe_gate_up_silu8_dev_q8_variant(&self, variant: &str, table: &CudaSlice<u64>,
3464                                            sel: &cudarc::driver::CudaView<i32>,
3465                                            aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3466                                            in_f: usize, n_ff: usize, n_used: usize,
3467                                            n_expert: usize, qt_g: i32, qt_u: i32,
3468                                            rb_g: usize, rb_u: usize)
3469                                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3470        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
3471        let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3472                                        rb_g as i64, rb_u as i64);
3473        let f = self.func(if variant == "v" { "moe_gate_up_silu8_dev_q8_v" }
3474                          else { "moe_gate_up_silu8_dev_q8" });
3475        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3476                                 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3477        let __s_b = self.gpu.stream();
3478        let mut b = __s_b.launch_builder(&f);
3479        b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3480         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
3481        unsafe { b.launch(cfg)?; }
3482        Ok(act)
3483    }
3484
3485    pub fn moe_gate_up_silu8_dev(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3486                                 x: &cudarc::driver::CudaView<f32>,
3487                                 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3488                                 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3489                                 macros: &CudaSlice<f32>)
3490                                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3491        let f = self.func("moe_gate_up_silu8_dev");
3492        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;  // fully overwritten
3493        let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3494                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3495        let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3496                                        rb_g as i64, rb_u as i64);
3497        let __s_b = self.gpu.stream();
3498        let mut b = __s_b.launch_builder(&f);
3499        b.arg(table).arg(sel).arg(x).arg(&mut act)
3500         .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(macros);
3501        unsafe { b.launch(cfg)?; }
3502        Ok(act)
3503    }
3504
3505    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
3506    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
3507    #[allow(clippy::too_many_arguments)]
3508    pub fn moe_down8_fma_dev(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3509                             w: &cudarc::driver::CudaView<f32>, act: &CudaSlice<f32>,
3510                             dst: &mut cudarc::driver::CudaViewMut<f32>,
3511                             in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3512                             qt: i32, rb: usize)
3513                             -> Result<(), Box<dyn std::error::Error>> {
3514        let f = self.func("moe_down8_fma_dev");
3515        let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3516                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3517        let (inf, outf, nu, ne, rbv) = (in_f as i32, out_f as i32, n_used as i32,
3518                                        n_expert as i32, rb as i64);
3519        let __s_b = self.gpu.stream();
3520        let mut b = __s_b.launch_builder(&f);
3521        b.arg(table).arg(sel).arg(w).arg(act).arg(dst)
3522         .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbv);
3523        unsafe { b.launch(cfg)?; }
3524        Ok(())
3525    }
3526
3527    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
3528    pub fn axpy_into(&self, src: &CudaSlice<f32>, alpha: f32,
3529                     dst: &mut cudarc::driver::CudaViewMut<f32>, n: usize)
3530                     -> Result<(), Box<dyn std::error::Error>> {
3531        let f = self.func("axpy_f32");
3532        let cfg = LaunchConfig::for_num_elems(n as u32);
3533        let (a, ni) = (alpha, n as i32);
3534        let __s_b = self.gpu.stream();
3535        let mut b = __s_b.launch_builder(&f);
3536        b.arg(src).arg(dst).arg(&a).arg(&ni);
3537        unsafe { b.launch(cfg)?; }
3538        Ok(())
3539    }
3540
3541    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
3542    pub fn add_scaled_rows(&self, src: &CudaSlice<f32>, scale: &CudaSlice<f32>,
3543                           dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize)
3544                           -> Result<(), Box<dyn std::error::Error>> {
3545        let f = self.func("add_scaled_rows_f32");
3546        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
3547        let (nc, nr) = (ncols as i32, nrows as i32);
3548        let __s_b = self.gpu.stream();
3549        let mut b = __s_b.launch_builder(&f);
3550        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
3551        unsafe { b.launch(cfg)?; }
3552        Ok(())
3553    }
3554
3555    // ======== A2 GROUPED MoE PREFILL KERNELS ========
3556
3557    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
3558    pub fn gather_rows(&self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>,
3559                       dst: &mut CudaSlice<f32>, ncols: usize, m_e: usize)
3560                       -> Result<(), Box<dyn std::error::Error>> {
3561        let f = self.func("gather_rows_f32");
3562        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
3563        let (nc, me) = (ncols as i32, m_e as i32);
3564        let __s_b = self.gpu.stream();
3565        let mut b = __s_b.launch_builder(&f);
3566        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
3567        unsafe { b.launch(cfg)?; }
3568        Ok(())
3569    }
3570
3571    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
3572    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
3573    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
3574    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
3575    pub fn scatter_slot(&self, src: &CudaSlice<f32>, tok_idx: &CudaSlice<i32>,
3576                        slot_idx: &CudaSlice<i32>, weight: &CudaSlice<f32>,
3577                        dst: &mut CudaSlice<f32>, wbuf: &mut CudaSlice<f32>,
3578                        ncols: usize, n_used: usize, m_e: usize)
3579                        -> Result<(), Box<dyn std::error::Error>> {
3580        let f = self.func("scatter_add_slot_f32");
3581        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
3582        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
3583        let __s_b = self.gpu.stream();
3584        let mut b = __s_b.launch_builder(&f);
3585        b.arg(src).arg(tok_idx).arg(slot_idx).arg(weight).arg(dst).arg(wbuf).arg(&nc).arg(&nu).arg(&me);
3586        unsafe { b.launch(cfg)?; }
3587        Ok(())
3588    }
3589
3590    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
3591    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
3592    /// Uses FMA for bit-identity with the sequential axpy path.
3593    pub fn reduce_slots(&self, slots: &CudaSlice<f32>, wbuf: &CudaSlice<f32>,
3594                        dst: &mut CudaSlice<f32>, ncols: usize, n_used: usize, t: usize)
3595                        -> Result<(), Box<dyn std::error::Error>> {
3596        let f = self.func("reduce_slots_f32");
3597        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
3598        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
3599        let __s_b = self.gpu.stream();
3600        let mut b = __s_b.launch_builder(&f);
3601        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
3602        unsafe { b.launch(cfg)?; }
3603        Ok(())
3604    }
3605
3606    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
3607    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
3608    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
3609    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
3610    /// GPU time, ~half of it redundant re-quantization of the same row.
3611    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
3612    pub fn quantize_q8_1_view(&self, x: &cudarc::driver::CudaView<f32>, m: usize, in_f: usize)
3613                     -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3614        let f = self.func("quantize_q8_1");
3615        let nblk = in_f / 32;
3616        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
3617        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
3618        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
3619        let (inf, mi) = (in_f as i32, m as i32);
3620        let __s_b = self.gpu.stream();
3621        let mut b = __s_b.launch_builder(&f);
3622        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
3623        unsafe { b.launch(cfg)?; }
3624        Ok((q, d))
3625    }
3626
3627    pub fn quantize_q8_1(&self, x: &CudaSlice<f32>, m: usize, in_f: usize)
3628                     -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3629        let nblk = in_f / 32;
3630        let mut q = self.alloc_uninit::<i8>(m * in_f)?;  // full-overwrite output: skip memset
3631        let mut d = self.alloc_uninit::<f32>(m * nblk)?;  // full-overwrite output: skip memset
3632        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
3633        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
3634        let (inf, mi) = (in_f as i32, m as i32);
3635        if Self::pdl_on() && Self::pdl_wb_on() {
3636            {
3637            use cudarc::driver::{DevicePtr, DevicePtrMut};
3638            let s = &self.gpu.stream();
3639            let (px, _g0) = x.device_ptr(s);
3640            let (pq, _g1) = q.device_ptr_mut(s); let (pd, _g2) = d.device_ptr_mut(s);
3641            let mut ps = [
3642                &px as *const _ as *mut std::ffi::c_void, &pq as *const _ as *mut _,
3643                &pd as *const _ as *mut _, &inf as *const _ as *mut _,
3644                &mi as *const _ as *mut _,
3645            ];
3646            unsafe { self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?; }
3647            }
3648            return Ok((q, d));
3649        }
3650        let f = self.func("quantize_q8_1");
3651        let __s_b = self.gpu.stream();
3652        let mut b = __s_b.launch_builder(&f);
3653        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
3654        unsafe { b.launch(cfg)?; }
3655        Ok((q, d))
3656    }
3657
3658    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
3659    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
3660    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
3661    pub fn quantize_fp4_act(&self, x: &CudaSlice<f32>, m: usize, in_f: usize)
3662                     -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
3663        let f = self.func("quantize_fp4_act");
3664        let nb16 = in_f / 16;
3665        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?;  // full-overwrite output: skip memset
3666        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?;  // full-overwrite output: skip memset
3667        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
3668        let (inf, mi) = (in_f as i32, m as i32);
3669        let __s_b = self.gpu.stream();
3670        let mut b = __s_b.launch_builder(&f);
3671        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
3672        unsafe { b.launch(cfg)?; }
3673        Ok((aq4, ad4))
3674    }
3675
3676    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
3677    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
3678    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
3679    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
3680    pub fn qmatvec_gemm_nvfp4_fp4(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3681                                  in_f: usize, out_f: usize, row_bytes: usize, scale: f32)
3682                                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3683        assert!(in_f % 64 == 0, "FP4 GEMM requires in_f % 64 == 0, got {in_f}");
3684        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
3685        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
3686        if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
3687        Ok(y)
3688    }
3689
3690    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
3691    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
3692    fn fp4_gemm_launch(&self, bytes: &CudaSlice<u8>, aq4: &CudaSlice<u32>, ad4: &CudaSlice<u8>,
3693                       m: usize, in_f: usize, out_f: usize, row_bytes: usize)
3694                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3695        let f = self.func("qmatvec_gemm_nvfp4_fp4");
3696        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
3697        const BM: u32 = 64; const BN: u32 = 256;
3698        let cfg = LaunchConfig {
3699            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
3700            block_dim: (32, 4, 1), shared_mem_bytes: 0,
3701        };
3702        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3703        let __s_b = self.gpu.stream();
3704        let mut b = __s_b.launch_builder(&f);
3705        b.arg(bytes).arg(aq4).arg(ad4).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3706        unsafe { b.launch(cfg)?; }
3707        Ok(y)
3708    }
3709
3710    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
3711    pub fn qmatvec_gemm_nvfp4_fp4_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3712                                      in_f: usize, out_f: usize, row_bytes: usize)
3713                                      -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3714        assert!(in_f % 64 == 0, "FP4 GEMM requires in_f % 64 == 0, got {in_f}");
3715        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
3716        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
3717    }
3718
3719    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
3720    pub fn qmatvec_q8_0_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3721                             out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3722        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3723        let f = self.func("qmatvec_q8_0_dp4a");
3724        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
3725        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3726        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3727        let __s_b = self.gpu.stream();
3728        let mut b = __s_b.launch_builder(&f);
3729        b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3730        unsafe { b.launch(cfg)?; }
3731        Ok(y)
3732    }
3733
3734    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
3735    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
3736    pub fn qmatvec_q4_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3737                             out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3738        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3739        let f = self.func("qmatvec_q4_K_dp4a");
3740        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
3741        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3742        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3743        let __s_b = self.gpu.stream();
3744        let mut b = __s_b.launch_builder(&f);
3745        b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3746        unsafe { b.launch(cfg)?; }
3747        Ok(y)
3748    }
3749
3750    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
3751    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
3752    pub fn qmatvec_q6_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3753                             out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3754        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3755        let f = self.func("qmatvec_q6_K_dp4a");
3756        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
3757        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3758        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3759        let __s_b = self.gpu.stream();
3760        let mut b = __s_b.launch_builder(&f);
3761        b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3762        unsafe { b.launch(cfg)?; }
3763        Ok(y)
3764    }
3765
3766    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
3767    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
3768    pub fn qmatvec_q5_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3769                             out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3770        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
3771    }
3772    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
3773    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
3774    pub fn qmatvec_q3_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3775                             out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3776        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
3777    }
3778    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
3779    pub fn qmatvec_nvfp4_fast_rp(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3780                                 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3781        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}");
3782        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
3783    }
3784    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
3785    pub fn qmatvec_nvfp4_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3786                              out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3787        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
3788        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
3789        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}");
3790        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
3791    }
3792    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
3793    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
3794    pub fn qmatvec_iq4_XS_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3795                               out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3796        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
3797    }
3798
3799    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
3800    fn qmatvec_dp4a_named(&self, name: &str, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3801                          in_f: usize, out_f: usize, row_bytes: usize)
3802                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3803        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3804        let f = self.func(name);
3805        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite output: skip memset
3806        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3807        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3808        let __s_b = self.gpu.stream();
3809        let mut b = __s_b.launch_builder(&f);
3810        b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3811        unsafe { b.launch(cfg)?; }
3812        Ok(y)
3813    }
3814
3815    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3816        Ok(self.gpu.stream().clone_htod(v)?)
3817    }
3818    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3819        Ok(self.gpu.stream().clone_htod(v)?)
3820    }
3821    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
3822    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3823        Ok(self.gpu.stream().clone_htod(v)?)
3824    }
3825    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
3826        Ok(self.gpu.stream().clone_htod(v)?)
3827    }
3828    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
3829    pub fn dtoh_view(&self, d: &cudarc::driver::CudaView<f32>)
3830                     -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3831        let v = self.gpu.stream().clone_dtoh(d)?;
3832        self.gpu.stream().synchronize()?;
3833        Ok(v)
3834    }
3835    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3836        let v = self.gpu.stream().clone_dtoh(d)?;
3837                self.gpu.stream().synchronize()?;
3838        Ok(v)
3839    }
3840    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
3841    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
3842    /// issuing them together avoids a second stream synchronization in every trunk layer.
3843    pub fn dtoh_pair(
3844        &self,
3845        a: &CudaSlice<f32>,
3846        b: &CudaSlice<f32>,
3847    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
3848        let av = self.gpu.stream().clone_dtoh(a)?;
3849        let bv = self.gpu.stream().clone_dtoh(b)?;
3850        self.gpu.stream().synchronize()?;
3851        Ok((av, bv))
3852    }
3853    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
3854    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
3855        let v = self.gpu.stream().clone_dtoh(d)?;
3856        self.gpu.stream().synchronize()?;
3857        Ok(v)
3858    }
3859    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
3860    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
3861        let v = self.gpu.stream().clone_dtoh(d)?;
3862        self.gpu.stream().synchronize()?;
3863        Ok(v)
3864    }
3865    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3866        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
3867        self.keep_if_capturing(&s);
3868        Ok(s)
3869    }
3870
3871    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
3872    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
3873    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
3874    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
3875    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
3876    /// back (or kept resident for graph replay). Returns the device token buffer.
3877    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
3878    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
3879    pub fn prob_of_token_device(&self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>, n_vocab: usize)
3880                                -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3881        let nb = ARGMAX_NB;
3882        let mut part = self.alloc_uninit::<f32>(nb)?;
3883        let mut p = self.alloc_uninit::<f32>(1)?;
3884        let f1 = self.func("prob_of_token_partial_f32");
3885        let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3886        let nv = n_vocab as i32;
3887        let __s_b1 = self.gpu.stream();
3888        let mut b1 = __s_b1.launch_builder(&f1);
3889        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
3890        unsafe { b1.launch(cfg1)?; }
3891        let f2 = self.func("prob_of_token_final_f32");
3892        let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3893        let nbi = nb as i32;
3894        let __s_b2 = self.gpu.stream();
3895        let mut b2 = __s_b2.launch_builder(&f2);
3896        b2.arg(&part).arg(&mut p).arg(&nbi);
3897        unsafe { b2.launch(cfg2)?; }
3898        Ok(p)
3899    }
3900
3901    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
3902    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
3903    /// where the host reads the p-min confidence between replays. Same kernels, same math.
3904    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
3905    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
3906    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
3907    pub fn prob_of_token_device_col(&self, logits: &CudaSlice<f32>,
3908                                    tok_all: &CudaSlice<u32>, tok_idx: usize,
3909                                    p_out: &mut CudaSlice<f32>, p_idx: usize, n_vocab: usize)
3910                                    -> Result<(), Box<dyn std::error::Error>> {
3911        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
3912        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
3913        let nb = ARGMAX_NB;
3914        let mut part = self.alloc_uninit::<f32>(nb)?;
3915        let f1 = self.func("prob_of_token_partial_f32");
3916        let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3917        let nv = n_vocab as i32;
3918        let __s_b1 = self.gpu.stream();
3919        let mut b1 = __s_b1.launch_builder(&f1);
3920        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
3921        unsafe { b1.launch(cfg1)?; }
3922        let f2 = self.func("prob_of_token_final_f32");
3923        let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3924        let nbi = nb as i32;
3925        let __s_b2 = self.gpu.stream();
3926        let mut b2 = __s_b2.launch_builder(&f2);
3927        b2.arg(&part).arg(&mut p_v).arg(&nbi);
3928        unsafe { b2.launch(cfg2)?; }
3929        Ok(())
3930    }
3931
3932    pub fn prob_of_token_device_into(&self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>,
3933                                     p_out: &mut CudaSlice<f32>, n_vocab: usize)
3934                                     -> Result<(), Box<dyn std::error::Error>> {
3935        let nb = ARGMAX_NB;
3936        let mut part = self.alloc_uninit::<f32>(nb)?;
3937        let f1 = self.func("prob_of_token_partial_f32");
3938        let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3939        let nv = n_vocab as i32;
3940        let __s_b1 = self.gpu.stream();
3941        let mut b1 = __s_b1.launch_builder(&f1);
3942        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
3943        unsafe { b1.launch(cfg1)?; }
3944        let f2 = self.func("prob_of_token_final_f32");
3945        let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3946        let nbi = nb as i32;
3947        let __s_b2 = self.gpu.stream();
3948        let mut b2 = __s_b2.launch_builder(&f2);
3949        b2.arg(&part).arg(p_out).arg(&nbi);
3950        unsafe { b2.launch(cfg2)?; }
3951        Ok(())
3952    }
3953
3954    pub fn argmax_token_device(&self, logits: &CudaSlice<f32>, n_vocab: usize)
3955                               -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
3956        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
3957        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
3958        Ok(tok)
3959    }
3960    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
3961    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
3962    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
3963    /// pointer is baked once and the token id never round-trips to host inside steady state. The
3964    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
3965    /// captured passes bake fixed addresses.
3966    pub fn argmax_token_device_into(&self, logits: &CudaSlice<f32>, tok: &mut CudaSlice<u32>,
3967                                    n_vocab: usize) -> Result<(), Box<dyn std::error::Error>> {
3968        let nb = ARGMAX_NB;
3969        let f1 = self.func("argmax_partial_f32");
3970        let f2 = self.func("argmax_final_f32");
3971        let mut guard = self.argmax_partials.lock().unwrap();
3972        if guard.is_none() {
3973            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
3974            // buffers carry no cudarc events (illegal inside capture).
3975            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
3976            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
3977            *guard = Some((pv, pi));
3978        }
3979        let (part_v, part_i) = guard.as_mut().unwrap();
3980        let nv = n_vocab as i32;
3981        let nbi = nb as i32;
3982        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
3983        let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3984        let __s_b1 = self.gpu.stream();
3985        let mut b1 = __s_b1.launch_builder(&f1);
3986        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
3987        unsafe { b1.launch(cfg1)?; }
3988        // pass 2: one block reduces NB partials -> token_out[0].
3989        let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3990        let __s_b2 = self.gpu.stream();
3991        let mut b2 = __s_b2.launch_builder(&f2);
3992        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
3993        unsafe { b2.launch(cfg2)?; }
3994        Ok(())
3995    }
3996    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
3997    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
3998    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
3999    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
4000    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
4001    pub fn argmax_token_device_col(&self, logits: &CudaSlice<f32>, col: usize, n_vocab: usize,
4002                                   toks: &mut CudaSlice<u32>, out_idx: usize)
4003                                   -> Result<(), Box<dyn std::error::Error>> {
4004        let nb = ARGMAX_NB;
4005        let f1 = self.func("argmax_partial_f32");
4006        let f2 = self.func("argmax_final_f32");
4007        let mut guard = self.argmax_partials.lock().unwrap();
4008        if guard.is_none() {
4009            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
4010            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
4011            *guard = Some((pv, pi));
4012        }
4013        let (part_v, part_i) = guard.as_mut().unwrap();
4014        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
4015        let nv = n_vocab as i32;
4016        let nbi = nb as i32;
4017        let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4018        let __s_b1 = self.gpu.stream();
4019        let mut b1 = __s_b1.launch_builder(&f1);
4020        b1.arg(&col_view).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
4021        unsafe { b1.launch(cfg1)?; }
4022        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
4023        let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4024        let __s_b2 = self.gpu.stream();
4025        let mut b2 = __s_b2.launch_builder(&f2);
4026        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
4027        unsafe { b2.launch(cfg2)?; }
4028        Ok(())
4029    }
4030    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
4031    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
4032        Ok(self.gpu.stream().clone_htod(v)?)
4033    }
4034    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4035        let v = self.gpu.stream().clone_dtoh(d)?;
4036        self.gpu.stream().synchronize()?;
4037        Ok(v)
4038    }
4039    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
4040    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
4041    /// contents change every step, the address must not, so a captured graph can read it).
4042    pub fn htod_u32_into(&self, dst: &mut CudaSlice<u32>, src: &[u32])
4043                         -> Result<(), Box<dyn std::error::Error>> {
4044        let mut view = dst.slice_mut(0..src.len());
4045        self.gpu.stream().memcpy_htod(src, &mut view)?;
4046        Ok(())
4047    }
4048
4049    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
4050        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
4051        self.keep_if_capturing(&s);
4052        Ok(s)
4053    }
4054    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
4055    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
4056    pub fn embed_gather_device_into(&self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>,
4057                                    x_out: &mut CudaSlice<f32>, n_embd: usize, qtype: i32,
4058                                    row_bytes: usize) -> Result<(), Box<dyn std::error::Error>> {
4059        let f = self.func("embed_gather_u32");
4060        let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
4061                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4062        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
4063        let __s_b = self.gpu.stream();
4064        let mut b = __s_b.launch_builder(&f);
4065        b.arg(embd).arg(token_d).arg(x_out).arg(&ne).arg(&qt).arg(&rb);
4066        unsafe { b.launch(cfg)?; }
4067        Ok(())
4068    }
4069    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
4070    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
4071        let v = self.gpu.stream().clone_dtoh(d)?;
4072        self.gpu.stream().synchronize()?;
4073        Ok(v[0])
4074    }
4075    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
4076    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
4077    /// the counter value after the throwaway capture warmups corrupt it.
4078    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
4079    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
4080    /// copy (fine at stream-idle boundaries, poison mid-round).
4081    pub fn i32_set_k(&self, dst: &mut CudaSlice<i32>, v: i32)
4082                     -> Result<(), Box<dyn std::error::Error>> {
4083        let f = self.func("i32_set_k");
4084        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
4085        let idx = 0i32;
4086        let __s_b = self.gpu.stream();
4087        let mut b = __s_b.launch_builder(&f);
4088        b.arg(dst).arg(&v).arg(&idx);
4089        unsafe { b.launch(cfg)?; }
4090        Ok(())
4091    }
4092
4093    pub fn set_i32_one(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
4094        self.gpu.stream().memcpy_htod(&[v], d)?;
4095        Ok(())
4096    }
4097    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
4098    /// during priming / capture-state restore.
4099    pub fn set_u32_one(&self, d: &mut CudaSlice<u32>, v: u32) -> Result<(), Box<dyn std::error::Error>> {
4100        self.gpu.stream().memcpy_htod(&[v], d)?;
4101        Ok(())
4102    }
4103    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
4104    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
4105        let v = self.gpu.stream().clone_dtoh(d)?;
4106        self.gpu.stream().synchronize()?;
4107        Ok(v[0])
4108    }
4109    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
4110    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4111        Ok(self.gpu.stream().clone_htod(bytes)?)
4112    }
4113    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
4114    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
4115    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
4116    pub fn embed_gather_device(&self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>,
4117                               n_embd: usize, qtype: i32, row_bytes: usize)
4118                               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4119        let f = self.func("embed_gather_u32");
4120        let mut x = self.alloc_uninit::<f32>(n_embd)?;
4121        let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
4122                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4123        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
4124        let __s_b = self.gpu.stream();
4125        let mut b = __s_b.launch_builder(&f);
4126        b.arg(embd).arg(token_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb);
4127        unsafe { b.launch(cfg)?; }
4128        Ok(x)
4129    }
4130
4131
4132    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
4133    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
4134    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
4135    pub fn embed_gather_device_t(&self, embd: &CudaSlice<u8>, tokens: &[u32],
4136                                 n_embd: usize, qtype: i32, row_bytes: usize)
4137                                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4138        let t = tokens.len();
4139        let tok_d = self.gpu.stream().clone_htod(tokens)?;
4140        let f = self.func("embed_gather_u32_t");
4141        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4142        let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4143                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4144        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4145        let __s_b = self.gpu.stream();
4146        let mut b = __s_b.launch_builder(&f);
4147        b.arg(embd).arg(&tok_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4148        unsafe { b.launch(cfg)?; }
4149        Ok(x)
4150    }
4151
4152    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
4153    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
4154    /// as embed_gather_device_t — bit-identical rows.
4155    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
4156    pub fn embed_gather_device_tv(&self, embd: &CudaSlice<u8>, tok_v: &cudarc::driver::CudaView<u32>,
4157                                  t: usize, n_embd: usize, qtype: i32, row_bytes: usize)
4158                                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4159        let f = self.func("embed_gather_u32_t");
4160        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4161        let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4162                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4163        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4164        let __s_b = self.gpu.stream();
4165        let mut b = __s_b.launch_builder(&f);
4166        b.arg(embd).arg(tok_v).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4167        unsafe { b.launch(cfg)?; }
4168        Ok(x)
4169    }
4170
4171    pub fn embed_gather_device_td(&self, embd: &CudaSlice<u8>, tok_d: &CudaSlice<u32>, t: usize,
4172                                  n_embd: usize, qtype: i32, row_bytes: usize)
4173                                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4174        let f = self.func("embed_gather_u32_t");
4175        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4176        let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4177                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4178        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4179        let __s_b = self.gpu.stream();
4180        let mut b = __s_b.launch_builder(&f);
4181        b.arg(embd).arg(tok_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4182        unsafe { b.launch(cfg)?; }
4183        Ok(x)
4184    }
4185
4186    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
4187    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
4188    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
4189    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
4190    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
4191    #[inline]
4192    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
4193    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
4194        if self.capture_keep_on.load(std::sync::atomic::Ordering::Relaxed) {
4195            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
4196        }
4197    }
4198
4199    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, n: usize)
4200            -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
4201        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
4202        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
4203        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
4204        // not cover engine-internal buffers). Debug-only: massive launch overhead.
4205        {
4206            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4207            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
4208                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
4209                use cudarc::driver::DevicePtrMut;
4210                let n_bytes = s.len() * std::mem::size_of::<T>();
4211                let stream = self.gpu.stream();
4212                let (p_, _g) = s.device_ptr_mut(&stream);
4213                unsafe {
4214                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
4215                        .result()?;
4216                }
4217            }
4218        }
4219        self.keep_if_capturing(&s);
4220        Ok(s)
4221    }
4222
4223    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
4224    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
4225    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
4226    /// consumers alloc through this (m=1 decode arms).
4227    pub fn uninit_q8_pair(&self, n: usize)
4228        -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4229        Ok((self.alloc_uninit::<i8>(n)?, self.alloc_uninit::<f32>(n / 32)?))
4230    }
4231
4232    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4233        self.alloc_uninit::<f32>(n)
4234    }
4235
4236    /// i8 uninitialized scratch (same contract as `uninit`).
4237    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4238        self.alloc_uninit::<i8>(n)
4239    }
4240
4241    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
4242    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
4243    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
4244    #[allow(clippy::too_many_arguments)]
4245    pub fn rms_norm3(&self, x: &CudaSlice<f32>, w0: &CudaSlice<f32>, w1: &CudaSlice<f32>,
4246                     w2: &CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>,
4247                     d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4248                     -> Result<(), Box<dyn std::error::Error>> {
4249        let f = self.func("rms_norm3_f32");
4250        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4251        let (nc, e) = (ncols as i32, eps);
4252        let __s_b = self.gpu.stream();
4253        let mut b = __s_b.launch_builder(&f);
4254        b.arg(x).arg(w0).arg(w1).arg(w2).arg(d0).arg(d1).arg(d2).arg(&nc).arg(&e);
4255        unsafe { b.launch(cfg)?; }
4256        Ok(())
4257    }
4258
4259    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
4260    #[allow(clippy::too_many_arguments)]
4261    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
4262    /// piggybacks on the same conditions.
4263    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
4264        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4265        *WARP_ON.get_or_init(|| {
4266            std::env::var("MEMRA_QKVNORM_W").map(|v| v != "0").unwrap_or(true)
4267        }) && ncols % 4 == 0 && rows >= 64
4268    }
4269
4270    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
4271    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
4272    #[allow(clippy::too_many_arguments)]
4273    pub fn rms_norm_qkv_w4b(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
4274                        wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4275                        dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>,
4276                        dvb: &mut CudaSlice<u8>,
4277                        ncols: usize, rq: usize, rk: usize, eps: f32, vf16: bool)
4278                        -> Result<(), Box<dyn std::error::Error>> {
4279        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
4280        let f = self.func("rms_norm_qkv_w4b_f32");
4281        let rows = (rq + 2 * rk) as u32;
4282        let cfg = LaunchConfig {
4283            grid_dim: (rows.div_ceil(8), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0,
4284        };
4285        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
4286        let vf = vf16 as i32;
4287        let __s_b = self.gpu.stream();
4288        let mut b = __s_b.launch_builder(&f);
4289        b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv).arg(&mut *dvb)
4290         .arg(&nc).arg(&rqi).arg(&rki).arg(&rvi).arg(&e).arg(&vf);
4291        unsafe { b.launch(cfg)?; }
4292        Ok(())
4293    }
4294
4295    pub fn rms_norm_qkv(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
4296                        wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4297                        dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>,
4298                        ncols: usize, rq: usize, rk: usize, eps: f32)
4299                        -> Result<(), Box<dyn std::error::Error>> {
4300        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
4301        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
4302        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
4303        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4304        let warp_on = *WARP_ON.get_or_init(|| {
4305            std::env::var("MEMRA_QKVNORM_W").map(|v| v != "0").unwrap_or(true)
4306        });
4307        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
4308        // replay numerics are untouched on every model; only prefill depth takes the new config.
4309        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
4310            let f = self.func("rms_norm_qkv_w4_f32");
4311            let rows = (rq + 2 * rk) as u32;
4312            let cfg = LaunchConfig {
4313                grid_dim: (rows.div_ceil(8), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0,
4314            };
4315            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
4316            let __s_b = self.gpu.stream();
4317            let mut b = __s_b.launch_builder(&f);
4318            b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv)
4319             .arg(&nc).arg(&rqi).arg(&rki).arg(&rvi).arg(&e);
4320            unsafe { b.launch(cfg)?; }
4321            return Ok(());
4322        }
4323        let f = self.func("rms_norm_qkv_f32");
4324        let grid = (rq + 2 * rk) as u32;
4325        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4326        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
4327        let __s_b = self.gpu.stream();
4328        let mut b = __s_b.launch_builder(&f);
4329        b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv)
4330         .arg(&nc).arg(&rqi).arg(&rki).arg(&e);
4331        unsafe { b.launch(cfg)?; }
4332        Ok(())
4333    }
4334
4335    /// gemma4 fused pair of rms_norms over two different inputs (same width).
4336    #[allow(clippy::too_many_arguments)]
4337    pub fn rms_norm2x(&self, a: &CudaSlice<f32>, bb: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4338                      wb: &CudaSlice<f32>, da: &mut CudaSlice<f32>, db: &mut CudaSlice<f32>,
4339                      ncols: usize, nrows: usize, eps: f32)
4340                      -> Result<(), Box<dyn std::error::Error>> {
4341        let f = self.func("rms_norm2x_f32");
4342        let cfg = LaunchConfig { grid_dim: (2 * nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4343        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
4344        let __s_b = self.gpu.stream();
4345        let mut b = __s_b.launch_builder(&f);
4346        b.arg(a).arg(bb).arg(wa).arg(wb).arg(da).arg(db).arg(&nc).arg(&nr).arg(&e);
4347        unsafe { b.launch(cfg)?; }
4348        Ok(())
4349    }
4350
4351    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
4352    pub fn softcap(&self, y: &mut CudaSlice<f32>, cap: f32, n: usize)
4353                   -> Result<(), Box<dyn std::error::Error>> {
4354        let f = self.func("softcap_f32");
4355        let cfg = LaunchConfig::for_num_elems(n as u32);
4356        let ni = n as i32;
4357        let __s_b = self.gpu.stream();
4358        let mut b = __s_b.launch_builder(&f);
4359        b.arg(y).arg(&cap).arg(&ni);
4360        unsafe { b.launch(cfg)?; }
4361        Ok(())
4362    }
4363
4364    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
4365    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
4366    pub fn mask_ids_rows(&self, y: &mut CudaSlice<f32>, ids: &CudaSlice<i32>, n_ids: usize,
4367                         n_vocab: usize, t: usize)
4368                         -> Result<(), Box<dyn std::error::Error>> {
4369        let f = self.func("mask_ids_rows_f32");
4370        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
4371        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
4372        let __s_b = self.gpu.stream();
4373        let mut b = __s_b.launch_builder(&f);
4374        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
4375        unsafe { b.launch(cfg)?; }
4376        Ok(())
4377    }
4378
4379    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
4380    #[allow(clippy::too_many_arguments)]
4381    pub fn add_scale_rms_norm(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4382                              w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4383                              ncols: usize, nrows: usize, eps: f32)
4384                              -> Result<(), Box<dyn std::error::Error>> {
4385        let f = self.func("add_scale_rms_norm_f32");
4386        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4387        let (nc, e2) = (ncols as i32, eps);
4388        let __s_b = self.gpu.stream();
4389        let mut b = __s_b.launch_builder(&f);
4390        b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(dst).arg(&nc).arg(&e2);
4391        unsafe { b.launch(cfg)?; }
4392        Ok(())
4393    }
4394
4395    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
4396    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
4397    #[allow(clippy::too_many_arguments)]
4398    pub fn add_scale_rms_norm_q8_1(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4399                                   w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4400                                   ncols: usize, nrows: usize, eps: f32)
4401                                   -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4402        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4403        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4404        let (nc, e2) = (ncols as i32, eps);
4405        if Self::pdl_on() && Self::pdl_wb_on() {
4406            {
4407            use cudarc::driver::{DevicePtr, DevicePtrMut};
4408            let s = &self.gpu.stream();
4409            let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b_in.device_ptr(s);
4410            let (pw, _g2) = w.device_ptr(s); let (pr, _g3) = res.device_ptr_mut(s);
4411            let (pq, _g4) = out_q.device_ptr_mut(s); let (pd, _g5) = out_d.device_ptr_mut(s);
4412            let mut ps = [
4413                &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4414                &c as *const _ as *mut _, &pw as *const _ as *mut _,
4415                &pr as *const _ as *mut _, &pq as *const _ as *mut _,
4416                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4417                &e2 as *const _ as *mut _,
4418            ];
4419            unsafe { self.launch_pdl("add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4420                                     (rms_block(), 1, 1), &mut ps)?; }
4421            }
4422            return Ok((out_q, out_d));
4423        }
4424        let f = self.func("add_scale_rms_norm_q8_1");
4425        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4426        let __s_b = self.gpu.stream();
4427        let mut b = __s_b.launch_builder(&f);
4428        b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e2);
4429        unsafe { b.launch(cfg)?; }
4430        Ok((out_q, out_d))
4431    }
4432
4433    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
4434    #[allow(clippy::too_many_arguments)]
4435    pub fn add_scale_rms_norm_q8_1_into(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4436                                        w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4437                                        ncols: usize, nrows: usize, eps: f32,
4438                                        out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>)
4439                                        -> Result<(), Box<dyn std::error::Error>> {
4440        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
4441        let (nc, e2) = (ncols as i32, eps);
4442        if Self::pdl_on() && Self::pdl_wb_on() {
4443            use cudarc::driver::{DevicePtr, DevicePtrMut};
4444            let s = &self.gpu.stream();
4445            let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b_in.device_ptr(s);
4446            let (pw, _g2) = w.device_ptr(s); let (pr, _g3) = res.device_ptr_mut(s);
4447            let (pq, _g4) = out_q.device_ptr_mut(s); let (pd, _g5) = out_d.device_ptr_mut(s);
4448            let mut ps = [
4449                &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4450                &c as *const _ as *mut _, &pw as *const _ as *mut _,
4451                &pr as *const _ as *mut _, &pq as *const _ as *mut _,
4452                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4453                &e2 as *const _ as *mut _,
4454            ];
4455            unsafe { self.launch_pdl("add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4456                                     (rms_block(), 1, 1), &mut ps)?; }
4457            return Ok(());
4458        }
4459        let f = self.func("add_scale_rms_norm_q8_1");
4460        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4461        let __s_b = self.gpu.stream();
4462        let mut b = __s_b.launch_builder(&f);
4463        b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut *out_q).arg(&mut *out_d).arg(&nc).arg(&e2);
4464        unsafe { b.launch(cfg)?; }
4465        Ok(())
4466    }
4467
4468    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
4469    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
4470    #[allow(clippy::too_many_arguments)]
4471    pub fn rms_pre_add_scale_rms_norm_q8_1(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4472                                           b_in: &CudaSlice<f32>, c: f32,
4473                                           w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4474                                           ncols: usize, nrows: usize, eps: f32)
4475                                           -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4476        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4477        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4478        let (nc, e2) = (ncols as i32, eps);
4479        if Self::pdl_on() {
4480            {
4481            use cudarc::driver::{DevicePtr, DevicePtrMut};
4482            let s = &self.gpu.stream();
4483            let (pa, _g0) = a.device_ptr(s); let (pwa, _g1) = wa.device_ptr(s);
4484            let (pb, _g2) = b_in.device_ptr(s); let (pw, _g3) = w.device_ptr(s);
4485            let (pr, _g4) = res.device_ptr_mut(s);
4486            let (pq, _g5) = out_q.device_ptr_mut(s); let (pd, _g6) = out_d.device_ptr_mut(s);
4487            let mut ps = [
4488                &pa as *const _ as *mut std::ffi::c_void, &pwa as *const _ as *mut _,
4489                &pb as *const _ as *mut _, &c as *const _ as *mut _,
4490                &pw as *const _ as *mut _, &pr as *const _ as *mut _,
4491                &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4492                &nc as *const _ as *mut _, &e2 as *const _ as *mut _,
4493            ];
4494            unsafe { self.launch_pdl("rms_pre_add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4495                                     (rms_block(), 1, 1), &mut ps)?; }
4496            }
4497            return Ok((out_q, out_d));
4498        }
4499        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
4500        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4501        let __s_b = self.gpu.stream();
4502        let mut b = __s_b.launch_builder(&f);
4503        b.arg(a).arg(wa).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e2);
4504        unsafe { b.launch(cfg)?; }
4505        Ok((out_q, out_d))
4506    }
4507
4508    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
4509    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
4510    pub fn gelu_tanh_mul_q8_1(&self, gate: &CudaSlice<f32>, up: &cudarc::driver::CudaView<f32>,
4511                              act: &mut CudaSlice<f32>, ncols: usize, nrows: usize)
4512                              -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4513        debug_assert!(ncols % 128 == 0);
4514        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4515        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4516        let nc = ncols as i32;
4517        if Self::pdl_on() {
4518            {
4519            use cudarc::driver::{DevicePtr, DevicePtrMut};
4520            let s = &self.gpu.stream();
4521            let (pg, _g0) = gate.device_ptr(s); let (pu, _g1) = up.device_ptr(s);
4522            let (pact, _g2) = act.device_ptr_mut(s);
4523            let (pq, _g3) = out_q.device_ptr_mut(s); let (pd, _g4) = out_d.device_ptr_mut(s);
4524            let mut ps = [
4525                &pg as *const _ as *mut std::ffi::c_void, &pu as *const _ as *mut _,
4526                &pact as *const _ as *mut _, &pq as *const _ as *mut _,
4527                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4528            ];
4529            unsafe { self.launch_pdl("gelu_tanh_mul_q8_1", (nrows as u32, 1, 1),
4530                                     (rms_block(), 1, 1), &mut ps)?; }
4531            }
4532            return Ok((out_q, out_d));
4533        }
4534        let f = self.func("gelu_tanh_mul_q8_1");
4535        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4536        let __s_b = self.gpu.stream();
4537        let mut b = __s_b.launch_builder(&f);
4538        b.arg(gate).arg(up).arg(act).arg(&mut out_q).arg(&mut out_d).arg(&nc);
4539        unsafe { b.launch(cfg)?; }
4540        Ok((out_q, out_d))
4541    }
4542
4543    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
4544    #[allow(clippy::too_many_arguments)]
4545    pub fn gelu_tanh_mul_q8_1_into(&self, gate: &CudaSlice<f32>, up: &cudarc::driver::CudaView<f32>,
4546                                   act: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
4547                                   out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>)
4548                                   -> Result<(), Box<dyn std::error::Error>> {
4549        debug_assert!(ncols % 128 == 0);
4550        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
4551        let nc = ncols as i32;
4552        if Self::pdl_on() {
4553            use cudarc::driver::{DevicePtr, DevicePtrMut};
4554            let s = &self.gpu.stream();
4555            let (pg, _g0) = gate.device_ptr(s); let (pu, _g1) = up.device_ptr(s);
4556            let (pact, _g2) = act.device_ptr_mut(s);
4557            let (pq, _g3) = out_q.device_ptr_mut(s); let (pd, _g4) = out_d.device_ptr_mut(s);
4558            let mut ps = [
4559                &pg as *const _ as *mut std::ffi::c_void, &pu as *const _ as *mut _,
4560                &pact as *const _ as *mut _, &pq as *const _ as *mut _,
4561                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4562            ];
4563            unsafe { self.launch_pdl("gelu_tanh_mul_q8_1", (nrows as u32, 1, 1),
4564                                     (rms_block(), 1, 1), &mut ps)?; }
4565            return Ok(());
4566        }
4567        let f = self.func("gelu_tanh_mul_q8_1");
4568        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4569        let __s_b = self.gpu.stream();
4570        let mut b = __s_b.launch_builder(&f);
4571        b.arg(gate).arg(up).arg(&mut *act).arg(&mut *out_q).arg(&mut *out_d).arg(&nc);
4572        unsafe { b.launch(cfg)?; }
4573        Ok(())
4574    }
4575
4576    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
4577    #[allow(clippy::too_many_arguments)]
4578    pub fn add_rms_norm3_q8z(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>,
4579                             w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>,
4580                             res: &mut CudaSlice<f32>, out1: &mut CudaSlice<f32>,
4581                             ncols: usize, nrows: usize, eps: f32)
4582                             -> Result<((CudaSlice<i8>, CudaSlice<f32>), (CudaSlice<i8>, CudaSlice<f32>)), Box<dyn std::error::Error>> {
4583        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
4584        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4585        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
4586        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4587        let f = self.func("add_rms_norm3_q8z_f32");
4588        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4589        let (nc, e2) = (ncols as i32, eps);
4590        let __s_b = self.gpu.stream();
4591        let mut b = __s_b.launch_builder(&f);
4592        b.arg(a).arg(b_in).arg(w0).arg(w1).arg(w2).arg(res)
4593         .arg(&mut q0).arg(&mut d0).arg(out1).arg(&mut q2).arg(&mut d2).arg(&nc).arg(&e2);
4594        unsafe { b.launch(cfg)?; }
4595        Ok(((q0, d0), (q2, d2)))
4596    }
4597
4598    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
4599    #[allow(clippy::too_many_arguments)]
4600    pub fn add_rms_norm3(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>,
4601                         w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>,
4602                         res: &mut CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>,
4603                         d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4604                         -> Result<(), Box<dyn std::error::Error>> {
4605        let f = self.func("add_rms_norm3_f32");
4606        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4607        let (nc, e2) = (ncols as i32, eps);
4608        let __s_b = self.gpu.stream();
4609        let mut b = __s_b.launch_builder(&f);
4610        b.arg(a).arg(b_in).arg(w0).arg(w1).arg(w2).arg(res).arg(d0).arg(d1).arg(d2).arg(&nc).arg(&e2);
4611        unsafe { b.launch(cfg)?; }
4612        Ok(())
4613    }
4614
4615    /// dst = (a + b) * c (residual add + layer scale, one launch).
4616    pub fn add_scale(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4617                     dst: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
4618        let f = self.func("add_scale_f32");
4619        let cfg = LaunchConfig::for_num_elems(n as u32);
4620        let ni = n as i32;
4621        let __s_b = self.gpu.stream();
4622        let mut b = __s_b.launch_builder(&f);
4623        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
4624        unsafe { b.launch(cfg)?; }
4625        Ok(())
4626    }
4627
4628    pub fn rms_norm(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4629                    ncols: usize, nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4630        let (nc, e) = (ncols as i32, eps);
4631        if Self::pdl_on() && Self::pdl_wb_on() {
4632            use cudarc::driver::{DevicePtr, DevicePtrMut};
4633            let s = &self.gpu.stream();
4634            let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4635            let (pd, _g2) = dst.device_ptr_mut(s);
4636            let mut ps = [
4637                &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4638                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4639                &e as *const _ as *mut _,
4640            ];
4641            unsafe { self.launch_pdl("rms_norm_f32", (nrows as u32, 1, 1),
4642                                     (rms_block(), 1, 1), &mut ps)?; }
4643            return Ok(());
4644        }
4645        let f = self.func("rms_norm_f32");
4646        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4647        let __s_b = self.gpu.stream();
4648        let mut b = __s_b.launch_builder(&f);
4649        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
4650        unsafe { b.launch(cfg)?; }
4651        Ok(())
4652    }
4653
4654    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
4655    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
4656    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
4657    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
4658    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
4659    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
4660    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
4661    pub fn rms_norm_decode(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4662                           ncols: usize, nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4663        let f = self.func("rms_norm_f32");
4664        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4665        let (nc, e) = (ncols as i32, eps);
4666        let __s_b = self.gpu.stream();
4667        let mut b = __s_b.launch_builder(&f);
4668        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
4669        unsafe { b.launch(cfg)?; }
4670        Ok(())
4671    }
4672
4673    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
4674    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
4675    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
4676    pub fn rms_norm_q8_1(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize, nrows: usize,
4677                         eps: f32) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4678        let nblk = ncols / 32;
4679        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
4680        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
4681        let (nc, e) = (ncols as i32, eps);
4682        if Self::pdl_on() {
4683            {
4684            use cudarc::driver::{DevicePtr, DevicePtrMut};
4685            let s = &self.gpu.stream();
4686            let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4687            let (pq, _g2) = q.device_ptr_mut(s); let (pd, _g3) = d.device_ptr_mut(s);
4688            let mut ps = [
4689                &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4690                &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4691                &nc as *const _ as *mut _, &e as *const _ as *mut _,
4692            ];
4693            unsafe { self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1),
4694                                     &mut ps)?; }
4695            }
4696            return Ok((q, d));
4697        }
4698        let f = self.func("rms_norm_q8_1");
4699        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
4700        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
4701        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4702        let __s_b = self.gpu.stream();
4703        let mut b = __s_b.launch_builder(&f);
4704        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
4705        unsafe { b.launch(cfg)?; }
4706        Ok((q, d))
4707    }
4708
4709    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
4710    /// PDL arm), caller-owned outputs.
4711    pub fn rms_norm_q8_1_into(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize,
4712                              nrows: usize, eps: f32,
4713                              q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>)
4714                              -> Result<(), Box<dyn std::error::Error>> {
4715        let nblk = ncols / 32;
4716        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
4717        let (nc, e) = (ncols as i32, eps);
4718        if Self::pdl_on() {
4719            use cudarc::driver::{DevicePtr, DevicePtrMut};
4720            let s = &self.gpu.stream();
4721            let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4722            let (pq, _g2) = q.device_ptr_mut(s); let (pd, _g3) = d.device_ptr_mut(s);
4723            let mut ps = [
4724                &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4725                &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4726                &nc as *const _ as *mut _, &e as *const _ as *mut _,
4727            ];
4728            unsafe { self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1),
4729                                     &mut ps)?; }
4730            return Ok(());
4731        }
4732        let f = self.func("rms_norm_q8_1");
4733        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4734        let __s_b = self.gpu.stream();
4735        let mut b = __s_b.launch_builder(&f);
4736        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
4737        unsafe { b.launch(cfg)?; }
4738        Ok(())
4739    }
4740
4741    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
4742    pub fn quantize_q8_1_into(&self, x: &CudaSlice<f32>, m: usize, in_f: usize,
4743                              q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>)
4744                              -> Result<(), Box<dyn std::error::Error>> {
4745        let nblk = in_f / 32;
4746        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
4747        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
4748        let (inf, mi) = (in_f as i32, m as i32);
4749        if Self::pdl_on() && Self::pdl_wb_on() {
4750            use cudarc::driver::{DevicePtr, DevicePtrMut};
4751            let s = &self.gpu.stream();
4752            let (px, _g0) = x.device_ptr(s);
4753            let (pq, _g1) = q.device_ptr_mut(s); let (pd, _g2) = d.device_ptr_mut(s);
4754            let mut ps = [
4755                &px as *const _ as *mut std::ffi::c_void, &pq as *const _ as *mut _,
4756                &pd as *const _ as *mut _, &inf as *const _ as *mut _,
4757                &mi as *const _ as *mut _,
4758            ];
4759            unsafe { self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?; }
4760            return Ok(());
4761        }
4762        let f = self.func("quantize_q8_1");
4763        let __s_b = self.gpu.stream();
4764        let mut b = __s_b.launch_builder(&f);
4765        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
4766        unsafe { b.launch(cfg)?; }
4767        Ok(())
4768    }
4769
4770    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
4771    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
4772    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
4773    pub fn add_rms_norm_q8_1(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>,
4774                             res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4775                             -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4776        let nblk = ncols / 32;
4777        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
4778        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
4779        let f = self.func("add_rms_norm_q8_1");
4780        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
4781        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4782        let (nc, e) = (ncols as i32, eps);
4783        let __s_bld = self.gpu.stream();
4784        let mut bld = __s_bld.launch_builder(&f);
4785        bld.arg(a).arg(b_in).arg(w).arg(res).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
4786        unsafe { bld.launch(cfg)?; }
4787        Ok((q, d))
4788    }
4789
4790    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
4791    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
4792    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
4793    pub fn add_rms_norm(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4794                        res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
4795                        eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4796        let (nc, e) = (ncols as i32, eps);
4797        if Self::pdl_on() && Self::pdl_wb_on() {
4798            use cudarc::driver::{DevicePtr, DevicePtrMut};
4799            let s = &self.gpu.stream();
4800            let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b.device_ptr(s);
4801            let (pw, _g2) = w.device_ptr(s);
4802            let (pr, _g3) = res.device_ptr_mut(s); let (pd, _g4) = dst.device_ptr_mut(s);
4803            let mut ps = [
4804                &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4805                &pw as *const _ as *mut _, &pr as *const _ as *mut _,
4806                &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4807                &e as *const _ as *mut _,
4808            ];
4809            unsafe { self.launch_pdl("add_rms_norm_f32", (nrows as u32, 1, 1),
4810                                     (rms_block(), 1, 1), &mut ps)?; }
4811            return Ok(());
4812        }
4813        let f = self.func("add_rms_norm_f32");
4814        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4815        let __s_b2 = self.gpu.stream();
4816        let mut b2 = __s_b2.launch_builder(&f);
4817        b2.arg(a).arg(b).arg(w).arg(&mut *res).arg(&mut *dst).arg(&nc).arg(&e);
4818        unsafe { b2.launch(cfg)?; }
4819        Ok(())
4820    }
4821
4822    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
4823    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
4824    #[allow(clippy::too_many_arguments)]
4825    pub fn rms_pre_add_rms_norm(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4826                                b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4827                                res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4828                                ncols: usize, nrows: usize, eps: f32)
4829                                -> Result<(), Box<dyn std::error::Error>> {
4830        let f = self.func("rms_pre_add_rms_norm_f32");
4831        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4832        let (nc, e) = (ncols as i32, eps);
4833        let __s_b2 = self.gpu.stream();
4834        let mut b2 = __s_b2.launch_builder(&f);
4835        b2.arg(a).arg(wa).arg(b).arg(w).arg(&mut *res).arg(&mut *dst).arg(&nc).arg(&e);
4836        unsafe { b2.launch(cfg)?; }
4837        Ok(())
4838    }
4839
4840    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
4841    #[allow(clippy::too_many_arguments)]
4842    pub fn rms_pre_add_rms_norm_q8z(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4843                                    b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4844                                    res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4845                                    ncols: usize, nrows: usize, eps: f32)
4846                                    -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4847        debug_assert!(ncols % 128 == 0);
4848        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4849        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4850        let (nc, e) = (ncols as i32, eps);
4851        if Self::pdl_on() {
4852            {
4853            use cudarc::driver::{DevicePtr, DevicePtrMut};
4854            let s = &self.gpu.stream();
4855            let (pa, _g0) = a.device_ptr(s); let (pwa, _g1) = wa.device_ptr(s);
4856            let (pb, _g2) = b.device_ptr(s); let (pw, _g3) = w.device_ptr(s);
4857            let (pr, _g4) = res.device_ptr_mut(s); let (pdst, _g5) = dst.device_ptr_mut(s);
4858            let (pq, _g6) = out_q.device_ptr_mut(s); let (pd, _g7) = out_d.device_ptr_mut(s);
4859            let mut ps = [
4860                &pa as *const _ as *mut std::ffi::c_void, &pwa as *const _ as *mut _,
4861                &pb as *const _ as *mut _, &pw as *const _ as *mut _,
4862                &pr as *const _ as *mut _, &pdst as *const _ as *mut _,
4863                &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4864                &nc as *const _ as *mut _, &e as *const _ as *mut _,
4865            ];
4866            unsafe { self.launch_pdl("rms_pre_add_rms_norm_q8z_f32", (nrows as u32, 1, 1),
4867                                     (rms_block(), 1, 1), &mut ps)?; }
4868            }
4869            return Ok((out_q, out_d));
4870        }
4871        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
4872        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4873        let __s_b2 = self.gpu.stream();
4874        let mut b2 = __s_b2.launch_builder(&f);
4875        b2.arg(a).arg(wa).arg(b).arg(w).arg(&mut *res).arg(&mut *dst)
4876          .arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e);
4877        unsafe { b2.launch(cfg)?; }
4878        Ok((out_q, out_d))
4879    }
4880
4881    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
4882    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
4883    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
4884    pub fn build_q4_out_concat3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
4885                                w2: &crate::model::GpuTensor)
4886                                -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
4887        use crate::model::GpuTensor;
4888        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
4889            match w {
4890                GpuTensor::Quant { qtype, row_bytes, rp, .. }
4891                    if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
4892                _ => None,
4893            }
4894        };
4895        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
4896        else { return Ok(None) };
4897        if rb0 != rb1 || rb0 != rb2
4898            || w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
4899            return Ok(None);
4900        }
4901        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
4902            match w { crate::model::GpuTensor::Quant { bytes, .. } => bytes, _ => unreachable!() }
4903        }
4904        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
4905        let total = rb0 * (o0 + o1 + o2);
4906        let mut cat = self.alloc_u8(total)?;
4907        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
4908        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
4909        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
4910        Ok(Some(GpuTensor::Quant {
4911            bytes: cat, qtype: QT_Q4_0, row_bytes: rb0,
4912            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64], scale: 1.0, rp: false,
4913            #[cfg(memra_cutlass)]
4914            cutlass: None,
4915            fp8: None, rp4: None, f16: None,
4916        }))
4917    }
4918
4919    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
4920    #[allow(clippy::too_many_arguments)]
4921    pub fn rms_norm_qkv_rope_cat(&self, qkv: &CudaSlice<f32>,
4922                                 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4923                                 q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
4924                                 head_dim: usize, rq: usize, rk: usize,
4925                                 pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
4926                                 base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32)
4927                                 -> Result<(), Box<dyn std::error::Error>> {
4928        let rows = rq + rk + rk;
4929        let theta_scale = base.powf(-2.0 / head_dim as f32);
4930        let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
4931        if Self::pdl_on() {
4932            use cudarc::driver::{DevicePtr, DevicePtrMut};
4933            let s = &self.gpu.stream();
4934            let (pqkv, _g0) = qkv.device_ptr(s);
4935            let (pwq, _g1) = wq.device_ptr(s); let (pwk, _g2) = wk.device_ptr(s);
4936            let (pwv, _g3) = wv.device_ptr(s);
4937            let (pq, _g4) = q.device_ptr_mut(s); let (pk, _g5) = k.device_ptr_mut(s);
4938            let (pv, _g6) = v.device_ptr_mut(s);
4939            let (ppos, _g7) = pos.device_ptr(s);
4940            let (pff, _g8) = match ff {
4941                Some(t) => { let (p, g) = t.device_ptr(s); (p, Some(g)) }
4942                None => (0, None),
4943            };
4944            let mut ps = [
4945                &pqkv as *const _ as *mut std::ffi::c_void,
4946                &pwq as *const _ as *mut _, &pwk as *const _ as *mut _,
4947                &pwv as *const _ as *mut _,
4948                &pq as *const _ as *mut _, &pk as *const _ as *mut _,
4949                &pv as *const _ as *mut _,
4950                &nc as *const _ as *mut _, &rqi as *const _ as *mut _,
4951                &rki as *const _ as *mut _, &ppos as *const _ as *mut _,
4952                &nhq as *const _ as *mut _, &nhk as *const _ as *mut _,
4953                &theta_scale as *const _ as *mut _, &freq_scale as *const _ as *mut _,
4954                &pff as *const _ as *mut _, &eps as *const _ as *mut _,
4955            ];
4956            unsafe { self.launch_pdl("rms_norm_qkv_rope_cat_f32", (rows as u32, 1, 1),
4957                                     (rms_block(), 1, 1), &mut ps)?; }
4958            return Ok(());
4959        }
4960        let f = self.func("rms_norm_qkv_rope_cat_f32");
4961        let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4962        let __s_b = self.gpu.stream();
4963        let mut b = __s_b.launch_builder(&f);
4964        match ff {
4965            Some(t) => { b.arg(qkv).arg(wq).arg(wk).arg(wv)
4966                          .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4967                          .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
4968                          .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps);
4969                         unsafe { b.launch(cfg)?; } }
4970            None => { let null: u64 = 0;
4971                      b.arg(qkv).arg(wq).arg(wk).arg(wv)
4972                       .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4973                       .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
4974                       .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps);
4975                      unsafe { b.launch(cfg)?; } }
4976        }
4977        Ok(())
4978    }
4979
4980    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
4981    #[allow(clippy::too_many_arguments)]
4982    pub fn rms_norm_qkv_rope(&self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>, v0: &CudaSlice<f32>,
4983                             wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4984                             q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
4985                             head_dim: usize, rq: usize, rk: usize,
4986                             pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
4987                             base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32)
4988                             -> Result<(), Box<dyn std::error::Error>> {
4989        let f = self.func("rms_norm_qkv_rope_f32");
4990        let rows = rq + rk + rk;   // q rows + k rows + v rows (rk == rv)
4991        let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4992        let theta_scale = base.powf(-2.0 / head_dim as f32);
4993        let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
4994        let __s_b = self.gpu.stream();
4995        let mut b = __s_b.launch_builder(&f);
4996        match ff {
4997            Some(t) => { b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
4998                          .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4999                          .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5000                          .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps);
5001                         unsafe { b.launch(cfg)?; } }
5002            None => { let null: u64 = 0;
5003                      b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5004                       .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5005                       .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5006                       .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps);
5007                      unsafe { b.launch(cfg)?; } }
5008        }
5009        Ok(())
5010    }
5011
5012    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
5013    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
5014    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
5015    #[allow(clippy::too_many_arguments)]
5016    pub fn rms_norm_qkv_rope_append_dc(&self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>,
5017                             v0: &CudaSlice<f32>,
5018                             wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
5019                             q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
5020                             head_dim: usize, rq: usize, rk: usize,
5021                             pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
5022                             base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32,
5023                             kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
5024                             t_dev: &CudaSlice<i32>, k_tok_bytes: usize, v_tok_bytes: usize,
5025                             g: bool)
5026                             -> Result<(), Box<dyn std::error::Error>> {
5027        let rows = rq + rk + rk;
5028        let theta_scale = base.powf(-2.0 / head_dim as f32);
5029        let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
5030        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5031        if Self::pdl_on() && Self::pdl_wb_on() {
5032            use cudarc::driver::{DevicePtr, DevicePtrMut};
5033            let s = &self.gpu.stream();
5034            let (p0, _a0) = q0.device_ptr(s); let (p1, _a1) = k0.device_ptr(s);
5035            let (p2, _a2) = v0.device_ptr(s);
5036            let (pwq, _a3) = wq.device_ptr(s); let (pwk, _a4) = wk.device_ptr(s);
5037            let (pwv, _a5) = wv.device_ptr(s);
5038            let (pq, _a6) = q.device_ptr_mut(s); let (pk, _a7) = k.device_ptr_mut(s);
5039            let (pv, _a8) = v.device_ptr_mut(s);
5040            let (pp, _a9) = pos.device_ptr(s);
5041            let pff: u64 = match ff { Some(t) => { let (p, _gg) = t.device_ptr(s); p as u64 }
5042                                      None => 0 };
5043            let (pkc, _a10) = kc.device_ptr_mut(s); let (pvc, _a11) = vc.device_ptr_mut(s);
5044            let (pt, _a12) = t_dev.device_ptr(s);
5045            let mut ps = [
5046                &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
5047                &p2 as *const _ as *mut _, &pwq as *const _ as *mut _,
5048                &pwk as *const _ as *mut _, &pwv as *const _ as *mut _,
5049                &pq as *const _ as *mut _, &pk as *const _ as *mut _,
5050                &pv as *const _ as *mut _, &nc as *const _ as *mut _,
5051                &rqi as *const _ as *mut _, &rki as *const _ as *mut _,
5052                &pp as *const _ as *mut _, &nhq as *const _ as *mut _,
5053                &nhk as *const _ as *mut _, &theta_scale as *const _ as *mut _,
5054                &freq_scale as *const _ as *mut _, &pff as *const _ as *mut _,
5055                &eps as *const _ as *mut _, &pkc as *const _ as *mut _,
5056                &pvc as *const _ as *mut _, &pt as *const _ as *mut _,
5057                &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
5058            ];
5059            unsafe { self.launch_pdl_flash(g, "rms_norm_qkv_rope_append_dc_f32",
5060                                           (rows as u32, 1, 1), (rms_block(), 1, 1), 0, &mut ps)?; }
5061            return Ok(());
5062        }
5063        let f = if g { self.func_g("rms_norm_qkv_rope_append_dc_f32") }
5064                else { self.func("rms_norm_qkv_rope_append_dc_f32") };
5065        let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
5066        let __s_b = self.gpu.stream();
5067        let mut b = __s_b.launch_builder(&f);
5068        match ff {
5069            Some(t) => { b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5070                          .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5071                          .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5072                          .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps)
5073                          .arg(&mut *kc).arg(&mut *vc).arg(t_dev).arg(&ktb).arg(&vtb);
5074                         unsafe { b.launch(cfg)?; } }
5075            None => { let null: u64 = 0;
5076                      b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5077                       .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5078                       .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5079                       .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps)
5080                       .arg(&mut *kc).arg(&mut *vc).arg(t_dev).arg(&ktb).arg(&vtb);
5081                      unsafe { b.launch(cfg)?; } }
5082        }
5083        Ok(())
5084    }
5085
5086    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
5087    pub fn add_q8_1(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
5088                    ncols: usize, nrows: usize)
5089                    -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5090        debug_assert!(ncols % 128 == 0);
5091        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
5092        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
5093        let f = self.func("add_q8_1_f32");
5094        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
5095        let nc = ncols as i32;
5096        let __s_b2 = self.gpu.stream();
5097        let mut b2 = __s_b2.launch_builder(&f);
5098        b2.arg(a).arg(b).arg(&mut *res).arg(&mut out_q).arg(&mut out_d).arg(&nc);
5099        unsafe { b2.launch(cfg)?; }
5100        Ok((out_q, out_d))
5101    }
5102
5103    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
5104    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
5105    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
5106    pub fn rms_pre_add_q8_1(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b: &CudaSlice<f32>,
5107                            res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
5108                            -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5109        debug_assert!(ncols % 128 == 0);
5110        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
5111        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
5112        let f = self.func("rms_pre_add_q8_1_f32");
5113        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1),
5114                                 shared_mem_bytes: 0 };
5115        let (nc, ep) = (ncols as i32, eps);
5116        let __s_b2 = self.gpu.stream();
5117        let mut b2 = __s_b2.launch_builder(&f);
5118        b2.arg(a).arg(wa).arg(b).arg(&mut *res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&ep);
5119        unsafe { b2.launch(cfg)?; }
5120        Ok((out_q, out_d))
5121    }
5122
5123    /// L2 norm per row (head_dim), no weight.
5124    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
5125    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
5126    pub fn l2_v2_on(ncols: usize) -> bool {
5127        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
5128    }
5129
5130    pub fn l2_norm_pp(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
5131                      dst16: Option<&mut CudaSlice<u8>>, ncols: usize, nrows: usize,
5132                      eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5133        if Self::l2_v2_on(ncols) {
5134            let f = self.func("l2_norm_pp_v2_f32");
5135            let rows_per_block = 8u32;   // 256 threads = 8 warps = 8 rows
5136            let cfg = LaunchConfig { grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
5137            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
5138            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
5139            let d16: u64 = match dst16 { Some(d) => self.addr_u8(d), None => 0 };
5140            let __s_b = self.gpu.stream();
5141            let mut b = __s_b.launch_builder(&f);
5142            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
5143            unsafe { b.launch(cfg)?; }
5144            return Ok(());
5145        }
5146        self.l2_norm(x, dst, ncols, nrows, eps)
5147    }
5148
5149    pub fn l2_norm(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
5150                   eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5151        let f = self.func("l2_norm_f32");
5152        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
5153        let (nc, e) = (ncols as i32, eps);
5154        let __s_b = self.gpu.stream();
5155        let mut b = __s_b.launch_builder(&f);
5156        b.arg(x).arg(dst).arg(&nc).arg(&e);
5157        unsafe { b.launch(cfg)?; }
5158        Ok(())
5159    }
5160
5161    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
5162    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
5163    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
5164    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
5165    /// propagate through gdn_scan and flip argmax on marginal logits.
5166    pub fn l2_norm_decode(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize,
5167                          nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5168        let f = self.func("l2_norm_f32");
5169        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
5170        let (nc, e) = (ncols as i32, eps);
5171        let __s_b = self.gpu.stream();
5172        let mut b = __s_b.launch_builder(&f);
5173        b.arg(x).arg(dst).arg(&nc).arg(&e);
5174        unsafe { b.launch(cfg)?; }
5175        Ok(())
5176    }
5177
5178    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
5179    pub fn rope_neox(&self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize,
5180                     n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32, freq_scale: f32)
5181                     -> Result<(), Box<dyn std::error::Error>> {
5182        let f = self.func("rope_neox_f32");
5183        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5184        let grid = (n_heads * n_tokens) as u32;
5185        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5186        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
5187        let __s_b = self.gpu.stream();
5188        let mut b = __s_b.launch_builder(&f);
5189        b.arg(x).arg(pos).arg(&hd).arg(&nd).arg(&nh).arg(&theta_scale).arg(&freq_scale);
5190        unsafe { b.launch(cfg)?; }
5191        Ok(())
5192    }
5193
5194    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
5195    pub fn rope_neox_ff(&self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize,
5196                        n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32,
5197                        freq_scale: f32, ff: &CudaSlice<f32>)
5198                        -> Result<(), Box<dyn std::error::Error>> {
5199        let f = self.func("rope_neox_ff_f32");
5200        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5201        let grid = (n_heads * n_tokens) as u32;
5202        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5203        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
5204        let __s_b = self.gpu.stream();
5205        let mut b = __s_b.launch_builder(&f);
5206        b.arg(x).arg(pos).arg(&hd).arg(&nd).arg(&nh).arg(&theta_scale).arg(&freq_scale).arg(ff);
5207        unsafe { b.launch(cfg)?; }
5208        Ok(())
5209    }
5210
5211    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
5212    #[allow(clippy::too_many_arguments)]
5213    pub fn rope_neox2(&self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>,
5214                      pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize,
5215                      nh_q: usize, nh_k: usize, n_tokens: usize, freq_base: f32,
5216                      freq_scale: f32, ff: Option<&CudaSlice<f32>>)
5217                      -> Result<(), Box<dyn std::error::Error>> {
5218        let f = self.func("rope_neox2_f32");
5219        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5220        let grid = ((nh_q + nh_k) * n_tokens) as u32;
5221        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5222        let (hd, nd, nq, nk, nt) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32, n_tokens as i32);
5223        let __s_b = self.gpu.stream();
5224        let mut b = __s_b.launch_builder(&f);
5225        b.arg(q).arg(k).arg(pos).arg(&hd).arg(&nd).arg(&nq).arg(&nk).arg(&nt)
5226         .arg(&theta_scale).arg(&freq_scale);
5227        match ff {
5228            Some(ffv) => { b.arg(ffv); unsafe { b.launch(cfg)?; } }
5229            None => {
5230                let null: u64 = 0;
5231                b.arg(&null);
5232                unsafe { b.launch(cfg)?; }
5233            }
5234        }
5235        Ok(())
5236    }
5237
5238    /// gemma4 R1: dst = GELU_tanh(gate) * up.
5239    pub fn gelu_tanh_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5240                         -> Result<(), Box<dyn std::error::Error>> {
5241        let f = self.func("gelu_tanh_mul_f32");
5242        let cfg = LaunchConfig::for_num_elems(n as u32);
5243        let ni = n as i32;
5244        let __s_b = self.gpu.stream();
5245        let mut b = __s_b.launch_builder(&f);
5246        b.arg(gate).arg(up).arg(dst).arg(&ni);
5247        unsafe { b.launch(cfg)?; }
5248        Ok(())
5249    }
5250
5251    pub fn silu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5252                    -> Result<(), Box<dyn std::error::Error>> {
5253        let f = self.func("silu_mul_f32");
5254        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
5255        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5256        let ni = n as i32;
5257        let __s_b = self.gpu.stream();
5258        let mut b = __s_b.launch_builder(&f);
5259        b.arg(gate).arg(up).arg(dst).arg(&ni);
5260        unsafe { b.launch(cfg)?; }
5261        Ok(())
5262    }
5263
5264    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
5265    /// for the down projection — kills the standalone convert pass. Bit-identical class.
5266    pub fn silu_mul_f16out(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
5267                           dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize)
5268                           -> Result<(), Box<dyn std::error::Error>> {
5269        let f = self.func("silu_mul_f16out_f32");
5270        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5271        let ni = n as i32;
5272        let __s_b = self.gpu.stream();
5273        let mut b = __s_b.launch_builder(&f);
5274        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
5275        unsafe { b.launch(cfg)?; }
5276        Ok(())
5277    }
5278
5279    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
5280    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
5281    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
5282    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
5283    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
5284    /// launches per dense FFN layer (the gate+up post-matmul scales).
5285    pub fn silu_mul_scaled(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5286                           dst: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
5287        let f = self.func("silu_mul_scaled_f32");
5288        let cfg = LaunchConfig::for_num_elems(n as u32);
5289        let ni = n as i32;
5290        let (gsf, usf) = (gs, us);
5291        let __s_b = self.gpu.stream();
5292        let mut b = __s_b.launch_builder(&f);
5293        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
5294        unsafe { b.launch(cfg)?; }
5295        Ok(())
5296    }
5297
5298    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
5299    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
5300    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
5301    #[allow(clippy::too_many_arguments)]
5302    pub fn swigluoai_mul_scaled(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5303                                alpha: f32, limit: f32, dst: &mut CudaSlice<f32>, n: usize)
5304                                -> Result<(), Box<dyn std::error::Error>> {
5305        let f = self.func("swigluoai_mul_scaled_f32");
5306        let cfg = LaunchConfig::for_num_elems(n as u32);
5307        let ni = n as i32;
5308        let __s_b = self.gpu.stream();
5309        let mut b = __s_b.launch_builder(&f);
5310        b.arg(gate).arg(up).arg(&gs).arg(&us).arg(&alpha).arg(&limit).arg(dst).arg(&ni);
5311        unsafe { b.launch(cfg)?; }
5312        Ok(())
5313    }
5314
5315    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
5316    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
5317    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
5318    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
5319    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
5320    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
5321    /// n must be a multiple of 32 (n_ff always is).
5322    pub fn silu_mul_scaled_q8_1(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5323                                n: usize)
5324                                -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5325        let f = self.func("silu_mul_scaled_q8_1");
5326        let nblk = n / 32;
5327        let mut aq = self.alloc_uninit::<i8>(n)?;       // full-overwrite output
5328        let mut ad = self.alloc_uninit::<f32>(nblk)?;   // full-overwrite output
5329        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
5330        let cfg = LaunchConfig::for_num_elems(n as u32);
5331        let (gsf, usf, ni) = (gs, us, n as i32);
5332        let __s_b = self.gpu.stream();
5333        let mut b = __s_b.launch_builder(&f);
5334        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(&mut aq).arg(&mut ad).arg(&ni);
5335        unsafe { b.launch(cfg)?; }
5336        Ok((aq, ad))
5337    }
5338
5339    pub fn add(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5340               -> Result<(), Box<dyn std::error::Error>> {
5341        let f = self.func("add_f32");
5342        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
5343        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5344        let ni = n as i32;
5345        let __s_bld = self.gpu.stream();
5346        let mut bld = __s_bld.launch_builder(&f);
5347        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
5348        unsafe { bld.launch(cfg)?; }
5349        Ok(())
5350    }
5351
5352    pub fn mul(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5353               -> Result<(), Box<dyn std::error::Error>> {
5354        let f = self.func("mul_f32");
5355        let cfg = LaunchConfig::for_num_elems(n as u32);
5356        let ni = n as i32;
5357        let __s_bld = self.gpu.stream();
5358        let mut bld = __s_bld.launch_builder(&f);
5359        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
5360        unsafe { bld.launch(cfg)?; }
5361        Ok(())
5362    }
5363
5364    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
5365    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
5366    pub fn matmul(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize)
5367                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5368        use crate::model::GpuTensor;
5369        let in_f = w.in_features();
5370        let out_f = w.out_features();
5371        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
5372        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
5373        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
5374        // gives nothing). Quantize the activation once here then call the GEMM.
5375        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
5376        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
5377        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
5378        #[allow(non_snake_case)]
5379        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
5380        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
5381        let GEMM_M_THRESHOLD = if self.verify_exact_on() { usize::MAX } else { 16usize };
5382
5383        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
5384        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
5385        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
5386        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
5387        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
5388        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
5389        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
5390        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
5391        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
5392        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
5393        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
5394        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
5395        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
5396        const GEMM_MIN_OUT_F: usize = 128;   // 2*BM; below this the GEMM grid.x starves the 82 SMs
5397        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
5398        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
5399        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
5400        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
5401        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
5402        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
5403        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
5404        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
5405        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
5406        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
5407        if m >= GEMM_M_THRESHOLD {
5408            if let Some(y) = self.try_fp8_gemm(w, x, m)? { return Ok(y); }
5409            // PER-BLOCK FP8 MMQ (MEMRA_FP8_MMQ=1, default OFF; lane/fp8-mmq): the block-128 class
5410            // try_fp8_gemm skips (cuBLASLt takes no block grid on sm_120). Exact per block — the
5411            // checkpoint's e4m3 bytes and its f32 grid go into the tile unchanged.
5412            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? { return Ok(y); }
5413            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
5414            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
5415            if let Some(y) = self.try_f16_gemm(w, x, m)? { return Ok(y); }
5416        }
5417        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
5418            return self.qmatvec_mmq(w, x, m);
5419        }
5420        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
5421            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5422            return self.qmatvec_gemm(w, &aq, &ad, m);
5423        }
5424        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
5425        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
5426        if m >= GEMM_M_THRESHOLD {
5427            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? { return Ok(y); }
5428        }
5429        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
5430        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
5431        // to Stage-A f32-dequant (the correctness oracle path).
5432        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
5433        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
5434        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
5435        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
5436        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
5437        if m == 1 && fast {
5438            if let GpuTensor::Quant { bytes, qtype, row_bytes, rp, rp4, scale, .. } = w {
5439                if self.mmvq_supports(*qtype) {
5440                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
5441                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
5442                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
5443                    let (bytes, rp) = match rp4 { Some(m4) => (m4, true), None => (bytes, *rp) };
5444                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5445                    return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp);
5446                }
5447            }
5448        }
5449        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
5450        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
5451        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
5452        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
5453        // block below. MEMRA_NO_BATCHED -> per-m path.
5454        //
5455        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
5456        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
5457        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
5458        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
5459        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
5460        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
5461        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
5462        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
5463        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
5464        if (2..=16).contains(&m) && fast && std::env::var("MEMRA_NO_BATCHED").is_err()
5465            && (m <= 4 || Self::b8_enabled()) {
5466            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
5467            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
5468            // is present (rp4) — the mirror pick below then routes to the _rp family.
5469            let m_ok = m <= 8 || matches!(w, GpuTensor::Quant { qtype, rp4, .. }
5470                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || (*qtype == QT_Q8_0 && rp4.is_some()));
5471            if m_ok {
5472            if let GpuTensor::Quant { bytes, qtype, row_bytes, rp, rp4, .. } = w {
5473                if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
5474                    let (bytes, rp) = match rp4 { Some(m4) => (m4, true), None => (bytes, *rp) };
5475                    let mcols = Self::batched_mcols(m);
5476                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5477                    let mut y = self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp)?;
5478                    if let GpuTensor::Quant { scale, .. } = w {
5479                        if *scale != 1.0 { self.scale_inplace(&mut y, *scale, m * out_f)?; }
5480                    }
5481                    return Ok(y);
5482                }
5483            }
5484        }
5485        }
5486        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
5487        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
5488        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
5489        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
5490        // for this dtype, so the generic match below must never see it under `fast`.
5491        if fast {
5492            if let GpuTensor::Quant { bytes, qtype, row_bytes, scale, .. } = w {
5493                if *qtype == QT_F8_E4M3 {
5494                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5495                    return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes,
5496                                             *scale, false);
5497                }
5498            }
5499        }
5500        let mut y = match w {
5501            GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q8_0 =>
5502                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5503            GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q4_K =>
5504                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5505            GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q6_K =>
5506                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5507            GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q5_K =>
5508                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5509            GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q3_K =>
5510                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5511            GpuTensor::Quant { bytes, qtype, row_bytes, rp, .. } if fast && *qtype == QT_NVFP4 =>
5512                self.qmatvec_dp4a_named(
5513                    if *rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
5514                    bytes, x, m, in_f, out_f, *row_bytes)?,
5515            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
5516            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
5517            // anomaly (research/kat-anomaly-20260802/).
5518            GpuTensor::Quant { bytes, qtype, row_bytes, .. }
5519                if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() =>
5520                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5521            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
5522            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
5523            // without first writing the matching kernel, or func() will panic
5524            // "kernel ... not in any fatbin".
5525            GpuTensor::Quant { bytes, qtype, row_bytes, rp, .. } =>
5526                // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
5527                // deq(row,j) form cannot address the planes; same value/product order).
5528                self.qmatvec(bytes, x, m, in_f, out_f,
5529                             if *rp && *qtype == QT_NVFP4 { QT_NVFP4_RP } else { *qtype },
5530                             *row_bytes)?,
5531            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
5532            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
5533            // cuBLASLt f32 GEMV as the Float arm.
5534            GpuTensor::FloatBf16 { data, .. } =>
5535                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?,
5536        };
5537        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
5538        if let GpuTensor::Quant { scale, .. } = w {
5539            if *scale != 1.0 { self.scale_inplace(&mut y, *scale, m * out_f)?; }
5540        }
5541        Ok(y)
5542    }
5543
5544    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
5545    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
5546    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
5547        use crate::model::GpuTensor;
5548        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") { return false; }
5549        match w {
5550            GpuTensor::Quant { qtype, .. } => matches!(*qtype,
5551                QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q3_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0)
5552                || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled()),
5553            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
5554        }
5555    }
5556
5557    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
5558    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
5559    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
5560    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
5561    pub fn matmul_pre(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
5562                      x_fallback: &CudaSlice<f32>, m: usize)
5563                      -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5564        use crate::model::GpuTensor;
5565        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
5566        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
5567        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
5568        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
5569        // rc=30013 dig, 2026-07-31).
5570        let x_raw_ok = x_fallback.len() >= m * w.in_features();
5571        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
5572        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
5573        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
5574            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? { return Ok(y); }
5575            // PER-BLOCK FP8 MMQ (MEMRA_FP8_MMQ=1) — same arm as `matmul`; its own quantizer wants
5576            // the RAW f32 activation, so x_fallback not aq/ad.
5577            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? { return Ok(y); }
5578            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
5579            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? { return Ok(y); }
5580        }
5581        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
5582        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
5583        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
5584        // aq/ad.
5585        if m >= 16 && w.out_features() >= 128 && self.mmq_supports(w) && !self.verify_exact_on()
5586            && x_raw_ok {
5587            return self.qmatvec_mmq(w, x_fallback, m);
5588        }
5589        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
5590        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
5591        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
5592            if let Some(y) = self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())? {
5593                return Ok(y);
5594            }
5595        }
5596        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
5597        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
5598        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
5599            return self.qmatvec_gemm(w, aq, ad, m);
5600        }
5601        if !self.uses_q8_1_fast(w) { return self.matmul(w, x_fallback, m); }
5602        let in_f = w.in_features();
5603        let out_f = w.out_features();
5604        let (bytes, qtype, row_bytes, scale, rp) = match w {
5605            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5606            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
5607        };
5608        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
5609        // the dp4a/oracle tails below keep the raw GGUF bytes.
5610        let (mbytes, mrp) = match w {
5611            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5612            _ => (bytes, rp),
5613        };
5614        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
5615        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
5616        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
5617        if m == 1 && self.mmvq_supports(qtype) {
5618            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
5619        }
5620        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
5621        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
5622        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
5623        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
5624        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
5625        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
5626        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
5627        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
5628        // m=5..8 on the old per-m path (b8-tier-only seam).
5629        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
5630        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
5631        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
5632        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5633            && std::env::var("MEMRA_NO_BATCHED").is_err()
5634            && (m <= 4 || Self::b8_enabled())
5635            // b16 tier: Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's b16 exists ONLY as the
5636            // split-plane _rp twin (qmatvec_q8_0_mmvq_b16_rp — the q8rp mirror lane was built
5637            // for the m<=16 family, hybrid.rs), so Q8_0 joins iff the mirror is present (mrp).
5638            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && mrp)) {
5639            let mcols = Self::batched_mcols(m);
5640            return self.qmatvec_mmvq_batched(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp);
5641        }
5642        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
5643        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
5644        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
5645        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
5646        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
5647        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
5648            let (b2, r2) = if qtype == QT_Q4_0 { (mbytes, mrp) } else { (bytes, rp) };
5649            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
5650        }
5651        let name = match qtype {
5652            QT_Q8_0 => "qmatvec_q8_0_dp4a", QT_Q4_K => "qmatvec_q4_K_dp4a",
5653            QT_Q6_K => "qmatvec_q6_K_dp4a", QT_Q5_K => "qmatvec_q5_K_dp4a",
5654            QT_Q3_K => "qmatvec_q3_K_dp4a",
5655            QT_NVFP4 => if rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
5656            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
5657            _ => unreachable!(),
5658        };
5659        let f = self.func(name);
5660        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite GEMM output: skip memset
5661        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
5662        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5663        let __s_b = self.gpu.stream();
5664        let mut b = __s_b.launch_builder(&f);
5665        b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
5666        unsafe { b.launch(cfg)?; }
5667        if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
5668        Ok(y)
5669    }
5670
5671    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
5672    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
5673    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
5674    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
5675    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
5676    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
5677    /// reduce as m=1); this method just forces that path unconditionally.
5678    pub fn matmul_decode_exact(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize)
5679                               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5680        use crate::model::GpuTensor;
5681        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
5682        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
5683        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
5684        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
5685        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
5686        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
5687        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
5688        if let GpuTensor::Float { data, .. } = w {
5689            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
5690        }
5691        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
5692        // float linear (same n-independent reduction contract as the Float arm above).
5693        if let GpuTensor::FloatBf16 { data, .. } = w {
5694            let (in_f, out_f) = (w.in_features(), w.out_features());
5695            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
5696        }
5697        if !self.uses_q8_1_fast(w) { return self.matmul(w, x, m); }
5698        let in_f = w.in_features();
5699        let out_f = w.out_features();
5700        let (bytes, qtype, row_bytes, scale, rp) = match w {
5701            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5702            _ => return self.matmul(w, x, m),
5703        };
5704        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
5705        // which does its own mirror pick).
5706        let (bytes, rp) = match w {
5707            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5708            _ => (bytes, rp),
5709        };
5710        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5711        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
5712        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
5713        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
5714        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
5715        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
5716        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
5717        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
5718        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
5719        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5720            && std::env::var("MEMRA_NO_BATCHED").is_err()
5721            && (m <= 4 || Self::b8_enabled())
5722            // Q8_0 b16 exists only as the split-plane _rp twin (see matmul_pre's note).
5723            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && rp)) {
5724            let mcols = Self::batched_mcols(m);
5725            return self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp);
5726        }
5727        if self.mmvq_supports(qtype) {
5728            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
5729            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
5730            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
5731        }
5732        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
5733        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
5734        self.matmul_pre(w, &aq, &ad, x, m)
5735    }
5736
5737    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
5738    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
5739    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
5740    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
5741    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
5742    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
5743    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
5744    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
5745    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
5746    pub fn matmul_decode_exact_pre(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>,
5747                                   ad: &CudaSlice<f32>, m: usize)
5748                                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5749        use crate::model::GpuTensor;
5750        debug_assert!(self.uses_q8_1_fast(w),
5751                      "matmul_decode_exact_pre: caller must guarantee q8_1-fast");
5752        let in_f = w.in_features();
5753        let out_f = w.out_features();
5754        let (bytes, qtype, row_bytes, scale, rp) = match w {
5755            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
5756                (bytes, *qtype, *row_bytes, *scale, *rp),
5757            _ => return Err("matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into()),
5758        };
5759        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
5760        let (bytes, rp) = match w {
5761            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5762            _ => (bytes, rp),
5763        };
5764        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
5765        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5766            && std::env::var("MEMRA_NO_BATCHED").is_err()
5767            && (m <= 4 || Self::b8_enabled())
5768            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && rp)) {
5769            let mcols = Self::batched_mcols(m);
5770            return self.qmatvec_mmvq_batched(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp);
5771        }
5772        if self.mmvq_supports(qtype) {
5773            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
5774        }
5775        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
5776        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
5777        let x0 = self.zeros(0)?;
5778        self.matmul_pre(w, aq, ad, &x0, m)
5779    }
5780
5781    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
5782    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
5783    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
5784    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
5785    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
5786    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
5787    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
5788    /// per-tensor path.
5789    pub fn matmul_decode_exact_dual_pre(&self, w0: &crate::model::GpuTensor,
5790                                        w1: &crate::model::GpuTensor,
5791                                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
5792        -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>> {
5793        use crate::model::GpuTensor;
5794        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5795        let on = *ON.get_or_init(|| {
5796            std::env::var("MEMRA_SPEC_DUAL_T").map(|v| v != "0").unwrap_or(true)
5797        });
5798        if !on || !(2..=7).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok()
5799            || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
5800            return Ok(None);
5801        }
5802        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
5803        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
5804        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
5805        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
5806        if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5807        let (in_f, out_f) = (w0.in_features(), w0.out_features());
5808        if w1.in_features() != in_f || w1.out_features() != out_f {
5809            return Ok(None);
5810        }
5811        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
5812            (GpuTensor::Quant { bytes: b0, qtype: q0, row_bytes: rb0, scale: s0, rp: rp0, rp4: None, .. },
5813             GpuTensor::Quant { bytes: b1, qtype: q1, row_bytes: rb1, scale: s1, rp: rp1, rp4: None, .. })
5814                if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 =>
5815                (b0, b1, *rb0, *s0, *s1, *rp0),
5816            _ => return Ok(None),
5817        };
5818        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
5819        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
5820        if m > 4 && !(rp && Self::b8_enabled()
5821            && std::env::var("MEMRA_B567").as_deref() != Ok("0")) {
5822            return Ok(None);
5823        }
5824        let (y0, y1) = self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
5825        Ok(Some(((y0, s0), (y1, s1))))
5826    }
5827
5828    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
5829    /// launch computes both FFN projections of a verify batch — same activation, same shape,
5830    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
5831    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
5832    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
5833    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
5834    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
5835    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
5836    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
5837    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
5838    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
5839    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
5840    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
5841    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
5842    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
5843    pub fn matmul_decode_exact_dual(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
5844                                    x: &CudaSlice<f32>, m: usize)
5845        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5846        use crate::model::GpuTensor;
5847        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5848        let on = *ON.get_or_init(|| {
5849            std::env::var("MEMRA_SPEC_DUAL_T").map(|v| v != "0").unwrap_or(true)
5850        });
5851        if !on || !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok()
5852            || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
5853            return Ok(None);
5854        }
5855        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
5856        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
5857        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
5858        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
5859        if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5860        let (in_f, out_f) = (w0.in_features(), w0.out_features());
5861        if w1.in_features() != in_f || w1.out_features() != out_f {
5862            return Ok(None);
5863        }
5864        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
5865            (GpuTensor::Quant { bytes: b0, qtype: q0, row_bytes: rb0, scale: s0, rp: rp0, rp4: None, .. },
5866             GpuTensor::Quant { bytes: b1, qtype: q1, row_bytes: rb1, scale: s1, rp: rp1, rp4: None, .. })
5867                if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 =>
5868                (b0, b1, *rb0, *s0, *s1, *rp0),
5869            _ => return Ok(None),
5870        };
5871        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
5872        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
5873        if std::env::var("MEMRA_DEBUG").is_ok() {
5874            static ONCE: std::sync::Once = std::sync::Once::new();
5875            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
5876        }
5877        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5878        let (y0, y1) = self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
5879        let mut y0 = y0;
5880        let mut y1 = y1;
5881        if s0 != 1.0 { self.scale_inplace(&mut y0, s0, m * out_f)?; }
5882        if s1 != 1.0 { self.scale_inplace(&mut y1, s1, m * out_f)?; }
5883        Ok(Some((y0, y1)))
5884    }
5885
5886    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
5887    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
5888    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
5889    /// twins (both buffers must be the repacked layout).
5890    #[allow(clippy::too_many_arguments)]
5891    pub fn qmatvec_batched_dual_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
5892                                    aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
5893                                    m: usize, in_f: usize, out_f: usize, row_bytes: usize, rp: bool)
5894        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5895        const ROWS_PER_BLOCK: u32 = 4;
5896        let mcols = Self::batched_mcols(m);
5897        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
5898        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
5899        let (name, rows_per_block) = match (mcols, rp, m) {
5900            (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
5901            (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
5902            (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
5903            (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
5904            (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
5905            (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
5906            (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
5907            _ => return Err(format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into()),
5908        };
5909        let f = self.func(name);
5910        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5911        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5912        let cfg = LaunchConfig {
5913            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
5914            block_dim: (32, ROWS_PER_BLOCK, 1),
5915            shared_mem_bytes: 0,
5916        };
5917        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5918        let __s_b = self.gpu.stream();
5919        let mut b = __s_b.launch_builder(&f);
5920        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
5921            .arg(&inf).arg(&outf).arg(&mi).arg(&rb);
5922        unsafe { b.launch(cfg)?; }
5923        Ok((y0, y1))
5924    }
5925
5926    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
5927    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
5928    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
5929    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
5930    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
5931    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
5932    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
5933    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
5934    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
5935    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
5936    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
5937    pub fn matmul_pre_dual_noscale(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
5938                                   aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
5939        -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>> {
5940        use crate::model::GpuTensor;
5941        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) { return Ok(None); }
5942        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
5943        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
5944        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
5945        // would mix dispatch families across the pair — the exact class `q8_fused_params`
5946        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
5947        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
5948        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
5949        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
5950        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
5951        if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5952        let (in_f, out_f) = (w0.in_features(), w0.out_features());
5953        if w1.in_features() != in_f || w1.out_features() != out_f { return Ok(None); }
5954        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
5955        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
5956        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
5957        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
5958        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
5959        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
5960        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
5961        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
5962        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
5963        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
5964        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
5965        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
5966        let no_mirror = |w: &crate::model::GpuTensor| {
5967            !matches!(w, GpuTensor::Quant { rp4: Some(_), .. })
5968        };
5969        if self.q8_ffn_fuse2_on()
5970            && no_mirror(w0) && no_mirror(w1)
5971            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
5972        {
5973            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
5974            return Ok(Some(((y0, 1.0), (y1, 1.0))));
5975        }
5976        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
5977        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
5978        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
5979        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
5980        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
5981        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
5982        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
5983        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
5984        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
5985        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
5986            let (y0, y1) = self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2,
5987                                                 1.0, 1.0)?;
5988            return Ok(Some(((y0, p0.3), (y1, p1.3))));
5989        }
5990        let (b0, q0, rb0, s0, rp0) = match w0 {
5991            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5992            _ => return Ok(None),
5993        };
5994        let (b1, q1, rb1, s1, rp1) = match w1 {
5995            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5996            _ => return Ok(None),
5997        };
5998        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 { return Ok(None); }
5999        const ROWS_PER_BLOCK: u32 = 4;   // matches MEMRA_MMVQ_ROWS in qmatvec.cu
6000        const RPW: u32 = 2;
6001        let rows_per_block = ROWS_PER_BLOCK * RPW;
6002        let f = self.func(if rp0 { "qmatvec_nvfp4_mmvq_dual_mr2_rp" } else { "qmatvec_nvfp4_mmvq_dual_mr2" });
6003        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
6004        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
6005        let cfg = LaunchConfig {
6006            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
6007            block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0,
6008        };
6009        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
6010        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
6011        // yscale args stay 1.0 here (they exist for the single-tensor callers).
6012        let one = 1.0f32;
6013        let __s_b = self.gpu.stream();
6014        let mut b = __s_b.launch_builder(&f);
6015        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6016         .arg(&inf).arg(&outf).arg(&mi).arg(&rb).arg(&one).arg(&one);
6017        unsafe { b.launch(cfg)?; }
6018        Ok(Some(((y0, s0), (y1, s1))))
6019    }
6020
6021    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
6022    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
6023    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
6024    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
6025    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
6026    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
6027    /// back to the per-tensor path.
6028    pub fn matmul_q8_fused2(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6029                            aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6030        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6031        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
6032        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
6033        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
6034        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
6035        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
6036        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6037            return Ok(Some(self.e4m3_fused2_core(p0.0, p1.0, aq, ad, w0.in_features(),
6038                                                 p0.1, p1.1, p0.2, p0.3, p1.3)?));
6039        }
6040        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6041        Ok(Some(self.q8_fused2_core(p0.0, p1.0, aq, ad, w0.in_features(), p0.1, p1.1, p0.2)?))
6042    }
6043
6044    #[allow(clippy::too_many_arguments)]
6045    fn q8_fused2_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6046                      aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6047                      in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6048        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6049        const ROWS_PER_BLOCK: u32 = 4;   // matches MEMRA_MMVQ_ROWS in qmatvec.cu
6050        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6051        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6052        let f = self.func("qmatvec_q8_0_mmvq_fused2");
6053        let mut y0 = self.alloc_uninit::<f32>(out0)?;
6054        let mut y1 = self.alloc_uninit::<f32>(out1)?;
6055        let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6056                                 shared_mem_bytes: 0 };
6057        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
6058        let __s_b = self.gpu.stream();
6059        let mut b = __s_b.launch_builder(&f);
6060        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6061         .arg(&inf).arg(&o0).arg(&o1).arg(&rbl);
6062        unsafe { b.launch(cfg)?; }
6063        Ok((y0, y1))
6064    }
6065
6066    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
6067    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
6068    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
6069    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
6070    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
6071    pub fn matmul_q8_fused2_x(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6072                              x: &CudaSlice<f32>)
6073        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6074        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) { return Ok(None); }
6075        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6076            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
6077            return Ok(Some(self.e4m3_fused2_core(p0.0, p1.0, &aq, &ad, w0.in_features(),
6078                                                 p0.1, p1.1, p0.2, p0.3, p1.3)?));
6079        }
6080        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6081        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
6082        Ok(Some(self.q8_fused2_core(p0.0, p1.0, &aq, &ad, w0.in_features(), p0.1, p1.1, p0.2)?))
6083    }
6084
6085    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
6086    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
6087    #[allow(clippy::too_many_arguments)]
6088    pub fn qmatvec_q8_fused2_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>,
6089                                 in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6090        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6091        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6092        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
6093    }
6094
6095    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
6096    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
6097    /// (tensor,row) to three separate m=1 MMVQ launches.
6098    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
6099    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
6100    pub fn matmul_q4_fused3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6101                            w2: &crate::model::GpuTensor,
6102                            aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6103        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6104        use crate::model::GpuTensor;
6105        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6106            match w {
6107                GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6108                    Some((*row_bytes, w.out_features())),
6109                _ => None,
6110            }
6111        };
6112        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2))
6113        else { return Ok(None) };
6114        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6115            return Ok(None);
6116        }
6117        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
6118        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
6119        // the separate matvecs (each routes its own rp).
6120        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6121            match w {
6122                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6123                    Some(m) => (m, true),
6124                    None => (bytes, *rp),
6125                },
6126                _ => unreachable!(),
6127            }
6128        }
6129        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6130        if rp0 != rp1 || rp1 != rp2 { return Ok(None); }
6131        let rp = rp0;
6132        let rpb: u32 = 4;
6133        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
6134        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
6135        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
6136        let mr1 = rp && Self::q40_mr1_on();
6137        let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6138                            else { (o as u32).div_ceil(2).div_ceil(rpb) };
6139        let grid = nb(o0) + nb(o1) + nb(o2);
6140        let mut y0 = self.alloc_uninit::<f32>(o0)?;
6141        let mut y1 = self.alloc_uninit::<f32>(o1)?;
6142        let mut y2 = self.alloc_uninit::<f32>(o2)?;
6143        let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused3_mr1_rp" }
6144                          else if rp { "qmatvec_q4_0_mmvq_fused3_rp" }
6145                          else { "qmatvec_q4_0_mmvq_fused3" });
6146        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6147        let inf = w0.in_features() as i32;
6148        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
6149        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
6150        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
6151        // variant may take the programmatic-serialization launch.
6152        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6153            {
6154            use cudarc::driver::{DevicePtr, DevicePtrMut};
6155            let s = &self.gpu.stream();
6156            let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6157            let (p2, _g2) = b2.device_ptr(s); let (paq, _g3) = aq.device_ptr(s);
6158            let (pad, _g4) = ad.device_ptr(s);
6159            let (py0, _g5) = y0.device_ptr_mut(s); let (py1, _g6) = y1.device_ptr_mut(s);
6160            let (py2, _g7) = y2.device_ptr_mut(s);
6161            let mut ps = [
6162                &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6163                &p2 as *const _ as *mut _, &paq as *const _ as *mut _,
6164                &pad as *const _ as *mut _, &py0 as *const _ as *mut _,
6165                &py1 as *const _ as *mut _, &py2 as *const _ as *mut _,
6166                &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6167                &oo1 as *const _ as *mut _, &oo2 as *const _ as *mut _,
6168                &r0 as *const _ as *mut _, &r1 as *const _ as *mut _,
6169                &r2 as *const _ as *mut _,
6170            ];
6171            unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused3_mr1_rp",
6172                                     (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6173            }
6174            return Ok(Some((y0, y1, y2)));
6175        }
6176        let __s_b = self.gpu.stream();
6177        let mut b = __s_b.launch_builder(&f);
6178        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6179         .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&r0).arg(&r1).arg(&r2);
6180        unsafe { b.launch(cfg)?; }
6181        Ok(Some((y0, y1, y2)))
6182    }
6183
6184    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
6185    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
6186    #[allow(clippy::too_many_arguments)]
6187    pub fn matmul_q4_fused3_into(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6188                                 w2: &crate::model::GpuTensor,
6189                                 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6190                                 y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>,
6191                                 y2: &mut CudaSlice<f32>)
6192        -> Result<bool, Box<dyn std::error::Error>> {
6193        use crate::model::GpuTensor;
6194        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6195            match w {
6196                GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6197                    Some((*row_bytes, w.out_features())),
6198                _ => None,
6199            }
6200        };
6201        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2))
6202        else { return Ok(false) };
6203        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6204            return Ok(false);
6205        }
6206        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6207            match w {
6208                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6209                    Some(m) => (m, true),
6210                    None => (bytes, *rp),
6211                },
6212                _ => unreachable!(),
6213            }
6214        }
6215        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6216        if rp0 != rp1 || rp1 != rp2 { return Ok(false); }
6217        let rp = rp0;
6218        let rpb: u32 = 4;
6219        let mr1 = rp && Self::q40_mr1_on();
6220        let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6221                            else { (o as u32).div_ceil(2).div_ceil(rpb) };
6222        let grid = nb(o0) + nb(o1) + nb(o2);
6223        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
6224        let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused3_mr1_rp" }
6225                          else if rp { "qmatvec_q4_0_mmvq_fused3_rp" }
6226                          else { "qmatvec_q4_0_mmvq_fused3" });
6227        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6228        let inf = w0.in_features() as i32;
6229        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
6230        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
6231        // PDL wave-A: identical to the owned twin (capture-lane parity).
6232        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6233            use cudarc::driver::{DevicePtr, DevicePtrMut};
6234            let s = &self.gpu.stream();
6235            let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6236            let (p2, _g2) = b2.device_ptr(s); let (paq, _g3) = aq.device_ptr(s);
6237            let (pad, _g4) = ad.device_ptr(s);
6238            let (py0, _g5) = y0.device_ptr_mut(s); let (py1, _g6) = y1.device_ptr_mut(s);
6239            let (py2, _g7) = y2.device_ptr_mut(s);
6240            let mut ps = [
6241                &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6242                &p2 as *const _ as *mut _, &paq as *const _ as *mut _,
6243                &pad as *const _ as *mut _, &py0 as *const _ as *mut _,
6244                &py1 as *const _ as *mut _, &py2 as *const _ as *mut _,
6245                &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6246                &oo1 as *const _ as *mut _, &oo2 as *const _ as *mut _,
6247                &r0 as *const _ as *mut _, &r1 as *const _ as *mut _,
6248                &r2 as *const _ as *mut _,
6249            ];
6250            unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused3_mr1_rp",
6251                                     (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6252            return Ok(true);
6253        }
6254        let __s_b = self.gpu.stream();
6255        let mut b = __s_b.launch_builder(&f);
6256        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut *y0).arg(&mut *y1).arg(&mut *y2)
6257         .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&r0).arg(&r1).arg(&r2);
6258        unsafe { b.launch(cfg)?; }
6259        Ok(true)
6260    }
6261
6262    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
6263    pub fn matmul_q4_fused2(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6264                            aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6265        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6266        use crate::model::GpuTensor;
6267        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6268            match w {
6269                GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6270                    Some((*row_bytes, w.out_features())),
6271                _ => None,
6272            }
6273        };
6274        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(None) };
6275        if w0.in_features() != w1.in_features() { return Ok(None); }
6276        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
6277        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6278            match w {
6279                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6280                    Some(m) => (m, true),
6281                    None => (bytes, *rp),
6282                },
6283                _ => unreachable!(),
6284            }
6285        }
6286        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6287        if rp0 != rp1 { return Ok(None); }
6288        let rp = rp0;
6289        let rpb: u32 = 4;
6290        // mr1 twin — see matmul_q4_fused3.
6291        let mr1 = rp && Self::q40_mr1_on();
6292        let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6293                            else { (o as u32).div_ceil(2).div_ceil(rpb) };
6294        let grid = nb(o0) + nb(o1);
6295        let mut y0 = self.alloc_uninit::<f32>(o0)?;
6296        let mut y1 = self.alloc_uninit::<f32>(o1)?;
6297        let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused2_mr1_rp" }
6298                          else if rp { "qmatvec_q4_0_mmvq_fused2_rp" }
6299                          else { "qmatvec_q4_0_mmvq_fused2" });
6300        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6301        let inf = w0.in_features() as i32;
6302        let (oo0, oo1) = (o0 as i32, o1 as i32);
6303        let (r0, r1) = (rb0 as i64, rb1 as i64);
6304        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
6305        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6306            {
6307            use cudarc::driver::{DevicePtr, DevicePtrMut};
6308            let s = &self.gpu.stream();
6309            let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6310            let (paq, _g2) = aq.device_ptr(s); let (pad, _g3) = ad.device_ptr(s);
6311            let (py0, _g4) = y0.device_ptr_mut(s); let (py1, _g5) = y1.device_ptr_mut(s);
6312            let mut ps = [
6313                &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6314                &paq as *const _ as *mut _, &pad as *const _ as *mut _,
6315                &py0 as *const _ as *mut _, &py1 as *const _ as *mut _,
6316                &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6317                &oo1 as *const _ as *mut _, &r0 as *const _ as *mut _,
6318                &r1 as *const _ as *mut _,
6319            ];
6320            unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused2_mr1_rp",
6321                                     (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6322            }
6323            return Ok(Some((y0, y1)));
6324        }
6325        let __s_b = self.gpu.stream();
6326        let mut b = __s_b.launch_builder(&f);
6327        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6328         .arg(&inf).arg(&oo0).arg(&oo1).arg(&r0).arg(&r1);
6329        unsafe { b.launch(cfg)?; }
6330        Ok(Some((y0, y1)))
6331    }
6332
6333    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
6334    pub fn matmul_q4_fused2_into(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6335                                 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6336                                 y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>)
6337        -> Result<bool, Box<dyn std::error::Error>> {
6338        use crate::model::GpuTensor;
6339        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6340            match w {
6341                GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6342                    Some((*row_bytes, w.out_features())),
6343                _ => None,
6344            }
6345        };
6346        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(false) };
6347        if w0.in_features() != w1.in_features() { return Ok(false); }
6348        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6349            match w {
6350                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6351                    Some(m) => (m, true),
6352                    None => (bytes, *rp),
6353                },
6354                _ => unreachable!(),
6355            }
6356        }
6357        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6358        if rp0 != rp1 { return Ok(false); }
6359        let rp = rp0;
6360        let rpb: u32 = 4;
6361        let mr1 = rp && Self::q40_mr1_on();
6362        let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6363                            else { (o as u32).div_ceil(2).div_ceil(rpb) };
6364        let grid = nb(o0) + nb(o1);
6365        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
6366        let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused2_mr1_rp" }
6367                          else if rp { "qmatvec_q4_0_mmvq_fused2_rp" }
6368                          else { "qmatvec_q4_0_mmvq_fused2" });
6369        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6370        let inf = w0.in_features() as i32;
6371        let (oo0, oo1) = (o0 as i32, o1 as i32);
6372        let (r0, r1) = (rb0 as i64, rb1 as i64);
6373        // PDL wave-A: identical to the owned twin (capture-lane parity).
6374        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6375            use cudarc::driver::{DevicePtr, DevicePtrMut};
6376            let s = &self.gpu.stream();
6377            let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6378            let (paq, _g2) = aq.device_ptr(s); let (pad, _g3) = ad.device_ptr(s);
6379            let (py0, _g4) = y0.device_ptr_mut(s); let (py1, _g5) = y1.device_ptr_mut(s);
6380            let mut ps = [
6381                &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6382                &paq as *const _ as *mut _, &pad as *const _ as *mut _,
6383                &py0 as *const _ as *mut _, &py1 as *const _ as *mut _,
6384                &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6385                &oo1 as *const _ as *mut _, &r0 as *const _ as *mut _,
6386                &r1 as *const _ as *mut _,
6387            ];
6388            unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused2_mr1_rp",
6389                                     (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6390            return Ok(true);
6391        }
6392        let __s_b = self.gpu.stream();
6393        let mut b = __s_b.launch_builder(&f);
6394        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut *y0).arg(&mut *y1)
6395         .arg(&inf).arg(&oo0).arg(&oo1).arg(&r0).arg(&r1);
6396        unsafe { b.launch(cfg)?; }
6397        Ok(true)
6398    }
6399
6400    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
6401    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
6402    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
6403    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
6404    pub fn matmul_q4_fused2_batched(&self, w0: &crate::model::GpuTensor,
6405                                    w1: &crate::model::GpuTensor,
6406                                    aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6407        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6408        use crate::model::GpuTensor;
6409        if m < 2 || m > 8 { return Ok(None); }
6410        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6411            match w {
6412                GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6413                    Some((*row_bytes, w.out_features())),
6414                _ => None,
6415            }
6416        };
6417        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(None) };
6418        if w0.in_features() != w1.in_features() { return Ok(None); }
6419        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6420            match w {
6421                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6422                    Some(mr) => (mr, true),
6423                    None => (bytes, *rp),
6424                },
6425                _ => unreachable!(),
6426            }
6427        }
6428        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6429        if !rp0 || !rp1 { return Ok(None); }
6430        let mcols = Self::batched_mcols(m);
6431        let rpb: u32 = 4;
6432        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
6433        let grid = nb(o0) + nb(o1);
6434        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
6435        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
6436        let f = self.func(match mcols { 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
6437                                        4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
6438                                        _ => "qmatvec_q4_0_mmvq_b8_f2_rp" });
6439        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1),
6440                                 shared_mem_bytes: 0 };
6441        let inf = w0.in_features() as i32;
6442        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
6443        let rb = rb0 as i64;
6444        let __s_b = self.gpu.stream();
6445        let mut b = __s_b.launch_builder(&f);
6446        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6447         .arg(&inf).arg(&oo0).arg(&oo1).arg(&mi).arg(&rb);
6448        unsafe { b.launch(cfg)?; }
6449        Ok(Some((y0, y1)))
6450    }
6451
6452    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
6453    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
6454    #[allow(clippy::too_many_arguments)]
6455    pub fn matmul_q4_fused3_batched(&self, w0: &crate::model::GpuTensor,
6456                                    w1: &crate::model::GpuTensor, w2: &crate::model::GpuTensor,
6457                                    aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6458        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6459        use crate::model::GpuTensor;
6460        if m < 2 || m > 8 { return Ok(None); }
6461        let q4 = |w: &GpuTensor| -> Option<usize> {
6462            match w {
6463                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
6464                _ => None,
6465            }
6466        };
6467        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else { return Ok(None) };
6468        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6469            return Ok(None);
6470        }
6471        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6472            match w {
6473                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6474                    Some(mr) => (mr, true),
6475                    None => (bytes, *rp),
6476                },
6477                _ => unreachable!(),
6478            }
6479        }
6480        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6481        if !rp0 || !rp1 || !rp2 { return Ok(None); }
6482        let mcols = Self::batched_mcols(m);
6483        let rpb: u32 = 4;
6484        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
6485        let grid = nb(o0) + nb(o1) + nb(o2);
6486        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
6487        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
6488        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
6489        let f = self.func(match mcols { 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
6490                                        4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
6491                                        _ => "qmatvec_q4_0_mmvq_b8_f3_rp" });
6492        let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1),
6493                                 shared_mem_bytes: 0 };
6494        let inf = w0.in_features() as i32;
6495        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
6496        let rb = 0i64;
6497        let __s_b = self.gpu.stream();
6498        let mut b = __s_b.launch_builder(&f);
6499        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6500         .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&mi).arg(&rb);
6501        unsafe { b.launch(cfg)?; }
6502        Ok(Some((y0, y1, y2)))
6503    }
6504
6505    pub fn matmul_q8_fused3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6506                            w2: &crate::model::GpuTensor,
6507                            aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6508        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6509        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
6510        // are per-tensor FP8, so native residency without this arm meant three separate launches.
6511        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
6512            return Ok(Some(self.e4m3_fused3_core(p0.0, p1.0, p2.0, aq, ad, w0.in_features(),
6513                                                 p0.1, p1.1, p2.1, p0.2,
6514                                                 p0.3, p1.3, p2.3)?));
6515        }
6516        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else { return Ok(None) };
6517        Ok(Some(self.q8_fused3_core(p0.0, p1.0, p2.0, aq, ad, w0.in_features(),
6518                                    p0.1, p1.1, p2.1, p0.2)?))
6519    }
6520
6521    #[allow(clippy::too_many_arguments)]
6522    fn q8_fused3_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6523                      aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6524                      in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize)
6525        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6526        const ROWS_PER_BLOCK: u32 = 4;
6527        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6528        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6529        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6530        let f = self.func("qmatvec_q8_0_mmvq_fused3");
6531        let mut y0 = self.alloc_uninit::<f32>(out0)?;
6532        let mut y1 = self.alloc_uninit::<f32>(out1)?;
6533        let mut y2 = self.alloc_uninit::<f32>(out2)?;
6534        let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6535                                 shared_mem_bytes: 0 };
6536        let (inf, o0, o1, o2, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32, row_bytes as i64);
6537        let __s_b = self.gpu.stream();
6538        let mut b = __s_b.launch_builder(&f);
6539        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6540         .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&rbl);
6541        unsafe { b.launch(cfg)?; }
6542        Ok((y0, y1, y2))
6543    }
6544
6545    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
6546    #[allow(clippy::too_many_arguments)]
6547    pub fn qmatvec_q8_fused3_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6548                                 x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize,
6549                                 out2: usize, row_bytes: usize)
6550        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6551        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6552        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
6553    }
6554
6555    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
6556    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
6557    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
6558    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
6559    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
6560    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
6561    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
6562    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
6563    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
6564    /// twin must not introduce a batched program the reference path would not run).
6565    pub fn matmul_q8_fused2_t(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6566                              aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6567        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6568        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
6569        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
6570        // fuses too — same template body, still bit-identical to the two _b8 launches.
6571        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() { return Ok(None); }
6572        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
6573        // so the fused b8 launch would introduce a batched program the reference path would not run.
6574        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6575            if m > 4 && !Self::b8_enabled() { return Ok(None); }
6576            return Ok(Some(self.e4m3_fused2_t_core(p0.0, p1.0, aq, ad, m, w0.in_features(),
6577                                                   p0.1, p1.1, p0.2, p0.3, p1.3)?));
6578        }
6579        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6580        Ok(Some(self.q8_fused2_t_core(p0.0, p1.0, aq, ad, m, w0.in_features(), p0.1, p1.1, p0.2)?))
6581    }
6582
6583    #[allow(clippy::too_many_arguments)]
6584    fn q8_fused2_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6585                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6586                        in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6587        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6588        const ROWS_PER_BLOCK: u32 = 4;   // matches MEMRA_MMVQ_ROWS in qmatvec.cu
6589        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6590        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6591        let f = self.func(match Self::batched_mcols(m) {
6592            2 => "qmatvec_q8_0_mmvq_fused2_b2",
6593            4 => "qmatvec_q8_0_mmvq_fused2_b4",
6594            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
6595            _ => "qmatvec_q8_0_mmvq_fused2_b8",
6596        });
6597        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6598        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6599        let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6600                                 shared_mem_bytes: 0 };
6601        let (inf, o0, o1, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, m as i32, row_bytes as i64);
6602        let __s_b = self.gpu.stream();
6603        let mut b = __s_b.launch_builder(&f);
6604        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6605         .arg(&inf).arg(&o0).arg(&o1).arg(&mi).arg(&rbl);
6606        unsafe { b.launch(cfg)?; }
6607        Ok((y0, y1))
6608    }
6609
6610    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
6611    /// q8_1 quant of the [m, in_f] activation), no env gating.
6612    #[allow(clippy::too_many_arguments)]
6613    pub fn qmatvec_q8_fused2_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6614                                   x: &CudaSlice<f32>, m: usize,
6615                                   in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6616        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6617        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6618        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
6619    }
6620
6621    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
6622    /// `matmul_q8_fused2_t` with three ranges.
6623    #[allow(clippy::too_many_arguments)]
6624    pub fn matmul_q8_fused3_t(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6625                              w2: &crate::model::GpuTensor,
6626                              aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6627        -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6628        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() { return Ok(None); }
6629        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
6630            return Ok(Some(self.e4m3_fused3_t_core(p0.0, p1.0, p2.0, aq, ad, m, w0.in_features(),
6631                                                   p0.1, p1.1, p2.1, p0.2,
6632                                                   p0.3, p1.3, p2.3)?));
6633        }
6634        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else { return Ok(None) };
6635        Ok(Some(self.q8_fused3_t_core(p0.0, p1.0, p2.0, aq, ad, m, w0.in_features(),
6636                                      p0.1, p1.1, p2.1, p0.2)?))
6637    }
6638
6639    #[allow(clippy::too_many_arguments)]
6640    fn q8_fused3_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6641                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6642                        in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize)
6643        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6644        const ROWS_PER_BLOCK: u32 = 4;
6645        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6646        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6647        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6648        let f = self.func(if Self::batched_mcols(m) == 2 { "qmatvec_q8_0_mmvq_fused3_b2" }
6649                          else { "qmatvec_q8_0_mmvq_fused3_b4" });
6650        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6651        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6652        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
6653        let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6654                                 shared_mem_bytes: 0 };
6655        let (inf, o0, o1, o2, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6656                                          m as i32, row_bytes as i64);
6657        let __s_b = self.gpu.stream();
6658        let mut b = __s_b.launch_builder(&f);
6659        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6660         .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&mi).arg(&rbl);
6661        unsafe { b.launch(cfg)?; }
6662        Ok((y0, y1, y2))
6663    }
6664
6665    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
6666    #[allow(clippy::too_many_arguments)]
6667    pub fn qmatvec_q8_fused3_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6668                                   x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize,
6669                                   out1: usize, out2: usize, row_bytes: usize)
6670        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6671        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6672        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
6673    }
6674
6675    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
6676    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
6677    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
6678    pub fn q8_ffn_fuse2_on(&self) -> bool {
6679        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6680        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
6681    }
6682
6683    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
6684    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
6685    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
6686    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
6687    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
6688    #[allow(clippy::type_complexity)]
6689    fn q8_fused_params<'w, const N: usize>(&self, ws: &[&'w crate::model::GpuTensor; N])
6690        -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
6691        use crate::model::GpuTensor;
6692        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") { return None; }
6693        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") { return None; }
6694        let in_f = ws[0].in_features();
6695        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
6696        for (i, w) in ws.iter().enumerate() {
6697            match w {
6698                GpuTensor::Quant { bytes, qtype, row_bytes, scale, .. }
6699                    if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f =>
6700                        out[i] = Some((bytes, w.out_features(), *row_bytes)),
6701                _ => return None,
6702            }
6703        }
6704        Some(out.map(|o| o.unwrap()))
6705    }
6706
6707    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
6708    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
6709    pub fn e4m3_dual_on(&self) -> bool {
6710        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6711        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
6712    }
6713
6714    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
6715    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
6716    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
6717    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
6718    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
6719    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
6720    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
6721    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
6722    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
6723    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
6724    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
6725    #[allow(clippy::type_complexity)]
6726    fn e4m3_fused_params<'w, const N: usize>(&self, ws: &[&'w crate::model::GpuTensor; N])
6727        -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
6728        use crate::model::GpuTensor;
6729        if !self.e4m3_dual_on() { return None; }
6730        let in_f = ws[0].in_features();
6731        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
6732        for (i, w) in ws.iter().enumerate() {
6733            match w {
6734                GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, rp4, .. }
6735                    if *qtype == QT_F8_E4M3 && w.in_features() == in_f
6736                        && *row_bytes == in_f && !*rp && rp4.is_none() =>
6737                        out[i] = Some((bytes, w.out_features(), *row_bytes, *scale)),
6738                _ => return None,
6739            }
6740        }
6741        Some(out.map(|o| o.unwrap()))
6742    }
6743
6744    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
6745    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
6746    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
6747    #[allow(clippy::too_many_arguments)]
6748    fn e4m3_fused2_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6749                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6750                        in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6751                        ws0: f32, ws1: f32)
6752        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6753        const ROWS_PER_BLOCK: u32 = 4;   // matches MEMRA_MMVQ_ROWS in qmatvec.cu
6754        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6755        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6756        let f = self.func("qmatvec_e4m3_mmvq_fused2");
6757        let mut y0 = self.alloc_uninit::<f32>(out0)?;
6758        let mut y1 = self.alloc_uninit::<f32>(out1)?;
6759        let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6760                                 shared_mem_bytes: 0 };
6761        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
6762        let __s_b = self.gpu.stream();
6763        let mut b = __s_b.launch_builder(&f);
6764        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6765         .arg(&inf).arg(&o0).arg(&o1).arg(&rbl).arg(&ws0).arg(&ws1);
6766        unsafe { b.launch(cfg)?; }
6767        Ok((y0, y1))
6768    }
6769
6770    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
6771    #[allow(clippy::too_many_arguments)]
6772    fn e4m3_fused3_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6773                        aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6774                        in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize,
6775                        ws0: f32, ws1: f32, ws2: f32)
6776        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6777        const ROWS_PER_BLOCK: u32 = 4;
6778        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6779        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6780        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6781        let f = self.func("qmatvec_e4m3_mmvq_fused3");
6782        let mut y0 = self.alloc_uninit::<f32>(out0)?;
6783        let mut y1 = self.alloc_uninit::<f32>(out1)?;
6784        let mut y2 = self.alloc_uninit::<f32>(out2)?;
6785        let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6786                                 shared_mem_bytes: 0 };
6787        let (inf, o0, o1, o2, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6788                                      row_bytes as i64);
6789        let __s_b = self.gpu.stream();
6790        let mut b = __s_b.launch_builder(&f);
6791        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6792         .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&rbl).arg(&ws0).arg(&ws1).arg(&ws2);
6793        unsafe { b.launch(cfg)?; }
6794        Ok((y0, y1, y2))
6795    }
6796
6797    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
6798    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
6799    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
6800    #[allow(clippy::too_many_arguments)]
6801    fn e4m3_fused2_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6802                          aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6803                          in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6804                          ws0: f32, ws1: f32)
6805        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6806        const ROWS_PER_BLOCK: u32 = 4;
6807        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6808        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6809        let f = self.func(match Self::batched_mcols(m) {
6810            2 => "qmatvec_e4m3_mmvq_fused2_b2",
6811            4 => "qmatvec_e4m3_mmvq_fused2_b4",
6812            _ => "qmatvec_e4m3_mmvq_fused2_b8",
6813        });
6814        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6815        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6816        let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6817                                 shared_mem_bytes: 0 };
6818        let (inf, o0, o1, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, m as i32,
6819                                      row_bytes as i64);
6820        let __s_b = self.gpu.stream();
6821        let mut b = __s_b.launch_builder(&f);
6822        b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6823         .arg(&inf).arg(&o0).arg(&o1).arg(&mi).arg(&rbl);
6824        unsafe { b.launch(cfg)?; }
6825        if ws0 != 1.0 { self.scale_inplace(&mut y0, ws0, m * out0)?; }
6826        if ws1 != 1.0 { self.scale_inplace(&mut y1, ws1, m * out1)?; }
6827        Ok((y0, y1))
6828    }
6829
6830    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
6831    #[allow(clippy::too_many_arguments)]
6832    fn e4m3_fused3_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6833                          aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6834                          in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize,
6835                          ws0: f32, ws1: f32, ws2: f32)
6836        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6837        const ROWS_PER_BLOCK: u32 = 4;
6838        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6839        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6840        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6841        let f = self.func(if Self::batched_mcols(m) == 2 { "qmatvec_e4m3_mmvq_fused3_b2" }
6842                          else { "qmatvec_e4m3_mmvq_fused3_b4" });
6843        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6844        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6845        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
6846        let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6847                                 shared_mem_bytes: 0 };
6848        let (inf, o0, o1, o2, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6849                                          m as i32, row_bytes as i64);
6850        let __s_b = self.gpu.stream();
6851        let mut b = __s_b.launch_builder(&f);
6852        b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6853         .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&mi).arg(&rbl);
6854        unsafe { b.launch(cfg)?; }
6855        if ws0 != 1.0 { self.scale_inplace(&mut y0, ws0, m * out0)?; }
6856        if ws1 != 1.0 { self.scale_inplace(&mut y1, ws1, m * out1)?; }
6857        if ws2 != 1.0 { self.scale_inplace(&mut y2, ws2, m * out2)?; }
6858        Ok((y0, y1, y2))
6859    }
6860
6861    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
6862    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
6863    #[allow(clippy::too_many_arguments)]
6864    pub fn qmatvec_e4m3_fused2_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>,
6865                                   in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6866                                   ws0: f32, ws1: f32)
6867        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6868        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6869        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
6870    }
6871
6872    #[allow(clippy::too_many_arguments)]
6873    pub fn qmatvec_e4m3_fused3_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6874                                   x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize,
6875                                   out2: usize, row_bytes: usize, ws0: f32, ws1: f32, ws2: f32)
6876        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6877        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6878        self.e4m3_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes,
6879                              ws0, ws1, ws2)
6880    }
6881
6882    #[allow(clippy::too_many_arguments)]
6883    pub fn qmatvec_e4m3_fused2_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6884                                     x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize,
6885                                     out1: usize, row_bytes: usize, ws0: f32, ws1: f32)
6886        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6887        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6888        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
6889    }
6890
6891    #[allow(clippy::too_many_arguments)]
6892    pub fn qmatvec_e4m3_fused3_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6893                                     b2: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
6894                                     in_f: usize, out0: usize, out1: usize, out2: usize,
6895                                     row_bytes: usize, ws0: f32, ws1: f32, ws2: f32)
6896        -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6897        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6898        self.e4m3_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes,
6899                                ws0, ws1, ws2)
6900    }
6901
6902    pub fn matmul_pre_noscale(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6903                              m: usize) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
6904        use crate::model::GpuTensor;
6905        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
6906        if m != 1 || !self.uses_q8_1_fast(w) { return Ok(None); }
6907        let in_f = w.in_features();
6908        let out_f = w.out_features();
6909        let (bytes, qtype, row_bytes, scale, rp) = match w {
6910            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
6911            _ => return Ok(None),
6912        };
6913        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
6914        if self.mmvq_supports(qtype) {
6915            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
6916            let (mbytes, mrp) = match w {
6917                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
6918                _ => (bytes, rp),
6919            };
6920            let y = self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp)?;
6921            return Ok(Some((y, scale)));
6922        }
6923        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
6924        let name = match qtype {
6925            QT_Q8_0 => "qmatvec_q8_0_dp4a", QT_Q4_K => "qmatvec_q4_K_dp4a",
6926            QT_Q6_K => "qmatvec_q6_K_dp4a", QT_Q5_K => "qmatvec_q5_K_dp4a",
6927            QT_Q3_K => "qmatvec_q3_K_dp4a",
6928            QT_NVFP4 => if rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
6929            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
6930            _ => return Ok(None),
6931        };
6932        let f = self.func(name);
6933        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
6934        let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
6935        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6936        let __s_b = self.gpu.stream();
6937        let mut b = __s_b.launch_builder(&f);
6938        b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
6939        unsafe { b.launch(cfg)?; }
6940        Ok(Some((y, scale)))
6941    }
6942
6943    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
6944    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
6945    pub fn mmvq_supports(&self, qtype: i32) -> bool {
6946        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
6947        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
6948        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
6949        // is a pure function of the dtype — the decode-parity law holds under every env.
6950        if qtype == QT_F8_E4M3 { return true; }
6951        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") { return false; }
6952        matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0)
6953    }
6954
6955    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
6956    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
6957    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
6958    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
6959    pub fn qmatvec_mmvq(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6960                        m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32,
6961                        rp: bool)
6962                        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6963        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite GEMM output: skip memset
6964        self.qmatvec_mmvq_into(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y)?;
6965        Ok(y)
6966    }
6967
6968    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
6969    #[allow(clippy::too_many_arguments)]
6970    pub fn qmatvec_mmvq_into(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6971                        m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32,
6972                        rp: bool, y: &mut CudaSlice<f32>)
6973                        -> Result<(), Box<dyn std::error::Error>> {
6974        debug_assert!(y.len() >= m * out_f);
6975        const ROWS_PER_BLOCK: u32 = 4;   // matches MEMRA_MMVQ_ROWS in qmatvec.cu
6976        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
6977        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
6978        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
6979        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
6980        if qtype == QT_Q8_0 && rp && m == 1 && out_f >= 64
6981            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
6982            && {
6983                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6984                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
6985            }
6986        {
6987            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
6988            let cfg = LaunchConfig {
6989                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
6990                block_dim: (32, 2, 1),
6991                shared_mem_bytes: 0,
6992            };
6993            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
6994            let __s_b = self.gpu.stream();
6995            let mut b = __s_b.launch_builder(&f);
6996            b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
6997            unsafe { b.launch(cfg)?; }
6998            if scale != 1.0 { self.scale_inplace(y, scale, out_f)?; }
6999            return Ok(());
7000        }
7001        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
7002        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
7003        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
7004        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
7005        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
7006        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
7007        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
7008        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
7009        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) { 2 } else { 1 };
7010        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
7011        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
7012        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
7013        // valid-window interleaved, bit-identical per row — same dot program).
7014        if m == 1 && qtype == QT_Q4_0 {
7015            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
7016            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
7017            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
7018            mr = *Q40MR.get_or_init(|| std::env::var("MEMRA_Q40_MR").ok()
7019                .and_then(|v| v.parse().ok()).unwrap_or(1));
7020        }
7021        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
7022        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
7023        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
7024        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
7025        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
7026        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
7027        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
7028        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
7029        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
7030        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
7031        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
7032        let q5_force = q5_mode.as_deref() == Some("2");
7033        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
7034        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
7035        let q5_il = qtype == QT_Q5_K && m == 1
7036            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
7037        if q5_il && !q5_force && out_f > 65536 { mr = 1; }
7038        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
7039        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
7040        if qtype == QT_Q4_0 && rp && mr != 1 { mr = 2; }
7041        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
7042        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
7043        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
7044        if qtype == QT_Q8_0 && rp {
7045            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
7046            mr = *Q80MR.get_or_init(|| std::env::var("MEMRA_Q80_MR").ok()
7047                .and_then(|v| v.parse().ok()).unwrap_or(1));
7048        }
7049        let name = match (qtype, mr, rp) {
7050            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
7051            (QT_NVFP4, 2, true)  => "qmatvec_nvfp4_mmvq_mr2_rp",
7052            (QT_NVFP4, _, true)  => "qmatvec_nvfp4_mmvq_rp",
7053            (QT_Q4_0, 1, true)   => "qmatvec_q4_0_mmvq_rp",
7054            (QT_Q4_0, _, true)   => "qmatvec_q4_0_mmvq_mr2_rp",
7055            (QT_Q5_K, 2, _) => if q5_il { "qmatvec_q5_K_mmvq_mr2_il" } else { "qmatvec_q5_K_mmvq_mr2" },
7056            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
7057            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
7058            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
7059            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
7060            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
7061            (QT_Q8_0, _, true) if in_f % 1024 == 0 && {
7062                static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7063                *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
7064            } => "qmatvec_q8_0_mmvq_rpca",
7065            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
7066            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
7067            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
7068            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
7069            // reach a GGUF-layout kernel or vice versa.
7070            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
7071            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
7072            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
7073            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
7074            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
7075            (QT_Q5_K, _, _) => if q5_il { "qmatvec_q5_K_mmvq_il" } else { "qmatvec_q5_K_mmvq" },
7076            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
7077            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
7078            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
7079            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
7080        };
7081        let f = self.func(name);
7082        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
7083        let rows_per_block = ROWS_PER_BLOCK * mr;
7084        let cfg = LaunchConfig {
7085            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, m as u32, 1),
7086            block_dim: (32, ROWS_PER_BLOCK, 1),   // warp-per-row (x mr rows each)
7087            shared_mem_bytes: 0,                  // warp-only reduce at m=1
7088        };
7089        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7090        let __s_b = self.gpu.stream();
7091        let mut b = __s_b.launch_builder(&f);
7092        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
7093        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
7094        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
7095        // weight_scale). Other mmvq kernels keep the 8-arg signature.
7096        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
7097            b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb).arg(&scale);
7098            unsafe { b.launch(cfg)?; }
7099        } else if Self::pdl_on() && Self::pdl_mmvq_on()
7100            && matches!(name, "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq"
7101                              | "qmatvec_q6_K_mmvq_rp") {
7102            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
7103            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
7104            // names may take this launch (unmarked kernels would read unordered).
7105            {
7106            use cudarc::driver::{DevicePtr, DevicePtrMut};
7107            let s = &self.gpu.stream();
7108            let (pw, _g0) = bytes.device_ptr(s); let (paq, _g1) = aq.device_ptr(s);
7109            let (pad, _g2) = ad.device_ptr(s); let (py, _g3) = y.device_ptr_mut(s);
7110            let mut ps = [
7111                &pw as *const _ as *mut std::ffi::c_void, &paq as *const _ as *mut _,
7112                &pad as *const _ as *mut _, &py as *const _ as *mut _,
7113                &inf as *const _ as *mut _, &outf as *const _ as *mut _,
7114                &mi as *const _ as *mut _, &rb as *const _ as *mut _,
7115            ];
7116            unsafe { self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?; }
7117            }
7118            if scale != 1.0 { self.scale_inplace(y, scale, m * out_f)?; }
7119        } else {
7120            b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7121            unsafe { b.launch(cfg)?; }
7122            if scale != 1.0 { self.scale_inplace(y, scale, m * out_f)?; }
7123        }
7124        Ok(())
7125    }
7126
7127    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
7128    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
7129    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
7130    pub fn qmatvec_mmvq_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
7131                            out_f: usize, qtype: i32, row_bytes: usize, rp: bool)
7132                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7133        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7134        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
7135    }
7136
7137    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
7138    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
7139    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
7140    pub fn batched_supports(&self, qtype: i32) -> bool {
7141        matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0)
7142    }
7143
7144    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
7145    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
7146    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
7147    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
7148    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
7149    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
7150    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
7151    pub fn iq_fast_enabled() -> bool {
7152        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7153        *ON.get_or_init(|| std::env::var("MEMRA_IQ_FAST").map(|v| v != "0").unwrap_or(true))
7154    }
7155
7156    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
7157    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
7158    pub fn b8_enabled() -> bool {
7159        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7160        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
7161    }
7162
7163    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
7164    pub fn batched_mcols(m: usize) -> usize {
7165        if m == 2 { 2 } else if m <= 4 { 4 } else if m <= 8 { 8 } else { 16 }
7166    }
7167
7168    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
7169    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
7170    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
7171    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
7172    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
7173        Some(match (qtype, mcols) {
7174            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2", (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
7175            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
7176            // rp-ONLY tier: qmatvec_q8_0_mmvq_b16 has no base twin — the mcols==16 dispatch
7177            // appends _rp, and every caller gates Q8_0 m>8 on the q8rp mirror being present.
7178            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
7179            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2", (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
7180            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
7181            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2", (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
7182            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
7183            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2", (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
7184            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8", (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
7185            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2", (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
7186            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
7187            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2", (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
7188            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
7189            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2", (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
7190            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8", (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
7191            _ => return None,
7192        })
7193    }
7194
7195    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
7196    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
7197    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
7198    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
7199    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
7200    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
7201    ///
7202    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
7203    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
7204    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
7205    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
7206    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
7207    /// msweep on all six 27B shapes (2026-07-03):
7208    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
7209    ///          it applies for b4 (-3..-14%), never loses;
7210    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
7211    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
7212    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
7213    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
7214    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
7215    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
7216    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
7217    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
7218    /// b2: in_f>=6144 -> r2, else base.
7219    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
7220    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
7221    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
7222    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
7223    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
7224    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
7225    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
7226    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
7227    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
7228    /// Device SM count (cached) — grid-fill policy input.
7229    pub fn sm_count(&self) -> i32 {
7230        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
7231        *SMS.get_or_init(|| {
7232            use cudarc::driver::sys::CUdevice_attribute_enum as A;
7233            self.gpu.ctx.attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT).unwrap_or(82)
7234        })
7235    }
7236
7237    pub fn batched_variant(&self, _m: usize, in_f: usize, out_f: usize, qtype: i32,
7238                           row_bytes: usize, mcols: usize, rp: bool) -> &'static str {
7239        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
7240        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
7241        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
7242        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
7243        if qtype == QT_Q8_0 {
7244            return if rp { "rp" } else { "base" };
7245        }
7246        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7247        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
7248            Ok("base") => "base", Ok("pf") => "pf", Ok("r2") => "r2", Ok("r2w8") => "r2w8",
7249            Ok("pfr2") => "pfr2", Ok("ca") => "ca", Ok("car2") => "car2",
7250            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
7251            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
7252            Ok("rp") => "rp", Ok("rpr2") => "rpr2", Ok("rpr2w8") => "rpr2w8",
7253            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
7254            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
7255            Ok("rpca") => "rpca", Ok("rpcar2") => "rpcar2",
7256            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
7257            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
7258            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
7259            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
7260            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
7261            // bit-identical to the decode path — measurement corpus ONLY, never auto).
7262            Ok("rpsc") => "rpsc", Ok("rpms") => "rpms", Ok("rpmsc") => "rpmsc",
7263            Ok("rpks") => "rpks", Ok("rpksc") => "rpksc",
7264            _ => "auto",
7265        });
7266        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
7267        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
7268        // shapes qualify; anything else falls back to the register variants.
7269        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
7270        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
7271        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
7272        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
7273        // forced MEMRA_MMVQ_BV values still work).
7274        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7275        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
7276        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
7277        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
7278        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
7279        let sms = *SMS.get_or_init(|| {
7280            use cudarc::driver::sys::CUdevice_attribute_enum as A;
7281            self.gpu.ctx.attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT).unwrap_or(82)
7282        });
7283        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
7284        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
7285        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
7286        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
7287        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
7288        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
7289        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
7290        // AUTO RULE = the measured winners table (differs from NVFP4's!):
7291        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
7292        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
7293        //     r2 1258us) — kernels kept behind the force seam for the corpus;
7294        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
7295        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
7296        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
7297        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
7298        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
7299        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
7300        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
7301        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
7302        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
7303        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
7304        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
7305        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7306        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
7307            Ok("base") => "base", Ok("r2") => "r2", Ok("r2w8") => "r2w8",
7308            _ => "auto",
7309        });
7310        let variant: &'static str = if qtype == QT_Q4_0 {
7311            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
7312            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
7313            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
7314            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7315            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
7316                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
7317                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
7318                // + syncs cost more than the stalls, bank-pad made no difference);
7319                // register load-ahead flat (nvcc already reorders). The b-tier limiter
7320                // is still unidentified — see the jsonl row.
7321                Ok("base") => "base", Ok("r2") => "r2", Ok("ms") => "ms", Ok("sm") => "sm",
7322                Ok("la") => "la", _ => "auto",
7323            });
7324            let v = if q40 != "auto" { q40 }
7325            else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 { "r2" } else { "base" };
7326            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
7327            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
7328            // and the limiter is the per-column activation load chain (long_scoreboard
7329            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
7330            if rp { match v { "ms" => "r2ms_rp", "sm" => "r2sm_rp", "la" => "r2la_rp",
7331                              "r2" => "r2_rp", _ => "rp" } }
7332            else if matches!(v, "ms" | "sm" | "la") { "r2" } else { v }
7333        } else if qtype != QT_NVFP4 && !kq_r2 {
7334            "base"
7335        } else if kq_r2 && rp {
7336            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
7337            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
7338            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
7339            "rp"
7340        } else if kq_r2 {
7341            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
7342            // mcols != 4 forced r2w8 falls to unbounded r2.
7343            if kq_bv != "auto" {
7344                if kq_bv == "r2w8" && mcols != 4 { "r2" } else { kq_bv }
7345            } else if bv != "auto" {
7346                match bv {
7347                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
7348                    "r2w8" | "rpr2w8" => if mcols != 4 { "r2" } else { "r2w8" },
7349                    _ => "base",   // base/pf/ca/rp forced -> base (no such k-quant kernels)
7350                }
7351            } else {
7352                let blocks = (out_f + 7) / 8;
7353                let waves = blocks as f64 / (7 * sms as usize) as f64;
7354                let filled = blocks >= 4 * sms as usize;
7355                let use_r2 = if qtype == QT_Q4_K { filled } else { waves >= 2.0 };
7356                if use_r2 { "r2" } else { "base" }
7357            }
7358        } else if bv != "auto" {
7359            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
7360            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
7361            // unsupported (shape, mcols) combos fall back to pf/r2.
7362            // On rp buffers, forced legacy names map to their rp twins (layout law).
7363            let v = if bv == "r2w8" && mcols == 2 { "r2" }
7364                else if bv == "ca" && (!ca_ok || mcols == 8) { "pf" }
7365                else if bv == "car2" && (!ca_ok || mcols == 8) { "r2" }
7366                else if bv == "pfr2" && mcols == 8 { "r2" }
7367                else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 { "rpr2" }
7368                // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
7369                else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
7370                    if mcols == 8 { "rpr2w8" } else { "rpr2" }
7371                }
7372                else if bv == "rpcar2" && mcols == 2 { "rpca" }
7373                // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
7374                // (rpms has no smem and no alignment need — always valid on rp buffers).
7375                else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok { "rpr2" }
7376                else if (bv == "rpks" || bv == "rpksc") && !ks_ok { "rpr2" }
7377                else { bv };
7378            if rp {
7379                match v {
7380                    "base" | "pf" | "ca" | "rp" => "rp",
7381                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
7382                    "r2w8" | "rpr2w8" => if mcols == 2 { "rpr2" } else { "rpr2w8" },
7383                    other => other,   // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
7384                }
7385            } else { v }
7386        } else if mcols == 8 {
7387            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
7388            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
7389            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
7390            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
7391            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
7392            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
7393            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
7394            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
7395            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
7396            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
7397            if rp { if sc_ok { "rpsc" } else { "rpr2w8" } } else { "r2w8" }
7398        } else if mcols >= 4 {
7399            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
7400            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
7401            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
7402            let blocks = (out_f + 7) / 8;
7403            let r7 = 7 * sms as usize;
7404            let r8 = 8 * sms as usize;
7405            let waves = blocks as f64 / r7 as f64;
7406            let filled = blocks >= 4 * sms as usize;
7407            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
7408            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
7409            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
7410            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
7411                // the extra residency drops the INTEGER wave count -> the straggler wave a
7412                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
7413                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
7414                if rp { "rpr2w8" } else { "r2w8" }
7415            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
7416                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
7417                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
7418                if rp { "rpr2" } else { "r2" }
7419            } else {
7420                // fractional straggler-wave window with no crossing, or grid too small to fill
7421                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
7422                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
7423                if rp { "rp" } else { "pf" }
7424            }
7425        } else if in_f >= 6144 {
7426            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
7427            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
7428            // stays.
7429            if rp { "rpr2" } else { "r2" }
7430        }
7431        else if rp {
7432            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
7433            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
7434            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
7435            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
7436            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
7437            if sc_ok && waves >= 0.9 && waves <= 1.1 { "rpsc" } else { "rp" }
7438        } else { "base" };
7439        variant
7440    }
7441
7442    pub fn qmatvec_mmvq_batched(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
7443                                m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize,
7444                                mcols: usize, scale: f32, rp: bool)
7445                                -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7446        const ROWS_PER_BLOCK: u32 = 4;
7447        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
7448        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
7449        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
7450        // weight keeps its rp-layout kernel family regardless of the override.
7451        let forced: Option<&'static str> = {
7452            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
7453            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
7454                .as_deref()
7455                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
7456        };
7457        let variant = match forced {
7458            Some(v) if !rp || v.contains("rp") => v,
7459            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
7460        };
7461        let base_name = Self::batched_kernel_name(qtype, mcols)
7462            .ok_or_else(|| format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}"))?;
7463        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
7464        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
7465        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
7466        let variant = if mcols == 16 { if rp { "rp" } else { "base" } } else { variant };
7467        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
7468        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
7469        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
7470        // per-(token,row) chain (columns c >= m never execute in either form) ->
7471        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
7472        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
7473        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7474        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
7475        if b567 && qtype == QT_NVFP4 && rp && mcols == 8 && (5..=7).contains(&m)
7476            && matches!(variant, "rpsc" | "rpr2w8") {
7477            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
7478            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
7479            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7480            let cfg = LaunchConfig {
7481                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
7482                block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0 };
7483            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7484            let __s_b = self.gpu.stream();
7485            let mut b = __s_b.launch_builder(&f);
7486            b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7487            unsafe { b.launch(cfg)?; }
7488            if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7489            return Ok(y);
7490        }
7491        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
7492            "base" => (base_name.into(), ROWS_PER_BLOCK),
7493            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
7494            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
7495            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
7496            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
7497            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
7498            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
7499            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
7500            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
7501            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
7502            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
7503            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
7504            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
7505            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
7506            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
7507        };
7508        debug_assert!(!rp || name.contains("_rp"), "rp weight dispatched to a GGUF-layout kernel");
7509        let f = self.func(&name);
7510        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7511        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
7512        let smem = if name.contains("_r2sm_rp") { (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32 }
7513                   else { 0 };
7514        let cfg = LaunchConfig {
7515            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
7516            block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: smem };
7517        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7518        let __s_b = self.gpu.stream();
7519        let mut b = __s_b.launch_builder(&f);
7520        b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7521        unsafe { b.launch(cfg)?; }
7522        if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7523        Ok(y)
7524    }
7525
7526    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
7527    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
7528    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
7529    pub fn qmatvec_batched_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
7530                               in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize,
7531                               rp: bool)
7532                               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7533        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7534        self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp)
7535    }
7536
7537    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
7538    pub fn qmatvec_nvfp4_batched_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
7539                                     in_f: usize, out_f: usize, row_bytes: usize, mcols: usize,
7540                                     rp: bool)
7541                                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7542        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
7543    }
7544
7545    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
7546    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
7547    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
7548    fn try_fp4_gemm(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize,
7549                    in_f: usize, out_f: usize)
7550                    -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7551        use crate::model::GpuTensor;
7552        if cfg!(memra_portable_cuda) { return Ok(None); }
7553        if std::env::var("MEMRA_FP4").is_err() { return Ok(None); }
7554        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
7555        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
7556        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
7557        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
7558        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
7559        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
7560        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
7561        // for the common no-macro-scale case.
7562        #[cfg(memra_cutlass)]
7563        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
7564            if let GpuTensor::Quant { bytes, qtype, scale, row_bytes, cutlass, .. } = w {
7565                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
7566                    if let Some(cw) = cutlass {
7567                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
7568                        let y = self.cutlass_fp4_gemm(&cw.b_packed, &cw.sfb_swizzled, x, *scale,
7569                                                      m, out_f, in_f)?;
7570                        return Ok(Some(y));
7571                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
7572                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
7573                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
7574                        // (the load-time repack ~doubles it) — needed for models that don't fit the
7575                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
7576                        let (b_packed, sfb_sw) = self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
7577                        let y = self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
7578                        return Ok(Some(y));
7579                    }
7580                }
7581            }
7582        }
7583        if let GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } = w {
7584            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
7585            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
7586            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
7587                let y = self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
7588                return Ok(Some(y));
7589            }
7590        }
7591        Ok(None)
7592    }
7593
7594    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
7595    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
7596    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
7597    pub fn rms_norm_f16out(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>,
7598                           dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
7599                           ncols: usize, nrows: usize, eps: f32)
7600                           -> Result<(), Box<dyn std::error::Error>> {
7601        let f = self.func("rms_norm_f16out_f32");
7602        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
7603        let (nc, e) = (ncols as i32, eps);
7604        let __s_b = self.gpu.stream();
7605        let mut b = __s_b.launch_builder(&f);
7606        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
7607        unsafe { b.launch(cfg)?; }
7608        Ok(())
7609    }
7610
7611    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
7612    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
7613    #[allow(clippy::too_many_arguments)]
7614    pub fn add_rms_norm_f16out(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>,
7615                               res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
7616                               dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32)
7617                               -> Result<(), Box<dyn std::error::Error>> {
7618        let f = self.func("add_rms_norm_f16out_f32");
7619        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
7620        let (nc, e) = (ncols as i32, eps);
7621        let __s_lb = self.gpu.stream();
7622        let mut lb = __s_lb.launch_builder(&f);
7623        lb.arg(a).arg(b).arg(w).arg(res).arg(dst).arg(dst16).arg(&nc).arg(&e);
7624        unsafe { lb.launch(cfg)?; }
7625        Ok(())
7626    }
7627
7628    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
7629    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
7630    pub fn matmul_group_xh(&self, ws: &[&crate::model::GpuTensor], x: &CudaSlice<f32>,
7631                           xh: &CudaSlice<u8>, m: usize)
7632                           -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7633        let mut out = Vec::with_capacity(ws.len());
7634        let in_f = ws[0].in_features();
7635        for w in ws {
7636            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
7637                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
7638                    out.push(y);
7639                    continue;
7640                }
7641            }
7642            out.push(self.matmul(w, x, m)?);
7643        }
7644        Ok(out)
7645    }
7646
7647    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
7648    /// GDN steps). Layouts [T, H].
7649    pub fn gdn_pad_mask(&self, beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
7650                        len_d: &CudaSlice<i32>, h: usize, t: usize)
7651                        -> Result<(), Box<dyn std::error::Error>> {
7652        let f = self.func("gdn_pad_mask_f32");
7653        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
7654        let (hi, ti) = (h as i32, t as i32);
7655        let __s_b = self.gpu.stream();
7656        let mut b = __s_b.launch_builder(&f);
7657        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
7658        unsafe { b.launch(cfg)?; }
7659        Ok(())
7660    }
7661
7662    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
7663    /// gather for the padded prime graph's h_seed/hlast.
7664    pub fn row_gather_dev(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
7665                          len_d: &CudaSlice<i32>, ncols: usize)
7666                          -> Result<(), Box<dyn std::error::Error>> {
7667        let f = self.func("row_gather_dev_f32");
7668        let cfg = LaunchConfig::for_num_elems(ncols as u32);
7669        let nc = ncols as i32;
7670        let __s_b = self.gpu.stream();
7671        let mut b = __s_b.launch_builder(&f);
7672        b.arg(src).arg(dst).arg(len_d).arg(&nc);
7673        unsafe { b.launch(cfg)?; }
7674        Ok(())
7675    }
7676
7677    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
7678    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
7679    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
7680    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
7681    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
7682    /// different in_f) falls back to its own `matmul` — behavior unchanged.
7683    pub fn matmul_group(&self, ws: &[&crate::model::GpuTensor], x: &CudaSlice<f32>, m: usize)
7684                        -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7685        use crate::model::GpuTensor;
7686        let mut out = Vec::with_capacity(ws.len());
7687        let any_mirror = ws.iter().any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
7688        if m >= 16 && any_mirror && !self.verify_exact_on() {
7689            let in_f = ws[0].in_features();
7690            let xh = self.f16_act(x, m * in_f, in_f)?;
7691            for w in ws {
7692                if w.in_features() == in_f {
7693                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
7694                        out.push(y);
7695                        continue;
7696                    }
7697                }
7698                out.push(self.matmul(w, x, m)?);
7699            }
7700            return Ok(out);
7701        }
7702        for w in ws {
7703            out.push(self.matmul(w, x, m)?);
7704        }
7705        Ok(out)
7706    }
7707
7708    /// Cross-request grouped matmul (task #13): run ONE projection group over the
7709    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
7710    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
7711    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
7712    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
7713    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
7714    pub fn matmul_group_multi(&self, ws: &[&crate::model::GpuTensor],
7715                              xs: &[&CudaSlice<f32>], ms: &[usize])
7716                              -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
7717        assert_eq!(xs.len(), ms.len());
7718        let in_f = ws[0].in_features();
7719        let total: usize = ms.iter().sum();
7720        let mut xcat = self.uninit(total * in_f)?;
7721        let mut off = 0usize;
7722        for (x, &m) in xs.iter().zip(ms) {
7723            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
7724            off += m;
7725        }
7726        let ys = self.matmul_group(ws, &xcat, total)?;
7727        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
7728        for (w, y) in ws.iter().zip(ys) {
7729            let out_f = w.out_features();
7730            let mut off = 0usize;
7731            for (s, &m) in ms.iter().enumerate() {
7732                let mut ys_s = self.uninit(m * out_f)?;
7733                let src = y.slice(off * out_f..(off + m) * out_f);
7734                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
7735                out[s].push(ys_s);
7736                off += m;
7737            }
7738        }
7739        Ok(out)
7740    }
7741
7742    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
7743    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
7744    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
7745    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
7746    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
7747    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
7748    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
7749    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
7750    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
7751    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
7752        use crate::model::GpuTensor;
7753        if !legacy_quant_gemm_allowed(
7754            cfg!(memra_portable_cuda),
7755            cfg!(memra_hopper_mma),
7756            std::env::var_os("MEMRA_NO_GEMM").is_some(),
7757        ) {
7758            return false;
7759        }
7760        match w {
7761            GpuTensor::Quant { qtype, .. } =>
7762                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
7763                || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0),
7764            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
7765        }
7766    }
7767
7768    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
7769    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
7770    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
7771    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
7772    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
7773    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
7774    pub fn qmatvec_gemm(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
7775                        m: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7776        use crate::model::GpuTensor;
7777        let in_f = w.in_features();
7778        let out_f = w.out_features();
7779        let (bytes, qtype, row_bytes, scale, rp) = match w {
7780            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
7781            _ => unreachable!("gemm_supports guaranteed Quant"),
7782        };
7783        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
7784        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
7785        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
7786        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
7787        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
7788        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
7789            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
7790                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
7791                if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7792                return Ok(y);
7793            }
7794        }
7795        let name = match qtype {
7796            QT_Q8_0 => "qmatvec_gemm_q8_0", QT_Q4_K => "qmatvec_gemm_q4_K",
7797            QT_Q4_0 => if rp { "qmatvec_gemm_q4_0_rp" } else { "qmatvec_gemm_q4_0" },
7798            QT_Q5_K => "qmatvec_gemm_q5_K",
7799            QT_Q6_K => "qmatvec_gemm_q6_K",
7800            QT_NVFP4 => if rp { "qmatvec_gemm_nvfp4_rp" } else { "qmatvec_gemm_nvfp4" },
7801            _ => unreachable!(),
7802        };
7803        let f = self.func(name);
7804        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite GEMM output: skip memset
7805        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
7806        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
7807        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
7808        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
7809        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
7810        let k1_tile = if is_k1 { k1_launch_override().unwrap_or((128, 128, 8)) } else { (128, 128, 8) };
7811        let (bm, bn): (u32, u32) = if is_k1 { (k1_tile.0, k1_tile.1) } else { (64, 256) };
7812        let warps: u32 = if is_k1 { k1_tile.2 } else {
7813            match qtype { QT_NVFP4 => 8, _ => 4 }
7814        };
7815        let cfg = LaunchConfig {
7816            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
7817            block_dim: (32, warps, 1),
7818            shared_mem_bytes: 0,
7819        };
7820        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7821        let __s_b = self.gpu.stream();
7822        let mut b = __s_b.launch_builder(&f);
7823        b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7824        unsafe { b.launch(cfg)?; }
7825        if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7826        Ok(y)
7827    }
7828
7829    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
7830    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
7831    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
7832    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
7833    pub fn qmatvec_gemm_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
7834                            out_f: usize, qtype: i32, row_bytes: usize)
7835                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7836        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7837        let name = match qtype {
7838            QT_Q8_0 => "qmatvec_gemm_q8_0", QT_Q4_K => "qmatvec_gemm_q4_K",
7839            QT_Q4_0 => "qmatvec_gemm_q4_0",
7840            QT_Q5_K => "qmatvec_gemm_q5_K",
7841            QT_Q6_K => "qmatvec_gemm_q6_K", QT_NVFP4 => "qmatvec_gemm_nvfp4",
7842            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
7843            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
7844        };
7845        let f = self.func(name);
7846        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite GEMM output: skip memset
7847        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
7848        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
7849        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
7850        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
7851        let k1_tile = if is_k1 { k1_launch_override().unwrap_or((128, 128, 8)) } else { (128, 128, 8) };
7852        let (bm, bn): (u32, u32) = if is_k1 { (k1_tile.0, k1_tile.1) } else { (64, 256) };
7853        let warps: u32 = if is_k1 { k1_tile.2 } else {
7854            match qtype { QT_NVFP4 | QT_NVFP4_RP => 8, _ => 4 }
7855        };
7856        let cfg = LaunchConfig {
7857            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
7858            block_dim: (32, warps, 1), shared_mem_bytes: 0,
7859        };
7860        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7861        let __s_b = self.gpu.stream();
7862        let mut b = __s_b.launch_builder(&f);
7863        b.arg(bytes).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7864        unsafe { b.launch(cfg)?; }
7865        Ok(y)
7866    }
7867
7868    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
7869    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
7870    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
7871    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
7872    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
7873    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
7874    pub fn qmatvec_gemm_q8_0_wgmma_raw(&self, rp4: &CudaSlice<u8>, aq: &CudaSlice<i8>,
7875                                       ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize)
7876                                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7877        assert!(out_f % 64 == 0 && in_f % 32 == 0, "wgmma GEMM needs out_f%64==0, in_f%32==0");
7878        let f = self.func("qmatvec_gemm_q8_0_wgmma");
7879        let mut y = self.alloc_uninit::<f32>(m * out_f)?;  // full-overwrite GEMM output
7880        let cfg = LaunchConfig {
7881            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
7882            block_dim: (128, 1, 1), shared_mem_bytes: 0,
7883        };
7884        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
7885        let __s_b = self.gpu.stream();
7886        let mut b = __s_b.launch_builder(&f);
7887        b.arg(rp4).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi);
7888        unsafe { b.launch(cfg)?; }
7889        Ok(y)
7890    }
7891
7892    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
7893    pub fn scale_inplace(&self, y: &mut CudaSlice<f32>, s: f32, n: usize)
7894                         -> Result<(), Box<dyn std::error::Error>> {
7895        let f = self.func("scale_f32");
7896        let cfg = LaunchConfig::for_num_elems(n as u32);
7897        let (sf, ni) = (s, n as i32);
7898        let __s_b = self.gpu.stream();
7899        let mut b = __s_b.launch_builder(&f);
7900        b.arg(y).arg(&sf).arg(&ni);
7901        unsafe { b.launch(cfg)?; }
7902        Ok(())
7903    }
7904
7905    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
7906    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
7907    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
7908    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
7909    pub fn bf16_to_f32(&self, data: &cudarc::driver::CudaView<'_, u8>, n: usize)
7910                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7911        let mut out = self.alloc_uninit::<f32>(n)?;
7912        let f = self.func("bf16_to_f32");
7913        let cfg = LaunchConfig::for_num_elems(n as u32);
7914        let ni = n as i32;
7915        let __s_b = self.gpu.stream();
7916        let mut b = __s_b.launch_builder(&f);
7917        b.arg(data).arg(&mut out).arg(&ni);
7918        unsafe { b.launch(cfg)?; }
7919        Ok(out)
7920    }
7921
7922    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
7923    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
7924    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
7925    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
7926    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
7927    /// calls, the spec-verify contract) vs plain linear.
7928    fn linear_bf16_chunked(&self, x: &CudaSlice<f32>, data: &CudaSlice<u8>, m: usize,
7929                           in_f: usize, out_f: usize, exact: bool)
7930                           -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7931        const CHUNK_BYTES: usize = 256 << 20;
7932        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
7933        if chunk_rows >= out_f {
7934            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
7935            return if exact { self.linear_decode_exact(x, &wf32, m, in_f, out_f) }
7936                   else { self.linear(x, &wf32, m, in_f, out_f) };
7937        }
7938        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7939        let mut r0 = 0usize;
7940        while r0 < out_f {
7941            let rows = chunk_rows.min(out_f - r0);
7942            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
7943            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
7944            let yc = if exact { self.linear_decode_exact(x, &wf32, m, in_f, rows)? }
7945                     else { self.linear(x, &wf32, m, in_f, rows)? };
7946            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
7947            for mi in 0..m {
7948                let src = yc.slice(mi * rows..(mi + 1) * rows);
7949                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
7950                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
7951            }
7952            r0 += rows;
7953        }
7954        Ok(y)
7955    }
7956
7957    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
7958    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
7959    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
7960    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
7961    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
7962    /// router/shexp sites and matmul_decode_exact's Float arm.
7963    pub fn linear_decode_exact(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize,
7964                               in_f: usize, out_f: usize)
7965                               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7966        if m_tokens == 1 { return self.linear(x, w, 1, in_f, out_f); }
7967        let xv = self.view(x, m_tokens * in_f);
7968        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
7969        for t in 0..m_tokens {
7970            let row = xv.slice(t * in_f..(t + 1) * in_f);
7971            let mut xr = self.alloc_uninit::<f32>(in_f)?;
7972            self.copy_view_into(&mut xr, 0, &row, in_f)?;
7973            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
7974            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
7975        }
7976        Ok(y)
7977    }
7978
7979    pub fn linear(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize, in_f: usize, out_f: usize)
7980                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7981        use cudarc::cublaslt::{Matmul, MatmulConfig};
7982        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?;  // cuBLASLt beta=0: C fully written
7983        let cfg = MatmulConfig {
7984            transa: true, transb: false, transc: false,
7985            m: out_f as u64, n: m_tokens as u64, k: in_f as u64,
7986            alpha: 1.0, lda: in_f as i64, ldb: in_f as i64, beta: 0.0, ldc: out_f as i64,
7987            stride_a: None, stride_b: None, stride_c: None, stride_bias: None, batch_size: None,
7988        };
7989        unsafe { self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?; }
7990        Ok(c)
7991    }
7992
7993    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
7994    pub fn sdpa_naive(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
7995                      o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
7996                      t: usize, t_kv: usize, scale: f32, causal: bool)
7997                      -> Result<(), Box<dyn std::error::Error>> {
7998        let f = self.func("sdpa_naive_f32");
7999        let cfg = LaunchConfig {
8000            grid_dim: (n_head as u32, t as u32, 1),
8001            block_dim: (128, 1, 1),
8002            shared_mem_bytes: (t_kv * 4) as u32,
8003        };
8004        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8005        let __s_b = self.gpu.stream();
8006        let mut b = __s_b.launch_builder(&f);
8007        b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8008        unsafe { b.launch(cfg)?; }
8009        Ok(())
8010    }
8011
8012    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
8013    #[allow(clippy::too_many_arguments)]
8014    pub fn sdpa_naive_w(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8015                        o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8016                        t: usize, t_kv: usize, scale: f32, causal: bool, window: usize)
8017                        -> Result<(), Box<dyn std::error::Error>> {
8018        let f = self.func("sdpa_naive_w_f32");
8019        let cfg = LaunchConfig {
8020            grid_dim: (n_head as u32, t as u32, 1),
8021            block_dim: (128, 1, 1),
8022            shared_mem_bytes: (t_kv * 4) as u32,
8023        };
8024        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8025                                                t as i32, t_kv as i32, causal as i32, window as i32);
8026        let __s_b = self.gpu.stream();
8027        let mut b = __s_b.launch_builder(&f);
8028        b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8029         .arg(&scale).arg(&cz).arg(&wi);
8030        unsafe { b.launch(cfg)?; }
8031        Ok(())
8032    }
8033
8034    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
8035    pub fn sdpa_naive_view(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<f32>,
8036                           v: &cudarc::driver::CudaView<f32>, o: &mut CudaSlice<f32>,
8037                           head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize,
8038                           scale: f32, causal: bool) -> Result<(), Box<dyn std::error::Error>> {
8039        let f = self.func("sdpa_naive_f32");
8040        let cfg = LaunchConfig {
8041            grid_dim: (n_head as u32, t as u32, 1), block_dim: (128, 1, 1),
8042            shared_mem_bytes: (t_kv * 4) as u32,
8043        };
8044        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8045        let __s_b = self.gpu.stream();
8046        let mut b = __s_b.launch_builder(&f);
8047        b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8048        unsafe { b.launch(cfg)?; }
8049        Ok(())
8050    }
8051
8052    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
8053    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
8054    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
8055    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
8056    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
8057    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
8058    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
8059    #[allow(clippy::too_many_arguments)]
8060    pub fn fa_dequant_kv_view_f32(&self, k: &cudarc::driver::CudaView<u8>,
8061                                  v: &cudarc::driver::CudaView<u8>,
8062                                  kf: &mut CudaSlice<f32>, vf: &mut CudaSlice<f32>,
8063                                  kv_dim_k: usize, kv_dim_v: usize, t_kv: usize,
8064                                  k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8065                                  -> Result<(), Box<dyn std::error::Error>> {
8066        let f = if g { self.func_g("fa_dequant_kv_ws_f32") } else { self.func("fa_dequant_kv_ws_f32") };
8067        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
8068        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8069        let cfg = LaunchConfig { grid_dim: (nblk.max(1), 1, 1), block_dim: (256, 1, 1),
8070                                 shared_mem_bytes: 0 };
8071        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
8072        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8073        let __s_b = self.gpu.stream();
8074        let mut b = __s_b.launch_builder(&f);
8075        b.arg(k).arg(v).arg(&mut *kf).arg(&mut *vf).arg(&kdk).arg(&kdv).arg(&tkvi).arg(&ktb).arg(&vtb);
8076        unsafe { b.launch(cfg)?; }
8077        Ok(())
8078    }
8079
8080    #[allow(clippy::too_many_arguments)]
8081    pub fn sdpa_naive_quantized_view(
8082        &self,
8083        q: &CudaSlice<f32>,
8084        k: &cudarc::driver::CudaView<u8>,
8085        v: &cudarc::driver::CudaView<u8>,
8086        o: &mut CudaSlice<f32>,
8087        head_dim: usize,
8088        n_head: usize,
8089        n_head_kv: usize,
8090        t: usize,
8091        t_kv: usize,
8092        scale: f32,
8093        causal: bool,
8094        k_tok_bytes: usize,
8095        v_tok_bytes: usize,
8096    ) -> Result<(), Box<dyn std::error::Error>> {
8097        let kv_dim = n_head_kv * head_dim;
8098        let mut kf = self.uninit(t_kv * kv_dim)?;
8099        let mut vf = self.uninit(t_kv * kv_dim)?;
8100        let f = self.func("fa_dequant_kv_ws_f32");
8101        let total = (2 * t_kv * kv_dim) as u64;
8102        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8103        let cfg = LaunchConfig {
8104            grid_dim: (nblk.max(1), 1, 1),
8105            block_dim: (256, 1, 1),
8106            shared_mem_bytes: 0,
8107        };
8108        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
8109        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
8110        let __s_b = self.gpu.stream();
8111        let mut b = __s_b.launch_builder(&f);
8112        b.arg(k)
8113            .arg(v)
8114            .arg(&mut kf)
8115            .arg(&mut vf)
8116            .arg(&kv_dim_i)
8117            .arg(&kv_dim_i)
8118            .arg(&t_kv_i)
8119            .arg(&k_tok_bytes_i)
8120            .arg(&v_tok_bytes_i);
8121        unsafe { b.launch(cfg)? };
8122        self.sdpa_naive(
8123            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
8124        )
8125    }
8126
8127    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
8128    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
8129    /// Q/K/V/O [head_dim, n_head(_kv), T].
8130    pub fn fa_prefill(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8131                      o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8132                      t: usize, t_kv: usize, scale: f32, causal: bool)
8133                      -> Result<(), Box<dyn std::error::Error>> {
8134        if portable_mma_gated() {
8135            return self.sdpa_naive(q, k, v, o, head_dim, n_head, n_head_kv,
8136                                   t, t_kv, scale, causal);
8137        }
8138        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
8139        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
8140        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
8141        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
8142        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
8143        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
8144        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
8145        let fa3_on = head_dim == 256 && causal && t == t_kv
8146            && match std::env::var("MEMRA_FA3").as_deref() {
8147                Ok("0") => false,
8148                Ok("1") => true,
8149                _ => cfg!(memra_hopper_mma),
8150            };
8151        if fa3_on {
8152            let n = t * n_head * head_dim;
8153            let nkv = t * n_head_kv * head_dim;
8154            let mut q16 = self.alloc_u8_uninit(n * 2)?;
8155            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
8156            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
8157            self.f32_to_bf16_into(q, &mut q16, n)?;
8158            self.f32_to_bf16_into(k, &mut k16, nkv)?;
8159            self.f32_to_bf16_into(v, &mut v16, nkv)?;
8160            let rc = {
8161                use cudarc::driver::{DevicePtr, DevicePtrMut};
8162                let stream = self.gpu.stream();
8163                let (qp, _g1) = q16.device_ptr(&stream);
8164                let (kp, _g2) = k16.device_ptr(&stream);
8165                let (vp, _g3) = v16.device_ptr(&stream);
8166                let (op, _g4) = o.device_ptr_mut(&stream);
8167                unsafe {
8168                    memra_fa3_prefill(qp as *const core::ffi::c_void,
8169                                     kp as *const core::ffi::c_void,
8170                                     vp as *const core::ffi::c_void,
8171                                     op as *mut f32,
8172                                     t as i32, n_head as i32, n_head_kv as i32,
8173                                     head_dim as i32, scale,
8174                                     stream.cu_stream() as *mut core::ffi::c_void)
8175                }
8176            };
8177            if rc != 0 {
8178                return Err(format!("memra_fa3_prefill rc={rc}").into());
8179            }
8180            return Ok(());
8181        }
8182        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
8183        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
8184        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
8185        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
8186        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8187        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
8188        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
8189            const BLOCK_Q: usize = 64; const BKX: usize = 32;
8190            let f = self.func("fa_prefill_bf16_p1");
8191            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
8192                       + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
8193            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8194            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8195            let cfg = LaunchConfig {
8196                grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8197                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8198            };
8199            let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32,
8200                n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8201            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8202            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8203            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8204            let __s_b = self.gpu.stream();
8205            let mut b = __s_b.launch_builder(&f);
8206            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti)
8207             .arg(&tkvi).arg(&scale).arg(&cz);
8208            unsafe { b.launch(cfg)?; }
8209            return Ok(());
8210        }
8211        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
8212        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
8213        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
8214        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
8215        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
8216        const BK: usize = 32;
8217        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
8218        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
8219        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
8220        let (block_q, warps, w2_sfx): (usize, u32, &str) =
8221            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
8222        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
8223        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
8224        // other head_dims to sdpa_naive before reaching here.
8225        let hd_sfx = fa_hd_suffix(head_dim)?;
8226        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
8227        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
8228        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
8229        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
8230        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
8231        let bf16kv = !floor && !w2
8232            && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
8233        let (kb16, vb16) = if bf16kv {
8234            let n = t_kv * n_head_kv * head_dim;
8235            let mut kb = self.alloc_u8_uninit(n * 2)?;
8236            let mut vb = self.alloc_u8_uninit(n * 2)?;
8237            let fcv = self.func("f32_to_bf16_bulk");
8238            let ni = n as i64;
8239            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
8240            let __s_b = self.gpu.stream();
8241            let mut b = __s_b.launch_builder(&fcv);
8242            b.arg(k).arg(&mut kb).arg(&ni);
8243            unsafe { b.launch(cfgc)?; }
8244            let __s_b = self.gpu.stream();
8245            let mut b = __s_b.launch_builder(&fcv);
8246            b.arg(v).arg(&mut vb).arg(&ni);
8247            unsafe { b.launch(cfgc)?; }
8248            (Some(kb), Some(vb))
8249        } else {
8250            (None, None)
8251        };
8252        let f = self.func(&if bf16kv {
8253            format!("fa_prefill_bf16kv_pp{hd_sfx}")
8254        } else {
8255            format!("fa_prefill_f32{}{}{hd_sfx}",
8256                    if floor { "" } else { "_pp" },
8257                    if floor { "" } else { w2_sfx })
8258        });
8259        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
8260        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
8261        let kv_stages = if bf16kv { 2 } else { 1 };
8262        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
8263                   + 4 * (block_q * BK + 2 * block_q)) as u32;
8264        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8265        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8266        let cfg = LaunchConfig {
8267            grid_dim: ((t as u32 + block_q as u32 - 1) / block_q as u32, n_head as u32, 1),
8268            block_dim: (32, warps, 1), shared_mem_bytes: shmem,
8269        };
8270        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8271        let __s_b = self.gpu.stream();
8272        let mut b = __s_b.launch_builder(&f);
8273        b.arg(q);
8274        match (&kb16, &vb16) {
8275            (Some(kb), Some(vb)) => { b.arg(kb).arg(vb); }
8276            _ => { b.arg(k).arg(v); }
8277        }
8278        b.arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8279        unsafe { b.launch(cfg)?; }
8280        Ok(())
8281    }
8282
8283    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
8284    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
8285    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
8286    #[allow(clippy::too_many_arguments)]
8287    pub fn fa_prefill_w(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8288                        o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8289                        t: usize, t_kv: usize, scale: f32, causal: bool, window: usize)
8290                        -> Result<(), Box<dyn std::error::Error>> {
8291        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
8292        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
8293        if portable_mma_gated() {
8294            return self.sdpa_naive_w(q, k, v, o, head_dim, n_head, n_head_kv,
8295                                     t, t_kv, scale, causal, window);
8296        }
8297        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
8298        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
8299        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
8300        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8301        let faw_f32 = *FAW_F32.get_or_init(|| {
8302            std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32")
8303        });
8304        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
8305        self.fa_prefill_w_arm(q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
8306                              window, floor || faw_f32, floor)
8307    }
8308
8309    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
8310    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
8311    #[allow(clippy::too_many_arguments)]
8312    pub fn fa_prefill_w_pre(&self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>,
8313                            o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8314                            n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8315                            window: usize, v_f16: bool)
8316                            -> Result<(), Box<dyn std::error::Error>> {
8317        const BLOCK_Q: usize = 64; const BK: usize = 32;
8318        debug_assert_eq!(head_dim, 256);
8319        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0
8320            && (n_head / n_head_kv) % 2 == 0;
8321        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
8322        if hp {
8323            const BLOCK_QH: usize = 32;
8324            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
8325            // else re-encode through the pooled scratch (stream-ordered reuse).
8326            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
8327            let vh: &CudaSlice<u8> = if v_f16 { vb } else {
8328                let n = t_kv * n_head_kv * head_dim;
8329                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
8330                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
8331                }
8332                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
8333                vguard.as_ref().unwrap()
8334            };
8335            let f = self.func("fa_prefill_w_bf16_p1h2");
8336            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK)
8337                       + 4 * (2 * BLOCK_QH)) as u32;
8338            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8339            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8340            let cfg = LaunchConfig {
8341                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
8342                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8343            };
8344            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8345                n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8346            let __s_b = self.gpu.stream();
8347            let mut b = __s_b.launch_builder(&f);
8348            b.arg(qb).arg(kb).arg(vh).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8349             .arg(&scale).arg(&cz).arg(&wi);
8350            unsafe { b.launch(cfg)?; }
8351            return Ok(());
8352        }
8353        let f = self.func("fa_prefill_w_bf16_p1");
8354        let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8355                   + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8356        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8357        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8358        let cfg = LaunchConfig {
8359            grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8360            block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8361        };
8362        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8363            n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8364        let __s_b = self.gpu.stream();
8365        let mut b = __s_b.launch_builder(&f);
8366        b.arg(qb).arg(kb).arg(vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8367         .arg(&scale).arg(&cz).arg(&wi);
8368        unsafe { b.launch(cfg)?; }
8369        Ok(())
8370    }
8371
8372    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
8373    #[allow(clippy::too_many_arguments)]
8374    pub fn fa_prefill_w_arm(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8375                            o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8376                            n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8377                            window: usize, f32_stage: bool, floor: bool)
8378                            -> Result<(), Box<dyn std::error::Error>> {
8379        const BLOCK_Q: usize = 64; const BK: usize = 32;
8380        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
8381        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
8382        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
8383        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
8384        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8385        let p1 = !floor && !f32_stage
8386            && *P1_ON.get_or_init(|| {
8387                std::env::var("MEMRA_FAW_P1").map(|v| v != "0").unwrap_or(true)
8388            });
8389        let hp = p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0
8390            && (n_head / n_head_kv) % 2 == 0;
8391        if hp {
8392            const BLOCK_QH: usize = 32;
8393            let f = self.func("fa_prefill_w_bf16_p1h2");
8394            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK)
8395                       + 4 * (2 * BLOCK_QH)) as u32;
8396            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8397            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8398            let cfg = LaunchConfig {
8399                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
8400                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8401            };
8402            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8403                n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8404            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8405            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8406            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
8407            let __s_b = self.gpu.stream();
8408            let mut b = __s_b.launch_builder(&f);
8409            b.arg(&qb).arg(&kb).arg(&vh).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8410             .arg(&scale).arg(&cz).arg(&wi);
8411            unsafe { b.launch(cfg)?; }
8412            return Ok(());
8413        }
8414        if p1 {
8415            let f = self.func("fa_prefill_w_bf16_p1");
8416            let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8417                       + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8418            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8419            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8420            let cfg = LaunchConfig {
8421                grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8422                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8423            };
8424            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8425                n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8426            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8427            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8428            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8429            let __s_b = self.gpu.stream();
8430            let mut b = __s_b.launch_builder(&f);
8431            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8432             .arg(&scale).arg(&cz).arg(&wi);
8433            unsafe { b.launch(cfg)?; }
8434            return Ok(());
8435        }
8436        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
8437        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
8438        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8439        let g4 = !floor && !f32_stage && n_head_kv == 1 && n_head % 4 == 0
8440            && *G4_ON.get_or_init(|| {
8441                std::env::var("MEMRA_FAW_G4").map(|v| v != "0").unwrap_or(true)
8442            });
8443        if g4 {
8444            const SP_M: usize = 16;
8445            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
8446            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
8447            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8448            let o2 = *O2_ON.get_or_init(|| {
8449                std::env::var("MEMRA_FAW_O2").map(|v| v != "0").unwrap_or(true)
8450            });
8451            let f = self.func(if o2 { "fa_prefill_w_bf16_g4o2" } else { "fa_prefill_w_bf16_g4" });
8452            let shmem = if o2 {
8453                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
8454            } else {
8455                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK)
8456                    + 4 * (4 * SP_M)) as u32
8457            };
8458            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8459            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8460            let cfg = LaunchConfig {
8461                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
8462                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8463            };
8464            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8465                n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8466            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8467            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8468            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8469            let __s_b = self.gpu.stream();
8470            let mut b = __s_b.launch_builder(&f);
8471            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8472             .arg(&scale).arg(&cz).arg(&wi);
8473            unsafe { b.launch(cfg)?; }
8474            return Ok(());
8475        }
8476        let f = self.func(if floor { "fa_prefill_w_f32" }
8477                          else if f32_stage { "fa_prefill_w_f32_pp" }
8478                          else { "fa_prefill_w_bf16_pp" });
8479        let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8480                   + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8481        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8482        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8483        let cfg = LaunchConfig {
8484            grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8485            block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8486        };
8487        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8488                                                t as i32, t_kv as i32, causal as i32, window as i32);
8489        if f32_stage {
8490            let __s_b = self.gpu.stream();
8491            let mut b = __s_b.launch_builder(&f);
8492            b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8493             .arg(&scale).arg(&cz).arg(&wi);
8494            unsafe { b.launch(cfg)?; }
8495        } else {
8496            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8497            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8498            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8499            let __s_b = self.gpu.stream();
8500            let mut b = __s_b.launch_builder(&f);
8501            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8502             .arg(&scale).arg(&cz).arg(&wi);
8503            unsafe { b.launch(cfg)?; }
8504        }
8505        Ok(())
8506    }
8507
8508    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
8509    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
8510    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
8511    #[allow(clippy::too_many_arguments)]
8512    pub fn fa_prefill_hd512(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8513                            o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8514                            n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool)
8515                            -> Result<(), Box<dyn std::error::Error>> {
8516        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
8517        if portable_mma_gated() {
8518            return self.sdpa_naive(q, k, v, o, head_dim, n_head, n_head_kv,
8519                                   t, t_kv, scale, causal);
8520        }
8521        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
8522        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
8523        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
8524        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
8525        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
8526        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8527        let f32_stage = *F32_STAGE.get_or_init(|| {
8528            std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32")
8529        });
8530        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
8531        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
8532        // Own numeric config (partial-sum order) — battery-gated.
8533        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8534        let sp = !f32_stage
8535            && *SP_ON.get_or_init(|| {
8536                std::env::var("MEMRA_FA512_SP").map(|v| v != "0").unwrap_or(true)
8537            });
8538        self.fa_prefill_hd512_arm(q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale,
8539                                  causal, f32_stage, sp, sp && fa_f16pv_on())
8540    }
8541
8542    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
8543    #[allow(clippy::too_many_arguments)]
8544    pub fn fa_prefill_hd512_pre(&self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>,
8545                                o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8546                                n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8547                                v_f16: bool)
8548                                -> Result<(), Box<dyn std::error::Error>> {
8549        debug_assert_eq!(head_dim, 512);
8550        const SP_M: usize = 16; const BKS: usize = 32;
8551        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
8552        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
8553        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
8554        let f16pv = fa_f16pv_on();
8555        let nw = if f16pv { fa512_wide_warps() } else { 2 };
8556        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
8557        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
8558        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
8559        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
8560            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
8561            let n = t_kv * n_head_kv * head_dim;
8562            let need = n * 2;
8563            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
8564                *vguard = Some(self.alloc_uninit::<u8>(need)?);
8565            }
8566            let dst = vguard.as_mut().unwrap();
8567            self.bf16_to_f16_into(vb, n, dst)?;
8568            vguard.as_ref().unwrap()
8569        } else { vb };
8570        let f = self.func(if hp { "fa_prefill_bf16_hd512_sp16h2" }
8571                          else { match (f16pv, nw) {
8572                              (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
8573                              (true, _) => "fa_prefill_bf16_hd512_sp16",
8574                              _ => "fa_prefill_bf16_hd512_sp",
8575                          } });
8576        let (nwarp, npart) = if hp { (4usize, 4usize) } else if nw > 2 { (nw, nw) } else { (2, 1) };
8577        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
8578        let shmem = if hp {
8579            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
8580               + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
8581        } else {
8582            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
8583               + 4 * (npart * SP_M * BKS + SP_M)) as u32
8584        };
8585        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8586        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8587        let grid_y = if hp { (n_head / 2) as u32 } else { n_head as u32 };
8588        let cfg = LaunchConfig {
8589            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
8590            block_dim: (32, nwarp as u32, 1), shared_mem_bytes: shmem,
8591        };
8592        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8593                                            t as i32, t_kv as i32, causal as i32);
8594        let __s_b = self.gpu.stream();
8595        let mut b = __s_b.launch_builder(&f);
8596        b.arg(qb).arg(kb).arg(vref).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8597         .arg(&scale).arg(&cz);
8598        unsafe { b.launch(cfg)?; }
8599        Ok(())
8600    }
8601
8602    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
8603    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
8604    #[allow(clippy::too_many_arguments)]
8605    pub fn fa_prefill_hd512_arm(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8606                                o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8607                                n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8608                                f32_stage: bool, sp: bool, f16pv: bool)
8609                                -> Result<(), Box<dyn std::error::Error>> {
8610        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
8611        if sp && !f32_stage {
8612            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
8613            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
8614            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
8615            const SP_M: usize = 16; const BKS: usize = 32;
8616            let nw = if f16pv { fa512_wide_warps() } else { 2 };
8617            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
8618            let f = self.func(if hp { "fa_prefill_bf16_hd512_sp16h2" }
8619                              else { match (f16pv, nw) {
8620                                  (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
8621                                  (true, _) => "fa_prefill_bf16_hd512_sp16",
8622                                  _ => "fa_prefill_bf16_hd512_sp",
8623                              } });
8624            let (nwarp, npart) = if hp { (4usize, 4usize) } else if nw > 2 { (nw, nw) } else { (2, 1) };
8625            let shmem = if hp {
8626                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
8627                   + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
8628            } else {
8629                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
8630                   + 4 * (npart * SP_M * BKS + SP_M)) as u32
8631            };
8632            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8633            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8634            let grid_y = if hp { (n_head / 2) as u32 } else { n_head as u32 };
8635            let cfg = LaunchConfig {
8636                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
8637                block_dim: (32, nwarp as u32, 1), shared_mem_bytes: shmem,
8638            };
8639            let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8640                                                t as i32, t_kv as i32, causal as i32);
8641            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8642            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8643            let vb = if f16pv { self.f32_to_f16(v, t_kv * n_head_kv * head_dim)? }
8644                     else { self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)? };
8645            let __s_b = self.gpu.stream();
8646            let mut b = __s_b.launch_builder(&f);
8647            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8648             .arg(&scale).arg(&cz);
8649            unsafe { b.launch(cfg)?; }
8650            return Ok(());
8651        }
8652        const BLOCK_Q: usize = 32; const BK: usize = 32; const HALF: usize = 256;
8653        let f = self.func(if f32_stage { "fa_prefill_f32_hd512" } else { "fa_prefill_bf16_hd512" });
8654        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
8655        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
8656                   + 4 * BLOCK_Q) as u32;
8657        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8658        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8659        let cfg = LaunchConfig {
8660            grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 2),
8661            block_dim: (32, 2, 1), shared_mem_bytes: shmem,
8662        };
8663        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8664                                            t as i32, t_kv as i32, causal as i32);
8665        if f32_stage {
8666            let __s_b = self.gpu.stream();
8667            let mut b = __s_b.launch_builder(&f);
8668            b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8669             .arg(&scale).arg(&cz);
8670            unsafe { b.launch(cfg)?; }
8671        } else {
8672            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8673            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8674            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8675            let __s_b = self.gpu.stream();
8676            let mut b = __s_b.launch_builder(&f);
8677            b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8678             .arg(&scale).arg(&cz);
8679            unsafe { b.launch(cfg)?; }
8680        }
8681        Ok(())
8682    }
8683
8684    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
8685    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
8686    /// separate f32_to_bf16 the FA entries would run).
8687    #[allow(clippy::too_many_arguments)]
8688    pub fn rope_neox2_bf16e(&self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>,
8689                            qb: &mut CudaSlice<u8>, kb: &mut CudaSlice<u8>,
8690                            pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize,
8691                            nh_q: usize, nh_k: usize, n_tokens: usize, base: f32,
8692                            freq_scale: f32, ff: Option<&CudaSlice<f32>>)
8693                            -> Result<(), Box<dyn std::error::Error>> {
8694        let f = self.func("rope_neox2_bf16e_f32");
8695        let rows = ((nh_q + nh_k) * n_tokens) as u32;
8696        let cfg = LaunchConfig { grid_dim: (rows, 1, 1),
8697                                 block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
8698        let theta_scale = base.powf(-2.0 / n_dims as f32);
8699        let (hd, nd, nhq, nhk, nt) = (head_dim as i32, n_dims as i32, nh_q as i32,
8700                                      nh_k as i32, n_tokens as i32);
8701        let __s_b = self.gpu.stream();
8702        let mut b = __s_b.launch_builder(&f);
8703        match ff {
8704            Some(t) => { b.arg(&mut *q).arg(&mut *k).arg(&mut *qb).arg(&mut *kb).arg(pos)
8705                          .arg(&hd).arg(&nd).arg(&nhq).arg(&nhk).arg(&nt)
8706                          .arg(&theta_scale).arg(&freq_scale).arg(t);
8707                         unsafe { b.launch(cfg)?; } }
8708            None => { let null: u64 = 0;
8709                      b.arg(&mut *q).arg(&mut *k).arg(&mut *qb).arg(&mut *kb).arg(pos)
8710                       .arg(&hd).arg(&nd).arg(&nhq).arg(&nhk).arg(&nt)
8711                       .arg(&theta_scale).arg(&freq_scale).arg(&null);
8712                      unsafe { b.launch(cfg)?; } }
8713        }
8714        Ok(())
8715    }
8716
8717    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
8718    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
8719    pub fn f32_to_bf16(&self, x: &CudaSlice<f32>, n: usize)
8720                       -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8721        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
8722        let mut y = self.alloc_uninit::<u8>(n * 2)?;
8723        let f = self.func("f32_to_bf16_flat");
8724        let n_i = n as i64;
8725        let cfg = LaunchConfig {
8726            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
8727            block_dim: (256, 1, 1), shared_mem_bytes: 0,
8728        };
8729        let __s_b = self.gpu.stream();
8730        let mut b = __s_b.launch_builder(&f);
8731        b.arg(x).arg(&mut y).arg(&n_i);
8732        unsafe { b.launch(cfg)?; }
8733        Ok(y)
8734    }
8735
8736    pub fn f32_to_f16(&self, x: &CudaSlice<f32>, n: usize)
8737                      -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8738        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
8739        let mut y = self.alloc_uninit::<u8>(n * 2)?;
8740        let f = self.func("f32_to_f16_flat");
8741        let n_i = n as i64;
8742        let cfg = LaunchConfig {
8743            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
8744            block_dim: (256, 1, 1), shared_mem_bytes: 0,
8745        };
8746        let __s_b = self.gpu.stream();
8747        let mut b = __s_b.launch_builder(&f);
8748        b.arg(x).arg(&mut y).arg(&n_i);
8749        unsafe { b.launch(cfg)?; }
8750        Ok(y)
8751    }
8752
8753    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
8754    pub fn bf16_to_f16(&self, xb: &CudaSlice<u8>, n: usize)
8755                       -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8756        let mut y = self.alloc_uninit::<u8>(n * 2)?;
8757        self.bf16_to_f16_into(xb, n, &mut y)?;
8758        Ok(y)
8759    }
8760
8761    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
8762    pub fn bf16_to_f16_into(&self, xb: &CudaSlice<u8>, n: usize, y: &mut CudaSlice<u8>)
8763                            -> Result<(), Box<dyn std::error::Error>> {
8764        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
8765        assert!(y.len() >= n * 2);
8766        let f = self.func("bf16_to_f16_flat");
8767        let n2 = (n / 2) as i64;
8768        let cfg = LaunchConfig {
8769            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
8770            block_dim: (256, 1, 1), shared_mem_bytes: 0,
8771        };
8772        let __s_b = self.gpu.stream();
8773        let mut b = __s_b.launch_builder(&f);
8774        b.arg(xb).arg(y).arg(&n2);
8775        unsafe { b.launch(cfg)?; }
8776        Ok(())
8777    }
8778
8779    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
8780    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
8781    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
8782    /// head_dim in {256, 128}, bf16kv lane on.
8783    #[allow(clippy::too_many_arguments)]
8784    pub fn fa_prefill_vl8(&self, seqs: &[FaSeqVl], head_dim: usize, n_head: usize,
8785                          n_head_kv: usize, scale: f32)
8786                          -> Result<(), Box<dyn std::error::Error>> {
8787        const BK: usize = 32;
8788        let b = seqs.len();
8789        assert!(b >= 1 && b <= 8);
8790        let mut packed = [FaSeqVl::default(); 8];
8791        packed[..b].copy_from_slice(seqs);
8792        let v = FaVl8(packed);
8793        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
8794        let ept = (n_head_kv * head_dim) as i32;
8795        {
8796            let f = self.func("fa_mirror_vl");
8797            let max_n = (max_t as i64) * ept as i64;
8798            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
8799            for which in 0..2i32 {
8800                let cfg = LaunchConfig { grid_dim: (blocks, 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8801                let __s_lb = self.gpu.stream();
8802                let mut lb = __s_lb.launch_builder(&f);
8803                lb.arg(&v).arg(&ept).arg(&which);
8804                unsafe { lb.launch(cfg)?; }
8805            }
8806        }
8807        let hd_sfx = fa_hd_suffix(head_dim)?;
8808        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
8809        let block_q = 64usize;
8810        let kv_stages = 2usize;
8811        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
8812                   + 4 * (block_q * BK + 2 * block_q)) as u32;
8813        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8814        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8815        let cfg = LaunchConfig {
8816            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
8817            block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8818        };
8819        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
8820        let __s_lb = self.gpu.stream();
8821        let mut lb = __s_lb.launch_builder(&f);
8822        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
8823        unsafe { lb.launch(cfg)?; }
8824        Ok(())
8825    }
8826
8827    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
8828    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
8829    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
8830    #[allow(clippy::too_many_arguments)]
8831    pub fn attn_pre_vl8(&self, seqs: &[AttnPreVl], wq: &CudaSlice<f32>, wk: &CudaSlice<f32>,
8832                        head_dim: usize, rope_dims: usize, n_head: usize, n_head_kv: usize,
8833                        eps: f32, freq_base: f32, freq_scale: f32,
8834                        kv_dim_k: usize, kv_dim_v: usize,
8835                        k_tok_bytes: usize, v_tok_bytes: usize)
8836                        -> Result<(), Box<dyn std::error::Error>> {
8837        let b = seqs.len();
8838        assert!(b >= 1 && b <= 8);
8839        let mut packed = [AttnPreVl::default(); 8];
8840        packed[..b].copy_from_slice(seqs);
8841        let v = AttnPreVl8(packed);
8842        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
8843        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
8844        {
8845            let f = self.func("q_gate_split_vl");
8846            let n = max_t * (n_head * head_dim) as u32;
8847            let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8848            let __s_lb = self.gpu.stream();
8849            let mut lb = __s_lb.launch_builder(&f);
8850            lb.arg(&v).arg(&hd).arg(&nh);
8851            unsafe { lb.launch(cfg)?; }
8852        }
8853        {
8854            let f = self.func("attn_rms_vl");
8855            let cfg = LaunchConfig { grid_dim: (max_t * n_head as u32, 2, b as u32), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
8856            let __s_lb = self.gpu.stream();
8857            let mut lb = __s_lb.launch_builder(&f);
8858            lb.arg(&v).arg(wq).arg(wk).arg(&hd).arg(&nh).arg(&nhkv).arg(&eps);
8859            unsafe { lb.launch(cfg)?; }
8860        }
8861        {
8862            let f = self.func("attn_rope_vl");
8863            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
8864            let nd = rope_dims as i32;
8865            let cfg = LaunchConfig { grid_dim: (max_t * n_head as u32, 2, b as u32), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
8866            let __s_lb = self.gpu.stream();
8867            let mut lb = __s_lb.launch_builder(&f);
8868            lb.arg(&v).arg(&hd).arg(&nd).arg(&nh).arg(&nhkv).arg(&theta_scale).arg(&freq_scale);
8869            unsafe { lb.launch(cfg)?; }
8870        }
8871        {
8872            let f = self.func("append_kv_vl");
8873            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
8874            let cfg = LaunchConfig { grid_dim: (nblk, max_t, b as u32), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
8875            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
8876            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8877            let __s_lb = self.gpu.stream();
8878            let mut lb = __s_lb.launch_builder(&f);
8879            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
8880            unsafe { lb.launch(cfg)?; }
8881        }
8882        Ok(())
8883    }
8884
8885    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
8886    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
8887    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
8888    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
8889    pub fn fa_prefill_view(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
8890                           v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
8891                           head_dim: usize, n_head: usize, n_head_kv: usize,
8892                           t: usize, t_kv: usize, scale: f32, causal: bool,
8893                           k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8894                           -> Result<(), Box<dyn std::error::Error>> {
8895        if portable_mma_gated() {
8896            return self.sdpa_naive_quantized_view(q, k, v, o, head_dim, n_head, n_head_kv,
8897                                                  t, t_kv, scale, causal,
8898                                                  k_tok_bytes, v_tok_bytes);
8899        }
8900        const BLOCK_Q: usize = 64; const BK: usize = 32;
8901        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
8902        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
8903        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
8904        let f = if g { self.func_g(&name) } else { self.func(&name) };
8905        let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8906                   + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8907        use cudarc::driver::sys::CUfunction_attribute_enum as A;
8908        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8909        let cfg = LaunchConfig {
8910            grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8911            block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8912        };
8913        let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8914        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8915        let __s_b = self.gpu.stream();
8916        let mut b = __s_b.launch_builder(&f);
8917        b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz)
8918         .arg(&ktb).arg(&vtb);
8919        unsafe { b.launch(cfg)?; }
8920        Ok(())
8921    }
8922
8923    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
8924    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
8925    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
8926    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
8927    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
8928    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
8929    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
8930    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
8931    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
8932    #[allow(clippy::too_many_arguments)]
8933    pub fn fa_prefill_view_ws(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
8934                              v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
8935                              head_dim: usize, n_head: usize, n_head_kv: usize,
8936                              t: usize, t_kv: usize, scale: f32, causal: bool,
8937                              k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8938                              -> Result<(), Box<dyn std::error::Error>> {
8939        if portable_mma_gated() {
8940            return self.sdpa_naive_quantized_view(q, k, v, o, head_dim, n_head, n_head_kv,
8941                                                  t, t_kv, scale, causal,
8942                                                  k_tok_bytes, v_tok_bytes);
8943        }
8944        const BLOCK_Q: usize = 64; const BK: usize = 32;
8945        let kv_dim_k = n_head_kv * head_dim;
8946        let kv_dim_v = n_head_kv * head_dim;
8947        let k_ws_bytes = t_kv * kv_dim_k * 2;   // bf16
8948        let v_ws_bytes = t_kv * kv_dim_v * 2;
8949        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
8950        let mut guard = self.prime_deqw_ws.lock().unwrap();
8951        let need_grow = match guard.as_ref() {
8952            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
8953            None => true,
8954        };
8955        if need_grow {
8956            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
8957            let (ck, cv) = guard.as_ref().map(|(a, b)| (a.len(), b.len())).unwrap_or((0, 0));
8958            *guard = Some((self.alloc_u8(grow(ck, k_ws_bytes))?, self.alloc_u8(grow(cv, v_ws_bytes))?));
8959        }
8960        let (kw, vw) = guard.as_mut().unwrap();
8961        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
8962        {
8963            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
8964            let f = if g { self.func_g("fa_dequant_kv_ws_bf16") } else { self.func("fa_dequant_kv_ws_bf16") };
8965            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
8966            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8967            let cfg = LaunchConfig { grid_dim: (nblk.max(1), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8968            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
8969            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8970            let __s_b = self.gpu.stream();
8971            let mut b = __s_b.launch_builder(&f);
8972            b.arg(k).arg(v).arg(&mut *kw).arg(&mut *vw).arg(&kdk).arg(&kdv).arg(&tkvi).arg(&ktb).arg(&vtb);
8973            unsafe { b.launch(cfg)?; }
8974        }
8975        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
8976        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
8977        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
8978        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
8979        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
8980        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
8981        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
8982        let db = std::env::var("MEMRA_PRIME_DEQW_DB").map(|v| v != "0").unwrap_or(true);
8983        {
8984            let hd_sfx = fa_hd_suffix(head_dim)?;
8985            let f = self.func(&format!("fa_prefill_qw{}{hd_sfx}", if db { "_db" } else { "" }));
8986            let shmem = if db {
8987                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
8988                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
8989            } else {
8990                (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8991                       + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
8992            };
8993            use cudarc::driver::sys::CUfunction_attribute_enum as A;
8994            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8995            let cfg = LaunchConfig {
8996                grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8997                block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8998            };
8999            let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
9000            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
9001            let __s_b = self.gpu.stream();
9002            let mut b = __s_b.launch_builder(&f);
9003            b.arg(q).arg(&*kw).arg(&*vw).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz)
9004             .arg(&kdk).arg(&kdv);
9005            unsafe { b.launch(cfg)?; }
9006        }
9007        Ok(())
9008    }
9009
9010    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
9011    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
9012    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
9013    pub fn fa_decode(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9014                     v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9015                     head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32,
9016                     k_tok_bytes: usize, v_tok_bytes: usize)
9017                     -> Result<(), Box<dyn std::error::Error>> {
9018        self.fa_decode_kvmod(q, k, v, o, head_dim, n_head, n_head_kv, t_kv, scale,
9019                             k_tok_bytes, v_tok_bytes, false)
9020    }
9021
9022    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
9023    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
9024    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
9025    #[allow(clippy::too_many_arguments)]
9026    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
9027    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
9028    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
9029    #[allow(clippy::too_many_arguments)]
9030    #[allow(clippy::too_many_arguments)]
9031    fn fa_decode_scalar_unified(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9032                                v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9033                                head_dim: usize, n_head: usize, n_head_kv: usize,
9034                                t_kv_host: usize, t_kv_dev: Option<&CudaSlice<i32>>,
9035                                scale: f32, n_splits: usize, split_keys: usize,
9036                                k_tok_bytes: usize, v_tok_bytes: usize, g: bool,
9037                                part_o: &mut CudaSlice<f32>, part_m: &mut CudaSlice<f32>,
9038                                part_l: &mut CudaSlice<f32>,
9039                                q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9040                                -> Result<(), Box<dyn std::error::Error>> {
9041        let f = if g { self.func_g("fa_decode_f32") } else { self.fa_func("fa_decode_f32", head_dim) };
9042        let cfg = LaunchConfig { grid_dim: (n_head as u32, n_splits as u32, 1),
9043            block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: (4 * (head_dim + 32)) as u32 };
9044        let (hd, nh, nhkv, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, n_splits as i32);
9045        let (ktb, vtb, tkvi, ski) = (k_tok_bytes as i64, v_tok_bytes as i64, t_kv_host as i32,
9046                                     split_keys as i32);
9047        let __s_b = self.gpu.stream();
9048        let mut b = __s_b.launch_builder(&f);
9049        match t_kv_dev {
9050            Some(d) => { b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9051                          .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(d).arg(&scale).arg(&nsp)
9052                          .arg(&ski).arg(&ktb).arg(&vtb);
9053                         unsafe { b.launch(cfg)?; } }
9054            None => { let null: u64 = 0;
9055                      b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9056                       .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(&null).arg(&scale).arg(&nsp)
9057                       .arg(&ski).arg(&ktb).arg(&vtb);
9058                      unsafe { b.launch(cfg)?; } }
9059        }
9060        let cfg2 = LaunchConfig { grid_dim: (n_head as u32, 1, 1),
9061            block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9062        if let Some((oq, od)) = q8_out {
9063            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
9064            let fc = if g { self.func_g("fa_decode_combine_q8_1") }
9065                     else { self.fa_func("fa_decode_combine_q8_1", head_dim) };
9066            let __s_b2 = self.gpu.stream();
9067            let mut b2 = __s_b2.launch_builder(&fc);
9068            b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh).arg(&nsp);
9069            unsafe { b2.launch(cfg2)?; }
9070            return Ok(());
9071        }
9072        let fc = if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) };
9073        let __s_b2 = self.gpu.stream();
9074        let mut b2 = __s_b2.launch_builder(&fc);
9075        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
9076        unsafe { b2.launch(cfg2)?; }
9077        Ok(())
9078    }
9079
9080    pub fn fa_decode_kvmod(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9081                     v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9082                     head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32,
9083                     k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
9084                     -> Result<(), Box<dyn std::error::Error>> {
9085        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
9086        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
9087        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
9088        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
9089        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
9090        //
9091        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
9092        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
9093        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
9094        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
9095        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
9096        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
9097        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
9098        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
9099        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
9100        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
9101        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
9102        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
9103        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
9104        // fall to the exact scalar there instead of the broken register arm.
9105        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
9106        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
9107        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
9108        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
9109        if g && head_dim == 256 && !fa_v4_at(t_kv) { fa_vec = false; }
9110        let sp = fa_split_keys(t_kv, n_head_kv);
9111        let n_splits = if fa_vec { ((t_kv + sp - 1) / sp).max(1) } else { ((t_kv + 255) / 256).max(1) };
9112        let o_len = n_head * n_splits * head_dim;
9113        let ml_len = n_head * n_splits;
9114        let mut part_guard = self.fa_part_pool.lock().unwrap();
9115        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9116            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9117            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9118            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9119            // later live allocations land at those addresses, and the next graph REPLAY writes
9120            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9121            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9122            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9123            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9124            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
9125            // (total retired < final size).
9126            let old = part_guard.take();
9127            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9128            if let Some(old) = old {
9129                self.fa_part_retired.lock().unwrap().push(old);
9130            }
9131            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9132                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9133            }
9134            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9135                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9136                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9137        }
9138        let pg = part_guard.as_mut().unwrap();
9139        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9140        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9141        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9142        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9143        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
9144        let (hd, nh, nhkv, tkvi, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, t_kv as i32, n_splits as i32);
9145        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9146        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
9147        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
9148        // silently truncating the accumulator.
9149        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
9150        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
9151        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
9152        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
9153        // 178.4 -> 173.7 when 512 rode vec unconditionally).
9154        let fa512_min = fa512_min_tkv();
9155        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
9156        // g-module keeps the v4 pick (its class is not the depth-decay class).
9157        let deep = fa_vec && head_dim == 256 && fa_v4_at(t_kv) && !g
9158            && fa_deep_at(t_kv) && !matches!(fa_v4_mode(), "noB3" | "stage");
9159        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
9160            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
9161            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
9162            let gqa = (n_head / n_head_kv).max(1) as u32;
9163            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
9164            (fv, LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9165                 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
9166        } else if fa_vec && head_dim <= 256 {
9167            let gqa = (n_head / n_head_kv).max(1) as u32;
9168            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
9169            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
9170            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
9171            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
9172            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
9173            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
9174            // dequant each tile ONCE per block.
9175            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
9176            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
9177            // there by 12x — latency, not bandwidth, rules small KV).
9178            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9179            let smem_tkv = *SMEM_TKV.get_or_init(|| {
9180                std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9181                    .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9182            });
9183            if fa_v4_at(t_kv) && head_dim == 256 {
9184                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
9185                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
9186                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
9187                let v4name = match fa_v4_mode() {
9188                    "noB3" => "fa_decode_vec_q_v4_noB3",     // phase probe (WRONG OUTPUT)
9189                    "stage" => "fa_decode_vec_q_v4_stage",   // phase probe (WRONG OUTPUT)
9190                    _ if deep => "fa_decode_vec_q_v4_deep",
9191                    _ => "fa_decode_vec_q_v4",
9192                };
9193                let fv = if g { self.func_g(v4name) } else { self.func(v4name) };
9194                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
9195                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
9196                let shmem = (if deep { 12160 } else { 11520 }
9197                             + 32 * head_dim * if g { 1 } else { 2 }) as u32;
9198                use cudarc::driver::sys::CUfunction_attribute_enum as A;
9199                fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9200                (fv,
9201                 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9202                     block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9203            } else if fa_v3_active(head_dim) {
9204                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
9205                // smem = sV only (half of v2's).
9206                let fv = if g { self.func_g("fa_decode_vec_q_v3") } else { self.func("fa_decode_vec_q_v3") };
9207                let shmem = (32 * head_dim * 2) as u32;      // sV bf16 [FA_DEC_TILE=32][hd]
9208                (fv,
9209                 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9210                     block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9211            } else if fa_v2_on() {
9212                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
9213                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
9214                // partials; same 32KB sK+sV tile as the smem twin.
9215                let fv = if g { self.func_g("fa_decode_vec_q_v2") } else { self.func("fa_decode_vec_q_v2") };
9216                let shmem = (2 * 32 * head_dim * 2) as u32;   // sK+sV bf16 [FA_DEC_TILE=32][hd]
9217                (fv,
9218                 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9219                     block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9220            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g
9221                && !(head_dim == 512 && Self::gkv_on()) {
9222                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
9223                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
9224                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
9225                let fv = if g { self.func_g("fa_decode_vec_q_smem") } else { self.func("fa_decode_vec_q_smem") };
9226                let shmem = (2 * 32 * head_dim * 2) as u32;   // sK+sV bf16 [FA_DEC_TILE=32][hd]
9227                use cudarc::driver::sys::CUfunction_attribute_enum as A;
9228                fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9229                (fv,
9230                 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9231                     block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9232            } else {
9233                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
9234                // dequant, zero dynamic shared memory.
9235                let fv = if g { self.func_g("fa_decode_vec_q") } else { self.func("fa_decode_vec_q") };
9236                (fv,
9237                 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9238                     block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
9239            }
9240        } else {
9241            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
9242            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
9243            return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
9244                                                 t_kv, None, scale, n_splits,
9245                                                 if fa_vec { sp } else { 256 },
9246                                                 k_tok_bytes, v_tok_bytes, g,
9247                                                 part_o, part_m, part_l, None);
9248        };
9249        let __s_b = self.gpu.stream();
9250        let mut b = __s_b.launch_builder(&f);
9251        b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9252         .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(&scale).arg(&nsp).arg(&ktb).arg(&vtb);
9253        unsafe { b.launch(cfg)?; }
9254        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
9255        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
9256        let (fc, cfg2) = (if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) },
9257            LaunchConfig { grid_dim: (n_head as u32, 1, 1), block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 });
9258        let __s_b2 = self.gpu.stream();
9259        let mut b2 = __s_b2.launch_builder(&fc);
9260        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
9261        unsafe { b2.launch(cfg2)?; }
9262        Ok(())
9263    }
9264
9265    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
9266    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
9267    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
9268    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
9269    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
9270    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
9271    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
9272    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
9273    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
9274    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
9275    #[allow(clippy::too_many_arguments)]
9276    pub fn fa_decode_batch_seqs_v4(&self, q: &CudaSlice<f32>,
9277                                   kv_ptrs: &cudarc::driver::CudaView<u64>,
9278                                   pos_seq: &CudaSlice<i32>, o: &mut CudaSlice<f32>,
9279                                   head_dim: usize, n_head: usize, n_head_kv: usize,
9280                                   b_n: usize, t_kv_max: usize, scale: f32,
9281                                   split_keys: usize, k_tok_bytes: usize, v_tok_bytes: usize)
9282                                   -> Result<(), Box<dyn std::error::Error>> {
9283        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
9284        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
9285        let o_len = b_n * n_head * n_splits_max * head_dim;
9286        let ml_len = b_n * n_head * n_splits_max;
9287        let mut part_guard = self.fa_part_pool.lock().unwrap();
9288        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9289            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9290            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9291            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9292            // later live allocations land at those addresses, and the next graph REPLAY writes
9293            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9294            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9295            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9296            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9297            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
9298            // (total retired < final size).
9299            let old = part_guard.take();
9300            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9301            if let Some(old) = old {
9302                self.fa_part_retired.lock().unwrap().push(old);
9303            }
9304            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9305                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9306            }
9307            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9308                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9309                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9310        }
9311        let pg = part_guard.as_mut().unwrap();
9312        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9313        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9314        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9315        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9316        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9317        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
9318        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9319        let gqa = (n_head / n_head_kv).max(1) as u32;
9320        let f = self.func("fa_decode_vec_q_seqs_v4");
9321        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
9322        let shmem = (11520 + 32 * head_dim * 2) as u32;
9323        use cudarc::driver::sys::CUfunction_attribute_enum as A;
9324        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9325        let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
9326            block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9327        {
9328            let __s_b = self.gpu.stream();
9329            let mut b = __s_b.launch_builder(&f);
9330            b.arg(q).arg(kv_ptrs).arg(pos_seq).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9331             .arg(&hd).arg(&nh).arg(&nhkv).arg(&scale).arg(&nspm).arg(&spk).arg(&ktb).arg(&vtb);
9332            unsafe { b.launch(cfg)?; }
9333        }
9334        let fc = self.func("fa_decode_combine_seqs");
9335        let cfg2 = LaunchConfig { grid_dim: (n_head as u32, b_n as u32, 1),
9336            block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9337        let __s_b2 = self.gpu.stream();
9338        let mut b2 = __s_b2.launch_builder(&fc);
9339        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9340          .arg(pos_seq).arg(&nspm).arg(&spk);
9341        unsafe { b2.launch(cfg2)?; }
9342        Ok(())
9343    }
9344
9345    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
9346    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
9347    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
9348    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
9349    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
9350    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
9351    #[allow(clippy::too_many_arguments)]
9352    pub fn append_kv_quantized_seqs(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
9353                                    kv_ptrs: &cudarc::driver::CudaView<u64>,
9354                                    pos_seq: &CudaSlice<i32>, b_n: usize,
9355                                    kv_dim_k: usize, kv_dim_v: usize,
9356                                    k_tok_bytes: usize, v_tok_bytes: usize)
9357                                    -> Result<(), Box<dyn std::error::Error>> {
9358        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
9359        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
9360        let cfg = LaunchConfig { grid_dim: (nblk, b_n as u32, 1),
9361            block_dim: (32, 1, 1), shared_mem_bytes: 0 };
9362        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
9363        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9364        let __s_b = self.gpu.stream();
9365        let mut b = __s_b.launch_builder(&f);
9366        b.arg(k_rows).arg(v_rows).arg(kv_ptrs).arg(pos_seq)
9367         .arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
9368        unsafe { b.launch(cfg)?; }
9369        Ok(())
9370    }
9371
9372    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
9373    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
9374    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
9375    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
9376    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
9377    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
9378        std::env::var("MEMRA_NO_FA_VEC").is_err()
9379            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
9380            && base_len + 1 >= fa_vec_min_tkv()
9381            && head_dim <= 256 && head_dim % 32 == 0
9382    }
9383
9384    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
9385    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
9386    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
9387    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
9388    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
9389    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
9390    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
9391    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
9392    #[allow(clippy::too_many_arguments)]
9393    pub fn fa_decode_rows(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9394                          v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9395                          head_dim: usize, n_head: usize, n_head_kv: usize,
9396                          base_len: usize, t: usize, scale: f32,
9397                          k_tok_bytes: usize, v_tok_bytes: usize,
9398                          // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
9399                          // kernel; host base_len keeps sizing the splits/partials. hd256 twins
9400                          // keep the host arg. None is a bug for hd512 (asserted below).
9401                          base_dev: Option<(&CudaSlice<i32>, i32)>,
9402                          // K and V planes hold the same values (gemma globals, wv:=wk): pick
9403                          // the _kv twin — V plane never read, value rides the q8_0 key dq.
9404                          kv_shared: bool,
9405                          // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
9406                          // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
9407                          // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
9408                          g: bool,
9409                          // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
9410                          // (hd512 path) — the standalone quantize launch folds away.
9411                          mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9412                          -> Result<(), Box<dyn std::error::Error>> {
9413        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
9414        let t_kv_max = base_len + t;                       // LAST row's key bound
9415        let mut sp = fa_split_keys(t_kv_max, n_head_kv);   // env/default — same value every row
9416        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
9417        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
9418        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
9419        // (parity law), so the partition is freely tunable — verify and decode move together.
9420        if head_dim == 512 {
9421            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9422            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
9423            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
9424            let v = *SP512.get_or_init(|| std::env::var("MEMRA_FA_SP512").ok()
9425                .and_then(|x| x.parse().ok()).unwrap_or(0));
9426            sp = if v >= 8 { v } else { FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed) };
9427        }
9428        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9429        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9430        let gqa = (n_head / n_head_kv).max(1) as u32;
9431        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
9432        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
9433        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
9434        // the different partition changes the combine's FP order (greedy tie flips at depth;
9435        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
9436        // consecutive rows by their OWN ladder value and launch once per group — each row then
9437        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
9438        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
9439        // sp override is t_kv-independent by construction).
9440        let mut groups: Vec<(usize, usize, usize)> = Vec::new();   // (row0, t_g, sp_g)
9441        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
9442            groups.push((0, t, sp));
9443        } else {
9444            let mut r0 = 0usize;
9445            while r0 < t {
9446                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
9447                let mut r1 = r0 + 1;
9448                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g { r1 += 1; }
9449                groups.push((r0, r1 - r0, sp_g));
9450                r0 = r1;
9451            }
9452        }
9453        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
9454        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
9455        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
9456        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9457        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
9458            std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9459                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9460        });
9461        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
9462        let v3 = fa_v3_active(head_dim);
9463        let smem_rows = head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
9464        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
9465        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
9466        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
9467        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
9468        let _ = kv_shared;
9469        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
9470        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
9471        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
9472        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
9473        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
9474        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
9475        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
9476        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
9477        // (kv_head, split) stages its tile once and loops the rows over it — kills the
9478        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
9479        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
9480        // shared by every hd512 caller through this wrapper (decode+verify flip together;
9481        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
9482        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
9483        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
9484        // not unpack-bound; jsonl 2026-07-14.
9485        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9486        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
9487        let tb512 = head_dim == 512 && sp <= 32 && n_head / n_head_kv.max(1) <= 16
9488            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
9489        let fname = if tb512 { "fa_decode_vec_q_rows_v4_512_tb" }
9490                    else if i2 { "fa_decode_vec_q_rows_dpl16_i2" }
9491                    else if head_dim == 512 { "fa_decode_vec_q_rows_dpl16" }   // gemma globals (parity law)
9492                    else if v4 { "fa_decode_vec_q_rows_v4" }
9493                    else if v3 { "fa_decode_vec_q_rows_v3" }
9494                    else if fa_v2_on() { "fa_decode_vec_q_rows_v2" }
9495                    else if smem_rows { "fa_decode_vec_q_rows_smem" }
9496                    else { "fa_decode_vec_q_rows" };
9497        let f = if head_dim == 512 { self.fa_func(fname, head_dim) }
9498                else if g {
9499                    // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
9500                    // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
9501                    // g-module rows against decode's g-module v4 — different programs, short-VG
9502                    // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
9503                    // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
9504                    // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
9505                    // dq macros are format-aware.
9506                    self.func_g(if smem_rows { "fa_decode_vec_q_rows" } else { fname })
9507                }
9508                else { self.func(fname) };
9509        let shmem = if tb512 {
9510            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
9511            let gk = Self::gkv_on();
9512            let sh = (8192 + 1024 + 32 * 512 + 32 * 64
9513                      + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
9514            use cudarc::driver::sys::CUfunction_attribute_enum as A;
9515            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9516            sh
9517        } else if v4 || v3 || smem_rows || fa_v2_on() {
9518            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
9519            let sh = (if v4 { 11520 + 32 * head_dim * if g { 1 } else { 2 } }
9520                      else if v3 { 32 * head_dim * 2 } else { 2 * 32 * head_dim * 2 }) as u32;
9521            use cudarc::driver::sys::CUfunction_attribute_enum as A;
9522            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9523            sh
9524        } else { 0 };
9525        // Per-GROUP launches (single group in the common case — identical to the pre-fix
9526        // single launch there): each group gets its own partials (the rows kernel indexes
9527        // partials by its LOCAL grid.z row) and q/o row-offset views.
9528        for &(r0, t_g, sp_g) in &groups {
9529            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
9530            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
9531            let base_i = (base_len + r0) as i32;
9532            let o_len = t_g * n_head * n_splits_g * head_dim;
9533            let ml_len = t_g * n_head * n_splits_g;
9534            let mut part_guard = self.fa_part_pool.lock().unwrap();
9535        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9536            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9537            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9538            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9539            // later live allocations land at those addresses, and the next graph REPLAY writes
9540            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9541            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9542            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9543            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9544            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
9545            // (total retired < final size).
9546            let old = part_guard.take();
9547            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9548            if let Some(old) = old {
9549                self.fa_part_retired.lock().unwrap().push(old);
9550            }
9551            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9552                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9553            }
9554            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9555                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9556                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9557        }
9558        let pg = part_guard.as_mut().unwrap();
9559        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9560        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9561        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9562        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9563            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
9564            let qv = self.view(q, t * n_head * head_dim);
9565            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
9566            let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
9567                block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9568            {
9569                let __s_b = self.gpu.stream();
9570                let mut b = __s_b.launch_builder(&f);
9571                if tb512 {
9572                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
9573                    let (bd, plus) = base_dev.expect("hd512 rows twin requires a device base counter");
9574                    let plus_g = plus + r0 as i32;
9575                    let nr = t_g as i32;
9576                    if Self::pdl_on() && Self::pdl_wb_on() {
9577                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
9578                        use cudarc::driver::{DevicePtr, DevicePtrMut};
9579                        let s = &self.gpu.stream();
9580                        let (pq, _b0) = q_g.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9581                        let (pv, _b2) = v.device_ptr(s);
9582                        let (po, _b3) = part_o.device_ptr_mut(s);
9583                        let (pm, _b4) = part_m.device_ptr_mut(s);
9584                        let (pl, _b5) = part_l.device_ptr_mut(s);
9585                        let (pb, _b6) = bd.device_ptr(s);
9586                        let mut ps = [
9587                            &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9588                            &pv as *const _ as *mut _, &po as *const _ as *mut _,
9589                            &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9590                            &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9591                            &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9592                            &plus_g as *const _ as *mut _, &scale as *const _ as *mut _,
9593                            &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9594                            &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9595                            &nr as *const _ as *mut _,
9596                        ];
9597                        unsafe { self.launch_pdl_flash(Self::gkv_on(),
9598                            "fa_decode_vec_q_rows_v4_512_tb",
9599                            (n_head_kv as u32, n_splits_g as u32, 1), (32, gqa, 1),
9600                            shmem, &mut ps)?; }
9601                    } else {
9602                    let cfg_tb = LaunchConfig {
9603                        grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
9604                        block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9605                    b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9606                     .arg(&hd).arg(&nh).arg(&nhkv).arg(bd).arg(&plus_g).arg(&scale).arg(&nspm).arg(&spk)
9607                     .arg(&ktb).arg(&vtb).arg(&nr);
9608                    unsafe { b.launch(cfg_tb)?; }
9609                    }
9610                } else if head_dim == 512 {
9611                    let (bd, plus) = base_dev.expect("hd512 rows twin requires a device base counter");
9612                    let plus_g = plus + r0 as i32;
9613                    b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9614                     .arg(&hd).arg(&nh).arg(&nhkv).arg(bd).arg(&plus_g).arg(&scale).arg(&nspm).arg(&spk)
9615                     .arg(&ktb).arg(&vtb);
9616                    unsafe { b.launch(cfg)?; }
9617                } else {
9618                    b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9619                     .arg(&hd).arg(&nh).arg(&nhkv).arg(&base_i).arg(&scale).arg(&nspm).arg(&spk)
9620                     .arg(&ktb).arg(&vtb);
9621                    unsafe { b.launch(cfg)?; }
9622                }
9623            }
9624            let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t_g as u32, 1),
9625                    block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9626            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
9627            if head_dim == 512 {
9628                // device-len combine (shared by verify/eager/graph — parity by symbol): the
9629                // per-row n_splits derives from the SAME counter the rows kernel read.
9630                let (bd, plus) = base_dev.unwrap();
9631                let plus_g = plus + r0 as i32;
9632                if let Some((oq, od)) = q8_out.as_mut() {
9633                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
9634                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
9635                    if Self::pdl_on() && Self::pdl_wb_on() {
9636                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
9637                        use cudarc::driver::{DevicePtr, DevicePtrMut};
9638                        let s = &self.gpu.stream();
9639                        let (po, _g0) = part_o.device_ptr(s); let (pm, _g1) = part_m.device_ptr(s);
9640                        let (pl, _g2) = part_l.device_ptr(s);
9641                        let (pq, _g3) = oq.device_ptr_mut(s); let (pd, _g4) = od.device_ptr_mut(s);
9642                        let (pb, _g5) = bd.device_ptr(s);
9643                        let mut ps = [
9644                            &po as *const _ as *mut std::ffi::c_void, &pm as *const _ as *mut _,
9645                            &pl as *const _ as *mut _, &pq as *const _ as *mut _,
9646                            &pd as *const _ as *mut _, &hd as *const _ as *mut _,
9647                            &nh as *const _ as *mut _, &pb as *const _ as *mut _,
9648                            &plus_g as *const _ as *mut _, &nspm as *const _ as *mut _,
9649                            &spk as *const _ as *mut _,
9650                        ];
9651                        unsafe { self.launch_pdl_flash(Self::gkv_on(),
9652                            "fa_decode_combine_rows_dc_q8_1",
9653                            cfg2.grid_dim, cfg2.block_dim, 0, &mut ps)?; }
9654                        continue;
9655                    }
9656                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
9657                    let __s_b2 = self.gpu.stream();
9658                    let mut b2 = __s_b2.launch_builder(&fc);
9659                    b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut **oq).arg(&mut **od)
9660                      .arg(&hd).arg(&nh).arg(bd).arg(&plus_g).arg(&nspm).arg(&spk);
9661                    unsafe { b2.launch(cfg2)?; }
9662                    continue;
9663                }
9664                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
9665                let __s_b2 = self.gpu.stream();
9666                let mut b2 = __s_b2.launch_builder(&fc);
9667                b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut o_g).arg(&hd).arg(&nh)
9668                  .arg(bd).arg(&plus_g).arg(&nspm).arg(&spk);
9669                unsafe { b2.launch(cfg2)?; }
9670            } else {
9671                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
9672                // leave the caller's pair unwritten (consumer would read garbage).
9673                assert!(q8_out.is_none(), "rows q8 emit requires the hd512 dc combine");
9674                let fc = self.func("fa_decode_combine_rows");
9675                let __s_b2 = self.gpu.stream();
9676                let mut b2 = __s_b2.launch_builder(&fc);
9677                b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut o_g).arg(&hd).arg(&nh)
9678                  .arg(&base_i).arg(&nspm).arg(&spk);
9679                unsafe { b2.launch(cfg2)?; }
9680            }
9681        }
9682        Ok(())
9683    }
9684
9685    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
9686    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
9687    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
9688    #[allow(clippy::too_many_arguments)]
9689    pub fn fa_decode_rows_w(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9690                            v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9691                            head_dim: usize, n_head: usize, n_head_kv: usize,
9692                            base_dev: &CudaSlice<i32>, base_plus: i32, t: usize, scale: f32,
9693                            window: usize, k_tok_bytes: usize, v_tok_bytes: usize,
9694                            q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9695                            -> Result<(), Box<dyn std::error::Error>> {
9696        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
9697        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
9698        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
9699        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
9700        debug_assert!(head_dim == 256);
9701        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
9702        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
9703        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
9704        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
9705        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
9706        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
9707        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
9708        let sp = {
9709            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9710            let v = *SPW.get_or_init(|| std::env::var("MEMRA_FA_SPW").ok()
9711                .and_then(|x| x.parse().ok()).unwrap_or(0));
9712            if v >= 8 { v } else { FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed) }
9713        };
9714        let n_splits_max = (window + sp - 1) / sp;
9715        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9716        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
9717        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9718        let gqa = (n_head / n_head_kv).max(1) as u32;
9719        let o_len = t * n_head * n_splits_max * head_dim;
9720        let ml_len = t * n_head * n_splits_max;
9721        let mut part_guard = self.fa_part_pool.lock().unwrap();
9722        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9723            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9724            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9725            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9726            // later live allocations land at those addresses, and the next graph REPLAY writes
9727            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9728            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9729            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9730            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9731            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
9732            // (total retired < final size).
9733            let old = part_guard.take();
9734            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9735            if let Some(old) = old {
9736                self.fa_part_retired.lock().unwrap().push(old);
9737            }
9738            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9739                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9740            }
9741            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9742                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9743                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9744        }
9745        let pg = part_guard.as_mut().unwrap();
9746        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9747        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9748        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9749        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9750        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
9751        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
9752        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
9753        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
9754        // floor (deep-ctx broadcast win); register twin between.
9755        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9756        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
9757            std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9758                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9759        });
9760        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
9761        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
9762        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
9763        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
9764        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
9765        use cudarc::driver::sys::CUfunction_attribute_enum as A;
9766        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
9767        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
9768        // per (lane, format-module) keeps parity structural; the old register-i2 detour
9769        // (-33%) is retired.
9770        let wg = Self::wkv_on();
9771        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
9772        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
9773        let sp2 = gqa <= 4 && fa_v4_at(window)
9774            && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
9775        if sp2 {
9776            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
9777            if Self::pdl_on() && Self::pdl_wb_on() {
9778                // wave-B2b: flavor mirrors wg.
9779                use cudarc::driver::{DevicePtr, DevicePtrMut};
9780                let s = &self.gpu.stream();
9781                let (pq, _b0) = q.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9782                let (pv, _b2) = v.device_ptr(s);
9783                let (po, _b3) = part_o.device_ptr_mut(s);
9784                let (pm, _b4) = part_m.device_ptr_mut(s);
9785                let (pl, _b5) = part_l.device_ptr_mut(s);
9786                let (pb, _b6) = base_dev.device_ptr(s);
9787                let mut ps = [
9788                    &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9789                    &pv as *const _ as *mut _, &po as *const _ as *mut _,
9790                    &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9791                    &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9792                    &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9793                    &base_plus as *const _ as *mut _, &scale as *const _ as *mut _,
9794                    &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9795                    &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9796                    &wini as *const _ as *mut _,
9797                ];
9798                unsafe { self.launch_pdl_flash(wg, "fa_decode_vec_q_rows_v4_w_sp",
9799                    (n_head_kv as u32, n_splits_max as u32, t as u32), (32, gqa + 1, 1),
9800                    sh, &mut ps)?; }
9801            } else {
9802            let f = if wg { self.func_g("fa_decode_vec_q_rows_v4_w_sp") }
9803                    else { self.func("fa_decode_vec_q_rows_v4_w_sp") };
9804            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9805            let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9806                block_dim: (32, gqa + 1, 1), shared_mem_bytes: sh };
9807            let __s_b = self.gpu.stream();
9808            let mut b = __s_b.launch_builder(&f);
9809            b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9810             .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale).arg(&nspm).arg(&spk)
9811             .arg(&ktb).arg(&vtb).arg(&wini);
9812            unsafe { b.launch(cfg)?; }
9813            }
9814        } else {
9815        if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
9816            // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
9817            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
9818            use cudarc::driver::{DevicePtr, DevicePtrMut};
9819            let s = &self.gpu.stream();
9820            let (pq, _b0) = q.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9821            let (pv, _b2) = v.device_ptr(s);
9822            let (po, _b3) = part_o.device_ptr_mut(s);
9823            let (pm, _b4) = part_m.device_ptr_mut(s);
9824            let (pl, _b5) = part_l.device_ptr_mut(s);
9825            let (pb, _b6) = base_dev.device_ptr(s);
9826            let mut ps = [
9827                &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9828                &pv as *const _ as *mut _, &po as *const _ as *mut _,
9829                &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9830                &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9831                &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9832                &base_plus as *const _ as *mut _, &scale as *const _ as *mut _,
9833                &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9834                &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9835                &wini as *const _ as *mut _,
9836            ];
9837            unsafe { self.launch_pdl_flash(wg, "fa_decode_vec_q_rows_v4_w",
9838                (n_head_kv as u32, n_splits_max as u32, t as u32), (32, gqa, 1),
9839                sh, &mut ps)?; }
9840        } else {
9841        let pick = |name: &str| if wg { self.func_g(name) } else { self.func(name) };
9842        let (f, sh) = if fa_v4_at(window) {
9843            let f = pick("fa_decode_vec_q_rows_v4_w");
9844            (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
9845        } else if smem_tkv > 0 && window >= smem_tkv {
9846            // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
9847            // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
9848            (pick("fa_decode_vec_q_rows_smem_w"), (2 * 32 * head_dim * 2) as u32)
9849        } else {
9850            (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
9851        };
9852        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9853        let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9854            block_dim: (32, gqa, 1), shared_mem_bytes: sh };
9855        let __s_b = self.gpu.stream();
9856        let mut b = __s_b.launch_builder(&f);
9857        b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9858         .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale).arg(&nspm).arg(&spk)
9859         .arg(&ktb).arg(&vtb).arg(&wini);
9860        unsafe { b.launch(cfg)?; }
9861        }
9862        }
9863        let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
9864                block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9865        if let Some((oq, od)) = q8_out {
9866            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
9867            // consumes the pair directly; the standalone quantize launch folds away.
9868            if Self::pdl_on() && Self::pdl_wb_on() {
9869                // wave-B2: flavor mirrors the builder's wg choice.
9870                use cudarc::driver::{DevicePtr, DevicePtrMut};
9871                let s = &self.gpu.stream();
9872                let (po, _g0) = part_o.device_ptr(s); let (pm, _g1) = part_m.device_ptr(s);
9873                let (pl, _g2) = part_l.device_ptr(s);
9874                let (pq, _g3) = oq.device_ptr_mut(s); let (pd, _g4) = od.device_ptr_mut(s);
9875                let mut ps = [
9876                    &po as *const _ as *mut std::ffi::c_void, &pm as *const _ as *mut _,
9877                    &pl as *const _ as *mut _, &pq as *const _ as *mut _,
9878                    &pd as *const _ as *mut _, &hd as *const _ as *mut _,
9879                    &nh as *const _ as *mut _, &nspm as *const _ as *mut _,
9880                    &spk as *const _ as *mut _, &wini as *const _ as *mut _,
9881                ];
9882                unsafe { self.launch_pdl_flash(wg, "fa_decode_combine_rows_w_q8_1",
9883                                               cfg2.grid_dim, cfg2.block_dim, 0, &mut ps)?; }
9884                return Ok(());
9885            }
9886            let fc = if wg { self.func_g("fa_decode_combine_rows_w_q8_1") }
9887                     else { self.func("fa_decode_combine_rows_w_q8_1") };
9888            let __s_b2 = self.gpu.stream();
9889            let mut b2 = __s_b2.launch_builder(&fc);
9890            b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh)
9891              .arg(&nspm).arg(&spk).arg(&wini);
9892            unsafe { b2.launch(cfg2)?; }
9893            return Ok(());
9894        }
9895        let fc = if wg { self.func_g("fa_decode_combine_rows_w") }
9896                 else { self.func("fa_decode_combine_rows_w") };
9897        let __s_b2 = self.gpu.stream();
9898        let mut b2 = __s_b2.launch_builder(&fc);
9899        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9900          .arg(&nspm).arg(&spk).arg(&wini);
9901        unsafe { b2.launch(cfg2)?; }
9902        Ok(())
9903    }
9904
9905    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
9906    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
9907    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
9908    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
9909    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
9910    #[allow(clippy::too_many_arguments)]
9911    pub fn fa_decode_rows_dc(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9912                             v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9913                             head_dim: usize, n_head: usize, n_head_kv: usize,
9914                             base_dev: &CudaSlice<i32>, t_kv_upper: usize, t: usize, scale: f32,
9915                             k_tok_bytes: usize, v_tok_bytes: usize, base_plus: i32, g: bool)
9916                             -> Result<(), Box<dyn std::error::Error>> {
9917        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
9918        assert!(v4 || fa_v3_active(head_dim), "stream fa rows requires the v3 or v4 lane");
9919        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
9920        if v4 {
9921            let sp = fa_split_keys(t_kv_upper, n_head_kv);
9922            let n_splits_max = (t_kv_upper + sp - 1) / sp;
9923            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9924            let (nspm, spk) = (n_splits_max as i32, sp as i32);
9925            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9926            let gqa = (n_head / n_head_kv).max(1) as u32;
9927            let o_len = t * n_head * n_splits_max * head_dim;
9928            let ml_len = t * n_head * n_splits_max;
9929            let mut part_guard = self.fa_part_pool.lock().unwrap();
9930            if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9931                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9932            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9933            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9934            // later live allocations land at those addresses, and the next graph REPLAY writes
9935            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9936            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9937            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9938            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9939            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
9940            // (total retired < final size).
9941                let old = part_guard.take();
9942                let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9943                if let Some(old) = old {
9944                    self.fa_part_retired.lock().unwrap().push(old);
9945                }
9946                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9947                    eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9948                }
9949                *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9950                                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9951                                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9952            }
9953            let pg = part_guard.as_mut().unwrap();
9954            self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9955            self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9956            self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9957            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9958            let f = if g { self.func_g("fa_decode_vec_q_rows_v4_dc") }
9959                    else { self.func("fa_decode_vec_q_rows_v4_dc") };
9960            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
9961            use cudarc::driver::sys::CUfunction_attribute_enum as A;
9962            f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9963            let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9964                block_dim: (32, gqa, 1), shared_mem_bytes: sh };
9965            let __s_b = self.gpu.stream();
9966            let mut b = __s_b.launch_builder(&f);
9967            b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9968             .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale)
9969             .arg(&nspm).arg(&spk).arg(&ktb).arg(&vtb);
9970            unsafe { b.launch(cfg)?; }
9971            let fc = self.func("fa_decode_combine_rows_dc");
9972            let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
9973                block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9974            let __s_b2 = self.gpu.stream();
9975            let mut b2 = __s_b2.launch_builder(&fc);
9976            b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9977              .arg(base_dev).arg(&base_plus).arg(&nspm).arg(&spk);
9978            unsafe { b2.launch(cfg2)?; }
9979            return Ok(());
9980        }
9981        let sp = fa_split_keys(t_kv_upper, n_head_kv);
9982        let n_splits_max = (t_kv_upper + sp - 1) / sp;
9983        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9984        let (nspm, spk) = (n_splits_max as i32, sp as i32);
9985        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9986        let gqa = (n_head / n_head_kv).max(1) as u32;
9987        let o_len = t * n_head * n_splits_max * head_dim;
9988        let ml_len = t * n_head * n_splits_max;
9989        let mut part_guard = self.fa_part_pool.lock().unwrap();
9990        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9991            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
9992            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
9993            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
9994            // later live allocations land at those addresses, and the next graph REPLAY writes
9995            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
9996            // output corruption began the burst after the trunk's t_kv growth first realloc'd
9997            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
9998            // the baked addresses alive (single-stream: eager writes the new buffers, replays
9999            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
10000            // (total retired < final size).
10001            let old = part_guard.take();
10002            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
10003            if let Some(old) = old {
10004                self.fa_part_retired.lock().unwrap().push(old);
10005            }
10006            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
10007                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
10008            }
10009            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
10010                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
10011                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
10012        }
10013        let pg = part_guard.as_mut().unwrap();
10014        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
10015        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
10016        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
10017        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
10018        let f = self.func("fa_decode_vec_q_rows_v3_dc");
10019        let sh = (32 * head_dim * 2) as u32;
10020        use cudarc::driver::sys::CUfunction_attribute_enum as A;
10021        f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
10022        let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
10023            block_dim: (32, gqa, 1), shared_mem_bytes: sh };
10024        let __s_b = self.gpu.stream();
10025        let mut b = __s_b.launch_builder(&f);
10026        b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
10027         .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&scale).arg(&nspm).arg(&spk)
10028         .arg(&ktb).arg(&vtb);
10029        unsafe { b.launch(cfg)?; }
10030        let fc = self.func("fa_decode_combine_rows_dc");
10031        let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
10032            block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
10033        let plus0 = 0i32;
10034        let __s_b2 = self.gpu.stream();
10035        let mut b2 = __s_b2.launch_builder(&fc);
10036        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
10037          .arg(base_dev).arg(&plus0).arg(&nspm).arg(&spk);
10038        unsafe { b2.launch(cfg2)?; }
10039        Ok(())
10040    }
10041
10042    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
10043    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
10044    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
10045    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
10046    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
10047    ///
10048    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
10049    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
10050    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
10051    /// grouping (different but mathematically-equal log-sum-exp merge).
10052    pub fn fa_decode_dc(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
10053                        v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
10054                        head_dim: usize, n_head: usize, n_head_kv: usize,
10055                        t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32,
10056                        k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
10057                        -> Result<(), Box<dyn std::error::Error>> {
10058        self.fa_decode_dc_q8(q, k, v, o, head_dim, n_head, n_head_kv, t_kv_dev, bucket_max,
10059                             scale, k_tok_bytes, v_tok_bytes, g, None)
10060    }
10061
10062    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
10063    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
10064    #[allow(clippy::too_many_arguments)]
10065    pub fn fa_decode_dc_q8(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
10066                        v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
10067                        head_dim: usize, n_head: usize, n_head_kv: usize,
10068                        t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32,
10069                        k_tok_bytes: usize, v_tok_bytes: usize, g: bool,
10070                        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
10071                        -> Result<(), Box<dyn std::error::Error>> {
10072        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
10073        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
10074        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
10075        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
10076        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
10077        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
10078        // 2026-07-12).
10079        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
10080        if g && head_dim == 256 && !fa_v4_at(bucket_max) { fa_vec = false; }   // mirror kvmod/geom
10081        let sp = fa_split_keys(bucket_max, n_head_kv);
10082        let n_splits = if fa_vec { ((bucket_max + sp - 1) / sp).max(1) } else { ((bucket_max + 255) / 256).max(1) };
10083        let o_len = n_head * n_splits * head_dim;
10084        let ml_len = n_head * n_splits;
10085        let mut part_guard = self.fa_part_pool.lock().unwrap();
10086        if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
10087            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
10088            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
10089            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
10090            // later live allocations land at those addresses, and the next graph REPLAY writes
10091            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
10092            // output corruption began the burst after the trunk's t_kv growth first realloc'd
10093            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
10094            // the baked addresses alive (single-stream: eager writes the new buffers, replays
10095            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
10096            // (total retired < final size).
10097            let old = part_guard.take();
10098            let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
10099            if let Some(old) = old {
10100                self.fa_part_retired.lock().unwrap().push(old);
10101            }
10102            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
10103                eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
10104            }
10105            *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
10106                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
10107                                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
10108        }
10109        let pg = part_guard.as_mut().unwrap();
10110        self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
10111        self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
10112        self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
10113        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
10114        let (hd, nh, nhkv, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, n_splits as i32);
10115        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10116        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
10117        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
10118        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
10119        let deep = fa_vec && head_dim == 256 && fa_v4_at(bucket_max) && !g
10120            && fa_deep_at(bucket_max) && !matches!(fa_v4_mode(), "noB3" | "stage");
10121        let (f, cfg) = if fa_vec && head_dim == 512 && bucket_max >= {
10122            static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
10123            *FA512_MIN_DC.get_or_init(|| std::env::var("MEMRA_FA512_MIN").ok()
10124                .and_then(|v| v.parse().ok()).unwrap_or(512))
10125        } {
10126            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
10127            let gqa = (n_head / n_head_kv).max(1) as u32;
10128            (self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
10129             LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10130                 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
10131        } else if fa_vec && head_dim == 512 {
10132            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
10133            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
10134            return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
10135                                                 0, Some(t_kv_dev), scale, n_splits, sp,
10136                                                 k_tok_bytes, v_tok_bytes, g,
10137                                                 &mut *part_o, &mut *part_m, &mut *part_l, q8_out);
10138        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
10139            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
10140            // incl the g-module route + raw-e4m3 sV sizing.
10141            let gqa = (n_head / n_head_kv).max(1) as u32;
10142            let fv = if g { self.func_g("fa_decode_vec_q_v4_dc") }
10143                     else if deep { self.func("fa_decode_vec_q_v4_deep_dc") }
10144                     else { self.func("fa_decode_vec_q_v4_dc") };
10145            let shmem = (if deep { 12160 } else { 11520 }
10146                         + 32 * head_dim * if g { 1 } else { 2 }) as u32;
10147            use cudarc::driver::sys::CUfunction_attribute_enum as A;
10148            fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
10149            (fv, LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10150                 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10151        } else if fa_vec && fa_v3_active(head_dim) {
10152            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
10153            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
10154            let gqa = (n_head / n_head_kv).max(1) as u32;
10155            let fv = if g { self.func_g("fa_decode_vec_q_v3_dc") } else { self.func("fa_decode_vec_q_v3_dc") };
10156            let shmem = (32 * head_dim * 2) as u32;       // sV bf16 [FA_DEC_TILE=32][hd]
10157            (fv,
10158             LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10159                 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10160        } else if fa_vec && fa_v2_on() {
10161            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
10162            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
10163            // a numeric config; eager, rows-verify and graph all switch together).
10164            let gqa = (n_head / n_head_kv).max(1) as u32;
10165            let fv = if g { self.func_g("fa_decode_vec_q_v2_dc") } else { self.func("fa_decode_vec_q_v2_dc") };
10166            let shmem = (2 * 32 * head_dim * 2) as u32;   // sK+sV bf16 [FA_DEC_TILE=32][hd]
10167            (fv,
10168             LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10169                 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10170        } else if fa_vec {
10171            let gqa = (n_head / n_head_kv).max(1) as u32;
10172            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
10173            let fv = if g { self.func_g("fa_decode_vec_q_dc") } else { self.func("fa_decode_vec_q_dc") };
10174            (fv,
10175             LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10176                 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
10177        } else {
10178            return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
10179                                                 0, Some(t_kv_dev), scale, n_splits,
10180                                                 if fa_vec { sp } else { 256 },
10181                                                 k_tok_bytes, v_tok_bytes, g,
10182                                                 &mut *part_o, &mut *part_m, &mut *part_l, q8_out);
10183        };
10184        let ski = sp as i32;   // one-partition law: the twins derive ns_eff from (T_kv, ski)
10185        let __s_b = self.gpu.stream();
10186        let mut b = __s_b.launch_builder(&f);
10187        b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
10188         .arg(&hd).arg(&nh).arg(&nhkv).arg(t_kv_dev).arg(&scale).arg(&nsp).arg(&ski)
10189         .arg(&ktb).arg(&vtb);
10190        unsafe { b.launch(cfg)?; }
10191        let cfg2 = LaunchConfig { grid_dim: (n_head as u32, 1, 1), block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
10192        if let Some((oq, od)) = q8_out {
10193            let fc = if g { self.func_g("fa_decode_combine_q8_1") }
10194                     else { self.fa_func("fa_decode_combine_q8_1", head_dim) };
10195            let __s_b2 = self.gpu.stream();
10196            let mut b2 = __s_b2.launch_builder(&fc);
10197            b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh).arg(&nsp);
10198            unsafe { b2.launch(cfg2)?; }
10199            return Ok(());
10200        }
10201        let fc = if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) };
10202        let __s_b2 = self.gpu.stream();
10203        let mut b2 = __s_b2.launch_builder(&fc);
10204        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
10205        unsafe { b2.launch(cfg2)?; }
10206        Ok(())
10207    }
10208
10209    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
10210    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
10211    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
10212    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
10213    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
10214    pub fn fa_geom_eager(&self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool) -> (bool, usize) {
10215        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
10216        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
10217        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
10218        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
10219        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
10220        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
10221        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
10222        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
10223        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
10224        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
10225        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
10226        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
10227        // family; everything else falls to the g-module scalar.
10228        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
10229        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
10230        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
10231        if g && head_dim == 256 && !fa_v4_at(t_kv) { fa_vec = false; }
10232        let sp = fa_split_keys(t_kv, n_head_kv);
10233        let n_splits = if fa_vec { ((t_kv + sp - 1) / sp).max(1) } else { ((t_kv + 255) / 256).max(1) };
10234        (fa_vec, n_splits)
10235    }
10236
10237    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
10238    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
10239    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
10240    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
10241    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
10242    pub fn fa_bucket_key(&self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool) -> (bool, usize) {
10243        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
10244    }
10245
10246    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
10247    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
10248    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
10249    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
10250    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
10251    /// device data) — every per-step varying scalar must come from a device counter. Returns the
10252    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
10253    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
10254    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
10255    /// replays (transients returning to the pool get reused by unrelated work and corrupt
10256    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
10257    pub fn capture_graph_retained<F>(&self, step: F)
10258        -> Result<(cudarc::driver::CudaGraph, Vec<Box<dyn std::any::Any + Send>>), Box<dyn std::error::Error>>
10259        where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10260    {
10261        use cudarc::driver::sys::CUgraphInstantiate_flags;
10262        self.capture_graph_retained_flags(
10263            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, step)
10264    }
10265
10266    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
10267    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
10268    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
10269    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
10270    pub fn capture_graph_retained_flags<F>(&self,
10271        flags: cudarc::driver::sys::CUgraphInstantiate_flags, mut step: F)
10272        -> Result<(cudarc::driver::CudaGraph, Vec<Box<dyn std::any::Any + Send>>), Box<dyn std::error::Error>>
10273        where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10274    {
10275        use cudarc::driver::sys::CUstreamCaptureMode;
10276        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
10277        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
10278        // while the capture region is open become dead copy NODES replayed every launch
10279        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
10280        // warmup runs allocate the same transient sequence at the same pool addresses, so
10281        // retaining the warmup clones preserves the draft-graph fix without polluting the
10282        // captured graph.
10283        self.capture_keep.lock().unwrap().clear();
10284        let was_tracking = self.gpu.ctx.is_event_tracking();
10285        if was_tracking { unsafe { self.gpu.ctx.disable_event_tracking(); } }
10286        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
10287            self.capture_keep_on.store(true, std::sync::atomic::Ordering::Relaxed);
10288            let w = (|| { step(self)?; step(self) })();
10289            self.capture_keep_on.store(false, std::sync::atomic::Ordering::Relaxed);
10290            w?;
10291            self.gpu.stream().synchronize()?;
10292            self.gpu.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
10293            let r = step(self);
10294            let g = self.gpu.stream().end_capture(flags);
10295            r?;
10296            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
10297            graph.upload()?;
10298            Ok(graph)
10299        };
10300        let result = run();
10301        self.capture_keep_on.store(false, std::sync::atomic::Ordering::Relaxed);
10302        if was_tracking { unsafe { self.gpu.ctx.enable_event_tracking(); } }
10303        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
10304        Ok((result?, keeper))
10305    }
10306
10307    pub fn capture_graph<F>(&self, mut step: F) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
10308        where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10309    {
10310        use cudarc::driver::sys::{CUstreamCaptureMode, CUgraphInstantiate_flags};
10311        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
10312        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
10313        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
10314        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
10315        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
10316        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
10317        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
10318        let was_tracking = self.gpu.ctx.is_event_tracking();
10319        if was_tracking { unsafe { self.gpu.ctx.disable_event_tracking(); } }
10320        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
10321            // warmup: two inline runs (no capture) so allocator pointers + kernel attrs are stable.
10322            step(self)?;
10323            step(self)?;
10324            self.gpu.stream().synchronize()?;
10325            // capture the third run.
10326            self.gpu.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
10327            // If the body errors mid-capture, end the capture before propagating so the stream isn't
10328            // left in a capturing state.
10329            let r = step(self);
10330            let g = self.gpu.stream().end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
10331            r?;
10332            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
10333            graph.upload()?;
10334            Ok(graph)
10335        };
10336        let result = run();
10337        if was_tracking { unsafe { self.gpu.ctx.enable_event_tracking(); } }
10338        result
10339    }
10340
10341    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
10342    pub fn gdn_scan_s128_view(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10343                              g: &CudaSlice<f32>, beta: &CudaSlice<f32>,
10344                              state_in: &cudarc::driver::CudaView<f32>,
10345                              state_out: &mut cudarc::driver::CudaViewMut<f32>,
10346                              o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32)
10347                              -> Result<(), Box<dyn std::error::Error>> {
10348        let f = self.func("gdn_scan_s128");
10349        const S_V: u32 = 128; const WARP: u32 = 32; const COLS: u32 = 4;
10350        let cfg = LaunchConfig { grid_dim: (n_head as u32, 1, S_V / COLS), block_dim: (WARP, COLS, 1), shared_mem_bytes: 0 };
10351        let (h, ti) = (n_head as i32, t as i32);
10352        let __s_b = self.gpu.stream();
10353        let mut b = __s_b.launch_builder(&f);
10354        b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o).arg(&h).arg(&ti).arg(&scale);
10355        unsafe { b.launch(cfg)?; }
10356        Ok(())
10357    }
10358
10359    /// conv1d where the input is a CudaView (resident conv state assembled in place).
10360    pub fn ssm_conv1d_view(&self, x: &cudarc::driver::CudaView<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10361                           conv_dim: usize, t: usize, d_conv: usize, silu: bool)
10362                           -> Result<(), Box<dyn std::error::Error>> {
10363        let f = self.func("ssm_conv1d_silu_f32");
10364        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
10365        let cfg = LaunchConfig { grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
10366                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10367        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
10368        let __s_b = self.gpu.stream();
10369        let mut b = __s_b.launch_builder(&f);
10370        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
10371        unsafe { b.launch(cfg)?; }
10372        Ok(())
10373    }
10374
10375    /// Depthwise causal conv1d + optional SiLU.
10376    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
10377    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
10378    /// FUSED prefill conv (token-major input, zero left-state): replaces
10379    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
10380    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
10381    pub fn ssm_conv1d_tm(&self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10382                         conv_dim: usize, t: usize, d_conv: usize)
10383                         -> Result<(), Box<dyn std::error::Error>> {
10384        let f = self.func("ssm_conv1d_tm_f32");
10385        let cfg = LaunchConfig {
10386            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10387            block_dim: (256, 1, 1), shared_mem_bytes: 0,
10388        };
10389        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10390        let __s_b = self.gpu.stream();
10391        let mut b = __s_b.launch_builder(&f);
10392        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10393        unsafe { b.launch(cfg)?; }
10394        Ok(())
10395    }
10396
10397    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
10398    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
10399    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
10400    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
10401    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
10402    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
10403    /// columns; the final ring == what T sequential decode ring rolls leave).
10404    pub fn ssm_conv1d_tm_state(&self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
10405                               w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10406                               conv_dim: usize, t: usize, d_conv: usize)
10407                               -> Result<(), Box<dyn std::error::Error>> {
10408        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
10409    }
10410
10411    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
10412    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
10413    #[allow(clippy::too_many_arguments)]
10414    pub fn ssm_conv1d_tm_state_pad(&self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
10415                               w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10416                               conv_dim: usize, t: usize, d_conv: usize,
10417                               pad_len: Option<&CudaSlice<i32>>)
10418                               -> Result<(), Box<dyn std::error::Error>> {
10419        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
10420        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
10421        // the window kernel both read the pre-roll ring; the roll launches after both) — but
10422        // cloning first keeps the ordering trivially correct under any future stream split.
10423        let ring_old = if t < d_conv - 1 { Some(self.clone_dtod(conv_state)?) } else { None };
10424        {
10425            let f = self.func("ssm_conv1d_tm_state_f32");
10426            let cfg = LaunchConfig {
10427                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10428                block_dim: (256, 1, 1), shared_mem_bytes: 0,
10429            };
10430            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10431            let __s_b = self.gpu.stream();
10432            let mut b = __s_b.launch_builder(&f);
10433            b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10434            unsafe { b.launch(cfg)?; }
10435        }
10436        match (ring_old, pad_len) {
10437            (None, Some(len_d)) => {
10438                let f = self.func("ssm_conv_ring_update_dev_f32");
10439                let n = conv_dim * (d_conv - 1);
10440                let cfg = LaunchConfig::for_num_elems(n as u32);
10441                let (cd, dc) = (conv_dim as i32, d_conv as i32);
10442                let __s_b = self.gpu.stream();
10443                let mut b = __s_b.launch_builder(&f);
10444                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10445                unsafe { b.launch(cfg)?; }
10446            }
10447            (None, None) => {
10448                let f = self.func("ssm_conv_ring_update_f32");
10449                let n = conv_dim * (d_conv - 1);
10450                let cfg = LaunchConfig::for_num_elems(n as u32);
10451                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10452                let __s_b = self.gpu.stream();
10453                let mut b = __s_b.launch_builder(&f);
10454                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10455                unsafe { b.launch(cfg)?; }
10456            }
10457            (Some(old), _) => self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?,
10458        }
10459        Ok(())
10460    }
10461
10462    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
10463    pub fn ssm_conv1d_tm_state_pad_v(&self, qkv_tm: &cudarc::driver::CudaView<f32>, conv_state: &mut CudaSlice<f32>,
10464                               w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10465                               conv_dim: usize, t: usize, d_conv: usize,
10466                               pad_len: Option<&CudaSlice<i32>>)
10467                               -> Result<(), Box<dyn std::error::Error>> {
10468        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
10469        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
10470        // the window kernel both read the pre-roll ring; the roll launches after both) — but
10471        // cloning first keeps the ordering trivially correct under any future stream split.
10472        let ring_old = if t < d_conv - 1 { Some(self.clone_dtod(conv_state)?) } else { None };
10473        {
10474            let f = self.func("ssm_conv1d_tm_state_f32");
10475            let cfg = LaunchConfig {
10476                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10477                block_dim: (256, 1, 1), shared_mem_bytes: 0,
10478            };
10479            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10480            let __s_b = self.gpu.stream();
10481            let mut b = __s_b.launch_builder(&f);
10482            b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10483            unsafe { b.launch(cfg)?; }
10484        }
10485        match (ring_old, pad_len) {
10486            (None, Some(len_d)) => {
10487                let f = self.func("ssm_conv_ring_update_dev_f32");
10488                let n = conv_dim * (d_conv - 1);
10489                let cfg = LaunchConfig::for_num_elems(n as u32);
10490                let (cd, dc) = (conv_dim as i32, d_conv as i32);
10491                let __s_b = self.gpu.stream();
10492                let mut b = __s_b.launch_builder(&f);
10493                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10494                unsafe { b.launch(cfg)?; }
10495            }
10496            (None, None) => {
10497                let f = self.func("ssm_conv_ring_update_f32");
10498                let n = conv_dim * (d_conv - 1);
10499                let cfg = LaunchConfig::for_num_elems(n as u32);
10500                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10501                let __s_b = self.gpu.stream();
10502                let mut b = __s_b.launch_builder(&f);
10503                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10504                unsafe { b.launch(cfg)?; }
10505            }
10506            (Some(_), _) => unreachable!(
10507                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"),
10508        }
10509        Ok(())
10510    }
10511
10512    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
10513    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
10514    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
10515    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
10516    pub fn ssm_conv_ring_rebuild(&self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>,
10517                                 conv_state: &mut CudaSlice<f32>,
10518                                 conv_dim: usize, tc: usize, d_conv: usize)
10519                                 -> Result<(), Box<dyn std::error::Error>> {
10520        let f = self.func("ssm_conv_ring_rebuild_f32");
10521        let n = conv_dim * (d_conv - 1);
10522        let cfg = LaunchConfig::for_num_elems(n as u32);
10523        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
10524        let __s_b = self.gpu.stream();
10525        let mut b = __s_b.launch_builder(&f);
10526        b.arg(qkv_tm).arg(ring_old).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10527        unsafe { b.launch(cfg)?; }
10528        Ok(())
10529    }
10530
10531    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
10532    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
10533    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
10534    /// the argmax + run-spec gates are the authority.
10535    #[allow(clippy::too_many_arguments)]
10536    pub fn gdn_prep_decode(&self, conv_out: &CudaSlice<f32>, beta_raw: &CudaSlice<f32>,
10537                           alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10538                           q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10539                           beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
10540                           d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32)
10541                           -> Result<(), Box<dyn std::error::Error>> {
10542        let f = self.func("gdn_prep_decode_f32");
10543        let cfg = LaunchConfig { grid_dim: (num_v as u32, 1, 1), block_dim: (32, 4, 1), shared_mem_bytes: 0 };
10544        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
10545        let __s_b = self.gpu.stream();
10546        let mut b = __s_b.launch_builder(&f);
10547        b.arg(conv_out).arg(beta_raw).arg(alpha).arg(dt_bias).arg(a)
10548         .arg(q_l2).arg(k_l2).arg(v_g).arg(beta).arg(g_log)
10549         .arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&eps);
10550        unsafe { b.launch(cfg)?; }
10551        Ok(())
10552    }
10553
10554    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
10555    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
10556    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
10557    #[allow(clippy::too_many_arguments)]
10558    pub fn ssm_conv1d_gdn(&self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>,
10559                          q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10560                          conv_dim: usize, t: usize, d_conv: usize,
10561                          d_state: usize, num_v: usize, num_k: usize, key_dim: usize)
10562                          -> Result<(), Box<dyn std::error::Error>> {
10563        let f = self.func("ssm_conv1d_gdn_f32");
10564        let cfg = LaunchConfig {
10565            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10566            block_dim: (256, 1, 1), shared_mem_bytes: 0,
10567        };
10568        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10569        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
10570        let __s_b = self.gpu.stream();
10571        let mut b = __s_b.launch_builder(&f);
10572        b.arg(qkv_tm).arg(w).arg(q_g).arg(k_g).arg(v_g)
10573         .arg(&cd).arg(&ti).arg(&dc).arg(&ds).arg(&nv).arg(&nk).arg(&kd);
10574        unsafe { b.launch(cfg)?; }
10575        Ok(())
10576    }
10577
10578    pub fn ssm_conv1d(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10579                      conv_dim: usize, t: usize, d_conv: usize, silu: bool)
10580                      -> Result<(), Box<dyn std::error::Error>> {
10581        let f = self.func("ssm_conv1d_silu_f32");
10582        let cfg = LaunchConfig { grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
10583                                 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10584        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
10585        let __s_b = self.gpu.stream();
10586        let mut b = __s_b.launch_builder(&f);
10587        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
10588        unsafe { b.launch(cfg)?; }
10589        Ok(())
10590    }
10591
10592    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
10593    /// o:[128,H,T]. Single sequence.
10594    pub fn gdn_scan_s128(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10595                         g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
10596                         state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
10597                         n_head: usize, t: usize, scale: f32)
10598                         -> Result<(), Box<dyn std::error::Error>> {
10599        let f = self.func("gdn_scan_s128");
10600        const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
10601        let cfg = LaunchConfig {
10602            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
10603            block_dim: (WARP, COLS_PER_BLOCK, 1),
10604            shared_mem_bytes: 0,
10605        };
10606        let (h, ti) = (n_head as i32, t as i32);
10607        let __s_b = self.gpu.stream();
10608        let mut b = __s_b.launch_builder(&f);
10609        b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o).arg(&h).arg(&ti).arg(&scale);
10610        unsafe { b.launch(cfg)?; }
10611        Ok(())
10612    }
10613
10614    // ==== B2' batched decode state ops (decode_batch.rs) ====
10615    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
10616    // Bodies are the single-seq kernels per sequence — bit-identical per row.
10617
10618    #[allow(clippy::too_many_arguments)]
10619    pub fn ssm_conv1d_fused_decode_b(
10620        &self, qkv_cols: &CudaSlice<f32>, conv_state_ptrs: &cudarc::driver::CudaView<u64>,
10621        w: &CudaSlice<f32>, conv_outs: &mut CudaSlice<f32>, conv_dim: usize, d_conv: usize,
10622        b_n: usize) -> Result<(), Box<dyn std::error::Error>> {
10623        let f = self.func("ssm_conv1d_fused_decode_b_f32");
10624        let cfg = LaunchConfig {
10625            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
10626            block_dim: (256, 1, 1), shared_mem_bytes: 0,
10627        };
10628        let (cd, dc) = (conv_dim as i32, d_conv as i32);
10629        let __s_b = self.gpu.stream();
10630        let mut b = __s_b.launch_builder(&f);
10631        b.arg(qkv_cols).arg(conv_state_ptrs).arg(w).arg(conv_outs).arg(&cd).arg(&dc);
10632        unsafe { b.launch(cfg)?; }
10633        Ok(())
10634    }
10635
10636    #[allow(clippy::too_many_arguments)]
10637    pub fn gdn_prep_decode_b(
10638        &self, conv_outs: &CudaSlice<f32>, beta_raws: &CudaSlice<f32>, alphas: &CudaSlice<f32>,
10639        dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10640        q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10641        beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
10642        d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32,
10643        conv_dim: usize, b_n: usize) -> Result<(), Box<dyn std::error::Error>> {
10644        let f = self.func("gdn_prep_decode_b_f32");
10645        let cfg = LaunchConfig {
10646            grid_dim: (num_v as u32, 1, b_n as u32),
10647            block_dim: (32, 4, 1), shared_mem_bytes: 0,
10648        };
10649        let (ds, nv, nk, kd, cd) =
10650            (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, conv_dim as i32);
10651        let __s_b = self.gpu.stream();
10652        let mut b = __s_b.launch_builder(&f);
10653        b.arg(conv_outs).arg(beta_raws).arg(alphas).arg(dt_bias).arg(a)
10654         .arg(q_l2).arg(k_l2).arg(v_g).arg(beta).arg(g_log)
10655         .arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&eps).arg(&cd);
10656        unsafe { b.launch(cfg)?; }
10657        Ok(())
10658    }
10659
10660    #[allow(clippy::too_many_arguments)]
10661    pub fn gdn_scan_s128_batched(
10662        &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10663        g: &CudaSlice<f32>, beta: &CudaSlice<f32>,
10664        state_in_ptrs: &cudarc::driver::CudaView<u64>,
10665        state_out_ptrs: &cudarc::driver::CudaView<u64>,
10666        o: &mut CudaSlice<f32>, n_head: usize, b_n: usize, scale: f32)
10667        -> Result<(), Box<dyn std::error::Error>> {
10668        let f = self.func("gdn_scan_s128_b");
10669        const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
10670        let cfg = LaunchConfig {
10671            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
10672            block_dim: (WARP, COLS_PER_BLOCK, 1), shared_mem_bytes: 0,
10673        };
10674        let h = n_head as i32;
10675        let __s_b = self.gpu.stream();
10676        let mut b = __s_b.launch_builder(&f);
10677        b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in_ptrs).arg(state_out_ptrs)
10678         .arg(o).arg(&h).arg(&scale);
10679        unsafe { b.launch(cfg)?; }
10680        Ok(())
10681    }
10682
10683    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
10684    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
10685    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
10686    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
10687    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
10688    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
10689    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
10690    /// identity law); prime_cache/forward/forward_last are the only callers.
10691    pub fn gdn_chunked_enabled() -> bool {
10692        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10693        *E.get_or_init(|| std::env::var("MEMRA_GDN_CHUNKED").map(|v| v != "0").unwrap_or(true))
10694    }
10695
10696    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
10697    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
10698    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
10699    /// of 32 in [32, 128] (kernel row mappings require it).
10700    pub fn gdn_chunk_size() -> usize {
10701        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
10702        *C.get_or_init(|| {
10703            let c: usize = std::env::var("MEMRA_GDN_CHUNK").ok()
10704                .and_then(|v| v.parse().ok()).unwrap_or(32);
10705            c.clamp(32, 128) / 32 * 32
10706        })
10707    }
10708
10709    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
10710    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
10711    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
10712    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
10713    #[allow(clippy::too_many_arguments)]
10714    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
10715    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
10716    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
10717    #[allow(clippy::too_many_arguments)]
10718    pub fn gdn_chunk_k123(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10719                          g: &CudaSlice<f32>, beta: &CudaSlice<f32>, wb16: Option<&mut CudaSlice<u8>>,
10720                          n_head: usize, t: usize, c: usize, hk: usize,
10721                          k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>)
10722                          -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10723        const D: usize = 128;
10724        let h = n_head;
10725        let nc = (t + c - 1) / c;
10726        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
10727        let mut gcum = self.uninit(t * h)?;
10728        let mut a = self.uninit(nc * h * c * c)?;
10729        let mut p = self.uninit(nc * h * c * c)?;
10730        let mut u = self.uninit(nc * h * c * D)?;
10731        let mut w = self.uninit(nc * h * c * D)?;
10732        {   // K1
10733            let f = self.func("gdn_chunk_cumgate_f32");
10734            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
10735            let __s_b = self.gpu.stream();
10736            let mut b = __s_b.launch_builder(&f);
10737            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
10738            unsafe { b.launch(cfg)?; }
10739        }
10740        if let Some((qb, kb, pb)) = k2w {
10741            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
10742            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
10743            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
10744            let f = self.func("gdn_k2_wgmma");
10745            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
10746            let hki = hk as i32;
10747            let __s_b = self.gpu.stream();
10748            let mut b = __s_b.launch_builder(&f);
10749            b.arg(qb).arg(kb).arg(&gcum).arg(beta).arg(&mut a).arg(&mut *pb).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
10750            unsafe { b.launch(cfg)?; }
10751        } else if c <= 64 && !portable_mma_gated() {   // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
10752            let f = self.func("gdn_chunk_attn_f32");
10753            let jt = ((c + 31) / 32) as u32;
10754            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10755            let hki = hk as i32;
10756            let __s_b = self.gpu.stream();
10757            let mut b = __s_b.launch_builder(&f);
10758            b.arg(q).arg(k).arg(&gcum).arg(beta).arg(&mut a).arg(&mut p).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
10759            unsafe { b.launch(cfg)?; }
10760        } else {       // K2 generic (C = 128, or the portable target's low-smem fallback)
10761            assert!(hk == h, "generic K2 is broadcast-only (de-broadcast rides C==32)");
10762            let f = self.func("gdn_chunk_attn_g_f32");
10763            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (32, 8, 1), shared_mem_bytes: 0 };
10764            let __s_b = self.gpu.stream();
10765            let mut b = __s_b.launch_builder(&f);
10766            b.arg(q).arg(k).arg(&gcum).arg(beta).arg(&mut a).arg(&mut p).arg(&hi).arg(&ti).arg(&ci);
10767            unsafe { b.launch(cfg)?; }
10768        }
10769        {   // K3 (register-history templates for C=32/64; local-memory generic otherwise)
10770            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10771            match c {
10772                32 | 64 => {
10773                    let f = self.func(if c == 32 { "gdn_chunk_solve32_f32" } else { "gdn_chunk_solve64_f32" });
10774                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
10775                    let wb: u64 = match wb16 { Some(d) => self.addr_u8(d), None => 0 };
10776                    let hki = hk as i32;
10777                    let __s_b = self.gpu.stream();
10778                    let mut b = __s_b.launch_builder(&f);
10779                    b.arg(v).arg(k).arg(&a).arg(&gcum).arg(&mut u).arg(&mut w).arg(&wb).arg(&hi).arg(&ti).arg(&hki);
10780                    unsafe { b.launch(cfg)?; }
10781                }
10782                _ => {
10783                    assert!(hk == h, "generic K3 is broadcast-only");
10784                    let f = self.func("gdn_chunk_solve_f32");
10785                    let __s_b = self.gpu.stream();
10786                    let mut b = __s_b.launch_builder(&f);
10787                    b.arg(v).arg(k).arg(&a).arg(&gcum).arg(&mut u).arg(&mut w).arg(&hi).arg(&ti).arg(&ci);
10788                    unsafe { b.launch(cfg)?; }
10789                }
10790            }
10791        }
10792        Ok((gcum, p, u, w))
10793    }
10794
10795    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
10796    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
10797    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
10798    pub fn gdn_db_on() -> bool {
10799        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
10800    }
10801
10802    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
10803    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
10804    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
10805        !portable_mma_gated() && c == 32
10806            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
10807                Ok("1") => true,
10808                Ok("0") => false,
10809                _ => cfg!(memra_hopper_mma),
10810            }
10811    }
10812
10813    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
10814    /// mma config; same per-call env read discipline).
10815    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
10816        self.gdn_mma_enabled(c)
10817            && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
10818                Ok("0") => false,
10819                Ok("1") => true,
10820                _ => cfg!(memra_hopper_mma),
10821            }
10822    }
10823
10824    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
10825    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
10826    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
10827    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
10828    #[allow(clippy::too_many_arguments)]
10829    pub fn ssm_conv1d_gdn_state_pad(&self, qkv_tm: &cudarc::driver::CudaView<f32>,
10830                               conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>,
10831                               q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>,
10832                               v_g: &mut CudaSlice<f32>,
10833                               conv_dim: usize, t: usize, d_conv: usize,
10834                               d_state: usize, num_v: usize, num_k: usize, key_dim: usize,
10835                               hk: usize,
10836                               pad_len: Option<&CudaSlice<i32>>)
10837                               -> Result<(), Box<dyn std::error::Error>> {
10838        assert!(t >= d_conv - 1, "fused state conv requires T >= pad (PRIME_MIN_T gates)");
10839        {
10840            let f = self.func("ssm_conv1d_gdn_state_f32");
10841            let cfg = LaunchConfig {
10842                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10843                block_dim: (256, 1, 1), shared_mem_bytes: 0,
10844            };
10845            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10846            let (ds, nv, nk, kd, hki) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, hk as i32);
10847            let __s_b = self.gpu.stream();
10848            let mut b = __s_b.launch_builder(&f);
10849            b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(q_g).arg(k_g).arg(v_g)
10850             .arg(&cd).arg(&ti).arg(&dc).arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&hki);
10851            unsafe { b.launch(cfg)?; }
10852        }
10853        match pad_len {
10854            Some(len_d) => {
10855                let f = self.func("ssm_conv_ring_update_dev_f32");
10856                let n = conv_dim * (d_conv - 1);
10857                let cfg = LaunchConfig::for_num_elems(n as u32);
10858                let (cd, dc) = (conv_dim as i32, d_conv as i32);
10859                let __s_b = self.gpu.stream();
10860                let mut b = __s_b.launch_builder(&f);
10861                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10862                unsafe { b.launch(cfg)?; }
10863            }
10864            None => {
10865                let f = self.func("ssm_conv_ring_update_f32");
10866                let n = conv_dim * (d_conv - 1);
10867                let cfg = LaunchConfig::for_num_elems(n as u32);
10868                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10869                let __s_b = self.gpu.stream();
10870                let mut b = __s_b.launch_builder(&f);
10871                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10872                unsafe { b.launch(cfg)?; }
10873            }
10874        }
10875        Ok(())
10876    }
10877
10878    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
10879    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
10880    /// K2/K3 can write them.
10881    pub fn gdn_chunk_alloc(&self, n_head: usize, t: usize, c: usize, hk: usize)
10882                           -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
10883        const D: usize = 128;
10884        assert!(c == 32, "gdn_chunk_alloc: varlen chain is the C==32 mma pair");
10885        let h = n_head;
10886        let nc = (t + c - 1) / c;
10887        Ok(GdnChunkBufs {
10888            gcum: self.uninit(t * h)?,
10889            a: self.uninit(nc * h * c * c)?,
10890            p: self.uninit(nc * h * c * c)?,
10891            u: self.uninit(nc * h * c * D)?,
10892            w: self.uninit(nc * h * c * D)?,
10893            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
10894            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
10895            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
10896            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
10897            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
10898            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
10899            o: self.uninit(D * h * t)?,
10900            t, nc,
10901        })
10902    }
10903
10904    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
10905    pub fn f32_to_bf16_v(&self, x: &cudarc::driver::CudaView<f32>, dst: &mut CudaSlice<u8>, n: usize)
10906                         -> Result<(), Box<dyn std::error::Error>> {
10907        let f = self.func("f32_to_bf16_bulk");
10908        let ni = n as i64;
10909        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10910        let __s_b = self.gpu.stream();
10911        let mut b = __s_b.launch_builder(&f);
10912        b.arg(x).arg(dst).arg(&ni);
10913        unsafe { b.launch(cfg)?; }
10914        Ok(())
10915    }
10916
10917    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
10918    pub fn f32_to_bf16_into(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<u8>, n: usize)
10919                       -> Result<(), Box<dyn std::error::Error>> {
10920        let f = self.func("f32_to_bf16_bulk");
10921        let ni = n as i64;
10922        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10923        let __s_b = self.gpu.stream();
10924        let mut b = __s_b.launch_builder(&f);
10925        b.arg(x).arg(dst).arg(&ni);
10926        unsafe { b.launch(cfg)?; }
10927        Ok(())
10928    }
10929
10930    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
10931    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
10932    pub fn gdn_chunk_k123_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, hk: usize,
10933                              wq: Option<&GdnWVl8>)
10934                              -> Result<(), Box<dyn std::error::Error>> {
10935        let b = seqs.len();
10936        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
10937        let mut packed = [GdnSeqVl::default(); 8];
10938        packed[..b].copy_from_slice(seqs);
10939        let v = GdnVl8(packed);
10940        let (hi, ci) = (n_head as i32, 32i32);
10941        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
10942        {
10943            let f = self.func("gdn_chunk_cumgate_vl");
10944            let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
10945            let __s_lb = self.gpu.stream();
10946            let mut lb = __s_lb.launch_builder(&f);
10947            lb.arg(&v).arg(&hi).arg(&ci);
10948            unsafe { lb.launch(cfg)?; }
10949        }
10950        let hki = hk as i32;
10951        if let Some(w) = wq {   // K2-wgmma vl twin (writes A + pre-masked Pb16)
10952            let f = self.func("gdn_k2_wgmma_vl");
10953            let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
10954            let __s_lb = self.gpu.stream();
10955            let mut lb = __s_lb.launch_builder(&f);
10956            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
10957            unsafe { lb.launch(cfg)?; }
10958        } else {
10959            let f = self.func("gdn_chunk_attn_vl");
10960            let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10961            let __s_lb = self.gpu.stream();
10962            let mut lb = __s_lb.launch_builder(&f);
10963            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
10964            unsafe { lb.launch(cfg)?; }
10965        }
10966        {
10967            let f = self.func("gdn_chunk_solve32_vl");
10968            let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10969            let __s_lb = self.gpu.stream();
10970            let mut lb = __s_lb.launch_builder(&f);
10971            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
10972            unsafe { lb.launch(cfg)?; }
10973        }
10974        Ok(())
10975    }
10976
10977    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
10978    /// fused gate-prep, 5 launches for every sequence (per-element math identical
10979    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
10980    #[allow(clippy::too_many_arguments)]
10981    pub fn gdn_prep_vl8(&self, seqs: &[GdnPrepVl], conv_w: &CudaSlice<f32>,
10982                        dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10983                        conv_dim: usize, d_conv: usize, d_state: usize,
10984                        num_v: usize, num_k: usize, key_dim: usize, hk: usize, eps: f32)
10985                        -> Result<(), Box<dyn std::error::Error>> {
10986        let b = seqs.len();
10987        assert!(b >= 1 && b <= 8);
10988        let mut packed = [GdnPrepVl::default(); 8];
10989        packed[..b].copy_from_slice(seqs);
10990        let v = GdnPrepVl8(packed);
10991        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
10992        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
10993        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
10994        assert!(conv_fuse || hk == num_v, "de-broadcast requires the fused conv");
10995        if conv_fuse {
10996            let f = self.func("ssm_conv1d_gdn_state_vl");
10997            let cfg = LaunchConfig { grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10998            let (dsi, nvi, nki, kdi, hki) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, hk as i32);
10999            let __s_lb = self.gpu.stream();
11000            let mut lb = __s_lb.launch_builder(&f);
11001            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi).arg(&hki);
11002            unsafe { lb.launch(cfg)?; }
11003        } else {
11004            let f = self.func("ssm_conv1d_tm_state_vl");
11005            let cfg = LaunchConfig { grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11006            let __s_lb = self.gpu.stream();
11007            let mut lb = __s_lb.launch_builder(&f);
11008            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
11009            unsafe { lb.launch(cfg)?; }
11010        }
11011        {
11012            let f = self.func("ssm_conv_ring_update_vl");
11013            let n = (conv_dim * (d_conv - 1)) as u32;
11014            let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11015            let __s_lb = self.gpu.stream();
11016            let mut lb = __s_lb.launch_builder(&f);
11017            lb.arg(&v).arg(&cdi).arg(&dci);
11018            unsafe { lb.launch(cfg)?; }
11019        }
11020        if !conv_fuse {
11021            let f = self.func("qkv_to_gdn_repack_vl");
11022            let n = max_t * (num_v * d_state) as u32;
11023            let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11024            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
11025            let __s_lb = self.gpu.stream();
11026            let mut lb = __s_lb.launch_builder(&f);
11027            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
11028            unsafe { lb.launch(cfg)?; }
11029        }
11030        if Self::l2_v2_on(d_state) {
11031            let f = self.func("gdn_l2_v2_vl");
11032            let cfg = LaunchConfig { grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11033            let (dsi, nvi) = (d_state as i32, hk as i32);
11034            let __s_lb = self.gpu.stream();
11035            let mut lb = __s_lb.launch_builder(&f);
11036            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
11037            unsafe { lb.launch(cfg)?; }
11038        } else {
11039            let f = self.func("gdn_l2_vl");
11040            let cfg = LaunchConfig { grid_dim: (max_t * hk as u32, 2, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11041            let (dsi, nvi) = (d_state as i32, hk as i32);
11042            let __s_lb = self.gpu.stream();
11043            let mut lb = __s_lb.launch_builder(&f);
11044            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
11045            unsafe { lb.launch(cfg)?; }
11046        }
11047        {
11048            let f = self.func("gdn_gate_prep_vl");
11049            let n = max_t * num_v as u32;
11050            let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11051            let nvi = num_v as i32;
11052            let __s_lb = self.gpu.stream();
11053            let mut lb = __s_lb.launch_builder(&f);
11054            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
11055            unsafe { lb.launch(cfg)?; }
11056        }
11057        Ok(())
11058    }
11059
11060    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
11061    pub fn gdn_mirror_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, which: i32, hk: usize)
11062                          -> Result<(), Box<dyn std::error::Error>> {
11063        let b = seqs.len();
11064        assert!(b >= 1 && b <= 8);
11065        let mut packed = [GdnSeqVl::default(); 8];
11066        packed[..b].copy_from_slice(seqs);
11067        let v = GdnVl8(packed);
11068        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
11069        let max_n = seqs.iter().map(|s| if which == 0 { s.t as i64 * ept as i64 }
11070                                        else { s.nc as i64 * ept as i64 * 32 }).max().unwrap();
11071        let f = self.func("gdn_mirror_vl");
11072        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
11073        let cfg = LaunchConfig { grid_dim: (blocks, 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11074        let __s_lb = self.gpu.stream();
11075        let mut lb = __s_lb.launch_builder(&f);
11076        lb.arg(&v).arg(&ept).arg(&which);
11077        unsafe { lb.launch(cfg)?; }
11078        Ok(())
11079    }
11080
11081    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
11082    pub fn gdn_tail_vl8(&self, seqs: &[GdnPrepVl], norm_w: &CudaSlice<f32>,
11083                        d_state: usize, num_v: usize, eps: f32)
11084                        -> Result<(), Box<dyn std::error::Error>> {
11085        let b = seqs.len();
11086        assert!(b >= 1 && b <= 8);
11087        let mut packed = [GdnPrepVl::default(); 8];
11088        packed[..b].copy_from_slice(seqs);
11089        let v = GdnPrepVl8(packed);
11090        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
11091        let f = self.func("gated_rmsnorm_f16out_vl");
11092        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
11093        let cfg = LaunchConfig { grid_dim: (max_t * num_v as u32, 1, b as u32), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11094        let (dsi, nvi) = (d_state as i32, num_v as i32);
11095        let __s_lb = self.gpu.stream();
11096        let mut lb = __s_lb.launch_builder(&f);
11097        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
11098        unsafe { lb.launch(cfg)?; }
11099        Ok(())
11100    }
11101
11102    /// Raw device address helpers for the varlen by-value arg struct (single-stream
11103    /// launches; every buffer outlives the call — the f16 FFI discipline).
11104    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
11105        use cudarc::driver::DevicePtr;
11106        let s = self.gpu.stream();
11107        let (p, _g) = x.device_ptr(&s);
11108        p as u64
11109    }
11110    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
11111        use cudarc::driver::DevicePtrMut;
11112        let s = self.gpu.stream();
11113        let (p, _g) = x.device_ptr_mut(&s);
11114        p as u64
11115    }
11116    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
11117        use cudarc::driver::DevicePtr;
11118        let s = self.gpu.stream();
11119        let (p, _g) = x.device_ptr(&s);
11120        p as u64
11121    }
11122    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
11123        use cudarc::driver::DevicePtr;
11124        let s = self.gpu.stream();
11125        let (p, _g) = x.device_ptr(&s);
11126        p as u64
11127    }
11128
11129    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
11130    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
11131    /// launches, so this is strictly bit-gateable against them).
11132    pub fn gdn_chunk_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, scale: f32, hk: usize,
11133                         wq: Option<&GdnWVl8>)
11134                         -> Result<(), Box<dyn std::error::Error>> {
11135        const NSPLIT: u32 = 4;
11136        let b = seqs.len();
11137        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
11138        let mut packed = [GdnSeqVl::default(); 8];
11139        packed[..b].copy_from_slice(seqs);
11140        let v = GdnVl8(packed);
11141        let (hi, ci) = (n_head as i32, 32i32);
11142        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
11143        let hki = hk as i32;
11144        if let Some(w) = wq {
11145            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
11146            let f = self.func("gdn_k45_wgmma_vl");
11147            let cfg = LaunchConfig { grid_dim: (n_head as u32, NSPLIT, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11148            let __s_lb = self.gpu.stream();
11149            let mut lb = __s_lb.launch_builder(&f);
11150            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
11151            unsafe { lb.launch(cfg)?; }
11152            let _ = max_nc;
11153            return Ok(());
11154        }
11155        {
11156            let f = self.func("gdn_chunk_state_mma_vl");
11157            let cfg = LaunchConfig { grid_dim: (n_head as u32, NSPLIT, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11158            let __s_lb = self.gpu.stream();
11159            let mut lb = __s_lb.launch_builder(&f);
11160            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
11161            unsafe { lb.launch(cfg)?; }
11162        }
11163        {
11164            let f = self.func("gdn_chunk_output_mma_vl");
11165            let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11166            let __s_lb = self.gpu.stream();
11167            let mut lb = __s_lb.launch_builder(&f);
11168            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
11169            unsafe { lb.launch(cfg)?; }
11170        }
11171        Ok(())
11172    }
11173    pub fn gdn_scan_chunked(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11174                            g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>,
11175                            qb16_pre: Option<&CudaSlice<u8>>,
11176                            state_in: &CudaSlice<f32>,
11177                            state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11178                            n_head: usize, t: usize, scale: f32, c: usize, hk: usize)
11179                            -> Result<(), Box<dyn std::error::Error>> {
11180        const D: usize = 128;
11181        const NSPLIT: u32 = 4;
11182        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
11183        let h = n_head;
11184        let nc = (t + c - 1) / c;
11185        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
11186        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
11187        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
11188        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
11189        let gdn_mma_pre = !portable_mma_gated() && c == 32
11190            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
11191                Ok("1") => true,
11192                Ok("0") => false,
11193                _ => cfg!(memra_hopper_mma),
11194            };
11195        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
11196            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
11197        } else { None };
11198        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
11199        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
11200        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
11201        let gdn_wgmma_pre = gdn_mma_pre
11202            && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
11203                Ok("0") => false,
11204                Ok("1") => true,
11205                _ => cfg!(memra_hopper_mma),
11206            };
11207        let nk = t * hk * D;
11208        let mut kb16_local: Option<CudaSlice<u8>> = None;
11209        if gdn_mma_pre && kb16_pre.is_none() {
11210            let mut kb = self.alloc_u8_uninit(nk * 2)?;
11211            let f = self.func("f32_to_bf16_bulk");
11212            let n2 = nk as i64;
11213            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
11214            let __s_b = self.gpu.stream();
11215            let mut b = __s_b.launch_builder(&f);
11216            b.arg(k).arg(&mut kb).arg(&n2);
11217            unsafe { b.launch(cfg2)?; }
11218            kb16_local = Some(kb);
11219        }
11220        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
11221        if let Some(kb) = kb16_pre { assert!(kb.len() >= nk * 2, "kb16_pre too small"); }
11222        let mut qb16: Option<CudaSlice<u8>> = None;
11223        let mut pb16: Option<CudaSlice<u8>> = None;
11224        if gdn_wgmma_pre {
11225            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
11226            // the standalone bulk cvt only serves callers without the prep mirror.
11227            if qb16_pre.is_none() {
11228                let mut qb = self.alloc_u8_uninit(nk * 2)?;
11229                let f = self.func("f32_to_bf16_bulk");
11230                let n2 = nk as i64;
11231                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
11232                let __s_b = self.gpu.stream();
11233                let mut b = __s_b.launch_builder(&f);
11234                b.arg(q).arg(&mut qb).arg(&n2);
11235                unsafe { b.launch(cfg2)?; }
11236                qb16 = Some(qb);
11237            } else if let Some(qb) = qb16_pre {
11238                assert!(qb.len() >= nk * 2, "qb16_pre too small");
11239            }
11240            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
11241        }
11242        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
11243        let k2w = if gdn_wgmma_pre {
11244            Some((*qb16_ref0.as_ref().unwrap(),
11245                  *kb16_ref0.as_ref().unwrap(),
11246                  pb16.as_mut().unwrap()))
11247        } else { None };
11248        let (gcum, p, u, w) = self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
11249        let _ = &w;
11250        let mut y = self.uninit(nc * h * c * D)?;
11251        let mut ssnap = self.uninit(nc * h * D * D)?;   // chunk-start state snapshots (K5 phase 1)
11252        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
11253        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
11254        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
11255        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
11256        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
11257        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
11258        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
11259        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
11260        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
11261        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
11262        let gdn_mma = !portable_mma_gated() && c == 32
11263            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
11264                Ok("1") => true,
11265                Ok("0") => false,
11266                _ => cfg!(memra_hopper_mma),
11267            };
11268        if gdn_mma {
11269            let wb16 = wb16_pre.take().expect("mma path pre-allocates wb16 (K3 store fold)");
11270            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
11271            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
11272            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
11273            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
11274            // pass runs inside the persistent-M kernel; Y and Ssnap are never
11275            // materialized. New numeric class (gk folds into k^T instead of ys) —
11276            // explicit opt-in until the state-carry battery promotes it. Env read per
11277            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
11278            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
11279            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
11280            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
11281            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
11282            if gdn_wgmma_pre {
11283                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
11284                let qb16 = qb16_ref0.unwrap();
11285                let pb16 = pb16.as_ref().unwrap();
11286                {
11287                    let f = self.func("gdn_k45_wgmma");
11288                    let cfg = LaunchConfig { grid_dim: (h as u32, 4, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11289                    let hki = hk as i32;
11290                    let __s_b = self.gpu.stream();
11291                    let mut b = __s_b.launch_builder(&f);
11292                    b.arg(kb16_ref).arg(&gcum).arg(beta).arg(&u).arg(&wb16).arg(qb16).arg(pb16)
11293                     .arg(o).arg(&scale).arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
11294                    unsafe { b.launch(cfg)?; }
11295                }
11296                return Ok(());
11297            }
11298            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
11299            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
11300            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
11301            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
11302            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
11303            {
11304                let f = self.func("gdn_chunk_state_mma");
11305                let cfg = LaunchConfig { grid_dim: (h as u32, NSPLIT, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11306                let hki = hk as i32;
11307                let __s_b = self.gpu.stream();
11308                let mut b = __s_b.launch_builder(&f);
11309                b.arg(kb16_ref).arg(&gcum).arg(beta).arg(&u).arg(&wb16).arg(&mut y16).arg(&mut ssnap16)
11310                 .arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
11311                unsafe { b.launch(cfg)?; }
11312            }
11313            {   // K5-mma (bf16 St/Y consumers)
11314                let f = self.func("gdn_chunk_output_mma");
11315                let jt = ((c + 31) / 32) as u32;
11316                let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11317                let hki = hk as i32;
11318                let __s_b = self.gpu.stream();
11319                let mut b = __s_b.launch_builder(&f);
11320                b.arg(q).arg(&gcum).arg(&p).arg(&y16).arg(&ssnap16).arg(o).arg(&hi).arg(&ti).arg(&ci).arg(&scale).arg(&hki);
11321                unsafe { b.launch(cfg)?; }
11322            }
11323            return Ok(());
11324        }
11325        {   // K4 (sequential over chunks inside; blocks col-partition the state)
11326            let f = self.func("gdn_chunk_state_f32");
11327            let cfg = LaunchConfig { grid_dim: (h as u32, NSPLIT, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11328            let __s_b = self.gpu.stream();
11329            let mut b = __s_b.launch_builder(&f);
11330            b.arg(k).arg(&gcum).arg(beta).arg(&u).arg(&w).arg(&mut y).arg(&mut ssnap)
11331             .arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci);
11332            unsafe { b.launch(cfg)?; }
11333        }
11334        {   // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
11335            let f = self.func("gdn_chunk_output_f32");
11336            let jt = ((c + 31) / 32) as u32;
11337            let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11338            let __s_b = self.gpu.stream();
11339            let mut b = __s_b.launch_builder(&f);
11340            b.arg(q).arg(&gcum).arg(&p).arg(&y).arg(&ssnap).arg(o).arg(&hi).arg(&ti).arg(&ci).arg(&scale);
11341            unsafe { b.launch(cfg)?; }
11342        }
11343        Ok(())
11344    }
11345
11346    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
11347    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
11348    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
11349    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
11350    ///
11351    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
11352    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
11353    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
11354    #[allow(clippy::too_many_arguments)]
11355    #[allow(clippy::too_many_arguments)]
11356    pub fn gdn_scan_prefill(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11357                            g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>,
11358                            qb16_pre: Option<&CudaSlice<u8>>,
11359                            state_in: &CudaSlice<f32>,
11360                            state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11361                            n_head: usize, t: usize, scale: f32, hk: usize)
11362                            -> Result<(), Box<dyn std::error::Error>> {
11363        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
11364            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
11365            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
11366        }
11367        if Self::gdn_chunked_enabled() && t >= 16 {
11368            self.gdn_scan_chunked(q, k, v, g, beta, kb16_pre, qb16_pre, state_in, state_out, o, n_head, t, scale,
11369                                  Self::gdn_chunk_size(), hk)
11370        } else {
11371            assert!(hk == n_head, "s128 scan is broadcast-only (prep guarantees by predicate)");
11372            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
11373        }
11374    }
11375
11376    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
11377    #[allow(clippy::too_many_arguments)]
11378    fn gdn_scan_diff(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11379                     g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
11380                     state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11381                     n_head: usize, t: usize, scale: f32)
11382                     -> Result<(), Box<dyn std::error::Error>> {
11383        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
11384        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11385        let mut o_c = self.uninit(o.len())?;
11386        let mut st_c = self.uninit(state_out.len())?;
11387        self.gdn_scan_chunked(q, k, v, g, beta, None, None, state_in, &mut st_c, &mut o_c,
11388                              n_head, t, scale, Self::gdn_chunk_size(), n_head)?;
11389        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
11390        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
11391        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
11392        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
11393            let mut max_abs = 0f32; let mut max_rel = 0f32; let mut sum_rel = 0f64;
11394            for (x, y) in a.iter().zip(b) {
11395                let ad = (x - y).abs();
11396                let rel = ad / x.abs().max(y.abs()).max(1e-3);
11397                if ad > max_abs { max_abs = ad; }
11398                if rel > max_rel { max_rel = rel; }
11399                sum_rel += rel as f64;
11400            }
11401            (max_abs, max_rel, sum_rel / a.len() as f64)
11402        };
11403        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
11404        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
11405        println!("[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
11406                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
11407                 Self::gdn_chunk_size());
11408        Ok(())
11409    }
11410
11411    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
11412    pub fn gdn_glog(&self, alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
11413                    g_log: &mut CudaSlice<f32>, n_head: usize, t: usize)
11414                    -> Result<(), Box<dyn std::error::Error>> {
11415        let f = self.func("gdn_glog_f32");
11416        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
11417        let (h, ti) = (n_head as i32, t as i32);
11418        let __s_b = self.gpu.stream();
11419        let mut b = __s_b.launch_builder(&f);
11420        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
11421        unsafe { b.launch(cfg)?; }
11422        Ok(())
11423    }
11424
11425    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
11426    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
11427    pub fn sigmoid_v(&self, x: &cudarc::driver::CudaView<f32>, y: &mut CudaSlice<f32>, n: usize)
11428                     -> Result<(), Box<dyn std::error::Error>> {
11429        let f = self.func("sigmoid_f32");
11430        let cfg = LaunchConfig::for_num_elems(n as u32);
11431        let ni = n as i32;
11432        let __s_b = self.gpu.stream();
11433        let mut b = __s_b.launch_builder(&f);
11434        b.arg(x).arg(y).arg(&ni);
11435        unsafe { b.launch(cfg)?; }
11436        Ok(())
11437    }
11438
11439    pub fn gdn_glog_v(&self, alpha: &cudarc::driver::CudaView<f32>, dt_bias: &CudaSlice<f32>,
11440                      a: &CudaSlice<f32>, g_log: &mut CudaSlice<f32>, n_head: usize, t: usize)
11441                      -> Result<(), Box<dyn std::error::Error>> {
11442        let f = self.func("gdn_glog_f32");
11443        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
11444        let (h, ti) = (n_head as i32, t as i32);
11445        let __s_b = self.gpu.stream();
11446        let mut b = __s_b.launch_builder(&f);
11447        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
11448        unsafe { b.launch(cfg)?; }
11449        Ok(())
11450    }
11451
11452    pub fn sigmoid(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize)
11453                   -> Result<(), Box<dyn std::error::Error>> {
11454        let f = self.func("sigmoid_f32");
11455        let cfg = LaunchConfig::for_num_elems(n as u32);
11456        let ni = n as i32;
11457        let __s_b = self.gpu.stream();
11458        let mut b = __s_b.launch_builder(&f);
11459        b.arg(x).arg(y).arg(&ni);
11460        unsafe { b.launch(cfg)?; }
11461        Ok(())
11462    }
11463
11464    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
11465    /// (replaces sigmoid + mul + convert). Bit-identical class.
11466    pub fn sig_mul_f16out(&self, a: &CudaSlice<f32>, g: &CudaSlice<f32>,
11467                          dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize)
11468                          -> Result<(), Box<dyn std::error::Error>> {
11469        let f = self.func("sig_mul_f16out_f32");
11470        let cfg = LaunchConfig::for_num_elems(n as u32);
11471        let ni = n as i32;
11472        let __s_b = self.gpu.stream();
11473        let mut b = __s_b.launch_builder(&f);
11474        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
11475        unsafe { b.launch(cfg)?; }
11476        Ok(())
11477    }
11478
11479    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
11480    pub fn gated_rmsnorm(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11481                         dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
11482                         -> Result<(), Box<dyn std::error::Error>> {
11483        let f = self.func("gated_rmsnorm_f32");
11484        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11485        let (nc, e) = (ncols as i32, eps);
11486        let __s_b = self.gpu.stream();
11487        let mut b = __s_b.launch_builder(&f);
11488        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
11489        unsafe { b.launch(cfg)?; }
11490        Ok(())
11491    }
11492
11493    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
11494    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
11495    pub fn gated_rmsnorm_f16out(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11496                                dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
11497                                ncols: usize, nrows: usize, eps: f32)
11498                                -> Result<(), Box<dyn std::error::Error>> {
11499        let f = self.func("gated_rmsnorm_f16out_f32");
11500        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
11501        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11502        let (nc, e) = (ncols as i32, eps);
11503        let __s_b = self.gpu.stream();
11504        let mut b = __s_b.launch_builder(&f);
11505        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
11506        unsafe { b.launch(cfg)?; }
11507        Ok(())
11508    }
11509
11510    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
11511    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
11512    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
11513    #[allow(clippy::too_many_arguments)]
11514    pub fn add_rms_norm_zq8(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>,
11515                            res: &mut CudaSlice<f32>, z: &mut CudaSlice<f32>,
11516                            ncols: usize, nrows: usize, eps: f32)
11517                            -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11518        assert!(ncols % 32 == 0);
11519        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
11520        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11521        let f = self.func("add_rms_norm_zq8");
11522        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
11523        let (nc, ep) = (ncols as i32, eps);
11524        let __s_b = self.gpu.stream();
11525        let mut b = __s_b.launch_builder(&f);
11526        b.arg(a).arg(b_in).arg(w).arg(res).arg(z).arg(&mut q).arg(&mut d).arg(&nc).arg(&ep);
11527        unsafe { b.launch(cfg)?; }
11528        Ok((q, d))
11529    }
11530
11531    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
11532    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
11533    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
11534    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
11535    pub fn gated_rmsnorm_zv(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>,
11536                            z: &cudarc::driver::CudaView<f32>,
11537                            dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
11538                            -> Result<(), Box<dyn std::error::Error>> {
11539        let f = self.func("gated_rmsnorm_f32");
11540        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11541        let (nc, e) = (ncols as i32, eps);
11542        let __s_b = self.gpu.stream();
11543        let mut b = __s_b.launch_builder(&f);
11544        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
11545        unsafe { b.launch(cfg)?; }
11546        Ok(())
11547    }
11548
11549    pub fn gated_rmsnorm_f16out_zv(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>,
11550                                   z: &cudarc::driver::CudaView<f32>,
11551                                   dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
11552                                   ncols: usize, nrows: usize, eps: f32)
11553                                   -> Result<(), Box<dyn std::error::Error>> {
11554        let f = self.func("gated_rmsnorm_f16out_f32");
11555        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
11556        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11557        let (nc, e) = (ncols as i32, eps);
11558        let __s_b = self.gpu.stream();
11559        let mut b = __s_b.launch_builder(&f);
11560        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
11561        unsafe { b.launch(cfg)?; }
11562        Ok(())
11563    }
11564
11565    pub fn gated_rmsnorm_q8_1(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11566                              ncols: usize, nrows: usize, eps: f32)
11567                              -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11568        assert!(ncols % 32 == 0);
11569        let f = self.func("gated_rmsnorm_q8_1");
11570        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11571        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11572        let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11573        let (nc, ep) = (ncols as i32, eps);
11574        let __s_b = self.gpu.stream();
11575        let mut b = __s_b.launch_builder(&f);
11576        b.arg(o).arg(w).arg(z).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&ep);
11577        unsafe { b.launch(cfg)?; }
11578        Ok((out_q, out_d))
11579    }
11580
11581    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
11582    pub fn transpose(&self, inp: &CudaSlice<f32>, rows: usize, cols: usize)
11583                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11584        let f = self.func("transpose_f32");
11585        let mut out = self.zeros(rows * cols)?;
11586        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
11587        let (r, c) = (rows as i32, cols as i32);
11588        let __s_b = self.gpu.stream();
11589        let mut b = __s_b.launch_builder(&f);
11590        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
11591        unsafe { b.launch(cfg)?; }
11592        Ok(out)
11593    }
11594
11595    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
11596    pub fn repeat_heads(&self, inp: &CudaSlice<f32>, out: &mut CudaSlice<f32>,
11597                        head_dim: usize, n_in: usize, n_out: usize, t: usize)
11598                        -> Result<(), Box<dyn std::error::Error>> {
11599        let f = self.func("repeat_heads_f32");
11600        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
11601        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
11602        let __s_b = self.gpu.stream();
11603        let mut b = __s_b.launch_builder(&f);
11604        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
11605        unsafe { b.launch(cfg)?; }
11606        Ok(())
11607    }
11608
11609    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
11610    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
11611    pub fn q_gate_split(&self, qf: &CudaSlice<f32>, q_out: &mut CudaSlice<f32>,
11612                        gate_out: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, t: usize)
11613                        -> Result<(), Box<dyn std::error::Error>> {
11614        let f = self.func("q_gate_split_f32");
11615        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
11616        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
11617        let __s_b = self.gpu.stream();
11618        let mut b = __s_b.launch_builder(&f);
11619        b.arg(qf).arg(q_out).arg(gate_out).arg(&hd).arg(&nh).arg(&ti);
11620        unsafe { b.launch(cfg)?; }
11621        Ok(())
11622    }
11623
11624    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
11625    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
11626    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
11627    pub fn qkv_to_gdn_repack(&self, conv_out: &CudaSlice<f32>, q_g: &mut CudaSlice<f32>,
11628                             k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
11629                             d_state: usize, num_v: usize, num_k: usize, key_dim: usize, t: usize)
11630                             -> Result<(), Box<dyn std::error::Error>> {
11631        let f = self.func("qkv_to_gdn_repack_f32");
11632        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
11633        let (ds, nv, nk, kd, ti) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, t as i32);
11634        let __s_b = self.gpu.stream();
11635        let mut b = __s_b.launch_builder(&f);
11636        b.arg(conv_out).arg(q_g).arg(k_g).arg(v_g).arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&ti);
11637        unsafe { b.launch(cfg)?; }
11638        Ok(())
11639    }
11640
11641    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
11642    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
11643    pub fn conv_left_pad(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
11644                         conv_dim: usize, t: usize, pad: usize)
11645                         -> Result<(), Box<dyn std::error::Error>> {
11646        let f = self.func("conv_left_pad_f32");
11647        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
11648        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
11649        let __s_b = self.gpu.stream();
11650        let mut b = __s_b.launch_builder(&f);
11651        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
11652        unsafe { b.launch(cfg)?; }
11653        Ok(())
11654    }
11655
11656    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
11657    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
11658    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
11659    pub fn conv_assemble_and_roll(&self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
11660                                  conv_in: &mut CudaSlice<f32>, conv_dim: usize, pad: usize)
11661                                  -> Result<(), Box<dyn std::error::Error>> {
11662        let f = self.func("conv_assemble_and_roll_f32");
11663        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
11664        let (cd, p) = (conv_dim as i32, pad as i32);
11665        let __s_b = self.gpu.stream();
11666        let mut b = __s_b.launch_builder(&f);
11667        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
11668        unsafe { b.launch(cfg)?; }
11669        Ok(())
11670    }
11671
11672    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
11673    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
11674    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
11675    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
11676    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
11677    pub fn ssm_conv1d_fused_decode(&self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
11678                                   w: &CudaSlice<f32>, conv_out: &mut CudaSlice<f32>,
11679                                   conv_dim: usize, d_conv: usize)
11680                                   -> Result<(), Box<dyn std::error::Error>> {
11681        let f = self.func("ssm_conv1d_fused_decode_f32");
11682        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
11683        let (cd, dc) = (conv_dim as i32, d_conv as i32);
11684        let __s_b = self.gpu.stream();
11685        let mut b = __s_b.launch_builder(&f);
11686        b.arg(qkv_col).arg(conv_state).arg(w).arg(conv_out).arg(&cd).arg(&dc);
11687        unsafe { b.launch(cfg)?; }
11688        Ok(())
11689    }
11690
11691    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
11692    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
11693    pub fn slice_range(&self, src: &CudaSlice<f32>, start: usize, len: usize)
11694                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11695        let host = self.gpu.stream().clone_dtoh(src)?;
11696        self.gpu.stream().synchronize()?;
11697        Ok(self.htod(&host[start..start + len])?)
11698    }
11699}
11700
11701#[cfg(test)]
11702mod target_dispatch_tests {
11703    use super::legacy_quant_gemm_allowed;
11704
11705    #[test]
11706    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
11707        // sm_120a native lane
11708        assert!(legacy_quant_gemm_allowed(false, false, false));
11709        assert!(!legacy_quant_gemm_allowed(false, false, true));
11710        // pure portable lane (sm_89): gated
11711        assert!(!legacy_quant_gemm_allowed(true, false, false));
11712        assert!(!legacy_quant_gemm_allowed(true, false, true));
11713        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
11714        assert!(legacy_quant_gemm_allowed(true, true, false));
11715        assert!(!legacy_quant_gemm_allowed(true, true, true));
11716    }
11717
11718    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
11719    #[test]
11720    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
11721        assert!(!legacy_quant_gemm_allowed(cfg!(memra_portable_cuda), cfg!(memra_hopper_mma), false));
11722    }
11723
11724    #[cfg(memra_hopper_mma)]
11725    #[test]
11726    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
11727        assert!(legacy_quant_gemm_allowed(cfg!(memra_portable_cuda), cfg!(memra_hopper_mma), false));
11728        assert!(super::portable_mma_gated() == false);
11729    }
11730}
11731
11732/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
11733/// inherent methods (inherent methods win name resolution, so no recursion).
11734impl memra_kv::KvDev for Engine {
11735    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11736        Engine::zeros(self, n)
11737    }
11738    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11739        Engine::uninit(self, n)
11740    }
11741    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
11742        Engine::alloc_u8(self, n)
11743    }
11744    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
11745        Engine::htod_i32(self, v)
11746    }
11747    fn clone_dtod(&self, src: &CudaSlice<f32>) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11748        Engine::clone_dtod(self, src)
11749    }
11750    fn copy_into(&self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize)
11751                 -> Result<(), Box<dyn std::error::Error>> {
11752        Engine::copy_into(self, dst, off, src, len)
11753    }
11754    fn set_i32_one(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
11755        Engine::set_i32_one(self, d, v)
11756    }
11757}